Introduce reduce_over_axis0 kernel for contiguous inputs

This achieves additional savings over the prior commit:

```
In [1]: import dpctl.tensor as dpt

In [2]: x = dpt.reshape(dpt.asarray(1, dtype="f4")/dpt.square(dpt.arange(1, 1282200*128 + 1, dtype="f4")), (1282200, 128))

In [3]: %time y = dpt.sum(x, axis=0)
CPU times: user 136 ms, sys: 9.52 ms, total: 145 ms
Wall time: 158 ms

In [4]: %time y = dpt.sum(x, axis=0)
CPU times: user 18.8 ms, sys: 17.3 ms, total: 36.1 ms
Wall time: 42 ms

In [5]: %time y = dpt.sum(x, axis=0)
CPU times: user 19.2 ms, sys: 16.9 ms, total: 36.1 ms
Wall time: 38.4 ms

In [6]: %time y = dpt.sum(x, axis=0)
CPU times: user 1.69 ms, sys: 35.2 ms, total: 36.9 ms
Wall time: 39.4 ms

In [7]: quit
```

Prior to this the wall time stood at 49 ms.

Author: oleksandr-pavlyk

dpctl/tensor/libtensor/include/kernels/reductions.hpp

@@ -234,7 +234,10 @@ template <typename T1, typename T2, typename T3, typename T4, typename T5>
 class sum_reduction_seq_contig_krn;
 
 template <typename T1, typename T2, typename T3, typename T4, typename T5>
-class sum_reduction_over_group_with_atomics_contig_krn;
+class sum_reduction_axis0_over_group_with_atomics_contig_krn;
+
+template <typename T1, typename T2, typename T3, typename T4, typename T5>
+class sum_reduction_axis1_over_group_with_atomics_contig_krn;
 
 using dpctl::tensor::sycl_utils::choose_workgroup_size;
 
@@ -390,7 +393,7 @@ typedef sycl::event (*sum_reduction_contig_impl_fn_ptr)(
 
 /* @brief Reduce rows in a matrix */
 template <typename argTy, typename resTy>
-sycl::event sum_reduction_over_group_with_atomics_contig_impl(
+sycl::event sum_reduction_axis1_over_group_with_atomics_contig_impl(
     sycl::queue exec_q,
     size_t iter_nelems, // number of reductions    (num. of rows in a matrix
                         // when reducing over rows)
@@ -458,11 +461,11 @@ sycl::event sum_reduction_over_group_with_atomics_contig_impl(
                     RowsIndexerT, NoOpIndexerT>;
             using ReductionIndexerT = NoOpIndexerT;
 
-            RowsIndexerT columns_indexer{
+            RowsIndexerT rows_indexer{
                 0, static_cast<py::ssize_t>(iter_nelems),
                 static_cast<py::ssize_t>(reduction_nelems)};
             NoOpIndexerT result_indexer{};
-            InputOutputIterIndexerT in_out_iter_indexer{columns_indexer,
+            InputOutputIterIndexerT in_out_iter_indexer{rows_indexer,
                                                         result_indexer};
             ReductionIndexerT reduction_indexer{};
 
@@ -495,7 +498,102 @@ sycl::event sum_reduction_over_group_with_atomics_contig_impl(
             auto localRange = sycl::range<1>{wg};
 
             using KernelName =
-                class sum_reduction_over_group_with_atomics_contig_krn<
+                class sum_reduction_axis1_over_group_with_atomics_contig_krn<
+                    argTy, resTy, ReductionOpT, InputOutputIterIndexerT,
+                    ReductionIndexerT>;
+
+            cgh.parallel_for<KernelName>(
+                sycl::nd_range<1>(globalRange, localRange),
+                ReductionOverGroupWithAtomicFunctor<argTy, resTy, ReductionOpT,
+                                                    InputOutputIterIndexerT,
+                                                    ReductionIndexerT>(
+                    arg_tp, res_tp, ReductionOpT(), identity_val,
+                    in_out_iter_indexer, reduction_indexer, reduction_nelems,
+                    iter_nelems, reductions_per_wi));
+        });
+
+        return comp_ev;
+    }
+}
+
+/* @brief Reduce rows in a matrix */
+template <typename argTy, typename resTy>
+sycl::event sum_reduction_axis0_over_group_with_atomics_contig_impl(
+    sycl::queue exec_q,
+    size_t iter_nelems, // number of reductions    (num. of cols in a matrix
+                        // when reducing over cols)
+    size_t reduction_nelems, // size of each reduction  (length of cols, i.e.
+                             // number of rows)
+    const char *arg_cp,
+    char *res_cp,
+    py::ssize_t iter_arg_offset,
+    py::ssize_t iter_res_offset,
+    py::ssize_t reduction_arg_offset,
+    const std::vector<sycl::event> &depends)
+{
+    const argTy *arg_tp = reinterpret_cast<const argTy *>(arg_cp) +
+                          iter_arg_offset + reduction_arg_offset;
+    resTy *res_tp = reinterpret_cast<resTy *>(res_cp) + iter_res_offset;
+
+    using ReductionOpT = sycl::plus<resTy>;
+    constexpr resTy identity_val = resTy{0};
+
+    const sycl::device &d = exec_q.get_device();
+    const auto &sg_sizes = d.get_info<sycl::info::device::sub_group_sizes>();
+    size_t wg = choose_workgroup_size<4>(reduction_nelems, sg_sizes);
+
+    {
+        sycl::event res_init_ev = exec_q.fill<resTy>(
+            res_tp, resTy(identity_val), iter_nelems, depends);
+
+        sycl::event comp_ev = exec_q.submit([&](sycl::handler &cgh) {
+            cgh.depends_on(res_init_ev);
+
+            using NoOpIndexerT = dpctl::tensor::offset_utils::NoOpIndexer;
+            using ColsIndexerT = dpctl::tensor::offset_utils::Strided1DIndexer;
+            using InputOutputIterIndexerT =
+                dpctl::tensor::offset_utils::TwoOffsets_CombinedIndexer<
+                    NoOpIndexerT, NoOpIndexerT>;
+            using ReductionIndexerT = ColsIndexerT;
+
+            NoOpIndexerT columns_indexer{};
+            NoOpIndexerT result_indexer{};
+            InputOutputIterIndexerT in_out_iter_indexer{columns_indexer,
+                                                        result_indexer};
+            ReductionIndexerT reduction_indexer{
+                0, /* size */ static_cast<py::ssize_t>(reduction_nelems),
+                /* step */ static_cast<py::ssize_t>(iter_nelems)};
+
+            constexpr size_t preferrered_reductions_per_wi = 8;
+            size_t reductions_per_wi =
+                (reduction_nelems < preferrered_reductions_per_wi * wg)
+                    ? std::max<size_t>(1, (reduction_nelems + wg - 1) / wg)
+                    : preferrered_reductions_per_wi;
+
+            size_t reduction_groups =
+                (reduction_nelems + reductions_per_wi * wg - 1) /
+                (reductions_per_wi * wg);
+
+            if (reduction_groups > 1) {
+                const size_t &max_wg =
+                    d.get_info<sycl::info::device::max_work_group_size>();
+
+                if (reduction_nelems < preferrered_reductions_per_wi * max_wg) {
+                    wg = max_wg;
+                    reductions_per_wi =
+                        std::max<size_t>(1, (reduction_nelems + wg - 1) / wg);
+                    reduction_groups =
+                        (reduction_nelems + reductions_per_wi * wg - 1) /
+                        (reductions_per_wi * wg);
+                }
+            }
+
+            auto globalRange =
+                sycl::range<1>{iter_nelems * reduction_groups * wg};
+            auto localRange = sycl::range<1>{wg};
+
+            using KernelName =
+                class sum_reduction_axis0_over_group_with_atomics_contig_krn<
                     argTy, resTy, ReductionOpT, InputOutputIterIndexerT,
                     ReductionIndexerT>;
 
@@ -1075,15 +1173,34 @@ struct SumOverAxisTempsStridedFactory
 };
 
 template <typename fnT, typename srcTy, typename dstTy>
-struct SumOverAxisAtomicContigFactory
+struct SumOverAxis1AtomicContigFactory
+{
+    fnT get() const
+    {
+        if constexpr (TypePairSupportDataForSumReductionAtomic<
+                          srcTy, dstTy>::is_defined)
+        {
+            return dpctl::tensor::kernels::
+                sum_reduction_axis1_over_group_with_atomics_contig_impl<srcTy,
+                                                                        dstTy>;
+        }
+        else {
+            return nullptr;
+        }
+    }
+};
+
+template <typename fnT, typename srcTy, typename dstTy>
+struct SumOverAxis0AtomicContigFactory
 {
     fnT get() const
     {
         if constexpr (TypePairSupportDataForSumReductionAtomic<
                           srcTy, dstTy>::is_defined)
         {
             return dpctl::tensor::kernels::
-                sum_reduction_over_group_with_atomics_contig_impl<srcTy, dstTy>;
+                sum_reduction_axis0_over_group_with_atomics_contig_impl<srcTy,
+                                                                        dstTy>;
         }
         else {
             return nullptr;

dpctl/tensor/libtensor/source/sum_reductions.cpp

@@ -88,8 +88,11 @@ static sum_reduction_strided_impl_fn_ptr
 
 using dpctl::tensor::kernels::sum_reduction_contig_impl_fn_ptr;
 static sum_reduction_contig_impl_fn_ptr
-    sum_over_axis_contig_atomic_dispatch_table[td_ns::num_types]
-                                              [td_ns::num_types];
+    sum_over_axis1_contig_atomic_dispatch_table[td_ns::num_types]
+                                               [td_ns::num_types];
+static sum_reduction_contig_impl_fn_ptr
+    sum_over_axis0_contig_atomic_dispatch_table[td_ns::num_types]
+                                               [td_ns::num_types];
 
 std::pair<sycl::event, sycl::event> py_sum_over_axis(
     dpctl::tensor::usm_ndarray src,
@@ -194,8 +197,30 @@ std::pair<sycl::event, sycl::event> py_sum_over_axis(
         if ((is_src_c_contig && is_dst_c_contig) ||
             (is_src_f_contig && dst_nelems == 1))
         {
-            auto fn = sum_over_axis_contig_atomic_dispatch_table[src_typeid]
-                                                                [dst_typeid];
+            auto fn = sum_over_axis1_contig_atomic_dispatch_table[src_typeid]
+                                                                 [dst_typeid];
+            if (fn != nullptr) {
+                size_t iter_nelems = dst_nelems;
+
+                constexpr py::ssize_t zero_offset = 0;
+
+                sycl::event sum_over_axis_contig_ev =
+                    fn(exec_q, iter_nelems, reduction_nelems, src.get_data(),
+                       dst.get_data(),
+                       zero_offset, // iteration_src_offset
+                       zero_offset, // iteration_dst_offset
+                       zero_offset, // reduction_src_offset
+                       depends);
+
+                sycl::event keep_args_event = dpctl::utils::keep_args_alive(
+                    exec_q, {src, dst}, {sum_over_axis_contig_ev});
+
+                return std::make_pair(keep_args_event, sum_over_axis_contig_ev);
+            }
+        }
+        else if (is_src_f_contig & is_dst_c_contig) {
+            auto fn = sum_over_axis0_contig_atomic_dispatch_table[src_typeid]
+                                                                 [dst_typeid];
             if (fn != nullptr) {
                 size_t iter_nelems = dst_nelems;
 
@@ -271,27 +296,58 @@ std::pair<sycl::event, sycl::event> py_sum_over_axis(
                                  iteration_src_offset, iteration_dst_offset);
     }
 
-    if (supports_atomics && (reduction_nd == 1) &&
-        (simplified_reduction_src_strides[0] == 1) && (iteration_nd == 1) &&
-        ((simplified_iteration_shape[0] == 1) ||
-         ((simplified_iteration_dst_strides[0] == 1) &&
-          (static_cast<size_t>(simplified_iteration_src_strides[0]) ==
-           reduction_nelems))))
-    {
-        auto fn =
-            sum_over_axis_contig_atomic_dispatch_table[src_typeid][dst_typeid];
-        if (fn != nullptr) {
-            size_t iter_nelems = dst_nelems;
+    if (supports_atomics && (reduction_nd == 1) && (iteration_nd == 1)) {
+        bool mat_reduce_over_axis1 = false;
+        bool mat_reduce_over_axis0 = false;
+        bool array_reduce_all_elems = false;
+        size_t iter_nelems = dst_nelems;
+
+        if (simplified_reduction_src_strides[0] == 1) {
+            array_reduce_all_elems = (simplified_iteration_shape[0] == 1);
+            mat_reduce_over_axis1 =
+                (simplified_iteration_dst_strides[0] == 1) &&
+                (static_cast<size_t>(simplified_iteration_src_strides[0]) ==
+                 reduction_nelems);
+        }
+        else if (static_cast<size_t>(simplified_reduction_src_strides[0]) ==
+                 iter_nelems)
+        {
+            mat_reduce_over_axis0 =
+                (simplified_iteration_dst_strides[0] == 1) &&
+                (simplified_iteration_src_strides[0] == 1);
+        }
+
+        if (mat_reduce_over_axis1 || array_reduce_all_elems) {
+            auto fn = sum_over_axis1_contig_atomic_dispatch_table[src_typeid]
+                                                                 [dst_typeid];
+            if (fn != nullptr) {
+                sycl::event sum_over_axis1_contig_ev =
+                    fn(exec_q, iter_nelems, reduction_nelems, src.get_data(),
+                       dst.get_data(), iteration_src_offset,
+                       iteration_dst_offset, reduction_src_offset, depends);
 
-            sycl::event sum_over_axis_contig_ev =
-                fn(exec_q, iter_nelems, reduction_nelems, src.get_data(),
-                   dst.get_data(), iteration_src_offset, iteration_dst_offset,
-                   reduction_src_offset, depends);
+                sycl::event keep_args_event = dpctl::utils::keep_args_alive(
+                    exec_q, {src, dst}, {sum_over_axis1_contig_ev});
+
+                return std::make_pair(keep_args_event,
+                                      sum_over_axis1_contig_ev);
+            }
+        }
+        else if (mat_reduce_over_axis0) {
+            auto fn = sum_over_axis0_contig_atomic_dispatch_table[src_typeid]
+                                                                 [dst_typeid];
+            if (fn != nullptr) {
+                sycl::event sum_over_axis0_contig_ev =
+                    fn(exec_q, iter_nelems, reduction_nelems, src.get_data(),
+                       dst.get_data(), iteration_src_offset,
+                       iteration_dst_offset, reduction_src_offset, depends);
 
-            sycl::event keep_args_event = dpctl::utils::keep_args_alive(
-                exec_q, {src, dst}, {sum_over_axis_contig_ev});
+                sycl::event keep_args_event = dpctl::utils::keep_args_alive(
+                    exec_q, {src, dst}, {sum_over_axis0_contig_ev});
 
-            return std::make_pair(keep_args_event, sum_over_axis_contig_ev);
+                return std::make_pair(keep_args_event,
+                                      sum_over_axis0_contig_ev);
+            }
         }
     }
 
@@ -451,11 +507,17 @@ void populate_sum_over_axis_dispatch_table(void)
         dtb2;
     dtb2.populate_dispatch_table(sum_over_axis_strided_temps_dispatch_table);
 
-    using dpctl::tensor::kernels::SumOverAxisAtomicContigFactory;
+    using dpctl::tensor::kernels::SumOverAxis1AtomicContigFactory;
     DispatchTableBuilder<sum_reduction_contig_impl_fn_ptr,
-                         SumOverAxisAtomicContigFactory, num_types>
+                         SumOverAxis1AtomicContigFactory, num_types>
         dtb3;
-    dtb3.populate_dispatch_table(sum_over_axis_contig_atomic_dispatch_table);
+    dtb3.populate_dispatch_table(sum_over_axis1_contig_atomic_dispatch_table);
+
+    using dpctl::tensor::kernels::SumOverAxis0AtomicContigFactory;
+    DispatchTableBuilder<sum_reduction_contig_impl_fn_ptr,
+                         SumOverAxis0AtomicContigFactory, num_types>
+        dtb4;
+    dtb4.populate_dispatch_table(sum_over_axis0_contig_atomic_dispatch_table);
 }
 
 namespace py = pybind11;