264 lines · cpp
1//===- ReductionTreePass.cpp - ReductionTreePass Implementation -----------===//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 defines the Reduction Tree Pass class. It provides a framework for10// the implementation of different reduction passes in the MLIR Reduce tool. It11// allows for custom specification of the variant generation behavior. It12// implements methods that define the different possible traversals of the13// reduction tree.14//15//===----------------------------------------------------------------------===//16 17#include "mlir/IR/DialectInterface.h"18#include "mlir/Reducer/Passes.h"19#include "mlir/Reducer/ReductionNode.h"20#include "mlir/Reducer/ReductionPatternInterface.h"21#include "mlir/Reducer/Tester.h"22#include "mlir/Rewrite/FrozenRewritePatternSet.h"23#include "mlir/Transforms/GreedyPatternRewriteDriver.h"24 25#include "llvm/ADT/ArrayRef.h"26#include "llvm/Support/Allocator.h"27 28namespace mlir {29#define GEN_PASS_DEF_REDUCTIONTREEPASS30#include "mlir/Reducer/Passes.h.inc"31} // namespace mlir32 33using namespace mlir;34 35/// We implicitly number each operation in the region and if an operation's36/// number falls into rangeToKeep, we need to keep it and apply the given37/// rewrite patterns on it.38static void applyPatterns(Region ®ion,39 const FrozenRewritePatternSet &patterns,40 ArrayRef<ReductionNode::Range> rangeToKeep,41 bool eraseOpNotInRange) {42 std::vector<Operation *> opsNotInRange;43 std::vector<Operation *> opsInRange;44 size_t keepIndex = 0;45 for (const auto &op : enumerate(region.getOps())) {46 int index = op.index();47 if (keepIndex < rangeToKeep.size() &&48 index == rangeToKeep[keepIndex].second)49 ++keepIndex;50 if (keepIndex == rangeToKeep.size() || index < rangeToKeep[keepIndex].first)51 opsNotInRange.push_back(&op.value());52 else53 opsInRange.push_back(&op.value());54 }55 56 // `applyOpPatternsGreedily` with folding may erase the ops so we can't do the57 // pattern matching in above iteration. Besides, erase op not-in-range may end58 // up in invalid module, so `applyOpPatternsGreedily` with folding should come59 // before that transform.60 for (Operation *op : opsInRange) {61 // `applyOpPatternsGreedily` with folding returns whether the op is62 // converted. Omit it because we don't have expectation this reduction will63 // be success or not.64 (void)applyOpPatternsGreedily(op, patterns,65 GreedyRewriteConfig().setStrictness(66 GreedyRewriteStrictness::ExistingOps));67 }68 69 if (eraseOpNotInRange)70 for (Operation *op : opsNotInRange) {71 op->dropAllUses();72 op->erase();73 }74}75 76/// We will apply the reducer patterns to the operations in the ranges specified77/// by ReductionNode. Note that we are not able to remove an operation without78/// replacing it with another valid operation. However, The validity of module79/// reduction is based on the Tester provided by the user and that means certain80/// invalid module is still interested by the use. Thus we provide an81/// alternative way to remove operations, which is using `eraseOpNotInRange` to82/// erase the operations not in the range specified by ReductionNode.83template <typename IteratorType>84static LogicalResult findOptimal(ModuleOp module, Region ®ion,85 const FrozenRewritePatternSet &patterns,86 const Tester &test, bool eraseOpNotInRange) {87 std::pair<Tester::Interestingness, size_t> initStatus =88 test.isInteresting(module);89 // While exploring the reduction tree, we always branch from an interesting90 // node. Thus the root node must be interesting.91 if (initStatus.first != Tester::Interestingness::True)92 return module.emitWarning() << "uninterested module will not be reduced";93 94 llvm::SpecificBumpPtrAllocator<ReductionNode> allocator;95 96 std::vector<ReductionNode::Range> ranges{97 {0, std::distance(region.op_begin(), region.op_end())}};98 99 ReductionNode *root = allocator.Allocate();100 new (root) ReductionNode(nullptr, ranges, allocator);101 // Duplicate the module for root node and locate the region in the copy.102 if (failed(root->initialize(module, region)))103 llvm_unreachable("unexpected initialization failure");104 root->update(initStatus);105 106 ReductionNode *smallestNode = root;107 IteratorType iter(root);108 109 while (iter != IteratorType::end()) {110 ReductionNode ¤tNode = *iter;111 Region &curRegion = currentNode.getRegion();112 113 applyPatterns(curRegion, patterns, currentNode.getRanges(),114 eraseOpNotInRange);115 currentNode.update(test.isInteresting(currentNode.getModule()));116 117 if (currentNode.isInteresting() == Tester::Interestingness::True &&118 currentNode.getSize() < smallestNode->getSize())119 smallestNode = ¤tNode;120 121 ++iter;122 }123 124 // At here, we have found an optimal path to reduce the given region. Retrieve125 // the path and apply the reducer to it.126 SmallVector<ReductionNode *> trace;127 ReductionNode *curNode = smallestNode;128 trace.push_back(curNode);129 while (curNode != root) {130 curNode = curNode->getParent();131 trace.push_back(curNode);132 }133 134 // Reduce the region through the optimal path.135 while (!trace.empty()) {136 ReductionNode *top = trace.pop_back_val();137 applyPatterns(region, patterns, top->getStartRanges(), eraseOpNotInRange);138 }139 140 if (test.isInteresting(module).first != Tester::Interestingness::True)141 llvm::report_fatal_error("Reduced module is not interesting");142 if (test.isInteresting(module).second != smallestNode->getSize())143 llvm::report_fatal_error(144 "Reduced module doesn't have consistent size with smallestNode");145 return success();146}147 148template <typename IteratorType>149static LogicalResult findOptimal(ModuleOp module, Region ®ion,150 const FrozenRewritePatternSet &patterns,151 const Tester &test) {152 // We separate the reduction process into 2 steps, the first one is to erase153 // redundant operations and the second one is to apply the reducer patterns.154 155 // In the first phase, we don't apply any patterns so that we only select the156 // range of operations to keep to the module stay interesting.157 if (failed(findOptimal<IteratorType>(module, region, /*patterns=*/{}, test,158 /*eraseOpNotInRange=*/true)))159 return failure();160 // In the second phase, we suppose that no operation is redundant, so we try161 // to rewrite the operation into simpler form.162 return findOptimal<IteratorType>(module, region, patterns, test,163 /*eraseOpNotInRange=*/false);164}165 166namespace {167 168//===----------------------------------------------------------------------===//169// Reduction Pattern Interface Collection170//===----------------------------------------------------------------------===//171 172class ReductionPatternInterfaceCollection173 : public DialectInterfaceCollection<DialectReductionPatternInterface> {174public:175 using Base::Base;176 177 // Collect the reduce patterns defined by each dialect.178 void populateReductionPatterns(RewritePatternSet &pattern,179 Tester &tester) const {180 for (const DialectReductionPatternInterface &interface : *this) {181 interface.populateReductionPatterns(pattern);182 interface.populateReductionPatternsWithTester(pattern, tester);183 }184 }185};186 187//===----------------------------------------------------------------------===//188// ReductionTreePass189//===----------------------------------------------------------------------===//190 191/// This class defines the Reduction Tree Pass. It provides a framework to192/// to implement a reduction pass using a tree structure to keep track of the193/// generated reduced variants.194class ReductionTreePass195 : public impl::ReductionTreePassBase<ReductionTreePass> {196public:197 using Base::Base;198 199 LogicalResult initialize(MLIRContext *context) override;200 201 /// Runs the pass instance in the pass pipeline.202 void runOnOperation() override;203 204private:205 LogicalResult reduceOp(ModuleOp module, Region ®ion);206 207 Tester tester;208 FrozenRewritePatternSet reducerPatterns;209};210 211} // namespace212 213LogicalResult ReductionTreePass::initialize(MLIRContext *context) {214 tester.setTestScript(testerName);215 tester.setTestScriptArgs(testerArgs);216 217 RewritePatternSet patterns(context);218 219 ReductionPatternInterfaceCollection reducePatternCollection(context);220 reducePatternCollection.populateReductionPatterns(patterns, tester);221 222 reducerPatterns = std::move(patterns);223 return success();224}225 226void ReductionTreePass::runOnOperation() {227 Operation *topOperation = getOperation();228 while (topOperation->getParentOp() != nullptr)229 topOperation = topOperation->getParentOp();230 ModuleOp module = dyn_cast<ModuleOp>(topOperation);231 if (!module) {232 emitError(getOperation()->getLoc())233 << "top-level op must be 'builtin.module'";234 return signalPassFailure();235 }236 237 SmallVector<Operation *, 8> workList;238 workList.push_back(getOperation());239 240 do {241 Operation *op = workList.pop_back_val();242 243 for (Region ®ion : op->getRegions())244 if (!region.empty())245 if (failed(reduceOp(module, region)))246 return signalPassFailure();247 248 for (Region ®ion : op->getRegions())249 for (Operation &op : region.getOps())250 if (op.getNumRegions() != 0)251 workList.push_back(&op);252 } while (!workList.empty());253}254 255LogicalResult ReductionTreePass::reduceOp(ModuleOp module, Region ®ion) {256 switch (traversalModeId) {257 case TraversalMode::SinglePath:258 return findOptimal<ReductionNode::iterator<TraversalMode::SinglePath>>(259 module, region, reducerPatterns, tester);260 default:261 return module.emitError() << "unsupported traversal mode detected";262 }263}264