[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> ®ions) {
+ // 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