From 20f25d151fb4c3fca03d34f7adec26a3ec4f378a Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:53:31 -0500 Subject: [PATCH 1/6] Implement cuda.bench.State.set_stream(provider) The stream provide type is expected to implement __cuda_stream__ protocol, and not be an instance of cuda.bench.CudaStream. The reason is that cuda.bench.CudaStream type is not user-constructible. It is returns by cuda.bench.State.get_stream method. It stores reference to stream stored in nvbench::state referenced from cuda.bench.State. Disallowing cuda.bench.CudaStream as argument to set_stream avoid footguns. --- python/cuda/bench/__init__.pyi | 5 +++ python/src/py_nvbench.cpp | 77 ++++++++++++++++++++++++++++++++++ python/test/test_cuda_bench.py | 50 ++++++++++++++++++++++ 3 files changed, 132 insertions(+) diff --git a/python/cuda/bench/__init__.pyi b/python/cuda/bench/__init__.pyi index 5be55327..ba134edc 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,9 @@ from typing import ( _F = TypeVar("_F", bound=Callable[..., Any]) +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 +84,7 @@ class State: def has_printers(self) -> bool: ... def get_device(self) -> Union[int, None]: ... def get_stream(self) -> CudaStream: ... + 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..b432ea67 100644 --- a/python/src/py_nvbench.cpp +++ b/python/src/py_nvbench.cpp @@ -24,6 +24,7 @@ #include +#include #include #include #include @@ -332,6 +333,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 +944,29 @@ 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) + { + 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 owns the stream. NVBench stores a non-owning view and keeps +the provider object alive while this State wrapper is alive. + +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"), + py::keep_alive<1, 2>()); + // method State.get_int64 auto method_get_int64_impl = &nvbench::state::get_int64; static constexpr const char *method_get_int64_doc = R"XXXX( diff --git a/python/test/test_cuda_bench.py b/python/test/test_cuda_bench.py index de8a1dfe..67ba22fd 100644 --- a/python/test/test_cuda_bench.py +++ b/python/test/test_cuda_bench.py @@ -42,6 +42,16 @@ def test_api_ctor(cls): def test_cpu_only(): saved_timers = [] observed = {} + 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,6 +96,44 @@ 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) + 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 + + 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 @@ -97,6 +145,7 @@ def cold_warmup_state_probe(state: bench.State): "benchmark_walltime": 0.5, "state_runs": 3, "state_walltime": 0.125, + "external_stream_handle": external_stream_handle, } @@ -307,6 +356,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) From fbfc158911a32319962c4a9d9f60c81e0fa08469 Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com> Date: Fri, 31 Jul 2026 16:37:28 -0500 Subject: [PATCH 2/6] Store only current Python stream provider Replace pybind11 keep_alive with an explicit per-state provider cache so repeated State.set_stream calls release previously installed providers instead of retaining every historical stream provider. --- python/cuda/bench/__init__.pyi | 3 +++ python/src/py_nvbench.cpp | 17 +++++++++++++---- python/test/test_cuda_bench.py | 10 ++++++++++ 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/python/cuda/bench/__init__.pyi b/python/cuda/bench/__init__.pyi index ba134edc..c8ede9ef 100644 --- a/python/cuda/bench/__init__.pyi +++ b/python/cuda/bench/__init__.pyi @@ -42,6 +42,8 @@ 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]: ... @@ -84,6 +86,7 @@ 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: ... diff --git a/python/src/py_nvbench.cpp b/python/src/py_nvbench.cpp index b432ea67..8baac792 100644 --- a/python/src/py_nvbench.cpp +++ b/python/src/py_nvbench.cpp @@ -31,6 +31,7 @@ #include #include #include +#include #include #include @@ -333,6 +334,13 @@ 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{}; +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; +} + cudaStream_t extract_cuda_stream_from_provider(const py::handle &stream_provider) { if (py::isinstance(stream_provider)) @@ -946,7 +954,9 @@ Get `CudaStream` object from this configuration // 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 stream_handle = extract_cuda_stream_from_provider(stream_provider); + cuda_stream_provider_cache()[&state] = py::reinterpret_borrow(stream_provider); + const auto ¤t_stream = state.get_cuda_stream_optional(); if (!current_stream.has_value() || current_stream->get_stream() != stream_handle) { @@ -957,15 +967,14 @@ Get `CudaStream` object from this configuration Set this configuration's CUDA stream from an object implementing __cuda_stream__. The stream provider owns the stream. NVBench stores a non-owning view and keeps -the provider object alive while this State wrapper is alive. +the current provider object alive until another provider replaces it. 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"), - py::keep_alive<1, 2>()); + py::arg("stream_provider")); // method State.get_int64 auto method_get_int64_impl = &nvbench::state::get_int64; diff --git a/python/test/test_cuda_bench.py b/python/test/test_cuda_bench.py index 67ba22fd..80f77fe3 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 @@ -106,6 +108,14 @@ def external_stream_state_probe(state: bench.State): state.set_stream(stream_provider) assert stream_provider.protocol_calls == 2 + 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 From 34cc57dc1d63162a16646611fe271d6b3134cfee Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com> Date: Sun, 2 Aug 2026 05:23:06 -0500 Subject: [PATCH 3/6] Clean up stream-provider cache in the benchmark_wrapper_t call operator --- python/src/py_nvbench.cpp | 33 +++++++++++++++++++++++++-------- python/test/test_cuda_bench.py | 5 +++++ 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/python/src/py_nvbench.cpp b/python/src/py_nvbench.cpp index 8baac792..3ed085e5 100644 --- a/python/src/py_nvbench.cpp +++ b/python/src/py_nvbench.cpp @@ -40,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) @@ -89,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); @@ -334,13 +357,6 @@ 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{}; -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; -} - cudaStream_t extract_cuda_stream_from_provider(const py::handle &stream_provider) { if (py::isinstance(stream_provider)) @@ -967,7 +983,8 @@ Get `CudaStream` object from this configuration Set this configuration's CUDA stream from an object implementing __cuda_stream__. The stream provider owns the stream. NVBench stores a non-owning view and keeps -the current provider object alive until another provider replaces it. +the current provider object alive until another provider replaces it or this +benchmark callback returns. cuda.bench.CudaStream instances are not accepted. )XXXX"; diff --git a/python/test/test_cuda_bench.py b/python/test/test_cuda_bench.py index 80f77fe3..c91c76e4 100644 --- a/python/test/test_cuda_bench.py +++ b/python/test/test_cuda_bench.py @@ -44,6 +44,7 @@ def test_api_ctor(cls): def test_cpu_only(): saved_timers = [] observed = {} + stream_provider_refs = {} external_stream_handle = 0x1234 class ExternalStreamProvider: @@ -102,6 +103,7 @@ def cold_warmup_state_probe(state: bench.State): @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 @@ -150,6 +152,9 @@ def __cuda_stream__(self): 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, From bb96a909cc7e7976e8449f346b1bb4f4634a526d Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com> Date: Sun, 2 Aug 2026 06:08:48 -0500 Subject: [PATCH 4/6] Only save stream-provider to cache if it is being set Add test for repeated set_stream of the same external owning stream Python object, with expectation that it does not double-free. --- python/src/py_nvbench.cpp | 4 ++-- python/test/test_cuda_bench.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/python/src/py_nvbench.cpp b/python/src/py_nvbench.cpp index 3ed085e5..0d58c96d 100644 --- a/python/src/py_nvbench.cpp +++ b/python/src/py_nvbench.cpp @@ -970,12 +970,12 @@ Get `CudaStream` object from this configuration // 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); - cuda_stream_provider_cache()[&state] = py::reinterpret_borrow(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)); } }; diff --git a/python/test/test_cuda_bench.py b/python/test/test_cuda_bench.py index c91c76e4..64300c38 100644 --- a/python/test/test_cuda_bench.py +++ b/python/test/test_cuda_bench.py @@ -109,6 +109,21 @@ def external_stream_state_probe(state: bench.State): 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) From a3fb83fcf0a5bd0a4e21f056ce6c6bf8d641d246 Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com> Date: Sun, 2 Aug 2026 06:13:21 -0500 Subject: [PATCH 5/6] Make SupportsCudaStream protocol private _SupportsCudaStream --- python/cuda/bench/__init__.pyi | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/cuda/bench/__init__.pyi b/python/cuda/bench/__init__.pyi index c8ede9ef..0e1ad7d1 100644 --- a/python/cuda/bench/__init__.pyi +++ b/python/cuda/bench/__init__.pyi @@ -44,7 +44,7 @@ _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): +class _SupportsCudaStream(Protocol): def __cuda_stream__(self) -> tuple[int, int]: ... class CudaStream: @@ -86,8 +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: ... + # 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: ... From 4287b99b87d3d426ed802e2c1b4d1fd56f8c322f Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com> Date: Sun, 2 Aug 2026 06:46:21 -0500 Subject: [PATCH 6/6] A pass at docstrings Docstring should not require that external stream provider must own stream, only that it keeps it alive. Added notes that both run_all_benchmarks and State.exec evaluate with GIL held. Add note about intended way of using State.set_stream --- python/src/py_nvbench.cpp | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/python/src/py_nvbench.cpp b/python/src/py_nvbench.cpp index 0d58c96d..2f307977 100644 --- a/python/src/py_nvbench.cpp +++ b/python/src/py_nvbench.cpp @@ -982,9 +982,12 @@ Get `CudaStream` object from this configuration 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 owns the stream. NVBench stores a non-owning view and keeps -the current provider object alive until another provider replaces it or this -benchmark callback returns. +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"; @@ -1338,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 ---------- @@ -1518,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]