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

1326

Nit: Fold -> Erase?

1332

Typo here.

1339

does not ... a single use

1339

ops or single op that you are handling here?

1354

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

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 canoot remove `op`.
+    if (dest.isa<BlockArgument>())
+      return failure();
+    auto isDeallocLikeOp = [](Operation *op) {
+      auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+      return memOp && memOp.hasEffect<MemoryEffects::Free>();
+    };
+    // We can erase `op` iff `dest` has no othet use apart from its
+    // use by `op` and dealloc ops.
+    if (llvm::any_of(dest.getUsers(), [isDeallocLikeOp, op](Operation *user) {
+          return user != op && !isDeallocLikeOp(user);
+        }))
+      return failure();
+
+    ValueRange asyncDependencies = op.asyncDependencies();
+    // Check that the async token we are going to erase has no other uses.
+    if (!op.asyncToken() || op.asyncToken().use_empty())
+      rewriter.eraseOp(op);
+
+    // Remove redundant gpu.wait op. If `op` has a single async dependency
+    // token, and the token value has a single user (other than `op`, of type
+    // gpu.wait, we can erase the gpu.wait op, along with the op defining the
+    // async token.
+    if (asyncDependencies.size() == 1 && asyncDependencies[0].hasOneUse()) {
+      if (auto waitOp = dyn_cast<WaitOp>(*asyncDependencies[0].user_begin())) {
+        if (!waitOp.asyncToken() && waitOp.asyncDependencies().size() == 1) {
+          rewriter.eraseOp(waitOp);
+          rewriter.eraseOp(asyncDependencies[0].getDefiningOp());
+        }
+      }
+    }
+    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>();

 
 #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 3.343820e-05 : f16
+    %cst_0 = arith.constant 0.000000e+00 : f16
+    %1 = memref.alloc() : memref<400x1024x1024x1xf16>
+    %2 = gpu.wait async 
+    %memref, %asyncToken = gpu.alloc async [%2] () : memref<400x1024x1024x1xf16>
+    gpu.wait [%2]
+    affine.store %cst, %memref[0, 0, 0, 0] : memref<400x1024x1024x1xf16>
+    %3 = gpu.wait async 
+    %4 = gpu.memcpy async [%3] %1, %memref : memref<400x1024x1024x1xf16>, memref<400x1024x1024x1xf16>
+    gpu.wait [%3]
+    %5 = scf.if %arg0 -> (i1) {
+      memref.dealloc %1 : memref<400x1024x1024x1xf16>
+      scf.yield %arg0 : i1
+    } else {
+      memref.dealloc %1 : memref<400x1024x1024x1xf16>
+      scf.yield %arg0 : i1
+    }  
+    return
+}
+// CHECK-NOT: 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
 }
 
 // CHECK-LABEL: @memset_after_cast
 func @memset_after_cast(%arg0: memref<10xf32>, %arg1: f32) {