Introduced dedicated function for rolling with nd-shift

Function for flattened rolling is renamed:
   _copy_usm_ndarray_for_roll -> _copy_usm_ndarray_for_roll_1d

_copy_usm_ndarray_for_roll_1d has the same signature:
   _copy_usm_ndarray_for_roll_1d(
         src : usm_ndarray,
	 dst : usm_ndarray,
	 shift: Int,
	 sycl_queue: dpctl.SyclQueue) ->
	      Tuple[dpctl.SyclEvent, dpctl.SyclEvent]

Introduced
   _copy_usm_ndarray_for_roll_nd(
         src : usm_ndarray,
	 dst : usm_ndarray,
	 shifts: Tuple[Int],
	 sycl_queue: dpctl.SyclQueue) ->
	      Tuple[dpctl.SyclEvent, dpctl.SyclEvent]

The length of shifts tuple must be the same as the dimensionality
of src and dst arrays, which are supposed to have the same shape
and the same data type.

Author: oleksandr-pavlyk

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

@@ -746,7 +746,85 @@ template <typename fnT, typename Ty> struct CopyForReshapeGenericFactory
     }
 };
 
-// =============== Copying for reshape ================== //
+// ================== Copying for roll ================== //
+
+/*! @brief Functor to cyclically roll global_id to the left */
+struct LeftRolled1DTransformer
+{
+    LeftRolled1DTransformer(size_t offset, size_t size)
+        : offset_(offset), size_(size)
+    {
+    }
+
+    size_t operator()(size_t gid) const
+    {
+        const size_t shifted_gid =
+            ((gid < offset_) ? gid + size_ - offset_ : gid - offset_);
+        return shifted_gid;
+    }
+
+private:
+    size_t offset_ = 0;
+    size_t size_ = 1;
+};
+
+/*! @brief Indexer functor to compose indexer and transformer */
+template <typename IndexerT, typename TransformerT> struct CompositionIndexer
+{
+    CompositionIndexer(IndexerT f, TransformerT t) : f_(f), t_(t) {}
+
+    auto operator()(size_t gid) const
+    {
+        return f_(t_(gid));
+    }
+
+private:
+    IndexerT f_;
+    TransformerT t_;
+};
+
+/*! @brief Indexer functor to find offset for nd-shifted indices lifted from
+ * iteration id */
+struct RolledNDIndexer
+{
+    RolledNDIndexer(int nd,
+                    const py::ssize_t *shape,
+                    const py::ssize_t *strides,
+                    const py::ssize_t *ndshifts,
+                    py::ssize_t starting_offset)
+        : nd_(nd), shape_(shape), strides_(strides), ndshifts_(ndshifts),
+          starting_offset_(starting_offset)
+    {
+    }
+
+    py::ssize_t operator()(size_t gid) const
+    {
+        return compute_offset(gid);
+    }
+
+private:
+    int nd_ = -1;
+    const py::ssize_t *shape_ = nullptr;
+    const py::ssize_t *strides_ = nullptr;
+    const py::ssize_t *ndshifts_ = nullptr;
+    py::ssize_t starting_offset_ = 0;
+
+    py::ssize_t compute_offset(py::ssize_t gid) const
+    {
+        using dpctl::tensor::strides::CIndexer_vector;
+
+        CIndexer_vector _ind(nd_);
+        py::ssize_t relative_offset_(0);
+        _ind.get_left_rolled_displacement<const py::ssize_t *,
+                                          const py::ssize_t *>(
+            gid,
+            shape_,    // shape ptr
+            strides_,  // strides ptr
+            ndshifts_, // shifts ptr
+            relative_offset_);
+        return starting_offset_ + relative_offset_;
+    }
+};
 
 template <typename Ty, typename SrcIndexerT, typename DstIndexerT>
 class copy_for_roll_strided_kernel;
@@ -755,32 +833,26 @@ template <typename Ty, typename SrcIndexerT, typename DstIndexerT>
 class StridedCopyForRollFunctor
 {
 private:
-    size_t offset = 0;
-    size_t size = 1;
     const Ty *src_p = nullptr;
     Ty *dst_p = nullptr;
     SrcIndexerT src_indexer_;
     DstIndexerT dst_indexer_;
 
 public:
-    StridedCopyForRollFunctor(size_t shift,
-                              size_t nelems,
-                              const Ty *src_ptr,
+    StridedCopyForRollFunctor(const Ty *src_ptr,
                               Ty *dst_ptr,
                               SrcIndexerT src_indexer,
                               DstIndexerT dst_indexer)
-        : offset(shift), size(nelems), src_p(src_ptr), dst_p(dst_ptr),
-          src_indexer_(src_indexer), dst_indexer_(dst_indexer)
+        : src_p(src_ptr), dst_p(dst_ptr), src_indexer_(src_indexer),
+          dst_indexer_(dst_indexer)
     {
     }
 
     void operator()(sycl::id<1> wiid) const
     {
         const size_t gid = wiid.get(0);
-        const size_t shifted_gid =
-            ((gid < offset) ? gid + size - offset : gid - offset);
 
-        const py::ssize_t src_offset = src_indexer_(shifted_gid);
+        const py::ssize_t src_offset = src_indexer_(gid);
         const py::ssize_t dst_offset = dst_indexer_(gid);
 
         dst_p[dst_offset] = src_p[src_offset];
@@ -800,8 +872,6 @@ typedef sycl::event (*copy_for_roll_strided_fn_ptr_t)(
     py::ssize_t,         // dst_offset
     const std::vector<sycl::event> &);
 
-template <typename Ty> class copy_for_roll_contig_kernel;
-
 /*!
  * @brief Function to copy content of array with a shift.
  *
@@ -812,8 +882,8 @@ template <typename Ty> class copy_for_roll_contig_kernel;
  * @param  shift  The shift in flat indexing, must be non-negative.
  * @param  nelems The number of elements to copy
  * @param  nd     Array dimensionality of the destination and source arrays
- * @param  packed_shapes_and_strides Kernel accessible USM array of size
- * `3*nd` with content `[common_shape, src_strides, dst_strides]`.
+ * @param  packed_shapes_and_strides Kernel accessible USM array
+ * of size `3*nd` with content `[common_shape, src_strides, dst_strides]`.
  * @param  src_p  Typeless USM pointer to the buffer of the source array
  * @param  src_offset Displacement of first element of src relative src_p in
  * elements
@@ -849,21 +919,29 @@ copy_for_roll_strided_impl(sycl::queue q,
         //   [ common_shape; src_strides; dst_strides ]
 
         StridedIndexer src_indexer{nd, src_offset, packed_shapes_and_strides};
+        LeftRolled1DTransformer left_roll_transformer{shift, nelems};
+
+        using CompositeIndexerT =
+            CompositionIndexer<StridedIndexer, LeftRolled1DTransformer>;
+
+        CompositeIndexerT rolled_src_indexer(src_indexer,
+                                             left_roll_transformer);
+
         UnpackedStridedIndexer dst_indexer{nd, dst_offset,
                                            packed_shapes_and_strides,
                                            packed_shapes_and_strides + 2 * nd};
 
-        using KernelName = copy_for_roll_strided_kernel<Ty, StridedIndexer,
+        using KernelName = copy_for_roll_strided_kernel<Ty, CompositeIndexerT,
                                                         UnpackedStridedIndexer>;
 
         const Ty *src_tp = reinterpret_cast<const Ty *>(src_p);
         Ty *dst_tp = reinterpret_cast<Ty *>(dst_p);
 
         cgh.parallel_for<KernelName>(
             sycl::range<1>(nelems),
-            StridedCopyForRollFunctor<Ty, StridedIndexer,
+            StridedCopyForRollFunctor<Ty, CompositeIndexerT,
                                       UnpackedStridedIndexer>(
-                shift, nelems, src_tp, dst_tp, src_indexer, dst_indexer));
+                src_tp, dst_tp, rolled_src_indexer, dst_indexer));
     });
 
     return copy_for_roll_ev;
@@ -880,6 +958,8 @@ typedef sycl::event (*copy_for_roll_contig_fn_ptr_t)(
     py::ssize_t,  // dst_offset
     const std::vector<sycl::event> &);
 
+template <typename Ty> class copy_for_roll_contig_kernel;
+
 /*!
  * @brief Function to copy content of array with a shift.
  *
@@ -917,6 +997,10 @@ sycl::event copy_for_roll_contig_impl(sycl::queue q,
         cgh.depends_on(depends);
 
         NoOpIndexer src_indexer{};
+        LeftRolled1DTransformer roller{shift, nelems};
+
+        CompositionIndexer<NoOpIndexer, LeftRolled1DTransformer>
+            left_rolled_src_indexer{src_indexer, roller};
         NoOpIndexer dst_indexer{};
 
         using KernelName = copy_for_roll_contig_kernel<Ty>;
@@ -926,8 +1010,10 @@ sycl::event copy_for_roll_contig_impl(sycl::queue q,
 
         cgh.parallel_for<KernelName>(
             sycl::range<1>(nelems),
-            StridedCopyForRollFunctor<Ty, NoOpIndexer, NoOpIndexer>(
-                shift, nelems, src_tp, dst_tp, src_indexer, dst_indexer));
+            StridedCopyForRollFunctor<
+                Ty, CompositionIndexer<NoOpIndexer, LeftRolled1DTransformer>,
+                NoOpIndexer>(src_tp, dst_tp, left_rolled_src_indexer,
+                             dst_indexer));
     });
 
     return copy_for_roll_ev;
@@ -961,6 +1047,86 @@ template <typename fnT, typename Ty> struct CopyForRollContigFactory
     }
 };
 
+template <typename Ty, typename SrcIndexerT, typename DstIndexerT>
+class copy_for_roll_ndshift_strided_kernel;
+
+// define function type
+typedef sycl::event (*copy_for_roll_ndshift_strided_fn_ptr_t)(
+    sycl::queue,
+    size_t,              // num_elements
+    int,                 // common_nd
+    const py::ssize_t *, // packed shape, strides, shifts
+    const char *,        // src_data_ptr
+    py::ssize_t,         // src_offset
+    char *,              // dst_data_ptr
+    py::ssize_t,         // dst_offset
+    const std::vector<sycl::event> &);
+
+template <typename Ty>
+sycl::event copy_for_roll_ndshift_strided_impl(
+    sycl::queue q,
+    size_t nelems,
+    int nd,
+    const py::ssize_t *packed_shapes_and_strides_and_shifts,
+    const char *src_p,
+    py::ssize_t src_offset,
+    char *dst_p,
+    py::ssize_t dst_offset,
+    const std::vector<sycl::event> &depends)
+{
+    dpctl::tensor::type_utils::validate_type_for_device<Ty>(q);
+
+    sycl::event copy_for_roll_ev = q.submit([&](sycl::handler &cgh) {
+        cgh.depends_on(depends);
+
+        // packed_shapes_and_strides_and_shifts:
+        //   USM array of size 4 * nd
+        //   [ common_shape; src_strides; dst_strides; shifts ]
+
+        const py::ssize_t *shape_ptr = packed_shapes_and_strides_and_shifts;
+        const py::ssize_t *src_strides_ptr =
+            packed_shapes_and_strides_and_shifts + nd;
+        const py::ssize_t *dst_strides_ptr =
+            packed_shapes_and_strides_and_shifts + 2 * nd;
+        const py::ssize_t *shifts_ptr =
+            packed_shapes_and_strides_and_shifts + 3 * nd;
+
+        RolledNDIndexer src_indexer{nd, shape_ptr, src_strides_ptr, shifts_ptr,
+                                    src_offset};
+
+        UnpackedStridedIndexer dst_indexer{nd, dst_offset, shape_ptr,
+                                           dst_strides_ptr};
+
+        using KernelName = copy_for_roll_strided_kernel<Ty, RolledNDIndexer,
+                                                        UnpackedStridedIndexer>;
+
+        const Ty *src_tp = reinterpret_cast<const Ty *>(src_p);
+        Ty *dst_tp = reinterpret_cast<Ty *>(dst_p);
+
+        cgh.parallel_for<KernelName>(
+            sycl::range<1>(nelems),
+            StridedCopyForRollFunctor<Ty, RolledNDIndexer,
+                                      UnpackedStridedIndexer>(
+                src_tp, dst_tp, src_indexer, dst_indexer));
+    });
+
+    return copy_for_roll_ev;
+}
+
+/*!
+ * @brief Factory to get function pointer of type `fnT` for given array data
+ * type `Ty`.
+ * @ingroup CopyAndCastKernels
+ */
+template <typename fnT, typename Ty> struct CopyForRollNDShiftFactory
+{
+    fnT get()
+    {
+        fnT f = copy_for_roll_ndshift_strided_impl<Ty>;
+        return f;
+    }
+};
+
 } // namespace copy_and_cast
 } // namespace kernels
 } // namespace tensor