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

[McCowan Zhang]

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

Fixed incorrect bufferization interaction with cf.assert
- reordered bufferization condition checking
- fixed hasNeitherAllocateNorFreeSideEffect checking bug
- implemented memory interface for cf.assert

>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/2] 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/2] 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