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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-vector

Author: Hung-Kuan Tseng (Tim096)

<details>
<summary>Changes</summary>

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


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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp (+5-4) 
- (modified) mlir/test/Dialect/Vector/linearize.mlir (+23) 


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

``````````

</details>


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