From 030ae18c0ff55269663080027d989d1d578c0bb3 Mon Sep 17 00:00:00 2001
From: Bradley Dice <bdice@bradleydice.com>
Date: Thu, 15 Aug 2024 13:09:07 -0700
Subject: [PATCH] Make rmm::prefetch fault tolerant

---
 include/rmm/mr/device/prefetch_resource_adaptor.hpp |  5 +++++
 include/rmm/prefetch.hpp                            | 12 +++++++++---
 python/rmm/rmm/_lib/memory_resource.pyx             |  5 +++++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/include/rmm/mr/device/prefetch_resource_adaptor.hpp b/include/rmm/mr/device/prefetch_resource_adaptor.hpp
index 48a64a5b9..587b62979 100644
--- a/include/rmm/mr/device/prefetch_resource_adaptor.hpp
+++ b/include/rmm/mr/device/prefetch_resource_adaptor.hpp
@@ -33,6 +33,11 @@ namespace rmm::mr {
 /**
  * @brief Resource that prefetches all memory allocations.
  *
+ * Note that prefetching is a no-op if the upstream resource is not using
+ * managed memory or the device does not support managed memory. In those
+ * cases, this adapter does nothing. Therefore it is only recommended to use
+ * this adapter with a managed memory resource on a supported system.
+ *
  * @tparam Upstream Type of the upstream resource used for
  * allocation/deallocation.
  */
diff --git a/include/rmm/prefetch.hpp b/include/rmm/prefetch.hpp
index 205a4d241..d3274c298 100644
--- a/include/rmm/prefetch.hpp
+++ b/include/rmm/prefetch.hpp
@@ -34,6 +34,7 @@ namespace rmm {
  * @brief Prefetch memory to the specified device on the specified stream.
  *
  * This function is a no-op if the pointer is not to CUDA managed memory.
+ * This function is a no-op if the CUDA device does not support managed memory.
  *
  * @throw rmm::cuda_error if the prefetch fails.
  *
@@ -48,9 +49,14 @@ void prefetch(void const* ptr,
               rmm::cuda_stream_view stream)
 {
   auto result = cudaMemPrefetchAsync(ptr, size, device.value(), stream.value());
-  // InvalidValue error is raised when non-managed memory is passed to cudaMemPrefetchAsync
-  // We should treat this as a no-op
-  if (result != cudaErrorInvalidValue && result != cudaSuccess) { RMM_CUDA_TRY(result); }
+  // cudaErrorInvalidValue is raised when non-managed memory is passed to
+  // cudaMemPrefetchAsync. cudaErrorInvalidDevice is raised when the device
+  // does not support managed memory. We treat both cases as a no-op, and
+  // document that this function does nothing in those cases.
+  if (result != cudaSuccess && result != cudaErrorInvalidValue &&
+      result != cudaErrorInvalidDevice) {
+    RMM_CUDA_TRY(result);
+  }
 }
 
 /**
diff --git a/python/rmm/rmm/_lib/memory_resource.pyx b/python/rmm/rmm/_lib/memory_resource.pyx
index 5030c5d2d..d048d853f 100644
--- a/python/rmm/rmm/_lib/memory_resource.pyx
+++ b/python/rmm/rmm/_lib/memory_resource.pyx
@@ -1058,6 +1058,11 @@ cdef class PrefetchResourceAdaptor(UpstreamResourceAdaptor):
         """
         Memory resource that prefetches all allocations.
 
+        Note that prefetching is a no-op if the upstream resource is not using
+        managed memory or the device does not support managed memory. In those
+        cases, this adapter does nothing. Therefore it is only recommended to
+        use this adapter with a managed memory resource on a supported system.
+
         Parameters
         ----------
         upstream : DeviceMemoryResource