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

[arnab-oss]

Summary

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

[arnab-oss]

Set hasVerifier field in gpu.memcpy removed earlier by mistake.

[csigg]

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

1077

Nit: canoot -> cannot

1084

Nit: othet -> other

1096

Shouldn't this rather be a separate canonicalizer that removes %t = gpu.wait async + gpu.wait [%t] pairs?

mlir/test/Dialect/GPU/canonicalize.mlir

14

Wouldn't this more natually be %4, and the canonicalizer should rewriter.replaceOp(op, asyncDependencies[0])?

[arnab-oss]

Addressed review comments by @csigg.

[mehdi_amini]

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

1082

Strictly speaking it should also check that it does not have a read effect I think (it is unlikely but an operation could be "print_and_free(%mem)".

Also isn't there an API to check the effect per operand? What is an operation has two operands but frees one of them?

(One easy way out would be to match the GpuDealloc op specifically intead, or check that this op has a single operand)

Also, in general this transformation seems like "dead store elimination", isn't it something that can be implementation generically (dialect independent, purely effect based)?

[bondhugula]

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

1082

so, in general this transformation seems like "dead store elimination", isn't it something
that can be implementation generically (dialect independent, purely effect based)?

Reg. this last comment: this issue came up in the past too (one of my revisions where we had a discussion). While there is a general pattern here and semblance of a reuse, a dead store elimination pass can't by itself be a substitute for a folding hook here because:

1. a separate dse pass won't lead to early simplification like via canonicalization and folding hooks,
2. one will run into phase ordering issues b/w other folding/canonicalizations and such dse; so it's ideal to have such O(1) and O(use-def) chain length stuff in folding hook.

Moreover, in this case, there is a connection to the wait op and aysnc token here; so I think we do need the folding hook.

1082

+1 to what Mehdi says on associating the effect to the operand. But I think there is API missing here and one would need to get the value associated with each Free effect to see if it is dest. This check is currently conservative.

[bondhugula]

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

1082

Strictly speaking it should also check that it does not have a read effect I think (it is
unlikely but an operation could be "print_and_free(%mem)".

That's right. The way it is, this pattern's check is incorrect. It needs to check that Free is the only effect that this op can have on dest. However, the pattern is checking if the op has a Free effect on some Value.

[arnab-oss]

Addressed comments by @bondhugula

[bondhugula]

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

1085

You don't need mlir:: - you already have a using for that.

1321

Nit: Fold -> Erase?

1327

Typo here.

1334

does not ... a single use

1334

ops or single op that you are handling here?

1349

eraseOp(waitOp)

mlir/test/Dialect/GPU/canonicalize.mlir

24

Missing check that the gpu.wait has been folded away.

[arnab-oss]

Addressed comments by @bondhugula

[bondhugula]

LGTM - thanks.

mlir/test/Dialect/GPU/canonicalize.mlir

29

Nit: Punctuate correctly.

[arnab-oss]

Addressed comments by @bondhugula

[csigg]

I think the gpu.wait canonicalizer could be cleaning up more cases:

• %unused gpu.wait async ...: eraseOp(op)
• gpu.wait []: eraseOp(op)
• %t1 = gpu.wait async [%t0]: replaceOp(op, t0)
• %t = gpu.wait async ... + gpu.wait {async} [%t, ...]: drop %t from async dependencies.

[bondhugula]

In D121279#3378709, @csigg wrote:

I think the gpu.wait canonicalizer could be cleaning up more cases:

• %unused gpu.wait async ...: eraseOp(op)
• gpu.wait []: eraseOp(op)
• %t1 = gpu.wait async [%t0]: replaceOp(op, t0)
• %t = gpu.wait async ... + gpu.wait {async} [%t, ...]: drop %t from async dependencies.

Would this fit logically in this revision or a separate revision for gpu.wait canonicalizer? This revision is meant to erase away trivial gpu.memcpy and ancillary stuff. A full-fledged gpu.wait folder/canonicalizer should ideally go into a separate commit.

[csigg]

In D121279#3381610, @bondhugula wrote:

Would this fit logically in this revision or a separate revision for gpu.wait canonicalizer? This revision is meant to erase away trivial gpu.memcpy and ancillary stuff. A full-fledged gpu.wait folder/canonicalizer should ideally go into a separate commit.

Sorry, I should have been more clear, none of my comments were meant to gate the revision.

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

1077

Why can the memcpy not be removed if it's a block argument?

1080

Would 'OnVal' be better than 'OnDest'?

1101

Should this also handle the common case of a tokens threaded through all gpu ops?

if (op.asyncDependencies().size() > 1 ||
    op.asyncDependencies().empty() == op.asyncToken())
  return failure()
rewriter.replaceOp(op, op.asyncDependencies());

This would also take care of not updating the op but still returning success().

[arnab-oss]

Addressed comments by @csigg.

[csigg]

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

1100–1103

Really just an idea, would this be easier to read?

[arnab-oss]

Addressed review comments and rebased on latest master.

[arnab-oss]

Fixed build issues.

[arnab-oss]

Removed trailing white spaces.

[MaskRay]

I am sorry but we have clear evidence this miscompiles some tensorflow gpu tests, so I am applying a revert. I will check with them to provide a reproduce.

[arnab-oss]

Can you please provide the test case for which incorrect IR is generated?
I will provide the fix. Thanks in advance.

[jpienaar]

In D121279#3464917, @arnab-oss wrote:

Can you please provide the test case for which incorrect IR is generated?
I will provide the fix. Thanks in advance.

Folks'll try getting self-contained repro, but in interim test case that fails with this but passes without is bazel test --config=cuda --compilation_mode=opt --test_env=XLA_FLAGS=--xla_gpu_bef_executable --copt=-Wno-error //tensorflow/compiler/xla/tests:concat_test_gpu in TF repo. You'd need to probably follow the instructions for local llvm repo in tensorflow/compiler/mlir for simplicity in repro

[hanbinyoon]

You also need --//tensorflow/compiler/xla/service/gpu:enable_xlir=true
bazel test --config=cuda --compilation_mode=opt --//tensorflow/compiler/xla/service/gpu:enable_xlir=true --test_env=XLA_FLAGS=--xla_gpu_bef_executable --copt=-Wno-error //tensorflow/compiler/xla/tests:concat_test_gpu

[csigg]

Here is a repro:

func @copy(%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
}

mlir-opt --canonicalize removes the memcpy when it really shouldn't.

[bondhugula]

In D121279#3466770, @csigg wrote:

Here is a repro:

func @copy(%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
}

mlir-opt --canonicalize removes the memcpy when it really shouldn't.

This is a big miss. The pattern will have to be made very conservative or instead be implemented as part of a copy-removal pass where aliasing information can be used. However, strictly speaking, the aliasing check available with AliasAnalysis is often just O(use-def) chain -- so one should be able to use something like bool doMemRefsAlias(Value memRefA, Value memRefB) from a canonicalization pattern or a folding hook. There isn't really any caching needed or caching happening there.

[bondhugula]

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

1077

This is the condition that's completely arbitrary. You can change this check to iterate through all uses and check if alloc/dealloc and this memcpy are the only ops that use this op? A more powerful removal has to be done anyway in a separate pass.

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

let description = [{
     // The gpu.wait op blocks until gpu.foo and gpu.bar have completed.
     gpu.wait [%t0, %t1]
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
+  let hasCanonicalizer = 1;
+
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) attr-dict
   }];
 }
 
 def GPU_AllocOp : GPU_Op<"alloc", [
     GPU_AsyncOpInterface,
     AttrSizedOperandSegments
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 a block argument, we cannot remove `op`.
+    if (dest.isa<BlockArgument>())
+      return failure();
+    auto isDeallocLikeOpActingOnDest = [](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(), [isDeallocLikeOpActingOnDest, op,
+                                       dest](Operation *user) {
+          return user != op && !isDeallocLikeOpActingOnDest(user, dest);
+        }))
+      return failure();
+
+    // Check that the async token we are going to erase has no other uses.
+    if (!op.asyncToken() || op.asyncToken().use_empty())
+      rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+} // end anonymous namespace
+
+void MemcpyOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                           MLIRContext *context) {
+  results.add<EraseTrivialCopyOp>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // GPU_SubgroupMmaLoadMatrixOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult SubgroupMmaLoadMatrixOp::verify() {
   auto srcType = srcMemref().getType();
   auto resType = res().getType();
   auto resMatrixType = resType.cast<gpu::MMAMatrixType>();
if (size_t(srcMemref.getRank()) != srcIndices().size())
     return emitOpError() << "expected " << srcMemref.getRank()
                          << " source indices, got " << srcIndices().size();
   if (size_t(dstMemref.getRank()) != dstIndices().size())
     return emitOpError() << "expected " << dstMemref.getRank()
                          << " destination indices, got " << dstIndices().size();
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// GPU_WaitOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Fold away redundant gpu.wait ops of the following pattern.
+/// %t = gpu.wait async
+/// gpu.wait [%t]
+struct EraseRedundantGpuWait : public OpRewritePattern<WaitOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WaitOp op,
+                                PatternRewriter &rewriter) const final {
+    // We check whether `op` produce result `asyncToken`.
+    Value token = op.asyncToken();
+    if (!token)
+      return failure();
+    // We check whether `op` has any async dependencies or not.
+    if (!op.asyncDependencies().empty())
+      return failure();
+    // If token do not have single use, we cannot fold away gpu.wait ops.
+    if (!token.hasOneUse())
+      return failure();
+    // If the only op operating on `token` is not a gpu.wait op, we cannot fold
+    // away gpu.wait ops.
+    auto waitOp = dyn_cast<mlir::gpu::WaitOp>(*token.user_begin());
+    if (!waitOp)
+      return failure();
+    // If `waitOp` produces any token, we cannot fold away the gpu.wait ops.
+    if (waitOp.asyncToken())
+      return failure();
+    // `waitOp` should have only single async dependency.
+    if (!llvm::hasSingleElement(waitOp.asyncDependencies()))
+      return failure();
+
+    rewriter.eraseOp(waitOp);
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+} // end anonymous namespace
+
+void WaitOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                         MLIRContext *context) {
+  results.add<EraseRedundantGpuWait>(context);
+}
+
 #include "mlir/Dialect/GPU/GPUOpInterfaces.cpp.inc"
 #include "mlir/Dialect/GPU/GPUOpsEnums.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/GPU/GPUOpsAttributes.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"

mlir/test/Dialect/GPU/canonicalize.mlir

 // RUN: mlir-opt %s -canonicalize --split-input-file -allow-unregistered-dialect | FileCheck %s
 
+// 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: gpu.wait
+// CHECK: gpu.wait
+// CHECK-NOT: gpu.wait
+// CHECK-NOT: gpu.memcpy
+// CHECK-NOT: gpu.wait
+
+// We cannot fold memcpy here 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.wait
+// CHECK: gpu.wait
+// CHECK: gpu.wait
+// CHECK: gpu.memcpy
+// CHECK: gpu.wait
+
+// 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.wait
+// CHECK: gpu.wait
+// CHECK: gpu.wait
+// CHECK: gpu.memcpy
+// CHECK: gpu.wait
+
 // 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