Add TwoOffsets_CombinedIndexer, and UnpackedStridedIndexer

The TwoOffsets_CombinedIndexer takes two single offset indexers and
combines them into a TwoOffsets struct.

The UnpackedStridedIndexer is a relative of StridedIndexer, except
shapes and strides are provided as separate pointers. This increases
its size, but may be useful to produce indexer that only computes
the second offset using shape/strides stored in packed format for
two-offset indexers.

These are going to be used in reduction kernels.

Author: oleksandr-pavlyk

dpctl/tensor/libtensor/include/utils/offset_utils.hpp

@@ -133,6 +133,7 @@ struct NoOpIndexer
     }
 };
 
+/* @brief Indexer with shape and strides arrays of same size are packed */
 struct StridedIndexer
 {
     StridedIndexer(int _nd,
@@ -143,24 +144,76 @@ struct StridedIndexer
     {
     }
 
+    size_t operator()(py::ssize_t gid) const
+    {
+        return compute_offset(gid);
+    }
+
     size_t operator()(size_t gid) const
+    {
+        return compute_offset(static_cast<py::ssize_t>(gid));
+    }
+
+private:
+    int nd;
+    py::ssize_t starting_offset;
+    py::ssize_t const *shape_strides;
+
+    size_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_displacement<const py::ssize_t *, const py::ssize_t *>(
-            static_cast<py::ssize_t>(gid),
+            gid,
             shape_strides,      // shape ptr
             shape_strides + nd, // strides ptr
             relative_offset);
         return starting_offset + relative_offset;
     }
+};
+
+/* @brief Indexer with shape, strides provided separately */
+struct UnpackedStridedIndexer
+{
+    UnpackedStridedIndexer(int _nd,
+                           py::ssize_t _offset,
+                           py::ssize_t const *_shape,
+                           py::ssize_t const *_strides)
+        : nd(_nd), starting_offset(_offset), shape(_shape), strides(_strides)
+    {
+    }
+
+    size_t operator()(py::ssize_t gid) const
+    {
+        return compute_offset(gid);
+    }
+
+    size_t operator()(size_t gid) const
+    {
+        return compute_offset(static_cast<py::ssize_t>(gid));
+    }
 
 private:
     int nd;
     py::ssize_t starting_offset;
-    py::ssize_t const *shape_strides;
+    py::ssize_t const *shape;
+    py::ssize_t const *strides;
+
+    size_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_displacement<const py::ssize_t *, const py::ssize_t *>(
+            gid,
+            shape,   // shape ptr
+            strides, // strides ptr
+            relative_offset);
+        return starting_offset + relative_offset;
+    }
 };
 
 struct Strided1DIndexer
@@ -206,7 +259,8 @@ struct Strided1DCyclicIndexer
 template <typename displacementT> struct TwoOffsets
 {
     TwoOffsets() : first_offset(0), second_offset(0) {}
-    TwoOffsets(displacementT first_offset_, displacementT second_offset_)
+    TwoOffsets(const displacementT &first_offset_,
+               const displacementT &second_offset_)
         : first_offset(first_offset_), second_offset(second_offset_)
     {
     }
@@ -238,6 +292,22 @@ struct TwoOffsets_StridedIndexer
     }
 
     TwoOffsets<py::ssize_t> operator()(py::ssize_t gid) const
+    {
+        return compute_offsets(gid);
+    }
+
+    TwoOffsets<py::ssize_t> operator()(size_t gid) const
+    {
+        return compute_offsets(static_cast<py::ssize_t>(gid));
+    }
+
+private:
+    int nd;
+    py::ssize_t starting_first_offset;
+    py::ssize_t starting_second_offset;
+    py::ssize_t const *shape_strides;
+
+    TwoOffsets<py::ssize_t> compute_offsets(py::ssize_t gid) const
     {
         using dpctl::tensor::strides::CIndexer_vector;
 
@@ -254,12 +324,6 @@ struct TwoOffsets_StridedIndexer
             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
@@ -272,12 +336,33 @@ struct TwoZeroOffsets_Indexer
     }
 };
 
+template <typename FirstIndexerT, typename SecondIndexerT>
+struct TwoOffsets_CombinedIndexer
+{
+private:
+    FirstIndexerT first_indexer_;
+    SecondIndexerT second_indexer_;
+
+public:
+    TwoOffsets_CombinedIndexer(const FirstIndexerT &first_indexer,
+                               const SecondIndexerT &second_indexer)
+        : first_indexer_(first_indexer), second_indexer_(second_indexer)
+    {
+    }
+
+    TwoOffsets<py::ssize_t> operator()(py::ssize_t gid) const
+    {
+        return TwoOffsets<py::ssize_t>(first_indexer_(gid),
+                                       second_indexer_(gid));
+    }
+};
+
 template <typename displacementT> struct ThreeOffsets
 {
     ThreeOffsets() : first_offset(0), second_offset(0), third_offset(0) {}
-    ThreeOffsets(displacementT first_offset_,
-                 displacementT second_offset_,
-                 displacementT third_offset_)
+    ThreeOffsets(const displacementT &first_offset_,
+                 const displacementT &second_offset_,
+                 const displacementT &third_offset_)
         : first_offset(first_offset_), second_offset(second_offset_),
           third_offset(third_offset_)
     {
@@ -317,6 +402,23 @@ struct ThreeOffsets_StridedIndexer
     }
 
     ThreeOffsets<py::ssize_t> operator()(py::ssize_t gid) const
+    {
+        return compute_offsets(gid);
+    }
+
+    ThreeOffsets<py::ssize_t> operator()(size_t gid) const
+    {
+        return compute_offsets(static_cast<py::ssize_t>(gid));
+    }
+
+private:
+    int nd;
+    py::ssize_t starting_first_offset;
+    py::ssize_t starting_second_offset;
+    py::ssize_t starting_third_offset;
+    py::ssize_t const *shape_strides;
+
+    ThreeOffsets<py::ssize_t> compute_offsets(py::ssize_t gid) const
     {
         using dpctl::tensor::strides::CIndexer_vector;
 
@@ -337,13 +439,6 @@ struct ThreeOffsets_StridedIndexer
             starting_second_offset + relative_second_offset,
             starting_third_offset + relative_third_offset);
     }
-
-private:
-    int nd;
-    py::ssize_t starting_first_offset;
-    py::ssize_t starting_second_offset;
-    py::ssize_t starting_third_offset;
-    py::ssize_t const *shape_strides;
 };
 
 struct ThreeZeroOffsets_Indexer