1 //===- TestLoopParametricTiling.cpp --- Parametric loop tiling pass -------===// 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 to parametrically tile nests of standard loops. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #include "mlir/Dialect/SCF/IR/SCF.h" 14 #include "mlir/Dialect/SCF/Utils/Utils.h" 15 #include "mlir/IR/Builders.h" 16 #include "mlir/Pass/Pass.h" 17 18 using namespace mlir; 19 20 namespace { 21 22 // Extracts fixed-range loops for top-level loop nests with ranges defined in 23 // the pass constructor. Assumes loops are permutable. 24 class SimpleParametricLoopTilingPass 25 : public PassWrapper<SimpleParametricLoopTilingPass, OperationPass<>> { 26 public: 27 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SimpleParametricLoopTilingPass) 28 29 StringRef getArgument() const final { 30 return "test-extract-fixed-outer-loops"; 31 } 32 StringRef getDescription() const final { 33 return "test application of parametric tiling to the outer loops so that " 34 "the ranges of outer loops become static"; 35 } 36 SimpleParametricLoopTilingPass() = default; 37 SimpleParametricLoopTilingPass(const SimpleParametricLoopTilingPass &) {} 38 explicit SimpleParametricLoopTilingPass(ArrayRef<int64_t> outerLoopSizes) { 39 sizes = outerLoopSizes; 40 } 41 42 void runOnOperation() override { 43 if (sizes.empty()) { 44 emitError( 45 UnknownLoc::get(&getContext()), 46 "missing `test-outer-loop-sizes` pass-option for outer loop sizes"); 47 signalPassFailure(); 48 return; 49 } 50 getOperation()->walk([this](scf::ForOp op) { 51 // Ignore nested loops. 52 if (op->getParentRegion()->getParentOfType<scf::ForOp>()) 53 return; 54 extractFixedOuterLoops(op, sizes); 55 }); 56 } 57 58 ListOption<int64_t> sizes{ 59 *this, "test-outer-loop-sizes", 60 llvm::cl::desc( 61 "fixed number of iterations that the outer loops should have")}; 62 }; 63 } // namespace 64 65 namespace mlir { 66 namespace test { 67 void registerSimpleParametricTilingPass() { 68 PassRegistration<SimpleParametricLoopTilingPass>(); 69 } 70 } // namespace test 71 } // namespace mlir 72