[Mlir-commits] [mlir] [mlir][SPIR-V] Lower memref.copy to spirv.CopyMemory (PR #206016)
Arseniy Obolenskiy
llvmlistbot at llvm.org
Wed Jul 22 06:21:13 PDT 2026
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 1/2] [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)>
>From 0f324623758c3f3c962a560b0034b00630735aee Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Wed, 22 Jul 2026 15:21:00 +0200
Subject: [PATCH 2/2] Address comment
---
.../MemRefToSPIRV/MemRefToSPIRV.cpp | 14 +++++
.../MemRefToSPIRV/memref-to-spirv.mlir | 63 +++++++++++++++++++
2 files changed, 77 insertions(+)
diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
index de0356fe7b59e..7eca3d65052a3 100644
--- a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
+++ b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
@@ -1188,6 +1188,20 @@ StoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
LogicalResult
CopyOpPattern::matchAndRewrite(memref::CopyOp copyOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const {
+ auto memrefType = cast<MemRefType>(copyOp.getSource().getType());
+ if (!memrefType.hasStaticShape())
+ return rewriter.notifyMatchFailure(copyOp, "unsupported dynamic shape");
+
+ for (MemRefType type :
+ {memrefType, cast<MemRefType>(copyOp.getTarget().getType())}) {
+ auto memorySpaceAttr =
+ dyn_cast_if_present<spirv::StorageClassAttr>(type.getMemorySpace());
+ if (memorySpaceAttr &&
+ memorySpaceAttr.getValue() == spirv::StorageClass::Image)
+ return rewriter.notifyMatchFailure(
+ copyOp, "cannot lower memref.copy in image storage class");
+ }
+
// 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.
diff --git a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
index 8702e82154664..5163120a8339e 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
@@ -566,6 +566,69 @@ func.func @copy() {
// -----
+// Check memref.copy is not lowered when either operand is in Image storage
+// class, since spirv.CopyMemory does not apply to image memory.
+
+module attributes {
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [
+ Shader,
+ Image1D,
+ StorageImageExtendedFormats
+ ], [
+ SPV_KHR_storage_buffer_storage_class
+ ]>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: func @copy_image
+// CHECK: memref.copy
+func.func @copy_image(%arg0: memref<4xf32, #spirv.storage_class<Image>>, %arg1: memref<4xf32, #spirv.storage_class<Image>>) {
+ memref.copy %arg0, %arg1 : memref<4xf32, #spirv.storage_class<Image>> to memref<4xf32, #spirv.storage_class<Image>>
+ return
+}
+
+} // end module
+
+// -----
+
+// Check memref.copy is not lowered for dynamically shaped memrefs, since
+// spirv.CopyMemory requires the copied type to have a fixed size.
+
+module attributes {
+ spirv.target_env = #spirv.target_env<
+ #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: func @copy_dynamic_shape
+// CHECK: memref.copy
+func.func @copy_dynamic_shape(%arg0: memref<?xf32, #spirv.storage_class<StorageBuffer>>, %arg1: memref<?xf32, #spirv.storage_class<StorageBuffer>>) {
+ memref.copy %arg0, %arg1 : memref<?xf32, #spirv.storage_class<StorageBuffer>> to memref<?xf32, #spirv.storage_class<StorageBuffer>>
+ return
+}
+
+} // end module
+
+// -----
+
+// Check memref.copy is not lowered for dynamically shaped memrefs under the
+// Kernel addressing model either, since the converted pointer only refers to
+// a single element and does not carry the dynamic extent.
+
+module attributes {
+ spirv.target_env = #spirv.target_env<
+ #spirv.vce<v1.0, [Kernel, Addresses], []>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: func @copy_dynamic_shape_kernel
+// CHECK: memref.copy
+func.func @copy_dynamic_shape_kernel(%arg0: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>, %arg1: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>) {
+ memref.copy %arg0, %arg1 : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, #spirv.storage_class<CrossWorkgroup>>
+ return
+}
+
+} // end module
+
+// -----
+
// Check Image Support.
// CHECK: #[[$COLMAJMAP:.*]] = affine_map<(d0, d1) -> (d1, d0)>
More information about the Mlir-commits
mailing list