[Mlir-commits] [mlir] 104fad9 - [MLIR] Add folder for zero trip count affine.for

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2021-07-21T20:28:35+05:30
New Revision: 104fad99c9e378c09061f2d4c05e9a506b4ad58b

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

LOG: [MLIR] Add folder for zero trip count affine.for

AffineForOp's folding hook is expected to fold away trivially empty
affine.for.  This allows simplification to happen as part of the
canonicalizer and from wherever the folding hook is used. While more
complex analysis based zero trip count detection is available from other
passes in analysis and transforms, simple and inexpensive folding had
been missing.

Also, update/improve affine.for op documentation clarifying semantics of
the result values for zero trip count loops.

Differential Revision: https://reviews.llvm.org/D106123

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Affine/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index dda8d42674cb6..b5f96b3b49ff9 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -178,18 +178,19 @@ def AffineForOp : Affine_Op<"for",
       return
     }
     ```
-    `affine.for` can also operate on loop-carried variables and return the final
-    values after loop termination. The initial values of the variables are
-    passed as additional SSA operands to the "affine.for" following the 2 loop
-    control values lower bound, upper bound. The operation region has equivalent
-    arguments for each variable representing the value of the variable at the
-    current iteration.
+    `affine.for` can also operate on loop-carried variables (`iter_args`) and
+    return the final values after loop termination. The initial values of the
+    variables are passed as additional SSA operands to the `affine.for`
+    following the operands for the loop's lower and upper bounds. The
+    operation's region has equivalent arguments for each variable representing
+    the value of the variable at the current iteration.
 
     The region must terminate with an `affine.yield` that passes all the current
-    iteration variables to the next iteration, or to the `affine.for` result, if
-    at the last iteration.
+    iteration variables to the next iteration, or to the `affine.for`'s results
+    if at the last iteration. For `affine.for`'s that execute zero iterations, the
+    initial values of the loop-carried variables (corresponding to the SSA
+    operands) will be the op's results.
 
-    `affine.for` results hold the final values after the last iteration.
     For example, to sum-reduce a memref:
 
      ```mlir
@@ -208,6 +209,15 @@ def AffineForOp : Affine_Op<"for",
       return %sum : f32
     }
     ```
+
+    ```mlir
+    %res:2 = affine.for %i = 0 to 128 iter_args(%arg0 = %init0, %arg1 = %init1)
+               -> (index, index) {
+      %y0 = addi %arg0, %c1 : index
+      %y1 = addi %arg1, %c2 : index
+      affine.yield %y0, %y1 : index, index
+    }
+    ```
     If the `affine.for` defines any values, a yield terminator must be
     explicitly present. The number and types of the "affine.for" results must
     match the initial values in the `iter_args` binding and the yield operands.

diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 52200f44c02a9..b1721cd2f4640 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1648,10 +1648,25 @@ void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<AffineForEmptyLoopFolder>(context);
 }
 
+/// Returns true if the affine.for has zero iterations in trivial cases.
+static bool hasTrivialZeroTripCount(AffineForOp op) {
+  if (!op.hasConstantBounds())
+    return false;
+  int64_t lb = op.getConstantLowerBound();
+  int64_t ub = op.getConstantUpperBound();
+  return ub - lb <= 0;
+}
+
 LogicalResult AffineForOp::fold(ArrayRef<Attribute> operands,
                                 SmallVectorImpl<OpFoldResult> &results) {
   bool folded = succeeded(foldLoopBounds(*this));
   folded |= succeeded(canonicalizeLoopBounds(*this));
+  if (hasTrivialZeroTripCount(*this)) {
+    // The initial values of the loop-carried variables (iter_args) are the
+    // results of the op.
+    results.assign(getIterOperands().begin(), getIterOperands().end());
+    folded = true;
+  }
   return success(folded);
 }
 

diff  --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index ba15d2a570244..f12a0ca032d51 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -469,6 +469,23 @@ func @fold_empty_loop() {
 }
 // CHECK: return
 
+// -----
+
+// CHECK-LABEL:  func @fold_zero_iter_loops
+// CHECK-SAME: %[[ARG:.*]]: index
+func @fold_zero_iter_loops(%in : index) -> index {
+  %c1 = constant 1 : index
+  affine.for %i = 0 to 0 {
+    affine.for %j = 0 to -1 {
+    }
+  }
+  %res = affine.for %i = 0 to 0 iter_args(%loop_arg = %in) -> index {
+    %yield = addi %loop_arg, %c1 : index
+    affine.yield %yield : index
+  }
+  // CHECK-NEXT: return %[[ARG]]
+  return %res : index
+}
 
 // -----