Add RegionBranchOpInterface on affine.for op

[bondhugula]

Summary

Add RegionBranchOpInterface on affine.for op so that transforms relying
on RegionBranchOpInterface can support affine.for. E.g.:
buffer-deallocation pass.

[bondhugula]

Base it on the right commit.

[ftynse]

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

345–350

You can do DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getSuccessorEntryOperands"]> in the declaration to request declarations to be emitted for additional interface methods.

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

1746–1748

I am a bit rusty on what the interface expects here, so just pointing out that the loop may have zero iterations and branch into the body.

[herhut]

Thanks for adding this.

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

1746–1748

That is a great observation and this is wrong for the scf::ForOp, as well. In both cases (so independent of index value) the loop might branch into the body or return to the parent. We could be clever here and if the loop is known to have at least one trip, omit the immediate return to the body. IIRC there is some helper to reason about trip counts.

[bondhugula]

Thanks for the review, @ftynse, @herhut. Fixed and enhanced the check. Tested via SCCP as well.

[bondhugula]

@ftynse @herhut -- any other comments here? Everything's addressed.

[herhut]

Thanks!

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

def AffineApplyOp : Affine_Op<"apply", [NoSideEffect]> {
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def AffineForOp : Affine_Op<"for",
     [AutomaticAllocationScope, ImplicitAffineTerminator, RecursiveSideEffects,
      DeclareOpInterfaceMethods<LoopLikeOpInterface,
      ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
-      "getSingleUpperBound"]>]> {
+      "getSingleUpperBound"]>,
+     DeclareOpInterfaceMethods<RegionBranchOpInterface,
+     ["getSuccessorEntryOperands"]>]> {
   let summary = "for operation";
   let description = [{
     Syntax:
 
     ```
     operation   ::= `affine.for` ssa-id `=` lower-bound `to` upper-bound
                     (`step` integer-literal)? `{` op* `}`
 
let extraClassDeclaration = [{
     /// Sets the lower bound to the given constant value.
     void setConstantLowerBound(int64_t value);
     /// Sets the upper bound to the given constant value.
     void setConstantUpperBound(int64_t value);
 
     /// Returns true if both the lower and upper bound have the same operand
     /// lists (same operands in the same order).
     bool matchingBoundOperandList();
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
   let hasRegionVerifier = 1;
 }
 
 def AffineIfOp : Affine_Op<"if",
                            [ImplicitAffineTerminator, RecursiveSideEffects,
                             NoRegionArguments]> {
   let summary = "if-then-else operation";
   let description = [{
     Syntax:

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

 };
 } // namespace
 
 void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   results.add<AffineForEmptyLoopFolder>(context);
 }
 
+/// Return operands used when entering the region at 'index'. These operands
+/// correspond to the loop iterator operands, i.e., those excluding the
+/// induction variable. AffineForOp only has one region, so zero is the only
+/// valid value for `index`.
+OperandRange AffineForOp::getSuccessorEntryOperands(unsigned index) {
+  assert(index == 0 && "invalid region index");
+
+  // The initial operands map to the loop arguments after the induction
+  // variable.
+  return getIterOperands();
+}
+
+/// Given the region at `index`, or the parent operation if `index` is None,
+/// return the successor regions. These are the regions that may be selected
+/// during the flow of control. `operands` is a set of optional attributes that
+/// correspond to a constant value for each operand, or null if that operand is
+/// not a constant.
+void AffineForOp::getSuccessorRegions(
+    Optional<unsigned> index, ArrayRef<Attribute> operands,
+    SmallVectorImpl<RegionSuccessor> &regions) {
+  assert((!index.hasValue() || index.getValue() == 0) &&
+         "expected loop region");
+  // The loop may typically branch back to its body or to the parent operation.
+  // If the predecessor is the parent op and the trip count is known to be at
+  // least one, branch into the body using the iterator arguments. And in cases
+  // we know the trip count is zero, it can only branch back to its parent.
+  Optional<uint64_t> tripCount = getTrivialConstantTripCount(*this);
+  if (!index.hasValue() && tripCount.hasValue()) {
+    if (tripCount.getValue() > 0) {
+      regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+      return;
+    }
+    if (tripCount.getValue() == 0) {
+      regions.push_back(RegionSuccessor(getResults()));
+      return;
+    }
+  }
+
+  // From the loop body, if the trip count is one, we can only branch back to
+  // the parent.
+  if (index.hasValue() && tripCount.hasValue() && tripCount.getValue() == 1) {
+    regions.push_back(RegionSuccessor(getResults()));
+    return;
+  }
+
+  // In all other cases, the loop may branch back to itself or the parent
+  // operation.
+  regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+  regions.push_back(RegionSuccessor(getResults()));
+}
+
 /// Returns true if the affine.for has zero iterations in trivial cases.
 static bool hasTrivialZeroTripCount(AffineForOp op) {
   Optional<uint64_t> tripCount = getTrivialConstantTripCount(op);
   return tripCount.hasValue() && tripCount.getValue() == 0;
 }
 
 LogicalResult AffineForOp::fold(ArrayRef<Attribute> operands,
                                 SmallVectorImpl<OpFoldResult> &results) {

mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir

 // CHECK-NEXT: memref.dealloc %[[ALLOC3]]
 // CHECK-NEXT: scf.yield %[[ALLOC4]]
 
 //      CHECK: test.copy(%[[VAL_7]], %arg4)
 // CHECK-NEXT: memref.dealloc %[[VAL_7]]
 
 // -----
 
+// CHECK-LABEL: func @affine_loop
+func @affine_loop() {
+  %buffer = memref.alloc() : memref<1024xf32>
+  %sum_init_0 = arith.constant 0.0 : f32
+  %res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_init_0) -> f32 {
+    %t = affine.load %buffer[%i] : memref<1024xf32>
+    %sum_next = arith.addf %sum_iter, %t : f32
+    affine.yield %sum_next : f32
+  }
+  // CHECK: %[[M:.*]] = memref.alloc
+  // CHECK: affine.for
+  // CHECK: }
+  // CHECK-NEXT: memref.dealloc %[[M]]
+  return
+}
+
+// -----
+
 // Test Case: explicit control-flow loop with a dynamically allocated buffer.
 // The BufferDeallocation transformation should fail on this explicit
 // control-flow loop since they are not supported.
 
 // expected-error@+1 {{Only structured control-flow loops are supported}}
 func @loop_dynalloc(
   %arg0 : i32,
   %arg1 : i32,

mlir/test/Transforms/sccp-structured.mlir

%0 = scf.while (%arg2 = %c2_i32) : (i32) -> (i32) {
     %1 = arith.cmpi slt, %arg2, %arg1 : i32
     scf.condition(%1) %arg2 : i32
   } do {
   ^bb0(%arg2: i32):
     scf.yield %arg2 : i32
   }
   return
 }
+
+/// Check that propgation happens for affine.for -- tests its region branch op
+/// interface as well.
+
+// CHECK-LABEL: func @affine_loop_one_iter(
+func @affine_loop_one_iter(%arg0 : index, %arg1 : index, %arg2 : index) -> i32 {
+  // CHECK: %[[C1:.*]] = arith.constant 1 : i32
+  %s0 = arith.constant 0 : i32
+  %s1 = arith.constant 1 : i32
+  %result = affine.for %i = 0 to 1 iter_args(%si = %s0) -> (i32) {
+    %sn = arith.addi %si, %s1 : i32
+    affine.yield %sn : i32
+  }
+  // CHECK: return %[[C1]] : i32
+  return %result : i32
+}
+
+// CHECK-LABEL: func @affine_loop_zero_iter(
+func @affine_loop_zero_iter(%arg0 : index, %arg1 : index, %arg2 : index) -> i32 {
+  // This exposes a crash in sccp/forward data flow analysis: https://github.com/llvm/llvm-project/issues/54928
+  // CHECK: %[[C0:.*]] = arith.constant 0 : i32
+  %s0 = arith.constant 0 : i32
+  // %result = affine.for %i = 0 to 0 iter_args(%si = %s0) -> (i32) {
+  //  %sn = arith.addi %si, %si : i32
+  //  affine.yield %sn : i32
+  // }
+  // return %result : i32
+  // CHECK: return %[[C0]] : i32
+  return %s0 : i32
+}