[Mlir-commits] [mlir] [mlir][memref] Promote buffers accessed via memref.copy (PR #232078)

[Federico Bruzzone]

================
@@ -164,7 +168,120 @@ Value mlir::memref::buildDynamicSubViewMask(OpBuilder &builder, Location loc,
   return buildSubViewMask(builder, loc, subView);
 }
 
+/// Produces the value that a `memref.copy` stores into a slot: reads `mem` at
+/// the origin into `vecType` with an identity map, marking every dimension
+/// in-bounds and masking the read with `mask` when set.
+///
+/// The padding value does not affect the copy: both operands must have the same
+/// runtime shape. For a whole-buffer slot, the source covers the whole vector.
+/// For a dynamic-subview slot,
+/// `PromotableAliaserInterface::projectAliasValueToSlotValue` discards lanes
+/// outside the subview's extent with a masked select.
+static Value readMemRefAsVector(OpBuilder &builder, Location loc, Value mem,
+                                VectorType vecType, Value mask) {
+  assert(!vecType.isScalable() && "expected a fixed-size vector");
+  int64_t rank = vecType.getRank();
+  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
+  SmallVector<Value> indices(rank, zero);
+  Value padding = arith::ConstantOp::create(
+      builder, loc, builder.getZeroAttr(vecType.getElementType()));
+  SmallVector<bool> inBounds(rank, true);
+  return vector::TransferReadOp::create(
+      builder, loc, vecType, mem, indices,
+      AffineMapAttr::get(builder.getMultiDimIdentityMap(rank)), padding, mask,
+      builder.getBoolArrayAttr(inBounds));
+}
+
+/// Stores a slot's value to the target of a `memref.copy`: writes `vec` into
+/// `mem` at the origin, with the identity map and in-bounds rule of
+/// `readMemRefAsVector`.
+static void writeVectorToMemRef(OpBuilder &builder, Location loc, Value vec,
+                                Value mem, Value mask) {
+  auto vecType = cast<VectorType>(vec.getType());
+  assert(!vecType.isScalable() && "expected a fixed-size vector");
+  int64_t rank = vecType.getRank();
+  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
+  SmallVector<Value> indices(rank, zero);
+  SmallVector<bool> inBounds(rank, true);
+  vector::TransferWriteOp::create(
+      builder, loc, vec, mem, indices,
+      AffineMapAttr::get(builder.getMultiDimIdentityMap(rank)), mask,
+      builder.getBoolArrayAttr(inBounds));
+}
+
+//===----------------------------------------------------------------------===//
+//  PromotableMemOpInterface model for memref.copy
+//===----------------------------------------------------------------------===//
 namespace {
+/// `PromotableMemOpInterface` model for `memref.copy`.
+///
+/// A copy into the slot becomes a `vector.transfer_read` of the source in
+/// `getStored`. A copy out of the slot becomes a `vector.transfer_write` of the
+/// reaching definition in `removeBlockingUses`.
+///
+/// A copy between two slots is rewritten one slot at a time, in either order.
+/// For dynamic subviews,
+/// `PromotableAliaserInterface::projectAliasValueToSlotValue` masks writes into
+/// the slot.
+struct CopyOpMemOpModel
+    : public PromotableMemOpInterface::ExternalModel<CopyOpMemOpModel,
+                                                     memref::CopyOp> {
+  bool loadsFrom(Operation *op, const MemorySlot &slot) const {
+    return cast<memref::CopyOp>(op).getSource() == slot.ptr;
+  }
+
+  bool storesTo(Operation *op, const MemorySlot &slot) const {
+    return cast<memref::CopyOp>(op).getTarget() == slot.ptr;
+  }
+
+  // Promotion replaces the slot with a vector value, so a copy into the slot
+  // becomes the code producing that value: a read of the copy's source.
+  Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
+                  Value reachingDef, const DataLayout &dataLayout) const {
+    auto copyOp = cast<memref::CopyOp>(op);
+    Location loc = op->getLoc();
+    return readMemRefAsVector(
+        builder, loc, copyOp.getSource(), cast<VectorType>(slot.valueType),
+        memref::buildDynamicSubViewMask(builder, loc, slot.ptr));
+  }
+
+  bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
+                        const SmallPtrSetImpl<OpOperand *> &blockingUses,
+                        SmallVectorImpl<OpOperand *> &newBlockingUses,
+                        const DataLayout &dataLayout) const {
+    auto copyOp = cast<memref::CopyOp>(op);
+    auto vecType = dyn_cast<VectorType>(slot.valueType);
+    if (!vecType || vecType.isScalable())
----------------
FedericoBruzzone wrote:

` !vecType` also covers `memref<1xf32>`. Is that what we want?

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