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

343–348

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 hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def AffineForOp : Affine_Op<"for",
     [AutomaticAllocationScope, ImplicitAffineTerminator, RecursiveSideEffects,
      DeclareOpInterfaceMethods<LoopLikeOpInterface,
-     ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep"]>]> {
+     ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep"]>,
+     DeclareOpInterfaceMethods<RegionBranchOpInterface>]> {
   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();
+
+    /// 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 getSuccessorEntryOperands(unsigned index);
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
   let hasRegionVerifier = 1;
 }
 

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. LoopOp only has one region, so 0 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) {
+  // If the predecessor is the affine.for op, branch into the body using the
+  // iterator arguments.
+  if (!index.hasValue()) {
+    regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+    return;
+  }
+
+  // Otherwise, the loop may branch back to itself or the parent operation.
+  assert(index.getValue() == 0 && "expected loop region");
+  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);
   if (tripCount.hasValue() && tripCount.getValue() == 0)
     return true;
   return false;
 }
 

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,