diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -793,15 +793,13 @@ if (parser.parseOptionalKeyword("args")) return success(); - SmallVector args; - if (parser.parseArgumentList(args, OpAsmParser::Delimiter::Paren, - /*allowType=*/true)) - return failure(); - for (auto &arg : args) { - argNames.push_back(arg.ssaName); - argTypes.push_back(arg.type); - } - return success(); + auto parseElement = [&]() -> ParseResult { + return failure(parser.parseOperand(argNames.emplace_back()) || + parser.parseColonType(argTypes.emplace_back())); + }; + + return parser.parseCommaSeparatedList(OpAsmParser::Delimiter::Paren, + parseElement, " in argument list"); } static void printLaunchFuncOperands(OpAsmPrinter &printer, Operation *, diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir --- a/mlir/test/Dialect/GPU/ops.mlir +++ b/mlir/test/Dialect/GPU/ops.mlir @@ -121,6 +121,8 @@ } } + func.func private @two_value_generator() -> (f32, memref) + func.func @foo() { %0 = "op"() : () -> (f32) %1 = "op"() : () -> (memref) @@ -140,6 +142,11 @@ // CHECK: %{{.*}} = gpu.launch_func async [%{{.*}}] @kernels::@kernel_2 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) %t1 = gpu.launch_func async [%t0] @kernels::@kernel_2 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) + // CHECK: %[[VALUES:.*]]:2 = call + %values:2 = func.call @two_value_generator() : () -> (f32, memref) + // CHECK: gpu.launch_func @kernels::@kernel_1 {{.*}} args(%[[VALUES]]#0 : f32, %[[VALUES]]#1 : memref) + gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) args(%values#0 : f32, %values#1 : memref) + return }