1 /* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_OPTIMIZED_OPS_UTILS_H_
16 #define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_OPTIMIZED_OPS_UTILS_H_
17
18 #include "Eigen/Core"
19 #include "unsupported/Eigen/CXX11/Tensor"
20 #include "tensorflow/lite/kernels/internal/runtime_shape.h"
21 #include "tensorflow/lite/kernels/internal/types.h"
22
23 namespace tflite {
24 namespace optimized_ops {
25 // Make a local VectorMap typedef allowing to map a float array
26 // as a Eigen vector expression. The std::conditional here is to
27 // construct the suitable Eigen type for the constness of the
28 // data. Indeed, for const data, we need to produce
29 // Eigen::Map<const Eigen::Matrix<float, ...>>
30 // and not the more straightforward
31 // Eigen::Map<Eigen::Matrix<const float, ...>>
32 template <typename Scalar>
33 using VectorMap = typename std::conditional<
34 std::is_const<Scalar>::value,
35 Eigen::Map<const Eigen::Matrix<typename std::remove_const<Scalar>::type,
36 Eigen::Dynamic, 1>>,
37 Eigen::Map<Eigen::Matrix<Scalar, Eigen::Dynamic, 1>>>::type;
38
39 template <typename Scalar>
MapAsVector(Scalar * data,const RuntimeShape & shape)40 VectorMap<Scalar> MapAsVector(Scalar* data, const RuntimeShape& shape) {
41 const int size = shape.FlatSize();
42 return VectorMap<Scalar>(data, size, 1);
43 }
44
45 // Make a local VectorMap typedef allowing to map a float array
46 // as a Eigen matrix expression. The same explanation as for VectorMap
47 // above also applies here.
48 template <typename Scalar>
49 using MatrixMap = typename std::conditional<
50 std::is_const<Scalar>::value,
51 Eigen::Map<const Eigen::Matrix<typename std::remove_const<Scalar>::type,
52 Eigen::Dynamic, Eigen::Dynamic>>,
53 Eigen::Map<Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>>>::type;
54
55 template <typename Scalar>
MapAsMatrixWithLastDimAsRows(Scalar * data,const RuntimeShape & shape)56 MatrixMap<Scalar> MapAsMatrixWithLastDimAsRows(Scalar* data,
57 const RuntimeShape& shape) {
58 const int dims_count = shape.DimensionsCount();
59 const int rows = shape.Dims(dims_count - 1);
60 const int cols = FlatSizeSkipDim(shape, dims_count - 1);
61 return MatrixMap<Scalar>(data, rows, cols);
62 }
63
64 template <typename Scalar>
MapAsMatrixWithFirstDimAsCols(Scalar * data,const RuntimeShape & shape)65 MatrixMap<Scalar> MapAsMatrixWithFirstDimAsCols(Scalar* data,
66 const RuntimeShape& shape) {
67 const int cols = shape.Dims(0);
68 const int rows = FlatSizeSkipDim(shape, 0);
69 return MatrixMap<Scalar>(data, rows, cols);
70 }
71
72 template <typename Scalar>
73 using ArrayMap = typename std::conditional<
74 std::is_const<Scalar>::value,
75 Eigen::Map<const Eigen::Array<typename std::remove_const<Scalar>::type,
76 Eigen::Dynamic, Eigen::Dynamic>>,
77 Eigen::Map<Eigen::Array<Scalar, Eigen::Dynamic, Eigen::Dynamic>>>::type;
78
79 template <typename Scalar>
MapAsArrayWithLastDimAsRows(Scalar * data,const RuntimeShape & shape)80 ArrayMap<Scalar> MapAsArrayWithLastDimAsRows(Scalar* data,
81 const RuntimeShape& shape) {
82 const int dims_count = shape.DimensionsCount();
83 const int rows = shape.Dims(dims_count - 1);
84 const int cols = FlatSizeSkipDim(shape, dims_count - 1);
85 return ArrayMap<Scalar>(data, rows, cols);
86 }
87
88 } // namespace optimized_ops
89 } // namespace tflite
90
91 #endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_OPTIMIZED_OPS_UTILS_H_
92