Merge pull request #1142 from IntelPython/refactor-indexers-utilities

Author: oleksandr-pavlyk

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

@@ -31,7 +31,7 @@
 #include <utility>
 #include <vector>
 
-#include "utils/strided_iters.hpp"
+#include "utils/offset_utils.hpp"
 #include "utils/type_dispatch.hpp"
 
 namespace dpctl
@@ -45,6 +45,8 @@ namespace indexing
 
 namespace py = pybind11;
 
+using namespace dpctl::tensor::offset_utils;
+
 template <typename T> T ceiling_quotient(T n, T m)
 {
     return (n + m - 1) / m;
@@ -67,82 +69,6 @@ template <typename inputT,
           typename TransformerT>
 class inclusive_scan_rec_chunk_update_krn;
 
-struct NoOpIndexer
-{
-    size_t operator()(size_t gid) const
-    {
-        return gid;
-    }
-};
-
-struct StridedIndexer
-{
-    StridedIndexer(int _nd,
-                   py::ssize_t _offset,
-                   py::ssize_t const *_packed_shape_strides)
-        : nd(_nd), starting_offset(_offset),
-          shape_strides(_packed_shape_strides)
-    {
-    }
-
-    size_t operator()(size_t gid) const
-    {
-        CIndexer_vector _ind(nd);
-        py::ssize_t relative_offset(0);
-        _ind.get_displacement<const py::ssize_t *, const py::ssize_t *>(
-            static_cast<py::ssize_t>(gid),
-            shape_strides,      // shape ptr
-            shape_strides + nd, // strides ptr
-            relative_offset);
-        return starting_offset + relative_offset;
-    }
-
-private:
-    int nd;
-    py::ssize_t starting_offset;
-    py::ssize_t const *shape_strides;
-};
-
-struct Strided1DIndexer
-{
-    Strided1DIndexer(py::ssize_t _offset, py::ssize_t _size, py::ssize_t _step)
-        : offset(_offset), size(static_cast<size_t>(_size)), step(_step)
-    {
-    }
-
-    size_t operator()(size_t gid) const
-    {
-        // ensure 0 <= gid < size
-        return static_cast<size_t>(offset +
-                                   std::min<size_t>(gid, size - 1) * step);
-    }
-
-private:
-    py::ssize_t offset = 0;
-    size_t size = 1;
-    py::ssize_t step = 1;
-};
-
-struct Strided1DCyclicIndexer
-{
-    Strided1DCyclicIndexer(py::ssize_t _offset,
-                           py::ssize_t _size,
-                           py::ssize_t _step)
-        : offset(_offset), size(static_cast<size_t>(_size)), step(_step)
-    {
-    }
-
-    size_t operator()(size_t gid) const
-    {
-        return static_cast<size_t>(offset + (gid % size) * step);
-    }
-
-private:
-    py::ssize_t offset = 0;
-    size_t size = 1;
-    py::ssize_t step = 1;
-};
-
 template <typename inputT, typename outputT> struct NonZeroIndicator
 {
     NonZeroIndicator() {}
@@ -200,9 +126,9 @@ sycl::event inclusive_scan_rec(sycl::queue exec_q,
 
         slmT slm_iscan_tmp(lws, cgh);
 
-        cgh.parallel_for<class inclusive_scan_rec_local_scan_krn<inputT, outputT, n_wi, IndexerT, decltype(transformer)>>(
-            sycl::nd_range<1>(gws, lws),
-            [=](sycl::nd_item<1> it)
+        cgh.parallel_for<class inclusive_scan_rec_local_scan_krn<
+            inputT, outputT, n_wi, IndexerT, decltype(transformer)>>(
+            sycl::nd_range<1>(gws, lws), [=](sycl::nd_item<1> it)
         {
             auto chunk_gid = it.get_global_id(0);
             auto lid = it.get_local_id(0);
@@ -245,8 +171,7 @@ sycl::event inclusive_scan_rec(sycl::queue exec_q,
             for (size_t m_wi = 0; m_wi < n_wi && i + m_wi < n_elems; ++m_wi) {
                 output[i + m_wi] = local_isum[m_wi];
             }
-            }
-        );
+        });
     });
 
     sycl::event out_event = inc_scan_phase1_ev;
@@ -266,15 +191,14 @@ sycl::event inclusive_scan_rec(sycl::queue exec_q,
         // output[ chunk_size * (i + 1) + j] += temp[i]
         auto e3 = exec_q.submit([&](sycl::handler &cgh) {
             cgh.depends_on(e2);
-            cgh.parallel_for<class inclusive_scan_rec_chunk_update_krn<inputT, outputT, IndexerT, decltype(transformer)>>(
-                {n_elems},
-                [=](auto wiid)
+            cgh.parallel_for<class inclusive_scan_rec_chunk_update_krn<
+                inputT, outputT, IndexerT, decltype(transformer)>>(
+                {n_elems}, [=](auto wiid)
             {
                 auto gid = wiid[0];
                 auto i = (gid / chunk_size);
                 output[gid] += (i > 0) ? temp[i - 1] : 0;
-                }
-            );
+            });
         });
 
         sycl::event e4 = exec_q.submit([&](sycl::handler &cgh) {
@@ -289,73 +213,6 @@ sycl::event inclusive_scan_rec(sycl::queue exec_q,
     return out_event;
 }
 
-template <typename displacementT> struct TwoOffsets
-{
-    TwoOffsets() : first_offset(0), second_offset(0) {}
-    TwoOffsets(displacementT first_offset_, displacementT second_offset_)
-        : first_offset(first_offset_), second_offset(second_offset_)
-    {
-    }
-
-    displacementT get_first_offset() const
-    {
-        return first_offset;
-    }
-    displacementT get_second_offset() const
-    {
-        return second_offset;
-    }
-
-private:
-    displacementT first_offset = 0;
-    displacementT second_offset = 0;
-};
-
-struct TwoOffsets_StridedIndexer
-{
-    TwoOffsets_StridedIndexer(int common_nd,
-                              py::ssize_t first_offset_,
-                              py::ssize_t second_offset_,
-                              py::ssize_t const *_packed_shape_strides)
-        : nd(common_nd), starting_first_offset(first_offset_),
-          starting_second_offset(second_offset_),
-          shape_strides(_packed_shape_strides)
-    {
-    }
-
-    TwoOffsets<py::ssize_t> operator()(py::ssize_t gid) const
-    {
-        CIndexer_vector _ind(nd);
-        py::ssize_t relative_first_offset(0);
-        py::ssize_t relative_second_offset(0);
-        _ind.get_displacement<const py::ssize_t *, const py::ssize_t *>(
-            gid,
-            shape_strides,          // shape ptr
-            shape_strides + nd,     // src strides ptr
-            shape_strides + 2 * nd, // src strides ptr
-            relative_first_offset, relative_second_offset);
-        return TwoOffsets<py::ssize_t>(
-            starting_first_offset + relative_first_offset,
-            starting_second_offset + relative_second_offset);
-    }
-
-private:
-    int nd;
-    py::ssize_t starting_first_offset;
-    py::ssize_t starting_second_offset;
-    py::ssize_t const *shape_strides;
-};
-
-struct TwoZeroOffsets_Indexer
-{
-    TwoZeroOffsets_Indexer() {}
-
-    TwoOffsets<py::ssize_t> operator()(py::ssize_t) const
-    {
-        return TwoOffsets<py::ssize_t>();
-    }
-};
-
 template <typename OrthogIndexerT,
           typename MaskedSrcIndexerT,
           typename MaskedDstIndexerT,

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

@@ -24,6 +24,7 @@
 //===----------------------------------------------------------------------===//
 
 #pragma once
+#include "utils/offset_utils.hpp"
 #include "utils/strided_iters.hpp"
 #include "utils/type_utils.hpp"
 #include <CL/sycl.hpp>
@@ -48,6 +49,7 @@ template <typename Ty, typename wTy> class linear_sequence_affine_kernel;
 template <typename Ty> class eye_kernel;
 
 namespace py = pybind11;
+using namespace dpctl::tensor::offset_utils;
 
 /* =========== Unboxing Python scalar =============== */
 
@@ -601,16 +603,17 @@ sycl::event tri_impl(sycl::queue exec_q,
     sycl::event tri_ev = exec_q.submit([&](sycl::handler &cgh) {
         cgh.depends_on(depends);
         cgh.depends_on(additional_depends);
+
         cgh.parallel_for<tri_kernel<Ty, upper>>(
             sycl::range<1>(inner_range * outer_range), [=](sycl::id<1> idx) {
                 py::ssize_t outer_gid = idx[0] / inner_range;
                 py::ssize_t inner_gid = idx[0] - inner_range * outer_gid;
 
-                py::ssize_t src_inner_offset, dst_inner_offset;
-                bool to_copy;
+                py::ssize_t src_inner_offset = 0, dst_inner_offset = 0;
+                bool to_copy(true);
 
                 {
-                    // py::ssize_t inner_gid = idx.get_id(0);
+                    using dpctl::tensor::strides::CIndexer_array;
                     CIndexer_array<d2, py::ssize_t> indexer_i(
                         {shape_and_strides[nd_2], shape_and_strides[nd_1]});
                     indexer_i.set(inner_gid);
@@ -631,7 +634,7 @@ sycl::event tri_impl(sycl::queue exec_q,
                 py::ssize_t src_offset = 0;
                 py::ssize_t dst_offset = 0;
                 {
-                    // py::ssize_t outer_gid = idx.get_id(1);
+                    using dpctl::tensor::strides::CIndexer_vector;
                     CIndexer_vector<py::ssize_t> outer(nd - d2);
                     outer.get_displacement(
                         outer_gid, shape_and_strides, shape_and_strides + src_s,