1 /* Copyright 2020 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 16 #ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_COMMON_H 17 #define TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_COMMON_H 18 19 #include "mlir/IR/PatternMatch.h" // from @llvm-project 20 #include "mlir/Support/LLVM.h" // from @llvm-project 21 22 // This file contains legalizations common to mapping both TensorFlow and 23 // TensorFlow Lite to TOSA. 24 // 25 // Conversion functions return None on a failure or result value on success. 26 // Callers must check and return a LogicalResult failure on nullptr. 27 // 28 // For these functions, the framework-specific operands/attributes/defaults 29 // are already extracted and placed in a common form for lowering. 30 31 namespace mlir { 32 namespace tosa { 33 34 // Lowers the Pack operator to TOSA. 35 llvm::Optional<Value> convertPackOp(PatternRewriter& rewriter, Operation* op, 36 Value result_value, 37 SmallVector<Value, 8>& inputs, 38 int32_t axis); 39 40 // Lowers the Unpack operator to TOSA. 41 llvm::Optional<ValueRange> convertUnpackOp(PatternRewriter& rewriter, 42 Operation* op, Value input_value, 43 int32_t axis); 44 45 // Lowers the Select operator to TOSA. 46 llvm::Optional<Value> convertSelectOp(PatternRewriter& rewriter, Operation* op, 47 Value result_value, Value condition_value, 48 Value x_value, Value y_value); 49 50 // Lowers the ZerosLike operator to TOSA by creating a constant 51 // of the desired type and shape. 52 llvm::Optional<Value> convertZerosLikeOp(PatternRewriter& rewriter, 53 Operation* op, Value result, 54 Value input); 55 56 // Lowers the Mul operator to TOSA. For quantized types, this requires 57 // inserting rescale operators before and after the operation. 58 llvm::Optional<Value> convertMultiplyOp(PatternRewriter& rewriter, 59 Operation* op, Value output_val, 60 Value input_lhs_val, 61 Value input_rhs_val); 62 63 // Lowers the SquaredDifference operator to TOSA. 64 llvm::Optional<Value> convertSquaredDifferenceOp(PatternRewriter& rewriter, 65 Operation* op, Value result, 66 Value x, Value y); 67 68 // Lowers the Round operator to TOSA. 69 llvm::Optional<Value> convertRoundOp(PatternRewriter& rewriter, Operation* op, 70 Value result, Value input); 71 72 // Lowers ConcatV2 to TOSA. 73 llvm::Optional<Value> convertConcatV2Op(PatternRewriter& rewriter, 74 Operation* op, Value result_value, 75 SmallVector<Value, 8>& values, 76 int32_t axis); 77 78 // Lowers SpaceToBatchND to TOSA. 79 llvm::Optional<Value> convertSpaceToBatchNDOp(PatternRewriter& rewriter, 80 Operation* op, Value result_value, 81 Value input_value, 82 Value block_shape_value, 83 Value paddings_value); 84 85 // Lowers BatchToSpaceND to TOSA. 86 llvm::Optional<Value> convertBatchToSpaceNDOp(PatternRewriter& rewriter, 87 Operation* op, Value result_value, 88 Value input_value, 89 Value block_shape_value, 90 Value crops_value); 91 92 // Lowers ExpandDims to TOSA. 93 llvm::Optional<Value> convertExpandDimsOp(PatternRewriter& rewriter, 94 Operation* op, Value result_value, 95 Value input_value, Value dim_value); 96 97 // Lowers Squeeze to TOSA. 98 llvm::Optional<Value> convertSqueezeOp(PatternRewriter& rewriter, Operation* op, 99 Value result_value, Value input_value, 100 SmallVector<int32_t, 8>& squeeze_dims); 101 102 // Lowers ELU to a sequence of TOSA ops. 103 llvm::Optional<Value> convertEluOp(PatternRewriter& rewriter, Operation* op, 104 Value result_value, Value features_value); 105 106 // Lowers Softmax to a sequence of TOSA ops. 107 llvm::Optional<Value> convertSoftmaxOp(PatternRewriter& rewriter, Operation* op, 108 Value result_value, Value logits_value); 109 110 // Lowers LogSoftmax to a sequence of TOSA ops. 111 llvm::Optional<Value> convertLogSoftmaxOp(PatternRewriter& rewriter, 112 Operation* op, Value result_value, 113 Value logits_value); 114 115 // Lowers SpaceToDepth to a sequence of TOSA ops. Supports NHWC. 116 llvm::Optional<Value> convertSpaceToDepthOp(PatternRewriter& rewriter, 117 Operation* op, Value result_value, 118 Value input_value, 119 IntegerAttr block_size_attr, 120 StringAttr data_format); 121 122 // Lowers DepthToSpace to a sequence of TOSA ops. Supports NHWC. 123 llvm::Optional<Value> convertDepthToSpaceOp(PatternRewriter& rewriter, 124 Operation* op, Value result_value, 125 Value input_value, 126 IntegerAttr block_size_attr, 127 StringAttr data_format); 128 129 // Lowers Split to a sequence of TOSA ops. 130 llvm::Optional<ValueRange> convertSplitOp(PatternRewriter& rewriter, 131 Operation* op, Value result_value, 132 Value input_value, int32_t num_split, 133 int32_t axis); 134 135 // Lowers SplitV to a sequence of TOSA ops. 136 llvm::Optional<ValueRange> convertSplitVOp(PatternRewriter& rewriter, 137 Operation* op, Value result_value, 138 Value input_value, 139 SmallVector<int32_t, 4>& size_split, 140 int32_t axis); 141 142 // Lowers StridedSlice to a sequence of TOSA ops. 143 llvm::Optional<Value> convertStridedSliceOp( 144 PatternRewriter& rewriter, Operation* op, Value result_value, 145 Value input_value, Value begin_value, Value end_value, Value strides_value, 146 int32_t begin_mask, int32_t end_mask, int32_t ellipsis_mask, 147 int32_t new_axis_mask, int32_t shrink_axis_mask); 148 149 // Lowers FloorDiv to a sequence of TOSA operators. 150 llvm::Optional<Value> convertFloorDivOp(PatternRewriter& rewriter, 151 Operation* op, Value result_value, 152 Value lhs_value, Value rhs_value); 153 154 // Lowers FloorMod to a sequence of TOSA operators. 155 llvm::Optional<Value> convertFloorModOp(PatternRewriter& rewriter, 156 Operation* op, Value result_value, 157 Value lhs_value, Value rhs_value); 158 159 // Lowers FusedActivation to a sequence of TOSA ops. 160 llvm::Optional<Value> convertFusedActivation(PatternRewriter& rewriter, 161 Operation* op, Value input_value, 162 StringAttr fused_activation_fn); 163 164 // Helper function for implementing quantized divide by power-of-two in TOSA 165 // ops. 166 llvm::Optional<Value> convertRoundingDivideByPOT(PatternRewriter& rewriter, 167 Operation* op, 168 Value input_value, 169 Value rshift_value); 170 171 // Lowers ReduceAll to a sequence of TOSA ops. 172 llvm::Optional<Value> convertReduceAllOp( 173 PatternRewriter& rewriter, Operation* op, RankedTensorType output_type, 174 Value input_value, ElementsAttr axes_elems, bool keep_dims); 175 176 // Lowers ReduceAny to a sequence of TOSA ops. 177 llvm::Optional<Value> convertReduceAnyOp( 178 PatternRewriter& rewriter, Operation* op, RankedTensorType output_type, 179 Value input_value, ElementsAttr axes_elems, bool keep_dims); 180 181 // Lowers ReduceMin to a sequence of TOSA ops. 182 llvm::Optional<Value> convertReduceMinOp( 183 PatternRewriter& rewriter, Operation* op, RankedTensorType output_type, 184 Value input_value, ElementsAttr axes_elems, bool keep_dims); 185 186 // Lowers ReduceMax to a sequence of TOSA ops. 187 llvm::Optional<Value> convertReduceMaxOp( 188 PatternRewriter& rewriter, Operation* op, RankedTensorType output_type, 189 Value input_value, ElementsAttr axes_elems, bool keep_dims); 190 191 // Lowers ReduceProd to a sequence of TOSA ops. 192 llvm::Optional<Value> convertReduceProdOp( 193 PatternRewriter& rewriter, Operation* op, RankedTensorType output_type, 194 Value input_value, ElementsAttr axes_elems, bool keep_dims); 195 196 // Lowers ReduceSum to a sequence of TOSA ops. 197 llvm::Optional<Value> convertReduceSumOp( 198 PatternRewriter& rewriter, Operation* op, RankedTensorType output_type, 199 Value input_value, ElementsAttr axes_elems, bool keep_dims); 200 201 // Lowers ReduceMean to a sequence of TOSA ops. 202 llvm::Optional<Value> convertReduceMeanOp( 203 PatternRewriter& rewriter, Operation* op, RankedTensorType output_type, 204 Value input_value, ElementsAttr axes_elems, bool keep_dims); 205 206 // Lowers ResizeBilinear and ResizeNearestNeighbor to TOSA resize. 207 llvm::Optional<Value> convertResizeOp(PatternRewriter& rewriter, Operation* op, 208 RankedTensorType output_type, 209 Value input_value, StringRef mode); 210 211 // Lowers Quantize to a sequence of TOSA quantization ops. 212 llvm::Optional<Value> convertQuantizeOp(PatternRewriter& rewriter, 213 Operation* op, 214 RankedTensorType output_type, 215 Value input_value, double scale, 216 int64_t zeropoint); 217 218 // Lowers Dequantize to a sequence of TOSA dequantization ops. 219 llvm::Optional<Value> convertDequantizeOp(PatternRewriter& rewriter, 220 Operation* op, 221 RankedTensorType output_type, 222 Value input_value, double scale, 223 int64_t zeropoint); 224 225 // Lowers FakeQuant to a sequence of TOSA quantization ops. 226 llvm::Optional<Value> convertFakeQuantOp(PatternRewriter& rewriter, 227 Operation* op, 228 RankedTensorType output_type, 229 Value input_value, double min, 230 double max, int64_t num_bits, 231 bool narrow_range); 232 233 // Lowers TensorFlow Conv2D to a sequence of TOSA quantization ops. 234 llvm::Optional<Value> convertTFConv2DCommon( 235 PatternRewriter& rewriter, Operation* op, RankedTensorType output_type, 236 Value input, Value filter, Value bias, ArrayAttr strides_attr, 237 ArrayAttr dilations_attr, ArrayAttr explicit_padding_attr, 238 StringRef padding_ref, StringRef data_format_ref); 239 240 }; // namespace tosa 241 }; // namespace mlir 242 243 #endif // TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_COMMON_H 244