86 lines · cpp
1//=== TestParallelLoopUnrolling.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::ParallelOp>(currOp))28 depth++;29 }30 return depth;31}32 33struct TestParallelLoopUnrollingPass34 : public PassWrapper<TestParallelLoopUnrollingPass, OperationPass<>> {35 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestParallelLoopUnrollingPass)36 37 StringRef getArgument() const final { return "test-parallel-loop-unrolling"; }38 StringRef getDescription() const final {39 return "Tests parallel loop unrolling transformation";40 }41 TestParallelLoopUnrollingPass() = default;42 TestParallelLoopUnrollingPass(const TestParallelLoopUnrollingPass &) {}43 44 void getDependentDialects(DialectRegistry ®istry) const override {45 registry.insert<arith::ArithDialect>();46 }47 48 void runOnOperation() override {49 SmallVector<scf::ParallelOp, 4> loops;50 getOperation()->walk([&](scf::ParallelOp parLoop) {51 if (getNestingDepth(parLoop) == loopDepth)52 loops.push_back(parLoop);53 });54 auto annotateFn = [this](unsigned i, Operation *op, OpBuilder b) {55 if (annotateLoop) {56 op->setAttr("unrolled_iteration", b.getUI32IntegerAttr(i));57 }58 };59 PatternRewriter rewriter(getOperation()->getContext());60 for (auto loop : loops) {61 (void)parallelLoopUnrollByFactors(loop, unrollFactors, rewriter,62 annotateFn);63 }64 }65 66 ListOption<uint64_t> unrollFactors{67 *this, "unroll-factors",68 llvm::cl::desc(69 "Unroll factors for each parallel loop dim. If fewer factors than "70 "loop dims are provided, they are applied to the inner dims.")};71 Option<unsigned> loopDepth{*this, "loop-depth", llvm::cl::desc("Loop depth."),72 llvm::cl::init(0)};73 Option<bool> annotateLoop{*this, "annotate",74 llvm::cl::desc("Annotate unrolled iterations."),75 llvm::cl::init(false)};76};77} // namespace78 79namespace mlir {80namespace test {81void registerTestParallelLoopUnrollingPass() {82 PassRegistration<TestParallelLoopUnrollingPass>();83}84} // namespace test85} // namespace mlir86