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

python/src/py_nvbench.cpp

@@ -24,6 +24,7 @@
 
 #include <nvbench/nvbench.cuh>
 
+#include <cstdint>
 #include <cstdio>
 #include <cstdlib>
 #include <functional>
@@ -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<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 +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 &current_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(

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)