[Mlir-commits] [mlir] [mlir][SPIR-V] Lower memref.copy to spirv.CopyMemory (PR #206016)

[Arseniy Obolenskiy]

https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/206016

>From b3172e577a92b663badf1508cc0baa194b789724 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Fri, 26 Jun 2026 11:43:33 +0200
Subject: [PATCH] [mlir][SPIR-V] Lower memref.copy to spirv.CopyMemory

---
 .../MemRefToSPIRV/MemRefToSPIRV.cpp           | 48 +++++++++++++++++--
 .../MemRefToSPIRV/memref-to-spirv.mlir        | 22 +++++++++
 2 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
index 4674aa351315f..de0356fe7b59e 100644
--- a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
+++ b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
@@ -339,6 +339,16 @@ class StoreOpPattern final : public OpConversionPattern<memref::StoreOp> {
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Converts memref.copy to spirv.CopyMemory.
+class CopyOpPattern final : public OpConversionPattern<memref::CopyOp> {
+public:
+  using Base::Base;
+
+  LogicalResult
+  matchAndRewrite(memref::CopyOp copyOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 class ReinterpretCastPattern final
     : public OpConversionPattern<memref::ReinterpretCastOp> {
 public:
@@ -1171,6 +1181,34 @@ StoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// CopyOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+CopyOpPattern::matchAndRewrite(memref::CopyOp copyOp, OpAdaptor adaptor,
+                               ConversionPatternRewriter &rewriter) const {
+  // The converted operands are SPIR-V pointers to the source and target
+  // storage. spirv.CopyMemory copies the whole pointed-to object, so it only
+  // applies when both pointers point to the same fixed-size element type.
+  Value source = adaptor.getSource();
+  Value target = adaptor.getTarget();
+  auto sourcePtrType = dyn_cast<spirv::PointerType>(source.getType());
+  auto targetPtrType = dyn_cast<spirv::PointerType>(target.getType());
+  if (!sourcePtrType || !targetPtrType)
+    return rewriter.notifyMatchFailure(copyOp, "failed to convert memref type");
+
+  if (sourcePtrType.getPointeeType() != targetPtrType.getPointeeType())
+    return rewriter.notifyMatchFailure(
+        copyOp, "source and target pointee types do not match");
+
+  rewriter.replaceOpWithNewOp<spirv::CopyMemoryOp>(
+      copyOp, target, source, /*memory_access=*/spirv::MemoryAccessAttr{},
+      /*alignment=*/IntegerAttr{}, /*source_memory_access=*/
+      spirv::MemoryAccessAttr{}, /*source_alignment=*/IntegerAttr{});
+  return success();
+}
+
 LogicalResult ReinterpretCastPattern::matchAndRewrite(
     memref::ReinterpretCastOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
@@ -1239,10 +1277,10 @@ namespace mlir {
 void populateMemRefToSPIRVPatterns(const SPIRVTypeConverter &typeConverter,
                                    RewritePatternSet &patterns) {
   patterns.add<AllocaOpPattern, AllocOpPattern, AtomicRMWOpPattern,
-               DeallocOpPattern, IntLoadOpPattern, ImageLoadOpPattern,
-               IntStoreOpPattern, LoadOpPattern, MemorySpaceCastOpPattern,
-               StoreOpPattern, ReinterpretCastPattern, CastPattern,
-               ExtractAlignedPointerAsIndexOpPattern>(typeConverter,
-                                                      patterns.getContext());
+               CopyOpPattern, DeallocOpPattern, IntLoadOpPattern,
+               ImageLoadOpPattern, IntStoreOpPattern, LoadOpPattern,
+               MemorySpaceCastOpPattern, StoreOpPattern, ReinterpretCastPattern,
+               CastPattern, ExtractAlignedPointerAsIndexOpPattern>(
+      typeConverter, patterns.getContext());
 }
 } // namespace mlir
diff --git a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
index 931dd43be33c3..8702e82154664 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
@@ -544,6 +544,28 @@ module attributes {
 
 // -----
 
+// Check memref.copy lowering to spirv.CopyMemory.
+
+module attributes {
+  spirv.target_env = #spirv.target_env<
+    #spirv.vce<v1.0, [Kernel, Addresses], []>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: func @copy
+func.func @copy() {
+  // CHECK: %[[SRC:.+]] = spirv.Variable : !spirv.ptr<!spirv.array<4 x f32>, Function>
+  // CHECK: %[[DST:.+]] = spirv.Variable : !spirv.ptr<!spirv.array<4 x f32>, Function>
+  // CHECK: spirv.CopyMemory "Function" %[[DST]], "Function" %[[SRC]]{{.*}} : !spirv.array<4 x f32>
+  %0 = memref.alloca() : memref<4xf32, #spirv.storage_class<Function>>
+  %1 = memref.alloca() : memref<4xf32, #spirv.storage_class<Function>>
+  memref.copy %0, %1 : memref<4xf32, #spirv.storage_class<Function>> to memref<4xf32, #spirv.storage_class<Function>>
+  return
+}
+
+} // end module
+
+// -----
+
 // Check Image Support.
 
 // CHECK: #[[$COLMAJMAP:.*]] = affine_map<(d0, d1) -> (d1, d0)>