Convert memref copy to linalg copy (#1098)

Author: erwei-xilinx

mlir/include/air/Dialect/AIR/AIRTransformOps.td

@@ -375,6 +375,32 @@ def EliminateCascadeMemcpyOp : Op<Transform_Dialect, "air.eliminate_cascade_memc
   let assemblyFormat = "$target attr-dict";
 }
 
+def ConvertMemrefCopyToLinalgCopyOp : Op<Transform_Dialect, "air.convert_memref_copy_to_linalg_copy",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let summary = "Convert memref.copy operations to linalg.copy operations";
+  let description = [{
+    This transform converts `memref.copy` operations to `linalg.copy` operations.
+    This can be useful for enabling further linalg-based optimizations and transformations.
+    
+    The transformation replaces:
+    ```mlir
+    memref.copy %source, %dest : memref<...> to memref<...>
+    ```
+    
+    With:
+    ```mlir
+    linalg.copy ins(%source : memref<...>) outs(%dest : memref<...>)
+    ```
+    
+    Returns a handle to the modified operation containing the transformed copies.
+  }];
+  
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+  let assemblyFormat = "$target attr-dict";
+}
+
 // Ops implemented in mlir/lib/Transform/AIRLinalgBufferize.cpp
 //
 

mlir/include/air/Util/Util.h

@@ -316,6 +316,8 @@ Operation *cloneOpAndOperands(
       return !isa<LoopLikeOpInterface>(o) && !isa<air::HierarchyInterface>(o);
     });
 
+bool opOrAncestorIsDominantOver(Operation *a, Operation *b);
+
 } // namespace air
 } // namespace xilinx
 

mlir/lib/Transform/AIRDependency.cpp

@@ -790,18 +790,8 @@ class AIRDependency : public air::impl::AIRDependencyBase<AIRDependency> {
       operand.getDefiningOp()->emitOpError(
           "operand being traced is not a memref");
     }
-    auto opOrAncestorIsDominantOver = [](Operation *a, Operation *b) {
-      Region *commonRegion = air::findCommonRegionContainingAllAncestors(
-          SmallVector<Operation *>{a, b}, nullptr);
-      auto aAncestor = commonRegion->findAncestorOpInRegion(*a);
-      auto bAncestor = commonRegion->findAncestorOpInRegion(*b);
-      if (!aAncestor || !bAncestor)
-        return false;
-      DominanceInfo domInfo(aAncestor);
-      return domInfo.properlyDominates(aAncestor, bAncestor);
-    };
     for (auto &u : operand.getUses()) {
-      if (!opOrAncestorIsDominantOver(u.getOwner(), op))
+      if (!air::opOrAncestorIsDominantOver(u.getOwner(), op))
         continue;
       // If used in MemcpyInterface Op
       if (auto memcpy = dyn_cast<air::MemcpyInterface>(u.getOwner())) {

mlir/lib/Transform/AIRLinalgCodegen.cpp

@@ -567,6 +567,22 @@ struct RemoveAllocCopyLinalgOpCopyPattern
 //  #map1> memref.dealloc %11 : memref<1x32x16x16xf32, 2>
 //}
 
+struct ConvertMemrefCopyToLinalgCopyPattern
+    : public OpRewritePattern<memref::CopyOp> {
+  using OpRewritePattern<memref::CopyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::CopyOp copyOp,
+                                PatternRewriter &rewriter) const override {
+    Value source = copyOp.getSource();
+    Value target = copyOp.getTarget();
+
+    // Create linalg.copy operation
+    rewriter.replaceOpWithNewOp<linalg::CopyOp>(copyOp, source, target);
+
+    return success();
+  }
+};
+
 // Eliminate intermediate memref in cascaded DMA operations
 // Replace a pattern like this:
 //  air.dma_memcpy_nd (%intermediate[] [] [], %source[] [] []) : (memref<...>,
@@ -590,17 +606,6 @@ struct EliminateIntermediateMemrefPattern
     if (std::distance(intermediate.use_begin(), intermediate.use_end()) != 2)
       return failure();
 
-    auto opOrAncestorIsDominantOver = [](Operation *a, Operation *b) {
-      Region *commonRegion = air::findCommonRegionContainingAllAncestors(
-          SmallVector<Operation *>{a, b}, nullptr);
-      auto aAncestor = commonRegion->findAncestorOpInRegion(*a);
-      auto bAncestor = commonRegion->findAncestorOpInRegion(*b);
-      if (!aAncestor || !bAncestor)
-        return false;
-      DominanceInfo domInfo(aAncestor);
-      return domInfo.properlyDominates(aAncestor, bAncestor);
-    };
-
     // Find the second memcpy that uses the intermediate buffer as source
     air::DmaMemcpyNdOp secondMemcpy = nullptr;
     for (auto user : intermediate.getUsers()) {
@@ -609,7 +614,7 @@ struct EliminateIntermediateMemrefPattern
         continue;
       if (memcpyOp.getSrcMemref() != intermediate)
         continue;
-      if (opOrAncestorIsDominantOver(memcpyOp, firstMemcpy))
+      if (air::opOrAncestorIsDominantOver(memcpyOp, firstMemcpy))
         continue;
       secondMemcpy = memcpyOp;
       break;
@@ -2438,13 +2443,10 @@ static bool hasWritesBetween(memref::AllocOp allocOp, Operation *beforeOp) {
 
     // Only consider operations that are dominated by the allocation
     // and that dominate the beforeOp
-    if (!domInfo.properlyDominates(allocOp.getOperation(), op)) {
+    if (!xilinx::air::opOrAncestorIsDominantOver(allocOp.getOperation(), op))
       return;
-    }
-
-    if (!domInfo.properlyDominates(op, beforeOp)) {
+    if (!xilinx::air::opOrAncestorIsDominantOver(op, beforeOp))
       return;
-    }
 
     // Check if this operation writes to our allocation
     if (hasWriteEffectOn(op, allocResult)) {
@@ -2546,6 +2548,41 @@ DiagnosedSilenceableFailure transform::EliminateCascadeMemcpyOp::apply(
   return DiagnosedSilenceableFailure::success();
 }
 
+//===----------------------------------------------------------------------===//
+// ConvertMemrefCopyToLinalgCopyOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::ConvertMemrefCopyToLinalgCopyOp::apply(
+    transform::TransformRewriter &rewriter,
+    transform::TransformResults &results, transform::TransformState &state) {
+
+  SmallVector<Operation *> targets =
+      llvm::to_vector(state.getPayloadOps(getTarget()));
+
+  if (targets.empty()) {
+    results.set(llvm::cast<OpResult>(getResult()), ArrayRef<Operation *>());
+    return DiagnosedSilenceableFailure::success();
+  }
+
+  SmallVector<Operation *> transformedOps;
+
+  for (Operation *target : targets) {
+    MLIRContext *ctx = target->getContext();
+    RewritePatternSet patterns(ctx);
+
+    // Use the ConvertMemrefCopyToLinalgCopyPattern
+    patterns.insert<xilinx::air::ConvertMemrefCopyToLinalgCopyPattern>(ctx);
+
+    // Apply the pattern to convert memref.copy to linalg.copy operations
+    (void)applyPatternsGreedily(target, std::move(patterns));
+
+    transformedOps.push_back(target);
+  }
+
+  results.set(llvm::cast<OpResult>(getResult()), transformedOps);
+  return DiagnosedSilenceableFailure::success();
+}
+
 namespace xilinx {
 namespace air {
 

mlir/lib/Util/Util.cpp

@@ -2141,4 +2141,15 @@ Operation *air::cloneOpAndOperands(RewriterBase &rewriter, IRMapping &remap,
   return new_op;
 }
 
+bool air::opOrAncestorIsDominantOver(Operation *a, Operation *b) {
+  Region *commonRegion = air::findCommonRegionContainingAllAncestors(
+      SmallVector<Operation *>{a, b}, nullptr);
+  auto aAncestor = commonRegion->findAncestorOpInRegion(*a);
+  auto bAncestor = commonRegion->findAncestorOpInRegion(*b);
+  if (!aAncestor || !bAncestor)
+    return false;
+  DominanceInfo domInfo(aAncestor);
+  return domInfo.properlyDominates(aAncestor, bAncestor);
+}
+
 } // namespace xilinx

mlir/test/Transform/AIRTransform/AIRConvertMemrefCopyToLinalgCopy/air_transform.mlir

@@ -0,0 +1,20 @@
+//===- air_transform.mlir --------------------------------------*- MLIR -*-===//
+//
+// Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: air-opt %s | FileCheck %s
+
+// CHECK: transform.air.convert_memref_copy_to_linalg_copy
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+    transform.sequence %arg0 : !pdl.operation failures(propagate) {
+    ^bb1(%arg1: !pdl.operation):
+        // Convert memref.copy to linalg.copy
+        %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+        %func_op_updated = transform.air.convert_memref_copy_to_linalg_copy %func_op
+    }
+}