diff --git a/python/cuda/bench/__init__.pyi b/python/cuda/bench/__init__.pyi index 5be55327..0e1ad7d1 100644 --- a/python/cuda/bench/__init__.pyi +++ b/python/cuda/bench/__init__.pyi @@ -31,6 +31,7 @@ from typing import ( Any, Literal, Optional, + Protocol, Self, SupportsFloat, SupportsInt, @@ -41,6 +42,11 @@ from typing import ( _F = TypeVar("_F", bound=Callable[..., Any]) +# cuda.bench.CudaStream structurally satisfies this protocol, but State.set_stream +# rejects it at runtime to avoid rebinding a state to its own stream. +class _SupportsCudaStream(Protocol): + def __cuda_stream__(self) -> tuple[int, int]: ... + class CudaStream: def __cuda_stream__(self) -> tuple[int, int]: ... def addressof(self) -> int: ... @@ -80,6 +86,8 @@ class State: def has_printers(self) -> bool: ... def get_device(self) -> Union[int, None]: ... def get_stream(self) -> CudaStream: ... + # Runtime rejects cuda.bench.CudaStream even though it matches _SupportsCudaStream. + def set_stream(self, stream_provider: _SupportsCudaStream) -> None: ... def get_int64(self, name: str) -> int: ... def get_int64_or_default(self, name: str, default_value: SupportsInt) -> int: ... def get_float64(self, name: str) -> float: ... diff --git a/python/src/py_nvbench.cpp b/python/src/py_nvbench.cpp index 82894e1e..2f307977 100644 --- a/python/src/py_nvbench.cpp +++ b/python/src/py_nvbench.cpp @@ -24,12 +24,14 @@ #include +#include #include #include #include #include #include #include +#include #include #include @@ -38,6 +40,13 @@ namespace py = pybind11; namespace { +std::unordered_map &cuda_stream_provider_cache() +{ + // Avoid destroying py::object values after Python interpreter shutdown. + static auto *cache = new std::unordered_map{}; + return *cache; +} + struct PyObjectDeleter { void operator()(py::object *p) @@ -87,6 +96,22 @@ struct benchmark_wrapper_t { throw std::runtime_error("No function to execute"); } + + struct cleanup_cuda_stream_provider + { + nvbench::state *state; + + ~cleanup_cuda_stream_provider() noexcept + { + try + { + cuda_stream_provider_cache().erase(state); + } + catch (...) + {} + } + } cleanup{&state}; + // box as Python object, using reference semantics auto arg = py::cast(std::ref(state), py::return_value_policy::reference); @@ -332,6 +357,59 @@ py::dict py_get_axis_values(const nvbench::state &state) // essentially a global variable, but allocated on the heap during module initialization std::unique_ptr global_registry{}; +cudaStream_t extract_cuda_stream_from_provider(const py::handle &stream_provider) +{ + if (py::isinstance(stream_provider)) + { + throw py::type_error("State.set_stream does not accept cuda.bench.CudaStream instances"); + } + + if (!py::hasattr(stream_provider, "__cuda_stream__")) + { + throw py::type_error("State.set_stream expects an object implementing __cuda_stream__"); + } + + const py::object protocol_method = stream_provider.attr("__cuda_stream__"); + if (!PyCallable_Check(protocol_method.ptr())) + { + throw py::type_error("State.set_stream expects __cuda_stream__ to be callable"); + } + + const py::object protocol_result = protocol_method(); + if (!py::isinstance(protocol_result)) + { + throw py::type_error("State.set_stream expects __cuda_stream__ to return " + "(protocol_version, cuda_stream_handle)"); + } + + const auto stream_info = py::reinterpret_borrow(protocol_result); + if (stream_info.size() != 2) + { + throw py::type_error("State.set_stream expects __cuda_stream__ to return " + "(protocol_version, cuda_stream_handle)"); + } + + int protocol_version{}; + std::uintptr_t stream_handle{}; + try + { + protocol_version = stream_info[0].cast(); + stream_handle = stream_info[1].cast(); + } + catch (const py::cast_error &) + { + throw py::type_error("State.set_stream expects __cuda_stream__ to return " + "(protocol_version, cuda_stream_handle) integers"); + } + + if (protocol_version != 0) + { + throw py::value_error("State.set_stream only supports CUDA stream protocol version 0"); + } + + return reinterpret_cast(stream_handle); +} + // Definitions of Python API static void def_class_CudaStream(py::module_ m) { @@ -890,6 +968,34 @@ Get `CudaStream` object from this configuration method_get_stream_doc, py::return_value_policy::reference); + // method State.set_stream + auto method_set_stream_impl = [](nvbench::state &state, py::handle stream_provider) { + const auto stream_handle = extract_cuda_stream_from_provider(stream_provider); + + const auto ¤t_stream = state.get_cuda_stream_optional(); + if (!current_stream.has_value() || current_stream->get_stream() != stream_handle) + { + cuda_stream_provider_cache()[&state] = py::reinterpret_borrow(stream_provider); + state.set_cuda_stream(nvbench::make_cuda_stream_view(stream_handle)); + } + }; + static constexpr const char *method_set_stream_doc = R"XXXX( +Set this configuration's CUDA stream from an object implementing __cuda_stream__. + +The stream provider is expected to keep the returned stream valid. NVBench +stores a non-owning view and keeps the current provider object alive until +another provider replaces it or this benchmark callback returns. + +Call this method before State.exec, not from launcher callbacks. The returned +stream must be valid for this state configuration's CUDA device. + +cuda.bench.CudaStream instances are not accepted. +)XXXX"; + pystate_cls.def("set_stream", + method_set_stream_impl, + method_set_stream_doc, + py::arg("stream_provider")); + // method State.get_int64 auto method_get_int64_impl = &nvbench::state::get_int64; static constexpr const char *method_get_int64_doc = R"XXXX( @@ -1235,6 +1341,7 @@ Use argument True to disable use of blocking kernel by NVBench" The callable may be executed multiple times. The callable will be passed a `Launch` object argument by default. When `timer=True`, the callable will be passed `Launch` and `Timer` arguments. + The callable is invoked with the Python GIL held. Parameters ---------- @@ -1415,6 +1522,10 @@ Register benchmark function of type Callable[[nvbench.State], None] static constexpr const char *func_run_all_benchmarks_doc = R"XXXX( Run all benchmarks registered with NVBench. + This function currently runs with the Python GIL held. Benchmark launcher + callbacks are invoked with the GIL held; native or extension functions + called by a launcher may release the GIL according to their own behavior. + Parameters ---------- argv: List[str] diff --git a/python/test/test_cuda_bench.py b/python/test/test_cuda_bench.py index de8a1dfe..64300c38 100644 --- a/python/test/test_cuda_bench.py +++ b/python/test/test_cuda_bench.py @@ -14,7 +14,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gc import json +import weakref from typing import Union import cuda.bench as bench @@ -42,6 +44,17 @@ def test_api_ctor(cls): def test_cpu_only(): saved_timers = [] observed = {} + stream_provider_refs = {} + external_stream_handle = 0x1234 + + class ExternalStreamProvider: + def __init__(self, handle): + self.handle = handle + self.protocol_calls = 0 + + def __cuda_stream__(self): + self.protocol_calls += 1 + return (0, self.handle) @bench.register() @bench.option.set_is_cpu_only(True) @@ -86,17 +99,83 @@ def cold_warmup_state_probe(state: bench.State): state.exec(lambda launch: None) + @bench.register() + @bench.option.set_is_cpu_only(True) + def external_stream_state_probe(state: bench.State): + stream_provider = ExternalStreamProvider(external_stream_handle) + stream_provider_refs["current"] = weakref.ref(stream_provider) + assert state.set_stream(stream_provider) is None + assert stream_provider.protocol_calls == 1 + + state.set_stream(stream_provider) + assert stream_provider.protocol_calls == 2 + del stream_provider + gc.collect() + assert stream_provider_refs["current"]() is not None + + same_handle_stream_provider = ExternalStreamProvider(external_stream_handle) + same_handle_stream_provider_ref = weakref.ref(same_handle_stream_provider) + state.set_stream(same_handle_stream_provider) + assert same_handle_stream_provider.protocol_calls == 1 + del same_handle_stream_provider + gc.collect() + assert same_handle_stream_provider_ref() is None + assert stream_provider_refs["current"]() is not None + + stream_provider = stream_provider_refs["current"]() + assert stream_provider is not None + + replaced_stream_provider = ExternalStreamProvider(0x5678) + replaced_stream_provider_ref = weakref.ref(replaced_stream_provider) + state.set_stream(replaced_stream_provider) + state.set_stream(stream_provider) + del replaced_stream_provider + gc.collect() + assert replaced_stream_provider_ref() is None + + class NonCallableProtocol: + __cuda_stream__ = 1 + + class BadProtocolReturn: + def __cuda_stream__(self): + return (0,) + + class UnsupportedProtocolVersion: + def __cuda_stream__(self): + return (1, external_stream_handle) + + with pytest.raises(TypeError, match="__cuda_stream__"): + state.set_stream(object()) + with pytest.raises(TypeError, match="callable"): + state.set_stream(NonCallableProtocol()) + with pytest.raises(TypeError, match="protocol_version"): + state.set_stream(BadProtocolReturn()) + with pytest.raises(ValueError, match="version 0"): + state.set_stream(UnsupportedProtocolVersion()) + with pytest.raises(TypeError, match="CudaStream"): + state.set_stream(state.get_stream()) + + state.exec( + lambda launch: observed.update( + {"external_stream_handle": launch.get_stream().addressof()} + ) + ) + bench.run_all_benchmarks(["-q", "--profile"]) assert saved_timers with pytest.raises(RuntimeError, match="Timer is no longer valid"): saved_timers[0].start() + gc.collect() + assert stream_provider_refs["current"]() is None + assert observed == { "benchmark_runs": 13, "benchmark_walltime": 0.5, "state_runs": 3, "state_walltime": 0.125, + "external_stream_handle": external_stream_handle, } @@ -307,6 +386,7 @@ def test_State_doc(): cl = bench.State obj_has_docstring_check(cl) obj_has_docstring_check(cl.exec) + obj_has_docstring_check(cl.set_stream) obj_has_docstring_check(cl.get_int64) obj_has_docstring_check(cl.get_float64) obj_has_docstring_check(cl.get_string)