[Mlir-commits] [mlir] [mlir][mem2reg] Support whole-buffer vector.load/store promotion (PR #224182)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Sep 16 19:34:48 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: crZhao (cr-zhao)
<details>
<summary>Changes</summary>
This PR addresses the `vector.load` / `vector.store` limitation documented in #<!-- -->211880. Following that PR's implementation approach, it adds external `PromotableMemOpInterface` models for whole-buffer vector loads and stores, allowing eligible `memref.alloca` buffers to be promoted to vector SSA values.
Promotion requires an exact match with the slot's vector type and constant zero indices. The models work with mixed transfer/load-store accesses, existing static subview aliases, and supported scalable vector slots. Single-element scalar slots remain unsupported, as in the existing transfer models.
Tests cover straight-line, multidimensional, loop, conditional, mixed-access, subview, and scalable promotion, along with accesses that must block promotion.
Validation: C++ compilation checks and 13 test cases passed using a temporary dialect plugin with the existing `mlir-opt`; the existing transfer Mem2Reg test produced identical output with and without the plugin. Full source-build validation is left to CI.
---
Full diff: https://github.com/llvm/llvm-project/pull/224182.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp (+85-4)
- (added) mlir/test/Dialect/Vector/mem2reg-load-store.mlir (+229)
``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
index 1acebcb215439..9a7765a267879 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -12,9 +12,9 @@
// such a buffer via a subview). With these models, Mem2Reg replaces the memory
// slot with a vector value, threading it as the reaching definition:
//
-// * `vector.transfer_read` of the whole buffer becomes a use of the current
-// vector value; `vector.transfer_write` of the whole buffer becomes a new
-// definition of it (see the `PromotableMemOpInterface` models below).
+// * `vector.transfer_read` or `vector.load` of the whole buffer becomes a use
+// of the current vector value; `vector.transfer_write` or `vector.store`
+// becomes a new definition of it (see the models below).
//
// * a static, same-rank `memref.subview` is exposed as a promotable sub-slice
// alias of the buffer's slot (via `PromotableAliaserInterface`): a read of
@@ -25,7 +25,7 @@
// and overlapping sub-writes composing in program order.
//
// Accesses that are not whole-(sub-)buffer -- dynamic offsets, rank-reducing or
-// non-unit-stride subviews, masked or partial transfers, non-zero transfer
+// non-unit-stride subviews, masked or partial accesses, non-zero access
// indices -- are left untouched, so the buffer is not promoted.
//
//===----------------------------------------------------------------------===//
@@ -45,6 +45,24 @@ using namespace mlir::vector;
// Utilities
//===----------------------------------------------------------------------===//
+/// Load/store accesses must cover the slot from its origin with an exactly
+/// matching vector type. The allocation or alias model determines that type.
+template <typename LoadStoreOpTy>
+static bool
+isWholeBufferLoadStore(LoadStoreOpTy op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses) {
+ if (blockingUses.size() != 1 || (*blockingUses.begin())->get() != slot.ptr ||
+ op.getBase() != slot.ptr || op.getVectorType() != slot.elemType)
+ return false;
+
+ for (Value index : op.getIndices()) {
+ std::optional<int64_t> constIndex = getConstantIntValue(index);
+ if (!constIndex || *constIndex != 0)
+ return false;
+ }
+ return true;
+}
+
/// Returns whether `xferOp` accesses exactly the whole contents of `slot`, so
/// it can act as a plain whole-buffer load/store during Mem2Reg.
static bool
@@ -96,6 +114,67 @@ isWholeBufferTransfer(VectorTransferOpInterface xferOp, const MemorySlot &slot,
namespace {
+struct LoadOpMemOpModel
+ : public PromotableMemOpInterface::ExternalModel<LoadOpMemOpModel,
+ vector::LoadOp> {
+ bool loadsFrom(Operation *op, const MemorySlot &slot) const {
+ return cast<vector::LoadOp>(op).getBase() == slot.ptr;
+ }
+
+ bool storesTo(Operation *op, const MemorySlot &slot) const { return false; }
+
+ Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
+ Value reachingDef, const DataLayout &dataLayout) const {
+ llvm_unreachable("getStored should not be called on LoadOp");
+ }
+
+ bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) const {
+ return isWholeBufferLoadStore(cast<vector::LoadOp>(op), slot, blockingUses);
+ }
+
+ DeletionKind
+ removeBlockingUses(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ OpBuilder &builder, Value reachingDefinition,
+ const DataLayout &dataLayout) const {
+ cast<vector::LoadOp>(op).getResult().replaceAllUsesWith(reachingDefinition);
+ return DeletionKind::Delete;
+ }
+};
+
+struct StoreOpMemOpModel
+ : public PromotableMemOpInterface::ExternalModel<StoreOpMemOpModel,
+ vector::StoreOp> {
+ bool loadsFrom(Operation *op, const MemorySlot &slot) const { return false; }
+
+ bool storesTo(Operation *op, const MemorySlot &slot) const {
+ return cast<vector::StoreOp>(op).getBase() == slot.ptr;
+ }
+
+ Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
+ Value reachingDef, const DataLayout &dataLayout) const {
+ return cast<vector::StoreOp>(op).getValueToStore();
+ }
+
+ bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) const {
+ return isWholeBufferLoadStore(cast<vector::StoreOp>(op), slot, blockingUses);
+ }
+
+ DeletionKind
+ removeBlockingUses(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ OpBuilder &builder, Value reachingDefinition,
+ const DataLayout &dataLayout) const {
+ return DeletionKind::Delete;
+ }
+};
+
struct TransferReadOpMemOpModel
: public PromotableMemOpInterface::ExternalModel<TransferReadOpMemOpModel,
vector::TransferReadOp> {
@@ -315,6 +394,8 @@ struct SubViewOpPromotableModel
void mlir::vector::registerMemorySlotOpInterfaceExternalModels(
DialectRegistry ®istry) {
registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+ LoadOp::attachInterface<LoadOpMemOpModel>(*ctx);
+ StoreOp::attachInterface<StoreOpMemOpModel>(*ctx);
TransferReadOp::attachInterface<TransferReadOpMemOpModel>(*ctx);
TransferWriteOp::attachInterface<TransferWriteOpMemOpModel>(*ctx);
});
diff --git a/mlir/test/Dialect/Vector/mem2reg-load-store.mlir b/mlir/test/Dialect/Vector/mem2reg-load-store.mlir
new file mode 100644
index 0000000000000..9881cce44699c
--- /dev/null
+++ b/mlir/test/Dialect/Vector/mem2reg-load-store.mlir
@@ -0,0 +1,229 @@
+// RUN: mlir-opt %s --mem2reg --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @whole_buffer
+// CHECK-SAME: (%[[V:.*]]: vector<4xf32>)
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.store
+// CHECK-NOT: vector.load
+// CHECK: return %[[V]] : vector<4xf32>
+func.func @whole_buffer(%v: vector<4xf32>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<4xf32>
+ vector.store %v, %a[%c0] : memref<4xf32>, vector<4xf32>
+ %r = vector.load %a[%c0] : memref<4xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @whole_buffer_2d
+// CHECK-SAME: (%[[V:.*]]: vector<2x4xf32>)
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.store
+// CHECK-NOT: vector.load
+// CHECK: return %[[V]] : vector<2x4xf32>
+func.func @whole_buffer_2d(%v: vector<2x4xf32>) -> vector<2x4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<2x4xf32>
+ vector.store %v, %a[%c0, %c0] : memref<2x4xf32>, vector<2x4xf32>
+ %r = vector.load %a[%c0, %c0] : memref<2x4xf32>, vector<2x4xf32>
+ return %r : vector<2x4xf32>
+}
+
+// -----
+
+// The loop carries the current vector value instead of the buffer.
+// CHECK-LABEL: func.func @in_loop
+// CHECK-SAME: (%[[V:.*]]: vector<4xf32>,
+// CHECK-NOT: memref.alloca
+// CHECK: %[[R:.*]] = scf.for {{.*}} iter_args(%[[IT:.*]] = %[[V]]) -> (vector<4xf32>)
+// CHECK-NEXT: %[[NEXT:.*]] = arith.addf %[[IT]], %[[IT]] : vector<4xf32>
+// CHECK-NEXT: scf.yield %[[NEXT]] : vector<4xf32>
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[R]] : vector<4xf32>
+func.func @in_loop(%v: vector<4xf32>, %lb: index, %ub: index, %step: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<4xf32>
+ vector.store %v, %a[%c0] : memref<4xf32>, vector<4xf32>
+ scf.for %i = %lb to %ub step %step {
+ %old = vector.load %a[%c0] : memref<4xf32>, vector<4xf32>
+ %next = arith.addf %old, %old : vector<4xf32>
+ vector.store %next, %a[%c0] : memref<4xf32>, vector<4xf32>
+ }
+ %r = vector.load %a[%c0] : memref<4xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// Both kinds of vector access can use the same promoted slot.
+// CHECK-LABEL: func.func @mixed_transfers
+// CHECK-SAME: (%[[V:.*]]: vector<4xf32>,
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.transfer
+// CHECK-NOT: vector.store
+// CHECK-NOT: vector.load
+// CHECK: return %[[V]] : vector<4xf32>
+func.func @mixed_transfers(%v: vector<4xf32>, %pad: f32) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<4xf32>
+ vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<4xf32>, memref<4xf32>
+ %loaded = vector.load %a[%c0] : memref<4xf32>, vector<4xf32>
+ vector.store %loaded, %a[%c0] : memref<4xf32>, vector<4xf32>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// Load/store works with the existing subview alias projections.
+// CHECK-LABEL: func.func @subview
+// CHECK-SAME: (%[[INIT:.*]]: vector<8xf32>, %[[V:.*]]: vector<4xf32>)
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.subview
+// CHECK: %[[INS:.*]] = vector.insert_strided_slice %[[V]], %[[INIT]] {{.*}}offsets = [2], strides = [1]
+// CHECK: %[[EXT:.*]] = vector.extract_strided_slice %[[INS]] {{.*}}offsets = [2], sizes = [4], strides = [1]
+// CHECK-NEXT: return %[[EXT]] : vector<4xf32>
+func.func @subview(%init: vector<8xf32>, %v: vector<4xf32>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8xf32>
+ vector.store %init, %a[%c0] : memref<8xf32>, vector<8xf32>
+ %sv = memref.subview %a[2] [4] [1] : memref<8xf32> to memref<4xf32, strided<[1], offset: 2>>
+ vector.store %v, %sv[%c0] : memref<4xf32, strided<[1], offset: 2>>, vector<4xf32>
+ %r = vector.load %sv[%c0] : memref<4xf32, strided<[1], offset: 2>>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// The allocator proves that this dynamic buffer has exactly the vector width.
+// CHECK-LABEL: func.func @scalable
+// CHECK-SAME: (%[[V:.*]]: vector<[4]xf32>)
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.store
+// CHECK-NOT: vector.load
+// CHECK: return %[[V]] : vector<[4]xf32>
+func.func @scalable(%v: vector<[4]xf32>) -> vector<[4]xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %vs = vector.vscale
+ %size = arith.muli %vs, %c4 : index
+ %a = memref.alloca(%size) : memref<?xf32>
+ vector.store %v, %a[%c0] : memref<?xf32>, vector<[4]xf32>
+ %r = vector.load %a[%c0] : memref<?xf32>, vector<[4]xf32>
+ return %r : vector<[4]xf32>
+}
+
+// -----
+
+// A partial read blocks promotion even when the store covers the whole buffer.
+// CHECK-LABEL: func.func @partial_load
+// CHECK: memref.alloca
+// CHECK: vector.store
+// CHECK: vector.load
+func.func @partial_load(%v: vector<8xf32>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8xf32>
+ vector.store %v, %a[%c0] : memref<8xf32>, vector<8xf32>
+ %r = vector.load %a[%c0] : memref<8xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// A partial store also blocks promotion.
+// CHECK-LABEL: func.func @partial_store
+// CHECK: memref.alloca
+// CHECK: vector.store
+// CHECK: vector.load
+func.func @partial_store(%v: vector<4xf32>) -> vector<8xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8xf32>
+ vector.store %v, %a[%c0] : memref<8xf32>, vector<4xf32>
+ %r = vector.load %a[%c0] : memref<8xf32>, vector<8xf32>
+ return %r : vector<8xf32>
+}
+
+// -----
+
+// Exact vector type alone is insufficient: the access must start at zero.
+// CHECK-LABEL: func.func @nonzero_index
+// CHECK: memref.alloca
+// CHECK: vector.store
+// CHECK: vector.load
+func.func @nonzero_index(%v: vector<4xf32>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %a = memref.alloca() : memref<4xf32>
+ vector.store %v, %a[%c0] : memref<4xf32>, vector<4xf32>
+ %r = vector.load %a[%c1] : memref<4xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// A dynamic store index cannot be proven to cover the whole buffer.
+// CHECK-LABEL: func.func @dynamic_index
+// CHECK: memref.alloca
+// CHECK: vector.store
+// CHECK: vector.load
+func.func @dynamic_index(%v: vector<4xf32>, %i: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<4xf32>
+ vector.store %v, %a[%i] : memref<4xf32>, vector<4xf32>
+ %r = vector.load %a[%c0] : memref<4xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// Single-element scalar buffers still use scalar slots, not vector<1x...>.
+// CHECK-LABEL: func.func @scalar_slot
+// CHECK: memref.alloca
+// CHECK: vector.store
+// CHECK: vector.load
+func.func @scalar_slot(%v: vector<1xf32>) -> vector<1xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<1xf32>
+ vector.store %v, %a[%c0] : memref<1xf32>, vector<1xf32>
+ %r = vector.load %a[%c0] : memref<1xf32>, vector<1xf32>
+ return %r : vector<1xf32>
+}
+
+// -----
+
+// A conditional store preserves the previous value on the other path.
+// CHECK-LABEL: func.func @conditional_store
+// CHECK-SAME: (%[[INIT:.*]]: vector<4xf32>, %[[V:.*]]: vector<4xf32>, %[[COND:.*]]: i1)
+// CHECK-NOT: memref.alloca
+// CHECK: %[[R:.*]] = scf.if %[[COND]] -> (vector<4xf32>) {
+// CHECK-NEXT: scf.yield %[[V]] : vector<4xf32>
+// CHECK-NEXT: } else {
+// CHECK-NEXT: scf.yield %[[INIT]] : vector<4xf32>
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[R]] : vector<4xf32>
+func.func @conditional_store(%init: vector<4xf32>, %v: vector<4xf32>, %cond: i1) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<4xf32>
+ vector.store %init, %a[%c0] : memref<4xf32>, vector<4xf32>
+ scf.if %cond {
+ vector.store %v, %a[%c0] : memref<4xf32>, vector<4xf32>
+ }
+ %r = vector.load %a[%c0] : memref<4xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// An arbitrary dynamic extent cannot be proven to match the vector width.
+// CHECK-LABEL: func.func @dynamic_shape
+// CHECK: memref.alloca
+// CHECK: vector.store
+// CHECK: vector.load
+func.func @dynamic_shape(%v: vector<4xf32>, %size: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca(%size) : memref<?xf32>
+ vector.store %v, %a[%c0] : memref<?xf32>, vector<4xf32>
+ %r = vector.load %a[%c0] : memref<?xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/224182
More information about the Mlir-commits
mailing list