[MLIR][GPU] Add canonicalizer for gpu.memcpy

[arnab-oss]

Summary

Erase gpu.memcpy op when only uses of dest are
the memcpy op in question, its allocation and deallocation
ops.

[bondhugula]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

1082–1095

Can this be made or use a helper method in the mlir Analysis library to the effect:

template <typename EffectTy>
bool mlir::hasSingleEffect(Operation *op, Value value) {
 ...
}

I see numerous use cases for something like all over.

if (hasSingleEffect<MemoryEffects::Allocate>())
 ...

[arnab-oss]

Added hasSingleEffect() as a part of SideEffectInterface.

[arnab-oss]

Fixed linting errors.

[arnab-oss]

Addressed comments by @bondhugula.

[bondhugula]

mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td

90

This is unclearly phrased.

-> `/// Returns true if EffectTy is the only effect that this operation has on value.`

?

94–104

We shouldn't have so much C++ code in the .td file. Please define out of line.

[csigg]

Looks good from my side but please address Uday's comments.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

1094–1096

Just a suggestion, feel free to ignore.

[arnab-oss]

Addressed comments by @bondhugula.

[bondhugula]

Looking mostly good. Some minor comments and a question to under the async dep related guard.

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

251 ↗

(On Diff #427613)

this operation -> op (in backticks)

Nit:
only has the ... -> has only the ..

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

1090

Need a code comment here to understand the rationale.

mlir/lib/Interfaces/SideEffectInterfaces.cpp

95 ↗

(On Diff #427613)

You don't need llvm::.

95 ↗

(On Diff #427613)

Define this small vector below the check.

102 ↗

(On Diff #427613)

Need a code comment here.

[arnab-oss]

Addressed comments by @bondhugula.

[bondhugula]

LGTM

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

1090–1092

single -> a single
dependency, -> dependency
produce -> produces
dependency, and does -> dependency and does

[arnab-oss]

Addressed comments by @bondhugula.

[arnab-oss]

Rebased on latest master.

mlir/include/mlir/Dialect/GPU/GPUOps.td

def GPU_MemcpyOp : GPU_Op<"memcpy", [GPU_AsyncOpInterface]> {
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $dst`,` $src `:` type($dst)`,` type($src) attr-dict
   }];
   let hasFolder = 1;
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 def GPU_MemsetOp : GPU_Op<"memset",
   [GPU_AsyncOpInterface, AllElementTypesMatch<["dst", "value"]>]> {
 
   let summary = "GPU memset operation";
 
   let description = [{

mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td

let extraClassDeclaration = [{
     template <typename Effect> bool onlyHasEffect() {
       SmallVector<SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
       getEffects(effects);
       return !effects.empty() && llvm::all_of(effects, [](const auto &it) {
         return isa<Effect>(it.getEffect());
       });
     }
 
+    /// Returns true if this operation only has the given effect on `value`.
+    template <typename EffectTy> bool hasSingleEffect(Value value) {
+      SmallVector<SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
+      getEffects(effects);
+      bool doesOpOnlyHaveSingleEffectOnVal = false;
+      for (auto &effect : effects) {
+        if (effect.getValue() == value &&
+            isa<EffectTy>(effect.getEffect()))
+          doesOpOnlyHaveSingleEffectOnVal = true;
+        if (effect.getValue() == value &&
+            !isa<EffectTy>(effect.getEffect())) {
+          doesOpOnlyHaveSingleEffectOnVal = false;
+          break;
+        }
+      }
+      return doesOpOnlyHaveSingleEffectOnVal;
+    }
+
     /// Returns true if this operation has no effects.
     bool hasNoEffect() {
       SmallVector<::mlir::SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
       getEffects(effects);
       return effects.empty();
     }
 
     /// Returns true if the given operation has no effects for this interface.

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

 // `MemoryEffects`.
 def MemoryEffectsOpInterface
     : EffectOpInterfaceBase<"MemoryEffectOpInterface",
                             "::mlir::MemoryEffects::Effect"> {
   let description = [{
     An interface used to query information about the memory effects applied by
     an operation.
   }];
+  
   let cppNamespace = "::mlir";
 }
 
 // The base class for defining specific memory effects.
 class MemoryEffect<string effectName, Resource resource>
   : SideEffect<MemoryEffectsOpInterface, effectName, resource>;
 
 // This class represents the trait for memory effects that may be placed on

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

if (asyncTokenType)
     printer << "async ";
   if (asyncDependencies.empty())
     return;
   printer << "[";
   llvm::interleaveComma(asyncDependencies, printer);
   printer << "]";
 }
 
+namespace {
+
+/// 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> {
+  using OpRewritePattern<MemcpyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(MemcpyOp op,
+                                PatternRewriter &rewriter) const override {
+    Value dest = op.dst();
+    Operation *destDefOp = dest.getDefiningOp();
+    if (!destDefOp || !(isa<MemoryEffectOpInterface>(destDefOp) &&
+                        cast<MemoryEffectOpInterface>(destDefOp)
+                            .hasSingleEffect<MemoryEffects::Allocate>(dest)))
+      return failure();
+    // We can erase `op` iff `dest` has no other use apart from its
+    // use by `op` and dealloc ops.
+    if (llvm::any_of(dest.getUsers(), [op, dest](Operation *user) {
+          return user != op &&
+                 !(isa<MemoryEffectOpInterface>(user) &&
+                   cast<MemoryEffectOpInterface>(user)
+                       .hasSingleEffect<MemoryEffects::Free>(dest));
+        }))
+      return failure();
+
+    if (op.asyncDependencies().size() > 1 ||
+        ((op.asyncDependencies().empty() && op.asyncToken()) ||
+         (!op.asyncDependencies().empty() && !op.asyncToken())))
+      return failure();
+    rewriter.replaceOp(op, op.asyncDependencies());
+    return success();
+  }
+};
+
+} // end anonymous namespace
+
+void MemcpyOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                           MLIRContext *context) {
+  results.add<EraseTrivialCopyOp>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // GPU_SubgroupMmaLoadMatrixOp
 //===----------------------------------------------------------------------===//
 
 /// Return true if the last dimension of the MemRefType has unit stride. Also
 /// return true for memrefs with no strides.
 static bool isLastMemrefDimUnitStride(MemRefType type) {
   int64_t offset;

mlir/test/Dialect/GPU/canonicalize.mlir

func @fold_wait_op_test2(%arg0: i1) -> (memref<5xf16>, memref<5xf16>) {
   return %memref, %memref_0 : memref<5xf16>, memref<5xf16>
 }
 // CHECK-NEXT: %[[TOKEN0:.*]] = gpu.wait async
 // CHECK-NEXT: gpu.alloc async [%[[TOKEN0]]] ()
 // CHECK-NEXT: %[[TOKEN1:.*]] = gpu.wait async
 // CHECK-NEXT: gpu.alloc async [%[[TOKEN1]]] ()
 // CHECK-NEXT: return
 
+// CHECK-LABEL: func @fold_memcpy_op
+func @fold_memcpy_op(%arg0: i1) {
+    %cst = arith.constant 0.000000e+00 : f16
+    %1 = memref.alloc() : memref<2xf16>
+    %2 = gpu.wait async
+    %memref, %asyncToken = gpu.alloc async [%2] () : memref<2xf16>
+    gpu.wait [%2]
+    affine.store %cst, %memref[0] : memref<2xf16>
+    %3 = gpu.wait async
+    %4 = gpu.memcpy async [%3] %1, %memref : memref<2xf16>, memref<2xf16>
+    gpu.wait [%3]
+    %5 = scf.if %arg0 -> (i1) {
+      memref.dealloc %1 : memref<2xf16>
+      scf.yield %arg0 : i1
+    } else {
+      memref.dealloc %1 : memref<2xf16>
+      scf.yield %arg0 : i1
+    }
+    return
+}
+// CHECK-NOT: gpu.memcpy
+
+// We cannot fold memcpy here as dest is a block argument.
+// CHECK-LABEL: func @do_not_fold_memcpy_op1
+func @do_not_fold_memcpy_op1(%arg0: i1, %arg1: memref<2xf16>) {
+    %cst = arith.constant 0.000000e+00 : f16
+    %2 = gpu.wait async
+    %memref, %asyncToken = gpu.alloc async [%2] () : memref<2xf16>
+    gpu.wait [%2]
+    affine.store %cst, %memref[0] : memref<2xf16>
+    %3 = gpu.wait async
+    %4 = gpu.memcpy async [%3] %arg1, %memref : memref<2xf16>, memref<2xf16>
+    gpu.wait [%3]
+    return
+}
+// CHECK: gpu.memcpy
+
+// We cannot fold gpu.memcpy as it is used by an op having read effect on dest.
+// CHECK-LABEL: func @do_not_fold_memcpy_op2
+func @do_not_fold_memcpy_op2(%arg0: i1, %arg1: index) -> f16 {
+    %cst = arith.constant 0.000000e+00 : f16
+    %1 = memref.alloc() : memref<2xf16>
+    %2 = gpu.wait async
+    %memref, %asyncToken = gpu.alloc async [%2] () : memref<2xf16>
+    gpu.wait [%2]
+    affine.store %cst, %memref[0] : memref<2xf16>
+    %3 = gpu.wait async
+    %4 = gpu.memcpy async [%3] %1, %memref : memref<2xf16>, memref<2xf16>
+    gpu.wait [%3]
+    %5 = memref.load %1[%arg1] : memref<2xf16>
+    return %5 : f16
+}
+// CHECK: gpu.memcpy
+
+// We cannot fold gpu.memcpy, as the defining op if dest is not a alloc like op.
+// CHECK-LABEL: func @do_not_fold_memcpy_op3
+func @do_not_fold_memcpy_op3(%arg0: memref<1xi8>, %arg1: memref<i1>) {
+  %0 = arith.constant 0 : index
+  %1 = memref.view %arg0[%0][] : memref<1xi8> to memref<i1>
+  gpu.memcpy  %1, %arg1 : memref<i1>, memref<i1>
+  func.return
+}
+// CHECK: gpu.memcpy
+
 // CHECK-LABEL: @memcpy_after_cast
 func @memcpy_after_cast(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
   // CHECK-NOT: memref.cast
   // CHECK: gpu.memcpy
   %0 = memref.cast %arg0 : memref<10xf32> to memref<?xf32>
   %1 = memref.cast %arg1 : memref<10xf32> to memref<?xf32>
   gpu.memcpy %0, %1 : memref<?xf32>, memref<?xf32>
   return