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 &registry) 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