[Mlir-commits] [mlir] [mlir][vector] Keep alignment when folding masked and contiguous memory ops (PR #221317)
Hung-Kuan Tseng
llvmlistbot at llvm.org
Fri Sep 4 11:48:30 PDT 2026
https://github.com/Tim096 created https://github.com/llvm/llvm-project/pull/221317
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`.
>From 1b2aca62fe7ef1e5d0f0fc8aa369c4f02e2b6c97 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Sat, 5 Sep 2026 02:39:48 +0800
Subject: [PATCH] [mlir][vector] Keep alignment when folding masked and
contiguous memory ops
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`:
%0 = vector.maskedload %base[%c0], %mask, %pt alignment = 64 : ...
// -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 gather starts. Forward the
attribute through the builders.
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 21 +++++---
mlir/test/Dialect/Vector/canonicalize.mlir | 31 +++++++++++
.../Dialect/Vector/vector-mem-transforms.mlir | 54 +++++++++++++++++++
3 files changed, 98 insertions(+), 8 deletions(-)
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>) {
More information about the Mlir-commits
mailing list