[Mlir-commits] [mlir] [mlir][vector] Keep load/store attributes when linearizing (PR #221319)

[Hung-Kuan Tseng]

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

`LinearizeVectorLoad` and `LinearizeVectorStore` rebuild the op with the
flattened vector type but leave out `alignment` and `nontemporal`:

```mlir
%0 = vector.load %m[%c0, %c0] alignment = 16 nontemporal = true : memref<2x8xf32>, vector<1x4xf32>
// -test-vector-linearize today
%0 = vector.load %m[%c0, %c0] : memref<2x8xf32>, vector<4xf32>
```

Only the vector type changes; the base, the indices and the bytes accessed
are the same, so this forwards both attributes through the builder.

Tests: one load and one store case with both attributes in `linearize.mlir`.


>From 2bb26b72e617e182caa545e2d8f185c501aa11e7 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Sat, 5 Sep 2026 02:50:02 +0800
Subject: [PATCH] [mlir][vector] Keep load/store attributes when linearizing

`LinearizeVectorLoad` and `LinearizeVectorStore` rebuild the op with the
flattened vector type but leave out `alignment` and `nontemporal`:

  %0 = vector.load %m[%c0, %c0] alignment = 16 nontemporal = true : memref<2x8xf32>, vector<1x4xf32>
  // -test-vector-linearize today
  %0 = vector.load %m[%c0, %c0] : memref<2x8xf32>, vector<4xf32>

Only the vector type changes; the base, the indices and the bytes
accessed are the same. Forward both attributes.
---
 .../Vector/Transforms/VectorLinearize.cpp     |  9 ++++----
 mlir/test/Dialect/Vector/linearize.mlir       | 23 +++++++++++++++++++
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index e6c28036ea1c5..c2bf3f4755579 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -684,9 +684,9 @@ struct LinearizeVectorLoad final : public OpConversionPattern<vector::LoadOp> {
 
     auto linearTy = typeConverter->convertType<VectorType>(vecTy);
 
-    auto newLoad =
-        vector::LoadOp::create(rewriter, loadOp.getLoc(), linearTy,
-                               adaptor.getBase(), adaptor.getIndices());
+    auto newLoad = vector::LoadOp::create(
+        rewriter, loadOp.getLoc(), linearTy, adaptor.getBase(),
+        adaptor.getIndices(), loadOp.getNontemporal(), loadOp.getMaybeAlign());
     rewriter.replaceOp(loadOp, newLoad.getResult());
     return success();
   }
@@ -731,7 +731,8 @@ struct LinearizeVectorStore final
 
     rewriter.replaceOpWithNewOp<vector::StoreOp>(
         storeOp, adaptor.getValueToStore(), adaptor.getBase(),
-        adaptor.getIndices());
+        adaptor.getIndices(), storeOp.getNontemporal(),
+        storeOp.getMaybeAlign());
     return success();
   }
 };
diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index 52a5bc5879a33..c5c01c10ca2c0 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -516,6 +516,29 @@ func.func @linearize_vector_store(%arg0: memref<2x8xf32>, %arg1: vector<1x4xf32>
   return
 }
 
+// CHECK-LABEL: linearize_vector_load_attrs
+// CHECK-SAME: (%[[ARG0:.*]]: memref<2x8xf32>) -> vector<1x4xf32>
+func.func @linearize_vector_load_attrs(%arg0: memref<2x8xf32>) -> vector<1x4xf32> {
+  // CHECK: %[[CST0:.*]] = arith.constant 0 : index
+  // CHECK: %[[LOAD:.*]] = vector.load %[[ARG0]][%[[CST0]], %[[CST0]]] alignment = 16 nontemporal = true : memref<2x8xf32>, vector<4xf32>
+  // CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] : vector<4xf32> to vector<1x4xf32>
+  // CHECK: return %[[CAST]] : vector<1x4xf32>
+  %c0 = arith.constant 0 : index
+  %0 = vector.load %arg0[%c0, %c0] alignment = 16 nontemporal = true : memref<2x8xf32>, vector<1x4xf32>
+  return %0 : vector<1x4xf32>
+}
+
+// CHECK-LABEL: linearize_vector_store_attrs
+// CHECK-SAME: (%[[ARG0:.*]]: memref<2x8xf32>, %[[ARG1:.*]]: vector<1x4xf32>)
+func.func @linearize_vector_store_attrs(%arg0: memref<2x8xf32>, %arg1: vector<1x4xf32>) {
+  // CHECK: %[[CAST:.*]] = vector.shape_cast %arg1 : vector<1x4xf32> to vector<4xf32>
+  // CHECK: %[[CST0:.*]] = arith.constant 0 : index
+  // CHECK: vector.store %[[CAST]], %[[ARG0]][%[[CST0]], %[[CST0]]] alignment = 16 nontemporal = true : memref<2x8xf32>, vector<4xf32>
+  %c0 = arith.constant 0 : index
+  vector.store %arg1, %arg0[%c0, %c0] alignment = 16 nontemporal = true : memref<2x8xf32>, vector<1x4xf32>
+  return
+}
+
 // CHECK-LABEL: linearize_vector_load_scalable
 // CHECK-SAME: (%[[ARG0:.*]]: memref<2x8xf32>) -> vector<1x[4]xf32>
 func.func @linearize_vector_load_scalable(%arg0: memref<2x8xf32>) -> vector<1x[4]xf32> {