Skip to content

Commit a9c417c

Browse files
nujaabanach-space
andauthored
[MLIR][SCF] Fix LoopPeelOp documentation (NFC) (llvm#113179)
As an example, I added annotations to the peel_front unit test. ``` func.func @loop_peel_first_iter_op() { // CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C41:.+]] = arith.constant 41 // CHECK: %[[C5:.+]] = arith.constant 5 // CHECK: %[[C5_0:.+]] = arith.constant 5 // CHECK: scf.for %{{.+}} = %[[C0]] to %[[C5_0]] step %[[C5]] // CHECK: arith.addi // CHECK: scf.for %{{.+}} = %[[C5_0]] to %[[C41]] step %[[C5]] // CHECK: arith.addi %0 = arith.constant 0 : index %1 = arith.constant 41 : index %2 = arith.constant 5 : index scf.for %i = %0 to %1 step %2 { arith.addi %i, %i : index } return } module attributes {transform.with_named_sequence} { transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for"> %main_loop, %remainder = transform.loop.peel %1 {peel_front = true} : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">) transform.annotate %main_loop "main_loop" : !transform.op<"scf.for"> transform.annotate %remainder "remainder" : !transform.op<"scf.for"> transform.yield } } ``` Gives : ``` func.func @loop_peel_first_iter_op() { %c0 = arith.constant 0 : index %c41 = arith.constant 41 : index %c5 = arith.constant 5 : index %c5_0 = arith.constant 5 : index scf.for %arg0 = %c0 to %c5_0 step %c5 { %0 = arith.addi %arg0, %arg0 : index } {remainder} // The first iteration loop (second result) has been annotated remainder scf.for %arg0 = %c5_0 to %c41 step %c5 { %0 = arith.addi %arg0, %arg0 : index } {main_loop} // The main loop (first result) has been annotated main_loop return } ``` --------- Co-authored-by: Andrzej Warzyński <[email protected]>
1 parent bf6c483 commit a9c417c

File tree

2 files changed

+20
-19
lines changed

2 files changed

+20
-19
lines changed

mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
146146
let summary = "Peels the first or last iteration of the loop";
147147
let description = [{
148148
Rewrite the given loop with a main loop and a partial (first or last) loop.
149-
When the `peelFront` option is set as true, the first iteration is peeled off.
149+
When the `peelFront` option is set to true, the first iteration is peeled off.
150150
Otherwise, updates the given loop so that its step evenly divides its range and puts
151151
the remaining iteration into a separate loop or a conditional.
152152

@@ -156,18 +156,20 @@ def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
156156
#### Return modes
157157

158158
This operation ignores non-scf::ForOp ops and drops them in the return.
159-
160-
When `peelFront` is true, this operation returns two scf::ForOp Ops, the
161-
first scf::ForOp corresponds to the first iteration of the loop which can
162-
be canonicalized away in the following optimization. The second loop Op
163-
contains the remaining iteration, and the new lower bound is the original
164-
lower bound plus the number of steps.
165-
166-
When `peelFront` is not true, this operation returns two scf::ForOp Ops, with the first
167-
scf::ForOp satisfying: "the loop trip count is divisible by the step".
168-
The second loop Op contains the remaining iteration. Note that even though the
169-
Payload IR modification may be performed in-place, this operation consumes
170-
the operand handle and produces a new one.
159+
The op returns two loops, the peeled loop which has trip count divisible
160+
by the step, and the remainder loop.
161+
162+
When `peelFront` is true, the first result (remainder loop) executes all
163+
but the first iteration of the target loop. The second result (peeled
164+
loop) corresponds to the first iteration of the loop which can be
165+
canonicalized away in the following optimizations.
166+
167+
When `peelFront` is false, the first result (peeled loop) is the portion
168+
of the target loop with the highest upper bound that is divisible by the
169+
step. The second result (remainder loop) contains the remaining iterations.
170+
171+
Note that even though the Payload IR modification may be performed
172+
in-place, this operation consumes the operand handle and produces a new one.
171173

172174
#### Return Modes
173175

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -206,12 +206,11 @@ LogicalResult mlir::scf::peelForLoopAndSimplifyBounds(RewriterBase &rewriter,
206206
return success();
207207
}
208208

209-
/// When the `peelFront` option is set as true, the first iteration of the loop
210-
/// is peeled off. This function rewrites the original scf::ForOp as two
211-
/// scf::ForOp Ops, the first scf::ForOp corresponds to the first iteration of
212-
/// the loop which can be canonicalized away in the following optimization. The
213-
/// second loop Op contains the remaining iteration, and the new lower bound is
214-
/// the original lower bound plus the number of steps.
209+
/// Rewrites the original scf::ForOp as two scf::ForOp Ops, the first
210+
/// scf::ForOp corresponds to the first iteration of the loop which can be
211+
/// canonicalized away in the following optimizations. The second loop Op
212+
/// contains the remaining iterations, with a lower bound updated as the
213+
/// original lower bound plus the step (i.e. skips the first iteration).
215214
LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
216215
ForOp &firstIteration) {
217216
RewriterBase::InsertionGuard guard(b);

0 commit comments

Comments
 (0)