From b61dd3220588120b8b4026d79db41abccd4e5366 Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk Date: Fri, 7 Apr 2023 16:22:55 -0500 Subject: [PATCH] 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. --- .../libtensor/include/utils/offset_utils.hpp | 133 +++++++++++++++--- 1 file changed, 114 insertions(+), 19 deletions(-) diff --git a/dpctl/tensor/libtensor/include/utils/offset_utils.hpp b/dpctl/tensor/libtensor/include/utils/offset_utils.hpp index 18bbf45a7d..74c62a7346 100644 --- a/dpctl/tensor/libtensor/include/utils/offset_utils.hpp +++ b/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(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( - static_cast(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(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( + gid, + shape, // shape ptr + strides, // strides ptr + relative_offset); + return starting_offset + relative_offset; + } }; struct Strided1DIndexer @@ -206,7 +259,8 @@ struct Strided1DCyclicIndexer template 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 operator()(py::ssize_t gid) const + { + return compute_offsets(gid); + } + + TwoOffsets operator()(size_t gid) const + { + return compute_offsets(static_cast(gid)); + } + +private: + int nd; + py::ssize_t starting_first_offset; + py::ssize_t starting_second_offset; + py::ssize_t const *shape_strides; + + TwoOffsets 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 +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 operator()(py::ssize_t gid) const + { + return TwoOffsets(first_indexer_(gid), + second_indexer_(gid)); + } +}; + template 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 operator()(py::ssize_t gid) const + { + return compute_offsets(gid); + } + + ThreeOffsets operator()(size_t gid) const + { + return compute_offsets(static_cast(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 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