MLIR  16.0.0git
LoopRangeFolding.cpp
Go to the documentation of this file.
1 //===- LoopRangeFolding.cpp - Code to perform loop range folding-----------===//
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 loop range folding.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "PassDetail.h"
20 
21 using namespace mlir;
22 using namespace mlir::scf;
23 
24 namespace {
25 struct ForLoopRangeFolding
26  : public SCFForLoopRangeFoldingBase<ForLoopRangeFolding> {
27  void runOnOperation() override;
28 };
29 } // namespace
30 
31 void ForLoopRangeFolding::runOnOperation() {
32  getOperation()->walk([&](ForOp op) {
33  Value indVar = op.getInductionVar();
34 
35  auto canBeFolded = [&](Value value) {
36  return op.isDefinedOutsideOfLoop(value) || value == indVar;
37  };
38 
39  // Fold until a fixed point is reached
40  while (true) {
41 
42  // If the induction variable is used more than once, we can't fold its
43  // arith ops into the loop range
44  if (!indVar.hasOneUse())
45  break;
46 
47  Operation *user = *indVar.getUsers().begin();
48  if (!isa<arith::AddIOp, arith::MulIOp>(user))
49  break;
50 
51  if (!llvm::all_of(user->getOperands(), canBeFolded))
52  break;
53 
54  OpBuilder b(op);
56  lbMap.map(indVar, op.getLowerBound());
58  ubMap.map(indVar, op.getUpperBound());
59  BlockAndValueMapping stepMap;
60  stepMap.map(indVar, op.getStep());
61 
62  if (isa<arith::AddIOp>(user)) {
63  Operation *lbFold = b.clone(*user, lbMap);
64  Operation *ubFold = b.clone(*user, ubMap);
65 
66  op.setLowerBound(lbFold->getResult(0));
67  op.setUpperBound(ubFold->getResult(0));
68 
69  } else if (isa<arith::MulIOp>(user)) {
70  Operation *ubFold = b.clone(*user, ubMap);
71  Operation *stepFold = b.clone(*user, stepMap);
72 
73  op.setUpperBound(ubFold->getResult(0));
74  op.setStep(stepFold->getResult(0));
75  }
76 
77  ValueRange wrapIndvar(indVar);
78  user->replaceAllUsesWith(wrapIndvar);
79  user->erase();
80  }
81  });
82 }
83 
84 std::unique_ptr<Pass> mlir::createForLoopRangeFoldingPass() {
85  return std::make_unique<ForLoopRangeFolding>();
86 }
Include the generated interface declarations.
std::unique_ptr< Pass > createForLoopRangeFoldingPass()
Creates a pass which folds arith ops on induction variable into loop range.
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:28
operand_range getOperands()
Returns an iterator on the underlying Value&#39;s.
Definition: Operation.h:295
user_range getUsers() const
Definition: Value.h:213
static constexpr const bool value
void erase()
Remove this operation from its parent block and delete it.
Definition: Operation.cpp:414
void map(Block *from, Block *to)
Inserts a new mapping for &#39;from&#39; to &#39;to&#39;.
Operation * clone(BlockAndValueMapping &mapper, CloneOptions options=CloneOptions::all())
Create a deep copy of this operation, remapping any operands that use values outside of the operation...
Definition: Operation.cpp:554
bool hasOneUse() const
Returns true if this value has exactly one use.
Definition: Value.h:200
void replaceAllUsesWith(ValuesT &&values)
Replace all uses of results of this operation with the provided &#39;values&#39;.
Definition: Operation.h:203
OpResult getResult(unsigned idx)
Get the &#39;idx&#39;th result of this operation.
Definition: Operation.h:324
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:85
This class helps build Operations.
Definition: Builders.h:192
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:345