[MLIR][SCF] Fix LoopPeelOp documentation (NFC) (#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 <andrzej.warzynski@gmail.com>
This commit is contained in:
Hugo Trachino
2024-10-29 15:47:13 +00:00
committed by GitHub
parent bf6c483e47
commit a9c417c28a
2 changed files with 18 additions and 17 deletions

View File

@@ -146,7 +146,7 @@ def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
let summary = "Peels the first or last iteration of the loop";
let description = [{
Rewrite the given loop with a main loop and a partial (first or last) loop.
When the `peelFront` option is set as true, the first iteration is peeled off.
When the `peelFront` option is set to true, the first iteration is peeled off.
Otherwise, updates the given loop so that its step evenly divides its range and puts
the remaining iteration into a separate loop or a conditional.
@@ -156,18 +156,20 @@ def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
#### Return modes
This operation ignores non-scf::ForOp ops and drops them in the return.
The op returns two loops, the peeled loop which has trip count divisible
by the step, and the remainder loop.
When `peelFront` is true, this operation returns two scf::ForOp Ops, the
first scf::ForOp corresponds to the first iteration of the loop which can
be canonicalized away in the following optimization. The second loop Op
contains the remaining iteration, and the new lower bound is the original
lower bound plus the number of steps.
When `peelFront` is true, the first result (remainder loop) executes all
but the first iteration of the target loop. The second result (peeled
loop) corresponds to the first iteration of the loop which can be
canonicalized away in the following optimizations.
When `peelFront` is not true, this operation returns two scf::ForOp Ops, with the first
scf::ForOp satisfying: "the loop trip count is divisible by the step".
The second loop Op contains the remaining iteration. Note that even though the
Payload IR modification may be performed in-place, this operation consumes
the operand handle and produces a new one.
When `peelFront` is false, the first result (peeled loop) is the portion
of the target loop with the highest upper bound that is divisible by the
step. The second result (remainder loop) contains the remaining iterations.
Note that even though the Payload IR modification may be performed
in-place, this operation consumes the operand handle and produces a new one.
#### Return Modes

View File

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