Skip to content

Commit

Permalink
fix dtype repr and stream pass-through
Browse files Browse the repository at this point in the history
  • Loading branch information
leofang committed Sep 5, 2024
1 parent ab83c5b commit 48a305c
Showing 1 changed file with 13 additions and 6 deletions.
19 changes: 13 additions & 6 deletions cuda_py/cuda/py/_memoryview.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ cdef class GPUMemoryView:
return (f"GPUMemoryView(ptr={self.ptr},\n"
+ f" shape={self.shape},\n"
+ f" strides={self.strides},\n"
+ f" dtype={get_simple_repr(numpy.dtype(self.dtype))},\n"
+ f" dtype={self.dtype.__name__},\n"
+ f" device_id={self.device_id},\n"
+ f" device_accessible={self.device_accessible},\n"
+ f" readonly={self.readonly},\n"
Expand All @@ -39,7 +39,7 @@ cdef class GPUMemoryView:
cdef str get_simple_repr(obj):
cdef object obj_class = obj.__class__
cdef str obj_repr
if obj_class.__module__ in (None, "__builtin__"):
if obj_class.__module__ in (None, "builtins"):
obj_repr = obj_class.__name__
else:
obj_repr = f"{obj_class.__module__}.{obj_class.__name__}"
Expand Down Expand Up @@ -78,17 +78,24 @@ cdef GPUMemoryView view_as_dlpack(obj, stream_ptr):
if dldevice == _kDLCPU:
device_accessible = False
assert device_id == 0
stream_ptr = None
if stream_ptr is None:
raise BufferError("stream=None is ambiguous with view()")
elif stream_ptr == -1:
stream_ptr = None
elif dldevice == _kDLCUDA:
device_accessible = True
stream_ptr = -1
# no need to check other stream values, it's a pass-through
if stream_ptr is None:
raise BufferError("stream=None is ambiguous with view()")
elif dldevice == _kDLCUDAHost:
device_accessible = True
assert device_id == 0
stream_ptr = None
# just do a pass-through without any checks, as pinned memory can be
# accessed on both host and device
elif dldevice == _kDLCUDAManaged:
device_accessible = True
stream_ptr = -1
# just do a pass-through without any checks, as managed memory can be
# accessed on both host and device
else:
raise BufferError("device not supported")

Expand Down

0 comments on commit 48a305c

Please sign in to comment.