[Mlir-commits] [mlir] Bufferization with ControlFlow Asserts (PR #95868)

[McCowan Zhang]

https://github.com/mccowanzhang updated https://github.com/llvm/llvm-project/pull/95868

>From cf0fcd982075bb30a70dbcdeec96c705540a0579 Mon Sep 17 00:00:00 2001
From: McCowan Zhang <86526121+mccowanzhang at users.noreply.github.com>
Date: Mon, 17 Jun 2024 16:07:23 -0700
Subject: [PATCH 1/4] update assertop to implement memory interface

---
 mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
index f77b8cbbbc61d..2f10e4fe34448 100644
--- a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
+++ b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
@@ -37,7 +37,7 @@ class CF_Op<string mnemonic, list<Trait> traits = []> :
 // AssertOp
 //===----------------------------------------------------------------------===//
 
-def AssertOp : CF_Op<"assert"> {
+def AssertOp : CF_Op<"assert", [Pure]> {
   let summary = "Assert operation with message attribute";
   let description = [{
     Assert operation at runtime with single boolean operand and an error

>From 62a9268f37f28cf5f4cf977bbb42b8b1386cdbe2 Mon Sep 17 00:00:00 2001
From: McCowan Zhang <mccowan.z at ssi.samsung.com>
Date: Mon, 17 Jun 2024 16:47:20 -0700
Subject: [PATCH 2/4] reordered bufferization condition checks and fixed helper
 function bug

---
 .../OwnershipBasedBufferDeallocation.cpp           | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index bd5c4d4769216..ca5d0688b5b59 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -52,8 +52,8 @@ static bool isMemref(Value v) { return isa<BaseMemRefType>(v.getType()); }
 /// "Free" side effects.
 static bool hasNeitherAllocateNorFreeSideEffect(Operation *op) {
   if (isa<MemoryEffectOpInterface>(op))
-    return hasEffect<MemoryEffects::Allocate>(op) ||
-           hasEffect<MemoryEffects::Free>(op);
+    return !hasEffect<MemoryEffects::Allocate>(op) &&
+           !hasEffect<MemoryEffects::Free>(op);
   // If the op does not implement the MemoryEffectOpInterface but has has
   // recursive memory effects, then this op in isolation (without its body) does
   // not have any side effects. All the ops inside the regions of this op will
@@ -497,6 +497,11 @@ BufferDeallocation::verifyFunctionPreconditions(FunctionOpInterface op) {
 }
 
 LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op) {
+  // We do not care about ops that do not operate on buffers and have no
+  // Allocate/Free side effect.
+  if (!hasBufferSemantics(op) && hasNeitherAllocateNorFreeSideEffect(op))
+    return success();
+
   // (1) The pass does not work properly when deallocations are already present.
   // Alternatively, we could also remove all deallocations as a pre-pass.
   if (isa<DeallocOp>(op))
@@ -517,11 +522,6 @@ LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op) {
     return op->emitError(
         "ops with unknown memory side effects are not supported");
 
-  // We do not care about ops that do not operate on buffers and have no
-  // Allocate/Free side effect.
-  if (!hasBufferSemantics(op) && hasNeitherAllocateNorFreeSideEffect(op))
-    return success();
-
   // (3) Check that the control flow structures are supported.
   auto regions = op->getRegions();
   // Check that if the operation has at

>From 22e3bce00892fa3b4285445230b60d7ccd592d84 Mon Sep 17 00:00:00 2001
From: McCowan Zhang <mccowan.z at ssi.samsung.com>
Date: Tue, 18 Jun 2024 11:51:13 -0700
Subject: [PATCH 3/4] implemented memoryopsinterface with custom geteffects for
 cf.assert

---
 mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td | 3 ++-
 mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp         | 7 +++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
index 2f10e4fe34448..181bc5a2d07fa 100644
--- a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
+++ b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
@@ -37,7 +37,8 @@ class CF_Op<string mnemonic, list<Trait> traits = []> :
 // AssertOp
 //===----------------------------------------------------------------------===//
 
-def AssertOp : CF_Op<"assert", [Pure]> {
+def AssertOp : CF_Op<"assert",
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "Assert operation with message attribute";
   let description = [{
     Assert operation at runtime with single boolean operand and an error
diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index ede158db911ea..98b429de1fd85 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -89,6 +89,13 @@ LogicalResult AssertOp::canonicalize(AssertOp op, PatternRewriter &rewriter) {
   return failure();
 }
 
+// This side effect models "program termination". 
+void AssertOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  effects.emplace_back(MemoryEffects::Write::get());
+}
+
 //===----------------------------------------------------------------------===//
 // BranchOp
 //===----------------------------------------------------------------------===//

>From 4e24bf18548187a4f11e25f7ba5a8d76cde8e28a Mon Sep 17 00:00:00 2001
From: McCowan Zhang <mccowan.z at ssi.samsung.com>
Date: Mon, 24 Jun 2024 14:56:09 -0700
Subject: [PATCH 4/4] added small test case

---
 .../Bufferization/Transforms/one-shot-bufferize.mlir  | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
index 0ed3a9f077ce0..84c7bad9555d2 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
@@ -269,3 +269,14 @@ func.func @materialize_in_dest_raw(%f: f32, %f2: f32, %idx: index) -> (tensor<5x
 
   return %0, %r : tensor<5xf32>, f32
 }
+
+// -----
+
+// CHECK-LABEL: func @func_with_assert(
+//       CHECK: %0 = arith.cmpi slt, %arg0, %arg1 : index
+//       CHECK: cf.assert %0, "%arg0 must be less than %arg1"
+func.func @func_with_assert(%arg0: index, %arg1: index) {
+  %0 = arith.cmpi slt, %arg0, %arg1 : index
+  cf.assert %0, "%arg0 must be less than %arg1"
+  return
+}
\ No newline at end of file