From 870fa57d7169087055f4125675dc655d1d266bb8 Mon Sep 17 00:00:00 2001 From: jluitjens Date: Tue, 9 Aug 2022 11:13:16 -0700 Subject: [PATCH] adding permute operator --- docs_input/api/tensorops.rst | 1 + include/matx_tensor_ops.h | 148 ++++++++++++++++++++++++++--- test/00_operators/OperatorTests.cu | 36 +++++++ 3 files changed, 171 insertions(+), 14 deletions(-) diff --git a/docs_input/api/tensorops.rst b/docs_input/api/tensorops.rst index e960b272..f621937b 100644 --- a/docs_input/api/tensorops.rst +++ b/docs_input/api/tensorops.rst @@ -93,3 +93,4 @@ Advanced Operators .. doxygenfunction:: lcollapse .. doxygenfunction:: clone .. doxygenfunction:: slice +.. doxygenfunction:: permute diff --git a/include/matx_tensor_ops.h b/include/matx_tensor_ops.h index f3bbea87..767929c1 100644 --- a/include/matx_tensor_ops.h +++ b/include/matx_tensor_ops.h @@ -565,6 +565,25 @@ __MATX_INLINE__ } }; } + +/** + * @brief Helper function to select values from a predicate operator + * + * select() is used to index from a source operator using indices stored + * in another operator. This is commonly used with the find_idx executor + * which returns the indices of values meeting a selection criteria. + * + * @tparam T Input type + * @tparam IdxType Operator with indices + * @param t Input operator + * @param idx Index tensor + * @return Value in t from each location in idx + */ +template +auto __MATX_INLINE__ select(T t, IdxType idx) +{ + return detail::SelectOp(t, idx); +}; namespace detail { template @@ -704,7 +723,7 @@ __MATX_INLINE__ { return T::Rank(); } - constexpr __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ index_t Size(int dim) const + constexpr __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ index_t Size(int32_t dim) const { if(dim == DIM) return idx_.Size(0); @@ -787,7 +806,7 @@ auto __MATX_INLINE__ remap(Op t, Ind idx, Inds... inds) private: typename base_type::type op_; std::array sizes_; - std::array dims_; + std::array dims_; std::array starts_; std::array strides_; @@ -799,8 +818,8 @@ auto __MATX_INLINE__ remap(Op t, Ind idx, Inds... inds) static_assert(DIM<=T::Rank(), "SliceOp: DIM must be less than or equal to operator rank."); __MATX_INLINE__ SliceOp(T op, const shape_type (&starts)[T::Rank()], const shape_type (&ends)[T::Rank()], const shape_type (&strides)[T::Rank()]) : op_(op) { - int d = 0; - for(int i = 0; i < T::Rank(); i++) { + int32_t d = 0; + for(int32_t i = 0; i < T::Rank(); i++) { shape_type start = starts[i]; shape_type end = ends[i]; @@ -836,12 +855,12 @@ auto __MATX_INLINE__ remap(Op t, Ind idx, Inds... inds) std::array ind{indices...}; #pragma unroll - for(int i = 0; i < T::Rank(); i++) { + for(int32_t i = 0; i < T::Rank(); i++) { ind[i] = starts_[i]; } #pragma unroll - for(int i = 0; i < Rank(); i++) { + for(int32_t i = 0; i < Rank(); i++) { ind[dims_[i]] += inds[i] * strides_[i]; } @@ -856,8 +875,8 @@ auto __MATX_INLINE__ remap(Op t, Ind idx, Inds... inds) static_assert((std::is_convertible_v && ... )); // convert variadic type to tuple so we can read/update - std::array inds{indices...}; - std::array ind{indices...}; + std::array inds{indices...}; + std::array ind{indices...}; #pragma unroll for(int i = 0; i < T::Rank(); i++) { @@ -877,7 +896,7 @@ auto __MATX_INLINE__ remap(Op t, Ind idx, Inds... inds) { return DIM; } - constexpr __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ index_t Size(int dim) const + constexpr __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ shape_type Size(int32_t dim) const { return sizes_[dim]; } @@ -951,13 +970,114 @@ auto __MATX_INLINE__ remap(Op t, Ind idx, Inds... inds) * @param idx Index tensor * @return Value in t from each location in idx */ -template -auto __MATX_INLINE__ select(T t, IdxType idx) -{ - return detail::SelectOp(t, idx); -}; +/** + * permutes dimensions of a tensor/operator + */ + namespace detail { + template + class PermuteOp : public BaseOp> + { + public: + using scalar_type = typename T::scalar_type; + using shape_type = typename T::shape_type; + + private: + typename base_type::type op_; + std::array dims_; + + public: + using matxop = bool; + using matxoplvalue = bool; + + static __MATX_INLINE__ constexpr __MATX_HOST__ __MATX_DEVICE__ int32_t Rank() + { + return T::Rank(); + } + + static_assert(Rank() > 0, "PermuteOp: Rank of operator must be greater than 0."); + + __MATX_INLINE__ PermuteOp(T op, const int32_t (&dims)[Rank()]) : op_(op) { + + bool selected[Rank()] = {0}; + + for(int32_t i = 0; i < Rank(); i++) { + int32_t dim = dims[i]; + MATX_ASSERT_STR(dim < Rank() && dim >= 0, matxInvalidDim, "PermuteOp: Invalid permute index."); + MATX_ASSERT_STR(selected[dim] == false, matxInvalidDim, "PermuteOp: Dim selected more than once"); + selected[dim] = true; + + dims_[i] = dims[i]; + } + }; + + template + __MATX_INLINE__ __MATX_DEVICE__ __MATX_HOST__ auto operator()(Is... indices) const + { + static_assert(sizeof...(Is)==Rank()); + static_assert((std::is_convertible_v && ... )); + + // convert variadic type to tuple so we can read/update + std::array inds{indices...}; + std::array ind{indices...}; + +#pragma unroll + for(int32_t i = 0; i < Rank(); i++) { + ind[dims_[i]] = inds[i]; + //ind[i] = inds[dims_[i]]; + } + + //return op_(ind); + return mapply(op_, ind); + } + + template + __MATX_INLINE__ __MATX_DEVICE__ __MATX_HOST__ auto& operator()(Is... indices) + { + static_assert(sizeof...(Is)==Rank()); +// static_assert((std::is_convertible_v && ... )); + + // convert variadic type to tuple so we can read/update + std::array inds{indices...}; + std::array ind{indices...}; + +#pragma unroll + for(int i = 0; i < Rank(); i++) { + ind[dims_[i]] = inds[i]; + //ind[i] = inds[dims_[i]]; + } + + //return op_(ind); + return mapply(op_, ind); + } + + constexpr __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ shape_type Size(int32_t dim) const + { + return op_.Size(dims_[dim]); + } + + template __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); } + }; + } + +/** + * @brief Operator to permute the dimensions of a tensor or operator. + * + * The each dimension must appear in the dims array once. + + * This operator can appear as an rvalue or lvalue. + * + * @tparam T Input operator/tensor type + * @param Op Input operator + * @param dims the reordered dimensions of the operator. + * @return permuted operator + */ + template + __MATX_INLINE__ auto permute( const T op, + const int32_t (&dims)[T::Rank()]) { + return detail::PermuteOp(op, dims); + } /** * Casts the element of the tensor to a specified type diff --git a/test/00_operators/OperatorTests.cu b/test/00_operators/OperatorTests.cu index d88ad053..1fe082c1 100644 --- a/test/00_operators/OperatorTests.cu +++ b/test/00_operators/OperatorTests.cu @@ -107,6 +107,42 @@ TYPED_TEST(OperatorTestsComplex, BaseOp) MATX_EXIT_HANDLER(); } +TYPED_TEST(OperatorTestsNumericNonComplex, PermuteOp) +{ + MATX_ENTER_HANDLER(); + auto A = make_tensor({10,20,30}); + for(int i=0; i < A.Size(0); i++) { + for(int j=0; j < A.Size(1); j++) { + for(int k=0; k < A.Size(2); k++) { + A(i,j,k) = TypeParam( i * A.Size(1)*A.Size(2) + + j * A.Size(2) + k); + } + } + } + + auto op = permute(A, {2, 0, 1}); + auto At = A.Permute({2, 0, 1}); + + ASSERT_TRUE(op.Size(0) == A.Size(2)); + ASSERT_TRUE(op.Size(1) == A.Size(0)); + ASSERT_TRUE(op.Size(2) == A.Size(1)); + + ASSERT_TRUE(op.Size(0) == At.Size(0)); + ASSERT_TRUE(op.Size(1) == At.Size(1)); + ASSERT_TRUE(op.Size(2) == At.Size(2)); + + for(int i=0; i < op.Size(0); i++) { + for(int j=0; j < op.Size(1); j++) { + for(int k=0; k < op.Size(2); k++) { + ASSERT_TRUE( op(i,j,k) == A(j,k,i)); + ASSERT_TRUE( op(i,j,k) == At(i,j,k)); + } + } + } + + MATX_EXIT_HANDLER(); +} + TYPED_TEST(OperatorTestsFloatNonComplex, FMod) { MATX_ENTER_HANDLER();