[Mlir-commits] [mlir] [mlir][gpu] make launch_func op use SymbolUserOpInterface (PR #173277)
lonely eagle
llvmlistbot at llvm.org
Mon Dec 22 08:10:21 PST 2025
https://github.com/linuxlonelyeagle created https://github.com/llvm/llvm-project/pull/173277
None
>From b8f0f5f4b60b0008ee431165379890be2997e62f Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Thu, 18 Dec 2025 05:20:23 +0000
Subject: [PATCH 1/2] make launch_func use SymbolUserOpInterface.
---
mlir/include/mlir/Dialect/GPU/IR/GPUOps.td | 3 +-
mlir/lib/Dialect/GPU/IR/GPUDialect.cpp | 156 +++++++++++----------
2 files changed, 81 insertions(+), 78 deletions(-)
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index 5c7df25c58cde..ae9a6a24724ae 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -607,7 +607,8 @@ def LaunchIndx : AnyTypeOf<[Index, I32, I64]>;
def GPU_LaunchFuncOp :GPU_Op<"launch_func", [
GPU_AsyncOpInterface, AttrSizedOperandSegments,
AllTypesMatch<["gridSizeX", "gridSizeY", "gridSizeZ", "blockSizeX",
- "blockSizeY", "blockSizeZ"]>]>,
+ "blockSizeY", "blockSizeZ"]>,
+ DeclareOpInterfaceMethods<SymbolUserOpInterface>]>,
Arguments<(ins Variadic<GPU_AsyncToken>:$asyncDependencies,
SymbolRefAttr:$kernel,
LaunchIndx:$gridSizeX,
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 61a630aa88960..0fd9c1bdcf565 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -405,83 +405,7 @@ LogicalResult GPUDialect::verifyOperationAttribute(Operation *op,
return op->emitError("expected '")
<< getContainerModuleAttrName() << "' attribute to be attached to '"
<< ModuleOp::getOperationName() << '\'';
-
- auto walkResult = module.walk([&module](LaunchFuncOp launchOp) -> WalkResult {
- // Ignore launches that are nested more or less deep than functions in the
- // module we are currently checking.
- if (!launchOp->getParentOp() ||
- launchOp->getParentOp()->getParentOp() != module)
- return success();
-
- // Ignore launch ops with missing attributes here. The errors will be
- // reported by the verifiers of those ops.
- if (!launchOp->getAttrOfType<SymbolRefAttr>(
- LaunchFuncOp::getKernelAttrName(launchOp->getName())))
- return success();
-
- // Check that `launch_func` refers to a well-formed GPU kernel container.
- StringAttr kernelContainerName = launchOp.getKernelModuleName();
- Operation *kernelContainer = module.lookupSymbol(kernelContainerName);
- if (!kernelContainer)
- return launchOp.emitOpError()
- << "kernel container '" << kernelContainerName.getValue()
- << "' is undefined";
-
- // If the container is a GPU binary op return success.
- if (isa<BinaryOp>(kernelContainer))
- return success();
-
- auto kernelModule = dyn_cast<GPUModuleOp>(kernelContainer);
- if (!kernelModule)
- return launchOp.emitOpError()
- << "kernel module '" << kernelContainerName.getValue()
- << "' is undefined";
-
- // Check that `launch_func` refers to a well-formed kernel function.
- Operation *kernelFunc = module.lookupSymbol(launchOp.getKernelAttr());
- if (!kernelFunc)
- return launchOp.emitOpError("kernel function '")
- << launchOp.getKernel() << "' is undefined";
- auto kernelConvertedFunction = dyn_cast<FunctionOpInterface>(kernelFunc);
- if (!kernelConvertedFunction) {
- InFlightDiagnostic diag = launchOp.emitOpError()
- << "referenced kernel '" << launchOp.getKernel()
- << "' is not a function";
- diag.attachNote(kernelFunc->getLoc()) << "see the kernel definition here";
- return diag;
- }
-
- if (!kernelFunc->getAttrOfType<mlir::UnitAttr>(
- GPUDialect::getKernelFuncAttrName()))
- return launchOp.emitOpError("kernel function is missing the '")
- << GPUDialect::getKernelFuncAttrName() << "' attribute";
-
- // TODO: If the kernel isn't a GPU function (which happens during separate
- // compilation), do not check type correspondence as it would require the
- // verifier to be aware of the type conversion.
- auto kernelGPUFunction = dyn_cast<gpu::GPUFuncOp>(kernelFunc);
- if (!kernelGPUFunction)
- return success();
-
- unsigned actualNumArguments = launchOp.getNumKernelOperands();
- unsigned expectedNumArguments = kernelGPUFunction.getNumArguments();
- if (expectedNumArguments != actualNumArguments)
- return launchOp.emitOpError("got ")
- << actualNumArguments << " kernel operands but expected "
- << expectedNumArguments;
-
- auto functionType = kernelGPUFunction.getFunctionType();
- for (unsigned i = 0; i < expectedNumArguments; ++i) {
- if (launchOp.getKernelOperand(i).getType() != functionType.getInput(i)) {
- return launchOp.emitOpError("type of function argument ")
- << i << " does not match";
- }
- }
-
- return success();
- });
-
- return walkResult.wasInterrupted() ? failure() : success();
+ return success();
}
/// Parses an optional list of async operands with an optional leading keyword.
@@ -1381,6 +1305,84 @@ LogicalResult LaunchFuncOp::verify() {
return success();
}
+LogicalResult
+LaunchFuncOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+ LaunchFuncOp launchOp = *this;
+ ModuleOp module = (*this)->getParentOfType<ModuleOp>();
+ // Ignore launches that are nested more or less deep than functions in the
+ // module we are currently checking.
+ if (!launchOp->getParentOp() ||
+ launchOp->getParentOp()->getParentOp() != module)
+ return success();
+
+ // Ignore launch ops with missing attributes here. The errors will be
+ // reported by the verifiers of those ops.
+ if (!launchOp->getAttrOfType<SymbolRefAttr>(
+ LaunchFuncOp::getKernelAttrName(launchOp->getName())))
+ return success();
+
+ // Check that `launch_func` refers to a well-formed GPU kernel container.
+ StringAttr kernelContainerName = launchOp.getKernelModuleName();
+ Operation *kernelContainer = module.lookupSymbol(kernelContainerName);
+ if (!kernelContainer)
+ return launchOp.emitOpError()
+ << "kernel container '" << kernelContainerName.getValue()
+ << "' is undefined";
+
+ // If the container is a GPU binary op return success.
+ if (isa<BinaryOp>(kernelContainer))
+ return success();
+
+ auto kernelModule = dyn_cast<GPUModuleOp>(kernelContainer);
+ if (!kernelModule)
+ return launchOp.emitOpError()
+ << "kernel module '" << kernelContainerName.getValue()
+ << "' is undefined";
+
+ // Check that `launch_func` refers to a well-formed kernel function.
+ Operation *kernelFunc = module.lookupSymbol(launchOp.getKernelAttr());
+ if (!kernelFunc)
+ return launchOp.emitOpError("kernel function '")
+ << launchOp.getKernel() << "' is undefined";
+ auto kernelConvertedFunction = dyn_cast<FunctionOpInterface>(kernelFunc);
+ if (!kernelConvertedFunction) {
+ InFlightDiagnostic diag = launchOp.emitOpError()
+ << "referenced kernel '" << launchOp.getKernel()
+ << "' is not a function";
+ diag.attachNote(kernelFunc->getLoc()) << "see the kernel definition here";
+ return diag;
+ }
+
+ if (!kernelFunc->getAttrOfType<mlir::UnitAttr>(
+ GPUDialect::getKernelFuncAttrName()))
+ return launchOp.emitOpError("kernel function is missing the '")
+ << GPUDialect::getKernelFuncAttrName() << "' attribute";
+
+ // TODO: If the kernel isn't a GPU function (which happens during separate
+ // compilation), do not check type correspondence as it would require the
+ // verifier to be aware of the type conversion.
+ auto kernelGPUFunction = dyn_cast<gpu::GPUFuncOp>(kernelFunc);
+ if (!kernelGPUFunction)
+ return success();
+
+ unsigned actualNumArguments = launchOp.getNumKernelOperands();
+ unsigned expectedNumArguments = kernelGPUFunction.getNumArguments();
+ if (expectedNumArguments != actualNumArguments)
+ return launchOp.emitOpError("got ")
+ << actualNumArguments << " kernel operands but expected "
+ << expectedNumArguments;
+
+ auto functionType = kernelGPUFunction.getFunctionType();
+ for (unsigned i = 0; i < expectedNumArguments; ++i) {
+ if (launchOp.getKernelOperand(i).getType() != functionType.getInput(i)) {
+ return launchOp.emitOpError("type of function argument ")
+ << i << " does not match";
+ }
+ }
+
+ return success();
+}
+
static ParseResult
parseLaunchDimType(OpAsmParser &parser, Type &dimTy,
std::optional<OpAsmParser::UnresolvedOperand> clusterValue,
>From e07b1e80ad489ec67b79979b75acead936418044 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 22 Dec 2025 16:08:33 +0000
Subject: [PATCH 2/2] use clang-format and fix test.
---
mlir/include/mlir/Dialect/GPU/IR/GPUOps.td | 4 ++--
mlir/lib/Dialect/GPU/IR/GPUDialect.cpp | 6 ++++--
mlir/test/Dialect/GPU/invalid.mlir | 4 ++--
3 files changed, 8 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index ae9a6a24724ae..0215adf7b8105 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -606,9 +606,9 @@ def LaunchIndx : AnyTypeOf<[Index, I32, I64]>;
def GPU_LaunchFuncOp :GPU_Op<"launch_func", [
GPU_AsyncOpInterface, AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<SymbolUserOpInterface>,
AllTypesMatch<["gridSizeX", "gridSizeY", "gridSizeZ", "blockSizeX",
- "blockSizeY", "blockSizeZ"]>,
- DeclareOpInterfaceMethods<SymbolUserOpInterface>]>,
+ "blockSizeY", "blockSizeZ"]>]>,
Arguments<(ins Variadic<GPU_AsyncToken>:$asyncDependencies,
SymbolRefAttr:$kernel,
LaunchIndx:$gridSizeX,
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 0fd9c1bdcf565..46e6c8f9386cb 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -1323,7 +1323,8 @@ LaunchFuncOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
// Check that `launch_func` refers to a well-formed GPU kernel container.
StringAttr kernelContainerName = launchOp.getKernelModuleName();
- Operation *kernelContainer = module.lookupSymbol(kernelContainerName);
+ Operation *kernelContainer =
+ symbolTable.lookupNearestSymbolFrom(module, kernelContainerName);
if (!kernelContainer)
return launchOp.emitOpError()
<< "kernel container '" << kernelContainerName.getValue()
@@ -1340,7 +1341,8 @@ LaunchFuncOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
<< "' is undefined";
// Check that `launch_func` refers to a well-formed kernel function.
- Operation *kernelFunc = module.lookupSymbol(launchOp.getKernelAttr());
+ Operation *kernelFunc =
+ symbolTable.lookupNearestSymbolFrom(module, launchOp.getKernelAttr());
if (!kernelFunc)
return launchOp.emitOpError("kernel function '")
<< launchOp.getKernel() << "' is undefined";
diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
index 26bcf948bc85d..eaff7d8a0d5db 100644
--- a/mlir/test/Dialect/GPU/invalid.mlir
+++ b/mlir/test/Dialect/GPU/invalid.mlir
@@ -137,14 +137,14 @@ module attributes {gpu.container_module} {
// -----
module attributes {gpu.container_module} {
- module @kernels {
+ gpu.module @kernels_container {
gpu.func @kernel_1(%arg1 : !llvm.ptr) kernel {
gpu.return
}
}
func.func @launch_func_missing_kernel_attr(%sz : index, %arg : !llvm.ptr) {
- // expected-error at +1 {{kernel module 'kernels' is undefined}}
+ // expected-error at +1 {{kernel container 'kernels' is undefined}}
gpu.launch_func @kernels::@kernel_1 blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz) args(%arg : !llvm.ptr)
return
}
More information about the Mlir-commits
mailing list