From 109e50248dd70abdcb7e916feb8addef6f705b65 Mon Sep 17 00:00:00 2001
From: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>
Date: Wed, 8 Jan 2025 18:21:51 +0900
Subject: [PATCH] use kokkos_malloc instead of hipMalloc

---
 fft/src/KokkosFFT_ROCM_types.hpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/fft/src/KokkosFFT_ROCM_types.hpp b/fft/src/KokkosFFT_ROCM_types.hpp
index 05856d68..ddfaea1f 100644
--- a/fft/src/KokkosFFT_ROCM_types.hpp
+++ b/fft/src/KokkosFFT_ROCM_types.hpp
@@ -82,8 +82,7 @@ struct ScopedRocfftExecutionInfo {
   }
   ~ScopedRocfftExecutionInfo() noexcept {
     if (m_workbuffer != nullptr) {
-      hipError_t hip_status = hipFree(m_workbuffer);
-      if (hip_status != hipSuccess) Kokkos::abort("hipFree failed");
+      Kokkos::kokkos_free<Kokkos::HIP>(m_workbuffer);
     }
     rocfft_status status = rocfft_execution_info_destroy(m_execution_info);
     if (status != rocfft_status_success)
@@ -112,8 +111,9 @@ struct ScopedRocfftExecutionInfo {
 
     // Set work buffer
     if (workbuffersize > 0) {
-      hipError_t hip_status = hipMalloc(&m_workbuffer, workbuffersize);
-      KOKKOSFFT_THROW_IF(hip_status != hipSuccess, "hipMalloc failed");
+      m_workbuffer = Kokkos::kokkos_malloc<Kokkos::HIP>(
+          "kokkos_malloc workbuffer", workbuffersize);
+
       status = rocfft_execution_info_set_work_buffer(
           m_execution_info, m_workbuffer, workbuffersize);
       KOKKOSFFT_THROW_IF(status != rocfft_status_success,