add for iamax and iamin

codeplaysoftware · Jan 17, 2024 · 2f5fe4d · 2f5fe4d
1 parent a1bbfc0
commit 2f5fe4d
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/src/interface/blas1_interface.hpp b/src/interface/blas1_interface.hpp
@@ -321,9 +321,10 @@ typename sb_handle_t::event_t _iamax_iamin_impl(
         localMemSize == 0
             ? _nWG * (static_cast<index_t>(localSize) / min_sg_size)
             : _nWG;
-    auto gpu_res = blas::helper::allocate < is_usm ? helper::AllocType::usm
-                                                   : helper::AllocType::buffer,
-         tuple_t > (memory_size, q);
+    auto gpu_res = sb_handle.template acquire_temp_mem < is_usm
+                       ? helper::AllocType::usm
+                       : helper::AllocType::buffer,
+         tuple_t > (memory_size);
     auto gpu_res_vec =
         make_vector_view(gpu_res, static_cast<increment_t>(1), memory_size);
     auto step0 = make_index_max_min<is_max, true>(gpu_res_vec, tupOp);
@@ -355,7 +356,7 @@ typename sb_handle_t::event_t _iamax_iamin_impl(
                                  static_cast<index_t>(localSize),
                                  static_cast<index_t>(localMemSize), ret));
     }
-    blas::helper::enqueue_deallocate(ret, gpu_res, q);
+    sb_handle.template release_temp_mem({*ret.rbegin()}, gpu_res);
   }
   return ret;
 }