Add RegionBranchOpInterface for AffineIf Op

[akshaybaviskar]

Summary

Adds RegionBranchOpInterface for AffineIf Op and tests it
using buffer deallocation pass.

[bondhugula]

LGTM

[akshaybaviskar]

Hi @reviewers,
Can this patch be merged now?
Thanks.

[bondhugula]

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

2534

Nit: is -> is the
Comma after region.

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

1304

Nit: then in backticks to avoid confusion.

[akshaybaviskar]

Add RegionBranchOpInterface for AffineIf Op

Adds RegionBranchOpInterface for AffineIf Op and tests it
using buffer deallocation pass.

[akshaybaviskar]

Add RegionBranchOpInterface for AffineIf Op

Adds RegionBranchOpInterface for AffineIf Op and tests it
using buffer deallocation pass.

[bondhugula]

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

2523

This leads to compile warnings - please fix.

[akshaybaviskar]

Add RegionBranchOpInterface for AffineIf Op

Adds RegionBranchOpInterface for AffineIf Op and tests it
using buffer deallocation pass.

[bondhugula]

clang-format shows a change is needed.

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

def AffineForOp : Affine_Op<"for",
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
   let hasRegionVerifier = 1;
 }
 
 def AffineIfOp : Affine_Op<"if",
                            [ImplicitAffineTerminator, RecursivelySpeculatable,
-                            RecursiveMemoryEffects, NoRegionArguments]> {
+                            RecursiveMemoryEffects, NoRegionArguments,
+                            DeclareOpInterfaceMethods<RegionBranchOpInterface>
+                           ]> {
   let summary = "if-then-else operation";
   let description = [{
     Syntax:
 
     ```
     operation  ::= `affine.if` if-op-cond `{` op* `}` (`else` `{` op* `}`)?
     if-op-cond ::= integer-set-attr dim-and-symbol-use-list
     ```

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

 /// The operation is also expected to always produce a single result. Return an
 /// `OpFoldResult` containing the Attribute representing the folded constant if
 /// complete folding was possible and a Value produced by the created operation
 /// otherwise.
 template <typename OpTy, typename... Args>
 static std::enable_if_t<OpTy::template hasTrait<OpTrait::OneResult>(),
                         OpFoldResult>
 createOrFold(OpBuilder &b, Location loc, ValueRange operands,
-             Args &&...leadingArguments) {
+             Args &&... leadingArguments) {
   // Identify the constant operands and extract their values as attributes.
   // Note that we cannot use the original values directly because the list of
   // operands may have changed due to canonicalization and composition.
   SmallVector<Attribute> constantOperands;
   constantOperands.reserve(operands.size());
   for (Value operand : operands) {
     IntegerAttr attr;
     if (matchPattern(operand, m_Constant(&attr)))
LogicalResult matchAndRewrite(AffineIfOp op,
     // Erase the "blockToMoveTerminator" operation since it is now in the parent
     // operation block, which already has its own terminator.
     rewriter.eraseOp(blockToMoveTerminator);
     return success();
   }
 };
 } // namespace
 
+/// AffineIfOp has two regions -- `then` and `else`. The flow of data should be
+/// as follows: AffineIfOp -> `then`/`else` -> AffineIfOp
+void AffineIfOp::getSuccessorRegions(
+    Optional<unsigned> index, ArrayRef<Attribute> operands,
+    SmallVectorImpl<RegionSuccessor> &regions) {
+  // If the predecessor is an AffineIfOp, then branching into both `then` and
+  // `else` region is valid.
+  if (!index.has_value()) {
+    regions.reserve(2);
+    regions.push_back(
+        RegionSuccessor(&getThenRegion(), getThenRegion().getArguments()));
+    // Don't consider the else region if it is empty.
+    if (!getElseRegion().empty())
+      regions.push_back(
+          RegionSuccessor(&getElseRegion(), getElseRegion().getArguments()));
+    return;
+  }
+
+  // If the predecessor is the `else`/`then` region, then branching into parent
+  // op is valid.
+  regions.push_back(RegionSuccessor(getResults()));
+}
+
 LogicalResult AffineIfOp::verify() {
   // Verify that we have a condition attribute.
   // FIXME: This should be specified in the arguments list in ODS.
   auto conditionAttr =
       (*this)->getAttrOfType<IntegerSetAttr>(getConditionAttrStrName());
   if (!conditionAttr)
     return emitOpError("requires an integer set attribute named 'condition'");
 

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

 // CHECK-DAG: memref.dealloc %[[WHILE]]#1
 // CHECK-DAG: memref.dealloc %[[WHILE]]#2
 // CHECK-NEXT: return
   return
 }
 
 // -----
 
 func.func @select_aliases(%arg0: index, %arg1: memref<?xi8>, %arg2: i1) {
   // CHECK: memref.alloc
   // CHECK: memref.alloc
   // CHECK: arith.select
   // CHECK: test.copy
   // CHECK: memref.dealloc
   // CHECK: memref.dealloc
   %0 = memref.alloc(%arg0) : memref<?xi8>
   %1 = memref.alloc(%arg0) : memref<?xi8>
   %2 = arith.select %arg2, %0, %1 : memref<?xi8>
   test.copy(%2, %arg1) : (memref<?xi8>, memref<?xi8>)
   return
 }
+
+// -----
+
+// Memref allocated in `then` region and passed back to the parent if op.
+#set = affine_set<() : (0 >= 0)>
+// CHECK-LABEL:  func @test_affine_if_1
+// CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>) -> memref<10xf32> {
+func.func @test_affine_if_1(%arg0: memref<10xf32>) -> memref<10xf32> {
+  %0 = affine.if #set() -> memref<10xf32> {
+    %alloc = memref.alloc() : memref<10xf32>
+    affine.yield %alloc : memref<10xf32>
+  } else {
+    affine.yield %arg0 : memref<10xf32>
+  }
+  return %0 : memref<10xf32>
+}
+// CHECK-NEXT:    %[[IF:.*]] = affine.if
+// CHECK-NEXT:      %[[MEMREF:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT:      %[[CLONED:.*]] = bufferization.clone %[[MEMREF]] : memref<10xf32> to memref<10xf32>
+// CHECK-NEXT:      memref.dealloc %[[MEMREF]] : memref<10xf32>
+// CHECK-NEXT:      affine.yield %[[CLONED]] : memref<10xf32>
+// CHECK-NEXT:    } else {
+// CHECK-NEXT:      %[[ARG0_CLONE:.*]] = bufferization.clone %[[ARG0]] : memref<10xf32> to memref<10xf32>
+// CHECK-NEXT:      affine.yield %[[ARG0_CLONE]] : memref<10xf32>
+// CHECK-NEXT:    }
+// CHECK-NEXT:    return %[[IF]] : memref<10xf32>
+
+// -----
+
+// Memref allocated before parent IfOp and used in `then` region.
+// Expected result: deallocation should happen after affine.if op.
+#set = affine_set<() : (0 >= 0)>
+// CHECK-LABEL:  func @test_affine_if_2() -> memref<10xf32> {
+func.func @test_affine_if_2() -> memref<10xf32> {
+  %alloc0 = memref.alloc() : memref<10xf32>
+  %0 = affine.if #set() -> memref<10xf32> {
+    affine.yield %alloc0 : memref<10xf32>
+  } else {
+    %alloc = memref.alloc() : memref<10xf32>
+    affine.yield %alloc : memref<10xf32>
+  }
+  return %0 : memref<10xf32>
+}
+// CHECK-NEXT:    %[[ALLOC:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT:    %[[IF_RES:.*]] = affine.if {{.*}} -> memref<10xf32> {
+// CHECK-NEXT:      %[[ALLOC_CLONE:.*]] = bufferization.clone %[[ALLOC]] : memref<10xf32> to memref<10xf32>
+// CHECK-NEXT:      affine.yield %[[ALLOC_CLONE]] : memref<10xf32>
+// CHECK-NEXT:    } else {
+// CHECK-NEXT:      %[[ALLOC2:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT:      %[[ALLOC2_CLONE:.*]] = bufferization.clone %[[ALLOC2]] : memref<10xf32> to memref<10xf32>
+// CHECK-NEXT:      memref.dealloc %[[ALLOC2]] : memref<10xf32>
+// CHECK-NEXT:      affine.yield %[[ALLOC2_CLONE]] : memref<10xf32>
+// CHECK-NEXT:    }
+// CHECK-NEXT:    memref.dealloc %[[ALLOC]] : memref<10xf32>
+// CHECK-NEXT:    return %[[IF_RES]] : memref<10xf32>
+
+// -----
+
+// Memref allocated before parent IfOp and used in `else` region.
+// Expected result: deallocation should happen after affine.if op.
+#set = affine_set<() : (0 >= 0)>
+// CHECK-LABEL:  func @test_affine_if_3() -> memref<10xf32> {
+func.func @test_affine_if_3() -> memref<10xf32> {
+  %alloc0 = memref.alloc() : memref<10xf32>
+  %0 = affine.if #set() -> memref<10xf32> {
+    %alloc = memref.alloc() : memref<10xf32>
+    affine.yield %alloc : memref<10xf32>
+  } else {
+    affine.yield %alloc0 : memref<10xf32>
+  }
+  return %0 : memref<10xf32>
+}
+// CHECK-NEXT:    %[[ALLOC:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT:    %[[IFRES:.*]] = affine.if {{.*}} -> memref<10xf32> {
+// CHECK-NEXT:      memref.alloc
+// CHECK-NEXT:      bufferization.clone
+// CHECK-NEXT:      memref.dealloc
+// CHECK-NEXT:      affine.yield
+// CHECK-NEXT:    } else {
+// CHECK-NEXT:      bufferization.clone
+// CHECK-NEXT:      affine.yield
+// CHECK-NEXT:    }
+// CHECK-NEXT:    memref.dealloc %[[ALLOC]] : memref<10xf32>
+// CHECK-NEXT:    return %[[IFRES]] : memref<10xf32>
+
+// -----
+
+// Memref allocated before parent IfOp and not used later.
+// Expected result: deallocation should happen before affine.if op.
+#set = affine_set<() : (0 >= 0)>
+// CHECK-LABEL:  func @test_affine_if_4({{.*}}: memref<10xf32>) -> memref<10xf32> {
+func.func @test_affine_if_4(%arg0 : memref<10xf32>) -> memref<10xf32> {
+  %alloc0 = memref.alloc() : memref<10xf32>
+  %0 = affine.if #set() -> memref<10xf32> {
+    affine.yield %arg0 : memref<10xf32>
+  } else {
+    %alloc = memref.alloc() : memref<10xf32>
+    affine.yield %alloc : memref<10xf32>
+  }
+  return %0 : memref<10xf32>
+}
+// CHECK-NEXT:    %[[ALLOC:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT:    memref.dealloc %[[ALLOC]] : memref<10xf32>
+// CHECK-NEXT:    affine.if