From 5bf3358b2dc5d8b6a2fa7aa841f14e44ec458ecc Mon Sep 17 00:00:00 2001
From: Robert Maynard <rmaynard@nvidia.com>
Date: Tue, 29 Oct 2024 16:43:29 -0400
Subject: [PATCH] Don't presume pointers are mutually exclusive for
 device/host.

Here is the results of looking at the cudaPointerGetAttributes
of different allocation types. As we can see things like
cudaMallocManaged allow the same pointer to be both host and device.
```
cudaPointerGetAttributes attributes integer
  is_dev_ptr -> 0
  is_host_ptr -> 0

cudaPointerGetAttributes attributes std::vector<int> data
  is_dev_ptr -> 0
  is_host_ptr -> 0

cudaPointerGetAttributes attributes malloc ptr
  is_dev_ptr -> 0
  is_host_ptr -> 0

cudaPointerGetAttributes attributes cudaMalloc ptr
  is_dev_ptr -> 1
  is_host_ptr -> 0

cudaPointerGetAttributes attributes cudaMallocManaged cudaMemAttachGlobal ptr
  is_dev_ptr -> 1
  is_host_ptr -> 1

cudaPointerGetAttributes attributes cudaMallocManaged cudaMemAttachHost ptr
  is_dev_ptr -> 1
  is_host_ptr -> 1

cudaPointerGetAttributes attributes cudaMallocHost ptr
  is_dev_ptr -> 1
  is_host_ptr -> 1
```
---
 cpp/src/decisiontree/decisiontree.cuh | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/cpp/src/decisiontree/decisiontree.cuh b/cpp/src/decisiontree/decisiontree.cuh
index 50f8d8d3ac..e5514ffa71 100644
--- a/cpp/src/decisiontree/decisiontree.cuh
+++ b/cpp/src/decisiontree/decisiontree.cuh
@@ -58,7 +58,19 @@ inline bool is_dev_ptr(const void* p)
   cudaPointerAttributes pointer_attr;
   cudaError_t err = cudaPointerGetAttributes(&pointer_attr, p);
   if (err == cudaSuccess) {
-    return pointer_attr.devicePointer;
+    return (pointer_attr.devicePointer || pointer_attr.type == cudaMemoryTypeDevice);
+  } else {
+    err = cudaGetLastError();
+    return false;
+  }
+}
+
+inline bool is_host_ptr(const void* p)
+{
+  cudaPointerAttributes pointer_attr;
+  cudaError_t err = cudaPointerGetAttributes(&pointer_attr, p);
+  if (err == cudaSuccess) {
+    return (pointer_attr.hostPointer || pointer_attr.type == cudaMemoryTypeUnregistered);
   } else {
     err = cudaGetLastError();
     return false;
@@ -355,7 +367,7 @@ class DecisionTree {
                       int verbosity)
   {
     if (verbosity >= 0) { ML::Logger::get().setLevel(verbosity); }
-    ASSERT(!is_dev_ptr(rows) && !is_dev_ptr(predictions),
+    ASSERT(is_host_ptr(rows) && is_host_ptr(predictions),
            "DT Error: Current impl. expects both input and predictions to be CPU "
            "pointers.\n");