cuda.core: fix some type signatures that are too generic

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

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

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:

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:

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

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,

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

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

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)

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\.$",

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"""

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
 
 

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"