[Mlir-commits] [mlir] a9c417c - [MLIR][SCF] Fix LoopPeelOp documentation (NFC) (#113179)

[llvmlistbot at llvm.org]

Author: Hugo Trachino
Date: 2024-10-29T15:47:13Z
New Revision: a9c417c28a25c153aa0fdbe2eb5453a93820a3b1

URL: https://github.com/llvm/llvm-project/commit/a9c417c28a25c153aa0fdbe2eb5453a93820a3b1
DIFF: https://github.com/llvm/llvm-project/commit/a9c417c28a25c153aa0fdbe2eb5453a93820a3b1.diff

LOG: [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 at gmail.com>

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
    mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
index 20880d94a83cac..5dba8c5e57ba86 100644
--- a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
+++ b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
@@ -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.
-
-     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 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.
+     The op returns two loops, the peeled loop which has trip count divisible
+     by the step, and the remainder loop.
+
+     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 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
 

diff  --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index a30e349d49136c..5104ad4b3a3038 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -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);