1 //===- TosaToLinalgPass.cpp - Lowering Tosa to Linalg Dialect -------------===//
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 // This transformation pass legalizes Tosa operations to the Linalg dialect.
10 //
11 //===----------------------------------------------------------------------===//
12
13 #include "../PassDetail.h"
14 #include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
15 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
16 #include "mlir/Dialect/Linalg/IR/Linalg.h"
17 #include "mlir/Dialect/Math/IR/Math.h"
18 #include "mlir/Dialect/SCF/IR/SCF.h"
19 #include "mlir/Dialect/Tensor/IR/Tensor.h"
20 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
21 #include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
22 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
23 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
24 #include "mlir/IR/PatternMatch.h"
25 #include "mlir/Pass/PassManager.h"
26 #include "mlir/Transforms/DialectConversion.h"
27 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
28 #include "mlir/Transforms/Passes.h"
29
30 using namespace mlir;
31
32 namespace {
33 struct TosaToLinalg : public TosaToLinalgBase<TosaToLinalg> {
34 public:
getDependentDialects__anona46ef7560111::TosaToLinalg35 void getDependentDialects(DialectRegistry ®istry) const override {
36 registry
37 .insert<arith::ArithmeticDialect, linalg::LinalgDialect,
38 math::MathDialect, tensor::TensorDialect, scf::SCFDialect>();
39 }
40
runOnOperation__anona46ef7560111::TosaToLinalg41 void runOnOperation() override {
42 RewritePatternSet patterns(&getContext());
43 ConversionTarget target(getContext());
44 target.addLegalDialect<linalg::LinalgDialect, tensor::TensorDialect,
45 scf::SCFDialect>();
46 target.addIllegalDialect<tosa::TosaDialect>();
47
48 // Not every TOSA op can be legalized to linalg.
49 target.addLegalOp<tosa::ApplyScaleOp>();
50 target.addLegalOp<tosa::IfOp>();
51 target.addLegalOp<tosa::ConstOp>();
52 target.addLegalOp<tosa::WhileOp>();
53 target.addLegalOp<tosa::SliceOp>();
54
55 target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
56
57 FunctionOpInterface func = getOperation();
58 mlir::tosa::populateTosaToLinalgConversionPatterns(&patterns);
59 if (failed(applyFullConversion(func, target, std::move(patterns))))
60 signalPassFailure();
61 }
62 };
63 } // namespace
64
createTosaToLinalg()65 std::unique_ptr<Pass> mlir::tosa::createTosaToLinalg() {
66 return std::make_unique<TosaToLinalg>();
67 }
68
addTosaToLinalgPasses(OpPassManager & pm,bool disableTosaDecompositions)69 void mlir::tosa::addTosaToLinalgPasses(OpPassManager &pm,
70 bool disableTosaDecompositions) {
71 // Optional decompositions are designed to benefit linalg.
72 if (!disableTosaDecompositions)
73 pm.addNestedPass<func::FuncOp>(tosa::createTosaOptionalDecompositions());
74 pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
75
76 pm.addNestedPass<func::FuncOp>(tosa::createTosaMakeBroadcastablePass());
77 pm.addNestedPass<func::FuncOp>(tosa::createTosaToLinalgNamed());
78 pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
79 // TODO: Remove pass that operates on const tensor and enable optionality
80 pm.addNestedPass<func::FuncOp>(tosa::createTosaLayerwiseConstantFoldPass());
81 pm.addNestedPass<func::FuncOp>(tosa::createTosaMakeBroadcastablePass());
82 pm.addNestedPass<func::FuncOp>(tosa::createTosaToLinalg());
83 }
84