HIP updates: managed_memory_resource and missing header (#272)

* Add missing include * Add hip version of managed_memory_resource * Update copyright * Extra includes and copyrights, formatting * Minor fixes for the HIP code's benefit. --------- Co-authored-by: Stewart Martin-Haugh [email protected] <[email protected]> Co-authored-by: Attila Krasznahorkay <[email protected]>
acts-project · Mar 7, 2024 · 5f2d20e · 5f2d20e
1 parent e8a9614
commit 5f2d20e
Show file tree

Hide file tree

Showing 9 changed files with 147 additions and 17 deletions.
diff --git a/core/include/vecmem/containers/impl/device_array.ipp b/core/include/vecmem/containers/impl/device_array.ipp
@@ -1,11 +1,16 @@
 /* VecMem project, part of the ACTS project (R&D line)
  *
- * (c) 2021 CERN for the benefit of the ACTS project
+ * (c) 2021-2024 CERN for the benefit of the ACTS project
  *
  * Mozilla Public License Version 2.0
  */
 #pragma once
 
+// HIP include(s)
+#if defined(__HIP_DEVICE_COMPILE__)
+#include <hip/hip_runtime.h>
+#endif
+
 // System include(s).
 #include <cassert>
 

diff --git a/core/include/vecmem/memory/impl/atomic.ipp b/core/include/vecmem/memory/impl/atomic.ipp
@@ -1,12 +1,17 @@
 /*
  * VecMem project, part of the ACTS project (R&D line)
  *
- * (c) 2021 CERN for the benefit of the ACTS project
+ * (c) 2021-2024 CERN for the benefit of the ACTS project
  *
  * Mozilla Public License Version 2.0
  */
 #pragma once
 
+// HIP include(s).
+#if defined(__HIP_DEVICE_COMPILE__)
+#include <hip/hip_runtime.h>
+#endif
+
 // SYCL include(s).
 #if defined(CL_SYCL_LANGUAGE_VERSION) || defined(SYCL_LANGUAGE_VERSION)
 #include <CL/sycl.hpp>

diff --git a/core/include/vecmem/memory/impl/device_atomic_ref.ipp b/core/include/vecmem/memory/impl/device_atomic_ref.ipp
@@ -1,12 +1,17 @@
 /*
  * VecMem project, part of the ACTS project (R&D line)
  *
- * (c) 2022-2023 CERN for the benefit of the ACTS project
+ * (c) 2022-2024 CERN for the benefit of the ACTS project
  *
  * Mozilla Public License Version 2.0
  */
 #pragma once
 
+// HIP include
+#if defined(__HIP_DEVICE_COMPILE__)
+#include <hip/hip_runtime.h>
+#endif
+
 // SYCL include(s).
 #if defined(CL_SYCL_LANGUAGE_VERSION) || defined(SYCL_LANGUAGE_VERSION)
 #include <CL/sycl.hpp>

diff --git a/hip/CMakeLists.txt b/hip/CMakeLists.txt
@@ -1,6 +1,6 @@
 # VecMem project, part of the ACTS project (R&D line)
 #
-# (c) 2021-2023 CERN for the benefit of the ACTS project
+# (c) 2021-2024 CERN for the benefit of the ACTS project
 #
 # Mozilla Public License Version 2.0
 
@@ -17,6 +17,8 @@ vecmem_add_library( vecmem_hip hip
    "src/memory/device_memory_resource.cpp"
    "include/vecmem/memory/hip/host_memory_resource.hpp"
    "src/memory/host_memory_resource.cpp"
+   "include/vecmem/memory/hip/managed_memory_resource.hpp"
+   "src/memory/managed_memory_resource.cpp"
    # Utilities.
    "include/vecmem/utils/hip/copy.hpp"
    "src/utils/hip/copy.cpp"

diff --git a/hip/include/vecmem/memory/hip/managed_memory_resource.hpp b/hip/include/vecmem/memory/hip/managed_memory_resource.hpp
@@ -0,0 +1,53 @@
+/*
+ * VecMem project, part of the ACTS project (R&D line)
+ *
+ * (c) 2024 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Local include(s).
+#include "vecmem/memory/memory_resource.hpp"
+#include "vecmem/vecmem_hip_export.hpp"
+
+namespace vecmem::hip {
+
+/**
+ * @brief Memory resource that wraps managed HIP allocation.
+ *
+ * This is an allocator-type memory resource that allocates managed HIP
+ * memory, which is accessible directly to devices as well as to the host.
+ */
+class managed_memory_resource final : public memory_resource {
+
+public:
+    /// Default constructor
+    VECMEM_HIP_EXPORT
+    managed_memory_resource();
+    /// Destructor
+    VECMEM_HIP_EXPORT
+    ~managed_memory_resource();
+
+private:
+    /// @name Function(s) implementing @c vecmem::memory_resource
+    /// @{
+
+    /// Allocate HIP managed memory
+    VECMEM_HIP_EXPORT
+    virtual void* do_allocate(std::size_t, std::size_t) override final;
+    /// De-allocate a previously allocated managed memory block
+    VECMEM_HIP_EXPORT
+    virtual void do_deallocate(void* p, std::size_t,
+                               std::size_t) override final;
+    /// Compares @c *this for equality with @c other
+    VECMEM_HIP_EXPORT
+    virtual bool do_is_equal(
+        const memory_resource& other) const noexcept override final;
+
+    /// @}
+
+};  // class managed_memory_resource
+
+}  // namespace vecmem::hip
diff --git a/hip/src/memory/managed_memory_resource.cpp b/hip/src/memory/managed_memory_resource.cpp
@@ -0,0 +1,55 @@
+/*
+ * VecMem project, part of the ACTS project (R&D line)
+ *
+ * (c) 2024 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+// Local include(s).
+#include "vecmem/memory/hip/managed_memory_resource.hpp"
+
+#include "../utils/hip_error_handling.hpp"
+#include "vecmem/utils/debug.hpp"
+
+// HIP include(s).
+#include <hip/hip_runtime_api.h>
+
+namespace vecmem::hip {
+
+managed_memory_resource::managed_memory_resource() = default;
+
+managed_memory_resource::~managed_memory_resource() = default;
+
+void *managed_memory_resource::do_allocate(std::size_t bytes, std::size_t) {
+
+    if (bytes == 0) {
+        return nullptr;
+    }
+
+    // Allocate the memory.
+    void *res = nullptr;
+    VECMEM_HIP_ERROR_CHECK(hipMallocManaged(&res, bytes));
+    VECMEM_DEBUG_MSG(2, "Allocated %ld bytes at %p", bytes, res);
+    return res;
+}
+
+void managed_memory_resource::do_deallocate(void *p, std::size_t, std::size_t) {
+
+    if (p == nullptr) {
+        return;
+    }
+
+    // Free the memory.
+    VECMEM_DEBUG_MSG(2, "De-allocating memory at %p", p);
+    VECMEM_HIP_ERROR_CHECK(hipFree(p));
+}
+
+bool managed_memory_resource::do_is_equal(
+    const memory_resource &other) const noexcept {
+
+    // The two are equal if they are of the same type.
+    return (dynamic_cast<const managed_memory_resource *>(&other) != nullptr);
+}
+
+}  // namespace vecmem::hip
diff --git a/tests/hip/test_hip_containers_kernels.hip b/tests/hip/test_hip_containers_kernels.hip
@@ -1,14 +1,10 @@
 /* VecMem project, part of the ACTS project (R&D line)
  *
- * (c) 2021-2023 CERN for the benefit of the ACTS project
+ * (c) 2021-2024 CERN for the benefit of the ACTS project
  *
  * Mozilla Public License Version 2.0
  */
 
-// HIP include(s). Note that this needs to come first, as it affects how
-// other headers later on include/see system headers like <cassert>.
-#include <hip/hip_runtime.h>
-
 // Local include(s).
 #include "../../hip/src/utils/hip_error_handling.hpp"
 #include "test_hip_containers_kernels.hpp"
@@ -21,6 +17,9 @@
 #include "vecmem/memory/device_atomic_ref.hpp"
 #include "vecmem/utils/tuple.hpp"
 
+// HIP include(s).
+#include <hip/hip_runtime.h>
+
 /// Kernel performing a linear transformation using the vector helper types
 __global__ void linearTransformKernel(
     vecmem::data::vector_view<const int> constants,

diff --git a/tests/hip/test_hip_edm_kernels.hip b/tests/hip/test_hip_edm_kernels.hip
@@ -5,9 +5,6 @@
  * Mozilla Public License Version 2.0
  */
 
-// HIP include(s).
-#include <hip/hip_runtime.h>
-
 // Local include(s).
 #include "../common/jagged_soa_container_helpers.hpp"
 #include "../common/simple_soa_container_helpers.hpp"
@@ -16,6 +13,9 @@
 // Project include(s).
 #include "../../hip/src/utils/hip_error_handling.hpp"
 
+// HIP include(s).
+#include <hip/hip_runtime.h>
+
 __global__ void hipSimpleFillKernel(
     vecmem::testing::simple_soa_container::view view) {
 

diff --git a/tests/hip/test_hip_memory_resources.cpp b/tests/hip/test_hip_memory_resources.cpp
@@ -11,24 +11,30 @@
 #include "../common/memory_resource_test_host_accessible.hpp"
 #include "vecmem/memory/hip/device_memory_resource.hpp"
 #include "vecmem/memory/hip/host_memory_resource.hpp"
+#include "vecmem/memory/hip/managed_memory_resource.hpp"
 
 // GoogleTest include(s).
 #include <gtest/gtest.h>
 
 // Memory resources.
 static vecmem::hip::device_memory_resource device_resource;
 static vecmem::hip::host_memory_resource host_resource;
+static vecmem::hip::managed_memory_resource managed_resource;
 
 // Instantiate the allocation tests on all of the resources.
-INSTANTIATE_TEST_SUITE_P(hip_memory_resource_tests, memory_resource_test_basic,
-                         testing::Values(&device_resource, &host_resource),
+INSTANTIATE_TEST_SUITE_P(hip_memory_resource_tests_basic,
+                         memory_resource_test_basic,
+                         testing::Values(&device_resource, &host_resource,
+                                         &managed_resource),
                          vecmem::testing::memory_resource_name_gen(
                              {{&device_resource, "device_resource"},
-                              {&host_resource, "host_resource"}}));
+                              {&host_resource, "host_resource"},
+                              {&managed_resource, "managed_resource"}}));
 
 // Instantiate the full test suite on the host-accessible memory resources.
 INSTANTIATE_TEST_SUITE_P(hip_host_accessible_memory_resource_tests,
                          memory_resource_test_host_accessible,
-                         testing::Values(&host_resource),
+                         testing::Values(&host_resource, &managed_resource),
                          vecmem::testing::memory_resource_name_gen(
-                             {{&host_resource, "host_resource"}}));
+                             {{&host_resource, "host_resource"},
+                              {&managed_resource, "managed_resource"}}));