[Mlir-commits] [mlir] [mlir] [bufferize] Inplace the operands that must be inplaced. (PR #106220)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: donald chen (cxy-1993)

<details>
<summary>Changes</summary>

Return true when an op operand marked as "must inplace" in a bufferizable interface, and the operand is the same as the argument of the isInplace function in AnalysisState.

Some operations (e.g., atomic operations) have write effect on certain operands and require in-place bufferization. In such cases, we need to ensure the in-place semantics during the bufferization process.

---
Full diff: https://github.com/llvm/llvm-project/pull/106220.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td (+5) 
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+4-2) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index 1c70a4b8df9254..d8f9b8f8fd28d0 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -540,6 +540,11 @@ def Bufferization_ToMemrefOp : Bufferization_Op<"to_memref", [
       return !getReadOnly();
     }
 
+    bool mustBufferizeInPlace(OpOperand &opOperand,
+                              const AnalysisState &state) {
+      return true;
+    }
+
     AliasingValueList getAliasingValues(
         OpOperand &opOperand, const AnalysisState &state) const {
       return {};
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index d51d63f243ea0c..d8a3aab3cb6575 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -600,8 +600,10 @@ bool AnalysisState::canOmitTensorCopy(OpOperand &opOperand) const {
 }
 
 bool AnalysisState::isInPlace(OpOperand &opOperand) const {
-  // ToMemrefOps are always in-place.
-  if (isa<ToMemrefOp>(opOperand.getOwner()))
+  // We should always inplace operands that must be inplaced.
+  auto bufferizableOp =
+      getOptions().dynCastBufferizableOp(opOperand.getOwner());
+  if (bufferizableOp && bufferizableOp.mustBufferizeInPlace(opOperand, *this))
     return true;
 
   // In the absence of analysis information, OpOperands that bufferize to a

``````````

</details>


https://github.com/llvm/llvm-project/pull/106220