[Mlir-commits] [mlir] 28c672c - Fix fold of 0-result 0-trip-count `affine.for`

[Uday Bondhugula]

Author: Vinayaka Bandishti
Date: 2023-05-29T01:53:37+05:30
New Revision: 28c672cf83e81757023cd5ae1e044fcfcfacfc0f

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

LOG: Fix fold of 0-result 0-trip-count `affine.for`

Folding 0-result ops is tricky. An 0-trip-count `affine.for` that does
not return any results is not erased at the end of its usual fold
method. Attempting fold it to only results in an infinite loop of folds
on the same op. Fix this by attempting for fold away only those
`affine.for` ops that return a non-zero number of results.

Reviewed By: bondhugula

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 9153686ab2f52..c4c1aaba638d1 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2482,9 +2482,12 @@ LogicalResult AffineForOp::fold(FoldAdaptor adaptor,
                                 SmallVectorImpl<OpFoldResult> &results) {
   bool folded = succeeded(foldLoopBounds(*this));
   folded |= succeeded(canonicalizeLoopBounds(*this));
-  if (hasTrivialZeroTripCount(*this)) {
+  if (hasTrivialZeroTripCount(*this) && getNumResults() != 0) {
     // The initial values of the loop-carried variables (iter_args) are the
-    // results of the op.
+    // results of the op. But this must be avoided for an affine.for op that
+    // does not return any results. Since ops that do not return results cannot
+    // be folded away, we would enter an infinite loop of folds on the same
+    // affine.for op.
     results.assign(getIterOperands().begin(), getIterOperands().end());
     folded = true;
   }

diff  --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index ace9798c368dc..f734f37b0dc1a 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -582,14 +582,23 @@ func.func @empty_loops_not_folded_3() -> (index, index) {
 
 // -----
 
+// CHECK-LABEL:  func @zero_iter_loop_not_folded
+func.func @zero_iter_loop_not_folded() {
+  %A = memref.alloc() : memref<4xf32>
+  affine.for %i = 0 to 0 {
+      %load = affine.load %A[%i] : memref<4xf32>
+      affine.store %load, %A[%i] : memref<4xf32>
+  }
+  // CHECK:   affine.for {{.*}} = 0 to 0 {
+  return
+}
+
+// -----
+
 // CHECK-LABEL:  func @fold_zero_iter_loops
 // CHECK-SAME: %[[ARG:.*]]: index
 func.func @fold_zero_iter_loops(%in : index) -> index {
   %c1 = arith.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 = arith.addi %loop_arg, %c1 : index
     affine.yield %yield : index