From 9b76d366d2d971839d4997c437e2d20490d9d65e Mon Sep 17 00:00:00 2001
From: Robert Maynard <rmaynard@nvidia.com>
Date: Fri, 1 Nov 2024 13:47:08 -0400
Subject: [PATCH] Correct rmm tests for validity of device pointers (#1714)

The `is_host_memory` function has been updated to understand that `cudaMemoryTypeUnregistered` is returned when provided pointers allocated by `malloc` and other host side allocation functions.

The `is_device_memory` function has been restricted to report only when device pointer that is usable by the calling cuda context. For that reason the tests now also set the active cuda device for all calling threads.

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

Approvers:
  - Rong Ou (https://github.com/rongou)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: https://github.com/rapidsai/rmm/pull/1714
---
 .../mr/device/mr_ref_multithreaded_tests.cpp  | 19 +++++++++++++++++--
 tests/mr/device/test_utils.hpp                |  7 ++-----
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/tests/mr/device/mr_ref_multithreaded_tests.cpp b/tests/mr/device/mr_ref_multithreaded_tests.cpp
index 7d749efd1..944ba1807 100644
--- a/tests/mr/device/mr_ref_multithreaded_tests.cpp
+++ b/tests/mr/device/mr_ref_multithreaded_tests.cpp
@@ -109,8 +109,13 @@ TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRef_mt)
 {
   // single thread changes default resource, then multiple threads use it
   auto old = rmm::mr::set_current_device_resource_ref(this->ref);
+  test_get_current_device_resource_ref();
 
-  spawn([mr = this->ref]() {
+  int device;
+  RMM_CUDA_TRY(cudaGetDevice(&device));
+
+  spawn([device, mr = this->ref]() {
+    RMM_CUDA_TRY(cudaSetDevice(device));
     EXPECT_EQ(mr, rmm::mr::get_current_device_resource_ref());
     test_get_current_device_resource_ref();  // test allocating with the new default resource
   });
@@ -156,7 +161,17 @@ TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRefPerThread_mt)
   }
 }
 
-TEST_P(mr_ref_test_mt, Allocate) { spawn(test_various_allocations, this->ref); }
+TEST_P(mr_ref_test_mt, Allocate)
+{
+  int device;
+  RMM_CUDA_TRY(cudaGetDevice(&device));
+
+  auto mr = this->ref;
+  spawn([device, mr]() {
+    RMM_CUDA_TRY(cudaSetDevice(device));
+    test_various_allocations(mr);
+  });
+}
 
 TEST_P(mr_ref_test_mt, AllocateDefaultStream)
 {
diff --git a/tests/mr/device/test_utils.hpp b/tests/mr/device/test_utils.hpp
index 2b9513793..5b7ef197b 100644
--- a/tests/mr/device/test_utils.hpp
+++ b/tests/mr/device/test_utils.hpp
@@ -31,17 +31,14 @@ inline bool is_device_accessible_memory(void* ptr)
 {
   cudaPointerAttributes attributes{};
   if (cudaSuccess != cudaPointerGetAttributes(&attributes, ptr)) { return false; }
-  return (attributes.type == cudaMemoryTypeDevice) or (attributes.type == cudaMemoryTypeManaged) or
-         ((attributes.type == cudaMemoryTypeHost) and (attributes.devicePointer != nullptr)) or
-         ((attributes.type == cudaMemoryTypeUnregistered) and
-          (rmm::mr::detail::is_system_memory_supported(rmm::get_current_cuda_device())));
+  return attributes.devicePointer != nullptr;
 }
 
 inline bool is_host_memory(void* ptr)
 {
   cudaPointerAttributes attributes{};
   if (cudaSuccess != cudaPointerGetAttributes(&attributes, ptr)) { return false; }
-  return attributes.type == cudaMemoryTypeHost;
+  return attributes.hostPointer != nullptr || attributes.type == cudaMemoryTypeUnregistered;
 }
 
 inline bool is_properly_aligned(void* ptr)