1 /** 2 * Copyright (c) Huawei Technologies Co., Ltd. 2021-2022. All rights reserved. 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef AICPU_EIGENTENSOR_H 18 #define AICPU_EIGENTENSOR_H 19 20 #include <functional> 21 #include <vector> 22 #include "cpu_tensor.h" 23 #include "inc/kernel_log.h" 24 #include "unsupported/Eigen/CXX11/Tensor" 25 26 namespace aicpu { 27 // Helper to define Tensor types given that the scalar is of type T. 28 template <typename T, int NDIMS = 1, typename IndexType = Eigen::DenseIndex> 29 struct TTypes { 30 // Rank-<NDIMS> tensor of scalar type T. 31 typedef Eigen::TensorMap<Eigen::Tensor<T, NDIMS, Eigen::RowMajor, IndexType>, Eigen::Aligned> Tensor; 32 typedef Eigen::TensorMap<Eigen::Tensor<const T, NDIMS, Eigen::RowMajor, IndexType>, Eigen::Aligned> ConstTensor; 33 34 // Unaligned Rank-<NDIMS> tensor of scalar type T. 35 typedef Eigen::TensorMap<Eigen::Tensor<T, NDIMS, Eigen::RowMajor, IndexType> > UnalignedTensor; 36 typedef Eigen::TensorMap<Eigen::Tensor<const T, NDIMS, Eigen::RowMajor, IndexType> > UnalignedConstTensor; 37 38 typedef Eigen::TensorMap<Eigen::Tensor<T, NDIMS, Eigen::RowMajor, int>, Eigen::Aligned> Tensor32Bit; 39 40 // Scalar tensor (implemented as a rank-0 tensor) of scalar type T. 41 typedef Eigen::TensorMap<Eigen::TensorFixedSize<T, Eigen::Sizes<>, Eigen::RowMajor, IndexType>, Eigen::Aligned> 42 Scalar; 43 typedef Eigen::TensorMap<Eigen::TensorFixedSize<const T, Eigen::Sizes<>, Eigen::RowMajor, IndexType>, Eigen::Aligned> 44 ConstScalar; 45 46 // Unaligned Scalar tensor of scalar type T. 47 typedef Eigen::TensorMap<Eigen::TensorFixedSize<T, Eigen::Sizes<>, Eigen::RowMajor, IndexType> > UnalignedScalar; 48 typedef Eigen::TensorMap<Eigen::TensorFixedSize<const T, Eigen::Sizes<>, Eigen::RowMajor, IndexType> > 49 UnalignedConstScalar; 50 51 // Rank-1 tensor (vector) of scalar type T. 52 typedef Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType>, Eigen::Aligned> Flat; 53 typedef Eigen::TensorMap<Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType>, Eigen::Aligned> ConstFlat; 54 typedef Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType>, Eigen::Aligned> Vec; 55 typedef Eigen::TensorMap<Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType>, Eigen::Aligned> ConstVec; 56 57 // Unaligned Rank-1 tensor (vector) of scalar type T. 58 typedef Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType> > UnalignedFlat; 59 typedef Eigen::TensorMap<Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType> > UnalignedConstFlat; 60 typedef Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType> > UnalignedVec; 61 typedef Eigen::TensorMap<Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType> > UnalignedConstVec; 62 63 // Rank-2 tensor (matrix) of scalar type T. 64 typedef Eigen::TensorMap<Eigen::Tensor<T, 2, Eigen::RowMajor, IndexType>, Eigen::Aligned> Matrix; 65 typedef Eigen::TensorMap<Eigen::Tensor<const T, 2, Eigen::RowMajor, IndexType>, Eigen::Aligned> ConstMatrix; 66 67 // Unaligned Rank-2 tensor (matrix) of scalar type T. 68 typedef Eigen::TensorMap<Eigen::Tensor<T, 2, Eigen::RowMajor, IndexType> > UnalignedMatrix; 69 typedef Eigen::TensorMap<Eigen::Tensor<const T, 2, Eigen::RowMajor, IndexType> > UnalignedConstMatrix; 70 }; 71 } // namespace aicpu 72 73 namespace aicpu { 74 75 using ShapeVector = std::vector<int64_t>; 76 77 class EigenTensor { 78 public: 79 EigenTensor() = delete; EigenTensor(Tensor * tensor,void * data)80 EigenTensor(Tensor *tensor, void *data) 81 : tensor_(tensor), tensor_shape_(tensor->GetTensorShape()->GetDimSizes()), tensor_data_(data) {} EigenTensor(ShapeVector & shape,void * data_ptr)82 EigenTensor(ShapeVector &shape, void *data_ptr) : tensor_(nullptr), tensor_shape_(shape), tensor_data_(data_ptr) {} 83 ~EigenTensor() = default; 84 85 /* 86 * Get tensor 87 * @return succ: tensor, error : nullptr 88 */ 89 const Tensor *GetTensor() const; 90 91 /* 92 * Eigen vec 93 * @return Eigen vec 94 */ 95 template <typename T> vec()96 typename TTypes<T>::Vec vec() { 97 return tensor<T, 1>(); 98 } 99 100 /* 101 * Eigen matrix 102 * @return Eigen matrix 103 */ 104 template <typename T> matrix()105 typename TTypes<T>::Matrix matrix() { 106 return tensor<T, 2>(); 107 } 108 109 /* 110 * Eigen ConstMatrix 111 * @return Eigen ConstMatrix 112 */ 113 template <typename T> matrix()114 typename TTypes<T>::ConstMatrix matrix() const { 115 return tensor<T, 2>(); 116 } 117 118 /* 119 * Eigen tensor 120 * @return Eigen tensor 121 */ 122 template <typename T, size_t NDIMS> tensor()123 typename TTypes<T, NDIMS>::Tensor tensor() { 124 return typename TTypes<T, NDIMS>::Tensor(reinterpret_cast<T *>(tensor_data_), AsEigenDSizes<NDIMS>()); 125 } 126 127 /* 128 * Eigen ConstTensor 129 * @return Eigen ConstTensor 130 */ 131 template <typename T, size_t NDIMS> tensor()132 typename TTypes<T, NDIMS>::ConstTensor tensor() const { 133 return typename TTypes<T, NDIMS>::ConstTensor(reinterpret_cast<const T *>(tensor_data_), AsEigenDSizes<NDIMS>()); 134 } 135 136 /* 137 * Eigen Flat 138 * @return Eigen Flat 139 */ 140 template <typename T> flat()141 typename TTypes<T>::Flat flat() { 142 return typename TTypes<T>::Flat( 143 reinterpret_cast<T *>(tensor_data_), 144 {std::accumulate(tensor_shape_.begin(), tensor_shape_.end(), 1, std::multiplies<int64_t>())}); 145 } 146 147 /* 148 * which case we pad the rest of the sizes with 1. 149 * @return Eigen::DSizes: pad the rest of the sizes with 1 150 */ 151 template <int NDIMS, typename IndexType> AsEigenDSizesWithPadding()152 Eigen::DSizes<IndexType, NDIMS> AsEigenDSizesWithPadding() const { 153 Eigen::DSizes<IndexType, NDIMS> dsizes; 154 for (size_t d = 0; d < tensor_shape_.size(); d++) { 155 dsizes[d] = static_cast<IndexType>(tensor_shape_[d]); 156 } 157 for (size_t d = tensor_shape_.size(); d < NDIMS; d++) { 158 dsizes[d] = 1; 159 } 160 return dsizes; 161 } 162 163 /* 164 * Fill `*dsizes` from `*this` 165 * @return Eigen::DSizes: pad the rest of the sizes with 1 166 */ 167 template <int NDIMS, typename IndexType = Eigen::DenseIndex> AsEigenDSizes()168 Eigen::DSizes<IndexType, NDIMS> AsEigenDSizes() const { 169 return AsEigenDSizesWithPadding<NDIMS, IndexType>(); 170 } 171 172 private: 173 Tensor *tensor_{nullptr}; 174 ShapeVector tensor_shape_; 175 void *tensor_data_; 176 }; 177 } // namespace aicpu 178 179 #endif // AICPU_EIGENTENSOR_H 180