[Mlir-commits] [mlir] [mlir][vector] Keep alignment when folding masked and contiguous memory ops (PR #221317)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Sep 4 11:49:09 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Hung-Kuan Tseng (Tim096)
<details>
<summary>Changes</summary>
The all-true mask folders for `vector.maskedload`, `maskedstore`,
`expandload` and `compressstore`, and the contiguous-indices folders for
`vector.gather` and `scatter`, rebuild the op without its `alignment`:
```mlir
%0 = vector.maskedload %base[%c0], %mask, %pt alignment = 64 : memref<16xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
// -canonicalize today
%0 = vector.load %base[%c0] : memref<16xf32>, vector<16xf32>
```
The folded op touches the same address as the original: the mask folders
keep the base and indices, and the contiguous folders keep the base and
offsets, which is where the contiguous run starts. This forwards the
attribute through the builders. The folders predate the `alignment`
attribute on these ops (#<!-- -->144344, #<!-- -->151690), which is why it was never
copied.
Tests: one `alignment = 64` case per folder in `vector-mem-transforms.mlir`
and `canonicalize.mlir`.
---
Full diff: https://github.com/llvm/llvm-project/pull/221317.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+13-8)
- (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+31)
- (modified) mlir/test/Dialect/Vector/vector-mem-transforms.mlir (+54)
``````````diff
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index f8f3deb2e4789..9456dcebaa903 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6356,7 +6356,8 @@ class MaskedLoadFolder final : public OpRewritePattern<MaskedLoadOp> {
switch (getMaskFormat(load.getMask())) {
case MaskFormat::AllTrue:
rewriter.replaceOpWithNewOp<vector::LoadOp>(
- load, load.getType(), load.getBase(), load.getIndices());
+ load, load.getType(), load.getBase(), load.getIndices(),
+ /*nontemporal=*/false, load.getMaybeAlign());
return success();
case MaskFormat::AllFalse:
rewriter.replaceOp(load, load.getPassThru());
@@ -6423,7 +6424,8 @@ class MaskedStoreFolder final : public OpRewritePattern<MaskedStoreOp> {
switch (getMaskFormat(store.getMask())) {
case MaskFormat::AllTrue:
rewriter.replaceOpWithNewOp<vector::StoreOp>(
- store, store.getValueToStore(), store.getBase(), store.getIndices());
+ store, store.getValueToStore(), store.getBase(), store.getIndices(),
+ /*nontemporal=*/false, store.getMaybeAlign());
return success();
case MaskFormat::AllFalse:
rewriter.eraseOp(store);
@@ -6554,9 +6556,9 @@ class FoldContiguousGather final : public OpRewritePattern<GatherOp> {
if (failed(isZeroBasedContiguousSeq(op.getIndices())))
return failure();
- rewriter.replaceOpWithNewOp<MaskedLoadOp>(op, op.getType(), op.getBase(),
- op.getOffsets(), op.getMask(),
- op.getPassThru());
+ rewriter.replaceOpWithNewOp<MaskedLoadOp>(
+ op, op.getType(), op.getBase(), op.getOffsets(), op.getMask(),
+ op.getPassThru(), op.getMaybeAlign());
return success();
}
};
@@ -6654,7 +6656,8 @@ class FoldContiguousScatter final : public OpRewritePattern<ScatterOp> {
return failure();
rewriter.replaceOpWithNewOp<MaskedStoreOp>(
- op, op.getBase(), op.getOffsets(), op.getMask(), op.getValueToStore());
+ op, op.getBase(), op.getOffsets(), op.getMask(), op.getValueToStore(),
+ op.getMaybeAlign());
return success();
}
};
@@ -6714,7 +6717,8 @@ class ExpandLoadFolder final : public OpRewritePattern<ExpandLoadOp> {
switch (getMaskFormat(expand.getMask())) {
case MaskFormat::AllTrue:
rewriter.replaceOpWithNewOp<vector::LoadOp>(
- expand, expand.getType(), expand.getBase(), expand.getIndices());
+ expand, expand.getType(), expand.getBase(), expand.getIndices(),
+ /*nontemporal=*/false, expand.getMaybeAlign());
return success();
case MaskFormat::AllFalse:
rewriter.replaceOp(expand, expand.getPassThru());
@@ -6779,7 +6783,8 @@ class CompressStoreFolder final : public OpRewritePattern<CompressStoreOp> {
case MaskFormat::AllTrue:
rewriter.replaceOpWithNewOp<vector::StoreOp>(
compress, compress.getValueToStore(), compress.getBase(),
- compress.getIndices());
+ compress.getIndices(), /*nontemporal=*/false,
+ compress.getMaybeAlign());
return success();
case MaskFormat::AllFalse:
rewriter.eraseOp(compress);
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index f6ba206e7449e..8af99cb771533 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -4126,6 +4126,22 @@ func.func @contiguous_gather(%base: memref<?xf32>,
// -----
+// CHECK-LABEL: @contiguous_gather_alignment
+// CHECK-SAME: (%[[BASE:.*]]: memref<?xf32>, %[[MASK:.*]]: vector<16xi1>, %[[PASSTHRU:.*]]: vector<16xf32>)
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[R:.*]] = vector.maskedload %[[BASE]][%[[C0]]], %[[MASK]], %[[PASSTHRU]] alignment = 64 : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+// CHECK: return %[[R]]
+func.func @contiguous_gather_alignment(%base: memref<?xf32>,
+ %mask: vector<16xi1>, %passthru: vector<16xf32>) -> vector<16xf32> {
+ %c0 = arith.constant 0 : index
+ %indices = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : vector<16xi32>
+ %1 = vector.gather %base[%c0][%indices], %mask, %passthru alignment = 64 :
+ memref<?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+ return %1 : vector<16xf32>
+}
+
+// -----
+
// CHECK-LABEL: @contiguous_gather_non_zero_start(
// TODO: Non-zero start is not supported yet.
// CHECK: %[[R:.*]] = vector.gather
@@ -4233,6 +4249,21 @@ func.func @contiguous_scatter(%base: memref<?xf32>,
// -----
+// CHECK-LABEL: @contiguous_scatter_alignment
+// CHECK-SAME: (%[[BASE:.*]]: memref<?xf32>, %[[MASK:.*]]: vector<16xi1>, %[[VALUE:.*]]: vector<16xf32>)
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: vector.maskedstore %[[BASE]][%[[C0]]], %[[MASK]], %[[VALUE]] alignment = 64 : memref<?xf32>, vector<16xi1>, vector<16xf32>
+func.func @contiguous_scatter_alignment(%base: memref<?xf32>,
+ %mask: vector<16xi1>, %value: vector<16xf32>) {
+ %c0 = arith.constant 0 : index
+ %indices = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : vector<16xi32>
+ vector.scatter %base[%c0][%indices], %mask, %value alignment = 64 :
+ memref<?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @contiguous_scatter_const_mask
// CHECK-SAME: (%[[BASE:.*]]: memref<?xf32>, %[[VALUE:.*]]: vector<16xf32>)
// CHECK: %[[C0:.*]] = arith.constant 0 : index
diff --git a/mlir/test/Dialect/Vector/vector-mem-transforms.mlir b/mlir/test/Dialect/Vector/vector-mem-transforms.mlir
index cd7ea84cdb2eb..79bc2f77e7a22 100644
--- a/mlir/test/Dialect/Vector/vector-mem-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-mem-transforms.mlir
@@ -32,6 +32,20 @@ func.func @fold_maskedload_all_true_static(%base: memref<16xf32>, %pass_thru: ve
return %ld : vector<16xf32>
}
+// CHECK-LABEL: func @fold_maskedload_all_true_alignment(
+// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
+// CHECK-SAME: %[[PASS_THRU:.*]]: vector<16xf32>) -> vector<16xf32> {
+// CHECK-DAG: %[[IDX:.*]] = arith.constant 0 : index
+// CHECK-NEXT: %[[LOAD:.*]] = vector.load %[[BASE]][%[[IDX]]] alignment = 64 : memref<16xf32>, vector<16xf32>
+// CHECK-NEXT: return %[[LOAD]] : vector<16xf32>
+func.func @fold_maskedload_all_true_alignment(%base: memref<16xf32>, %pass_thru: vector<16xf32>) -> vector<16xf32> {
+ %c0 = arith.constant 0 : index
+ %mask = vector.constant_mask [16] : vector<16xi1>
+ %ld = vector.maskedload %base[%c0], %mask, %pass_thru alignment = 64
+ : memref<16xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+ return %ld : vector<16xf32>
+}
+
// CHECK-LABEL: func @fold_maskedload_all_false_static(
// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
// CHECK-SAME: %[[PASS_THRU:.*]]: vector<16xf32>) -> vector<16xf32> {
@@ -75,6 +89,19 @@ func.func @fold_maskedstore_all_true(%base: memref<16xf32>, %value: vector<16xf3
return
}
+// CHECK-LABEL: func @fold_maskedstore_all_true_alignment(
+// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
+// CHECK-SAME: %[[VALUE:.*]]: vector<16xf32>) {
+// CHECK-NEXT: %[[IDX:.*]] = arith.constant 0 : index
+// CHECK-NEXT: vector.store %[[VALUE]], %[[BASE]][%[[IDX]]] alignment = 64 : memref<16xf32>, vector<16xf32>
+// CHECK-NEXT: return
+func.func @fold_maskedstore_all_true_alignment(%base: memref<16xf32>, %value: vector<16xf32>) {
+ %c0 = arith.constant 0 : index
+ %mask = vector.constant_mask [16] : vector<16xi1>
+ vector.maskedstore %base[%c0], %mask, %value alignment = 64 : memref<16xf32>, vector<16xi1>, vector<16xf32>
+ return
+}
+
// CHECK-LABEL: func @fold_maskedstore_all_false(
// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
// CHECK-SAME: %[[VALUE:.*]]: vector<16xf32>) {
@@ -175,6 +202,20 @@ func.func @fold_expandload_all_true(%base: memref<16xf32>, %pass_thru: vector<16
return %ld : vector<16xf32>
}
+// CHECK-LABEL: func @fold_expandload_all_true_alignment(
+// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
+// CHECK-SAME: %[[PASS_THRU:.*]]: vector<16xf32>) -> vector<16xf32> {
+// CHECK-DAG: %[[C:.*]] = arith.constant 0 : index
+// CHECK-NEXT: %[[T:.*]] = vector.load %[[BASE]][%[[C]]] alignment = 64 : memref<16xf32>, vector<16xf32>
+// CHECK-NEXT: return %[[T]] : vector<16xf32>
+func.func @fold_expandload_all_true_alignment(%base: memref<16xf32>, %pass_thru: vector<16xf32>) -> vector<16xf32> {
+ %c0 = arith.constant 0 : index
+ %mask = vector.constant_mask [16] : vector<16xi1>
+ %ld = vector.expandload %base[%c0], %mask, %pass_thru alignment = 64
+ : memref<16xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+ return %ld : vector<16xf32>
+}
+
// CHECK-LABEL: func @fold_expandload_all_true_scalable(
// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
// CHECK-SAME: %[[PASS_THRU:.*]]: vector<[16]xf32>) -> vector<[16]xf32> {
@@ -230,6 +271,19 @@ func.func @fold_compressstore_all_true(%base: memref<16xf32>, %value: vector<16x
return
}
+// CHECK-LABEL: func @fold_compressstore_all_true_alignment(
+// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
+// CHECK-SAME: %[[VALUE:.*]]: vector<16xf32>) {
+// CHECK-NEXT: %[[C:.*]] = arith.constant 0 : index
+// CHECK-NEXT: vector.store %[[VALUE]], %[[BASE]][%[[C]]] alignment = 64 : memref<16xf32>, vector<16xf32>
+// CHECK-NEXT: return
+func.func @fold_compressstore_all_true_alignment(%base: memref<16xf32>, %value: vector<16xf32>) {
+ %c0 = arith.constant 0 : index
+ %mask = vector.constant_mask [16] : vector<16xi1>
+ vector.compressstore %base[%c0], %mask, %value alignment = 64 : memref<16xf32>, vector<16xi1>, vector<16xf32>
+ return
+}
+
// CHECK-LABEL: func @fold_compressstore_all_true_scalable(
// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
// CHECK-SAME: %[[VALUE:.*]]: vector<[16]xf32>) {
``````````
</details>
https://github.com/llvm/llvm-project/pull/221317
More information about the Mlir-commits
mailing list