From 55e419022f76d993d7d7088015e9c950fe7d4081 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 | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/cpp/src/decisiontree/decisiontree.cuh b/cpp/src/decisiontree/decisiontree.cuh
index 50f8d8d3ac..e3cee7e5e7 100644
--- a/cpp/src/decisiontree/decisiontree.cuh
+++ b/cpp/src/decisiontree/decisiontree.cuh
@@ -65,6 +65,19 @@ inline bool is_dev_ptr(const void* p)
   }
 }
 
+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 != nullptr ||
+            pointer_attr.type == cudaMemoryTypeUnregistered);
+  } else {
+    err = cudaGetLastError();
+    return false;
+  }
+}
+
 template <typename T>
 std::string to_string_high_precision(T x)
 {
@@ -355,7 +368,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");