[mlir][gpu] Move gpu.wait ops from async.execute regions to its dependencies.

[csigg]

Summary

This can prevent unnecessary host synchronization.

[csigg]

Rebase.

[herhut]

mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp

37

Use op.isKnownTerminator()? But you have to be careful because unknown operations, even if they are terminators, are reported as not being a terminator. So sometimes !op.isKnownNonTerminator is the better choice.

53

Yeah, I think using the !isKnownNonTerminator is the right approach here.

135

How could it be a terminator? Or is this just for reuse?

183

Why drop?

196

You can also do isa<async::ExecuteOp, async::AwaitOp>(user)

208

I am not sure why this is done. If the user of the token is a wait

212

Use .result()?

226

Should this not be the first operation after op, async or not?

[csigg]

Apply reviewer comments and adjust to latest async dialect changes.

[csigg]

mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp

135

I added a comment and split it in two separate functions now.

183

This should have been drop_back(count) to exclude the newly added gpu.async.tokens.

208

Say you have something like this:

%token, %async_gpu_token = async.execute() ...
async.await %token

This adds %gpu_token = async.await %async_gpu_token, and then further down we add gpu.await %gpu_token.

I added a comment.

226

I think the confusion came from that I reused op. I introduced it now.

This function adds gpu.wait between it and the first async/terminator following it.
Before find_if, it either points to the beginning of an async.execute body or just after the block of async.await ops (the first one synchronizing on the token, the others waiting for the !gpu.async.token).

[csigg]

Rest of reviewer comments.

[herhut]

Thanks, much clearer now. I am confused a bit by the types that get used but otherwise this looks great.

mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp

129

This explicitly walks the single block without the terminator. So this would already break if it gets remodeled. That is why I was confused below.

172

Is the getOperation needed here?

173

Does the execute op return an async.token or an async.value<async.token>? I assumed the latter, because then the body of the execute can unwrap it into an async.token or the async.wait can do the unwrapping.

The token is an async value because it is only created during the execution of the parent async region. It could be a stream that gets created in there, no?

180

If you do not map anything here, why the mapper?

213

This is nit picking, but should they not be inserted before? Because originally they also happened before.

[csigg]

Add some additional comments, drop getOperation().

mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp

173

The create semantics of ExecuteOp changed recently and now automatically wraps the result types in async.value<>s.
So yes, the execute op returns async.value<gpu.async.token>s.

180

The interface needs a mapper. Do I need to map anything?

213

If we insert them before, we won't advance below (because the async.await on the token has side-effects) and will never find a gpu async op to pair it up with.

[herhut]

Thanks!

mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp

180

Ah, I see. I assumed it had a default of {} for the mapper, which some other variants of this functionality have.

mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp

 //===- AsyncRegionRewriter.cpp - Implementation of GPU async rewriters ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the GPU dialect pattern rewriters that make GPU op
 // within a region execute asynchronously.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
+#include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/GPU/Utils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/RegionUtils.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 namespace {
 class GpuAsyncRegionPass : public GpuAsyncRegionPassBase<GpuAsyncRegionPass> {
-  struct Callback;
+  struct ThreadTokenCallback;
+  struct DeferWaitCallback;
   void runOnFunction() override;
 };
 } // namespace
 
+static bool isTerminator(Operation &op) {
+  return op.hasTrait<OpTrait::IsTerminator>();
+}
+static bool isTerminatorOrHasSideEffects(Operation &op) {
+  return isTerminator(op) || !MemoryEffectOpInterface::hasNoEffect(&op);
+}
+
 // Region walk callback which makes GPU ops implementing the AsyncOpInterface
 // execute asynchronously.
-struct GpuAsyncRegionPass::Callback {
+struct GpuAsyncRegionPass::ThreadTokenCallback {
+  ThreadTokenCallback(MLIRContext &context) : builder(&context) {}
+
   // If `op` implements the AsyncOpInterface, insert a `gpu.wait async` to
   // create a current token (unless it already exists), and 'thread' that token
   // through the `op` so that it executes asynchronously.
   //
   // If `op` is a terminator or an op with side-effects, insert a `gpu.wait` to
-  // host-synchronize execution.
+  // host-synchronize execution. A `!gpu.async.token` will therefore only be
+  // used inside of its block and GPU execution will always synchronize with
+  // the host at block boundaries.
   WalkResult operator()(Operation *op) {
     if (isa<gpu::LaunchOp>(op))
       return op->emitOpError("replace with gpu.launch_func first");
     if (isa<gpu::WaitOp>(op))
       return op->emitOpError("unexpected pre-existing gpu.wait");
     builder.setInsertionPoint(op);
     if (auto asyncOp = dyn_cast<gpu::AsyncOpInterface>(op))
       return rewriteAsyncOp(asyncOp); // Replace GPU op with async version.
     if (!currentToken)
       return success();
-    if (!op->hasTrait<OpTrait::IsTerminator>() &&
-        MemoryEffectOpInterface::hasNoEffect(op))
-      return success();
     // Insert host synchronization before terminator or op with side effects.
+    if (isTerminatorOrHasSideEffects(*op))
       currentToken = createWaitOp(op->getLoc(), Type(), {currentToken});
     return success();
   }
 
+private:
   // Replaces asyncOp with a clone that returns a token.
   LogicalResult rewriteAsyncOp(gpu::AsyncOpInterface asyncOp) {
     auto *op = asyncOp.getOperation();
     if (asyncOp.getAsyncToken())
       // TODO: Support ops that are already async.
       return op->emitOpError("is already async");
     if (op->getNumRegions() > 0)
       return op->emitOpError("regions are not supported");
private:
   const Type tokenType = builder.getType<gpu::AsyncTokenType>();
   // The token that represents the current asynchronous dependency. It's valid
   // range starts with a `gpu.wait async` op, and ends with a `gpu.wait` op.
   // In between, each gpu::AsyncOpInterface depends on the current token and
   // produces the new one.
   Value currentToken = {};
 };
 
+// Callback for `async.execute` ops which tries to push the contained
+// synchronous `gpu.wait` op to the dependencies of the `async.execute`.
+struct GpuAsyncRegionPass::DeferWaitCallback {
+  // If the `executeOp`s token is used only in `async.execute` or `async.await`
+  // ops, add the region's last `gpu.wait` op to the worklist if it is
+  // synchronous and is the last op with side effects.
+  void operator()(async::ExecuteOp executeOp) {
+    if (!areUsersSupported(executeOp))
+      return;
+    for (auto &op : llvm::reverse(executeOp.getBody()->without_terminator())) {
+      if (auto waitOp = dyn_cast<gpu::WaitOp>(op)) {
+        if (!waitOp.asyncToken())
+          worklist.push_back(waitOp);
+        return;
+      }
+      if (isTerminatorOrHasSideEffects(op))
+        return;
+    }
+  }
+
+  // The destructor performs the actual rewrite work.
+  ~DeferWaitCallback() {
+    for (size_t i = 0; i < worklist.size(); ++i) {
+      auto waitOp = worklist[i];
+      auto executeOp = waitOp.getParentOfType<async::ExecuteOp>();
+      auto numDependencies = waitOp.asyncDependencies().size();
+
+      // Erase `gpu.wait` and return async dependencies from region instead.
+      auto &yieldOp = executeOp.getBody()->getOperations().back();
+      yieldOp.insertOperands(yieldOp.getNumOperands(),
+                             waitOp.asyncDependencies());
+      waitOp.erase();
+      auto asyncTokens = addAsyncTokenResults(executeOp, numDependencies);
+
+      // Add the async dependency to each user of the `async.execute` token.
+      for (Operation *user : executeOp.token().getUsers())
+        addAsyncDependencyAfter(asyncTokens, user);
+    }
+  }
+
+private:
+  // Adds `count` `!async.value<!gpu.async.token>` results to `executeOp`.
+  static ValueRange addAsyncTokenResults(async::ExecuteOp &executeOp,
+                                         unsigned count) {
+    auto numResults = executeOp.getNumResults() + count;
+
+    // Construct new result type list with `count` additional types.
+    SmallVector<Type, 2> resultTypes;
+    resultTypes.reserve(numResults);
+    copy(executeOp.getResultTypes(), std::back_inserter(resultTypes));
+    OpBuilder builder(executeOp.getOperation());
+    auto tokenType = builder.getType<gpu::AsyncTokenType>();
+    resultTypes.resize(numResults, async::ValueType::get(tokenType));
+
+    // Clone executeOp with an extra `!gpu.async.token` results.
+    auto newOp = builder.create<async::ExecuteOp>(
+        executeOp.getLoc(), TypeRange{resultTypes}, executeOp.dependencies(),
+        executeOp.operands());
+    BlockAndValueMapping mapper;
+    executeOp.getRegion().cloneInto(&newOp.getRegion(), mapper);
+
+    // Replace executeOp with cloned one.
+    executeOp.getOperation()->replaceAllUsesWith(
+        newOp.getResults().drop_back());
+    executeOp.erase();
+    executeOp = newOp;
+
+    // Return the new result values.
+    return executeOp.getResults().take_back(count);
+  }
+
+  // Returns whether we support adding async dependencies to the users of
+  // `executeOp`. Specifically, we do not support terminator users which could
+  // mean that the `executeOp` is inside control flow code.
+  static bool areUsersSupported(async::ExecuteOp executeOp) {
+    return llvm::all_of(executeOp.token().getUsers(), [](Operation *user) {
+      return isa<async::ExecuteOp>(user) || isa<async::AwaitOp>(user);
+    });
+  }
+
+  // Add the `asyncToken` as dependency as needed after `op`.
+  void addAsyncDependencyAfter(ValueRange asyncTokens, Operation *op) {
+    OpBuilder builder(op->getContext());
+    auto loc = op->getLoc();
+
+    SmallVector<Value, 1> tokens;
+    tokens.reserve(asyncTokens.size());
+    TypeSwitch<Operation *>(op)
+        .Case<async::AwaitOp>([&](auto awaitOp) {
+          builder.setInsertionPointAfter(op);
+          for (auto asyncToken : asyncTokens)
+            tokens.push_back(
+                builder.create<async::AwaitOp>(loc, asyncToken).getResult(0));
+        })
+        .Case<async::ExecuteOp>([&](auto executeOp) {
+          // Set `op` to the beginning of the region and add asyncTokens to the
+          // async.execute operands.
+          op = &executeOp.getBody()->front();
+          executeOp.operandsMutable().append(asyncTokens);
+          SmallVector<Type, 1> tokenTypes(
+              asyncTokens.size(), builder.getType<gpu::AsyncTokenType>());
+          copy(executeOp.getBody()->addArguments(tokenTypes),
+               std::back_inserter(tokens));
+        });
+
+    // Find terminator or op with side-effects after `op`.
+    op = &*std::find_if(Block::iterator(op), Block::iterator(),
+                        isTerminatorOrHasSideEffects);
+
+    // If `op` implements the AsyncOpInterface, add `token` to the list of async
+    // dependencies.
+    if (auto asyncOp = dyn_cast<gpu::AsyncOpInterface>(op)) {
+      for (auto token : tokens)
+        asyncOp.addAsyncDependency(token);
+      return;
+    }
+
+    // Otherwise, insert a gpu.wait before 'op'.
+    builder.setInsertionPoint(op);
+    auto waitOp = builder.create<gpu::WaitOp>(loc, Type{}, tokens);
+
+    // If the new waitOp is at the end of an async.execute region, add it to the
+    // worklist.
+    auto executeOp = dyn_cast<async::ExecuteOp>(op->getParentOp());
+    if (executeOp && areUsersSupported(executeOp) && isTerminator(*op))
+      worklist.push_back(waitOp);
+  }
+
+  SmallVector<gpu::WaitOp, 8> worklist;
+};
+
 // Replaces synchronous GPU ops in the op's region with asynchronous ones and
 // inserts the necessary synchronization (as gpu.wait ops). Assumes sequential
 // execution semantics and that no GPU ops are asynchronous yet.
 void GpuAsyncRegionPass::runOnFunction() {
   if (getFunction()
           .getRegion()
-          .walk(Callback{OpBuilder(&getContext())})
+          .walk(ThreadTokenCallback(getContext()))
           .wasInterrupted())
     return signalPassFailure();
+
+  // Collect gpu.wait ops that we can move out of gpu.execute regions.
+  getFunction().getRegion().walk(DeferWaitCallback());
 }
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createGpuAsyncRegionPass() {
   return std::make_unique<GpuAsyncRegionPass>();
 }

mlir/test/Dialect/GPU/async-region.mlir

 // RUN: mlir-opt -gpu-async-region %s | FileCheck %s
 
 // CHECK: module attributes {gpu.container_module}
 module attributes {gpu.container_module} {
 
   gpu.module @kernels {
     gpu.func @kernel() kernel { gpu.return }
   }
 
   func @foo() -> ()
 
-  // CHECK-LABEL:func @async(%{{.*}}: index)
-  func @async(%sz : index) {
+  // CHECK-LABEL:func @thread_token(%{{.*}}: index)
+  func @thread_token(%sz : index) {
     // CHECK: %[[t0:.*]] = gpu.wait async
     // CHECK: %[[t1:.*]] = gpu.launch_func async [%[[t0]]]
     gpu.launch_func @kernels::@kernel
         blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
     // CHECK: %[[t2:.*]] = gpu.launch_func async [%[[t1]]]
     gpu.launch_func @kernels::@kernel
         blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
     // CHECK: gpu.wait [%[[t2]]]
     // CHECK: call @foo
     call @foo() : () -> ()
     return
   }
 
+  // CHECK-LABEL:func @defer_wait(%{{.*}}: index)
+  func @defer_wait(%sz : index) {
+    // CHECK: %[[a0:.*]], %[[f0:.*]] = async.execute
+    %a0 = async.execute {
+      // CHECK: %[[t:.*]] = gpu.launch_func async
+      gpu.launch_func @kernels::@kernel
+          blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
+      // CHECK-NOT: gpu.wait
+      // CHECK: async.yield %[[t]]
+      async.yield
+    }
+
+    // CHECK: %[[a1:.*]], %[[f1:.*]] = async.execute
+    // CHECK-SAME: %[[f0]]
+    %a1 = async.execute [%a0] {
+      // CHECK: %[[t:.*]] = gpu.launch_func async
+      gpu.launch_func @kernels::@kernel
+          blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
+      // CHECK-NOT: gpu.wait
+      // CHECK: async.yield %[[t]]
+      async.yield
+    }
+
+    // CHECK: async.await %[[a1]]
+    // CHECK: %[[t:.*]] = async.await %[[f1]]
+    // CHECK: gpu.wait [%[[t]]]
+    async.await %a1 : !async.token
+    return
+  }
+
+  // CHECK-LABEL:func @defer_wait_blocked_by_side_effect(%{{.*}}: index)
+  func @defer_wait_blocked_by_side_effect(%sz : index) {
+    // CHECK: %[[a:.*]] = async.execute
+    %a = async.execute {
+      // CHECK: %[[t:.*]] = gpu.launch_func async
+      gpu.launch_func @kernels::@kernel
+          blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
+      // CHECK: gpu.wait [%[[t]]]
+      call @foo() : () -> ()
+      async.yield
+    }
+
+    // CHECK: async.await %[[a]]
+    // CHECK-NOT: gpu.wait
+    async.await %a : !async.token
+    return
+  }
+
+  // CHECK-LABEL:func @defer_wait_pass_through(%{{.*}}: index)
+  func @defer_wait_pass_through(%sz : index) {
+    // CHECK: %[[a0:.*]], %[[f0:.*]] = async.execute
+    %a0 = async.execute {
+      // CHECK: %[[t:.*]] = gpu.launch_func async
+      gpu.launch_func @kernels::@kernel
+          blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
+      // CHECK-NOT: gpu.wait
+      // CHECK: async.yield %[[t]]
+      async.yield
+    }
+
+    // CHECK: %[[a1:.*]], %[[f1:.*]] = async.execute
+    // CHECK-SAME: %[[f0]]
+    %a1 = async.execute [%a0] {
+      // CHECK-NOT: gpu.wait
+      // CHECK: async.yield %{{.*}}
+      async.yield
+    }
+
+    // CHECK: async.await %[[a1]]
+    // CHECK: %[[t:.*]] = async.await %[[f1]]
+    // CHECK: gpu.wait [%[[t]]]
+    async.await %a1 : !async.token
+    return
+  }
 }