Reduced number of created iota and map_back kernels

oleksandr-pavlyk · oleksandr-pavlyk · commit 869faef6964b · 2025-01-08T11:06:31.000-06:00
Avoid using comparator type to form kernel name types for
iota and map_back kernels (as they do not depedent on
comparator). This reduces the number of kernels generated
during instantiation of template implementation functions.
diff --git a/dpctl/tensor/libtensor/include/kernels/sorting/merge_sort.hpp b/dpctl/tensor/libtensor/include/kernels/sorting/merge_sort.hpp
@@ -211,8 +211,6 @@ void merge_impl(const std::size_t offset,
     }
 }
 
-namespace
-{
 template <typename Iter, typename Compare>
 void insertion_sort_impl(Iter first,
                          const std::size_t begin,
@@ -259,7 +257,6 @@ void leaf_sort_impl(Iter first,
     return insertion_sort_impl<Iter, Compare>(
         std::move(first), std::move(begin), std::move(end), std::move(comp));
 }
-} // namespace
 
 template <typename Iter> struct GetValueType
 {
@@ -768,9 +765,9 @@ sycl::event stable_sort_axis1_contig_impl(
     }
 }
 
-template <typename T1, typename T2, typename T3> class populate_index_data_krn;
+template <typename T1, typename T2> class populate_index_data_krn;
 
-template <typename T1, typename T2, typename T3> class index_map_to_rows_krn;
+template <typename T1, typename T2> class index_map_to_rows_krn;
 
 template <typename IndexT, typename ValueT, typename ValueComp> struct IndexComp
 {
@@ -820,7 +817,7 @@ sycl::event stable_argsort_axis1_contig_impl(
 
     using dpctl::tensor::kernels::sort_utils_detail::iota_impl;
 
-    using IotaKernelName = populate_index_data_krn<argTy, IndexTy, ValueComp>;
+    using IotaKernelName = populate_index_data_krn<argTy, IndexTy>;
 
     sycl::event populate_indexed_data_ev = iota_impl<IotaKernelName, IndexTy>(
         exec_q, res_tp, total_nelems, depends);
@@ -838,7 +835,7 @@ sycl::event stable_argsort_axis1_contig_impl(
         exec_q, iter_nelems, sort_nelems, res_tp, index_comp, sorted_block_size,
         {base_sort_ev});
 
-    using MapBackKernelName = index_map_to_rows_krn<argTy, IndexTy, ValueComp>;
+    using MapBackKernelName = index_map_to_rows_krn<argTy, IndexTy>;
     using dpctl::tensor::kernels::sort_utils_detail::map_back_impl;
 
     sycl::event write_out_ev = map_back_impl<MapBackKernelName, IndexTy>(
diff --git a/dpctl/tensor/libtensor/include/kernels/sorting/radix_sort.hpp b/dpctl/tensor/libtensor/include/kernels/sorting/radix_sort.hpp
@@ -1759,6 +1759,8 @@ template <typename ValueT, typename IndexT> struct ValueProj
 
 template <typename IndexT, typename ValueT, typename ProjT> struct IndexedProj
 {
+    IndexedProj(const ValueT *arg_ptr) : ptr(arg_ptr), value_projector{} {}
+
     IndexedProj(const ValueT *arg_ptr, const ProjT &proj_op)
         : ptr(arg_ptr), value_projector(proj_op)
     {
@@ -1848,7 +1850,7 @@ radix_argsort_axis1_contig_impl(sycl::queue &exec_q,
     using IdentityProjT = radix_sort_details::IdentityProj;
     using IndexedProjT =
         radix_sort_details::IndexedProj<IndexTy, argTy, IdentityProjT>;
-    const IndexedProjT proj_op{arg_tp, IdentityProjT{}};
+    const IndexedProjT proj_op{arg_tp};
 
     using IotaKernelName = radix_argsort_iota_krn<argTy, IndexTy>;
 
diff --git a/dpctl/tensor/libtensor/include/kernels/sorting/topk.hpp b/dpctl/tensor/libtensor/include/kernels/sorting/topk.hpp
@@ -32,6 +32,7 @@
 #include <stdexcept>
 #include <vector>
 
+#include <sycl/ext/oneapi/sub_group_mask.hpp>
 #include <sycl/sycl.hpp>
 
 #include "kernels/dpctl_tensor_types.hpp"
@@ -40,7 +41,6 @@
 #include "kernels/sorting/search_sorted_detail.hpp"
 #include "kernels/sorting/sort_utils.hpp"
 #include "utils/sycl_alloc_utils.hpp"
-#include <sycl/ext/oneapi/sub_group_mask.hpp>
 
 namespace dpctl
 {
@@ -134,11 +134,9 @@ sycl::event write_out_impl(sycl::queue &exec_q,
 
 } // namespace topk_detail
 
-template <typename T1, typename T2, typename T3>
-class topk_populate_index_data_krn;
+template <typename T1, typename T2> class topk_populate_index_data_krn;
 
-template <typename T1, typename T2, typename T3>
-class topk_full_merge_map_back_krn;
+template <typename T1, typename T2> class topk_full_merge_map_back_krn;
 
 template <typename argTy, typename IndexTy, typename CompT>
 sycl::event
@@ -158,7 +156,7 @@ topk_full_merge_sort_impl(sycl::queue &exec_q,
     // extract USM pointer
     IndexTy *index_data = index_data_owner.get();
 
-    using IotaKernelName = topk_populate_index_data_krn<argTy, IndexTy, CompT>;
+    using IotaKernelName = topk_populate_index_data_krn<argTy, IndexTy>;
 
     using dpctl::tensor::kernels::sort_utils_detail::iota_impl;
 
@@ -179,8 +177,7 @@ topk_full_merge_sort_impl(sycl::queue &exec_q,
         exec_q, iter_nelems, axis_nelems, index_data, comp, sorted_block_size,
         {base_sort_ev});
 
-    using WriteOutKernelName =
-        topk_full_merge_map_back_krn<argTy, IndexTy, CompT>;
+    using WriteOutKernelName = topk_full_merge_map_back_krn<argTy, IndexTy>;
 
     sycl::event write_out_ev =
         topk_detail::write_out_impl<WriteOutKernelName, argTy, IndexTy>(
@@ -194,8 +191,7 @@ topk_full_merge_sort_impl(sycl::queue &exec_q,
     return cleanup_host_task_event;
 };
 
-template <typename T1, typename T2, typename T3>
-class topk_partial_merge_map_back_krn;
+template <typename T1, typename T2> class topk_partial_merge_map_back_krn;
 
 template <typename T1, typename T2, typename Comp>
 class topk_over_work_group_krn;
@@ -213,24 +209,15 @@ sycl::event topk_merge_impl(
     const char *arg_cp,
     char *vals_cp,
     char *inds_cp,
-    dpctl::tensor::ssize_t iter_arg_offset,
-    dpctl::tensor::ssize_t iter_vals_offset,
-    dpctl::tensor::ssize_t iter_inds_offset,
-    dpctl::tensor::ssize_t axis_arg_offset,
-    dpctl::tensor::ssize_t axis_vals_offset,
-    dpctl::tensor::ssize_t axis_inds_offset,
     const std::vector<sycl::event> &depends)
 {
     if (axis_nelems < k) {
         throw std::runtime_error("Invalid sort axis size for value of k");
     }
 
-    const argTy *arg_tp = reinterpret_cast<const argTy *>(arg_cp) +
-                          iter_arg_offset + axis_arg_offset;
-    argTy *vals_tp = reinterpret_cast<argTy *>(vals_cp) + iter_vals_offset +
-                     axis_vals_offset;
-    IndexTy *inds_tp = reinterpret_cast<IndexTy *>(inds_cp) + iter_inds_offset +
-                       axis_inds_offset;
+    const argTy *arg_tp = reinterpret_cast<const argTy *>(arg_cp);
+    argTy *vals_tp = reinterpret_cast<argTy *>(vals_cp);
+    IndexTy *inds_tp = reinterpret_cast<IndexTy *>(inds_cp);
 
     using dpctl::tensor::kernels::IndexComp;
     const IndexComp<IndexTy, argTy, ValueComp> index_comp{arg_tp, ValueComp{}};
@@ -434,7 +421,7 @@ sycl::event topk_merge_impl(
 
         // Write out top k of the merge-sorted memory
         using WriteOutKernelName =
-            topk_partial_merge_map_back_krn<argTy, IndexTy, ValueComp>;
+            topk_partial_merge_map_back_krn<argTy, IndexTy>;
 
         sycl::event write_topk_ev =
             topk_detail::write_out_impl<WriteOutKernelName, argTy, IndexTy>(
@@ -462,24 +449,15 @@ sycl::event topk_radix_impl(sycl::queue &exec_q,
                             const char *arg_cp,
                             char *vals_cp,
                             char *inds_cp,
-                            dpctl::tensor::ssize_t iter_arg_offset,
-                            dpctl::tensor::ssize_t iter_vals_offset,
-                            dpctl::tensor::ssize_t iter_inds_offset,
-                            dpctl::tensor::ssize_t axis_arg_offset,
-                            dpctl::tensor::ssize_t axis_vals_offset,
-                            dpctl::tensor::ssize_t axis_inds_offset,
                             const std::vector<sycl::event> &depends)
 {
     if (axis_nelems < k) {
         throw std::runtime_error("Invalid sort axis size for value of k");
     }
 
-    const argTy *arg_tp = reinterpret_cast<const argTy *>(arg_cp) +
-                          iter_arg_offset + axis_arg_offset;
-    argTy *vals_tp = reinterpret_cast<argTy *>(vals_cp) + iter_vals_offset +
-                     axis_vals_offset;
-    IndexTy *inds_tp = reinterpret_cast<IndexTy *>(inds_cp) + iter_inds_offset +
-                       axis_inds_offset;
+    const argTy *arg_tp = reinterpret_cast<const argTy *>(arg_cp);
+    argTy *vals_tp = reinterpret_cast<argTy *>(vals_cp);
+    IndexTy *inds_tp = reinterpret_cast<IndexTy *>(inds_cp);
 
     const std::size_t total_nelems = iter_nelems * axis_nelems;
     const std::size_t padded_total_nelems = ((total_nelems + 63) / 64) * 64;
@@ -494,7 +472,7 @@ sycl::event topk_radix_impl(sycl::queue &exec_q,
     using IdentityProjT = radix_sort_details::IdentityProj;
     using IndexedProjT =
         radix_sort_details::IndexedProj<IndexTy, argTy, IdentityProjT>;
-    const IndexedProjT proj_op{arg_tp, IdentityProjT{}};
+    const IndexedProjT proj_op{arg_tp};
 
     using IotaKernelName = topk_iota_krn<argTy, IndexTy>;
 
diff --git a/dpctl/tensor/libtensor/source/sorting/topk.cpp b/dpctl/tensor/libtensor/source/sorting/topk.cpp
@@ -46,15 +46,15 @@
 #include "rich_comparisons.hpp"
 #include "topk.hpp"
 
-namespace td_ns = dpctl::tensor::type_dispatch;
-
 namespace dpctl
 {
 namespace tensor
 {
 namespace py_internal
 {
 
+namespace td_ns = dpctl::tensor::type_dispatch;
+
 typedef sycl::event (*topk_impl_fn_ptr_t)(sycl::queue &,
                                           std::size_t,
                                           std::size_t,
@@ -63,12 +63,6 @@ typedef sycl::event (*topk_impl_fn_ptr_t)(sycl::queue &,
                                           const char *,
                                           char *,
                                           char *,
-                                          py::ssize_t,
-                                          py::ssize_t,
-                                          py::ssize_t,
-                                          py::ssize_t,
-                                          py::ssize_t,
-                                          py::ssize_t,
                                           const std::vector<sycl::event> &);
 
 static topk_impl_fn_ptr_t topk_dispatch_vector[td_ns::num_types];
@@ -102,21 +96,14 @@ sycl::event topk_caller(sycl::queue &exec_q,
                         const char *arg_cp,
                         char *vals_cp,
                         char *inds_cp,
-                        py::ssize_t iter_arg_offset,
-                        py::ssize_t iter_vals_offset,
-                        py::ssize_t iter_inds_offset,
-                        py::ssize_t axis_arg_offset,
-                        py::ssize_t axis_vals_offset,
-                        py::ssize_t axis_inds_offset,
                         const std::vector<sycl::event> &depends)
 {
     if constexpr (use_radix_sort<argTy>::value) {
         using dpctl::tensor::kernels::topk_radix_impl;
         auto ascending = !largest;
-        return topk_radix_impl<argTy, IndexTy>(
-            exec_q, iter_nelems, axis_nelems, k, ascending, arg_cp, vals_cp,
-            inds_cp, iter_arg_offset, iter_vals_offset, iter_inds_offset,
-            axis_arg_offset, axis_vals_offset, axis_inds_offset, depends);
+        return topk_radix_impl<argTy, IndexTy>(exec_q, iter_nelems, axis_nelems,
+                                               k, ascending, arg_cp, vals_cp,
+                                               inds_cp, depends);
     }
     else {
         using dpctl::tensor::kernels::topk_merge_impl;
@@ -126,16 +113,14 @@ sycl::event topk_caller(sycl::queue &exec_q,
                     argTy>::type;
             return topk_merge_impl<argTy, IndexTy, CompTy>(
                 exec_q, iter_nelems, axis_nelems, k, arg_cp, vals_cp, inds_cp,
-                iter_arg_offset, iter_vals_offset, iter_inds_offset,
-                axis_arg_offset, axis_vals_offset, axis_inds_offset, depends);
+                depends);
         }
         else {
             using CompTy = typename dpctl::tensor::py_internal::AscendingSorter<
                 argTy>::type;
             return topk_merge_impl<argTy, IndexTy, CompTy>(
                 exec_q, iter_nelems, axis_nelems, k, arg_cp, vals_cp, inds_cp,
-                iter_arg_offset, iter_vals_offset, iter_inds_offset,
-                axis_arg_offset, axis_vals_offset, axis_inds_offset, depends);
+                depends);
         }
     }
 }
@@ -268,14 +253,11 @@ py_topk(const dpctl::tensor::usm_ndarray &src,
     bool is_inds_c_contig = inds.is_c_contiguous();
 
     if (is_src_c_contig && is_vals_c_contig && is_inds_c_contig) {
-        static constexpr py::ssize_t zero_offset = py::ssize_t(0);
-
         auto fn = topk_dispatch_vector[src_typeid];
 
         sycl::event comp_ev =
             fn(exec_q, iter_nelems, axis_nelems, k, largest, src.get_data(),
-               vals.get_data(), inds.get_data(), zero_offset, zero_offset,
-               zero_offset, zero_offset, zero_offset, zero_offset, depends);
+               vals.get_data(), inds.get_data(), depends);
 
         sycl::event keep_args_alive_ev =
             dpctl::utils::keep_args_alive(exec_q, {src, vals, inds}, {comp_ev});

Original file line number	Diff line number	Diff line change
`@@ -211,8 +211,6 @@ void merge_impl(const std::size_t offset,`
`211`	`211`	`}`
`212`	`212`	`}`
`213`	`213`
`214`		`-namespace`
`215`		`-{`
`216`	`214`	`template <typename Iter, typename Compare>`
`217`	`215`	`void insertion_sort_impl(Iter first,`
`218`	`216`	`const std::size_t begin,`
`@@ -259,7 +257,6 @@ void leaf_sort_impl(Iter first,`
`259`	`257`	`return insertion_sort_impl<Iter, Compare>(`
`260`	`258`	`std::move(first), std::move(begin), std::move(end), std::move(comp));`
`261`	`259`	`}`
`262`		`-} // namespace`
`263`	`260`
`264`	`261`	`template <typename Iter> struct GetValueType`
`265`	`262`	`{`
`@@ -768,9 +765,9 @@ sycl::event stable_sort_axis1_contig_impl(`
`768`	`765`	`}`
`769`	`766`	`}`
`770`	`767`
`771`		`-template <typename T1, typename T2, typename T3> class populate_index_data_krn;`
	`768`	`+template <typename T1, typename T2> class populate_index_data_krn;`
`772`	`769`
`773`		`-template <typename T1, typename T2, typename T3> class index_map_to_rows_krn;`
	`770`	`+template <typename T1, typename T2> class index_map_to_rows_krn;`
`774`	`771`
`775`	`772`	`template <typename IndexT, typename ValueT, typename ValueComp> struct IndexComp`
`776`	`773`	`{`
`@@ -820,7 +817,7 @@ sycl::event stable_argsort_axis1_contig_impl(`
`820`	`817`
`821`	`818`	`using dpctl::tensor::kernels::sort_utils_detail::iota_impl;`
`822`	`819`
`823`		`- using IotaKernelName = populate_index_data_krn<argTy, IndexTy, ValueComp>;`
	`820`	`+ using IotaKernelName = populate_index_data_krn<argTy, IndexTy>;`
`824`	`821`
`825`	`822`	`sycl::event populate_indexed_data_ev = iota_impl<IotaKernelName, IndexTy>(`
`826`	`823`	`exec_q, res_tp, total_nelems, depends);`
`@@ -838,7 +835,7 @@ sycl::event stable_argsort_axis1_contig_impl(`
`838`	`835`	`exec_q, iter_nelems, sort_nelems, res_tp, index_comp, sorted_block_size,`
`839`	`836`	`{base_sort_ev});`
`840`	`837`
`841`		`- using MapBackKernelName = index_map_to_rows_krn<argTy, IndexTy, ValueComp>;`
	`838`	`+ using MapBackKernelName = index_map_to_rows_krn<argTy, IndexTy>;`
`842`	`839`	`using dpctl::tensor::kernels::sort_utils_detail::map_back_impl;`
`843`	`840`
`844`	`841`	`sycl::event write_out_ev = map_back_impl<MapBackKernelName, IndexTy>(`
Original file line number	Diff line number	Diff line change
`@@ -1759,6 +1759,8 @@ template <typename ValueT, typename IndexT> struct ValueProj`
`1759`	`1759`
`1760`	`1760`	`template <typename IndexT, typename ValueT, typename ProjT> struct IndexedProj`
`1761`	`1761`	`{`
	`1762`	`+ IndexedProj(const ValueT *arg_ptr) : ptr(arg_ptr), value_projector{} {}`
	`1763`	`+`
`1762`	`1764`	`IndexedProj(const ValueT *arg_ptr, const ProjT &proj_op)`
`1763`	`1765`	`: ptr(arg_ptr), value_projector(proj_op)`
`1764`	`1766`	`{`
`@@ -1848,7 +1850,7 @@ radix_argsort_axis1_contig_impl(sycl::queue &exec_q,`
`1848`	`1850`	`using IdentityProjT = radix_sort_details::IdentityProj;`
`1849`	`1851`	`using IndexedProjT =`
`1850`	`1852`	`radix_sort_details::IndexedProj<IndexTy, argTy, IdentityProjT>;`
`1851`		`- const IndexedProjT proj_op{arg_tp, IdentityProjT{}};`
	`1853`	`+ const IndexedProjT proj_op{arg_tp};`
`1852`	`1854`
`1853`	`1855`	`using IotaKernelName = radix_argsort_iota_krn<argTy, IndexTy>;`
`1854`	`1856`