1 //===- TestTilingInterface.cpp - Test tiling using `TilingInterface` -----===// 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 file implements a pass for testing tiling operations using 10 // `TilingInterface`. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Dialect/Affine/IR/AffineOps.h" 15 #include "mlir/Dialect/Func/IR/FuncOps.h" 16 #include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h" 17 #include "mlir/Dialect/Linalg/Transforms/Transforms.h" 18 #include "mlir/Dialect/MemRef/IR/MemRef.h" 19 #include "mlir/Dialect/SCF/IR/SCF.h" 20 #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h" 21 #include "mlir/Dialect/Tensor/IR/Tensor.h" 22 #include "mlir/Interfaces/TilingInterface.h" 23 #include "mlir/Pass/Pass.h" 24 #include "mlir/Pass/PassManager.h" 25 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 26 #include "llvm/ADT/TypeSwitch.h" 27 28 using namespace mlir; 29 30 namespace { 31 32 /// Pattern for testing `TileUsingSCFForOp` pattern (that tiles operations using 33 /// the `TilingInterface` with `scf.for` ops for iterating over the tiles) while 34 /// using a `filter` to avoid recursive application. 35 struct TestTileUsingSCFForOpWithFilter : public scf::TileUsingSCFForOp { 36 TestTileUsingSCFForOpWithFilter(MLIRContext *context, 37 scf::SCFTilingOptions options, 38 linalg::LinalgTransformationFilter filter = 39 linalg::LinalgTransformationFilter(), 40 PatternBenefit benefit = 1) 41 : scf::TileUsingSCFForOp(context, options, benefit), filter(filter) {} 42 43 /// Construct a generic pattern applied to `opName`. 44 TestTileUsingSCFForOpWithFilter(StringRef opName, MLIRContext *context, 45 scf::SCFTilingOptions options, 46 linalg::LinalgTransformationFilter filter = 47 linalg::LinalgTransformationFilter(), 48 PatternBenefit benefit = 1) 49 : scf::TileUsingSCFForOp(context, options, benefit), filter(filter) {} 50 51 LogicalResult matchAndRewrite(TilingInterface op, 52 PatternRewriter &rewriter) const override { 53 if (failed(filter.checkAndNotify(rewriter, op))) 54 return failure(); 55 56 auto tilingResult = returningMatchAndRewrite(op, rewriter); 57 if (failed(tilingResult)) { 58 return failure(); 59 } 60 filter.replaceLinalgTransformationFilter(rewriter, tilingResult->tiledOp); 61 return success(); 62 } 63 64 private: 65 linalg::LinalgTransformationFilter filter; 66 }; 67 68 /// Pattern for testing `TileConsumerAndFUseProducersUsingSCFForOp` pattern 69 /// (that tiles and fuses operations using the `TilingInterface` with `scf.for` 70 /// ops for iterating over the tiles) while using a `filter` to avoid recursive 71 /// application. 72 struct TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter 73 : public scf::TileConsumerAndFuseProducersUsingSCFForOp { 74 TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter( 75 MLIRContext *context, scf::SCFTilingOptions options, 76 linalg::LinalgTransformationFilter filter = 77 linalg::LinalgTransformationFilter(), 78 PatternBenefit benefit = 1) 79 : scf::TileConsumerAndFuseProducersUsingSCFForOp(context, options, 80 benefit), 81 filter(filter) {} 82 83 /// Construct a generic pattern applied to `opName`. 84 TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter( 85 StringRef opName, MLIRContext *context, scf::SCFTilingOptions options, 86 linalg::LinalgTransformationFilter filter = 87 linalg::LinalgTransformationFilter(), 88 PatternBenefit benefit = 1) 89 : scf::TileConsumerAndFuseProducersUsingSCFForOp(context, options, 90 benefit), 91 filter(filter) {} 92 93 LogicalResult matchAndRewrite(TilingInterface op, 94 PatternRewriter &rewriter) const override { 95 if (failed(filter.checkAndNotify(rewriter, op))) 96 return failure(); 97 98 auto tileAndFuseResult = returningMatchAndRewrite(op, rewriter); 99 if (failed(tileAndFuseResult)) { 100 return failure(); 101 } 102 filter.replaceLinalgTransformationFilter( 103 rewriter, tileAndFuseResult->tiledAndFusedOps.front()); 104 return success(); 105 } 106 107 private: 108 linalg::LinalgTransformationFilter filter; 109 }; 110 111 /// Test pass for testing the use of `TilingInterface`. 112 struct TestTilingInterfacePass 113 : public PassWrapper<TestTilingInterfacePass, OperationPass<func::FuncOp>> { 114 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTilingInterfacePass) 115 116 TestTilingInterfacePass() = default; 117 TestTilingInterfacePass(const TestTilingInterfacePass &pass) 118 : PassWrapper(pass) {} 119 void getDependentDialects(DialectRegistry ®istry) const override { 120 registry.insert<AffineDialect, memref::MemRefDialect, scf::SCFDialect, 121 tensor::TensorDialect>(); 122 linalg::registerTilingInterfaceExternalModels(registry); 123 } 124 StringRef getArgument() const final { return "test-tiling-interface"; } 125 StringRef getDescription() const final { 126 return "Test tiling using TilingInterface"; 127 } 128 129 Option<bool> testTiling{ 130 *this, "tile-using-scf-for", 131 llvm::cl::desc( 132 "Test tiling using TilingInterface with scf.for operations"), 133 llvm::cl::init(false)}; 134 135 Option<bool> testTileConsumerAndFuseProducer{ 136 *this, "tile-consumer-and-fuse-producer-using-scf-for", 137 llvm::cl::desc("Test tile and fuse transformation using TilingInterface " 138 "with scf.for operations"), 139 llvm::cl::init(false)}; 140 141 void runOnOperation() override; 142 143 private: 144 void addTestPatterns(MLIRContext *context, RewritePatternSet &patterns); 145 }; 146 } // namespace 147 148 template <class Pattern> 149 static void 150 addPatternForTiling(MLIRContext *context, ArrayRef<int64_t> tileSizes, 151 StringRef filterName, RewritePatternSet &patterns) { 152 scf::SCFTilingOptions tilingOptions; 153 tilingOptions.setTileSizes(tileSizes); 154 linalg::LinalgTransformationFilter filter( 155 StringAttr::get(context, filterName), StringAttr::get(context, "tiled")); 156 patterns.add<Pattern>(context, tilingOptions, filter); 157 } 158 159 void TestTilingInterfacePass::addTestPatterns(MLIRContext *context, 160 RewritePatternSet &patterns) { 161 if (testTiling) { 162 // 1. Tiling M and N dims of `linalg.matmul` on tensors. 163 addPatternForTiling<TestTileUsingSCFForOpWithFilter>( 164 context, {10, 20}, "simple_gemm", patterns); 165 // 2. Tiling M, N and K of `linalg.matmul` on buffers. 166 addPatternForTiling<TestTileUsingSCFForOpWithFilter>( 167 context, {10, 20, 30}, "simple_gemm_memref", patterns); 168 // 3. Tiling 3D parallel generic op which implements a transpose 169 addPatternForTiling<TestTileUsingSCFForOpWithFilter>( 170 context, {10, 0, 20}, "parallel_generic_transpose", patterns); 171 // 4. Tiling 2D conv op. 172 addPatternForTiling<TestTileUsingSCFForOpWithFilter>( 173 context, {0, 0, 0, 0, 10, 20, 30}, "simple_conv", patterns); 174 // 5. Tiling a simple op with `linalg.index` inside. 175 addPatternForTiling<TestTileUsingSCFForOpWithFilter>( 176 context, {10, 20}, "indexed_semantics", patterns); 177 return; 178 } 179 if (testTileConsumerAndFuseProducer) { 180 // 1. Tile and fuse of gemm with bias-add operation. 181 addPatternForTiling< 182 TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>( 183 context, {10, 20}, "fusion", patterns); 184 addPatternForTiling< 185 TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>( 186 context, {10}, "gemm_fusion", patterns); 187 return; 188 } 189 } 190 191 void TestTilingInterfacePass::runOnOperation() { 192 MLIRContext *context = &getContext(); 193 194 RewritePatternSet tilingPatterns(context); 195 addTestPatterns(context, tilingPatterns); 196 if (failed(applyPatternsAndFoldGreedily(getOperation(), 197 std::move(tilingPatterns)))) 198 return signalPassFailure(); 199 } 200 201 namespace mlir { 202 namespace test { 203 void registerTestTilingInterface() { 204 PassRegistration<TestTilingInterfacePass>(); 205 } 206 } // namespace test 207 } // namespace mlir 208