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

[bondhugula]

Summary

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.

[vinayaka-polymage]

Looks good, minor documentation nits.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

191–192

Nit: loop variables is used here, it would be good to stick to loop-carried varaibles (iter-args), as the same term is used elsewhere.

Nit: bethe -> be the

[bondhugula]

Typo fixes.

[nmostafa]

Minor comments.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

183

Replace " with `

190

surround affine.for with `

[bondhugula]

Address review comments.

[bondhugula]

Address comments.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

func @simple_example(%A: memref<?x?xf32>, %B: memref<?x?xf32>) {
           %0 = affine.apply #map57(%j)[%N]
           %tmp = call @F1(%A, %i, %0) : (memref<?x?xf32>, index, index)->(f32)
           call @F2(%tmp, %B, %i, %0) : (f32, memref<?x?xf32>, index, index)->()
         }
       }
       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
     func @reduce(%buffer: memref<1024xf32>) -> (f32) {
       // Initial sum set to 0.
       %sum_0 = constant 0.0 : f32
       // iter_args binds initial values to the loop's region arguments.
       %sum = affine.for %i = 0 to 10 step 2
           iter_args(%sum_iter = %sum_0) -> (f32) {
         %t = affine.load %buffer[%i] : memref<1024xf32>
         %sum_next = addf %sum_iter, %t : f32
         // Yield current iteration sum to next iteration %sum_iter or to %sum
         // if final iteration.
         affine.yield %sum_next : f32
       }
       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.
   }];
   let arguments = (ins Variadic<AnyType>);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 

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

 };
 } // end anonymous namespace
 
 void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   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);
 }
 
 AffineBound AffineForOp::getLowerBound() {
   auto lbMap = getLowerBoundMap();
   return AffineBound(AffineForOp(*this), 0, lbMap.getNumInputs(), lbMap);
 }
 

mlir/test/Dialect/Affine/canonicalize.mlir

 func @fold_empty_loop() {
   // CHECK-NOT: affine.for
   affine.for %i = 0 to 10 {
   }
   return
 }
 // 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
+}
 
 // -----
 
 // CHECK-DAG: #[[$SET:.*]] = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1022 >= 0, d1 >= 0, -d1 + s0 - 2 >= 0)>
 
 // CHECK-LABEL: func @canonicalize_affine_if
 //  CHECK-SAME:   %[[M:[0-9a-zA-Z]*]]: index,
 //  CHECK-SAME:   %[[N:[0-9a-zA-Z]*]]: index)