[MLIR][GPU] Add canonicalization patterns for folding simple gpu.wait ops.

[arnab-oss]

Summary

• Fold away redundant %t = gpu.wait async + gpu.wait [%t] pairs.

• Fold away %t = gpu.wait async ... ops when %t has no uses.

• Fold away gpu.wait [] ops.

• In case of %t1 = gpu.wait async [%t0], replace all uses of %t1 with %t0.

[arnab-oss]

Resolved build issues.

[csigg]

Thanks for taking care of this!

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

1179–1214

Untested, but would it work to simply remove the async dependency and let the other canonicalizer take care of erasing ops?

[arnab-oss]

ping @csigg, can you please take look at my comment?

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

1179–1214

%0 = gpu.wait async
 %memref, %asyncToken = gpu.alloc async [%0] () : memref<5xf16>
gpu.wait [%0]

In the above example, we cannot exactly remove the %0 async dependency from the 2nd gpu.wait right @csigg ?

[bondhugula]

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

1179–1214

@arnab-oss Does %memref have any uses in this example? If it doesn't, this is all dead code. Is the rule eliminating these now? Can you add this test case?

[bondhugula]

mlir/test/Dialect/GPU/canonicalize.mlir

12

Nit: You can just instead use:

CHECK-NEXT: return

for a stronger and more direct check.

[bondhugula]

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

1234–1244

These conditions are all completely wrong: you can't erase a gpu.wait op just because its result token has no uses. Also, you can't simply replace its uses (which could be empty) by its aysnc dependencies and replace the op.

Can you please read through the gpu.wait doc description (which is pretty straightforward and simple) before adding these patterns?

[mehdi_amini]

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

1216

Please remove.

1234–1244

I'm confused, the description seems to exactly call for this:

If the op contains the async keyword, it returns a new async token which
is synchronized with the op arguments. This new token is merely a shortcut
to the argument list, and one could replace the uses of the result with the
arguments for the same effect.

[bondhugula]

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

1219

In fact, 1) is just a trivial special case of 3), but anyway, both (1) and (3) are wrong. You can't erase the gpu.wait op if its token has no uses. The op synchronizes the device/host (internally, it destroys the stream) You can only replace it if its result token has at least one use that is another gpu.wait op.

[bondhugula]

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

1234–1244

If none of those result uses are another gpu.wait op, you've eliminated a device synchronization incorrectly.

[bondhugula]

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

1234–1244

Actually, the documentation is messed up here, it says:

If the op does not contain the `async` keyword, it does not return a new
    async token but blocks until all ops producing the async dependency tokens
    finished execution.

But the op can definitely return a new token even without the async keyword. That's how a fresh token is created in the first place!

[mehdi_amini]

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

1234–1244

How so? What is specific to gpu.wait in how it is using a token that does not apply to the other users?
Also: what kind of synchronization does an asynchronous gpu.wait without users providing?

(if you can also explain how you interpret what the doc says).

[bondhugula]

Dismissing change request.

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

1234–1244

Sorry, looks like I was mistaken on this part (whenever you return an async token and also have async dependencies).

[bondhugula]

LGTM - mostly doc/code comment fix suggestions.

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

1218

1. %t = gpu.wait ... ops where %t has no uses (regardless of async dependencies).

1218–1220

Nit: The order of these doesn't match the order in the implementation. Permute.

1219

gpu.wait ops that neither have any async dependencies nor return any token.

1228–1229

Erase gpu.wait ops that neither have any async dependencies nor return any async token.

1234–1235

Replace uses of %t1 = gpu.wait async [%t0] ops with %t0 and erase the op.

[mehdi_amini]

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

1239

Ideally this should be implemented as a folder I think.

[bondhugula]

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

1239

Yes, I think this pattern can be the folding hook since we aren't erasing any other ops.

[mehdi_amini]

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

1239

Unfortunately I think the folding hook can't implement conditional deletion, only replacements. Can it?
Otherwise the eraseOp() can't be moved there.

[csigg]

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

1179–1214

%0 = gpu.wait async
%memref, %asyncToken = gpu.alloc async [%0] () : memref<5xf16>
gpu.wait [%0]

In the above example, we cannot exactly remove the %0 async dependency from the 2nd gpu.wait right @csigg ?

I think we can. The second gpu.wait does not synchronize with anything. The order of the gpu.alloc and second gpu.wait can be swapped, unless gpu.wait also uses %asyncToken (which would still allow you to remove %0).

[bondhugula]

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

1179–1214

But if you remove the gpu.wait, you've left a stream "undestroyed". Also, we shouldn't be converting an async alloc to a sync alloc -- the latter won't even lower through gpu-to-llvm - so I'm confused here.

[csigg]

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

1179–1214

But if you remove the gpu.wait, you've left a stream "undestroyed".

I assume there is a gpu.wait [<dependency of %asyncToken>] somewhere. Otherwise the input is not really valid.

Also, we shouldn't be converting an async alloc to a sync alloc

That's not what's happening.

%0 = gpu.wait async
%memref, %asyncToken = gpu.alloc async [%0] () : memref<5xf16>
gpu.wait [%0]

should be folded to

%0 = gpu.wait async
%memref, %asyncToken = gpu.alloc async [%0] () : memref<5xf16>
gpu.wait []

should be folded to

%0 = gpu.wait async
%memref, %asyncToken = gpu.alloc async [%0] () : memref<5xf16>

[bondhugula]

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

1179–1214

I think I see where the difference in understanding is. I was going by how the lowering functions worked (to CUDA streams) as opposed to the op's semantics. In the snippet below, asyncToken doesn't need to have a use later: when the wait on %0 happens, the stream is synchronized and destroyed: this means the memcpy will be completed (even though %asyncToken isn't specified as a dep on it) and anything else attached to that stream. And one doesn't need another use of %asyncToken -- this IR currently lowers and executes correctly as intended AFAIU. In fact, deleting the second gpu.wait would eliminate a synchronization and violate semantics.

%0 = gpu.wait async
%memref, %asyncToken = gpu.alloc async [%0] () : memref<5xf16>
gpu.wait [%0]

However, per the doc, the second gpu.wait would really "block" on the completion of the first op which is a trivial block and so can be eliminated.

[bondhugula]

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

1179–1214

@csigg Furthermore, based on what you said, the following would be valid IR:

%0 = gpu.wait async
%memref, %asyncToken = gpu.alloc async [%0] () : memref<5xf16>
gpu.wait [%0]
gpu.wait %asyncToken

But won't this crash on lowering today since it would destroy the same stream twice?
https://github.com/llvm/llvm-project/blob/32f3633171aa9d7352e9507c12d219efb48899a0/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp#L560

[mehdi_amini]

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

1179–1214

The notion of stream being destroyed in a wait or being left undestroyed is a bit confusing to me, and I can't connect this to the concept of token and async dependencies. Seems like some lowering issues to me more than anything to fix at this level of abstraction.

[mehdi_amini]

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

1197

Why does multiple uses prevent optimizing it away?

1208

Seems like instead of handling by looking at the user, we should root the pattern on the user and look for the operands.
That way it would handle the "multiple uses" above, and it could operate by eliminating operands from the operands list of the gpu.wait.
An async gpu.wait without users could be just deleted.

[bondhugula]

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

1179–1214

I agree - this all makes sense to me and I think @csigg's change can safely be applied here.

1234–1244

A gpu.wait that returns a token doesn't provide any synchronization. We are talking about a gpu.wait that doesn't return an async token but has async deps. This as you know will block till those producers are complete. If one of those producers is a gpu.wait, we can erase it -- this is fine per the semantics of the op and this is why @csigg's suggestion above can be applied.

[bondhugula]

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

1208

An async gpu.wait without users could be just deleted.

The pattern further below already does this: see line 1218.

[arnab-oss]

Adressed comments by @csigg, @bondhugula

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

def GPU_WaitOp : GPU_Op<"wait", [GPU_AsyncOpInterface]> {
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) attr-dict
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def GPU_AllocOp : GPU_Op<"alloc", [
     GPU_AsyncOpInterface,
     AttrSizedOperandSegments
   ]> {
 
   let summary = "GPU memory allocation operation.";

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

 }
 
 LogicalResult MemsetOp::fold(ArrayRef<Attribute> operands,
                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
+// GPU_WaitOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Remove gpu.wait op use of gpu.wait op def without async dependencies.
+/// %t = gpu.wait async []       // No async dependencies.
+/// ...  gpu.wait ... [%t, ...]  // %t can be removed.
+struct EraseRedundantGpuWaitOpPairs : public OpRewritePattern<WaitOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WaitOp op,
+                                PatternRewriter &rewriter) const final {
+    auto predicate = [](Value value) {
+      auto wait_op = value.getDefiningOp<WaitOp>();
+      return wait_op && wait_op->getNumOperands() == 0;
+    };
+    if (llvm::none_of(op.asyncDependencies(), predicate))
+      return failure();
+    SmallVector<Value> validOperands;
+    for (Value operand : op->getOperands()) {
+      if (predicate(operand))
+        continue;
+      validOperands.push_back(operand);
+    }
+    op->setOperands(validOperands);
+    return success();
+  }
+};
+
+/// Simplify trivial gpu.wait ops for the following patterns.
+/// 1. %t = gpu.wait async ... ops, where %t has no uses (regardless of async
+/// dependencies).
+/// 2. %t1 = gpu.wait async [%t0], in this case, we can replace uses of %t1 with
+/// %t0.
+/// 3. gpu.wait [] ops, i.e gpu.wait ops that neither have any async
+/// dependencies nor return any token.
+struct SimplifyGpuWaitOp : public OpRewritePattern<WaitOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WaitOp op,
+                                PatternRewriter &rewriter) const final {
+    // Erase gpu.wait ops that neither have any async dependencies nor return
+    // any async token.
+    if (op.asyncDependencies().empty() && !op.asyncToken()) {
+      rewriter.eraseOp(op);
+      return success();
+    }
+    // Replace uses of %t1 = gpu.wait async [%t0] ops with %t0 and erase the op.
+    if (llvm::hasSingleElement(op.asyncDependencies()) && op.asyncToken()) {
+      rewriter.replaceOp(op, op.asyncDependencies());
+      return success();
+    }
+    // Erase %t = gpu.wait async ... ops, where %t has no uses.
+    if (op.asyncToken() && op.asyncToken().use_empty()) {
+      rewriter.eraseOp(op);
+      return success();
+    }
+    return failure();
+  }
+};
+
+} // end anonymous namespace
+
+void WaitOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                         MLIRContext *context) {
+  results.add<EraseRedundantGpuWaitOpPairs, SimplifyGpuWaitOp>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // GPU_AllocOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(AllocOp op) {
   auto memRefType = op.memref().getType().cast<MemRefType>();
 
   if (static_cast<int64_t>(op.dynamicSizes().size()) !=
       memRefType.getNumDynamicDims())

mlir/test/Dialect/GPU/canonicalize.mlir

 // RUN: mlir-opt %s -canonicalize --split-input-file -allow-unregistered-dialect | FileCheck %s
 
+// Fold all the gpu.wait ops as they are redundant.
+// CHECK-LABEL: func @fold_wait_op_test1
+func @fold_wait_op_test1() {
+  %1 = gpu.wait async
+  gpu.wait []
+  %3 = gpu.wait async 
+  gpu.wait [%3]
+  return
+}
+// CHECK-NOT: gpu.wait
+
+// Replace uses of gpu.wait op with its async dependency.
+// CHECK-LABEL: func @fold_wait_op_test2
+func @fold_wait_op_test2(%arg0: i1) -> (memref<5xf16>, memref<5xf16>) {
+  %0 = gpu.wait async
+  %memref, %asyncToken = gpu.alloc async [%0] () : memref<5xf16>
+  gpu.wait [%0]
+  %1 = gpu.wait async [%0]
+  %memref_0, %asyncToken_0 = gpu.alloc async [%1] () : memref<5xf16>
+  gpu.wait [%1]
+  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: @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