[mlir][gpu] Add pass to make GPU ops within a region execute asynchronously.

[csigg]

Summary

Do not use the pass yet, except in a test.

[herhut]

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

112

I find it somewhat hard to read that this is a carried state. A better name or comment would already help a lot.

134

Is it useful to formulate this as a pattern? It could equally well be just a walk over all operations, as there is no matching going on, at all.

[csigg]

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

134

I'm using notifyMatchFailure, and maybe/theoretically there could be different patterns in the same pass. But to be honest I mainly used RewritePattern because I vaguely remember that River didn't like the raw walk in the all-reduce lowering. If you think a raw walk is better here, I'm happy to change it.

[rriddle]

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

134

I don't exactly remember the details there. A normal walk is fine, but I tend to err more towards patterns because they compose better/have better logging/remove the need to write transformation application logic when things are more complicated than a simple walk. Not to say that you shouldn't just use a simple walk here, but I'd say it depends on how this pass is intended to evolve.

[mehdi_amini]

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

134

Because there is state in the pattern keeping a value, this is to me a strong indication that we should not use a pattern here. This is just hiding the traversal logic, while it seems core to the algorithm.

[csigg]

Use raw walk instead of a RewritePattern.

This does seem a little cleaner, at least at the moment:

• It's easier to make the pass fail when some preconditions aren't met.
• The rewriter.replaceOp() assumes that we replace the root, which is not what we did.

[sanjoy.google]

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

39

The more common term for this is to "thread the token through", WDYT?

80

Is this correct if the region has cycles or diamond control flow?

[csigg]

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

39

Thanks for the suggestion. Yes, that sound better. I will include that change in my next update (it's always a bit of a dance, so I'm waiting if there are more requested changes).

80

I thought control flow requires separate blocks, each with their own terminator (which triggers a host synchronization and clears the currentToken). Is that not correct?

134

I've removed the pattern, but it did not contain state. The Callback containing the state was instantiated inside the matchAndRewrite().

[csigg]

Error out on encountering pre-existing gpu.wait.

[sanjoy.google]

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

80

You're right, I didn't notice that bit.

[herhut]

Thanks!

[csigg]

Rebase.

[csigg]

Rebase.

[mehdi_amini]

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

82

Is there a reason why this is using Operation::create to clone the operation instead of one of the clone() routine or could we just switch it?

mlir/include/mlir/Dialect/GPU/Passes.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_GPU_PASSES_H_
 #define MLIR_DIALECT_GPU_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+/// Replaces `gpu.launch` with `gpu.launch_func` by moving the region into
+/// a separate kernel function.
 std::unique_ptr<OperationPass<ModuleOp>> createGpuKernelOutliningPass();
 
+/// Rewrites a function region so that GPU ops execute asynchronously.
+std::unique_ptr<OperationPass<FuncOp>> createGpuAsyncRegionPass();
+
 /// Collect a set of patterns to rewrite all-reduce ops within the GPU dialect.
 void populateGpuAllReducePatterns(MLIRContext *context,
                                   OwningRewritePatternList &patterns);
 
 /// Collect all patterns to rewrite ops within the GPU dialect.
 inline void populateGpuRewritePatterns(MLIRContext *context,
                                        OwningRewritePatternList &patterns) {
   populateGpuAllReducePatterns(context, patterns);

mlir/include/mlir/Dialect/GPU/Passes.td

 
 include "mlir/Pass/PassBase.td"
 
 def GpuKernelOutlining : Pass<"gpu-kernel-outlining", "ModuleOp"> {
   let summary = "Outline gpu.launch bodies to kernel functions";
   let constructor = "mlir::createGpuKernelOutliningPass()";
 }
 
+def GpuAsyncRegionPass : FunctionPass<"gpu-async-region"> {
+  let summary = "Make GPU ops async";
+  let constructor = "mlir::createGpuAsyncRegionPass()";
+}
+
 #endif // MLIR_DIALECT_GPU_PASSES

mlir/lib/Dialect/GPU/CMakeLists.txt

 add_mlir_dialect_library(MLIRGPU
   IR/GPUDialect.cpp
   Transforms/AllReduceLowering.cpp
+  Transforms/AsyncRegionRewriter.cpp
   Transforms/KernelOutlining.cpp
   Transforms/MemoryPromotion.cpp
   Transforms/ParallelLoopMapper.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GPU
 
   DEPENDS

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/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"
+
+using namespace mlir;
+namespace {
+class GpuAsyncRegionPass : public GpuAsyncRegionPassBase<GpuAsyncRegionPass> {
+  struct Callback;
+  void runOnFunction() override;
+};
+} // namespace
+
+// Region walk callback which makes GPU ops implementing the AsyncOpInterface
+// execute asynchronously.
+struct GpuAsyncRegionPass::Callback {
+  // 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.
+  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.
+    currentToken = createWaitOp(op->getLoc(), Type(), {currentToken});
+    return success();
+  }
+
+  // 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");
+
+    // If there is no current token, insert a `gpu.wait async` without
+    // dependencies to create one.
+    if (!currentToken)
+      currentToken = createWaitOp(op->getLoc(), tokenType, {});
+    asyncOp.addAsyncDependency(currentToken);
+
+    // Clone the op to return a token in addition to the other results.
+    SmallVector<Type, 1> resultTypes = {tokenType};
+    resultTypes.reserve(1 + op->getNumResults());
+    copy(op->getResultTypes(), std::back_inserter(resultTypes));
+    auto *newOp = Operation::create(op->getLoc(), op->getName(), resultTypes,
+                                    op->getOperands(), op->getMutableAttrDict(),
+                                    op->getSuccessors());
+
+    // Replace the op with the async clone.
+    auto results = newOp->getResults();
+    currentToken = results.front();
+    builder.insert(newOp);
+    op->replaceAllUsesWith(results.drop_front());
+    op->erase();
+
+    return success();
+  }
+
+  Value createWaitOp(Location loc, Type resultType, ValueRange operands) {
+    return builder.create<gpu::WaitOp>(loc, resultType, operands).asyncToken();
+  }
+
+  OpBuilder builder;
+  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 = {};
+};
+
+// 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())})
+          .wasInterrupted())
+    return signalPassFailure();
+}
+
+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: %[[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
+  }
+
+}