[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 ↗

(On Diff #426019)

This is unclearly phrased.

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

?

94–104 ↗

(On Diff #426019)

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/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();
+    // If `dest` is not defined by alloc op, we cannot remove `op`.
+    auto isAllocLikeOp = [](Operation *op, Value value) {
+      auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+      if (!memOp)
+        return false;
+      llvm::SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4>
+          memOpEffects;
+      memOp.getEffects(memOpEffects);
+      bool doesOpOnlyAllocateVal = false;
+      for (auto &effect : memOpEffects) {
+        if (effect.getValue() == value &&
+            isa<MemoryEffects::Allocate>(effect.getEffect()))
+          doesOpOnlyAllocateVal = true;
+        if (effect.getValue() == value &&
+            !isa<MemoryEffects::Allocate>(effect.getEffect())) {
+          doesOpOnlyAllocateVal = false;
+          break;
+        }
+      }
+      return doesOpOnlyAllocateVal;
+    };
+    Operation *destDefOp = dest.getDefiningOp();
+    if (!destDefOp || !isAllocLikeOp(destDefOp, dest))
+      return failure();
+    auto isDeallocLikeOpActingOnVal = [](Operation *op, Value val) {
+      auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+      if (!memOp)
+        return false;
+      llvm::SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4>
+          memOpEffects;
+      memOp.getEffects(memOpEffects);
+      return llvm::none_of(memOpEffects, [val](auto &effect) {
+        return effect.getValue() == val &&
+               !isa<MemoryEffects::Free>(effect.getEffect());
+      });
+    };
+    // 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(), [isDeallocLikeOpActingOnVal, op,
+                                       dest](Operation *user) {
+          return user != op && !isDeallocLikeOpActingOnVal(user, 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