[Mlir-commits] [mlir] [mlir][mem2reg] Promote whole-buffer memref to a vector SSA value (PR #211880)
Jianhui Li
llvmlistbot at llvm.org
Fri Jul 31 19:48:50 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/211880
>From 7b298c473292e8ea53c2b398df384713b908bafb Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 24 Jul 2026 16:54:30 +0000
Subject: [PATCH 1/6] [mlir][mem2reg] Promote whole-buffer memref to vector via
transfer ops
Mem2Reg previously only promoted memref allocas holding a single element
(to a scalar SSA value). The pass machinery itself is type-agnostic: a
MemorySlot carries an arbitrary elemType, the reaching definition is
threaded at that type, and scf.for already threads any slot type as an
iter_arg/result. The scalar-only behavior lived entirely in
memref.alloca's getPromotableSlots and the load/store type checks.
This teaches Mem2Reg to promote a statically-shaped multi-element memref
whose only accesses are whole-buffer vector.transfer_read/transfer_write
into a single vector SSA value:
- memref::AllocaOp::getPromotableSlots now additionally offers a
VectorType slot (shape+element type of the memref) for static
multi-element memrefs with a valid vector element type. This is purely
additive and self-guarding: any non-whole-buffer access fails its
canUsesBeRemoved check and promotion of the slot safely aborts.
- vector.transfer_read/transfer_write implement PromotableMemOpInterface
via an out-of-line ExternalModel (matching how the Vector dialect
attaches its other interfaces). A transfer is treated as a whole-buffer
load/store only when it targets the slot pointer with a matching vector
type, all-zero constant indices, an identity permutation map, all
in-bounds dims, and no mask -- which also rejects scalable vectors.
No changes to the Mem2Reg pass or SCF are required.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../Transforms/MemorySlotOpInterfaceImpl.h | 21 +++
.../Dialect/MemRef/IR/MemRefMemorySlot.cpp | 19 +-
.../Dialect/Vector/Transforms/CMakeLists.txt | 2 +
.../Transforms/MemorySlotOpInterfaceImpl.cpp | 165 +++++++++++++++++
mlir/lib/RegisterAllDialects.cpp | 2 +
mlir/test/Dialect/Vector/mem2reg.mlir | 172 ++++++++++++++++++
6 files changed, 377 insertions(+), 4 deletions(-)
create mode 100644 mlir/include/mlir/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.h
create mode 100644 mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
create mode 100644 mlir/test/Dialect/Vector/mem2reg.mlir
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.h
new file mode 100644
index 0000000000000..4c2e386728419
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.h
@@ -0,0 +1,21 @@
+//===- MemorySlotOpInterfaceImpl.h - Mem2Reg for vector ops -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_VECTOR_TRANSFORMS_MEMORYSLOTOPINTERFACEIMPL_H
+#define MLIR_DIALECT_VECTOR_TRANSFORMS_MEMORYSLOTOPINTERFACEIMPL_H
+
+namespace mlir {
+
+class DialectRegistry;
+
+namespace vector {
+void registerMemorySlotOpInterfaceExternalModels(DialectRegistry ®istry);
+} // namespace vector
+} // namespace mlir
+
+#endif // MLIR_DIALECT_VECTOR_TRANSFORMS_MEMORYSLOTOPINTERFACEIMPL_H
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index 6748e2cf71804..999abc7354106 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -66,11 +66,22 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
MemRefType type = getType();
if (!type.hasStaticShape())
return {};
- // Make sure the memref contains only a single element.
- if (type.getNumElements() != 1)
- return {};
- return {MemorySlot{getResult(), type.getElementType()}};
+ // A single-element memref is promoted to a scalar SSA value.
+ if (type.getNumElements() == 1)
+ return {MemorySlot{getResult(), type.getElementType()}};
+
+ // A multi-element memref can be promoted to a single vector SSA value when it
+ // is only ever accessed as a whole buffer (e.g. through whole-buffer
+ // `vector.transfer_read`/`vector.transfer_write`). Offering the slot is safe
+ // and purely additive: any access that is not a whole-buffer transfer will
+ // fail its `canUsesBeRemoved` check and abort promotion of this slot.
+ if (VectorType::isValidElementType(type.getElementType()))
+ return {MemorySlot{
+ getResult(),
+ VectorType::get(type.getShape(), type.getElementType())}};
+
+ return {};
}
Value memref::AllocaOp::getDefaultValue(const MemorySlot &slot,
diff --git a/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
index 112a1db6fe93b..dfe873f1a1b8d 100644
--- a/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
@@ -15,6 +15,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
LowerVectorToFromElementsToShuffleTree.cpp
LowerVectorTransfer.cpp
LowerVectorTranspose.cpp
+ MemorySlotOpInterfaceImpl.cpp
SubsetOpInterfaceImpl.cpp
VectorDistribute.cpp
VectorDropLeadUnitDim.cpp
@@ -44,6 +45,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
MLIRGPUUtils
MLIRIR
MLIRLinalgDialect
+ MLIRMemorySlotInterfaces
MLIRMemRefDialect
MLIRMemRefTransforms
MLIRMemRefUtils
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
new file mode 100644
index 0000000000000..e911c38542e7a
--- /dev/null
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -0,0 +1,165 @@
+//===- MemorySlotOpInterfaceImpl.cpp - Mem2Reg for vector ops -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements Mem2Reg-related interfaces for Vector dialect
+// operations. It allows a memref that is only ever accessed as a whole buffer
+// through `vector.transfer_read`/`vector.transfer_write` to be promoted into a
+// single vector SSA value.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
+
+using namespace mlir;
+using namespace mlir::vector;
+
+//===----------------------------------------------------------------------===//
+// Utilities
+//===----------------------------------------------------------------------===//
+
+/// Returns whether the transfer operation `xferOp` accesses exactly the whole
+/// contents of the memory slot `slot`, so that it can be treated as a plain
+/// whole-buffer load or store during Mem2Reg. `blockingUses` are the uses of
+/// the slot pointer that this operation must stop using for promotion.
+template <typename TransferOpTy>
+static bool
+isWholeBufferTransfer(TransferOpTy xferOp, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses) {
+ // The only blocking use must be the slot pointer itself.
+ if (blockingUses.size() != 1)
+ return false;
+ Value blockingUse = (*blockingUses.begin())->get();
+ if (blockingUse != slot.ptr || xferOp.getBase() != slot.ptr)
+ return false;
+
+ // Only the memref form can access a memref slot. This is already implied by
+ // `getBase() == slot.ptr` above (slot pointers are always memrefs), but guard
+ // defensively against the tensor form.
+ if (!isa<MemRefType>(xferOp.getBase().getType()))
+ return false;
+
+ // The transferred vector must match the slot type exactly. This pins the
+ // rank, per-dimension extent and element type, and rejects scalable vectors
+ // (which never equal the fixed slot type).
+ if (xferOp.getVectorType() != slot.elemType)
+ return false;
+
+ // All indices must be constant zero so the access starts at the buffer
+ // origin in every dimension.
+ for (Value index : xferOp.getIndices()) {
+ std::optional<int64_t> constIndex = getConstantIntValue(index);
+ if (!constIndex || *constIndex != 0)
+ return false;
+ }
+
+ // The permutation map must be the identity: no broadcast, no transpose.
+ if (!xferOp.getPermutationMap().isIdentity())
+ return false;
+
+ // Every dimension must be in bounds so no element lies outside the buffer and
+ // no padding takes effect.
+ if (xferOp.hasOutOfBoundsDim())
+ return false;
+
+ // A mask could disable some elements, making the access partial.
+ if (xferOp.getMask())
+ return false;
+
+ return true;
+}
+
+//===----------------------------------------------------------------------===//
+// Interface models
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+struct TransferReadOpMemOpModel
+ : public PromotableMemOpInterface::ExternalModel<TransferReadOpMemOpModel,
+ vector::TransferReadOp> {
+ bool loadsFrom(Operation *op, const MemorySlot &slot) const {
+ return cast<vector::TransferReadOp>(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 TransferReadOp");
+ }
+
+ bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) const {
+ return isWholeBufferTransfer(cast<vector::TransferReadOp>(op), slot,
+ blockingUses);
+ }
+
+ DeletionKind removeBlockingUses(
+ Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder,
+ Value reachingDefinition, const DataLayout &dataLayout) const {
+ // `canUsesBeRemoved` guaranteed a whole-buffer read of the slot.
+ cast<vector::TransferReadOp>(op).getVector().replaceAllUsesWith(
+ reachingDefinition);
+ return DeletionKind::Delete;
+ }
+};
+
+struct TransferWriteOpMemOpModel
+ : public PromotableMemOpInterface::ExternalModel<TransferWriteOpMemOpModel,
+ vector::TransferWriteOp> {
+ bool loadsFrom(Operation *op, const MemorySlot &slot) const { return false; }
+
+ bool storesTo(Operation *op, const MemorySlot &slot) const {
+ return cast<vector::TransferWriteOp>(op).getBase() == slot.ptr;
+ }
+
+ Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
+ Value reachingDef, const DataLayout &dataLayout) const {
+ return cast<vector::TransferWriteOp>(op).getValueToStore();
+ }
+
+ bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) const {
+ auto xferOp = cast<vector::TransferWriteOp>(op);
+ // The stored value must not be the slot pointer itself.
+ if (xferOp.getValueToStore() == slot.ptr)
+ return false;
+ return isWholeBufferTransfer(xferOp, slot, blockingUses);
+ }
+
+ DeletionKind removeBlockingUses(
+ Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder,
+ Value reachingDefinition, const DataLayout &dataLayout) const {
+ return DeletionKind::Delete;
+ }
+};
+
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Register external models
+//===----------------------------------------------------------------------===//
+
+void mlir::vector::registerMemorySlotOpInterfaceExternalModels(
+ DialectRegistry ®istry) {
+ registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+ TransferReadOp::attachInterface<TransferReadOpMemOpModel>(*ctx);
+ TransferWriteOp::attachInterface<TransferWriteOpMemOpModel>(*ctx);
+ });
+}
diff --git a/mlir/lib/RegisterAllDialects.cpp b/mlir/lib/RegisterAllDialects.cpp
index 974b5f533860a..62304ae88b43c 100644
--- a/mlir/lib/RegisterAllDialects.cpp
+++ b/mlir/lib/RegisterAllDialects.cpp
@@ -97,6 +97,7 @@
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Vector/Transforms/IndexedAccessOpInterfaceImpl.h"
+#include "mlir/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.h"
#include "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h"
#include "mlir/Dialect/WasmSSA/IR/WasmSSA.h"
#include "mlir/Dialect/X86/X86Dialect.h"
@@ -200,6 +201,7 @@ void mlir::registerAllDialects(DialectRegistry ®istry) {
tosa::registerShardingInterfaceExternalModels(registry);
vector::registerBufferizableOpInterfaceExternalModels(registry);
vector::registerIndexedAccessOpInterfaceExternalModels(registry);
+ vector::registerMemorySlotOpInterfaceExternalModels(registry);
vector::registerSubsetOpInterfaceExternalModels(registry);
vector::registerValueBoundsOpInterfaceExternalModels(registry);
NVVM::registerNVVMTargetInterfaceExternalModels(registry);
diff --git a/mlir/test/Dialect/Vector/mem2reg.mlir b/mlir/test/Dialect/Vector/mem2reg.mlir
new file mode 100644
index 0000000000000..ec6e19d0a5c19
--- /dev/null
+++ b/mlir/test/Dialect/Vector/mem2reg.mlir
@@ -0,0 +1,172 @@
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(mem2reg))' --split-input-file | FileCheck %s
+
+// A memref that is only ever accessed as a whole buffer through
+// vector.transfer_read / vector.transfer_write is promoted to a single vector
+// SSA value.
+
+// CHECK-LABEL: func.func @whole_buffer_write_read
+// CHECK-SAME: (%[[PAD:.*]]: f32)
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.transfer_write
+// CHECK-NOT: vector.transfer_read
+// CHECK: %[[CST:.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
+// CHECK: return %[[CST]] : vector<4xf32>
+func.func @whole_buffer_write_read(%pad: f32) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<4xf32>
+ %a = memref.alloca() : memref<4xf32>
+ vector.transfer_write %cst, %a[%c0] {in_bounds = [true]} : vector<4xf32>, memref<4xf32>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// A whole-buffer slot carried across scf.for is threaded as an iter_arg/result.
+
+// CHECK-LABEL: func.func @whole_buffer_in_loop
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.transfer_write
+// CHECK-NOT: vector.transfer_read
+// CHECK: %[[RES:.*]] = scf.for {{.*}} iter_args(%[[IT:.*]] = %{{.*}}) -> (vector<4xf32>)
+// CHECK: %[[NEXT:.*]] = arith.addf %[[IT]], %[[IT]] : vector<4xf32>
+// CHECK: scf.yield %[[NEXT]] : vector<4xf32>
+// CHECK: return %[[RES]] : vector<4xf32>
+func.func @whole_buffer_in_loop(%pad: f32, %lb: index, %ub: index, %step: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<4xf32>
+ %a = memref.alloca() : memref<4xf32>
+ vector.transfer_write %cst, %a[%c0] {in_bounds = [true]} : vector<4xf32>, memref<4xf32>
+ scf.for %i = %lb to %ub step %step {
+ %v = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<4xf32>
+ %n = arith.addf %v, %v : vector<4xf32>
+ vector.transfer_write %n, %a[%c0] {in_bounds = [true]} : vector<4xf32>, memref<4xf32>
+ }
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// A multi-dimensional whole-buffer memref is promoted to a matching vector.
+
+// CHECK-LABEL: func.func @whole_buffer_2d
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.transfer
+// CHECK: return %{{.*}} : vector<2x4xf32>
+func.func @whole_buffer_2d(%pad: f32) -> vector<2x4xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<2x4xf32>
+ %a = memref.alloca() : memref<2x4xf32>
+ vector.transfer_write %cst, %a[%c0, %c0] {in_bounds = [true, true]} : vector<2x4xf32>, memref<2x4xf32>
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]} : memref<2x4xf32>, vector<2x4xf32>
+ return %r : vector<2x4xf32>
+}
+
+// -----
+
+// Non-zero access offset: the transfer does not cover the whole buffer, so the
+// slot must NOT be promoted.
+
+// CHECK-LABEL: func.func @negative_nonzero_index
+// CHECK: memref.alloca
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+func.func @negative_nonzero_index(%pad: f32) -> vector<4xf32> {
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant dense<1.0> : vector<4xf32>
+ %a = memref.alloca() : memref<8xf32>
+ vector.transfer_write %cst, %a[%c1] {in_bounds = [true]} : vector<4xf32>, memref<8xf32>
+ %r = vector.transfer_read %a[%c1], %pad {in_bounds = [true]} : memref<8xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// A masked transfer only touches part of the buffer: must NOT be promoted.
+
+// CHECK-LABEL: func.func @negative_masked
+// CHECK: memref.alloca
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+func.func @negative_masked(%pad: f32, %m: vector<4xi1>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<4xf32>
+ %a = memref.alloca() : memref<4xf32>
+ vector.transfer_write %cst, %a[%c0], %m {in_bounds = [true]} : vector<4xf32>, memref<4xf32>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// -----
+
+// A partial (out-of-bounds) transfer must NOT be promoted.
+
+// CHECK-LABEL: func.func @negative_out_of_bounds
+// CHECK: memref.alloca
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+func.func @negative_out_of_bounds(%pad: f32) -> vector<8xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<8xf32>
+ %a = memref.alloca() : memref<4xf32>
+ vector.transfer_write %cst, %a[%c0] : vector<8xf32>, memref<4xf32>
+ %r = vector.transfer_read %a[%c0], %pad : memref<4xf32>, vector<8xf32>
+ return %r : vector<8xf32>
+}
+
+// -----
+
+// A non-identity (transposing) permutation map is not a whole-buffer identity
+// access: must NOT be promoted.
+
+// CHECK-LABEL: func.func @negative_transpose_map
+// CHECK: memref.alloca
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+func.func @negative_transpose_map(%pad: f32) -> vector<4x2xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<2x4xf32>
+ %a = memref.alloca() : memref<2x4xf32>
+ vector.transfer_write %cst, %a[%c0, %c0] {in_bounds = [true, true]} : vector<2x4xf32>, memref<2x4xf32>
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : memref<2x4xf32>, vector<4x2xf32>
+ return %r : vector<4x2xf32>
+}
+
+// -----
+
+// An alloca also accessed through a scalar memref.load cannot be promoted to a
+// vector: must NOT be promoted.
+
+// CHECK-LABEL: func.func @negative_mixed_scalar_access
+// CHECK: memref.alloca
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+// CHECK: memref.load
+func.func @negative_mixed_scalar_access(%pad: f32) -> (vector<4xf32>, f32) {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<4xf32>
+ %a = memref.alloca() : memref<4xf32>
+ vector.transfer_write %cst, %a[%c0] {in_bounds = [true]} : vector<4xf32>, memref<4xf32>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<4xf32>
+ %s = memref.load %a[%c0] : memref<4xf32>
+ return %r, %s : vector<4xf32>, f32
+}
+
+// -----
+
+// A scalable vector never equals the fixed-size slot type: must NOT be
+// promoted.
+
+// CHECK-LABEL: func.func @negative_scalable
+// CHECK: memref.alloca
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+func.func @negative_scalable(%pad: f32) -> vector<[4]xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<[4]xf32>
+ %a = memref.alloca() : memref<4xf32>
+ vector.transfer_write %cst, %a[%c0] {in_bounds = [true]} : vector<[4]xf32>, memref<4xf32>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<[4]xf32>
+ return %r : vector<[4]xf32>
+}
>From 24c11489338115c200a48d648a8787901c4bb604 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 24 Jul 2026 17:40:45 +0000
Subject: [PATCH 2/6] [mlir] Simplify comments and drop dead self-store guard
in vector mem2reg
Trim verbose comments in the whole-buffer transfer promotion logic and
remove the unreachable self-store check in TransferWriteOp's
canUsesBeRemoved: a transfer_write's stored value is always a vector and
the slot pointer always a memref, so they can never be the same Value.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../Dialect/MemRef/IR/MemRefMemorySlot.cpp | 4 +--
.../Transforms/MemorySlotOpInterfaceImpl.cpp | 36 +++++++------------
2 files changed, 14 insertions(+), 26 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index 999abc7354106..59798a65f72e8 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -73,9 +73,7 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
// A multi-element memref can be promoted to a single vector SSA value when it
// is only ever accessed as a whole buffer (e.g. through whole-buffer
- // `vector.transfer_read`/`vector.transfer_write`). Offering the slot is safe
- // and purely additive: any access that is not a whole-buffer transfer will
- // fail its `canUsesBeRemoved` check and abort promotion of this slot.
+ // `vector.transfer_read`/`vector.transfer_write`).
if (VectorType::isValidElementType(type.getElementType()))
return {MemorySlot{
getResult(),
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
index e911c38542e7a..1789c59da0720 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -27,51 +27,43 @@ using namespace mlir::vector;
// Utilities
//===----------------------------------------------------------------------===//
-/// Returns whether the transfer operation `xferOp` accesses exactly the whole
-/// contents of the memory slot `slot`, so that it can be treated as a plain
-/// whole-buffer load or store during Mem2Reg. `blockingUses` are the uses of
-/// the slot pointer that this operation must stop using for promotion.
+/// Returns whether `xferOp` accesses exactly the whole contents of `slot`, so
+/// it can act as a plain whole-buffer load/store during Mem2Reg.
template <typename TransferOpTy>
static bool
isWholeBufferTransfer(TransferOpTy xferOp, const MemorySlot &slot,
const SmallPtrSetImpl<OpOperand *> &blockingUses) {
- // The only blocking use must be the slot pointer itself.
+ // The sole blocking use must be the slot pointer as the transfer's base.
if (blockingUses.size() != 1)
return false;
Value blockingUse = (*blockingUses.begin())->get();
if (blockingUse != slot.ptr || xferOp.getBase() != slot.ptr)
return false;
- // Only the memref form can access a memref slot. This is already implied by
- // `getBase() == slot.ptr` above (slot pointers are always memrefs), but guard
- // defensively against the tensor form.
+ // Reject the tensor form (already implied, since slot pointers are memrefs).
if (!isa<MemRefType>(xferOp.getBase().getType()))
return false;
- // The transferred vector must match the slot type exactly. This pins the
- // rank, per-dimension extent and element type, and rejects scalable vectors
- // (which never equal the fixed slot type).
+ // Exact type match pins rank/extents/element type and rejects scalable vectors.
if (xferOp.getVectorType() != slot.elemType)
return false;
- // All indices must be constant zero so the access starts at the buffer
- // origin in every dimension.
+ // Access must start at the buffer origin in every dimension.
for (Value index : xferOp.getIndices()) {
std::optional<int64_t> constIndex = getConstantIntValue(index);
if (!constIndex || *constIndex != 0)
return false;
}
- // The permutation map must be the identity: no broadcast, no transpose.
+ // Identity map: no broadcast or transpose.
if (!xferOp.getPermutationMap().isIdentity())
return false;
- // Every dimension must be in bounds so no element lies outside the buffer and
- // no padding takes effect.
+ // All dimensions in bounds: no out-of-buffer element, no padding.
if (xferOp.hasOutOfBoundsDim())
return false;
- // A mask could disable some elements, making the access partial.
+ // A mask would make the access partial.
if (xferOp.getMask())
return false;
@@ -110,7 +102,7 @@ struct TransferReadOpMemOpModel
Operation *op, const MemorySlot &slot,
const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder,
Value reachingDefinition, const DataLayout &dataLayout) const {
- // `canUsesBeRemoved` guaranteed a whole-buffer read of the slot.
+ // Whole-buffer read: replace the loaded vector with the reaching definition.
cast<vector::TransferReadOp>(op).getVector().replaceAllUsesWith(
reachingDefinition);
return DeletionKind::Delete;
@@ -135,11 +127,9 @@ struct TransferWriteOpMemOpModel
const SmallPtrSetImpl<OpOperand *> &blockingUses,
SmallVectorImpl<OpOperand *> &newBlockingUses,
const DataLayout &dataLayout) const {
- auto xferOp = cast<vector::TransferWriteOp>(op);
- // The stored value must not be the slot pointer itself.
- if (xferOp.getValueToStore() == slot.ptr)
- return false;
- return isWholeBufferTransfer(xferOp, slot, blockingUses);
+ // No self-store guard needed: a vector value can never equal a memref slot.
+ return isWholeBufferTransfer(cast<vector::TransferWriteOp>(op), slot,
+ blockingUses);
}
DeletionKind removeBlockingUses(
>From 1075f1a57a3eff4effb6071b00574a1f9980efc4 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 24 Jul 2026 17:44:13 +0000
Subject: [PATCH 3/6] [mlir] clang-format vector mem2reg files
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../Dialect/MemRef/IR/MemRefMemorySlot.cpp | 3 +--
.../Transforms/MemorySlotOpInterfaceImpl.cpp | 24 +++++++++++--------
2 files changed, 15 insertions(+), 12 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index 59798a65f72e8..17b1417550524 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -76,8 +76,7 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
// `vector.transfer_read`/`vector.transfer_write`).
if (VectorType::isValidElementType(type.getElementType()))
return {MemorySlot{
- getResult(),
- VectorType::get(type.getShape(), type.getElementType())}};
+ getResult(), VectorType::get(type.getShape(), type.getElementType())}};
return {};
}
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
index 1789c59da0720..352983ba153fa 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -44,7 +44,8 @@ isWholeBufferTransfer(TransferOpTy xferOp, const MemorySlot &slot,
if (!isa<MemRefType>(xferOp.getBase().getType()))
return false;
- // Exact type match pins rank/extents/element type and rejects scalable vectors.
+ // Exact type match pins rank/extents/element type and rejects scalable
+ // vectors.
if (xferOp.getVectorType() != slot.elemType)
return false;
@@ -98,11 +99,13 @@ struct TransferReadOpMemOpModel
blockingUses);
}
- DeletionKind removeBlockingUses(
- Operation *op, const MemorySlot &slot,
- const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder,
- Value reachingDefinition, const DataLayout &dataLayout) const {
- // Whole-buffer read: replace the loaded vector with the reaching definition.
+ DeletionKind
+ removeBlockingUses(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ OpBuilder &builder, Value reachingDefinition,
+ const DataLayout &dataLayout) const {
+ // Whole-buffer read: replace the loaded vector with the reaching
+ // definition.
cast<vector::TransferReadOp>(op).getVector().replaceAllUsesWith(
reachingDefinition);
return DeletionKind::Delete;
@@ -132,10 +135,11 @@ struct TransferWriteOpMemOpModel
blockingUses);
}
- DeletionKind removeBlockingUses(
- Operation *op, const MemorySlot &slot,
- const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder,
- Value reachingDefinition, const DataLayout &dataLayout) const {
+ DeletionKind
+ removeBlockingUses(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ OpBuilder &builder, Value reachingDefinition,
+ const DataLayout &dataLayout) const {
return DeletionKind::Delete;
}
};
>From 9699ca7b9767d5a28cb79aa6aa4439bca9a86895 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 25 Jul 2026 02:45:38 +0000
Subject: [PATCH 4/6] add test
---
mlir/test/Dialect/Vector/mem2reg.mlir | 20 ++++++++++++++++++++
1 file changed, 20 insertions(+)
diff --git a/mlir/test/Dialect/Vector/mem2reg.mlir b/mlir/test/Dialect/Vector/mem2reg.mlir
index ec6e19d0a5c19..7b2f0bd75ec31 100644
--- a/mlir/test/Dialect/Vector/mem2reg.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg.mlir
@@ -170,3 +170,23 @@ func.func @negative_scalable(%pad: f32) -> vector<[4]xf32> {
%r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<[4]xf32>
return %r : vector<[4]xf32>
}
+
+// -----
+
+// The buffer has a use (subview) that mem2reg does not recognize as a
+// removable load or store, so the slot is left untouched: must NOT be promoted.
+
+// CHECK-LABEL: func.func @negative_subview
+// CHECK: memref.alloca
+// CHECK: vector.transfer_write
+// CHECK: memref.subview
+// CHECK: vector.transfer_read
+func.func @negative_subview(%pad: f32) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<8xf32>
+ %a = memref.alloca() : memref<8xf32>
+ vector.transfer_write %cst, %a[%c0] {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
+ %sv = memref.subview %a[0] [4] [1] : memref<8xf32> to memref<4xf32, strided<[1]>>
+ %r = vector.transfer_read %sv[%c0], %pad {in_bounds = [true]} : memref<4xf32, strided<[1]>>, vector<4xf32>
+ return %r : vector<4xf32>
+}
>From af81c61ceeff75cdf9fe3f8f7db12920bd0b6309 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 1 Aug 2026 01:40:03 +0000
Subject: [PATCH 5/6] [mlir][mem2reg] Address PR review feedback
- Replace the isWholeBufferTransfer template with the common
VectorTransferOpInterface (adam-smnk).
- Simplify the lit RUN line to --mem2reg (adam-smnk).
- Reword the mask comment ("would" -> "could") and expand the
in-bounds/padding comment (rengolin).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Transforms/MemorySlotOpInterfaceImpl.cpp | 14 ++++++++------
mlir/test/Dialect/Vector/mem2reg.mlir | 2 +-
2 files changed, 9 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
index 352983ba153fa..cace1f9042898 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -29,9 +29,8 @@ using namespace mlir::vector;
/// Returns whether `xferOp` accesses exactly the whole contents of `slot`, so
/// it can act as a plain whole-buffer load/store during Mem2Reg.
-template <typename TransferOpTy>
static bool
-isWholeBufferTransfer(TransferOpTy xferOp, const MemorySlot &slot,
+isWholeBufferTransfer(VectorTransferOpInterface xferOp, const MemorySlot &slot,
const SmallPtrSetImpl<OpOperand *> &blockingUses) {
// The sole blocking use must be the slot pointer as the transfer's base.
if (blockingUses.size() != 1)
@@ -60,11 +59,14 @@ isWholeBufferTransfer(TransferOpTy xferOp, const MemorySlot &slot,
if (!xferOp.getPermutationMap().isIdentity())
return false;
- // All dimensions in bounds: no out-of-buffer element, no padding.
+ // All dimensions must be in bounds. An out-of-bounds dimension means the
+ // transfer reaches past the buffer, so a read would materialize padding
+ // rather than buffer contents and a write would only cover part of the
+ // buffer: in neither case does the transfer stand in for the whole slot.
if (xferOp.hasOutOfBoundsDim())
return false;
- // A mask would make the access partial.
+ // A mask could make the access partial.
if (xferOp.getMask())
return false;
@@ -95,7 +97,7 @@ struct TransferReadOpMemOpModel
const SmallPtrSetImpl<OpOperand *> &blockingUses,
SmallVectorImpl<OpOperand *> &newBlockingUses,
const DataLayout &dataLayout) const {
- return isWholeBufferTransfer(cast<vector::TransferReadOp>(op), slot,
+ return isWholeBufferTransfer(cast<VectorTransferOpInterface>(op), slot,
blockingUses);
}
@@ -131,7 +133,7 @@ struct TransferWriteOpMemOpModel
SmallVectorImpl<OpOperand *> &newBlockingUses,
const DataLayout &dataLayout) const {
// No self-store guard needed: a vector value can never equal a memref slot.
- return isWholeBufferTransfer(cast<vector::TransferWriteOp>(op), slot,
+ return isWholeBufferTransfer(cast<VectorTransferOpInterface>(op), slot,
blockingUses);
}
diff --git a/mlir/test/Dialect/Vector/mem2reg.mlir b/mlir/test/Dialect/Vector/mem2reg.mlir
index 7b2f0bd75ec31..8578f6bae905d 100644
--- a/mlir/test/Dialect/Vector/mem2reg.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(mem2reg))' --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --mem2reg --split-input-file | FileCheck %s
// A memref that is only ever accessed as a whole buffer through
// vector.transfer_read / vector.transfer_write is promoted to a single vector
>From dcd635b859242d740de319225338f6c06c457d5a Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 1 Aug 2026 02:48:13 +0000
Subject: [PATCH 6/6] [mlir][mem2reg] Add negative test for dynamic-shape
memref
A dynamic-shape memref never yields a promotable slot (its extents are
not known statically, so it cannot map to a fixed-size vector), so the
buffer is left untouched.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
mlir/test/Dialect/Vector/mem2reg.mlir | 19 +++++++++++++++++++
1 file changed, 19 insertions(+)
diff --git a/mlir/test/Dialect/Vector/mem2reg.mlir b/mlir/test/Dialect/Vector/mem2reg.mlir
index 8578f6bae905d..96ea2bc952414 100644
--- a/mlir/test/Dialect/Vector/mem2reg.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg.mlir
@@ -190,3 +190,22 @@ func.func @negative_subview(%pad: f32) -> vector<4xf32> {
%r = vector.transfer_read %sv[%c0], %pad {in_bounds = [true]} : memref<4xf32, strided<[1]>>, vector<4xf32>
return %r : vector<4xf32>
}
+
+// -----
+
+// A dynamic-shape memref never yields a promotable slot (its extents are not
+// known statically, so it cannot map to a fixed-size vector): must NOT be
+// promoted.
+
+// CHECK-LABEL: func.func @negative_dynamic_shape
+// CHECK: memref.alloca
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+func.func @negative_dynamic_shape(%pad: f32, %d: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.0> : vector<4xf32>
+ %a = memref.alloca(%d) : memref<?xf32>
+ vector.transfer_write %cst, %a[%c0] {in_bounds = [true]} : vector<4xf32>, memref<?xf32>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<?xf32>, vector<4xf32>
+ return %r : vector<4xf32>
+}
More information about the Mlir-commits
mailing list