diff --git a/libcudacxx/include/cuda/__memory/is_pointer_accessible.h b/libcudacxx/include/cuda/__memory/is_pointer_accessible.h index 371dca490a6..e6ccb2be8dc 100644 --- a/libcudacxx/include/cuda/__memory/is_pointer_accessible.h +++ b/libcudacxx/include/cuda/__memory/is_pointer_accessible.h @@ -100,6 +100,53 @@ _CCCL_HOST_API inline bool is_host_accessible(const void* __p) return false; } +/** + * @brief Checks if a pointer is a device pointer. + * + * This internal-only function can be used when the device id is not known. + * The main difference between this function and is_device_accessible() is that this function does not check if the + * pointer is peer accessible from a specified device. + * + * @param __p The pointer to check. + * @return `true` if the pointer is a device pointer, `false` otherwise. + */ +[[nodiscard]] +_CCCL_HOST_API inline bool __is_device_or_managed_memory(const void* __p) noexcept +{ + if (__p == nullptr) + { + return false; + } + ::CUpointer_attribute __attrs[4] = { + ::CU_POINTER_ATTRIBUTE_MEMORY_TYPE, + ::CU_POINTER_ATTRIBUTE_IS_MANAGED, + ::CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, + ::CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE}; + auto __memory_type = static_cast<::CUmemorytype>(0); + int __is_managed = 0; + int __ptr_dev_id = 0; + ::CUmemoryPool __mempool = nullptr; + void* __results[4] = {&__memory_type, &__is_managed, &__ptr_dev_id, &__mempool}; + const auto __status = ::cuda::__driver::__pointerGetAttributesNoThrow(__attrs, __results, __p); + if (__status != ::cudaSuccess) + { + return false; + } + // (1) check if the pointer is unregistered + if (__memory_type == static_cast<::CUmemorytype>(0)) + { + return false; + } + // (2) check if a memory pool is associated with the pointer + if (__mempool != nullptr) + { + ::CUmemLocation __prop{::CU_MEM_LOCATION_TYPE_DEVICE, __ptr_dev_id}; + return static_cast(::cuda::__driver::__mempoolGetAccess(__mempool, &__prop)); + } + // (3) check if the pointer is a device accessible pointer or managed memory + return __is_managed || __memory_type == ::CU_MEMORYTYPE_DEVICE; +} + /** * @brief Checks if a pointer is a device accessible pointer. * diff --git a/libcudacxx/test/libcudacxx/cuda/memory/is_pointer_accessible.pass.cpp b/libcudacxx/test/libcudacxx/cuda/memory/is_pointer_accessible.pass.cpp index f171b3b177c..960f9214ba6 100644 --- a/libcudacxx/test/libcudacxx/cuda/memory/is_pointer_accessible.pass.cpp +++ b/libcudacxx/test/libcudacxx/cuda/memory/is_pointer_accessible.pass.cpp @@ -28,11 +28,13 @@ void test_accessible_pointer( { assert(cuda::is_host_accessible(ptr) == is_host_accessible); assert(cuda::is_device_accessible(ptr, device) == is_device_accessible); + assert(cuda::__is_device_or_managed_memory(ptr) == is_device_accessible); assert(cuda::is_managed(ptr) == is_managed_accessible); if constexpr (!cuda::std::is_same_v && !cuda::std::is_same_v) { assert(cuda::is_host_accessible(ptr + 1) == is_host_accessible); assert(cuda::is_device_accessible(ptr + 1, device) == is_device_accessible); + assert(cuda::__is_device_or_managed_memory(ptr + 1) == is_device_accessible); assert(cuda::is_managed(ptr + 1) == is_managed_accessible); } } @@ -149,6 +151,7 @@ bool test_multiple_devices() /// DEVICE 1 CONTEXT cuda::__ensure_current_context ctx1(dev1); + assert(cuda::__is_device_or_managed_memory(device_ptr0) == true); assert(cuda::is_device_accessible(device_ptr0, dev0) == true); assert(cuda::is_device_accessible(device_ptr0, dev1) == false); @@ -158,6 +161,7 @@ bool test_multiple_devices() { return true; } + assert(cuda::__is_device_or_managed_memory(device_ptr0) == true); assert(cuda::is_device_accessible(device_ptr0, dev1) == false); assert(cudaDeviceEnablePeerAccess(dev1.get(), 0) == cudaSuccess); @@ -186,6 +190,7 @@ bool test_multiple_devices_from_pool() return true; } assert(cuda::is_device_accessible(ptr, dev1) == false); + assert(cuda::__is_device_or_managed_memory(ptr) == true); assert(cudaDeviceEnablePeerAccess(dev1.get(), 0) == cudaSuccess); assert(cuda::is_device_accessible(ptr, dev0) == true);