[Mlir-commits] [mlir] 7fbdee3 - Add RegionBranchOpInterface for AffineIf Op

Uday Bondhugula llvmlistbot at llvm.org
Wed Nov 23 00:50:17 PST 2022


Author: Akshay Baviskar
Date: 2022-11-23T14:19:22+05:30
New Revision: 7fbdee3e29203f2ffd1996c2919096e0bfe7c93b

URL: https://github.com/llvm/llvm-project/commit/7fbdee3e29203f2ffd1996c2919096e0bfe7c93b
DIFF: https://github.com/llvm/llvm-project/commit/7fbdee3e29203f2ffd1996c2919096e0bfe7c93b.diff

LOG: Add RegionBranchOpInterface for AffineIf Op

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

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D130962

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 148e2a4174ac2..cef199f2385af 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -353,7 +353,9 @@ def AffineForOp : Affine_Op<"for",
 
 def AffineIfOp : Affine_Op<"if",
                            [ImplicitAffineTerminator, RecursivelySpeculatable,
-                            RecursiveMemoryEffects, NoRegionArguments]> {
+                            RecursiveMemoryEffects, NoRegionArguments,
+                            DeclareOpInterfaceMethods<RegionBranchOpInterface>
+                           ]> {
   let summary = "if-then-else operation";
   let description = [{
     Syntax:

diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 4a581727e190c..bff38658741ed 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -911,7 +911,7 @@ 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.
@@ -2513,6 +2513,29 @@ struct AlwaysTrueOrFalseIf : public OpRewritePattern<AffineIfOp> {
 };
 } // 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.

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
index 1b3be68dd5bdc..08696da3140e7 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
@@ -1314,3 +1314,107 @@ func.func @select_aliases(%arg0: index, %arg1: memref<?xi8>, %arg2: i1) {
   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


        


More information about the Mlir-commits mailing list