Add internal cuda::__is_device_or_managed_memory

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<bool>(::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.
  *

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<Pointer, const void*> && !cuda::std::is_same_v<Pointer, void*>)
   {
     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);