[flang-commits] [flang] [flang][CUDA] Safely verify registered kernel symbols (PR #230588)

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: nvptm

<details>
<summary>Changes</summary>

`cuf.register_kernel` resolved GPU module and kernel symbols during nested function verification, potentially racing with concurrent changes to the sibling GPU module. Add `SymbolUserOpInterface` and move these lookups to `verifySymbolUses()`, where they run as part of module-level symbol verification while the sibling IR is stable. 

---
Full diff: https://github.com/llvm/llvm-project/pull/230588.diff


2 Files Affected:

- (modified) flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td (+3-1) 
- (modified) flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp (+22-13) 


``````````diff
diff --git a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
index 71e0ee84d4d0b4..9a136df875a493 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
+++ b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
@@ -24,6 +24,7 @@ include "mlir/Dialect/GPU/IR/GPUBase.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/IR/BuiltinAttributes.td"
+include "mlir/IR/SymbolInterfaces.td"
 
 class cuf_Op<string mnemonic, list<Trait> traits>
     : Op<CUFDialect, mnemonic, traits>;
@@ -334,7 +335,8 @@ def cuf_RegisterModuleOp : cuf_Op<"register_module", []> {
   let results = (outs LLVM_AnyPointer:$modulePtr);
 }
 
-def cuf_RegisterKernelOp : cuf_Op<"register_kernel", []> {
+def cuf_RegisterKernelOp : cuf_Op<"register_kernel", [
+    DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
   let summary = "Register a CUDA kernel";
 
   let arguments = (ins
diff --git a/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp b/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
index 0c657d050c0572..6f5ef701b3c892 100644
--- a/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
+++ b/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
@@ -23,6 +23,7 @@
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
 #include "llvm/ADT/SmallVector.h"
 
 //===----------------------------------------------------------------------===//
@@ -319,31 +320,39 @@ mlir::LogicalResult cuf::RegisterKernelOp::verify() {
   if (getKernelName() == getKernelModuleName())
     return emitOpError("expect a module and a kernel name");
 
-  auto mod = getOperation()->getParentOfType<mlir::ModuleOp>();
-  if (!mod)
+  if (!getOperation()->getParentOfType<mlir::ModuleOp>())
     return emitOpError("expect to be in a module");
 
-  mlir::SymbolTable symTab(mod);
-  auto gpuMod = symTab.lookup<mlir::gpu::GPUModuleOp>(getKernelModuleName());
-  if (!gpuMod) {
-    // If already a gpu.binary then stop the check here.
-    if (symTab.lookup<mlir::gpu::BinaryOp>(getKernelModuleName()))
+  return mlir::success();
+}
+
+mlir::LogicalResult cuf::RegisterKernelOp::verifySymbolUses(
+    mlir::SymbolTableCollection &symbolTables) {
+  auto module = getOperation()->getParentOfType<mlir::ModuleOp>();
+  auto gpuModule = symbolTables.lookupSymbolIn<mlir::gpu::GPUModuleOp>(
+      module, getKernelModuleName());
+  if (!gpuModule) {
+    if (symbolTables.lookupSymbolIn<mlir::gpu::BinaryOp>(module,
+                                                         getKernelModuleName()))
       return mlir::success();
     return emitOpError("gpu module not found");
   }
 
-  mlir::SymbolTable gpuSymTab(gpuMod);
-  if (auto func = gpuSymTab.lookup<mlir::gpu::GPUFuncOp>(getKernelName())) {
-    if (!func.isKernel())
+  if (auto function = symbolTables.lookupSymbolIn<mlir::gpu::GPUFuncOp>(
+          gpuModule, getKernelName())) {
+    if (!function.isKernel())
       return emitOpError("only kernel gpu.func can be registered");
     return mlir::success();
-  } else if (auto func =
-                 gpuSymTab.lookup<mlir::LLVM::LLVMFuncOp>(getKernelName())) {
-    if (!func->getAttrOfType<mlir::UnitAttr>(
+  }
+
+  if (auto function = symbolTables.lookupSymbolIn<mlir::LLVM::LLVMFuncOp>(
+          gpuModule, getKernelName())) {
+    if (!function->getAttrOfType<mlir::UnitAttr>(
             mlir::gpu::GPUDialect::getKernelFuncAttrName()))
       return emitOpError("only gpu.kernel llvm.func can be registered");
     return mlir::success();
   }
+
   return emitOpError("device function not found");
 }
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/230588