Author: McCowan Zhang
Date: 2024-06-26T08:00:39+02:00
New Revision: a159b36724d9b7dc67702d742ffd63503f6a1417
URL: https://github.com/llvm/llvm-project/commit/a159b36724d9b7dc67702d742ffd63503f6a1417
DIFF: https://github.com/llvm/llvm-project/commit/a159b36724d9b7dc67702d742ffd63503f6a1417.diff
LOG: Bufferization with ControlFlow Asserts (#95868)
Fixed incorrect bufferization interaction with cf.assert
- reordered bufferization condition checking
- fixed hasNeitherAllocateNorFreeSideEffect checking bug
- implemented memory interface for cf.assert
---------
Co-authored-by: McCowan Zhang <mccowan.z at ssi.samsung.com>
Added:
mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/misc-other.mlir
Modified:
mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
index f77b8cbbbc61d..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"> {
+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/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
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
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/misc-other.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/misc-other.mlir
new file mode 100644
index 0000000000000..05e52848ca877
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/misc-other.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation -split-input-file %s
+
+// Test Case: ownership-based-buffer-deallocation should not fail
+// with cf.assert op
+
+// 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
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
index 0ed3a9f077ce0..dbf8d6563477b 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
@@ -268,4 +268,4 @@ func.func @materialize_in_dest_raw(%f: f32, %f2: f32, %idx: index) -> (tensor<5x
%r = tensor.extract %dest_filled[%idx] : tensor<5xf32>
return %0, %r : tensor<5xf32>, f32
-}
+}
\ No newline at end of file