[Mlir-commits] [mlir] [mlir][vector] Keep store attributes when sinking a broadcast into vector.store (PR #221383)

[Hung-Kuan Tseng]

https://github.com/Tim096 created https://github.com/llvm/llvm-project/pull/221383

`StoreOpFromBroadcast` rebuilds a one-element `vector.store` of a
broadcast as a `vector.store` or `memref.store` of the broadcast source,
at the same base and indices, but does not forward `nontemporal` or
`alignment`:

```mlir
%0 = vector.broadcast %v : f32 to vector<1xf32>
vector.store %0, %mem[%i] alignment = 16 nontemporal = true : memref<?xf32>, vector<1xf32>
// sink patterns today
memref.store %v, %mem[%i] : memref<?xf32>
```

The rebuilt store writes the same address, so both attributes still
hold. This passes them through the builders. Same shape of fix as
#221319 for `vector.linearize`.


>From 7f333326e7b07ab2c6799aa4998c8ce5bca28026 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Sat, 5 Sep 2026 11:18:32 +0800
Subject: [PATCH] [mlir][vector] Keep store attributes when sinking a broadcast
 into vector.store

`StoreOpFromBroadcast` rebuilds a one-element `vector.store` of a
broadcast as a `vector.store` or `memref.store` of the broadcast source,
at the same base and indices, but does not forward `nontemporal` or
`alignment`. The rebuilt store writes the same address, so both still
hold. Pass them through the builders.
---
 .../Vector/Transforms/VectorTransforms.cpp     |  6 ++++--
 mlir/test/Dialect/Vector/vector-sink.mlir      | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 123caf2c11b83..04ea7fbbe441b 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1351,9 +1351,11 @@ class StoreOpFromBroadcast final : public OpRewritePattern<vector::StoreOp> {
     ValueRange indices = op.getIndices();
 
     if (isa<VectorType>(source.getType())) {
-      rewriter.replaceOpWithNewOp<vector::StoreOp>(op, source, base, indices);
+      rewriter.replaceOpWithNewOp<vector::StoreOp>(
+          op, source, base, indices, op.getNontemporal(), op.getMaybeAlign());
     } else {
-      rewriter.replaceOpWithNewOp<memref::StoreOp>(op, source, base, indices);
+      rewriter.replaceOpWithNewOp<memref::StoreOp>(
+          op, source, base, indices, op.getNontemporal(), op.getMaybeAlign());
     }
     rewriter.eraseOp(broadcast);
     return success();
diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index ac1e4fa1ea1cf..466befa908618 100644
--- a/mlir/test/Dialect/Vector/vector-sink.mlir
+++ b/mlir/test/Dialect/Vector/vector-sink.mlir
@@ -931,6 +931,24 @@ func.func @store_broadcast_1d_to_2d(%arg0: memref<?x?xf32>, %arg1: index, %arg2:
   return
 }
 
+// CHECK-LABEL: @store_broadcast_attrs
+//  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: f32)
+func.func @store_broadcast_attrs(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) {
+// CHECK:   memref.store %[[ARG2]], %[[ARG0]][%[[ARG1]]] alignment(16) nontemporal(true) : memref<?xf32>
+  %0 = vector.broadcast %arg2 : f32 to vector<1xf32>
+  vector.store %0, %arg0[%arg1] alignment = 16 nontemporal = true : memref<?xf32>, vector<1xf32>
+  return
+}
+
+// CHECK-LABEL: @store_broadcast_1d_to_2d_attrs
+//  CHECK-SAME:   (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: vector<1xf32>)
+func.func @store_broadcast_1d_to_2d_attrs(%arg0: memref<?x?xf32>, %arg1: index, %arg2: index, %arg3: vector<1xf32>) {
+// CHECK:   vector.store %[[ARG3]], %[[ARG0]][%[[ARG1]], %[[ARG2]]] alignment = 16 nontemporal = true : memref<?x?xf32>, vector<1xf32>
+  %0 = vector.broadcast %arg3 : vector<1xf32> to vector<1x1xf32>
+  vector.store %0, %arg0[%arg1, %arg2] alignment = 16 nontemporal = true : memref<?x?xf32>, vector<1x1xf32>
+  return
+}
+
 // CHECK-LABEL: @negative_store_scalable
 //  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: f32)
 func.func @negative_store_scalable(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) {