Use unique_ptr as temporary owner of USM allocation

Until it is passed over to the host function, and
unique_ptr's ownership is released.

Also reduced allocation sizes, where too much was being
allocated.

Author: oleksandr-pavlyk

dpctl/tensor/libtensor/include/kernels/sorting/radix_sort.hpp

@@ -30,6 +30,7 @@
 #include <array>
 #include <cstdint>
 #include <limits>
+#include <memory>
 #include <stdexcept>
 #include <type_traits>
 #include <utility>
@@ -1590,10 +1591,13 @@ sycl::event parallel_radix_sort_impl(sycl::queue &exec_q,
         // memory for storing count and offset values
         CountT *count_ptr =
             sycl::malloc_device<CountT>(n_iters * n_counts, exec_q);
+
         if (nullptr == count_ptr) {
             throw std::runtime_error("Could not allocate USM-device memory");
         }
 
+        auto count_owner =
+            dpctl::tensor::alloc_utils::make_owner(count_ptr, exec_q);
         constexpr std::uint32_t zero_radix_iter{0};
 
         if constexpr (std::is_same_v<KeyT, bool>) {
@@ -1620,11 +1624,12 @@ sycl::event parallel_radix_sort_impl(sycl::queue &exec_q,
         ValueT *tmp_arr =
             sycl::malloc_device<ValueT>(n_iters * n_to_sort, exec_q);
         if (nullptr == tmp_arr) {
-            using dpctl::tensor::alloc_utils::sycl_free_noexcept;
-            sycl_free_noexcept(count_ptr, exec_q);
             throw std::runtime_error("Could not allocate USM-device memory");
         }
 
+        auto tmp_arr_owner =
+            dpctl::tensor::alloc_utils::make_owner(tmp_arr, exec_q);
+
         // iterations per each bucket
         assert("Number of iterations must be even" && radix_iters % 2 == 0);
         assert(radix_iters > 0);
@@ -1668,6 +1673,8 @@ sycl::event parallel_radix_sort_impl(sycl::queue &exec_q,
                 sycl_free_noexcept(count_ptr, ctx);
             });
         });
+        count_owner.release();
+        tmp_arr_owner.release();
     }
 
     return sort_ev;
@@ -1769,13 +1776,13 @@ radix_argsort_axis1_contig_impl(sycl::queue &exec_q,
         reinterpret_cast<IndexTy *>(res_cp) + iter_res_offset + sort_res_offset;
 
     const std::size_t total_nelems = iter_nelems * sort_nelems;
-    const std::size_t padded_total_nelems = ((total_nelems + 63) / 64) * 64;
-    IndexTy *workspace = sycl::malloc_device<IndexTy>(
-        padded_total_nelems + total_nelems, exec_q);
+    IndexTy *workspace = sycl::malloc_device<IndexTy>(total_nelems, exec_q);
 
     if (nullptr == workspace) {
         throw std::runtime_error("Could not allocate workspace on device");
     }
+    auto workspace_owner =
+        dpctl::tensor::alloc_utils::make_owner(workspace, exec_q);
 
     using IdentityProjT = radix_sort_details::IdentityProj;
     using IndexedProjT =
@@ -1829,6 +1836,7 @@ radix_argsort_axis1_contig_impl(sycl::queue &exec_q,
         cgh.host_task([ctx, workspace] { sycl_free_noexcept(workspace, ctx); });
     });
 
+    workspace_owner.release();
     return cleanup_ev;
 }
 

dpctl/tensor/libtensor/include/kernels/sorting/topk.hpp

@@ -29,10 +29,12 @@
 #include <cstdint>
 #include <iterator>
 #include <limits>
+#include <memory>
 #include <stdexcept>
-#include <sycl/sycl.hpp>
 #include <vector>
 
+#include <sycl/sycl.hpp>
+
 #include "kernels/dpctl_tensor_types.hpp"
 #include "kernels/sorting/merge_sort.hpp"
 #include "kernels/sorting/radix_sort.hpp"
@@ -95,6 +97,8 @@ topk_full_merge_sort_impl(sycl::queue &exec_q,
     if (index_data == nullptr) {
         throw std::runtime_error("Unable to allocate device_memory");
     }
+    auto index_data_owner =
+        dpctl::tensor::alloc_utils::make_owner(index_data, exec_q);
 
     using IotaKernelName = topk_populate_index_data_krn<argTy, IndexTy, CompT>;
 
@@ -161,6 +165,7 @@ topk_full_merge_sort_impl(sycl::queue &exec_q,
                 [ctx, index_data] { sycl_free_noexcept(index_data, ctx); });
         });
 
+    index_data_owner.release();
     return cleanup_host_task_event;
 };
 
@@ -287,6 +292,8 @@ sycl::event topk_merge_impl(
         if (index_data == nullptr) {
             throw std::runtime_error("Unable to allocate device_memory");
         }
+        auto index_data_owner =
+            dpctl::tensor::alloc_utils::make_owner(index_data, exec_q);
 
         // no need to populate index data: SLM will be populated with default
         // values
@@ -434,6 +441,7 @@ sycl::event topk_merge_impl(
                     [ctx, index_data] { sycl_free_noexcept(index_data, ctx); });
             });
 
+        index_data_owner.release();
         return cleanup_host_task_event;
     }
 }
@@ -479,15 +487,10 @@ sycl::event topk_radix_impl(sycl::queue &exec_q,
         throw std::runtime_error(
             "Not enough device memory for radix sort topk");
     }
+    auto workspace_owner =
+        dpctl::tensor::alloc_utils::make_owner(workspace, exec_q);
 
-    IndexTy *tmp_tp = sycl::malloc_device<IndexTy>(total_nelems, exec_q);
-
-    if (nullptr == tmp_tp) {
-        using dpctl::tensor::alloc_utils::sycl_free_noexcept;
-        sycl_free_noexcept(workspace, exec_q);
-        throw std::runtime_error(
-            "Not enough device memory for radix sort topk");
-    }
+    IndexTy *tmp_tp = workspace + padded_total_nelems;
 
     using IdentityProjT = radix_sort_details::IdentityProj;
     using IndexedProjT =
@@ -546,12 +549,10 @@ sycl::event topk_radix_impl(sycl::queue &exec_q,
         const sycl::context &ctx = exec_q.get_context();
 
         using dpctl::tensor::alloc_utils::sycl_free_noexcept;
-        cgh.host_task([ctx, workspace, tmp_tp] {
-            sycl_free_noexcept(workspace, ctx);
-            sycl_free_noexcept(tmp_tp, ctx);
-        });
+        cgh.host_task([ctx, workspace] { sycl_free_noexcept(workspace, ctx); });
     });
 
+    workspace_owner.release();
     return cleanup_ev;
 }
 

dpctl/tensor/libtensor/include/utils/sycl_alloc_utils.hpp

@@ -28,6 +28,7 @@
 
 #include <exception>
 #include <iostream>
+#include <memory>
 
 #include "sycl/sycl.hpp"
 
@@ -78,6 +79,28 @@ template <typename T> void sycl_free_noexcept(T *ptr, sycl::queue &q) noexcept
     sycl_free_noexcept(ptr, q.get_context());
 }
 
+class USMDeleter
+{
+private:
+    sycl::context ctx_;
+
+public:
+    USMDeleter(const sycl::queue &q) : ctx_(q.get_context()) {}
+    USMDeleter(const sycl::context &ctx) : ctx_(ctx) {}
+
+    template <typename T> void operator()(T *ptr)
+    {
+        sycl_free_noexcept(ptr, ctx_);
+    }
+};
+
+template <typename T>
+std::unique_ptr<T, USMDeleter> make_owner(T *ptr, const sycl::queue &q)
+{
+    auto usm_deleter = USMDeleter(q);
+    return std::unique_ptr<T, USMDeleter>(ptr, usm_deleter);
+}
+
 } // end of namespace alloc_utils
 } // end of namespace tensor
 } // end of namespace dpctl