diff --git a/cuda_bindings/cuda/bindings/_lib/utils.pxi b/cuda_bindings/cuda/bindings/_lib/utils.pxi index 7783afed97c..7e4b1185fc3 100644 --- a/cuda_bindings/cuda/bindings/_lib/utils.pxi +++ b/cuda_bindings/cuda/bindings/_lib/utils.pxi @@ -308,11 +308,12 @@ cdef class _HelperCUpointer_attribute: if self._attr in (cydriver.CUpointer_attribute_enum.CU_POINTER_ATTRIBUTE_CONTEXT,): return self._ctx elif self._attr in (cydriver.CUpointer_attribute_enum.CU_POINTER_ATTRIBUTE_MEMORY_TYPE, - cydriver.CUpointer_attribute_enum.CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, cydriver.CUpointer_attribute_enum.CU_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES, cydriver.CUpointer_attribute_enum.CU_POINTER_ATTRIBUTE_IS_GPU_DIRECT_RDMA_CAPABLE, cydriver.CUpointer_attribute_enum.CU_POINTER_ATTRIBUTE_ACCESS_FLAGS,): return self._uint + elif self._attr in (cydriver.CUpointer_attribute_enum.CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL,): + return self._int elif self._attr in (cydriver.CUpointer_attribute_enum.CU_POINTER_ATTRIBUTE_DEVICE_POINTER, cydriver.CUpointer_attribute_enum.CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,): return self._devptr diff --git a/cuda_bindings/tests/test_cuda.py b/cuda_bindings/tests/test_cuda.py index 7bef2b844aa..bd871bcd5eb 100644 --- a/cuda_bindings/tests/test_cuda.py +++ b/cuda_bindings/tests/test_cuda.py @@ -1335,3 +1335,35 @@ def test_dealloc_clears_array_field_in_external_struct(): f"external struct still holds a dangling pointer ({attrs_after:#x}) " "where attrs was, after the aliasing wrapper was destroyed" ) + + +def test_pointer_device_ordinal_matches_allocation_device(): + # Regression test for cuPointerGetAttribute(s)(DEVICE_ORDINAL): the binding + # must report the real device the pointer was allocated on, not always 0. + # Needs >= 2 GPUs: on device 0 a wrong 0 is indistinguishable from correct. + err, ndev = cudart.cudaGetDeviceCount() + assert err == cudart.cudaError_t.cudaSuccess + if ndev < 2: + pytest.skip("needs >= 2 GPUs to distinguish a wrong 0 from the correct ordinal") + + ord_attr = cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL + for dev in range(ndev): + (err,) = cudart.cudaSetDevice(dev) + assert err == cudart.cudaError_t.cudaSuccess + err, ptr = cudart.cudaMalloc(256) + assert err == cudart.cudaError_t.cudaSuccess + try: + err, attrs = cudart.cudaPointerGetAttributes(ptr) + assert err == cudart.cudaError_t.cudaSuccess + expected = attrs.device + + err, singular = cuda.cuPointerGetAttribute(ord_attr, cuda.CUdeviceptr(ptr)) + assert err == cuda.CUresult.CUDA_SUCCESS + assert singular == expected, f"singular ordinal {singular} != {expected} on device {dev}" + + err, plural = cuda.cuPointerGetAttributes(1, [ord_attr], cuda.CUdeviceptr(ptr)) + assert err == cuda.CUresult.CUDA_SUCCESS + assert plural[0] == expected, f"plural ordinal {plural[0]} != {expected} on device {dev}" + finally: + (err,) = cudart.cudaFree(ptr) + assert err == cudart.cudaError_t.cudaSuccess