Implement cuda.bench.State.set_stream(provider)

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: ...

python/src/py_nvbench.cpp

@@ -24,12 +24,14 @@
 
 #include <nvbench/nvbench.cuh>
 
+#include <cstdint>
 #include <cstdio>
 #include <cstdlib>
 #include <functional>
 #include <memory>
 #include <sstream>
 #include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -38,6 +40,13 @@ namespace py = pybind11;
 namespace
 {
 
+std::unordered_map<nvbench::state *, py::object> &cuda_stream_provider_cache()
+{
+  // Avoid destroying py::object values after Python interpreter shutdown.
+  static auto *cache = new std::unordered_map<nvbench::state *, py::object>{};
+  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<GlobalBenchmarkRegistry, py::nodelete> global_registry{};
 
+cudaStream_t extract_cuda_stream_from_provider(const py::handle &stream_provider)
+{
+  if (py::isinstance<nvbench::cuda_stream>(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<py::tuple>(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<py::tuple>(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<int>();
+    stream_handle    = stream_info[1].cast<std::uintptr_t>();
+  }
+  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<cudaStream_t>(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 &current_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<py::object>(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]

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)