diff --git a/cuda_core/cuda/core/_stream.pxd b/cuda_core/cuda/core/_stream.pxd index dc9a2da826c..abc4e396fd0 100644 --- a/cuda_core/cuda/core/_stream.pxd +++ b/cuda_core/cuda/core/_stream.pxd @@ -13,6 +13,7 @@ cdef class Stream: int _device_id int _nonblocking int _priority + int _synchronization_policy object __weakref__ @staticmethod diff --git a/cuda_core/cuda/core/_stream.pyi b/cuda_core/cuda/core/_stream.pyi index 99af5f9b15b..52a87299fbb 100644 --- a/cuda_core/cuda/core/_stream.pyi +++ b/cuda_core/cuda/core/_stream.pyi @@ -25,10 +25,15 @@ class StreamOptions: priority : int, optional Stream priority where lower number represents a higher priority. (Default to lowest priority) + synchronization_policy : SynchronizationPolicyType | None, optional + CPU wait policy applied when synchronizing this stream from the host. + Maps to ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY`` via + ``cuStreamSetAttribute``. (Default to None, leaving the driver default) """ nonblocking: cython.bint = True priority: int | None = None + synchronization_policy: object = None class IsStreamType(Protocol): @@ -110,6 +115,18 @@ class Stream: def priority(self) -> int: """Return the stream priority.""" + @property + def synchronization_policy(self): + """Return the stream's CPU wait policy for host-side synchronization. + + Maps to ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY``, queried via + ``cuStreamGetAttribute``. + """ + + @synchronization_policy.setter + def synchronization_policy(self, policy: object): + """Set the stream's CPU wait policy for subsequent work on this stream.""" + def sync(self) -> None: """Synchronize the stream.""" diff --git a/cuda_core/cuda/core/_stream.pyx b/cuda_core/cuda/core/_stream.pyx index c8c5faf74bc..066dbdfe6f5 100644 --- a/cuda_core/cuda/core/_stream.pyx +++ b/cuda_core/cuda/core/_stream.pyx @@ -49,6 +49,60 @@ if TYPE_CHECKING: __all__ = ['LEGACY_DEFAULT_STREAM', 'PER_THREAD_DEFAULT_STREAM', 'Stream', 'StreamOptions'] +cdef int _SYNC_POLICY_ATTR_ID = ( + cydriver.CUlaunchAttributeID_enum.CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY +) + + +cdef object _validate_synchronization_policy(object policy): + from cuda.core.typing import SynchronizationPolicyType + + if policy is None: + raise TypeError("synchronization_policy must not be None") + if isinstance(policy, SynchronizationPolicyType): + return policy + try: + value = int(policy) + except (TypeError, ValueError) as exc: + raise TypeError( + "Stream.synchronization_policy must be a SynchronizationPolicyType, " + f"cuda.bindings.driver.CUsynchronizationPolicy, or int; got {type(policy).__name__}" + ) from exc + try: + return SynchronizationPolicyType(value) + except ValueError as exc: + raise ValueError( + f"Stream.synchronization_policy must be one of " + f"{[member.name for member in SynchronizationPolicyType]}; got {policy!r}" + ) from exc + + +cdef int Stream_query_synchronization_policy(Stream self) except?-1 nogil: + cdef cydriver.CUstreamAttrValue value + HANDLE_RETURN(cydriver.cuStreamGetAttribute( + as_cu(self._h_stream), + _SYNC_POLICY_ATTR_ID, + &value, + )) + return value.syncPolicy + + +cdef int Stream_normalize_queried_sync_policy(int policy) nogil: + # cuStreamGetAttribute may return -1 when the attribute was never set. + if policy < 0: + return cydriver.CUsynchronizationPolicy_enum.CU_SYNC_POLICY_AUTO + return policy + + +cdef void Stream_apply_synchronization_policy(Stream self, int policy) except * nogil: + cdef cydriver.CUstreamAttrValue value + value.syncPolicy = policy + HANDLE_RETURN(cydriver.cuStreamSetAttribute( + as_cu(self._h_stream), + _SYNC_POLICY_ATTR_ID, + &value, + )) + @dataclass cdef class StreamOptions: @@ -61,11 +115,16 @@ cdef class StreamOptions: priority : int, optional Stream priority where lower number represents a higher priority. (Default to lowest priority) + synchronization_policy : SynchronizationPolicyType | None, optional + CPU wait policy applied when synchronizing this stream from the host. + Maps to ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY`` via + ``cuStreamSetAttribute``. (Default to None, leaving the driver default) """ nonblocking : cython.bint = True priority: int | None = None + synchronization_policy: object = None class IsStreamType(Protocol): @@ -110,6 +169,7 @@ cdef class Stream: s._device_id = -1 # lazy init'd (invalid sentinel) s._nonblocking = -1 # lazy init'd s._priority = INT32_MIN # lazy init'd + s._synchronization_policy = INT32_MIN # lazy init'd return s @classmethod @@ -146,12 +206,14 @@ cdef class Stream: cdef StreamOptions opts = check_or_create_options(StreamOptions, options, "Stream options") nonblocking = opts.nonblocking priority = opts.priority + sync_policy_opt = opts.synchronization_policy cdef unsigned int flags = (cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING if nonblocking else cydriver.CUstream_flags.CU_STREAM_DEFAULT) # TODO: we might want to consider memoizing high/low per CUDA context and avoid this call cdef int high, low cdef cydriver.CUresult res_code + cdef int sync_policy with nogil: res_code = cydriver.cuCtxGetStreamPriorityRange(&high, &low) if res_code != cydriver.CUresult.CUDA_SUCCESS: @@ -190,6 +252,12 @@ cdef class Stream: cdef Stream self = Stream._from_handle(cls, h_stream) self._nonblocking = int(nonblocking) self._priority = prio + if sync_policy_opt is not None: + validated = _validate_synchronization_policy(sync_policy_opt) + sync_policy = int(validated) + with nogil: + Stream_apply_synchronization_policy(self, sync_policy) + self._synchronization_policy = sync_policy if device_id is not None: self._device_id = device_id return self @@ -251,6 +319,36 @@ cdef class Stream: self._priority = prio return self._priority + @property + def synchronization_policy(self): + """Return the stream's CPU wait policy for host-side synchronization. + + Maps to ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY``, queried via + ``cuStreamGetAttribute``. + """ + from cuda.core.typing import SynchronizationPolicyType + + cdef int policy + if self._synchronization_policy == INT32_MIN or Stream_is_default_token(self): + with nogil: + policy = Stream_query_synchronization_policy(self) + policy = Stream_normalize_queried_sync_policy(policy) + if not Stream_is_default_token(self): + self._synchronization_policy = policy + else: + policy = self._synchronization_policy + return SynchronizationPolicyType(policy) + + @synchronization_policy.setter + def synchronization_policy(self, object policy): + """Set the stream's CPU wait policy for subsequent work on this stream.""" + cdef object validated = _validate_synchronization_policy(policy) + cdef int value = int(validated) + with nogil: + Stream_apply_synchronization_policy(self, value) + if not Stream_is_default_token(self): + self._synchronization_policy = value + def sync(self) -> None: """Synchronize the stream.""" with nogil: diff --git a/cuda_core/cuda/core/typing.py b/cuda_core/cuda/core/typing.py index 1bf9bb7c0d2..0d40d11d9a3 100644 --- a/cuda_core/cuda/core/typing.py +++ b/cuda_core/cuda/core/typing.py @@ -5,6 +5,7 @@ """Public type aliases, protocols, and enumerations used in cuda.core API signatures.""" import sys +from enum import IntEnum from typing import TYPE_CHECKING from typing import Literal as _Literal from typing import TypeAlias as _TypeAlias @@ -47,6 +48,7 @@ class StrEnum(str, Enum): "ProcessStateType", "ReadModeType", "SourceCodeType", + "SynchronizationPolicyType", "VirtualMemoryAccessType", "VirtualMemoryAllocationType", "VirtualMemoryGranularityType", @@ -124,6 +126,24 @@ class PCHStatusType(StrEnum): FAILED = "failed" +class SynchronizationPolicyType(IntEnum): + """CPU wait policy when synchronizing a stream from the host. + + Maps to ``cuda.bindings.driver.CUsynchronizationPolicy`` and is applied via + ``cuStreamSetAttribute`` with ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY``. + + * ``AUTO`` — driver selects the policy (default for new streams). + * ``SPIN`` — busy-wait on the CPU (lowest latency). + * ``YIELD`` — yield the CPU while waiting. + * ``BLOCKING_SYNC`` — block in the OS scheduler while waiting. + """ + + AUTO = driver.CUsynchronizationPolicy.CU_SYNC_POLICY_AUTO + SPIN = driver.CUsynchronizationPolicy.CU_SYNC_POLICY_SPIN + YIELD = driver.CUsynchronizationPolicy.CU_SYNC_POLICY_YIELD + BLOCKING_SYNC = driver.CUsynchronizationPolicy.CU_SYNC_POLICY_BLOCKING_SYNC + + class GraphConditionalType(StrEnum): """Conditional node flavor for :class:`~cuda.core.graph.GraphBuilder`. diff --git a/cuda_core/docs/source/api_private.rst b/cuda_core/docs/source/api_private.rst index 80675799c07..de161c96ae0 100644 --- a/cuda_core/docs/source/api_private.rst +++ b/cuda_core/docs/source/api_private.rst @@ -35,6 +35,7 @@ CUDA runtime typing.ProcessStateType typing.ReadModeType typing.SourceCodeType + typing.SynchronizationPolicyType typing.VirtualMemoryAccessType typing.VirtualMemoryAllocationType typing.VirtualMemoryGranularityType diff --git a/cuda_core/tests/test_stream.py b/cuda_core/tests/test_stream.py index 55f34bbc9ec..a89d8c963e1 100644 --- a/cuda_core/tests/test_stream.py +++ b/cuda_core/tests/test_stream.py @@ -469,3 +469,88 @@ def test_created_stream_keeps_its_own_context(): finally: dev0.set_current() stream.close() + + +# ============================================================================ +# Stream synchronization policy tests +# ============================================================================ + + +@pytest.mark.parametrize( + "policy", + [ + "AUTO", + "SPIN", + "YIELD", + "BLOCKING_SYNC", + ], +) +def test_stream_synchronization_policy_set_get(init_cuda, policy): + from cuda.core.typing import SynchronizationPolicyType + + stream = Device().create_stream() + enum_policy = SynchronizationPolicyType[policy] + stream.synchronization_policy = enum_policy + assert stream.synchronization_policy == enum_policy + stream.close() + + +@pytest.mark.parametrize("invalid_policy", ["spin", -1, 99]) +def test_stream_synchronization_policy_invalid(init_cuda, invalid_policy): + stream = Device().create_stream() + with pytest.raises((TypeError, ValueError)): + stream.synchronization_policy = invalid_policy + stream.close() + + +def test_stream_synchronization_policy_default(init_cuda): + from cuda.core.typing import SynchronizationPolicyType + + stream = Device().create_stream() + assert stream.synchronization_policy == SynchronizationPolicyType.AUTO + stream.close() + + +def test_stream_options_synchronization_policy(init_cuda): + from cuda.core.typing import SynchronizationPolicyType + + stream = Device().create_stream(options=StreamOptions(synchronization_policy=SynchronizationPolicyType.YIELD)) + assert stream.synchronization_policy == SynchronizationPolicyType.YIELD + stream.close() + + +@pytest.mark.parametrize( + "policy", + [ + "AUTO", + "SPIN", + "YIELD", + "BLOCKING_SYNC", + ], +) +@pytest.mark.agent_authored(model="composer-2.5") +def test_stream_synchronization_policy_launch_and_sync(init_cuda, policy): + """Stream-level sync policy is applied via cuStreamSetAttribute and survives launch+sync.""" + import cuda.bindings + from cuda.core import Device, LaunchConfig, Program, ProgramOptions, launch + from cuda.core._utils.version import driver_version + from cuda.core.typing import ObjectCodeFormatType, SourceCodeType, SynchronizationPolicyType + + if int(cuda.bindings.__version__.split(".")[0]) >= 13 and driver_version()[0] < 13: + pytest.skip("CUDA 13 bindings produce modules incompatible with CUDA 12 drivers") + + dev = Device() + dev.set_current() + stream = dev.create_stream() + stream.synchronization_policy = SynchronizationPolicyType[policy] + + code = 'extern "C" __global__ void noop() {}' + arch = "".join(f"{i}" for i in dev.compute_capability) + program = Program(code, SourceCodeType.CXX, options=ProgramOptions(arch=f"sm_{arch}")) + mod = program.compile(ObjectCodeFormatType.CUBIN) + ker = mod.get_kernel("noop") + + launch(stream, LaunchConfig(grid=1, block=1), ker) + stream.sync() + assert stream.synchronization_policy == SynchronizationPolicyType[policy] + stream.close()