[Mlir-commits] [mlir] Added FoldSelfCopy pattern to gpu memcpy (PR #221216)

Fri Sep 4 05:51:44 PDT 2026

llvmorg-github-actions[bot] wrote:




@llvm/pr-subscribers-mlir

Author: Edward-EJ

<details>
<summary>Changes</summary>

Hi, 

This is my first attempt at a PR for MLIR so I've decided to try implement something simple that I encountered while working on lowering my dialect to GPU targets, namely a canonicalization pattern for the gpu.memcpy operation. The purpose of the pattern is to fold self copies, i.e 
```
gpu.memcpy(%a, %a) : memref<2xf16>, memref<2xf16>
```
While I don't think IR like the above is going to be generated often, I also think it can be canonicalized away at this stage. Currently lowering the following IR: 
```
    func.func @main(%arg0: memref<100xf16, 1>)
    {
        %1 = gpu.wait async 
        %2 = gpu.memcpy async [%1]  %arg0, %arg0: memref<100xf16, 1>, memref<100xf16, 1>
        gpu.wait[%2]
        return 
    }
```
using "--canonicalize --gpu-to-llvm --reconcile-unrealized-casts --canonicalize" yields:
```
llvm.func @main(%arg0: !llvm.ptr<1>, %arg1: !llvm.ptr<1>, %arg2: i64, %arg3: i64, %arg4: i64) {
    %0 = llvm.mlir.zero : !llvm.ptr<1>
    %1 = llvm.call @mgpuStreamCreate() : () -> !llvm.ptr
    %2 = llvm.getelementptr %0[100] : (!llvm.ptr<1>) -> !llvm.ptr<1>, f16
    %3 = llvm.ptrtoint %2 : !llvm.ptr<1> to i64
    %4 = llvm.addrspacecast %arg1 : !llvm.ptr<1> to !llvm.ptr
    %5 = llvm.addrspacecast %arg1 : !llvm.ptr<1> to !llvm.ptr
    llvm.call @mgpuMemcpy(%5, %4, %3, %1) : (!llvm.ptr, !llvm.ptr, i64, !llvm.ptr) -> ()
    llvm.call @mgpuStreamSynchronize(%1) : (!llvm.ptr) -> ()
    llvm.call @mgpuStreamDestroy(%1) : (!llvm.ptr) -> ()
    llvm.return
  }
  llvm.func @mgpuStreamCreate() -> !llvm.ptr
  llvm.func @mgpuMemcpy(!llvm.ptr, !llvm.ptr, i64, !llvm.ptr)
  llvm.func @mgpuStreamSynchronize(!llvm.ptr)
  llvm.func @mgpuStreamDestroy(!llvm.ptr)
}
```
but I think this should just be a no-op.

The only complexity with this pattern is how to deal with the asynchronous case when the ```memcpy``` operation produces an asyncToken. In the case where the ```memcpy``` has just one async dependency, the pattern will simply replace the users of the result of the ```memcpy``` with its async dependency. When there are several async dependencies, the ```memcpy``` can instead be replaced by a ```wait``` which takes in the ```memcpy```'s dependencies.

I look forward to hearing your feedback and thank you in advance for your time.

--Edward 

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


2 Files Affected:

- (modified) mlir/lib/Dialect/GPU/IR/GPUDialect.cpp (+35-6) 
- (modified) mlir/test/Dialect/GPU/canonicalize.mlir (+55) 


``````````diff

diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index a39b43c9803eb..4792cad60e715 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -2089,6 +2089,35 @@ LogicalResult MemcpyOp::verify() {
 
 namespace {
 
+/// Fold gpu.memcpy(%x, %x).
+struct FoldSelfCopy : public OpRewritePattern<MemcpyOp> {
+  using OpRewritePattern<MemcpyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(MemcpyOp memcpyOp,
+                                PatternRewriter &rewriter) const override {
+    if (memcpyOp.getSrc() != memcpyOp.getDst())
+      return failure();
+
+    // Synchronous self-copy is a no-op.
+    if (!memcpyOp.getAsyncToken()) {
+      rewriter.eraseOp(memcpyOp);
+      return success();
+    }
+
+    ValueRange asyncDependencies = memcpyOp.getAsyncDependencies();
+
+    // If only 1 dependency, replace `memcpyOp` with the dependency.
+    if (asyncDependencies.size() == 1)
+      rewriter.replaceOp(memcpyOp, asyncDependencies.front());
+
+    // If several dependencies, replace `memcpyOp` with a `gpu.wait`.
+    rewriter.replaceOpWithNewOp<WaitOp>(
+        memcpyOp, memcpyOp.getAsyncToken().getType(), asyncDependencies);
+
+    return success();
+  }
+};
+
 /// Erases a common case of copy ops where a destination value is used only by
 /// the copy op, alloc and dealloc ops.
 struct EraseTrivialCopyOp : public OpRewritePattern<MemcpyOp> {
@@ -2126,7 +2155,12 @@ struct EraseTrivialCopyOp : public OpRewritePattern<MemcpyOp> {
 
 void MemcpyOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
-  results.add<EraseTrivialCopyOp>(context);
+  results.add<EraseTrivialCopyOp, FoldSelfCopy>(context);
+}
+
+LogicalResult MemcpyOp::fold(FoldAdaptor adaptor,
+                             SmallVectorImpl<::mlir::OpFoldResult> &results) {
+  return memref::foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
@@ -2198,11 +2232,6 @@ LogicalResult SubgroupMmaComputeOp::verify() {
   return success();
 }
 
-LogicalResult MemcpyOp::fold(FoldAdaptor adaptor,
-                             SmallVectorImpl<::mlir::OpFoldResult> &results) {
-  return memref::foldMemRefCast(*this);
-}
-
 LogicalResult MemsetOp::fold(FoldAdaptor adaptor,
                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
   return memref::foldMemRefCast(*this);
diff --git a/mlir/test/Dialect/GPU/canonicalize.mlir b/mlir/test/Dialect/GPU/canonicalize.mlir
index 8f2803b267490..e3811f3b87dee 100644
--- a/mlir/test/Dialect/GPU/canonicalize.mlir
+++ b/mlir/test/Dialect/GPU/canonicalize.mlir
@@ -157,6 +157,61 @@ func.func @fold_memcpy_op(%arg0: i1) {
 
 // -----
 
+// CHECK-LABEL: func @synhcronous_self_memcpy
+func.func @synhcronous_self_memcpy(%arg0: memref<2xf16>)
+{
+  gpu.memcpy %arg0, %arg0 : memref<2xf16>, memref<2xf16>
+  return
+}
+
+// CHECK-NOT: gpu.memcpy
+
+// -----
+
+// CHECK-LABEL: func @asynchronous_self_memcpy1
+func.func @asynchronous_self_memcpy1(%arg0: memref<2xf16>) {
+  %1 = gpu.wait async
+  %memref, %asyncToken = gpu.alloc async [%1] () : memref<2xf16>
+  gpu.wait [%1]
+  %2 = gpu.wait async
+  %3 = gpu.memcpy async [%2] %memref, %memref : memref<2xf16>, memref<2xf16>
+  gpu.wait [%2]
+  return
+}
+
+// CHECK-NOT: gpu.memcpy
+
+// -----
+
+// CHECK-LABEL: func @asynchronous_self_memcpy2(
+// CHECK-SAME:    %[[ARG0:.*]]: memref<2xf16>
+func.func @asynchronous_self_memcpy2(%arg0: memref<2xf16>) {
+    // CHECK: %[[T0:.*]] = gpu.wait async
+    %1 = gpu.wait async
+    
+    // CHECK: %{{.*}}, %[[T1:.*]] = gpu.alloc async [%[[T0]]] () : memref<2xf16>
+    // CHECK: %{{.*}}, %[[T2:.*]] = gpu.alloc async [%[[T0]]] () : memref<2xf16>
+    %memref, %asyncToken = gpu.alloc async [%1] () : memref<2xf16>
+    %memref1, %asyncToken1 = gpu.alloc async [%1] () : memref<2xf16>
+
+    // CHECK-NOT: gpu.memcpy {{.*}} %[[ARG0]], %[[ARG0]]
+    // CHECK: %[[T3:.*]] = gpu.wait async [%[[T1]], %[[T2]]]
+    %3 = gpu.memcpy async [%asyncToken, %asyncToken1] %arg0, %arg0 : memref<2xf16>, memref<2xf16>
+    
+    // CHECK: gpu.wait [%[[T3]]]
+    gpu.wait [%3]
+    
+    // CHECK: gpu.memcpy async [%[[T3]]] %[[ARG0]], %{{.*}}
+    // CHECK: gpu.memcpy async [%[[T3]]] %[[ARG0]], %{{.*}}
+    %4 = gpu.memcpy async [%3] %arg0, %memref : memref<2xf16>, memref<2xf16>
+    %5 = gpu.memcpy async [%3] %arg0, %memref1 : memref<2xf16>, memref<2xf16>
+
+    gpu.wait [%4, %5]
+    return
+}
+
+// -----
+
 // We cannot fold memcpy here as dest is a block argument.
 // CHECK-LABEL: func @do_not_fold_memcpy_op1
 func.func @do_not_fold_memcpy_op1(%arg0: i1, %arg1: memref<2xf16>) {

``````````

</details>


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