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

[donald chen]

https://github.com/cxy-1993 created https://github.com/llvm/llvm-project/pull/106220

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.

>From 68b335ff11c3cc7bf95c6712f1206ec4ec635399 Mon Sep 17 00:00:00 2001
From: donald chen <chenxunyu1993 at gmail.com>
Date: Tue, 27 Aug 2024 12:43:29 +0000
Subject: [PATCH] [mlir] [bufferize] Inplace the operands that must be
 inplaced.

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.
---
 .../mlir/Dialect/Bufferization/IR/BufferizationOps.td       | 5 +++++
 .../Dialect/Bufferization/IR/BufferizableOpInterface.cpp    | 6 ++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

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