From 56e5e622a950efce365562d8f17bb38c3608d980 Mon Sep 17 00:00:00 2001
From: Robert Maynard <rmaynard@nvidia.com>
Date: Tue, 12 Nov 2024 11:55:44 -0500
Subject: [PATCH] Don't presume pointers are mutually exclusive for
 device/host. (#6128)

A pointer can be usable on the device and host at the same time. We can't invert `is_dev_ptr()` to check that something is a host pointer.

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

cudaPointerGetAttributes attributes cudaMalloc ptr
  is_dev_ptr  -> 1
  is_host_ptr -> 0
  memory loc  -> device

cudaPointerGetAttributes attributes cudaMallocManaged cudaMemAttachGlobal ptr
  is_dev_ptr  -> 1
  is_host_ptr -> 1
  memory loc  -> managed

cudaPointerGetAttributes attributes cudaMallocManaged cudaMemAttachHost ptr
  is_dev_ptr  -> 1
  is_host_ptr -> 1
  memory loc  -> managed

cudaPointerGetAttributes attributes cudaMallocHost ptr
  is_dev_ptr  -> 1
  is_host_ptr -> 1
  memory loc  -> host
```

Authors:
  - Robert Maynard (https://github.com/robertmaynard)

Approvers:
  - Dante Gama Dessavre (https://github.com/dantegd)
  - William Hicks (https://github.com/wphicks)
  - Divye Gala (https://github.com/divyegala)

URL: https://github.com/rapidsai/cuml/pull/6128
---
 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");