261 lines · cpp
1//===- TestSCFUtils.cpp --- Pass to test independent SCF dialect utils ----===//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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a pass to test SCF dialect utils.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Arith/IR/Arith.h"14#include "mlir/Dialect/Func/IR/FuncOps.h"15#include "mlir/Dialect/MemRef/IR/MemRef.h"16#include "mlir/Dialect/SCF/IR/SCF.h"17#include "mlir/Dialect/SCF/Transforms/Patterns.h"18#include "mlir/Dialect/SCF/Utils/Utils.h"19#include "mlir/IR/Builders.h"20#include "mlir/IR/PatternMatch.h"21#include "mlir/Pass/Pass.h"22#include "mlir/Transforms/GreedyPatternRewriteDriver.h"23 24using namespace mlir;25 26namespace {27struct TestSCFForUtilsPass28 : public PassWrapper<TestSCFForUtilsPass, OperationPass<func::FuncOp>> {29 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSCFForUtilsPass)30 31 StringRef getArgument() const final { return "test-scf-for-utils"; }32 StringRef getDescription() const final { return "test scf.for utils"; }33 explicit TestSCFForUtilsPass() = default;34 TestSCFForUtilsPass(const TestSCFForUtilsPass &pass) : PassWrapper(pass) {}35 36 Option<bool> testReplaceWithNewYields{37 *this, "test-replace-with-new-yields",38 llvm::cl::desc("Test replacing a loop with a new loop that returns new "39 "additional yield values"),40 llvm::cl::init(false)};41 42 void runOnOperation() override {43 func::FuncOp func = getOperation();44 45 // Annotate every loop-like operation with the static trip count.46 func.walk([&](LoopLikeOpInterface loopOp) {47 std::optional<APInt> tripCount = loopOp.getStaticTripCount();48 if (tripCount.has_value())49 loopOp->setDiscardableAttr(50 "test.trip-count",51 IntegerAttr::get(IntegerType::get(&getContext(),52 tripCount.value().getBitWidth()),53 tripCount.value().getSExtValue()));54 else55 loopOp->setDiscardableAttr("test.trip-count",56 StringAttr::get(&getContext(), "none"));57 });58 59 if (testReplaceWithNewYields) {60 func.walk([&](scf::ForOp forOp) {61 if (forOp.getNumResults() == 0)62 return;63 auto newInitValues = forOp.getInitArgs();64 if (newInitValues.empty())65 return;66 SmallVector<Value> oldYieldValues =67 llvm::to_vector(forOp.getYieldedValues());68 NewYieldValuesFn fn = [&](OpBuilder &b, Location loc,69 ArrayRef<BlockArgument> newBBArgs) {70 SmallVector<Value> newYieldValues;71 for (auto yieldVal : oldYieldValues) {72 newYieldValues.push_back(73 arith::AddFOp::create(b, loc, yieldVal, yieldVal));74 }75 return newYieldValues;76 };77 IRRewriter rewriter(forOp.getContext());78 if (failed(forOp.replaceWithAdditionalYields(79 rewriter, newInitValues, /*replaceInitOperandUsesInLoop=*/true,80 fn)))81 signalPassFailure();82 });83 }84 }85};86 87struct TestSCFIfUtilsPass88 : public PassWrapper<TestSCFIfUtilsPass, OperationPass<ModuleOp>> {89 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSCFIfUtilsPass)90 91 StringRef getArgument() const final { return "test-scf-if-utils"; }92 StringRef getDescription() const final { return "test scf.if utils"; }93 explicit TestSCFIfUtilsPass() = default;94 95 void getDependentDialects(DialectRegistry ®istry) const override {96 registry.insert<func::FuncDialect>();97 }98 99 void runOnOperation() override {100 int count = 0;101 getOperation().walk([&](scf::IfOp ifOp) {102 auto strCount = std::to_string(count++);103 func::FuncOp thenFn, elseFn;104 OpBuilder b(ifOp);105 IRRewriter rewriter(b);106 if (failed(outlineIfOp(rewriter, ifOp, &thenFn,107 std::string("outlined_then") + strCount, &elseFn,108 std::string("outlined_else") + strCount))) {109 this->signalPassFailure();110 return WalkResult::interrupt();111 }112 return WalkResult::advance();113 });114 }115};116 117static const StringLiteral kTestPipeliningLoopMarker =118 "__test_pipelining_loop__";119static const StringLiteral kTestPipeliningStageMarker =120 "__test_pipelining_stage__";121/// Marker to express the order in which operations should be after122/// pipelining.123static const StringLiteral kTestPipeliningOpOrderMarker =124 "__test_pipelining_op_order__";125 126static const StringLiteral kTestPipeliningAnnotationPart =127 "__test_pipelining_part";128static const StringLiteral kTestPipeliningAnnotationIteration =129 "__test_pipelining_iteration";130 131struct TestSCFPipeliningPass132 : public PassWrapper<TestSCFPipeliningPass, OperationPass<func::FuncOp>> {133 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSCFPipeliningPass)134 135 TestSCFPipeliningPass() = default;136 TestSCFPipeliningPass(const TestSCFPipeliningPass &) {}137 StringRef getArgument() const final { return "test-scf-pipelining"; }138 StringRef getDescription() const final { return "test scf.forOp pipelining"; }139 140 Option<bool> annotatePipeline{141 *this, "annotate",142 llvm::cl::desc("Annote operations during loop pipelining transformation"),143 llvm::cl::init(false)};144 145 Option<bool> noEpiloguePeeling{146 *this, "no-epilogue-peeling",147 llvm::cl::desc("Use predicates instead of peeling the epilogue."),148 llvm::cl::init(false)};149 150 static void151 getSchedule(scf::ForOp forOp,152 std::vector<std::pair<Operation *, unsigned>> &schedule) {153 if (!forOp->hasAttr(kTestPipeliningLoopMarker))154 return;155 156 schedule.resize(forOp.getBody()->getOperations().size() - 1);157 forOp.walk([&schedule](Operation *op) {158 auto attrStage =159 op->getAttrOfType<IntegerAttr>(kTestPipeliningStageMarker);160 auto attrCycle =161 op->getAttrOfType<IntegerAttr>(kTestPipeliningOpOrderMarker);162 if (attrCycle && attrStage) {163 // TODO: Index can be out-of-bounds if ops of the loop body disappear164 // due to folding.165 schedule[attrCycle.getInt()] =166 std::make_pair(op, unsigned(attrStage.getInt()));167 }168 });169 }170 171 /// Helper to generate "predicated" version of `op`. For simplicity we just172 /// wrap the operation in a scf.ifOp operation.173 static Operation *predicateOp(RewriterBase &rewriter, Operation *op,174 Value pred) {175 Location loc = op->getLoc();176 auto ifOp =177 scf::IfOp::create(rewriter, loc, op->getResultTypes(), pred, true);178 // True branch.179 rewriter.moveOpBefore(op, &ifOp.getThenRegion().front(),180 ifOp.getThenRegion().front().begin());181 rewriter.setInsertionPointAfter(op);182 if (op->getNumResults() > 0)183 scf::YieldOp::create(rewriter, loc, op->getResults());184 // False branch.185 rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());186 SmallVector<Value> elseYieldOperands;187 elseYieldOperands.reserve(ifOp.getNumResults());188 if (auto viewOp = dyn_cast<memref::SubViewOp>(op)) {189 // For sub-views, just clone the op.190 // NOTE: This is okay in the test because we use dynamic memref sizes, so191 // the verifier will not complain. Otherwise, we may create a logically192 // out-of-bounds view and a different technique should be used.193 Operation *opClone = rewriter.clone(*op);194 elseYieldOperands.append(opClone->result_begin(), opClone->result_end());195 } else {196 // Default to assuming constant numeric values.197 for (Type type : op->getResultTypes()) {198 elseYieldOperands.push_back(arith::ConstantOp::create(199 rewriter, loc, rewriter.getZeroAttr(type)));200 }201 }202 if (op->getNumResults() > 0)203 scf::YieldOp::create(rewriter, loc, elseYieldOperands);204 return ifOp.getOperation();205 }206 207 static void annotate(Operation *op,208 mlir::scf::PipeliningOption::PipelinerPart part,209 unsigned iteration) {210 OpBuilder b(op);211 switch (part) {212 case mlir::scf::PipeliningOption::PipelinerPart::Prologue:213 op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("prologue"));214 break;215 case mlir::scf::PipeliningOption::PipelinerPart::Kernel:216 op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("kernel"));217 break;218 case mlir::scf::PipeliningOption::PipelinerPart::Epilogue:219 op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("epilogue"));220 break;221 }222 op->setAttr(kTestPipeliningAnnotationIteration,223 b.getI32IntegerAttr(iteration));224 }225 226 void getDependentDialects(DialectRegistry ®istry) const override {227 registry.insert<arith::ArithDialect, memref::MemRefDialect>();228 }229 230 void runOnOperation() override {231 RewritePatternSet patterns(&getContext());232 mlir::scf::PipeliningOption options;233 options.getScheduleFn = getSchedule;234 options.supportDynamicLoops = true;235 options.predicateFn = predicateOp;236 if (annotatePipeline)237 options.annotateFn = annotate;238 if (noEpiloguePeeling) {239 options.peelEpilogue = false;240 }241 scf::populateSCFLoopPipeliningPatterns(patterns, options);242 (void)applyPatternsGreedily(getOperation(), std::move(patterns));243 getOperation().walk([](Operation *op) {244 // Clean up the markers.245 op->removeAttr(kTestPipeliningStageMarker);246 op->removeAttr(kTestPipeliningOpOrderMarker);247 });248 }249};250} // namespace251 252namespace mlir {253namespace test {254void registerTestSCFUtilsPass() {255 PassRegistration<TestSCFForUtilsPass>();256 PassRegistration<TestSCFIfUtilsPass>();257 PassRegistration<TestSCFPipeliningPass>();258}259} // namespace test260} // namespace mlir261