[Mlir-commits] [mlir] [mlir][mem2reg] Promote whole-buffer memref to a vector SSA value (PR #211880)

[Jianhui Li]

================
@@ -0,0 +1,328 @@
+//===- 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 that let a statically-shaped
+// memref be promoted into a single vector SSA value, provided every access to
+// the buffer is a whole-buffer read or write (or a whole-sub-region access of
+// 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).
+//
+//   * a static, same-rank `memref.subview` is exposed as a promotable sub-slice
+//     alias of the buffer's slot (via `PromotableAliaserInterface`): a read of
+//     the subview projects out of the vector value with
+//     `vector.extract_strided_slice`, and a write into it composes back into
+//     the value with `vector.insert_strided_slice`. This lets a buffer that is
+//     only ever accessed through static subviews promote as well, with partial
+//     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
+// indices -- are left untouched, so the buffer is not promoted.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.h"
+
+#include "mlir/Dialect/MemRef/IR/MemRef.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 `xferOp` accesses exactly the whole contents of `slot`, so
+/// it can act as a plain whole-buffer load/store during Mem2Reg.
+static bool
+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)
+    return false;
+  Value blockingUse = (*blockingUses.begin())->get();
+  if (blockingUse != slot.ptr || xferOp.getBase() != slot.ptr)
+    return false;
+
+  // Reject the tensor form (already implied, since slot pointers are memrefs).
+  if (!isa<MemRefType>(xferOp.getBase().getType()))
+    return false;
+
+  // Exact type match pins rank/extents/element type and rejects scalable
+  // vectors.
+  if (xferOp.getVectorType() != slot.elemType)
+    return false;
+
+  // 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;
+  }
+
+  // Identity map: no broadcast or transpose.
+  if (!xferOp.getPermutationMap().isIdentity())
+    return false;
+
+  // 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 could make the access partial.
+  if (xferOp.getMask())
+    return false;
+
+  return true;
+}
+
+//===----------------------------------------------------------------------===//
+//  Interface models
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+struct TransferReadOpMemOpModel
+    : public PromotableMemOpInterface::ExternalModel<TransferReadOpMemOpModel,
----------------
Jianhui-Li wrote:

Those could be ODS impls (like the LLVM dialect does), but keeping them external means the core Vector dialect doesn't have to link the MemorySlot interfaces, so the whole promotion feature stays opt-in via registerMemorySlotOpInterfaceExternalModels.

MemRef also keeps its own mem2reg implementation behind registerMemorySlotExternalModels rather than in ODS.

https://github.com/llvm/llvm-project/pull/211880