From aafa18afe987741746388d78b79742c90fcc5461 Mon Sep 17 00:00:00 2001
From: Mark Harris <783069+harrism@users.noreply.github.com>
Date: Thu, 18 Jan 2024 21:17:03 +0000
Subject: [PATCH] Factor out mr test utilities.

---
 tests/mr/device/mr_ref_test.hpp | 28 +-----------------
 tests/mr/device/mr_test.hpp     | 28 +-----------------
 tests/mr/device/test_utils.hpp  | 50 +++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 54 deletions(-)
 create mode 100644 tests/mr/device/test_utils.hpp

diff --git a/tests/mr/device/mr_ref_test.hpp b/tests/mr/device/mr_ref_test.hpp
index b1af26934..9826c10be 100644
--- a/tests/mr/device/mr_ref_test.hpp
+++ b/tests/mr/device/mr_ref_test.hpp
@@ -17,6 +17,7 @@
 #pragma once
 
 #include "../../byte_literals.hpp"
+#include "test_utils.hpp"
 
 #include <rmm/aligned.hpp>
 #include <rmm/cuda_device.hpp>
@@ -35,8 +36,6 @@
 
 #include <gtest/gtest.h>
 
-#include <cuda_runtime_api.h>
-
 #include <cuda/memory_resource>
 
 #include <cstddef>
@@ -50,31 +49,6 @@ using async_resource_ref = cuda::mr::async_resource_ref<cuda::mr::device_accessi
 
 namespace rmm::test {
 
-/**
- * @brief Returns if a pointer points to a device memory or managed memory
- * allocation.
- */
-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));
-}
-
-inline bool is_host_memory(void* ptr)
-{
-  cudaPointerAttributes attributes{};
-  if (cudaSuccess != cudaPointerGetAttributes(&attributes, ptr)) { return false; }
-  return attributes.type == cudaMemoryTypeHost;
-}
-
-inline bool is_properly_aligned(void* ptr)
-{
-  if (is_host_memory(ptr)) { return rmm::is_pointer_aligned(ptr, rmm::RMM_DEFAULT_HOST_ALIGNMENT); }
-  return rmm::is_pointer_aligned(ptr, rmm::CUDA_ALLOCATION_ALIGNMENT);
-}
-
 enum size_in_bytes : size_t {};
 
 constexpr auto default_num_allocations{100};
diff --git a/tests/mr/device/mr_test.hpp b/tests/mr/device/mr_test.hpp
index 3c621df56..3808ec6f3 100644
--- a/tests/mr/device/mr_test.hpp
+++ b/tests/mr/device/mr_test.hpp
@@ -17,6 +17,7 @@
 #pragma once
 
 #include "../../byte_literals.hpp"
+#include "test_utils.hpp"
 
 #include <rmm/aligned.hpp>
 #include <rmm/cuda_device.hpp>
@@ -36,8 +37,6 @@
 
 #include <gtest/gtest.h>
 
-#include <cuda_runtime_api.h>
-
 #include <cstddef>
 #include <cstdint>
 #include <functional>
@@ -46,31 +45,6 @@
 
 namespace rmm::test {
 
-/**
- * @brief Returns if a pointer points to a device memory or managed memory
- * allocation.
- */
-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));
-}
-
-inline bool is_host_memory(void* ptr)
-{
-  cudaPointerAttributes attributes{};
-  if (cudaSuccess != cudaPointerGetAttributes(&attributes, ptr)) { return false; }
-  return attributes.type == cudaMemoryTypeHost;
-}
-
-inline bool is_properly_aligned(void* ptr)
-{
-  if (is_host_memory(ptr)) { return rmm::is_pointer_aligned(ptr, rmm::RMM_DEFAULT_HOST_ALIGNMENT); }
-  return rmm::is_pointer_aligned(ptr, rmm::CUDA_ALLOCATION_ALIGNMENT);
-}
-
 enum size_in_bytes : size_t {};
 
 constexpr auto default_num_allocations{100};
diff --git a/tests/mr/device/test_utils.hpp b/tests/mr/device/test_utils.hpp
new file mode 100644
index 000000000..932a72a7e
--- /dev/null
+++ b/tests/mr/device/test_utils.hpp
@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <rmm/aligned.hpp>
+
+#include <cuda_runtime_api.h>
+
+namespace rmm::test {
+
+/**
+ * @brief Returns if a pointer points to a device memory or managed memory
+ * allocation.
+ */
+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));
+}
+
+inline bool is_host_memory(void* ptr)
+{
+  cudaPointerAttributes attributes{};
+  if (cudaSuccess != cudaPointerGetAttributes(&attributes, ptr)) { return false; }
+  return attributes.type == cudaMemoryTypeHost;
+}
+
+inline bool is_properly_aligned(void* ptr)
+{
+  if (is_host_memory(ptr)) { return rmm::is_pointer_aligned(ptr, rmm::RMM_DEFAULT_HOST_ALIGNMENT); }
+  return rmm::is_pointer_aligned(ptr, rmm::CUDA_ALLOCATION_ALIGNMENT);
+}
+
+}  // namespace rmm::test