Optimization for in-place arithmetic operators between rows and matrices

- In-place operations from C-contiguous rows into C-contiguous matrices show significant performance improvements

Author: ndgrigorian

dpctl/tensor/libtensor/include/kernels/elementwise_functions/add.hpp

@@ -480,6 +480,61 @@ struct AddInplaceStridedFactory
     }
 };
 
+template <typename argT, typename resT>
+class add_inplace_row_matrix_broadcast_sg_krn;
+
+template <typename argT, typename resT>
+using AddInplaceRowMatrixBroadcastingFunctor =
+    elementwise_common::BinaryInplaceRowMatrixBroadcastingFunctor<
+        argT,
+        resT,
+        AddInplaceFunctor<argT, resT>>;
+
+template <typename argT, typename resT>
+sycl::event add_inplace_row_matrix_broadcast_impl(
+    sycl::queue exec_q,
+    std::vector<sycl::event> &host_tasks,
+    size_t n0,
+    size_t n1,
+    const char *vec_p, // typeless pointer to (n1,) contiguous row
+    py::ssize_t vec_offset,
+    char *mat_p, // typeless pointer to (n0, n1) C-contiguous matrix
+    py::ssize_t mat_offset,
+    const std::vector<sycl::event> &depends = {})
+{
+    return elementwise_common::binary_inplace_row_matrix_broadcast_impl<
+        argT, resT, AddInplaceRowMatrixBroadcastingFunctor,
+        add_inplace_row_matrix_broadcast_sg_krn>(exec_q, host_tasks, n0, n1,
+                                                 vec_p, vec_offset, mat_p,
+                                                 mat_offset, depends);
+}
+
+template <typename fnT, typename T1, typename T2>
+struct AddInplaceRowMatrixBroadcastFactory
+{
+    fnT get()
+    {
+        using resT = typename AddOutputType<T1, T2>::value_type;
+        if constexpr (std::is_same_v<resT, void>) {
+            fnT fn = nullptr;
+            return fn;
+        }
+        else {
+            if constexpr (dpctl::tensor::type_utils::is_complex<T1>::value ||
+                          dpctl::tensor::type_utils::is_complex<T2>::value ||
+                          dpctl::tensor::type_utils::is_complex<resT>::value)
+            {
+                fnT fn = nullptr;
+                return fn;
+            }
+            else {
+                fnT fn = add_inplace_row_matrix_broadcast_impl<T1, T2>;
+                return fn;
+            }
+        }
+    }
+};
+
 } // namespace add
 } // namespace kernels
 } // namespace tensor

dpctl/tensor/libtensor/include/kernels/elementwise_functions/common_inplace.hpp

@@ -191,6 +191,60 @@ struct BinaryInplaceStridedFunctor
     }
 };
 
+template <typename argT, typename resT, typename BinaryOperatorT>
+struct BinaryInplaceRowMatrixBroadcastingFunctor
+{
+private:
+    const argT *padded_vec;
+    resT *mat;
+    size_t n_elems;
+    size_t n1;
+
+public:
+    BinaryInplaceRowMatrixBroadcastingFunctor(const argT *row_tp,
+                                              resT *mat_tp,
+                                              size_t n_elems_in_mat,
+                                              size_t n_elems_in_row)
+        : padded_vec(row_tp), mat(mat_tp), n_elems(n_elems_in_mat),
+          n1(n_elems_in_row)
+    {
+    }
+
+    void operator()(sycl::nd_item<1> ndit) const
+    {
+        BinaryOperatorT op{};
+        static_assert(BinaryOperatorT::supports_sg_loadstore::value);
+
+        auto sg = ndit.get_sub_group();
+        size_t gid = ndit.get_global_linear_id();
+
+        std::uint8_t sgSize = sg.get_local_range()[0];
+        size_t base = gid - sg.get_local_id()[0];
+
+        if (base + sgSize < n_elems) {
+            using in_ptrT =
+                sycl::multi_ptr<const argT,
+                                sycl::access::address_space::global_space>;
+            using res_ptrT =
+                sycl::multi_ptr<resT,
+                                sycl::access::address_space::global_space>;
+
+            const argT vec_el = sg.load(in_ptrT(&padded_vec[base % n1]));
+            resT mat_el = sg.load(res_ptrT(&mat[base]));
+
+            op(mat_el, vec_el);
+
+            sg.store(res_ptrT(&mat[base]), mat_el);
+        }
+        else {
+            for (size_t k = base + sg.get_local_id()[0]; k < n_elems;
+                 k += sgSize) {
+                op(mat[k], padded_vec[k % n1]);
+            }
+        }
+    }
+};
+
 // Typedefs for function pointers
 
 typedef sycl::event (*binary_inplace_contig_impl_fn_ptr_t)(
@@ -214,6 +268,17 @@ typedef sycl::event (*binary_inplace_strided_impl_fn_ptr_t)(
     const std::vector<sycl::event> &,
     const std::vector<sycl::event> &);
 
+typedef sycl::event (*binary_inplace_row_matrix_broadcast_impl_fn_ptr_t)(
+    sycl::queue,
+    std::vector<sycl::event> &,
+    size_t,
+    size_t,
+    const char *,
+    py::ssize_t,
+    char *,
+    py::ssize_t,
+    const std::vector<sycl::event> &);
+
 template <typename argTy,
           typename resTy,
           template <typename T1, typename T2, unsigned int vs, unsigned int nv>
@@ -289,6 +354,79 @@ binary_inplace_strided_impl(sycl::queue exec_q,
     return comp_ev;
 }
 
+template <typename argT,
+          typename resT,
+          template <typename T1, typename T3>
+          class BinaryInplaceRowMatrixBroadcastFunctorT,
+          template <typename T1, typename T3>
+          class kernel_name>
+sycl::event binary_inplace_row_matrix_broadcast_impl(
+    sycl::queue exec_q,
+    std::vector<sycl::event> &host_tasks,
+    size_t n0,
+    size_t n1,
+    const char *vec_p, // typeless pointer to (n1,) contiguous row
+    py::ssize_t vec_offset,
+    char *mat_p, // typeless pointer to (n0, n1) C-contiguous matrix
+    py::ssize_t mat_offset,
+    const std::vector<sycl::event> &depends = {})
+{
+    const argT *vec = reinterpret_cast<const argT *>(vec_p) + vec_offset;
+    resT *mat = reinterpret_cast<resT *>(mat_p) + mat_offset;
+
+    const auto &dev = exec_q.get_device();
+    const auto &sg_sizes = dev.get_info<sycl::info::device::sub_group_sizes>();
+    // Get device-specific kernel info max_sub_group_size
+    size_t max_sgSize =
+        *(std::max_element(std::begin(sg_sizes), std::end(sg_sizes)));
+
+    size_t n1_padded = n1 + max_sgSize;
+    argT *padded_vec = sycl::malloc_device<argT>(n1_padded, exec_q);
+
+    if (padded_vec == nullptr) {
+        throw std::runtime_error("Could not allocate memory on the device");
+    }
+    sycl::event make_padded_vec_ev = exec_q.submit([&](sycl::handler &cgh) {
+        cgh.depends_on(depends); // ensure vec contains actual data
+        cgh.parallel_for({n1_padded}, [=](sycl::id<1> id) {
+            auto i = id[0];
+            padded_vec[i] = vec[i % n1];
+        });
+    });
+
+    // sub-group spans work-items [I, I + sgSize)
+    // base = ndit.get_global_linear_id() - sg.get_local_id()[0]
+    // Generically, sg.load( &mat[base]) may load arrays from
+    // different rows of mat. The start corresponds to row (base / n0)
+    // We read sg.load(&padded_vec[(base / n0)]). The vector is padded to
+    // ensure that reads are accessible
+
+    size_t lws = 64;
+
+    sycl::event comp_ev = exec_q.submit([&](sycl::handler &cgh) {
+        cgh.depends_on(make_padded_vec_ev);
+
+        auto lwsRange = sycl::range<1>(lws);
+        size_t n_elems = n0 * n1;
+        size_t n_groups = (n_elems + lws - 1) / lws;
+        auto gwsRange = sycl::range<1>(n_groups * lws);
+
+        cgh.parallel_for<class kernel_name<argT, resT>>(
+            sycl::nd_range<1>(gwsRange, lwsRange),
+            BinaryInplaceRowMatrixBroadcastFunctorT<argT, resT>(padded_vec, mat,
+                                                                n_elems, n1));
+    });
+
+    sycl::event tmp_cleanup_ev = exec_q.submit([&](sycl::handler &cgh) {
+        cgh.depends_on(comp_ev);
+        sycl::context ctx = exec_q.get_context();
+        cgh.host_task([ctx, padded_vec]() { sycl::free(padded_vec, ctx); });
+    });
+    host_tasks.push_back(tmp_cleanup_ev);
+
+    return comp_ev;
+}
+
 } // namespace elementwise_common
 } // namespace kernels
 } // namespace tensor

dpctl/tensor/libtensor/include/kernels/elementwise_functions/multiply.hpp

@@ -496,6 +496,61 @@ struct MultiplyInplaceStridedFactory
     }
 };
 
+template <typename argT, typename resT>
+class multiply_inplace_row_matrix_broadcast_sg_krn;
+
+template <typename argT, typename resT>
+using MultiplyInplaceRowMatrixBroadcastingFunctor =
+    elementwise_common::BinaryInplaceRowMatrixBroadcastingFunctor<
+        argT,
+        resT,
+        MultiplyInplaceFunctor<argT, resT>>;
+
+template <typename argT, typename resT>
+sycl::event multiply_inplace_row_matrix_broadcast_impl(
+    sycl::queue exec_q,
+    std::vector<sycl::event> &host_tasks,
+    size_t n0,
+    size_t n1,
+    const char *vec_p, // typeless pointer to (n1,) contiguous row
+    py::ssize_t vec_offset,
+    char *mat_p, // typeless pointer to (n0, n1) C-contiguous matrix
+    py::ssize_t mat_offset,
+    const std::vector<sycl::event> &depends = {})
+{
+    return elementwise_common::binary_inplace_row_matrix_broadcast_impl<
+        argT, resT, MultiplyInplaceRowMatrixBroadcastingFunctor,
+        multiply_inplace_row_matrix_broadcast_sg_krn>(
+        exec_q, host_tasks, n0, n1, vec_p, vec_offset, mat_p, mat_offset,
+        depends);
+}
+
+template <typename fnT, typename T1, typename T2>
+struct MultiplyInplaceRowMatrixBroadcastFactory
+{
+    fnT get()
+    {
+        using resT = typename MultiplyOutputType<T1, T2>::value_type;
+        if constexpr (std::is_same_v<resT, void>) {
+            fnT fn = nullptr;
+            return fn;
+        }
+        else {
+            if constexpr (dpctl::tensor::type_utils::is_complex<T1>::value ||
+                          dpctl::tensor::type_utils::is_complex<T2>::value ||
+                          dpctl::tensor::type_utils::is_complex<resT>::value)
+            {
+                fnT fn = nullptr;
+                return fn;
+            }
+            else {
+                fnT fn = multiply_inplace_row_matrix_broadcast_impl<T1, T2>;
+                return fn;
+            }
+        }
+    }
+};
+
 } // namespace multiply
 } // namespace kernels
 } // namespace tensor

dpctl/tensor/libtensor/include/kernels/elementwise_functions/subtract.hpp

@@ -509,6 +509,61 @@ struct SubtractInplaceStridedFactory
     }
 };
 
+template <typename argT, typename resT>
+class subtract_inplace_row_matrix_broadcast_sg_krn;
+
+template <typename argT, typename resT>
+using SubtractInplaceRowMatrixBroadcastingFunctor =
+    elementwise_common::BinaryInplaceRowMatrixBroadcastingFunctor<
+        argT,
+        resT,
+        SubtractInplaceFunctor<argT, resT>>;
+
+template <typename argT, typename resT>
+sycl::event subtract_inplace_row_matrix_broadcast_impl(
+    sycl::queue exec_q,
+    std::vector<sycl::event> &host_tasks,
+    size_t n0,
+    size_t n1,
+    const char *vec_p, // typeless pointer to (n1,) contiguous row
+    py::ssize_t vec_offset,
+    char *mat_p, // typeless pointer to (n0, n1) C-contiguous matrix
+    py::ssize_t mat_offset,
+    const std::vector<sycl::event> &depends = {})
+{
+    return elementwise_common::binary_inplace_row_matrix_broadcast_impl<
+        argT, resT, SubtractInplaceRowMatrixBroadcastingFunctor,
+        subtract_inplace_row_matrix_broadcast_sg_krn>(
+        exec_q, host_tasks, n0, n1, vec_p, vec_offset, mat_p, mat_offset,
+        depends);
+}
+
+template <typename fnT, typename T1, typename T2>
+struct SubtractInplaceRowMatrixBroadcastFactory
+{
+    fnT get()
+    {
+        using resT = typename SubtractOutputType<T1, T2>::value_type;
+        if constexpr (std::is_same_v<resT, void>) {
+            fnT fn = nullptr;
+            return fn;
+        }
+        else {
+            if constexpr (dpctl::tensor::type_utils::is_complex<T1>::value ||
+                          dpctl::tensor::type_utils::is_complex<T2>::value ||
+                          dpctl::tensor::type_utils::is_complex<resT>::value)
+            {
+                fnT fn = nullptr;
+                return fn;
+            }
+            else {
+                fnT fn = subtract_inplace_row_matrix_broadcast_impl<T1, T2>;
+                return fn;
+            }
+        }
+    }
+};
+
 } // namespace subtract
 } // namespace kernels
 } // namespace tensor