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

[Renato Golin]

================
@@ -0,0 +1,159 @@
+//===- 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 `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 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 in bounds: no out-of-buffer element, no padding.
+  if (xferOp.hasOutOfBoundsDim())
+    return false;
+
+  // A mask would make 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 {
+    // Whole-buffer read: replace the loaded vector with the reaching
+    // definition.
+    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; }
----------------
rengolin wrote:

why is this `return false` while the other is `llvm_unreachable`?

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