100 lines · cpp
1//===-------- TestLoopUnrolling.cpp --- loop unrolling test 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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a pass to unroll loops by a specified unroll factor.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Arith/IR/Arith.h"14#include "mlir/Dialect/SCF/IR/SCF.h"15#include "mlir/Dialect/SCF/Utils/Utils.h"16#include "mlir/IR/Builders.h"17#include "mlir/Pass/Pass.h"18 19using namespace mlir;20 21namespace {22 23static unsigned getNestingDepth(Operation *op) {24 Operation *currOp = op;25 unsigned depth = 0;26 while ((currOp = currOp->getParentOp())) {27 if (isa<scf::ForOp>(currOp))28 depth++;29 }30 return depth;31}32 33struct TestLoopUnrollingPass34 : public PassWrapper<TestLoopUnrollingPass, OperationPass<>> {35 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLoopUnrollingPass)36 37 StringRef getArgument() const final { return "test-loop-unrolling"; }38 StringRef getDescription() const final {39 return "Tests loop unrolling transformation";40 }41 TestLoopUnrollingPass() = default;42 TestLoopUnrollingPass(const TestLoopUnrollingPass &) {}43 explicit TestLoopUnrollingPass(uint64_t unrollFactorParam,44 unsigned loopDepthParam,45 bool annotateLoopParam) {46 unrollFactor = unrollFactorParam;47 loopDepth = loopDepthParam;48 annotateLoop = annotateLoopParam;49 }50 51 void getDependentDialects(DialectRegistry ®istry) const override {52 registry.insert<arith::ArithDialect>();53 }54 55 void runOnOperation() override {56 if (!(unrollFactor.getValue() > 0 || unrollFactor.getValue() == -1)) {57 emitError(UnknownLoc::get(&getContext()),58 "Invalid option: 'unroll-factor' should be greater than 0 or "59 "equal to -1");60 return signalPassFailure();61 }62 SmallVector<scf::ForOp, 4> loops;63 getOperation()->walk([&](scf::ForOp forOp) {64 if (getNestingDepth(forOp) == loopDepth)65 loops.push_back(forOp);66 });67 auto annotateFn = [this](unsigned i, Operation *op, OpBuilder b) {68 if (annotateLoop) {69 op->setAttr("unrolled_iteration", b.getUI32IntegerAttr(i));70 }71 };72 for (auto loop : loops) {73 if (unrollFactor.getValue() == -1)74 (void)loopUnrollFull(loop);75 else76 (void)loopUnrollByFactor(loop, unrollFactor, annotateFn);77 }78 }79 Option<int64_t> unrollFactor{*this, "unroll-factor",80 llvm::cl::desc("Loop unroll factor."),81 llvm::cl::init(1)};82 Option<bool> annotateLoop{*this, "annotate",83 llvm::cl::desc("Annotate unrolled iterations."),84 llvm::cl::init(false)};85 Option<bool> unrollUpToFactor{*this, "unroll-up-to-factor",86 llvm::cl::desc("Loop unroll up to factor."),87 llvm::cl::init(false)};88 Option<unsigned> loopDepth{*this, "loop-depth", llvm::cl::desc("Loop depth."),89 llvm::cl::init(0)};90};91} // namespace92 93namespace mlir {94namespace test {95void registerTestLoopUnrollingPass() {96 PassRegistration<TestLoopUnrollingPass>();97}98} // namespace test99} // namespace mlir100