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

[vinayaka-bandishti]

Summary

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.

[bondhugula]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2484

getNumResults() != 0 to make it more obvious.

mlir/test/Dialect/Affine/canonicalize.mlir

588

We should still be able to erase such ops via a canonicalization pattern. This could be added in another revision though.

[vinayaka-bandishti]

Address comments, use getNumResults.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 static bool hasTrivialZeroTripCount(AffineForOp op) {
   std::optional<uint64_t> tripCount = getTrivialConstantTripCount(op);
   return tripCount && *tripCount == 0;
 }
 
 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;
   }
   return success(folded);
 }
 
 AffineBound AffineForOp::getLowerBound() {
   auto lbMap = getLowerBoundMap();

mlir/test/Dialect/Affine/canonicalize.mlir

func.func @empty_loops_not_folded_3() -> (index, index) {
   %res:2 = affine.for %i = 0 to 10 iter_args(%arg0 = %c0, %arg1 = %c1) -> (index, index) {
     affine.yield %arg1, %arg0 : index, index
   }
   return %res#0, %res#1 : 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
   }
   // CHECK-NEXT: return %[[ARG]]
   return %res : index
 }