diff --git a/cuda_core/cuda/core/_launcher.pyi b/cuda_core/cuda/core/_launcher.pyi index 56f189fe72f..27f7b7842b2 100644 --- a/cuda_core/cuda/core/_launcher.pyi +++ b/cuda_core/cuda/core/_launcher.pyi @@ -1,8 +1,11 @@ # This file was generated by stubgen-pyx v0.2.23 from cuda_core/cuda/core/_launcher.pyx +from collections.abc import Callable + from cuda.core._launch_config import LaunchConfig from cuda.core._module import Kernel from cuda.core._stream import Stream +from cuda.core._utils.pycompat import BufferProtocol from cuda.core.graph import GraphBuilder from cuda.core.typing import IsStreamType @@ -27,7 +30,7 @@ def launch(stream: Stream | GraphBuilder | IsStreamType, config: LaunchConfig, k launching kernel. """ -def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=None) -> None: +def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn: Callable[..., object], *, user_data: int | BufferProtocol | None=None) -> None: """Enqueue a host function on a stream in stream order. The function runs on a driver-owned thread once all work enqueued on the diff --git a/cuda_core/cuda/core/_launcher.pyx b/cuda_core/cuda/core/_launcher.pyx index f5becd4b316..51397fee524 100644 --- a/cuda_core/cuda/core/_launcher.pyx +++ b/cuda_core/cuda/core/_launcher.pyx @@ -22,6 +22,9 @@ from math import prod from typing import TYPE_CHECKING if TYPE_CHECKING: + from collections.abc import Callable + + from cuda.core._utils.pycompat import BufferProtocol from cuda.core.graph import GraphBuilder from cuda.core.typing import IsStreamType @@ -70,7 +73,7 @@ def launch( HANDLE_RETURN(cydriver.cuLaunchKernelEx(&drv_cfg, func_handle, args_ptr, NULL)) -def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=None) -> None: +def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> None: """Enqueue a host function on a stream in stream order. The function runs on a driver-owned thread once all work enqueued on the diff --git a/cuda_core/cuda/core/graph/_graph_builder.pyi b/cuda_core/cuda/core/graph/_graph_builder.pyi index bf20e0cb3bf..106a926866f 100644 --- a/cuda_core/cuda/core/graph/_graph_builder.pyi +++ b/cuda_core/cuda/core/graph/_graph_builder.pyi @@ -1,10 +1,12 @@ # This file was generated by stubgen-pyx v0.2.23 from cuda_core/cuda/core/graph/_graph_builder.pyx +from collections.abc import Callable from dataclasses import dataclass from typing import overload from cuda.core._stream import Stream from cuda.core._utils.cuda_utils import driver +from cuda.core._utils.pycompat import BufferProtocol from cuda.core.graph._graph_definition import GraphCondition, GraphDefinition from cuda.core.graph._subclasses import (ChildGraphNode, EventRecordNode, EventWaitNode, @@ -408,7 +410,7 @@ class GraphBuilder: child : :obj:`~graph.GraphBuilder` The child graph builder. Must have finished building. """ - def callback(self, fn, *, user_data=None) -> None: + def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None=None) -> None: """Add a host callback to the graph during stream capture. The callback runs on the host CPU when the graph reaches this point diff --git a/cuda_core/cuda/core/graph/_graph_builder.pyx b/cuda_core/cuda/core/graph/_graph_builder.pyx index 2ac92caacbc..194fa98c9c2 100644 --- a/cuda_core/cuda/core/graph/_graph_builder.pyx +++ b/cuda_core/cuda/core/graph/_graph_builder.pyx @@ -62,8 +62,11 @@ from cuda.core._utils.cuda_utils import ( ) if TYPE_CHECKING: + from collections.abc import Callable from typing import overload + from cuda.core._utils.pycompat import BufferProtocol + from cuda.core.graph._graph_definition import GraphDefinition else: # Cython applies a method decorator by rebinding the type's dict entry, so @@ -947,7 +950,7 @@ cdef class GraphBuilder: ) ) - def callback(self, fn, *, user_data=None) -> None: + def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> None: """Add a host callback to the graph during stream capture. The callback runs on the host CPU when the graph reaches this point diff --git a/cuda_core/cuda/core/graph/_graph_definition.pyi b/cuda_core/cuda/core/graph/_graph_definition.pyi index 6a139eb7935..871201d4bb0 100644 --- a/cuda_core/cuda/core/graph/_graph_definition.pyi +++ b/cuda_core/cuda/core/graph/_graph_definition.pyi @@ -1,12 +1,15 @@ # This file was generated by stubgen-pyx v0.2.23 from cuda_core/cuda/core/graph/_graph_definition.pyx """GraphDefinition: explicit CUDA graph definition.""" +from collections.abc import Callable + from cuda.core._device import Device from cuda.core._event import Event from cuda.core._launch_config import LaunchConfig from cuda.core._memory._buffer import Buffer from cuda.core._module import Kernel from cuda.core._utils.cuda_utils import driver +from cuda.core._utils.pycompat import BufferProtocol from cuda.core.graph._graph_builder import (Graph, GraphCompleteOptions, GraphDebugPrintOptions) from cuda.core.graph._graph_node import GraphNode @@ -127,7 +130,7 @@ class GraphDefinition: See :meth:`GraphNode.wait` for full documentation. """ - def callback(self, fn, *, user_data=None) -> HostCallbackNode: + def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None=None) -> HostCallbackNode: """Add an entry-point host callback node (no dependencies). See :meth:`GraphNode.callback` for full documentation. diff --git a/cuda_core/cuda/core/graph/_graph_definition.pyx b/cuda_core/cuda/core/graph/_graph_definition.pyx index f9254e8c229..ea97c47bc15 100644 --- a/cuda_core/cuda/core/graph/_graph_definition.pyx +++ b/cuda_core/cuda/core/graph/_graph_definition.pyx @@ -31,11 +31,14 @@ from cuda.core._utils.cuda_utils import driver from cuda.core.typing import GraphMemoryType if TYPE_CHECKING: + from collections.abc import Callable + from cuda.core._device import Device from cuda.core._event import Event from cuda.core._launch_config import LaunchConfig from cuda.core._memory._buffer import Buffer from cuda.core._module import Kernel + from cuda.core._utils.pycompat import BufferProtocol from cuda.core.graph._graph_builder import ( Graph, GraphCompleteOptions, @@ -250,7 +253,7 @@ cdef class GraphDefinition: """ return self._entry.wait(event) - def callback(self, fn, *, user_data=None) -> HostCallbackNode: + def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> HostCallbackNode: """Add an entry-point host callback node (no dependencies). See :meth:`GraphNode.callback` for full documentation. diff --git a/cuda_core/cuda/core/graph/_graph_node.pyi b/cuda_core/cuda/core/graph/_graph_node.pyi index b22f3d64dcd..a92f09d1467 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyi +++ b/cuda_core/cuda/core/graph/_graph_node.pyi @@ -2,7 +2,7 @@ """GraphNode base class — factory, properties, and builder methods.""" import weakref -from collections.abc import Iterable +from collections.abc import Callable, Iterable from cuda.core._device import Device from cuda.core._event import Event @@ -10,6 +10,7 @@ from cuda.core._launch_config import LaunchConfig from cuda.core._memory._buffer import Buffer from cuda.core._module import Kernel from cuda.core._utils.cuda_utils import driver +from cuda.core._utils.pycompat import BufferProtocol from cuda.core.graph._adjacency_set_proxy import AdjacencySetProxy from cuda.core.graph._graph_definition import GraphCondition, GraphDefinition from cuda.core.graph._subclasses import (AllocNode, ChildGraphNode, EmptyNode, @@ -297,7 +298,7 @@ class GraphNode: EventWaitNode A new EventWaitNode representing the event wait operation. """ - def callback(self, fn, *, user_data=None) -> object: + def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None=None) -> object: """Add a host callback node depending on this node. The callback runs on the host CPU when the graph reaches this node. diff --git a/cuda_core/cuda/core/graph/_graph_node.pyx b/cuda_core/cuda/core/graph/_graph_node.pyx index 4abb1f7106c..d13413a576d 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/cuda_core/cuda/core/graph/_graph_node.pyx @@ -79,7 +79,10 @@ from cuda.core._utils.cuda_utils import driver from cuda.core.typing import GraphMemoryType if TYPE_CHECKING: + from collections.abc import Callable + from cuda.core._device import Device + from cuda.core._utils.pycompat import BufferProtocol __all__ = ['GraphNode'] @@ -495,7 +498,7 @@ cdef class GraphNode: """ return GN_wait_event(self, event) - def callback(self, fn, *, user_data=None) -> object: + def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> object: """Add a host callback node depending on this node. The callback runs on the host CPU when the graph reaches this node. diff --git a/cuda_core/examples/host_launch.py b/cuda_core/examples/host_launch.py index e70ec675b93..0e7bc2d855b 100644 --- a/cuda_core/examples/host_launch.py +++ b/cuda_core/examples/host_launch.py @@ -10,7 +10,7 @@ # ################################################################################ # /// script -# dependencies = ["cuda_bindings", "cuda_core"] +# dependencies = ["cuda_core"] # /// import ctypes