[flang-commits] [flang] [mlir] [mlir][gpu] Add stream-based async mode to gpu.launch (PR #213031)
Ivan R. Ivanov via flang-commits
flang-commits at lists.llvm.org
Thu Jul 30 07:54:45 PDT 2026
https://github.com/ivanradanov updated https://github.com/llvm/llvm-project/pull/213031
>From 6ff6c8e916fe8c5bdf69c4dfafc3ee0dac28f617 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <iivanov at nvidia.com>
Date: Thu, 30 Jul 2026 05:09:51 -0700
Subject: [PATCH 1/2] [mlir][gpu] Add stream-based async mode to gpu.launch
Currently, only gpu.launch_func supports stream based asynchronicity.
This patch adds the same mode to gpu.launch, and makes kernel outlining
preserve that when converting to gpu.launch_func.
This patch also documents the stream-based model in the operations
description, which was previously missing.
Assisted-by: Claude Code
---
flang/lib/Optimizer/CodeGen/TargetRewrite.cpp | 14 +--
.../Transforms/CUDA/CUFOpConversion.cpp | 18 +--
flang/test/Fir/CUDA/cuda-gpu-launch-func.mlir | 6 +-
flang/test/Fir/CUDA/cuda-launch.fir | 2 +-
flang/test/Fir/CUDA/cuda-stream.mlir | 4 +-
mlir/include/mlir/Dialect/GPU/IR/GPUOps.td | 40 ++++---
.../GPUCommon/GPUToLLVMConversion.cpp | 2 +-
mlir/lib/Dialect/GPU/IR/GPUDialect.cpp | 107 +++++++++---------
.../GPU/Transforms/KernelOutlining.cpp | 2 +-
mlir/test/Dialect/GPU/invalid.mlir | 55 ++++++++-
mlir/test/Dialect/GPU/ops.mlir | 11 ++
11 files changed, 163 insertions(+), 98 deletions(-)
diff --git a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
index 37f502d1aa1e0..22435aafe352d 100644
--- a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
@@ -536,13 +536,13 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
// TODO propagate/update call argument and result attributes.
if constexpr (std::is_same_v<std::decay_t<A>, mlir::gpu::LaunchFuncOp>) {
mlir::Value asyncToken = callOp.getAsyncToken();
- auto newCall = A::create(*rewriter, loc, callOp.getKernel(),
- callOp.getGridSizeOperandValues(),
- callOp.getBlockSizeOperandValues(),
- callOp.getDynamicSharedMemorySize(), newOpers,
- asyncToken ? asyncToken.getType() : nullptr,
- callOp.getAsyncDependencies(),
- /*clusterSize=*/std::nullopt);
+ auto newCall = A::create(
+ *rewriter, loc, callOp.getKernel(), callOp.getGridSizeOperandValues(),
+ callOp.getBlockSizeOperandValues(),
+ callOp.getDynamicSharedMemorySize(), newOpers,
+ asyncToken ? asyncToken.getType() : nullptr,
+ callOp.getAsyncDependencies(), callOp.getAsyncObject(),
+ /*clusterSize=*/std::nullopt);
if (callOp.getClusterSizeX())
newCall.getClusterSizeXMutable().assign(callOp.getClusterSizeX());
if (callOp.getClusterSizeY())
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
index 44f1b5154a2b0..8e16473eaf982 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
@@ -479,23 +479,23 @@ struct CUFLaunchOpConversion
args.push_back(arg);
}
mlir::Value dynamicShmemSize = op.getBytes() ? op.getBytes() : zero;
+ mlir::Type tokenType = nullptr;
+ SmallVector<Value, 1> tokens;
+ if (op.getStream()) {
+ tokens.push_back(
+ cuf::StreamCastOp::create(rewriter, loc, op.getStream()));
+ tokenType = tokens.front().getType();
+ }
auto gpuLaunchOp = mlir::gpu::LaunchFuncOp::create(
rewriter, loc, kernelName,
mlir::gpu::KernelDim3{gridSizeX, gridSizeY, gridSizeZ},
mlir::gpu::KernelDim3{blockSizeX, blockSizeY, blockSizeZ},
- dynamicShmemSize, args);
+ dynamicShmemSize, args, tokenType, tokens);
if (clusterDimX && clusterDimY && clusterDimZ) {
gpuLaunchOp.getClusterSizeXMutable().assign(clusterDimX);
gpuLaunchOp.getClusterSizeYMutable().assign(clusterDimY);
gpuLaunchOp.getClusterSizeZMutable().assign(clusterDimZ);
}
- if (op.getStream()) {
- mlir::OpBuilder::InsertionGuard guard(rewriter);
- rewriter.setInsertionPoint(gpuLaunchOp);
- mlir::Value stream =
- cuf::StreamCastOp::create(rewriter, loc, op.getStream());
- gpuLaunchOp.getAsyncDependenciesMutable().append(stream);
- }
if (procAttr)
gpuLaunchOp->setAttr(cuf::getProcAttrName(), procAttr);
else
@@ -503,7 +503,7 @@ struct CUFLaunchOpConversion
gpuLaunchOp->setAttr(cuf::getProcAttrName(),
cuf::ProcAttributeAttr::get(
op.getContext(), cuf::ProcAttribute::Global));
- rewriter.replaceOp(op, gpuLaunchOp);
+ rewriter.eraseOp(op);
return mlir::success();
}
diff --git a/flang/test/Fir/CUDA/cuda-gpu-launch-func.mlir b/flang/test/Fir/CUDA/cuda-gpu-launch-func.mlir
index 47e6dfd220081..03e441ada81c9 100644
--- a/flang/test/Fir/CUDA/cuda-gpu-launch-func.mlir
+++ b/flang/test/Fir/CUDA/cuda-gpu-launch-func.mlir
@@ -180,7 +180,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr<272>, d
%2 = llvm.mlir.constant(0 : i32) : i32
%3 = llvm.mlir.constant(10 : index) : i64
%token = cuf.stream_cast %stream : !llvm.ptr
- gpu.launch_func [%token] @cuda_device_mod::@_QMmod1Psub1 blocks in (%3, %3, %0) threads in (%3, %3, %0) : i64 dynamic_shared_memory_size %2 {cuf.proc_attr = #cuf.cuda_proc<global>}
+ %res_token = gpu.launch_func async [%token] @cuda_device_mod::@_QMmod1Psub1 blocks in (%3, %3, %0) threads in (%3, %3, %0) : i64 dynamic_shared_memory_size %2 {cuf.proc_attr = #cuf.cuda_proc<global>}
llvm.return
}
gpu.binary @cuda_device_mod [#gpu.object<#nvvm.target, "">]
@@ -213,7 +213,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i1, dense<8> : ve
%13 = llvm.mlir.addressof @_QQclX91d13f6e74caa2f03965d7a7c6a8fdd5 : !llvm.ptr
%14 = llvm.call @_FortranACUFMemAlloc(%2, %11, %13, %6) : (i64, i32, !llvm.ptr, i32) -> !llvm.ptr
%token = cuf.stream_cast %stream : !llvm.ptr
- gpu.launch_func [%token] @cuda_device_mod::@_QMmod1Psub1 blocks in (%7, %7, %7) threads in (%12, %7, %7) : i64 dynamic_shared_memory_size %11 args(%14 : !llvm.ptr) {cuf.proc_attr = #cuf.cuda_proc<grid_global>}
+ %res_token = gpu.launch_func async [%token] @cuda_device_mod::@_QMmod1Psub1 blocks in (%7, %7, %7) threads in (%12, %7, %7) : i64 dynamic_shared_memory_size %11 args(%14 : !llvm.ptr) {cuf.proc_attr = #cuf.cuda_proc<grid_global>}
llvm.return
}
llvm.func @_QMmod1Psub1(!llvm.ptr) -> ()
@@ -244,7 +244,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr<272>, d
%3 = llvm.mlir.constant(10 : index) : i64
%stream = llvm.alloca %0 x i64 : (i64) -> !llvm.ptr
%token = cuf.stream_cast %stream : !llvm.ptr
- %4 = gpu.launch_func async [%token] @cuda_device_mod::@_QMmod1Psub1 blocks in (%3, %3, %0) threads in (%3, %3, %0) : i64 dynamic_shared_memory_size %2 {cuf.proc_attr = #cuf.cuda_proc<global>}
+ %res_token = gpu.launch_func async [%token] @cuda_device_mod::@_QMmod1Psub1 blocks in (%3, %3, %0) threads in (%3, %3, %0) : i64 dynamic_shared_memory_size %2 {cuf.proc_attr = #cuf.cuda_proc<global>}
llvm.return
}
gpu.binary @cuda_device_mod [#gpu.object<#nvvm.target, "">]
diff --git a/flang/test/Fir/CUDA/cuda-launch.fir b/flang/test/Fir/CUDA/cuda-launch.fir
index 92db6ecaadc45..e5ea7bb2479ac 100644
--- a/flang/test/Fir/CUDA/cuda-launch.fir
+++ b/flang/test/Fir/CUDA/cuda-launch.fir
@@ -155,4 +155,4 @@ module attributes {gpu.container_module, dlti.dl_spec = #dlti.dl_spec<#dlti.dl_e
// CHECK: %[[STREAM:.*]] = fir.alloca i64 {bindc_name = "stream", uniq_name = "_QMtest_callFhostEstream"}
// CHECK: %[[DECL_STREAM:.*]]:2 = hlfir.declare %[[STREAM]] {uniq_name = "_QMtest_callFhostEstream"} : (!fir.ref<i64>) -> (!fir.ref<i64>, !fir.ref<i64>)
// CHECK: %[[TOKEN:.*]] = cuf.stream_cast %[[DECL_STREAM]]#0 : !fir.ref<i64>
-// CHECK: gpu.launch_func [%[[TOKEN]]] @cuda_device_mod::@_QMdevptrPtest
+// CHECK: %[[RES_TOKEN:.*]] = gpu.launch_func async [%[[TOKEN]]] @cuda_device_mod::@_QMdevptrPtest
diff --git a/flang/test/Fir/CUDA/cuda-stream.mlir b/flang/test/Fir/CUDA/cuda-stream.mlir
index a501603fd35d1..1c463ea39458e 100644
--- a/flang/test/Fir/CUDA/cuda-stream.mlir
+++ b/flang/test/Fir/CUDA/cuda-stream.mlir
@@ -10,7 +10,7 @@ module attributes {gpu.container_module} {
%0 = fir.alloca i64
%1 = arith.constant 1 : index
%asyncTok = cuf.stream_cast %0 : !fir.ref<i64>
- gpu.launch_func [%asyncTok] @cuda_device_mod::@_QMmod1Psub1 blocks in (%1, %1, %1) threads in (%1, %1, %1) args() {cuf.proc_attr = #cuf.cuda_proc<grid_global>}
+ %res_token = gpu.launch_func async [%asyncTok] @cuda_device_mod::@_QMmod1Psub1 blocks in (%1, %1, %1) threads in (%1, %1, %1) args() {cuf.proc_attr = #cuf.cuda_proc<grid_global>}
return
}
}
@@ -18,4 +18,4 @@ module attributes {gpu.container_module} {
// CHECK-LABEL: func.func @_QMmod1Phost_sub()
// CHECK: %[[STREAM:.*]] = fir.alloca i64
// CHECK: %[[TOKEN:.*]] = cuf.stream_cast %[[STREAM]] : !fir.ref<i64>
-// CHECK: gpu.launch_func [%[[TOKEN]]] @cuda_device_mod::@_QMmod1Psub1
+// CHECK: %[[RES_TOKEN:.*]] = gpu.launch_func async [%[[TOKEN]]] @cuda_device_mod::@_QMmod1Psub1
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index a9032820f5646..91edd1dbd5dad 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -637,13 +637,22 @@ def GPU_LaunchFuncOp :GPU_Op<"launch_func", [
operation has a symbol attribute named `kernel` to identify the fully
specified kernel function to launch (both the gpu.module and func).
- The `gpu.launch_func` supports async dependencies: the kernel does not start
+ By the default, the host implicitly blocks until kernel execution has
+ completed.
+
+ Otherwise, the operation supports two async models.
+
+ The first one is dependency-based and is enabled when the `async` keyword is
+ present in text form, and corresponds to when the operation produces the
+ optional token result of type `!gpu.async.token`. Other async GPU ops can
+ take this token as dependency. In this case, the `gpu.launch_func` does not
+ block, and supports specifying async dependencies: the kernel does not start
executing until the ops producing those async dependencies have completed.
- By the default, the host implicitly blocks until kernel execution has
- completed. If the `async` keyword is present, the host does not block but
- instead a `!gpu.async.token` is returned. Other async GPU ops can take this
- token as dependency.
+ The second async model is stream-based. When `asyncObject` is present, the
+ launch operation is queued to execute on the queue represented by it.
+
+ The two async models are mutually exclusive.
The operation requires at least the grid and block sizes along the x,y,z
dimensions as arguments. When a lower-dimensional kernel is required,
@@ -734,16 +743,13 @@ def GPU_LaunchFuncOp :GPU_Op<"launch_func", [
"ValueRange":$kernelOperands,
CArg<"Type", "nullptr">:$asyncTokenType,
CArg<"ValueRange", "{}">:$asyncDependencies,
+ CArg<"Value", "nullptr">:$asyncObject,
CArg<"std::optional<KernelDim3>", "std::nullopt">:$clusterSize)>,
OpBuilder<(ins "SymbolRefAttr":$kernel, "KernelDim3":$gridSize,
"KernelDim3":$blockSize, "Value":$dynamicSharedMemorySize,
"ValueRange":$kernelOperands,
- "Type":$asyncTokenType,
+ CArg<"Type", "nullptr">:$asyncTokenType,
CArg<"ValueRange", "{}">:$asyncDependencies,
- CArg<"std::optional<KernelDim3>", "std::nullopt">:$clusterSize)>,
- OpBuilder<(ins "SymbolRefAttr":$kernel, "KernelDim3":$gridSize,
- "KernelDim3":$blockSize, "Value":$dynamicSharedMemorySize,
- "ValueRange":$kernelOperands,
CArg<"Value", "nullptr">:$asyncObject,
CArg<"std::optional<KernelDim3>", "std::nullopt">:$clusterSize)>
];
@@ -813,20 +819,17 @@ def GPU_LaunchOp : GPU_Op<"launch", [
UnitAttr:$cooperative,
OptionalAttr<FlatSymbolRefAttr>:$module,
OptionalAttr<FlatSymbolRefAttr>:$function,
- OptionalAttr<ConfinedAttr<I64Attr, [IntNonNegative]>>:$workgroup_attributions)>,
+ OptionalAttr<ConfinedAttr<I64Attr, [IntNonNegative]>>:$workgroup_attributions,
+ Optional<AnyType>:$asyncObject)>,
Results<(outs Optional<GPU_AsyncToken>:$asyncToken)> {
let summary = "GPU kernel launch operation";
let description = [{
Launch a kernel on the specified grid of thread blocks. The body of the
- kernel is defined by the single region that this operation contains. The
- operation takes an optional list of async dependencies followed by six
- operands and an optional operand.
+ kernel is defined by the single region that this operation contains.
- The `async` keyword indicates the kernel should be launched asynchronously;
- the operation returns a new !gpu.async.token when the keyword is specified.
- The kernel launched does not start executing until the ops producing its
- async dependencies (optional operands) have completed.
+ The async execution model is equivalent to the `gpu.launch_func` op, refer
+ to its description.
The first three operands (following any async dependencies) are grid sizes
along the x,y,z dimensions and the following three are block sizes along the
@@ -945,6 +948,7 @@ def GPU_LaunchOp : GPU_Op<"launch", [
CArg<"Value", "nullptr">:$dynamicSharedMemorySize,
CArg<"Type", "nullptr">:$asyncTokenType,
CArg<"ValueRange", "{}">:$asyncDependencies,
+ CArg<"Value", "nullptr">:$asyncObject,
CArg<"TypeRange", "{}">:$workgroupAttributions,
CArg<"TypeRange", "{}">:$privateAttributions,
CArg<"Value", "nullptr">:$clusterSizeX,
diff --git a/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp b/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
index e0048876bf348..80b9a15ac9a41 100644
--- a/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
+++ b/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
@@ -1102,7 +1102,7 @@ LogicalResult LegalizeLaunchFuncOpPattern::matchAndRewrite(
adaptor.getBlockSizeZ()},
adaptor.getDynamicSharedMemorySize(),
llvmArgumentsWithSizes.empty() ? llvmArguments : llvmArgumentsWithSizes,
- stream, clusterSize);
+ nullptr, {}, stream, clusterSize);
if (launchOp.getCooperative())
newLaunchOp.setCooperative(true);
if (launchOp.getAsyncToken())
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index eaea0142e0438..22c5e3c9b86ad 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -715,7 +715,7 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
Value getBlockSizeX, Value getBlockSizeY,
Value getBlockSizeZ, Value dynamicSharedMemorySize,
Type asyncTokenType, ValueRange asyncDependencies,
- TypeRange workgroupAttributions,
+ Value asyncObject, TypeRange workgroupAttributions,
TypeRange privateAttributions, Value clusterSizeX,
Value clusterSizeY, Value clusterSizeZ,
FlatSymbolRefAttr module, FlatSymbolRefAttr function) {
@@ -742,6 +742,8 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
result.addOperands(clusterSizeZ);
if (dynamicSharedMemorySize)
result.addOperands(dynamicSharedMemorySize);
+ if (asyncObject)
+ result.addOperands(asyncObject);
// Add optional module and function attributes.
if (module)
@@ -763,12 +765,13 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
for (Type argTy : privateAttributions)
body->addArgument(argTy, result.location);
// Fill OperandSegmentSize Attribute.
- SmallVector<int32_t, 11> segmentSizes(11, 1);
+ SmallVector<int32_t, 12> segmentSizes(12, 1);
segmentSizes.front() = asyncDependencies.size();
- segmentSizes.back() = dynamicSharedMemorySize ? 1 : 0;
segmentSizes[7] = clusterSizeX ? 1 : 0;
segmentSizes[8] = clusterSizeY ? 1 : 0;
segmentSizes[9] = clusterSizeZ ? 1 : 0;
+ segmentSizes[10] = dynamicSharedMemorySize ? 1 : 0;
+ segmentSizes[11] = asyncObject ? 1 : 0;
result.addAttribute(getOperandSegmentSizeAttr(),
builder.getDenseI32ArrayAttr(segmentSizes));
}
@@ -830,7 +833,23 @@ std::optional<KernelDim3> LaunchOp::getClusterSizeOperandValues() {
return KernelDim3{operands[6], operands[7], operands[8]};
}
+template <typename OpTy>
+static LogicalResult verifyLaunchAsyncModel(OpTy op) {
+ if (!op.getAsyncDependencies().empty() && !op.getAsyncToken())
+ return op.emitOpError("dependency operands require the dependency-based "
+ "async model i.e. returning a token");
+ if (op.getAsyncToken() && op.getAsyncObject())
+ return op.emitOpError("stream-based and dependency-based async models are "
+ "mutually exclusive");
+ if (op.getNumResults() == 0 && op.getAsyncToken())
+ return op.emitOpError("needs to be named when async keyword is specified");
+ return success();
+}
+
LogicalResult LaunchOp::verify() {
+ if (verifyLaunchAsyncModel(*this).failed())
+ return failure();
+
if (!(hasClusterSize()) &&
(getClusterSizeX() || getClusterSizeY() || getClusterSizeZ()))
return emitOpError() << "cluster size must be all present";
@@ -877,9 +896,6 @@ LogicalResult LaunchOp::verifyRegions() {
}
}
- if (getNumResults() == 0 && getAsyncToken())
- return emitOpError("needs to be named when async keyword is specified");
-
return success();
}
@@ -896,6 +912,9 @@ static void printSizeAssignment(OpAsmPrinter &p, KernelDim3 size,
}
void LaunchOp::print(OpAsmPrinter &p) {
+ if (auto asyncObject = getAsyncObject()) {
+ p << " <" << asyncObject << " : " << asyncObject.getType() << ">";
+ }
if (getAsyncToken()) {
p << " async";
if (!getAsyncDependencies().empty())
@@ -986,8 +1005,9 @@ parseSizeAssignment(OpAsmParser &parser,
}
/// Parses a Launch operation.
-/// operation ::= `gpu.launch` (`async` `[` ssa-id-list `]`)?
-/// `clusters` `(` ssa-id-list `)` `in` ssa-reassignment (Optional)
+/// operation ::= `gpu.launch` (`<` ssa-use `:` type `>`)?
+/// (`async` `[` ssa-id-list `]`)?
+/// (`clusters` `(` ssa-id-list `)` `in` ssa-reassignment)?
/// `blocks` `(` ssa-id-list `)` `in` ssa-reassignment
/// `threads` `(` ssa-id-list `)` `in` ssa-reassignment
/// (`dynamic_shared_memory_size` ssa-use)?
@@ -1005,6 +1025,17 @@ ParseResult LaunchOp::parse(OpAsmParser &parser, OperationState &result) {
SmallVector<OpAsmParser::UnresolvedOperand, 16> regionArgs(
LaunchOp::kNumConfigRegionAttributes);
+ // Parse optional asyncObject: < value : type >
+ OpAsmParser::UnresolvedOperand asyncObjectOperand;
+ Type asyncObjectType;
+ bool hasAsyncObject = false;
+ if (succeeded(parser.parseOptionalLess())) {
+ hasAsyncObject = true;
+ if (parser.parseOperand(asyncObjectOperand) || parser.parseColon() ||
+ parser.parseType(asyncObjectType) || parser.parseGreater())
+ return failure();
+ }
+
// Parse optional async dependencies.
SmallVector<OpAsmParser::UnresolvedOperand, 4> asyncDependencies;
Type asyncTokenType;
@@ -1067,6 +1098,11 @@ ParseResult LaunchOp::parse(OpAsmParser &parser, OperationState &result) {
return failure();
}
+ // Resolve the asyncObject operand
+ if (hasAsyncObject && parser.resolveOperand(asyncObjectOperand,
+ asyncObjectType, result.operands))
+ return failure();
+
// Parse optional module attribute.
StringRef moduleAttrName = getModuleAttrName(result.name);
if (succeeded(parser.parseOptionalKeyword(moduleAttrName))) {
@@ -1135,7 +1171,7 @@ ParseResult LaunchOp::parse(OpAsmParser &parser, OperationState &result) {
parser.parseOptionalAttrDict(result.attributes))
return failure();
- SmallVector<int32_t, 11> segmentSizes(11, 1);
+ SmallVector<int32_t, 12> segmentSizes(12, 1);
segmentSizes.front() = asyncDependencies.size();
if (!hasCluster) {
@@ -1143,7 +1179,8 @@ ParseResult LaunchOp::parse(OpAsmParser &parser, OperationState &result) {
segmentSizes[8] = 0;
segmentSizes[9] = 0;
}
- segmentSizes.back() = hasDynamicSharedMemorySize ? 1 : 0;
+ segmentSizes[10] = hasDynamicSharedMemorySize ? 1 : 0;
+ segmentSizes[11] = hasAsyncObject ? 1 : 0;
result.addAttribute(LaunchOp::getOperandSegmentSizeAttr(),
parser.getBuilder().getDenseI32ArrayAttr(segmentSizes));
return success();
@@ -1216,7 +1253,7 @@ void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
SymbolRefAttr kernelSymbol, KernelDim3 gridSize,
KernelDim3 getBlockSize, Value dynamicSharedMemorySize,
ValueRange kernelOperands, Type asyncTokenType,
- ValueRange asyncDependencies,
+ ValueRange asyncDependencies, Value asyncObject,
std::optional<KernelDim3> clusterSize) {
assert(kernelSymbol.getNestedReferences().size() == 1 &&
"expected a symbol reference with a single nested reference");
@@ -1232,6 +1269,8 @@ void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
if (dynamicSharedMemorySize)
result.addOperands(dynamicSharedMemorySize);
result.addOperands(kernelOperands);
+ if (asyncObject)
+ result.addOperands(asyncObject);
Properties &prop = result.getOrAddProperties<Properties>();
prop.kernel = kernelSymbol;
@@ -1248,14 +1287,14 @@ void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
dynamicSharedMemorySize ? 1 : 0;
prop.operandSegmentSizes[segmentSizesLen - 2] =
static_cast<int32_t>(kernelOperands.size());
- prop.operandSegmentSizes[segmentSizesLen - 1] = 0;
+ prop.operandSegmentSizes[segmentSizesLen - 1] = asyncObject ? 1 : 0;
}
void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
GPUFuncOp kernelFunc, KernelDim3 gridSize,
KernelDim3 getBlockSize, Value dynamicSharedMemorySize,
ValueRange kernelOperands, Type asyncTokenType,
- ValueRange asyncDependencies,
+ ValueRange asyncDependencies, Value asyncObject,
std::optional<KernelDim3> clusterSize) {
auto kernelModule = kernelFunc->getParentOfType<GPUModuleOp>();
auto kernelSymbol =
@@ -1263,40 +1302,7 @@ void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
{SymbolRefAttr::get(kernelFunc.getNameAttr())});
build(builder, result, kernelSymbol, gridSize, getBlockSize,
dynamicSharedMemorySize, kernelOperands, asyncTokenType,
- asyncDependencies, clusterSize);
-}
-
-void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
- SymbolRefAttr kernel, KernelDim3 gridSize,
- KernelDim3 getBlockSize, Value dynamicSharedMemorySize,
- ValueRange kernelOperands, Value asyncObject,
- std::optional<KernelDim3> clusterSize) {
- // Add grid and block sizes as op operands, followed by the data operands.
- result.addOperands({gridSize.x, gridSize.y, gridSize.z, getBlockSize.x,
- getBlockSize.y, getBlockSize.z});
- if (clusterSize.has_value())
- result.addOperands({clusterSize->x, clusterSize->y, clusterSize->z});
- if (dynamicSharedMemorySize)
- result.addOperands(dynamicSharedMemorySize);
- result.addOperands(kernelOperands);
- if (asyncObject)
- result.addOperands(asyncObject);
- Properties &prop = result.getOrAddProperties<Properties>();
- prop.kernel = kernel;
- size_t segmentSizesLen = std::size(prop.operandSegmentSizes);
- // Initialize the segment sizes to 1.
- llvm::fill(prop.operandSegmentSizes, 1);
- prop.operandSegmentSizes[0] = 0;
- if (!clusterSize.has_value()) {
- prop.operandSegmentSizes[segmentSizesLen - 4] = 0;
- prop.operandSegmentSizes[segmentSizesLen - 5] = 0;
- prop.operandSegmentSizes[segmentSizesLen - 6] = 0;
- }
- prop.operandSegmentSizes[segmentSizesLen - 3] =
- dynamicSharedMemorySize ? 1 : 0;
- prop.operandSegmentSizes[segmentSizesLen - 2] =
- static_cast<int32_t>(kernelOperands.size());
- prop.operandSegmentSizes[segmentSizesLen - 1] = asyncObject ? 1 : 0;
+ asyncDependencies, asyncObject, clusterSize);
}
StringAttr LaunchFuncOp::getKernelModuleName() {
@@ -1333,6 +1339,9 @@ KernelDim3 LaunchFuncOp::getClusterSizeOperandValues() {
}
LogicalResult LaunchFuncOp::verify() {
+ if (verifyLaunchAsyncModel(*this).failed())
+ return failure();
+
auto module = (*this)->getParentOfType<ModuleOp>();
if (!module)
return emitOpError("expected to belong to a module");
@@ -1350,10 +1359,6 @@ LogicalResult LaunchFuncOp::verify() {
<< "expects types of the cluster dimensions must be the same";
}
- if (!getAsyncDependencies().empty() && getAsyncObject())
- return emitOpError(
- "cannot have both async dependencies and an explicit async object");
-
return success();
}
diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index 1c05dd4416aba..831ac7fb0354e 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
@@ -295,7 +295,7 @@ static void convertToLaunchFuncOp(gpu::LaunchOp launchOp,
launchOp.getGridSizeOperandValues(), launchOp.getBlockSizeOperandValues(),
launchOp.getDynamicSharedMemorySize(), operands,
asyncToken ? asyncToken.getType() : nullptr,
- launchOp.getAsyncDependencies(), clusterSize);
+ launchOp.getAsyncDependencies(), launchOp.getAsyncObject(), clusterSize);
if (launchOp.getCooperative())
launchFunc.setCooperative(true);
launchOp.replaceAllUsesWith(launchFunc);
diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
index e3a981c1c7afe..f91923d5fe503 100644
--- a/mlir/test/Dialect/GPU/invalid.mlir
+++ b/mlir/test/Dialect/GPU/invalid.mlir
@@ -3,7 +3,7 @@
func.func @empty_body(%sz : index) {
// expected-error at +1 {{'gpu.launch' op body region is empty}}
"gpu.launch"(%sz, %sz, %sz, %sz, %sz, %sz) ({
- }) {operandSegmentSizes = array<i32: 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0>} : (index, index, index, index, index, index) -> ()
+ }) {operandSegmentSizes = array<i32: 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0>} : (index, index, index, index, index, index) -> ()
return
}
@@ -13,7 +13,7 @@ func.func @not_enough_sizes(%sz : index) {
// expected-error at +1 {{expected 6 or more operands, but found 5}}
"gpu.launch"(%sz, %sz, %sz, %sz, %sz) ({
gpu.return
- }) {operandSegmentSizes = array<i32: 0, 1, 1, 1, 1, 1, 1, 0>} : (index, index, index, index, index) -> ()
+ }) {operandSegmentSizes = array<i32: 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0>} : (index, index, index, index, index) -> ()
return
}
@@ -25,7 +25,7 @@ func.func @no_region_attrs(%sz : index) {
^bb1(%bx: index, %by: index, %bz: index,
%tx: index, %ty: index, %tz: index):
gpu.terminator
- }) {operandSegmentSizes = array<i32: 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0>} : (index, index, index, index, index, index) -> ()
+ }) {operandSegmentSizes = array<i32: 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0>} : (index, index, index, index, index, index) -> ()
return
}
@@ -39,7 +39,7 @@ func.func @not_enough_cluster_region_attrs(%sz : index) {
%sbx: index, %sby: index, %sbz: index,
%stx: index, %sty: index, %stz: index):
gpu.terminator
- }) {operandSegmentSizes = array<i32: 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0>} : (index, index, index, index, index, index, index, index, index) -> ()
+ }) {operandSegmentSizes = array<i32: 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0>} : (index, index, index, index, index, index, index, index, index) -> ()
return
}
@@ -103,6 +103,32 @@ func.func @launch_wrong_threads(%sz : index) {
// -----
+func.func @launch_async_deps_no_token(%dep : !gpu.async.token, %sz : index) {
+ // expected-error at +1 {{'gpu.launch' op dependency operands require the dependency-based async model i.e. returning a token}}
+ "gpu.launch"(%dep, %sz, %sz, %sz, %sz, %sz, %sz) ({
+ ^bb0(%bx: index, %by: index, %bz: index,
+ %tx: index, %ty: index, %tz: index,
+ %nbx: index, %nby: index, %nbz: index,
+ %ntx: index, %nty: index, %ntz: index):
+ gpu.terminator
+ }) {operandSegmentSizes = array<i32: 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0>}
+ : (!gpu.async.token, index, index, index, index, index, index) -> ()
+ return
+}
+
+// -----
+
+func.func @launch_async_object_and_token(%sz : index, %stream : !llvm.ptr) {
+ // expected-error at +1 {{'gpu.launch' op stream-based and dependency-based async models are mutually exclusive}}
+ %t = gpu.launch <%stream : !llvm.ptr> async blocks(%bx, %by, %bz) in (%sbx = %sz, %sby = %sz, %sbz = %sz)
+ threads(%tx, %ty, %tz) in (%stx = %sz, %sty = %sz, %stz = %sz) {
+ gpu.terminator
+ }
+ return
+}
+
+// -----
+
func.func @launch_func_too_few_operands(%sz : index) {
// expected-error at +1 {{expected 6 or more operands}}
"gpu.launch_func"(%sz, %sz, %sz, %sz, %sz)
@@ -246,7 +272,7 @@ module attributes {gpu.container_module} {
func.func @launch_func_async_deps_and_async_object(%sz : index,
%stream : !llvm.ptr) {
%dep = gpu.wait async
- // expected-error at +1 {{cannot have both async dependencies and an explicit async object}}
+ // expected-error at +1 {{stream-based and dependency-based async models are mutually exclusive}}
%t = gpu.launch_func async [%dep] <%stream : !llvm.ptr> @kernels::@kernel_1
blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
return
@@ -255,6 +281,25 @@ module attributes {gpu.container_module} {
// -----
+module attributes {gpu.container_module} {
+ gpu.module @kernels {
+ gpu.func @kernel_1() kernel {
+ gpu.return
+ }
+ }
+
+ func.func @launch_func_async_deps_no_token(%dep : !gpu.async.token, %sz : index) {
+ // expected-error at +1 {{'gpu.launch_func' op dependency operands require the dependency-based async model i.e. returning a token}}
+ "gpu.launch_func"(%dep, %sz, %sz, %sz, %sz, %sz, %sz) {
+ kernel = @kernels::@kernel_1,
+ operandSegmentSizes = array<i32: 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0>}
+ : (!gpu.async.token, index, index, index, index, index, index) -> ()
+ return
+ }
+}
+
+// -----
+
module attributes {gpu.container_module} {
gpu.module @kernels {
gpu.func @kernel_1(%arg1 : !llvm.ptr) {
diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir
index 11cea6f82d7b5..cf496fef1fca8 100644
--- a/mlir/test/Dialect/GPU/ops.mlir
+++ b/mlir/test/Dialect/GPU/ops.mlir
@@ -65,6 +65,17 @@ module attributes {gpu.container_module} {
return
}
+ // CHECK-LABEL: func @launch_async_object(%{{.*}}: index, %{{.*}}: index) {
+ func.func @launch_async_object(%blk : index, %thrd : index) {
+ %stream = llvm.mlir.zero : !llvm.ptr
+ // CHECK: gpu.launch <%{{.*}} : !llvm.ptr> blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})
+ gpu.launch <%stream : !llvm.ptr> blocks(%arg0, %arg1, %arg2) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)
+ threads(%arg3, %arg4, %arg5) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {
+ gpu.terminator
+ }
+ return
+ }
+
// CHECK-LABEL:func @launch_async_no_deps(%{{.*}}: index, %{{.*}}: index) {
func.func @launch_async_no_deps(%blk : index, %thrd : index) {
// CHECK: %{{.*}} = gpu.launch async blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})
>From 3ef9e53dd0079d1999cb0e30417e39398f15dc53 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <iivanov at nvidia.com>
Date: Thu, 30 Jul 2026 07:54:14 -0700
Subject: [PATCH 2/2] typo
---
mlir/include/mlir/Dialect/GPU/IR/GPUOps.td | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index 91edd1dbd5dad..de145960692ad 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -637,8 +637,7 @@ def GPU_LaunchFuncOp :GPU_Op<"launch_func", [
operation has a symbol attribute named `kernel` to identify the fully
specified kernel function to launch (both the gpu.module and func).
- By the default, the host implicitly blocks until kernel execution has
- completed.
+ By default, the host implicitly blocks until kernel execution has completed.
Otherwise, the operation supports two async models.
More information about the flang-commits
mailing list