Skip to content

Commit 629305b

Browse files
ergawyergawy
committed
Simplify iv update op detection
1 parent 933d1e7 commit 629305b

File tree

1 file changed

+62
-95
lines changed

1 file changed

+62
-95
lines changed

‎flang/lib/Optimizer/OpenMP/DoConcurrentConversion.cpp‎

Lines changed: 62 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -28,108 +28,80 @@ namespace looputils {
2828
/// Stores info needed about the induction/iteration variable for each `do
2929
/// concurrent` in a loop nest.
3030
struct InductionVariableInfo {
31+
InductionVariableInfo(fir::DoLoopOp doLoop) { populateInfo(doLoop); }
32+
3133
/// The operation allocating memory for iteration variable.
3234
mlir::Operation *iterVarMemDef;
3335
/// the operation(s) updating the iteration variable with the current
3436
/// iteration number.
35-
llvm::SetVector<mlir::Operation *> indVarUpdateOps;
36-
};
37-
38-
using LoopNestToIndVarMap =
39-
llvm::MapVector<fir::DoLoopOp, InductionVariableInfo>;
37+
llvm::SmallVector<mlir::Operation *, 2> indVarUpdateOps;
4038

41-
/// For the \p doLoop parameter, find the operation that declares its iteration
42-
/// variable or allocates memory for it.
43-
///
44-
/// For example, give the following loop:
45-
/// ```
46-
/// ...
47-
/// %i:2 = hlfir.declare %0 {uniq_name = "_QFEi"} : ...
48-
/// ...
49-
/// fir.do_loop %ind_var = %lb to %ub step %s unordered {
50-
/// %ind_var_conv = fir.convert %ind_var : (index) -> i32
51-
/// fir.store %ind_var_conv to %i#1 : !fir.ref<i32>
52-
/// ...
53-
/// }
54-
/// ```
55-
///
56-
/// This function returns the `hlfir.declare` op for `%i`.
57-
///
58-
/// Note: The current implementation is dependent on how flang emits loop
59-
/// bodies; which is sufficient for the current simple test/use cases. If this
60-
/// proves to be insufficient, this should be made more generic.
61-
mlir::Operation *findLoopIterationVarMemDecl(fir::DoLoopOp doLoop) {
62-
mlir::Value result = nullptr;
63-
64-
// Checks if a StoreOp is updating the memref of the loop's iteration
65-
// variable.
66-
auto isStoringIV = [&](fir::StoreOp storeOp) {
67-
// Direct store into the IV memref.
68-
if (storeOp.getValue() == doLoop.getInductionVar())
69-
return true;
70-
71-
// Indirect store into the IV memref.
72-
if (auto convertOp = mlir::dyn_cast<fir::ConvertOp>(
73-
storeOp.getValue().getDefiningOp())) {
74-
if (convertOp.getOperand() == doLoop.getInductionVar())
39+
private:
40+
/// For the \p doLoop parameter, find the following:
41+
///
42+
/// 1. The operation that declares its iteration variable or allocates memory
43+
/// for it. For example, give the following loop:
44+
/// ```
45+
/// ...
46+
/// %i:2 = hlfir.declare %0 {uniq_name = "_QFEi"} : ...
47+
/// ...
48+
/// fir.do_loop %ind_var = %lb to %ub step %s unordered {
49+
/// %ind_var_conv = fir.convert %ind_var : (index) -> i32
50+
/// fir.store %ind_var_conv to %i#1 : !fir.ref<i32>
51+
/// ...
52+
/// }
53+
/// ```
54+
///
55+
/// This function sets the `iterVarMemDef` member to the `hlfir.declare` op
56+
/// for `%i`.
57+
///
58+
/// 2. The operation(s) that update the loop's iteration variable from its
59+
/// induction variable. For the above example, the `indVarUpdateOps` is
60+
/// populated with the first 2 ops in the loop's body.
61+
///
62+
/// Note: The current implementation is dependent on how flang emits loop
63+
/// bodies; which is sufficient for the current simple test/use cases. If this
64+
/// proves to be insufficient, this should be made more generic.
65+
void populateInfo(fir::DoLoopOp doLoop) {
66+
mlir::Value result = nullptr;
67+
68+
// Checks if a StoreOp is updating the memref of the loop's iteration
69+
// variable.
70+
auto isStoringIV = [&](fir::StoreOp storeOp) {
71+
// Direct store into the IV memref.
72+
if (storeOp.getValue() == doLoop.getInductionVar()) {
73+
indVarUpdateOps.push_back(storeOp);
7574
return true;
76-
}
77-
78-
return false;
79-
};
80-
81-
for (mlir::Operation &op : doLoop) {
82-
if (auto storeOp = mlir::dyn_cast<fir::StoreOp>(op))
83-
if (isStoringIV(storeOp)) {
84-
result = storeOp.getMemref();
85-
break;
8675
}
87-
}
8876

89-
assert(result != nullptr && result.getDefiningOp() != nullptr);
90-
return result.getDefiningOp();
91-
}
92-
93-
/// Collects the op(s) responsible for updating a loop's iteration variable with
94-
/// the current iteration number. For example, for the input IR:
95-
/// ```
96-
/// %i = fir.alloca i32 {bindc_name = "i"}
97-
/// %i_decl:2 = hlfir.declare %i ...
98-
/// ...
99-
/// fir.do_loop %i_iv = %lb to %ub step %step unordered {
100-
/// %1 = fir.convert %i_iv : (index) -> i32
101-
/// fir.store %1 to %i_decl#1 : !fir.ref<i32>
102-
/// ...
103-
/// }
104-
/// ```
105-
/// this function would return the first 2 ops in the `fir.do_loop`'s region.
106-
llvm::SetVector<mlir::Operation *>
107-
extractIndVarUpdateOps(fir::DoLoopOp doLoop) {
108-
mlir::Value indVar = doLoop.getInductionVar();
109-
llvm::SetVector<mlir::Operation *> indVarUpdateOps;
110-
111-
llvm::SmallVector<mlir::Value> toProcess;
112-
toProcess.push_back(indVar);
113-
114-
llvm::DenseSet<mlir::Value> done;
115-
116-
while (!toProcess.empty()) {
117-
mlir::Value val = toProcess.back();
118-
toProcess.pop_back();
119-
120-
if (!done.insert(val).second)
121-
continue;
77+
// Indirect store into the IV memref.
78+
if (auto convertOp = mlir::dyn_cast<fir::ConvertOp>(
79+
storeOp.getValue().getDefiningOp())) {
80+
if (convertOp.getOperand() == doLoop.getInductionVar()) {
81+
indVarUpdateOps.push_back(convertOp);
82+
indVarUpdateOps.push_back(storeOp);
83+
return true;
84+
}
85+
}
12286

123-
for (mlir::Operation *user : val.getUsers()) {
124-
indVarUpdateOps.insert(user);
87+
return false;
88+
};
12589

126-
for (mlir::Value result : user->getResults())
127-
toProcess.push_back(result);
90+
for (mlir::Operation &op : doLoop) {
91+
if (auto storeOp = mlir::dyn_cast<fir::StoreOp>(op))
92+
if (isStoringIV(storeOp)) {
93+
result = storeOp.getMemref();
94+
break;
95+
}
12896
}
97+
98+
assert(result != nullptr && result.getDefiningOp() != nullptr);
99+
iterVarMemDef = result.getDefiningOp();
129100
}
101+
};
130102

131-
return std::move(indVarUpdateOps);
132-
}
103+
using LoopNestToIndVarMap =
104+
llvm::MapVector<fir::DoLoopOp, InductionVariableInfo>;
133105

134106
/// Loop \p innerLoop is considered perfectly-nested inside \p outerLoop iff
135107
/// there are no operations in \p outerloop's body other than:
@@ -225,12 +197,7 @@ mlir::LogicalResult collectLoopNest(fir::DoLoopOp currentLoop,
225197
assert(currentLoop.getUnordered());
226198

227199
while (true) {
228-
loopNest.insert(
229-
{currentLoop,
230-
InductionVariableInfo{
231-
findLoopIterationVarMemDecl(currentLoop),
232-
std::move(looputils::extractIndVarUpdateOps(currentLoop))}});
233-
200+
loopNest.insert({currentLoop, InductionVariableInfo(currentLoop)});
234201
llvm::SmallVector<fir::DoLoopOp> unorderedLoops;
235202

236203
for (auto nestedLoop : currentLoop.getRegion().getOps<fir::DoLoopOp>())

0 commit comments

Comments
 (0)