[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

1124–1137

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

1136–1138

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

this operation -> op (in backticks)

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

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

1132

Need a code comment here to understand the rationale.

mlir/lib/Interfaces/SideEffectInterfaces.cpp

95

You don't need llvm::.

95

Define this small vector below the check.

102

Need a code comment here.

[arnab-oss]

Addressed comments by @bondhugula.

[bondhugula]

LGTM

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

1132–1134

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/SideEffectInterfaces.h

 /// dereference or read.
 struct Write : public Effect::Base<Write> {};
 } // namespace MemoryEffects
 
 //===----------------------------------------------------------------------===//
 // SideEffect Utilities
 //===----------------------------------------------------------------------===//
 
+/// Returns true if this operation only has the given effect on `value`.
+template <typename EffectTy>
+bool hasSingleEffect(Operation *op, Value value);
+
 /// Return true if the given operation is unused, and has no side effects on
 /// memory that prevent erasing.
 bool isOpTriviallyDead(Operation *op);
 
 /// Return true if the given operation would be dead if unused, and has no side
 /// effects on memory that would prevent erasing. This is equivalent to checking
 /// `isOpTriviallyDead` if `op` was unused.
 bool wouldOpBeTriviallyDead(Operation *op);

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

 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::gpu;
 
 #include "mlir/Dialect/GPU/GPUOpsDialect.cpp.inc"
 
if (getElementTypeOrSelf(srcType) != getElementTypeOrSelf(dstType))
     return emitOpError("arguments have incompatible element type");
 
   if (failed(verifyCompatibleShape(srcType, dstType)))
     return emitOpError("arguments have incompatible shape");
 
   return success();
 }
 
+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();
+    // `dest` must be defined by an op having Allocate memory effect in order to
+    // perform the folding.
+    if (!destDefOp ||
+        !hasSingleEffect<MemoryEffects::Allocate>(destDefOp, 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 &&
+                 !hasSingleEffect<MemoryEffects::Free>(user, dest);
+        }))
+      return failure();
+    // We can perform the folding if and only if op has a single async
+    // dependency and produces an async token as result, or if it does not have
+    // any async dependency and does not produce any async token result.
+    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/lib/Interfaces/SideEffectInterfaces.cpp

while (!effectingOps.empty()) {
     return false;
   }
 
   // If we get here, none of the operations had effects that prevented marking
   // 'op' as dead.
   return true;
 }
 
+template <typename EffectTy>
+bool mlir::hasSingleEffect(Operation *op, Value value) {
+  auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+  if (!memOp)
+    return false;
+  SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
+  memOp.getEffects(effects);
+  bool doesOpOnlyHaveSingleEffectOnVal = false;
+  // Iterate through `effects` and check if and only if effect of type
+  // `EffectTy` is present.
+  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;
+}
+
+template bool mlir::hasSingleEffect<MemoryEffects::Allocate>(Operation *,
+                                                             Value);
+template bool mlir::hasSingleEffect<MemoryEffects::Free>(Operation *, Value);
+template bool mlir::hasSingleEffect<MemoryEffects::Write>(Operation *, Value);
+template bool mlir::hasSingleEffect<MemoryEffects::Read>(Operation *, Value);
+
 bool mlir::wouldOpBeTriviallyDead(Operation *op) {
   if (op->mightHaveTrait<OpTrait::IsTerminator>())
     return false;
   return wouldOpBeTriviallyDeadImpl(op);
 }

mlir/test/Dialect/GPU/canonicalize.mlir

func.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.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.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.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.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.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