[mlir][bufferization] Don't clone on unknown ownership and verify function boundary ABI #66626

maerhart · 2023-09-18T10:27:48Z

Inserting clones requires a lot of assumptions to hold on the input IR, e.g.,
all writes to a buffer need to dominate all reads. This is not guaranteed by
one-shot bufferization and isn't easy to verify, thus it could quickly lead to
incorrect results that are hard to debug. This commit changes the mechanism of
how an ownership indicator is materialized when there is not already a unique
ownership present. Additionally, we don't create copies of returned memrefs
anymore when we don't have ownership. Instead, we insert assert operations to
make sure we have ownership at runtime, or otherwise report to the user that
correctness could not be guaranteed.

llvmbot · 2023-09-18T10:29:26Z

@llvm/pr-subscribers-mlir-cf
@llvm/pr-subscribers-mlir-bufferization
@llvm/pr-subscribers-mlir-arith
@llvm/pr-subscribers-mlir-scf
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-sparse
@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir-tensor

Changes

Adds a pass option to the ownership-based-buffer-deallocation pass to
forbid insertion of clone operations. This is necessary to support IR
that does not have the property that every buffer write dominates every
buffer read to the same buffer. Instead of silently producing invalid
IR, the pass would then emit an error.
This is a restriction in the old buffer-deallocation pass, but the new
function boundary ABI was not enforced in this old pass. Having this
option allows easier migration from the old to the new deallocation
pass because enabling this option allows the new deallocation pass to
fix IR that does not adhere to the function boundary ABI (in some
situations).

Depends on #66619
Only review top commit.

Patch is 168.54 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/66626.diff

74 Files Affected:

(modified) mlir/docs/Bufferization.md (+10-36)
(modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h (+11-2)
(modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td (+2-2)
(modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h (+1-15)
(modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td (-10)
(modified) mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h (+17)
(modified) mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td (+1-2)
(modified) mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h (+3-3)
(modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h (+5-3)
(modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td (+14-17)
(modified) mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h (+1-2)
(modified) mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp (+3-3)
(modified) mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp (+10-5)
(modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+4-57)
(modified) mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp (-35)
(modified) mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp (-11)
(modified) mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp (+13-1)
(modified) mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt (+1)
(modified) mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp (+2-3)
(modified) mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp (+3-10)
(modified) mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp (+1-1)
(modified) mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp (+1-7)
(modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp (-78)
(modified) mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp (+58-27)
(modified) mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp (-21)
(modified) mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp (+1-6)
(modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+1-1)
(modified) mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp (+7-16)
(modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp (-6)
(modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp (-8)
(modified) mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp (-3)
(modified) mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp (+11-29)
(modified) mlir/python/mlir/dialects/_bufferization_transform_ops_ext.py (+4-8)
(modified) mlir/test/Dialect/Arith/one-shot-bufferize.mlir (+5-5)
(modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir (+7-1)
(modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir (+9-1)
(modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir (+157-56)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-allow-return-allocs.mlir (+4-4)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis-empty-tensor-elimination.mlir (+1-1)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-compat.mlir (+6-16)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir (+1-1)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir (+7-9)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-pass-statistics.mlir (-1)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir (-6)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-allow-return-allocs.mlir (+6-6)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir (+4-4)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir (-147)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-out-params.mlir (+3-3)
(modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir (+6-15)
(modified) mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space.mlir (+3-3)
(modified) mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir (+9-18)
(modified) mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir (-3)
(modified) mlir/test/Dialect/Bufferization/invalid.mlir (-32)
(modified) mlir/test/Dialect/ControlFlow/one-shot-bufferize-analysis.mlir (+1-1)
(modified) mlir/test/Dialect/ControlFlow/one-shot-bufferize-invalid.mlir (+1-1)
(modified) mlir/test/Dialect/ControlFlow/one-shot-bufferize.mlir (+2-2)
(modified) mlir/test/Dialect/Linalg/one-shot-bufferize-analysis.mlir (+1-1)
(modified) mlir/test/Dialect/Linalg/one-shot-bufferize.mlir (+5-7)
(modified) mlir/test/Dialect/SCF/one-shot-bufferize-allow-return-allocs-no-deallocs.mlir (+1-4)
(modified) mlir/test/Dialect/SCF/one-shot-bufferize-analysis.mlir (+4-4)
(modified) mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir (+1-1)
(modified) mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir (+15-15)
(modified) mlir/test/Dialect/SCF/one-shot-bufferize.mlir (+5-5)
(modified) mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir (+12-12)
(modified) mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir (+6-6)
(modified) mlir/test/Dialect/Tensor/one-shot-bufferize-tensor-copy-insertion.mlir (+3-3)
(modified) mlir/test/Dialect/Tensor/one-shot-bufferize.mlir (+5-43)
(modified) mlir/test/Integration/Dialect/Linalg/CPU/test-collapse-tensor.mlir (+1-1)
(modified) mlir/test/Integration/Dialect/Linalg/CPU/test-expand-tensor.mlir (+1-1)
(modified) mlir/test/Integration/Dialect/Linalg/CPU/test-one-shot-bufferize.mlir (+1-1)
(modified) mlir/test/Integration/Dialect/Linalg/CPU/test-tensor-e2e.mlir (+1-1)
(modified) mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp (+4-9)
(modified) mlir/test/python/dialects/transform_bufferization_ext.py (+2-4)
(modified) utils/bazel/llvm-project-overlay/mlir/BUILD.bazel (+1)

diff --git a/mlir/docs/Bufferization.md b/mlir/docs/Bufferization.md
index 09bec06743c7a65..d9d0751cae8c9dd 100644
--- a/mlir/docs/Bufferization.md
+++ b/mlir/docs/Bufferization.md
@@ -266,42 +266,16 @@ must be inserted due to a RaW conflict. E.g.:
 In the above example, a buffer copy of buffer(`%another_tensor`) (with `%cst`
 inserted) is yielded from the "then" branch.
 
-In both examples, a buffer is allocated inside of a block and then yielded from
-the block. Deallocation of such buffers is tricky and not currently implemented
-in an efficient way. For this reason, One-Shot Bufferize must be explicitly
-configured with `allow-return-allocs` to support such IR.
-
-When running with `allow-return-allocs`, One-Shot Bufferize may introduce
-allocations that cannot be deallocated by One-Shot Bufferize yet. For that
-reason, `-buffer-deallocation` must be run after One-Shot Bufferize. This buffer
-deallocation pass resolves yields of newly allocated buffers with copies. E.g.,
-the `scf.if` example above would bufferize to IR similar to the following:
-
-```mlir
-%0 = scf.if %c -> (memref<?xf32>) {
-  %1 = memref.alloc(...) : memref<?xf32>
-  ...
-  scf.yield %1 : memref<?xf32>
-} else {
-  %2 = memref.alloc(...) : memref<?xf32>
-  memref.copy %another_memref, %2
-  scf.yield %2 : memref<?xf32>
-}
-```
-
-In the bufferized IR, both branches return a newly allocated buffer, so it does
-not matter which if-branch was taken. In both cases, the resulting buffer `%0`
-must be deallocated at some point after the `scf.if` (unless the `%0` is
-returned/yielded from its block).
-
-Note: Buffer allocations that are returned from a function are not deallocated,
-not even with `-buffer-deallocation`. It is the caller's responsibility to
-deallocate the buffer. In the future, this could be automated with allocation
-hoisting (across function boundaries) or reference counting.
-
-One-Shot Bufferize can be configured to leak all memory and not generate any
-buffer deallocations with `create-deallocs=0`. This can be useful for
-compatibility with legacy code that has its own method of deallocating buffers.
+Note: Buffer allocations that are returned from a function are not deallocated.
+It is the caller's responsibility to deallocate the buffer. For the full
+function boundary ABI for MemRefs w.r.t. buffer deallocation refer to the
+[*Function Boundary ABI*](#function-boundary-abi) section. In the future, this
+could be automated with allocation hoisting (across function boundaries) or
+reference counting.
+
+One-Shot Bufferize leaks all memory and does not generate any buffer
+deallocations. The `-buffer-deallocation-pipeline` has to be run afterwards to
+insert the deallocation operations.
 
 ## Ownership-based Buffer Deallocation
 
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
index 7ac4592de7875fb..3aa61fae8c6caee 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
@@ -96,6 +96,14 @@ struct DeallocationOptions {
   // pass the ownership of MemRef values instead of adhering to the function
   // boundary ABI.
   bool privateFuncDynamicOwnership = false;
+
+  // Allows the pass to insert `bufferization.clone` operations. This is useful
+  // for supporting IR that does not adhere to the function boundary ABI
+  // initially (excl. external functions) and to support operations with results
+  // with 'Unknown' ownership. However, it requires that all buffer writes
+  // dominate all buffer reads (i.e., only enable this option if your IR is
+  // guaranteed to have this property).
+  bool allowCloning = false;
 };
 
 /// This class collects all the state that we need to perform the buffer
@@ -142,8 +150,9 @@ class DeallocationState {
   /// a new SSA value, returned as the first element of the pair, which has
   /// 'Unique' ownership and can be used instead of the passed Value with the
   /// the ownership indicator returned as the second element of the pair.
-  std::pair<Value, Value>
-  getMemrefWithUniqueOwnership(OpBuilder &builder, Value memref, Block *block);
+  FailureOr<std::pair<Value, Value>>
+  getMemrefWithUniqueOwnership(const DeallocationOptions &options,
+                               OpBuilder &builder, Value memref, Block *block);
 
   /// Given two basic blocks and the values passed via block arguments to the
   /// destination block, compute the list of MemRefs that have to be retained in
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td
index 3e11432c65c5f08..3b9a9c3f4fef667 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td
@@ -56,7 +56,7 @@ def BufferDeallocationOpInterface :
           method (which is especially important if operations are created that
           cannot be easily canonicalized away anymore).
         }],
-        /*retType=*/"std::pair<Value, Value>",
+        /*retType=*/"FailureOr<std::pair<Value, Value>>",
         /*methodName=*/"materializeUniqueOwnershipForMemref",
         /*args=*/(ins "DeallocationState &":$state,
                       "const DeallocationOptions &":$options,
@@ -65,7 +65,7 @@ def BufferDeallocationOpInterface :
         /*methodBody=*/[{}],
         /*defaultImplementation=*/[{
           return state.getMemrefWithUniqueOwnership(
-            builder, memref, memref.getParentBlock());
+            options, builder, memref, memref.getParentBlock());
         }]>,
   ];
 }
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 9ec44dfd16a0c00..1c715f8b9a53ef3 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -361,10 +361,6 @@ struct BufferizationOptions {
   /// used.
   UnknownTypeConverterFn unknownTypeConverterFn = nullptr;
 
-  /// Specifies whether dealloc ops should be generated along with alloc ops. If
-  /// not, new memory allocations will leak.
-  bool createDeallocs = true;
-
   /// Seed for the analysis fuzzer. If set to `0`, the fuzzer is deactivated.
   /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
@@ -588,13 +584,9 @@ class AnalysisState {
 /// undefined contents is allocated.
 FailureOr<Value>
 allocateTensorForShapedValue(OpBuilder &b, Location loc, Value shapedValue,
-                             bool escape, const BufferizationOptions &options,
+                             const BufferizationOptions &options,
                              bool copy = true);
 
-/// Return `true` if the allocation of the given op is guaranteed to not escape
-/// the containing block.
-bool allocationDoesNotEscape(OpResult opResult);
-
 /// Lookup the buffer for the given value. If the value was not bufferized
 /// yet, wrap it in a ToMemrefOp. Otherwise, it is the result of a ToTensorOp,
 /// from which the memref operand is returned.
@@ -641,12 +633,6 @@ OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
   return newOp;
 }
 
-/// Return `true` if the buffer of given OpResult should be deallocated. This
-/// function should be called during `BufferizableOpInterface::bufferize`
-/// implementations that allocate a new buffer for the given OpResult.
-bool shouldDeallocateOpResult(OpResult opResult,
-                              const BufferizationOptions &options);
-
 /// Return a MemRefType to which the type of the given value can be bufferized.
 ///
 /// If possible, op bufferization implementations should not use this function
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td
index e9c140859344ef8..0d509e69349e918 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td
@@ -60,16 +60,6 @@ def Bufferization_Dialect : Dialect {
     /// arguments during One-Shot Module Bufferize.
     constexpr const static ::llvm::StringLiteral
         kBufferLayoutAttrName = "bufferization.buffer_layout";
-
-    /// Attribute name used to mark escaping behavior of buffer allocations.
-    /// Escaping allocations cannot be deallocated in the same block and must
-    /// be treated specially: They are currently deallocated with the
-    /// BufferDeallocation pass.
-    ///
-    /// Note: Only ops with at least one OpResult that bufferizes to a buffer
-    /// allocation (as per BufferizableOpInterface) may have this attribute.
-    constexpr const static ::llvm::StringLiteral
-        kEscapeAttrName = "bufferization.escape";
   }];
   let hasOperationAttrVerify = 1;
 }
diff --git a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
index 7acacb763cd2c18..7500257ed95eac8 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
@@ -17,6 +17,7 @@
 
 namespace mlir {
 namespace bufferization {
+struct DeallocationOptions;
 
 /// Options for the buffer deallocation pipeline.
 struct BufferDeallocationPipelineOptions
@@ -28,6 +29,22 @@ struct BufferDeallocationPipelineOptions
           "dynamically pass ownership of memrefs to callees. This can enable "
           "earlier deallocations."),
       llvm::cl::init(false)};
+  PassOptions::Option<bool> allowCloning{
+      *this, "allow-cloning",
+      llvm::cl::desc(
+          "Allows the pass to insert `bufferization.clone` operations. This is "
+          "useful for supporting IR that does not adhere to the function "
+          "boundary ABI initially (excl. external functions) and to support "
+          "operations with results with 'Unknown' ownership. However, it "
+          "requires that all buffer writes dominate all buffer reads (i.e., "
+          "only enable this option if your IR is guaranteed to have this "
+          "property)."),
+      llvm::cl::init(false)};
+
+  /// Convert this BufferDeallocationPipelineOptions struct to a
+  /// DeallocationOptions struct to be passed to the
+  /// OwnershipBasedBufferDeallocationPass.
+  DeallocationOptions asDeallocationOptions() const;
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td b/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
index 84bd047e6d51eed..a0eb5ff00cb9fea 100644
--- a/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
@@ -82,10 +82,9 @@ def OneShotBufferizeOp
   let arguments = (
       ins TransformHandleTypeInterface:$target,
       OptionalAttr<LayoutMapOption>:$function_boundary_type_conversion,
-      DefaultValuedAttr<BoolAttr, "false">:$allow_return_allocs,
+      DefaultValuedAttr<BoolAttr, "false">:$allow_return_allocs_from_loops,
       DefaultValuedAttr<BoolAttr, "false">:$allow_unknown_ops,
       DefaultValuedAttr<BoolAttr, "false">:$bufferize_function_boundaries,
-      DefaultValuedAttr<BoolAttr, "true">:$create_deallocs,
       DefaultValuedAttr<BoolAttr, "false">:$test_analysis_only,
       DefaultValuedAttr<BoolAttr, "false">:$print_conflicts,
       DefaultValuedAttr<StrAttr, "\"memref.copy\"">:$memcpy_op);
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
index 585c7ca92c71895..328aff07280a92b 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
@@ -28,9 +28,9 @@ struct OneShotBufferizationOptions : public BufferizationOptions {
 
   OneShotBufferizationOptions() = default;
 
-  /// Specifies whether returning newly allocated memrefs should be allowed.
-  /// Otherwise, a pass failure is triggered.
-  bool allowReturnAllocs = false;
+  /// Specifies whether returning newly allocated memrefs from loops should be
+  /// allowed.  Otherwise, a pass failure is triggered.
+  bool allowReturnAllocsFromLoops = false;
 
   /// Specifies whether the tensor IR should be annotated with alias sets.
   bool dumpAliasSets = false;
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
index 92520eb13da6875..37a3942f7bac6c5 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
@@ -1,6 +1,7 @@
 #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 
+#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
@@ -31,7 +32,7 @@ std::unique_ptr<Pass> createBufferDeallocationPass();
 /// Creates an instance of the OwnershipBasedBufferDeallocation pass to free all
 /// allocated buffers.
 std::unique_ptr<Pass> createOwnershipBasedBufferDeallocationPass(
-    bool privateFuncDynamicOwnership = false);
+    const DeallocationOptions &options = DeallocationOptions());
 
 /// Creates a pass that optimizes `bufferization.dealloc` operations. For
 /// example, it reduces the number of alias checks needed at runtime using
@@ -134,8 +135,9 @@ func::FuncOp buildDeallocationLibraryFunction(OpBuilder &builder, Location loc,
 LogicalResult deallocateBuffers(Operation *op);
 
 /// Run ownership basedbuffer deallocation.
-LogicalResult deallocateBuffersOwnershipBased(FunctionOpInterface op,
-                                              bool privateFuncDynamicOwnership);
+LogicalResult deallocateBuffersOwnershipBased(
+    FunctionOpInterface op,
+    const DeallocationOptions &options = DeallocationOptions());
 
 /// Creates a pass that moves allocations upwards to reduce the number of
 /// required copies that are inserted during the BufferDeallocation pass.
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index f3c2a29c0589f29..5b8af7a975c34b5 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -223,6 +223,14 @@ def OwnershipBasedBufferDeallocation : Pass<
            "Allows to add additional arguments to private functions to "
            "dynamically pass ownership of memrefs to callees. This can enable "
            "earlier deallocations.">,
+    Option<"allowCloning", "allow-cloning", "bool", /*default=*/"false",
+           "Allows the pass to insert `bufferization.clone` operations. This "
+           "is useful for supporting IR that does not adhere to the function "
+           "boundary ABI initially (excl. external functions) and to support "
+           "operations with results with 'Unknown' ownership. However, it "
+           "requires that all buffer writes dominate all buffer reads (i.e., "
+           "only enable this option if your IR is guaranteed to have this "
+           "property).">,
   ];
   let constructor = "mlir::bufferization::createOwnershipBasedBufferDeallocationPass()";
 
@@ -387,15 +395,9 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
     example, `tensor.generate` is not in destination-passing style and always
     results in a new buffer allocation.
 
-    One-Shot Bufferize deallocates all buffers that it allocates. Yielding newly
-    allocated buffers from a block can lead to bad performance because
-    additional buffer copies are often needed to make sure that every buffer
-    allocation is also deallocated again. By default, such IR is rejected by
-    One-Shot Bufferize. Such IR can be allowed with `allow-return-allocs`. In
-    that case, the `-buffer-deallocation` pass should be run after One-Shot
-    Bufferize. Note that new buffer allocations that are returned from a
-    function can currently not be deallocated by `-buffer-deallocation` and
-    leak.
+    One-Shot Bufferize does not deallocate any buffers that it allocates. The
+    `-buffer-deallocation` pass should be run after One-Shot Bufferize to insert
+    the deallocation operations necessary to eliminate memory leaks.
 
     One-Shot Bufferize will by default reject IR that contains non-bufferizable
     op, i.e., ops that do not implemement BufferizableOpInterface. Such IR can
@@ -462,9 +464,9 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
     `test-analysis-only`.
   }];
   let options = [
-    Option<"allowReturnAllocs", "allow-return-allocs", "bool",
-            /*default=*/"false",
-           "Allows returning/yielding new allocations from a block.">,
+    Option<"allowReturnAllocsFromLoops", "allow-return-allocs-from-loops",
+           "bool", /*default=*/"false",
+           "Allows returning/yielding new allocations from a loop.">,
     Option<"allowUnknownOps", "allow-unknown-ops", "bool",
            /*default=*/"false",
            "Allows unknown (not bufferizable) ops in the input IR.">,
@@ -479,9 +481,6 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
            "Bufferize function boundaries (experimental).">,
     Option<"copyBeforeWrite", "copy-before-write", "bool", /*default=*/"false",
            "Skip the analysis. Make a buffer copy on every write.">,
-    Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
-           "Specify if buffers should be deallocated. For compatibility with "
-           "core bufferization passes.">,
     ListOption<"dialectFilter", "dialect-filter", "std::string",
                "Restrict bufferization to ops from these dialects.">,
     Option<"dumpAliasSets", "dump-alias-sets", "bool", /*default=*/"false",
@@ -513,8 +512,6 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
   let statistics = [
     Statistic<"numBufferAlloc", "num-buffer-alloc",
               "Number of buffer allocations">,
-    Statistic<"numBufferDealloc", "num-buffer-dealloc",
-              "Number of buffer deallocations">,
     Statistic<"numTensorInPlace", "num-tensor-in-place",
               "Number of in-place tensor OpOperands">,
     Statistic<"numTensorOutOfPlace", "num-tensor-out-of-place",
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
index 71fcfc84919eb5c..c88963d399c4c92 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
@@ -93,8 +93,7 @@ struct SparseCompilerOptions
       desc("Specify if the temporary buffers created by the sparse "
            "compiler should be deallocated. For compatibility with core "
            "bufferization passes. "
-           "This option is only used when enable-runtime-library=false. "
-           "See also create-deallocs for BufferizationOption."),
+           "This option is only used when enable-runtime-library=false."),
       init(true)};
 
   PassOptions::Option<int32_t> vectorLength{
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index f2e7732e8ea4aa3..8ab4717739a7643 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -53,7 +53,7 @@ struct SelectOpInterface
     return op; // nothing to do
   }
 
-  std::pair<Value, Value>
+  FailureOr<std::pair<Value, Value>>
   materializeUniqueOwnershipForMemref(Operation *op, DeallocationState &state,
                                       const DeallocationOptions &options,
                                       OpBuilder &builder, Value value) const {
@@ -64,14 +64,14 @@ struct SelectOpInterface
     Block *b...
[truncated]

mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h

mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp

mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h

mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td

mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h

mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp

…ction boundary ABI Inserting clones requires a lot of assumptions to hold on the input IR, e.g., all writes to a buffer need to dominate all reads. This is not guaranteed by one-shot bufferization and isn't easy to verify, thus it could quickly lead to incorrect results that are hard to debug. This commit changes the mechanism of how an ownership indicator is materialized when there is not already a unique ownership present. Additionally, we don't create copies of returned memrefs anymore when we don't have ownership. Instead, we insert assert operations to make sure we have ownership at runtime, or otherwise report to the user that correctness could not be guaranteed.

…rify function boundary ABI (#66626)" This reverts commit aa9eb47. It introduced a double free in a test case. Reverting to have some time for fixing this and relanding later.

…ction boundary ABI (llvm#66626) Inserting clones requires a lot of assumptions to hold on the input IR, e.g., all writes to a buffer need to dominate all reads. This is not guaranteed by one-shot bufferization and isn't easy to verify, thus it could quickly lead to incorrect results that are hard to debug. This commit changes the mechanism of how an ownership indicator is materialized when there is not already a unique ownership present. Additionally, we don't create copies of returned memrefs anymore when we don't have ownership. Instead, we insert assert operations to make sure we have ownership at runtime, or otherwise report to the user that correctness could not be guaranteed.

…rify function boundary ABI (llvm#66626)" This reverts commit aa9eb47. It introduced a double free in a test case. Reverting to have some time for fixing this and relanding later.

amrami · 2025-01-15T10:30:56Z

Hi!
I have a test case where an argument and a result are aliased.
Today a clone is added and I started looking why.
So I found this commit that was reverted.
Is there a plan to push it? :)
Thanks in advance,
Maya

maerhart requested a review from matthias-springer September 18, 2023 10:27

llvmbot added mlir:linalg mlir:sparse Sparse compiler in MLIR mlir:python MLIR Python bindings mlir mlir:tensor mlir:bufferization Bufferization infrastructure mlir:scf mlir:arith mlir:cf labels Sep 18, 2023

maerhart force-pushed the merhart_cloning_option branch from ab01756 to 1647897 Compare September 18, 2023 14:59

maerhart mentioned this pull request Sep 18, 2023

[mlir][bufferization] OwnershipBasedBufferDeallocation: support unstructured control flow loops #66657

Open

matthias-springer reviewed Sep 20, 2023

View reviewed changes

maerhart force-pushed the merhart_cloning_option branch from 1647897 to 9bda5e8 Compare September 27, 2023 07:52

maerhart changed the title ~~[mlir][bufferization] Add OwnershipBasedBufferDeallocation pass option to forbid clones~~ [mlir][bufferization] Don't clone on unknown ownership and verify function boundary ABI Sep 27, 2023

maerhart requested a review from matthias-springer September 27, 2023 07:56

matthias-springer reviewed Sep 27, 2023

View reviewed changes

matthias-springer approved these changes Sep 27, 2023

View reviewed changes

maerhart mentioned this pull request Sep 27, 2023

[mlir][bufferization] Add deallocation option to remove existing dealloc operations, add option to specify the kind of alloc operations to consider #67556

Open

maerhart force-pushed the merhart_cloning_option branch from 1675d91 to 5b0f67b Compare September 28, 2023 07:43

maerhart merged commit aa9eb47 into llvm:main Sep 28, 2023

maerhart deleted the merhart_cloning_option branch September 28, 2023 08:45

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][bufferization] Don't clone on unknown ownership and verify function boundary ABI #66626

[mlir][bufferization] Don't clone on unknown ownership and verify function boundary ABI #66626

Uh oh!

maerhart commented Sep 18, 2023 •

edited

Loading

Uh oh!

llvmbot commented Sep 18, 2023 •

edited

Loading

Depends on #66619
Only review top commit.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

amrami commented Jan 15, 2025

Uh oh!

Uh oh!

[mlir][bufferization] Don't clone on unknown ownership and verify function boundary ABI #66626

[mlir][bufferization] Don't clone on unknown ownership and verify function boundary ABI #66626

Uh oh!

Conversation

maerhart commented Sep 18, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Sep 18, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Depends on #66619 Only review top commit.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

amrami commented Jan 15, 2025

Uh oh!

Uh oh!

maerhart commented Sep 18, 2023 •

edited

Loading

llvmbot commented Sep 18, 2023 •

edited

Loading

Depends on #66619
Only review top commit.