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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-spirv

@llvm/pr-subscribers-mlir

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/206016.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp (+40-1) 
- (modified) mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir (+22) 


``````````diff
diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
index 4674aa351315f..2e7ca0c8f2b61 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,7 +1277,8 @@ namespace mlir {
 void populateMemRefToSPIRVPatterns(const SPIRVTypeConverter &typeConverter,
                                    RewritePatternSet &patterns) {
   patterns.add<AllocaOpPattern, AllocOpPattern, AtomicRMWOpPattern,
-               DeallocOpPattern, IntLoadOpPattern, ImageLoadOpPattern,
+               CopyOpPattern, DeallocOpPattern, IntLoadOpPattern,
+               ImageLoadOpPattern,
                IntStoreOpPattern, LoadOpPattern, MemorySpaceCastOpPattern,
                StoreOpPattern, ReinterpretCastPattern, CastPattern,
                ExtractAlignedPointerAsIndexOpPattern>(typeConverter,
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)>

``````````

</details>


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