@@ -22,6 +22,7 @@ from cuda.core._resource_handles cimport (
2222 as_cu,
2323 set_deallocation_stream,
2424)
25+ from cuda.core.typing import DevicePointerType
2526
2627from cuda.core._stream cimport Stream, Stream_accept
2728from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value
3536 BufferProtocol = object
3637
3738from cuda.core._dlpack import classify_dl_device, make_py_capsule
38- from cuda.core._utils.cuda_utils import driver
3939from cuda.core._device import Device
4040
4141
@@ -65,11 +65,6 @@ register_mr_dealloc_callback(_mr_dealloc_callback)
6565__all__ = [' Buffer' , ' MemoryResource' ]
6666
6767
68- DevicePointerT = driver.CUdeviceptr | int | None
69- """
70- A type union of :obj:`~driver.CUdeviceptr`, `int` and `None` for hinting
71- :attr:`Buffer.handle`.
72- """
7368
7469cdef class Buffer:
7570 """ Represent a handle to allocated memory.
@@ -97,7 +92,7 @@ cdef class Buffer:
9792
9893 @classmethod
9994 def _init (
100- cls , ptr: DevicePointerT , size_t size , mr: MemoryResource | None = None ,
95+ cls , ptr: DevicePointerType , size_t size , mr: MemoryResource | None = None ,
10196 ipc_descriptor: IPCBufferDescriptor | None = None ,
10297 owner : object | None = None
10398 ):
@@ -132,14 +127,14 @@ cdef class Buffer:
132127
133128 @staticmethod
134129 def from_handle (
135- ptr: DevicePointerT , size_t size , mr: MemoryResource | None = None ,
130+ ptr: DevicePointerType , size_t size , mr: MemoryResource | None = None ,
136131 owner: object | None = None ,
137132 ) -> Buffer:
138133 """Create a new :class:`Buffer` object from a pointer.
139134
140135 Parameters
141136 ----------
142- ptr : :obj:`~_memory.DevicePointerT `
137+ ptr : :obj:`~_memory.DevicePointerType `
143138 Allocated buffer handle object
144139 size : int
145140 Memory size of the buffer
@@ -347,7 +342,7 @@ cdef class Buffer:
347342 return self._mem_attrs.device_id
348343
349344 @property
350- def handle(self ) -> DevicePointerT :
345+ def handle(self ) -> DevicePointerType :
351346 """Return the buffer handle object.
352347
353348 .. caution::
@@ -515,12 +510,12 @@ cdef class MemoryResource:
515510 """
516511 raise TypeError("MemoryResource.allocate must be implemented by subclasses.")
517512
518- def deallocate(self , ptr: DevicePointerT , size_t size , stream: Stream | GraphBuilder | None = None ):
513+ def deallocate(self , ptr: DevicePointerType , size_t size , stream: Stream | GraphBuilder | None = None ):
519514 """ Deallocate a buffer previously allocated by this resource.
520515
521516 Parameters
522517 ----------
523- ptr : :obj:`~_memory.DevicePointerT `
518+ ptr : :obj:`~_memory.DevicePointerType `
524519 The pointer or handle to the buffer to deallocate.
525520 size : int
526521 The size of the buffer to deallocate, in bytes.
0 commit comments