[Mlir-commits] [mlir] 75e6377 - [mlir][SPIR-V] Lower memref.copy to spirv.CopyMemory (#206016)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 22 07:04:01 PDT 2026
Author: Arseniy Obolenskiy
Date: 2026-07-22T16:03:55+02:00
New Revision: 75e637791bff7b6c0c11754fb29991b6b140ee2c
URL: https://github.com/llvm/llvm-project/commit/75e637791bff7b6c0c11754fb29991b6b140ee2c
DIFF: https://github.com/llvm/llvm-project/commit/75e637791bff7b6c0c11754fb29991b6b140ee2c.diff
LOG: [mlir][SPIR-V] Lower memref.copy to spirv.CopyMemory (#206016)
Added:
Modified:
mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
index 4674aa351315f..7eca3d65052a3 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,48 @@ StoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
return success();
}
+//===----------------------------------------------------------------------===//
+// CopyOp
+//===----------------------------------------------------------------------===//
+
+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.
+ 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 +1291,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..5163120a8339e 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
@@ -544,6 +544,91 @@ 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 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