[Mlir-commits] [mlir] [mlir][vector] Recompute load/store alignment when the index moves (PR #221385)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Sep 4 20:44:42 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Hung-Kuan Tseng (Tim096)
<details>
<summary>Changes</summary>
`ExtractOpFromLoad` rebuilds a `vector.load` as a narrower load at the
extract position, and `UnrollLoadPattern`/`UnrollStorePattern` rebuild a
`vector.load`/`vector.store` per tile at the tile offset. All of them
drop `nontemporal` and `alignment`:
```mlir
%0 = vector.load %mem[%c0, %c0] alignment = 16 nontemporal = true : memref<4x4xf16>, vector<4x4xf16>
// unrolled to 2x2 today
%1 = vector.load %mem[%c0, %c0] : memref<4x4xf16>, vector<2x2xf16>
%2 = vector.load %mem[%c0, %c2] : memref<4x4xf16>, vector<2x2xf16>
...
```
`nontemporal` still describes the access. The alignment does not carry
over as is: the new access is at a byte offset from the old one, and
only the alignment that both the old access and the offset satisfy holds
for it (`[%c0, %c2]` above is 4 bytes in, so it is 4-byte aligned, while
`[%c2, %c0]` is 16 bytes in and keeps the 16).
This adds `vector::getAlignmentAfterOffset`, which turns the element
offsets on the trailing dimensions into a byte offset through the static
strides of the memref and returns `llvm::commonAlignment` of the two, or
nothing when the strides or the offset are not static (a dynamic extract
position, a `memref<?x?xf32>` row stride, an `index` element type), and
uses it in the three patterns, forwarding `nontemporal` next to it.
Tests: for the sink pattern, a zero offset, a static offset in the
innermost dimension, a dynamic offset, and a static and a dynamic row
stride; for unrolling, a 2-D load and store where the tiles alternate
between the original and the reduced alignment.
The forwarding-only cases, where the address does not move, are
#<!-- -->221317, #<!-- -->221319, #<!-- -->221382 and #<!-- -->221383.
---
Full diff: https://github.com/llvm/llvm-project/pull/221385.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h (+10)
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp (+18-3)
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+10-3)
- (modified) mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp (+34)
- (modified) mlir/test/Dialect/Vector/vector-sink.mlir (+61)
- (modified) mlir/test/Dialect/Vector/vector-unroll-options.mlir (+35)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
index edfbbf4ba6da0..d8c9950500644 100644
--- a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
+++ b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
@@ -20,6 +20,7 @@
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Alignment.h"
namespace mlir {
@@ -217,6 +218,15 @@ struct MaskableOpRewritePattern : OpRewritePattern<SourceOp> {
/// are not linearizable.
bool isLinearizableVector(VectorType type);
+/// Returns the alignment of an access to `memRefType` at `indices + offsets`,
+/// given that the access at `indices` is aligned to `alignment`. `offsets`
+/// apply to the trailing dimensions of the memref. Returns an empty alignment
+/// when `alignment` is empty, or when the byte distance between the two
+/// accesses is not static.
+llvm::MaybeAlign getAlignmentAfterOffset(llvm::MaybeAlign alignment,
+ MemRefType memRefType,
+ ArrayRef<int64_t> offsets);
+
/// Creates a TransferReadOp from `source`.
///
/// If the shape of vector to read differs from the shape of the value being
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 123caf2c11b83..df449163aed86 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1284,21 +1284,36 @@ class ExtractOpFromLoad final : public OpRewritePattern<vector::ExtractOp> {
rewriter.setInsertionPoint(loadOp);
Location loc = loadOp.getLoc();
ArithIndexingBuilder idxBuilderf(rewriter, loc);
+ SmallVector<int64_t> offsets(indices.size(), 0);
+ bool hasDynamicOffset = false;
for (auto i : llvm::seq<int64_t>(rankOffset, indices.size() - finalRank)) {
OpFoldResult pos = extractPos[i - rankOffset];
if (isZeroInteger(pos))
continue;
+ if (std::optional<int64_t> staticPos = getConstantIntValue(pos))
+ offsets[i] = *staticPos;
+ else
+ hasDynamicOffset = true;
Value offset = getValueOrCreateConstantIndexOp(rewriter, loc, pos);
indices[i] = idxBuilderf.add(indices[i], offset);
}
+ // The new load is at an offset from the original one, so its alignment
+ // has to be recomputed, which needs a static offset.
+ llvm::MaybeAlign alignment;
+ if (!hasDynamicOffset)
+ alignment =
+ getAlignmentAfterOffset(loadOp.getMaybeAlign(), memType, offsets);
+ bool nontemporal = loadOp.getNontemporal();
+
Value base = loadOp.getBase();
if (extractVecType) {
- rewriter.replaceOpWithNewOp<vector::LoadOp>(op, extractVecType, base,
- indices);
+ rewriter.replaceOpWithNewOp<vector::LoadOp>(
+ op, extractVecType, base, indices, nontemporal, alignment);
} else {
- rewriter.replaceOpWithNewOp<memref::LoadOp>(op, base, indices);
+ rewriter.replaceOpWithNewOp<memref::LoadOp>(op, base, indices,
+ nontemporal, alignment);
}
// We checked for single use so we can safely erase the load op.
rewriter.eraseOp(loadOp);
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index a95bfba0814bd..edf2e8ec59686 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -14,6 +14,7 @@
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/Interfaces/VectorInterfaces.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
@@ -737,8 +738,11 @@ struct UnrollLoadPattern : public OpRewritePattern<vector::LoadOp> {
StaticTileOffsetRange(originalShape, *targetShape, loopOrder)) {
SmallVector<Value> indices =
sliceLoadStoreIndices(rewriter, loc, loadOp.getIndices(), offsets);
- Value slicedLoad = vector::LoadOp::create(rewriter, loc, targetVecType,
- loadOp.getBase(), indices);
+ Value slicedLoad = vector::LoadOp::create(
+ rewriter, loc, targetVecType, loadOp.getBase(), indices,
+ loadOp.getNontemporal(),
+ getAlignmentAfterOffset(loadOp.getMaybeAlign(),
+ loadOp.getMemRefType(), offsets));
result = rewriter.createOrFold<vector::InsertStridedSliceOp>(
loc, slicedLoad, result, offsets, strides);
}
@@ -780,7 +784,10 @@ struct UnrollStorePattern : public OpRewritePattern<vector::StoreOp> {
sliceLoadStoreIndices(rewriter, loc, storeOp.getIndices(), offsets);
Value slice = rewriter.createOrFold<vector::ExtractStridedSliceOp>(
loc, vector, offsets, *targetShape, strides);
- vector::StoreOp::create(rewriter, loc, slice, base, indices);
+ vector::StoreOp::create(
+ rewriter, loc, slice, base, indices, storeOp.getNontemporal(),
+ getAlignmentAfterOffset(storeOp.getMaybeAlign(),
+ storeOp.getMemRefType(), offsets));
}
rewriter.eraseOp(storeOp);
return success();
diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index c38213850c6ea..b57589c5f8150 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -29,6 +29,7 @@
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseSet.h"
+#include "llvm/Support/Alignment.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/InterleavedRange.h"
@@ -307,6 +308,39 @@ SmallVector<OpFoldResult> vector::getMixedSizesXfer(bool hasTensorSemantics,
return mixedSourceDims;
}
+llvm::MaybeAlign vector::getAlignmentAfterOffset(llvm::MaybeAlign alignment,
+ MemRefType memRefType,
+ ArrayRef<int64_t> offsets) {
+ if (!alignment)
+ return {};
+ assert(offsets.size() <= static_cast<size_t>(memRefType.getRank()) &&
+ "more offsets than memref dimensions");
+ if (llvm::all_of(offsets, [](int64_t offset) { return offset == 0; }))
+ return alignment;
+
+ Type elementType = memRefType.getElementType();
+ if (!elementType.isIntOrFloat() ||
+ elementType.getIntOrFloatBitWidth() % 8 != 0)
+ return {};
+ SmallVector<int64_t> strides;
+ int64_t memRefOffset;
+ if (failed(memRefType.getStridesAndOffset(strides, memRefOffset)))
+ return {};
+
+ int64_t elementOffset = 0;
+ size_t start = strides.size() - offsets.size();
+ for (auto [i, offset] : llvm::enumerate(offsets)) {
+ if (offset == 0)
+ continue;
+ if (ShapedType::isDynamic(strides[start + i]))
+ return {};
+ elementOffset += offset * strides[start + i];
+ }
+ uint64_t byteOffset =
+ std::abs(elementOffset) * (elementType.getIntOrFloatBitWidth() / 8);
+ return llvm::commonAlignment(*alignment, byteOffset);
+}
+
bool vector::isLinearizableVector(VectorType type) {
return (type.getRank() > 1) && (type.getNumScalableDims() <= 1);
}
diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index ac1e4fa1ea1cf..29777b16434b5 100644
--- a/mlir/test/Dialect/Vector/vector-sink.mlir
+++ b/mlir/test/Dialect/Vector/vector-sink.mlir
@@ -900,6 +900,67 @@ func.func @negative_extract_load_scalable(%arg0: memref<?xf32>, %arg1: index) ->
return %1 : f32
}
+// CHECK-LABEL: @extract_load_scalar_attrs
+// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)
+func.func @extract_load_scalar_attrs(%arg0: memref<?xf32>, %arg1: index) -> f32 {
+// CHECK: %[[RES:.*]] = memref.load %[[ARG0]][%[[ARG1]]] alignment(16) nontemporal(true) : memref<?xf32>
+// CHECK: return %[[RES]] : f32
+ %0 = vector.load %arg0[%arg1] alignment = 16 nontemporal = true : memref<?xf32>, vector<4xf32>
+ %1 = vector.extract %0[0] : f32 from vector<4xf32>
+ return %1 : f32
+}
+
+// The extracted element is 4 bytes past the 16-byte aligned load.
+// CHECK-LABEL: @extract_load_scalar_non_zero_off_alignment
+// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)
+func.func @extract_load_scalar_non_zero_off_alignment(%arg0: memref<?xf32>, %arg1: index) -> f32 {
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[OFF:.*]] = arith.addi %[[ARG1]], %[[C1]] overflow<nsw> : index
+// CHECK: %[[RES:.*]] = memref.load %[[ARG0]][%[[OFF]]] alignment(4) : memref<?xf32>
+// CHECK: return %[[RES]] : f32
+ %0 = vector.load %arg0[%arg1] alignment = 16 : memref<?xf32>, vector<4xf32>
+ %1 = vector.extract %0[1] : f32 from vector<4xf32>
+ return %1 : f32
+}
+
+// The offset is not static, so the alignment is dropped.
+// CHECK-LABEL: @extract_load_scalar_dyn_off_alignment
+// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
+func.func @extract_load_scalar_dyn_off_alignment(%arg0: memref<?xf32>, %arg1: index, %arg2: index) -> f32 {
+// CHECK: %[[OFF:.*]] = arith.addi %[[ARG1]], %[[ARG2]] overflow<nsw> : index
+// CHECK: %[[RES:.*]] = memref.load %[[ARG0]][%[[OFF]]] nontemporal(true) : memref<?xf32>
+// CHECK: return %[[RES]] : f32
+ %0 = vector.load %arg0[%arg1] alignment = 16 nontemporal = true : memref<?xf32>, vector<4xf32>
+ %1 = vector.extract %0[%arg2] : f32 from vector<4xf32>
+ return %1 : f32
+}
+
+// The extracted row is 32 bytes past the 16-byte aligned load.
+// CHECK-LABEL: @extract_load_vec_static_stride_alignment
+// CHECK-SAME: (%[[ARG0:.*]]: memref<8x8xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
+func.func @extract_load_vec_static_stride_alignment(%arg0: memref<8x8xf32>, %arg1: index, %arg2: index) -> vector<4xf32> {
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[OFF:.*]] = arith.addi %[[ARG1]], %[[C1]] overflow<nsw> : index
+// CHECK: %[[RES:.*]] = vector.load %[[ARG0]][%[[OFF]], %[[ARG2]]] alignment = 16 : memref<8x8xf32>, vector<4xf32>
+// CHECK: return %[[RES]] : vector<4xf32>
+ %0 = vector.load %arg0[%arg1, %arg2] alignment = 16 : memref<8x8xf32>, vector<2x4xf32>
+ %1 = vector.extract %0[1] : vector<4xf32> from vector<2x4xf32>
+ return %1 : vector<4xf32>
+}
+
+// The row stride is not static, so the alignment is dropped.
+// CHECK-LABEL: @extract_load_vec_dynamic_stride_alignment
+// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
+func.func @extract_load_vec_dynamic_stride_alignment(%arg0: memref<?x?xf32>, %arg1: index, %arg2: index) -> vector<4xf32> {
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[OFF:.*]] = arith.addi %[[ARG1]], %[[C1]] overflow<nsw> : index
+// CHECK: %[[RES:.*]] = vector.load %[[ARG0]][%[[OFF]], %[[ARG2]]] : memref<?x?xf32>, vector<4xf32>
+// CHECK: return %[[RES]] : vector<4xf32>
+ %0 = vector.load %arg0[%arg1, %arg2] alignment = 16 : memref<?x?xf32>, vector<2x4xf32>
+ %1 = vector.extract %0[1] : vector<4xf32> from vector<2x4xf32>
+ return %1 : vector<4xf32>
+}
+
//-----------------------------------------------------------------------------
// [Pattern: StoreOpFromBroadcast]
//-----------------------------------------------------------------------------
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 38bebdc24757c..13471bca5e5f7 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -467,6 +467,41 @@ func.func @vector_store_2D(%mem: memref<4x4xf16>, %v: vector<4x4xf16>) {
// -----
+// The tiles at offsets [0, 2] and [2, 2] are 4 bytes past a 16-byte boundary.
+func.func @vector_load_2D_alignment(%mem: memref<4x4xf16>) -> vector<4x4xf16> {
+ %c0 = arith.constant 0 : index
+ %0 = vector.load %mem[%c0, %c0] alignment = 16 nontemporal = true : memref<4x4xf16>, vector<4x4xf16>
+ return %0 : vector<4x4xf16>
+}
+
+// CHECK-LABEL: func.func @vector_load_2D_alignment(
+// CHECK-SAME: %[[ARG:.*]]: memref<4x4xf16>) -> vector<4x4xf16> {
+ // CHECK: %[[C2:.*]] = arith.constant 2 : index
+ // CHECK: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: vector.load %[[ARG]][%[[C0]], %[[C0]]] alignment = 16 nontemporal = true : memref<4x4xf16>, vector<2x2xf16>
+ // CHECK: vector.load %[[ARG]][%[[C0]], %[[C2]]] alignment = 4 nontemporal = true : memref<4x4xf16>, vector<2x2xf16>
+ // CHECK: vector.load %[[ARG]][%[[C2]], %[[C0]]] alignment = 16 nontemporal = true : memref<4x4xf16>, vector<2x2xf16>
+ // CHECK: vector.load %[[ARG]][%[[C2]], %[[C2]]] alignment = 4 nontemporal = true : memref<4x4xf16>, vector<2x2xf16>
+
+// -----
+
+func.func @vector_store_2D_alignment(%mem: memref<4x4xf16>, %v: vector<4x4xf16>) {
+ %c0 = arith.constant 0 : index
+ vector.store %v, %mem[%c0, %c0] alignment = 16 nontemporal = true : memref<4x4xf16>, vector<4x4xf16>
+ return
+}
+
+// CHECK-LABEL: func.func @vector_store_2D_alignment(
+// CHECK-SAME: %[[ARG0:.*]]: memref<4x4xf16>, %[[ARG1:.*]]: vector<4x4xf16>) {
+ // CHECK: %[[C2:.*]] = arith.constant 2 : index
+ // CHECK: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: vector.store %{{.*}}, %[[ARG0]][%[[C0]], %[[C0]]] alignment = 16 nontemporal = true : memref<4x4xf16>, vector<2x2xf16>
+ // CHECK: vector.store %{{.*}}, %[[ARG0]][%[[C0]], %[[C2]]] alignment = 4 nontemporal = true : memref<4x4xf16>, vector<2x2xf16>
+ // CHECK: vector.store %{{.*}}, %[[ARG0]][%[[C2]], %[[C0]]] alignment = 16 nontemporal = true : memref<4x4xf16>, vector<2x2xf16>
+ // CHECK: vector.store %{{.*}}, %[[ARG0]][%[[C2]], %[[C2]]] alignment = 4 nontemporal = true : memref<4x4xf16>, vector<2x2xf16>
+
+// -----
+
func.func @vector_step() -> vector<32xindex> {
%0 = vector.step : vector<32xindex>
return %0 : vector<32xindex>
``````````
</details>
https://github.com/llvm/llvm-project/pull/221385
More information about the Mlir-commits
mailing list