From 28c672cf83e81757023cd5ae1e044fcfcfacfc0f Mon Sep 17 00:00:00 2001
From: Vinayaka Bandishti <vinayaka@polymagelabs.com>
Date: Mon, 29 May 2023 01:53:31 +0530
Subject: [PATCH] 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
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   |  7 +++++--
 mlir/test/Dialect/Affine/canonicalize.mlir | 17 +++++++++++++----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 9153686..c4c1aab 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 ace9798..f734f37 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
-- 
2.7.4