Merge pull request #1162 from IntelPython/add-combined-indexer

Add TwoOffsets_CombinedIndexer, and UnpackedStridedIndexer

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