From 24418475774958df2a0381e0c3c73a712e01881c Mon Sep 17 00:00:00 2001
From: Mike S Wang <32841762+MikeSWang@users.noreply.github.com>
Date: Thu, 12 Sep 2024 14:26:58 +0100
Subject: [PATCH] FIX(cpp): (#75) Fix memory management in hipFFT array
 copying/implementation

---
 src/triumvirate/include/arrayops.hpp | 24 ++++++++++++++++++++----
 src/triumvirate/src/arrayops.cpp     | 24 ++++++++++++++++--------
 src/triumvirate/src/fftlog.cpp       |  4 ++--
 3 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/src/triumvirate/include/arrayops.hpp b/src/triumvirate/include/arrayops.hpp
index d82b4319..333dc802 100644
--- a/src/triumvirate/include/arrayops.hpp
+++ b/src/triumvirate/include/arrayops.hpp
@@ -239,12 +239,28 @@ std::vector<int> get_sorted_indices(std::vector<int> sorting_vector);
 // ***********************************************************************
 
 #ifdef TRV_USE_HIP
-void copy_array_value_dtoh(
-  const hipDoubleComplex* hiparr, fftw_complex* arr, size_t length
+/**
+ * @brief Copy complex array values from device to host with different
+ *        type definitions.
+ *
+ * @param d_arr Device array.
+ * @param arr Default-type array on host.
+ * @param length Array size.
+ */
+void copy_complex_array_dtoh(
+  const hipDoubleComplex* d_arr, fftw_complex* arr, size_t length
 );
 
-void copy_array_value_htod(
-  fftw_complex* arr, const hipDoubleComplex* hiparr, size_t length
+/**
+ * @brief Copy complex array values from host to device with different
+ *        type definitions.
+ *
+ * @param arr Default-type array on host.
+ * @param d_arr Device array.
+ * @param length Array size.
+ */
+void copy_complex_array_htod(
+  const fftw_complex* arr, hipDoubleComplex* d_arr, size_t length
 );
 #endif
 
diff --git a/src/triumvirate/src/arrayops.cpp b/src/triumvirate/src/arrayops.cpp
index e7835d7e..da4faad1 100644
--- a/src/triumvirate/src/arrayops.cpp
+++ b/src/triumvirate/src/arrayops.cpp
@@ -271,22 +271,30 @@ std::vector<int> get_sorted_indices(std::vector<int> sorting_vector) {
 // ***********************************************************************
 
 #ifdef TRV_USE_HIP
-void copy_array_value_dtoh(
-  const hipDoubleComplex* hiparr, fftw_complex* arr, size_t length
+void copy_complex_array_dtoh(
+  const hipDoubleComplex* d_arr, fftw_complex* arr, size_t length
 ) {
+  hipDoubleComplex* h_arr = new hipDoubleComplex[length];
+  hipMemcpy(
+    h_arr, d_arr, sizeof(hipDoubleComplex) * length, hipMemcpyDeviceToHost
+  );
   for (size_t i = 0; i < length; ++i) {
-    arr[i][0] = hiparr[i].x;
-    arr[i][1] = hiparr[i].y;
+    arr[i][0] = h_arr[i].x;
+    arr[i][1] = h_arr[i].y;
   }
 }
 
-void copy_array_value_htod(
-  fftw_complex* arr, const hipDoubleComplex* hiparr, size_t length
+void copy_complex_array_htod(
+  const fftw_complex* arr, hipDoubleComplex* d_arr, size_t length
 ) {
+  hipDoubleComplex* h_arr = new hipDoubleComplex[length];
   for (size_t i = 0; i < length; ++i) {
-    hiparr[i].x = arr[i][0];
-    hiparr[i].y = arr[i][1];
+    h_arr[i].x = arr[i][0];
+    h_arr[i].y = arr[i][1];
   }
+  hipMemcpy(
+    d_arr, h_arr, sizeof(hipDoubleComplex) * length, hipMemcpyHostToDevice
+  );
 }
 #endif
 
diff --git a/src/triumvirate/src/fftlog.cpp b/src/triumvirate/src/fftlog.cpp
index 064c79db..3ee8f2bd 100644
--- a/src/triumvirate/src/fftlog.cpp
+++ b/src/triumvirate/src/fftlog.cpp
@@ -411,7 +411,7 @@ void HankelTransform::biased_transform(
   fftw_execute(this->pre_plan);
 #else  // TRV_USE_HIP
   hipDoubleComplex* d_pre_buffer;
-  hipMallocManaged(
+  hipMalloc(
     &d_pre_buffer, sizeof(hipDoubleComplex) * this->nsamp_trans
   );
   trva::copy_array_value_htod(
@@ -437,7 +437,7 @@ void HankelTransform::biased_transform(
   fftw_execute(this->post_plan);
 #else  // TRV_USE_HIP
   hipDoubleComplex* d_post_buffer;
-  hipMallocManaged(
+  hipMalloc(
     &d_post_buffer, sizeof(hipDoubleComplex) * this->nsamp_trans
   );
   trva::copy_array_value_htod(