[Mlir-commits] [mlir] [mlir][vector] Fix attaching write effects on transfer_write's base (PR #142940)

[Simone Pellegrini]

https://github.com/simpel01 created https://github.com/llvm/llvm-project/pull/142940

This fixes an issue with `TransferWriteOp`'s implementation of the `MemoryEffectOpInterface` where the write effect was attached to the stored value rather than the base.

This had the effect that when asking for the memory effects for the input memref buffer using `getEffectsOnValue(...)`, the function would return no-effects (as the effect would have been attached to the stored value rather than the input buffer).

>From 68996d365dfe01d2ed89b87b8b61a5b1ecbbac26 Mon Sep 17 00:00:00 2001
From: Simone Pellegrini <simone.pellegrini at arm.com>
Date: Thu, 5 Jun 2025 09:21:08 +0200
Subject: [PATCH] [mlir][vector] Fix attaching write effects on
 transfer_write's base

This fixes an issue with `TransferWriteOp`'s implementation of the
`MemoryEffectOpInterface` where the write effect was attached to the
stored value rather than the base.

This had the effect that when asking for the memory effects for the
input memref buffer using `getEffectsOnValue(...)`, the function would
return no-effects (as the effect would have been attached to the stored
value rather than the input buffer).
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp        |  2 +-
 .../Dialect/Bufferization/side-effects.mlir     |  6 +++---
 mlir/test/Dialect/Vector/side-effects.mlir      | 17 +++++++++++++++++
 mlir/test/IR/test-side-effects.mlir             |  8 ++++----
 mlir/test/lib/IR/TestSideEffects.cpp            | 12 ++++++------
 5 files changed, 31 insertions(+), 14 deletions(-)
 create mode 100644 mlir/test/Dialect/Vector/side-effects.mlir

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index fcfb401fd9867..bd926b82cf14f 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5034,7 +5034,7 @@ void TransferWriteOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   if (llvm::isa<MemRefType>(getShapedType()))
-    effects.emplace_back(MemoryEffects::Write::get(), &getValueToStoreMutable(),
+    effects.emplace_back(MemoryEffects::Write::get(), &getBaseMutable(),
                          SideEffects::DefaultResource::get());
 }
 
diff --git a/mlir/test/Dialect/Bufferization/side-effects.mlir b/mlir/test/Dialect/Bufferization/side-effects.mlir
index 841490e9f3234..129fc8b32c270 100644
--- a/mlir/test/Dialect/Bufferization/side-effects.mlir
+++ b/mlir/test/Dialect/Bufferization/side-effects.mlir
@@ -1,9 +1,9 @@
 // RUN: mlir-opt %s --test-side-effects --verify-diagnostics
 
 func.func @test_side_effects(%arg0: memref<2xi32>) -> memref<2xi32> {
-  // expected-remark @below {{found an instance of 'read' on a op operand, on resource '<Default>'}}
-  // expected-remark @below {{found an instance of 'write' on a op result, on resource '<Default>'}}
-  // expected-remark @below {{found an instance of 'allocate' on a op result, on resource '<Default>'}}
+  // expected-remark @below {{found an instance of 'read' on op operand 0, on resource '<Default>'}}
+  // expected-remark @below {{found an instance of 'write' on op result 0, on resource '<Default>'}}
+  // expected-remark @below {{found an instance of 'allocate' on op result 0, on resource '<Default>'}}
   %0 = bufferization.clone %arg0 : memref<2xi32> to memref<2xi32>
   return %0 : memref<2xi32>
 }
diff --git a/mlir/test/Dialect/Vector/side-effects.mlir b/mlir/test/Dialect/Vector/side-effects.mlir
new file mode 100644
index 0000000000000..a1c205c2e3228
--- /dev/null
+++ b/mlir/test/Dialect/Vector/side-effects.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s --test-side-effects --verify-diagnostics
+
+func.func @test_side_effects(%arg0: memref<8xf32>) {
+  // expected-remark @below {{operation has no memory effects}}
+  %c0 = arith.constant 0 : index
+  // expected-remark @below {{operation has no memory effects}}
+  %c4 = arith.constant 4 : index
+  // expected-remark @below {{operation has no memory effects}}
+  %cst = arith.constant 0.0 : f32
+  // expected-remark @below {{operation has no memory effects}}
+  %vec_cst = arith.constant dense<0.0> : vector<4xf32>
+  // expected-remark @below {{found an instance of 'read' on op operand 0, on resource '<Default>'}}
+  %0 = vector.transfer_read %arg0[%c0], %cst : memref<8xf32>, vector<4xf32>
+  // expected-remark @below {{found an instance of 'write' on op operand 1, on resource '<Default>'}}
+  vector.transfer_write %vec_cst, %arg0[%c4] : vector<4xf32>, memref<8xf32>
+  return
+}
diff --git a/mlir/test/IR/test-side-effects.mlir b/mlir/test/IR/test-side-effects.mlir
index efce4856041a1..b652ecb7dad1d 100644
--- a/mlir/test/IR/test-side-effects.mlir
+++ b/mlir/test/IR/test-side-effects.mlir
@@ -15,7 +15,7 @@ func.func @side_effect(%arg : index) {
     {effect="write", test_resource}
   ]} : () -> i32
   
-  // expected-remark at +1 {{found an instance of 'allocate' on a op result, on resource '<Test>'}}
+  // expected-remark at +1 {{found an instance of 'allocate' on op result 0, on resource '<Test>'}}
   %3 = "test.side_effect_op"() {effects = [
     {effect="allocate", on_result, test_resource}
   ]} : () -> i32
@@ -38,19 +38,19 @@ func.func @side_effect(%arg : index) {
     effect_parameter = affine_map<(i, j) -> (j, i)>
   } : () -> i32
 
-  // expected-remark at +1 {{found an instance of 'allocate' on a op operand, on resource '<Test>'}}
+  // expected-remark at +1 {{found an instance of 'allocate' on op operand 0, on resource '<Test>'}}
   %6 = test.side_effect_with_region_op (%arg) {
   ^bb0(%arg0 : index):
     test.region_yield %arg0 : index 
   } {effects = [ {effect="allocate", on_operand, test_resource} ]} : index -> index
 
-  // expected-remark at +1 {{found an instance of 'allocate' on a op result, on resource '<Test>'}}
+  // expected-remark at +1 {{found an instance of 'allocate' on op result 0, on resource '<Test>'}}
   %7 = test.side_effect_with_region_op (%arg) {
   ^bb0(%arg0 : index):
     test.region_yield %arg0 : index 
   } {effects = [ {effect="allocate", on_result, test_resource} ]} : index -> index
 
-  // expected-remark at +1 {{found an instance of 'allocate' on a block argument, on resource '<Test>'}}
+  // expected-remark at +1 {{found an instance of 'allocate' on block argument 0, on resource '<Test>'}}
   %8 = test.side_effect_with_region_op (%arg) {
   ^bb0(%arg0 : index):
     test.region_yield %arg0 : index 
diff --git a/mlir/test/lib/IR/TestSideEffects.cpp b/mlir/test/lib/IR/TestSideEffects.cpp
index 7e01509d55685..000e7c204fd5f 100644
--- a/mlir/test/lib/IR/TestSideEffects.cpp
+++ b/mlir/test/lib/IR/TestSideEffects.cpp
@@ -52,12 +52,12 @@ struct SideEffectsPass
           diag << "'write'";
 
         if (instance.getValue()) {
-          if (instance.getEffectValue<OpOperand *>())
-            diag << " on a op operand,";
-          else if (instance.getEffectValue<OpResult>())
-            diag << " on a op result,";
-          else if (instance.getEffectValue<BlockArgument>())
-            diag << " on a block argument,";
+          if (auto *opOpd = instance.getEffectValue<OpOperand *>())
+            diag << " on op operand " << opOpd->getOperandNumber() << ",";
+          else if (auto opRes = instance.getEffectValue<OpResult>())
+            diag << " on op result " << opRes.getResultNumber() << ",";
+          else if (auto opBlk = instance.getEffectValue<BlockArgument>())
+            diag << " on block argument " << opBlk.getArgNumber() << ",";
         } else if (SymbolRefAttr symbolRef = instance.getSymbolRef())
           diag << " on a symbol '" << symbolRef << "',";