1 //===- TosaCanonicalizations.cpp - Canonicalization patterns & folders ----===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // \file
10 // TOSA canonicalization patterns and folders.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "mlir/Dialect/Quant/QuantOps.h"
15 #include "mlir/Dialect/Tensor/IR/Tensor.h"
16 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
17 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
18 #include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
19 #include "mlir/IR/BuiltinTypes.h"
20 #include "mlir/IR/DialectImplementation.h"
21 #include "mlir/IR/Matchers.h"
22 #include "mlir/IR/PatternMatch.h"
23 #include "mlir/Transforms/FoldUtils.h"
24 #include "mlir/Transforms/InliningUtils.h"
25 #include "mlir/Transforms/RegionUtils.h"
26 #include "llvm/ADT/DenseMap.h"
27 #include "llvm/ADT/TypeSwitch.h"
28 
29 using namespace mlir;
30 using namespace mlir::tosa;
31 
32 //===----------------------------------------------------------------------===//
33 // Operator Canonicalizers.
34 //===----------------------------------------------------------------------===//
35 
36 struct ConcatOptimization : public OpRewritePattern<tosa::ConcatOp> {
37   using OpRewritePattern<tosa::ConcatOp>::OpRewritePattern;
38 
matchAndRewriteConcatOptimization39   LogicalResult matchAndRewrite(tosa::ConcatOp op,
40                                 PatternRewriter &rewriter) const override {
41     if (op.getInput1().size() != 1)
42       return failure();
43     if (op.getInput1().front().getType() != op.getType()) {
44       rewriter
45           .replaceOpWithNewOp<tensor::CastOp>(op, op.getType(),
46                                               op.getInput1().front())
47           .getResult();
48       return success();
49     }
50 
51     rewriter.replaceOp(op, op.getInput1().front());
52     return success();
53   }
54 };
55 
getCanonicalizationPatterns(RewritePatternSet & results,MLIRContext * context)56 void ConcatOp::getCanonicalizationPatterns(RewritePatternSet &results,
57                                            MLIRContext *context) {
58   results.add<ConcatOptimization>(context);
59 }
60 
61 struct ReshapeReshapeOptimization : public OpRewritePattern<tosa::ReshapeOp> {
62   using OpRewritePattern<tosa::ReshapeOp>::OpRewritePattern;
63 
matchAndRewriteReshapeReshapeOptimization64   LogicalResult matchAndRewrite(tosa::ReshapeOp op,
65                                 PatternRewriter &rewriter) const override {
66     Value input = op.getInput1();
67     Operation *definingOp = input.getDefiningOp();
68     if (!definingOp)
69       return failure();
70 
71     if (tosa::ReshapeOp reshapeOp = dyn_cast<tosa::ReshapeOp>(definingOp)) {
72       rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
73           op, op.getType(), reshapeOp.getInput1(), op.getNewShape());
74       return success();
75     }
76 
77     return failure();
78   }
79 };
80 
81 struct ReshapeConstOptimization : public OpRewritePattern<tosa::ReshapeOp> {
82   using OpRewritePattern<tosa::ReshapeOp>::OpRewritePattern;
83 
matchAndRewriteReshapeConstOptimization84   LogicalResult matchAndRewrite(tosa::ReshapeOp op,
85                                 PatternRewriter &rewriter) const override {
86     Value input = op.getInput1();
87     ArrayAttr newShape = op.getNewShape();
88 
89     // Check if input is constant
90     DenseElementsAttr inputAttr;
91     if (!matchPattern(input, m_Constant(&inputAttr)))
92       return failure();
93 
94     // Check if has >1 consumer and is not splat
95     if (!input.hasOneUse() && !inputAttr.isSplat())
96       return failure();
97 
98     // Grab the new shape
99     SmallVector<int64_t> newShapeValues = llvm::to_vector<6>(
100         llvm::map_range(newShape.getValue(), [](const Attribute &val) {
101           return val.cast<IntegerAttr>().getValue().getSExtValue();
102         }));
103 
104     // Build new const op with correct output shape
105     ShapedType inputShape = input.getType().cast<ShapedType>();
106     DenseElementsAttr outputAttr =
107         inputAttr.reshape(inputShape.clone(newShapeValues));
108     rewriter.replaceOpWithNewOp<tosa::ConstOp>(op, outputAttr.getType(),
109                                                outputAttr);
110     return success();
111   }
112 };
113 
getCanonicalizationPatterns(RewritePatternSet & results,MLIRContext * context)114 void ReshapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
115                                             MLIRContext *context) {
116   results.add<ReshapeReshapeOptimization>(context);
117   results.add<ReshapeConstOptimization>(context);
118 }
119 
canonicalize(SelectOp op,PatternRewriter & rewriter)120 LogicalResult SelectOp::canonicalize(SelectOp op, PatternRewriter &rewriter) {
121   auto notOp = op.getPred().getDefiningOp<tosa::LogicalNotOp>();
122   if (!notOp)
123     return failure();
124   rewriter.updateRootInPlace(op, [&]() {
125     op.getOperation()->setOperands(
126         {notOp.getInput1(), op.getOnFalse(), op.getOnTrue()});
127   });
128   return success();
129 }
130 
131 struct NoOpOptimization : public OpRewritePattern<tosa::TransposeOp> {
132   using OpRewritePattern::OpRewritePattern;
133 
matchAndRewriteNoOpOptimization134   LogicalResult matchAndRewrite(tosa::TransposeOp op,
135                                 PatternRewriter &rewriter) const override {
136     auto perm = op.getPerms();
137 
138     DenseIntElementsAttr permAttr;
139     if (!matchPattern(perm, m_Constant(&permAttr))) {
140       return failure();
141     }
142 
143     SmallVector<int64_t> permValues = llvm::to_vector<6>(
144         llvm::map_range(permAttr.getValues<APInt>(),
145                         [](const APInt &val) { return val.getSExtValue(); }));
146 
147     for (int i = 0, s = permValues.size(); i < s; i++) {
148       if (i != permValues[i]) {
149         return failure();
150       }
151     }
152 
153     rewriter.replaceOp(op, op.getInput1());
154     return success();
155   }
156 };
157 
getCanonicalizationPatterns(RewritePatternSet & results,MLIRContext * context)158 void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
159                                               MLIRContext *context) {
160   results.add<NoOpOptimization>(context);
161 }
162 
163 struct AddZeroOptimization : public OpRewritePattern<tosa::AddOp> {
164   using OpRewritePattern::OpRewritePattern;
165 
matchAndRewriteAddZeroOptimization166   LogicalResult matchAndRewrite(tosa::AddOp op,
167                                 PatternRewriter &rewriter) const override {
168     auto input1 = op.getInput1();
169     auto input2 = op.getInput2();
170 
171     DenseElementsAttr input1Attr;
172     if (matchPattern(input1, m_Constant(&input1Attr)) && input1Attr.isSplat() &&
173         input2.getType() == op.getType()) {
174       if (input1Attr.getType().getElementType().isa<IntegerType>() &&
175           input1Attr.getSplatValue<APInt>().isZero()) {
176         rewriter.replaceOp(op, op.getInput2());
177         return success();
178       }
179     }
180 
181     DenseElementsAttr input2Attr;
182     if (matchPattern(input2, m_Constant(&input2Attr)) && input2Attr.isSplat() &&
183         input1.getType() == op.getType()) {
184       if (input2Attr.getType().getElementType().isa<IntegerType>() &&
185           input2Attr.getSplatValue<APInt>().isZero()) {
186         rewriter.replaceOp(op, op.getInput1());
187         return success();
188       }
189     }
190 
191     return failure();
192   }
193 };
194 
getCanonicalizationPatterns(RewritePatternSet & results,MLIRContext * context)195 void AddOp::getCanonicalizationPatterns(RewritePatternSet &results,
196                                         MLIRContext *context) {
197   results.add<AddZeroOptimization>(context);
198 }
199 
200 struct MulOneOptimization : public OpRewritePattern<tosa::MulOp> {
201   using OpRewritePattern::OpRewritePattern;
202 
matchAndRewriteMulOneOptimization203   LogicalResult matchAndRewrite(tosa::MulOp op,
204                                 PatternRewriter &rewriter) const override {
205     auto input1 = op.getInput1();
206     auto input2 = op.getInput2();
207 
208     DenseElementsAttr input1Attr;
209     if (matchPattern(input1, m_Constant(&input1Attr)) && input1Attr.isSplat() &&
210         input2.getType() == op.getType()) {
211       if (input1Attr.getType().getElementType().isa<FloatType>() &&
212           input1Attr.getSplatValue<APFloat>().isExactlyValue(1)) {
213         rewriter.replaceOp(op, op.getInput2());
214         return success();
215       }
216 
217       if (input1Attr.getType().getElementType().isa<IntegerType>() &&
218           matchPattern(input1, m_One())) {
219         rewriter.replaceOp(op, op.getInput2());
220         return success();
221       }
222     }
223 
224     DenseElementsAttr input2Attr;
225     if (matchPattern(input2, m_Constant(&input2Attr)) && input2Attr.isSplat() &&
226         input1.getType() == op.getType()) {
227       if (input2Attr.getType().getElementType().isa<FloatType>() &&
228           input2Attr.getSplatValue<APFloat>().isExactlyValue(1)) {
229         rewriter.replaceOp(op, op.getInput1());
230         return success();
231       }
232 
233       if (input2Attr.getType().getElementType().isa<IntegerType>() &&
234           matchPattern(input2, m_One())) {
235         rewriter.replaceOp(op, op.getInput1());
236         return success();
237       }
238     }
239 
240     return failure();
241   }
242 };
243 
getCanonicalizationPatterns(RewritePatternSet & results,MLIRContext * context)244 void MulOp::getCanonicalizationPatterns(RewritePatternSet &results,
245                                         MLIRContext *context) {
246   results.add<MulOneOptimization>(context);
247 }
248 
249 struct MaterializePadValue : public OpRewritePattern<tosa::PadOp> {
250   using OpRewritePattern::OpRewritePattern;
251 
matchAndRewriteMaterializePadValue252   LogicalResult matchAndRewrite(tosa::PadOp op,
253                                 PatternRewriter &rewriter) const override {
254     if (op.getPadConst())
255       return failure();
256 
257     auto input = op.getInput1();
258     auto padding = op.getPadding();
259 
260     ShapedType inputTy = input.getType().cast<ShapedType>();
261     Type elementTy = inputTy.getElementType();
262 
263     Attribute constantAttr;
264     if (elementTy.isa<FloatType>()) {
265       constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
266     } else if (elementTy.isa<IntegerType>() && !op.getQuantizationInfo()) {
267       constantAttr = rewriter.getIntegerAttr(elementTy, 0);
268     } else if (elementTy.isa<IntegerType>() && op.getQuantizationInfo()) {
269       auto value = op.getQuantizationInfo()->getInputZp();
270       constantAttr = rewriter.getIntegerAttr(elementTy, value);
271     }
272 
273     if (!constantAttr) {
274       return rewriter.notifyMatchFailure(
275           op,
276           "tosa.pad to linalg lowering encountered an unknown element type");
277     }
278 
279     auto denseAttr = DenseElementsAttr::get(
280         RankedTensorType::get({}, elementTy), constantAttr);
281     auto constantVal = rewriter.create<tosa::ConstOp>(
282         op.getLoc(), denseAttr.getType(), denseAttr);
283 
284     rewriter.replaceOpWithNewOp<tosa::PadOp>(
285         op, op.getType(), ValueRange{input, padding, constantVal},
286         op->getAttrs());
287     return success();
288   }
289 };
290 
getCanonicalizationPatterns(RewritePatternSet & results,MLIRContext * context)291 void PadOp::getCanonicalizationPatterns(RewritePatternSet &results,
292                                         MLIRContext *context) {
293   results.add<MaterializePadValue>(context);
294 }
295 
296 struct MaxPool2dIsNoOp : public OpRewritePattern<tosa::MaxPool2dOp> {
297   using OpRewritePattern::OpRewritePattern;
298 
matchAndRewriteMaxPool2dIsNoOp299   LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
300                                 PatternRewriter &rewriter) const override {
301     Value input = op.getInput();
302     Value output = op.getOutput();
303     ShapedType inputType = input.getType().cast<ShapedType>();
304     ShapedType outputType = output.getType().cast<ShapedType>();
305 
306     if (!inputType.hasStaticShape() || !outputType.hasStaticShape()) {
307       return failure();
308     }
309 
310     // If the output and input shapes are 1x1, then this is a no op.
311     ArrayRef<int64_t> outputShape = outputType.getShape();
312     if (outputShape[1] != 1 || outputShape[2] != 1) {
313       return failure();
314     }
315 
316     ArrayRef<int64_t> inputShape = inputType.getShape();
317     if (inputShape[1] != 1 || inputShape[2] != 1) {
318       return failure();
319     }
320 
321     rewriter.replaceOp(op, input);
322     return success();
323   }
324 };
325 
getCanonicalizationPatterns(RewritePatternSet & results,MLIRContext * context)326 void MaxPool2dOp::getCanonicalizationPatterns(RewritePatternSet &results,
327                                               MLIRContext *context) {
328   results.add<MaxPool2dIsNoOp>(context);
329 }
330 
331 struct ClampIsNoOp : public OpRewritePattern<tosa::ClampOp> {
332   using OpRewritePattern::OpRewritePattern;
333 
matchAndRewriteClampIsNoOp334   LogicalResult matchAndRewrite(tosa::ClampOp op,
335                                 PatternRewriter &rewriter) const override {
336     Value input = op.getInput();
337     auto inputType =
338         op.getInput().getType().template dyn_cast<RankedTensorType>();
339     auto inputElementType = inputType.getElementType();
340 
341     if (!inputType.hasStaticShape()) {
342       return failure();
343     }
344 
345     if (inputElementType.isF32()) {
346       auto minClamp = op.getMinFp();
347       auto maxClamp = op.getMaxFp();
348       bool isMin = (minClamp.isLargest() || minClamp.isInfinity()) &&
349                    minClamp.isNegative();
350       bool isMax = (maxClamp.isLargest() || maxClamp.isInfinity()) &&
351                    !maxClamp.isNegative();
352 
353       if (isMin && isMax) {
354         rewriter.replaceOp(op, input);
355         return success();
356       }
357       return failure();
358     }
359 
360     if (inputElementType.isUnsignedInteger()) {
361       int64_t minClamp = op.getMinInt();
362       int64_t maxClamp = op.getMaxInt();
363 
364       int64_t intMin =
365           APInt::getMinValue(inputElementType.getIntOrFloatBitWidth())
366               .getZExtValue();
367       int64_t intMax =
368           APInt::getMaxValue(inputElementType.getIntOrFloatBitWidth())
369               .getZExtValue();
370 
371       if (minClamp <= intMin && maxClamp >= intMax) {
372         rewriter.replaceOp(op, input);
373         return success();
374       }
375       return failure();
376     }
377 
378     if (inputElementType.isa<IntegerType>()) {
379       int64_t minClamp = op.getMinInt();
380       int64_t maxClamp = op.getMaxInt();
381 
382       int64_t intMin =
383           APInt::getSignedMinValue(inputElementType.getIntOrFloatBitWidth())
384               .getSExtValue();
385       int64_t intMax =
386           APInt::getSignedMaxValue(inputElementType.getIntOrFloatBitWidth())
387               .getSExtValue();
388 
389       if (minClamp <= intMin && maxClamp >= intMax) {
390         rewriter.replaceOp(op, input);
391         return success();
392       }
393       return failure();
394     }
395 
396     return failure();
397   }
398 };
399 
400 struct ClampClampOptimization : public OpRewritePattern<tosa::ClampOp> {
401   using OpRewritePattern<tosa::ClampOp>::OpRewritePattern;
402 
matchAndRewriteClampClampOptimization403   LogicalResult matchAndRewrite(tosa::ClampOp op,
404                                 PatternRewriter &rewriter) const override {
405     Value input = op.getInput();
406 
407     Operation *definingOp = input.getDefiningOp();
408     if (!definingOp)
409       return failure();
410 
411     if (tosa::ClampOp clampOp = dyn_cast<tosa::ClampOp>(definingOp)) {
412       auto minFp = std::max(op.getMinFp(), clampOp.getMinFp()).convertToFloat();
413       auto maxFp = std::min(op.getMaxFp(), clampOp.getMaxFp()).convertToFloat();
414 
415       auto minInt = std::max(op.getMinInt(), clampOp.getMinInt());
416       auto maxInt = std::min(op.getMaxInt(), clampOp.getMaxInt());
417 
418       rewriter.replaceOpWithNewOp<tosa::ClampOp>(
419           op, op.getType(), clampOp.getInput(),
420           rewriter.getI64IntegerAttr(minInt),
421           rewriter.getI64IntegerAttr(maxInt), rewriter.getF32FloatAttr(minFp),
422           rewriter.getF32FloatAttr(maxFp));
423       return success();
424     }
425 
426     return failure();
427   }
428 };
429 
getCanonicalizationPatterns(RewritePatternSet & results,MLIRContext * context)430 void ClampOp::getCanonicalizationPatterns(RewritePatternSet &results,
431                                           MLIRContext *context) {
432   results.add<ClampIsNoOp>(context);
433   results.add<ClampClampOptimization>(context);
434 }
435 
436 //===----------------------------------------------------------------------===//
437 // Operator Folders.
438 //===----------------------------------------------------------------------===//
439 
fold(ArrayRef<Attribute> operands)440 OpFoldResult CastOp::fold(ArrayRef<Attribute> operands) {
441   if (getInput().getType() == getType())
442     return getInput();
443   return {};
444 }
445 
fold(ArrayRef<Attribute> operands)446 OpFoldResult ConstOp::fold(ArrayRef<Attribute> operands) {
447   assert(operands.empty() && "constant has no operands");
448   return getValueAttr();
449 }
450 
451 #define REDUCE_FOLDER(OP)                                                      \
452   OpFoldResult OP::fold(ArrayRef<Attribute> operands) {                        \
453     ShapedType inputTy = input().getType().cast<ShapedType>();                 \
454     if (!inputTy.hasRank())                                                    \
455       return {};                                                               \
456     if (inputTy.getDimSize(axis()) == 1)                                       \
457       return input();                                                          \
458     return {};                                                                 \
459   }
460 
461 REDUCE_FOLDER(ReduceAllOp)
REDUCE_FOLDER(ReduceAnyOp)462 REDUCE_FOLDER(ReduceAnyOp)
463 REDUCE_FOLDER(ReduceMaxOp)
464 REDUCE_FOLDER(ReduceMinOp)
465 REDUCE_FOLDER(ReduceProdOp)
466 REDUCE_FOLDER(ReduceSumOp)
467 #undef REDUCE_FOLDER
468 
469 OpFoldResult ReshapeOp::fold(ArrayRef<Attribute> operands) {
470   auto inputTy = getInput1().getType().dyn_cast<RankedTensorType>();
471   auto outputTy = getType().dyn_cast<RankedTensorType>();
472 
473   if (!inputTy || !outputTy || inputTy != outputTy)
474     return {};
475   return getInput1();
476 }
477 
fold(ArrayRef<Attribute> operands)478 OpFoldResult PadOp::fold(ArrayRef<Attribute> operands) {
479   // If the pad is all zeros we can fold this operation away.
480   if (operands[1]) {
481     auto densePad = operands[1].cast<DenseElementsAttr>();
482     if (densePad.isSplat() && densePad.getSplatValue<APInt>().isZero()) {
483       return getInput1();
484     }
485   }
486 
487   return {};
488 }
489 
fold(ArrayRef<Attribute> operands)490 OpFoldResult SliceOp::fold(ArrayRef<Attribute> operands) {
491   auto inputTy = getInput().getType().dyn_cast<RankedTensorType>();
492   auto outputTy = getType().dyn_cast<RankedTensorType>();
493 
494   if (!inputTy || !outputTy || inputTy != outputTy)
495     return {};
496   if (inputTy.hasStaticShape())
497     return getInput();
498 
499   return {};
500 }
501 
fold(ArrayRef<Attribute> operands)502 OpFoldResult tosa::SelectOp::fold(ArrayRef<Attribute> operands) {
503   if (getOnTrue() == getOnFalse())
504     return getOnTrue();
505 
506   auto predicate = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
507   if (!predicate)
508     return {};
509 
510   if (!predicate.isSplat())
511     return {};
512   return predicate.getSplatValue<APInt>().getBoolValue() ? getOnTrue()
513                                                          : getOnFalse();
514 }
515 
fold(ArrayRef<Attribute> operands)516 OpFoldResult TileOp::fold(ArrayRef<Attribute> operands) {
517   bool allOnes = true;
518   for (Attribute val : getMultiples().getValue()) {
519     allOnes = allOnes && val.cast<IntegerAttr>().getValue().getSExtValue() == 1;
520   }
521 
522   if (allOnes && getInput1().getType() == getType())
523     return getInput1();
524   return {};
525 }
526 
fold(ArrayRef<Attribute> operands)527 OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
528   if (!operands[1])
529     return {};
530 
531   // Transposing splat values just means reshaping.
532   if (auto input = operands[0].dyn_cast_or_null<DenseElementsAttr>()) {
533     if (input.isSplat())
534       return input.reshape(getType().cast<ShapedType>());
535   }
536 
537   auto perms = llvm::to_vector<6>(llvm::map_range(
538       operands[1].cast<DenseIntElementsAttr>().getValues<APInt>(),
539       [](const APInt &val) { return val.getSExtValue(); }));
540 
541   if (llvm::equal(llvm::seq<int64_t>(0, perms.size()), perms) &&
542       getInput1().getType() == getType())
543     return getInput1();
544   return {};
545 }
546