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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hung-Kuan Tseng (Tim096)

<details>
<summary>Changes</summary>

`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`.


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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp (+4-2) 
- (modified) mlir/test/Dialect/Vector/vector-sink.mlir (+18) 


``````````diff
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) {

``````````

</details>


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