diff --git a/cuda_core/cuda/core/_context.pyi b/cuda_core/cuda/core/_context.pyi index afbc130882e..ff6aa5276e9 100644 --- a/cuda_core/cuda/core/_context.pyi +++ b/cuda_core/cuda/core/_context.pyi @@ -8,7 +8,7 @@ from dataclasses import dataclass import cuda.bindings.driver from cuda.core._device_resources import (DeviceResources, SMResource, WorkqueueResource) -from cuda.core._stream import Stream +from cuda.core._stream import Stream, StreamOptions class Context: @@ -47,7 +47,7 @@ class Context: Raises :class:`RuntimeError` if the context has been closed. """ - def create_stream(self, options: object=None) -> Stream: + def create_stream(self, options: StreamOptions | None=None) -> Stream: """Create a new stream bound to this green context. This method is only available on green contexts. For primary diff --git a/cuda_core/cuda/core/_context.pyx b/cuda_core/cuda/core/_context.pyx index 6da72addb0d..4e28481d0e8 100644 --- a/cuda_core/cuda/core/_context.pyx +++ b/cuda_core/cuda/core/_context.pyx @@ -21,7 +21,7 @@ from cuda.core._resource_handles cimport ( as_intptr, as_py, ) -from cuda.core._stream import Stream +from cuda.core._stream import Stream, StreamOptions from cuda.core._utils.cuda_utils cimport HANDLE_RETURN if TYPE_CHECKING: @@ -95,7 +95,7 @@ cdef class Context: raise RuntimeError("Cannot query resources on a closed context") return DeviceResources._init_from_ctx(self._h_context, self._device_id) - def create_stream(self, options: object = None) -> Stream: + def create_stream(self, options: StreamOptions | None = None) -> Stream: """Create a new stream bound to this green context. This method is only available on green contexts. For primary diff --git a/cuda_core/cuda/core/_device.pyi b/cuda_core/cuda/core/_device.pyi index e83aef8a8d0..4672d29d879 100644 --- a/cuda_core/cuda/core/_device.pyi +++ b/cuda_core/cuda/core/_device.pyi @@ -9,7 +9,7 @@ from cuda.core._context import Context, ContextOptions from cuda.core._device_resources import DeviceResources from cuda.core._event import Event, EventOptions from cuda.core._memory._buffer import Buffer, MemoryResource -from cuda.core._stream import IsStreamType, Stream +from cuda.core._stream import IsStreamType, Stream, StreamOptions from cuda.core._utils.cuda_utils import ComputeCapability from cuda.core.graph import GraphBuilder from cuda.core.texture import (MipmappedArray, MipmappedArrayOptions, @@ -821,7 +821,7 @@ class Device: """ - def create_stream(self, obj: IsStreamType | None=None, options: object=None) -> Stream: + def create_stream(self, obj: IsStreamType | None=None, options: StreamOptions | None=None) -> Stream: """Create a :obj:`~_stream.Stream` object. New stream objects can be created in two different ways: diff --git a/cuda_core/cuda/core/_device.pyx b/cuda_core/cuda/core/_device.pyx index a0a0f472f2b..d0f523e2d2a 100644 --- a/cuda_core/cuda/core/_device.pyx +++ b/cuda_core/cuda/core/_device.pyx @@ -28,7 +28,7 @@ from cuda.core._resource_handles cimport ( as_cu, ) -from cuda.core._stream import IsStreamType, Stream +from cuda.core._stream import IsStreamType, Stream, StreamOptions from cuda.core._utils.clear_error_support import assert_type from cuda.core._utils.cuda_utils import ( ComputeCapability, @@ -1376,7 +1376,7 @@ class Device: return Context._from_green_ctx(Context, h_green, self._device_id) - def create_stream(self, obj: IsStreamType | None = None, options: object = None) -> Stream: + def create_stream(self, obj: IsStreamType | None = None, options: StreamOptions | None = None) -> Stream: """Create a :obj:`~_stream.Stream` object. New stream objects can be created in two different ways: diff --git a/cuda_core/cuda/core/_memoryview.pyi b/cuda_core/cuda/core/_memoryview.pyi index e0ed0d3cf0d..240c72ca438 100644 --- a/cuda_core/cuda/core/_memoryview.pyi +++ b/cuda_core/cuda/core/_memoryview.pyi @@ -10,6 +10,7 @@ import numpy from cuda.core._layout import _StridedLayout from cuda.core._memory import Buffer from cuda.core._stream import Stream +from cuda.core._tensor_map import TensorMapDescriptorOptions from ._dlpack import * @@ -165,7 +166,7 @@ class StridedMemoryView: Same as calling :meth:`from_buffer` with the current buffer. """ - def as_tensor_map(self, box_dim: tuple[int, ...] | None=None, *, options: object=None, element_strides: tuple[int, ...] | None=None, data_type: object=None, interleave: object=None, swizzle: object=None, l2_promotion: object=None, oob_fill: object=None) -> object: + def as_tensor_map(self, box_dim: tuple[int, ...] | None=None, *, options: TensorMapDescriptorOptions | None=None, element_strides: tuple[int, ...] | None=None, data_type: object=None, interleave: object=None, swizzle: object=None, l2_promotion: object=None, oob_fill: object=None) -> object: """Create a tiled :obj:`TensorMapDescriptor` from this view. This is the public entry point for creating tiled tensor map diff --git a/cuda_core/cuda/core/_memoryview.pyx b/cuda_core/cuda/core/_memoryview.pyx index bbe5a887700..03cb6dcd45e 100644 --- a/cuda_core/cuda/core/_memoryview.pyx +++ b/cuda_core/cuda/core/_memoryview.pyx @@ -16,7 +16,10 @@ import functools import sys import warnings from collections.abc import Callable # no-cython-lint # used in string annotations below -from typing import Any # no-cython-lint # used in string annotations below +from typing import TYPE_CHECKING, Any # no-cython-lint # used in string annotations below + +if TYPE_CHECKING: + from cuda.core._tensor_map import TensorMapDescriptorOptions import numpy @@ -377,7 +380,7 @@ cdef class StridedMemoryView: self, box_dim: tuple[int, ...] | None = None, *, - options: object = None, + options: TensorMapDescriptorOptions | None = None, element_strides: tuple[int, ...] | None = None, data_type: object = None, interleave: object = None, diff --git a/cuda_core/cuda/core/_stream.pyi b/cuda_core/cuda/core/_stream.pyi index 99af5f9b15b..ec236dd21c7 100644 --- a/cuda_core/cuda/core/_stream.pyi +++ b/cuda_core/cuda/core/_stream.pyi @@ -77,7 +77,7 @@ class Stream: """Return the per-thread default stream (supports subclassing).""" @classmethod - def _init(cls, obj: IsStreamType | None=None, options: object=None, device_id: int | None=None, ctx: Context | None=None) -> Stream: + def _init(cls, obj: IsStreamType | None=None, options: StreamOptions | None=None, device_id: int | None=None, ctx: Context | None=None) -> Stream: ... def __cuda_stream__(self) -> tuple[int, int]: diff --git a/cuda_core/cuda/core/_stream.pyx b/cuda_core/cuda/core/_stream.pyx index a376db96e11..fde94ba14c5 100644 --- a/cuda_core/cuda/core/_stream.pyx +++ b/cuda_core/cuda/core/_stream.pyx @@ -123,7 +123,7 @@ cdef class Stream: return Stream._from_handle(cls, get_per_thread_stream()) @classmethod - def _init(cls, obj: IsStreamType | None = None, options: object = None, + def _init(cls, obj: IsStreamType | None = None, options: StreamOptions | None = None, device_id: int | None = None, ctx: Context | None = None) -> Stream: cdef StreamHandle h_stream cdef cydriver.CUstream borrowed diff --git a/cuda_core/tests/graph/test_graph_builder.py b/cuda_core/tests/graph/test_graph_builder.py index 0033cb99e2c..cfabde64de4 100644 --- a/cuda_core/tests/graph/test_graph_builder.py +++ b/cuda_core/tests/graph/test_graph_builder.py @@ -17,7 +17,7 @@ from packaging.version import Version import cuda.bindings -from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, Program, ProgramOptions, launch +from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, Program, ProgramOptions, StreamOptions, launch from cuda.core.graph import GraphBuilder, GraphDefinition from cuda.core.graph._graph_builder import ( _capture_callback_with_tail_failure_for_testing, @@ -826,7 +826,7 @@ def test_pdl_same_stream_primary_secondary_overlap_via_graph(init_cuda): primary = module.get_kernel("primary_kernel") secondary = module.get_kernel("secondary_kernel") - stream = dev.create_stream(options={"nonblocking": True}) + stream = dev.create_stream(options=StreamOptions(nonblocking=True)) mr = LegacyPinnedMemoryResource() secondary_started = np.from_dlpack(mr.allocate(4)).view(np.int32) overlapped = np.from_dlpack(mr.allocate(4)).view(np.int32) diff --git a/cuda_core/tests/memory_ipc/test_event_ipc.py b/cuda_core/tests/memory_ipc/test_event_ipc.py index e3cefe6a211..7608b9c8794 100644 --- a/cuda_core/tests/memory_ipc/test_event_ipc.py +++ b/cuda_core/tests/memory_ipc/test_event_ipc.py @@ -108,10 +108,10 @@ def test_event_is_monadic(ipc_device): """Check that IPC-enabled events are always bound and cannot be reset.""" device = ipc_device with pytest.raises(TypeError, match=r"^IPC-enabled events must be bound; use Stream.record for creation\.$"): - device.create_event({"ipc_enabled": True}) + device.create_event(EventOptions(ipc_enabled=True)) stream = device.create_stream() - e = stream.record(options={"ipc_enabled": True}) + e = stream.record(options=EventOptions(ipc_enabled=True)) with pytest.raises( TypeError, match=r"^IPC-enabled events should not be re-recorded, instead create a new event by supplying options\.$", diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index e5cf05b435d..c24fdcdf4a4 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -22,6 +22,7 @@ LegacyPinnedMemoryResource, Program, ProgramOptions, + StreamOptions, launch, ) from cuda.core._memory._legacy import _SynchronousMemoryResource @@ -218,7 +219,7 @@ def test_pdl_primary_secondary_overlap_same_stream(): if dev.compute_capability < (9, 0): pytest.skip("Programmatic Dependent Launch requires compute capability >= 9.0") dev.set_current() - stream = dev.create_stream(options={"nonblocking": True}) + stream = dev.create_stream(options=StreamOptions(nonblocking=True)) # clock64 budgets are in GPU cycles; keep the post-trigger window long enough # for the secondary to boot, but short enough for a unit test. @@ -483,7 +484,7 @@ def test_launch_scalar_argument(python_type, cpp_type, init_value): def test_cooperative_launch(): dev = Device() dev.set_current() - s = dev.create_stream(options={"nonblocking": True}) + s = dev.create_stream(options=StreamOptions(nonblocking=True)) # CUDA kernel templated on type T code = r""" diff --git a/cuda_core/tests/test_object_protocols.py b/cuda_core/tests/test_object_protocols.py index baf790abea8..865f434d409 100644 --- a/cuda_core/tests/test_object_protocols.py +++ b/cuda_core/tests/test_object_protocols.py @@ -23,6 +23,7 @@ Device, DeviceMemoryResource, DeviceMemoryResourceOptions, + EventOptions, Kernel, LaunchConfig, Program, @@ -243,7 +244,7 @@ def sample_ipc_buffer_descriptor(ipc_device): def sample_ipc_event_descriptor(ipc_device): """An IPCEventDescriptor.""" stream = ipc_device.create_stream() - e = stream.record(options={"ipc_enabled": True}) + e = stream.record(options=EventOptions(ipc_enabled=True)) return e.ipc_descriptor diff --git a/cuda_core/tests/test_program.py b/cuda_core/tests/test_program.py index c4e2e0c6428..64cefd8cbc9 100644 --- a/cuda_core/tests/test_program.py +++ b/cuda_core/tests/test_program.py @@ -368,7 +368,7 @@ def test_program_options_name_accepts_none(name): # This is tested against the current device's arch def test_program_compile_valid_target_type(init_cuda): code = 'extern "C" __global__ void my_kernel() {}' - program = Program(code, "c++", options={"name": "42"}) + program = Program(code, "c++", options=ProgramOptions(name="42")) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") @@ -380,7 +380,7 @@ def test_program_compile_valid_target_type(init_cuda): ptx_kernel = ptx_object_code.get_kernel("my_kernel") assert isinstance(ptx_kernel, Kernel) - program = Program(ptx_object_code.code.decode(), "ptx", options={"name": "24"}) + program = Program(ptx_object_code.code.decode(), "ptx", options=ProgramOptions(name="24")) cubin_object_code = program.compile("cubin") assert isinstance(cubin_object_code, ObjectCode) assert cubin_object_code.name == "24"