[flang][cuda] Adapt ExternalNameConversion to work in gpu module (#117039)

Author: clementval

flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp

@@ -12,6 +12,7 @@
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Support/InternalNames.h"
 #include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Pass/Pass.h"
@@ -58,26 +59,36 @@ void ExternalNameConversionPass::runOnOperation() {
   auto *context = &getContext();
 
   llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr> remappings;
-  // Update names of external Fortran functions and names of Common Block
-  // globals.
-  for (auto &funcOrGlobal : op->getRegion(0).front()) {
-    if (llvm::isa<mlir::func::FuncOp>(funcOrGlobal) ||
-        llvm::isa<fir::GlobalOp>(funcOrGlobal)) {
-      auto symName = funcOrGlobal.getAttrOfType<mlir::StringAttr>(
-          mlir::SymbolTable::getSymbolAttrName());
-      auto deconstructedName = fir::NameUniquer::deconstruct(symName);
-      if (fir::NameUniquer::isExternalFacingUniquedName(deconstructedName)) {
-        auto newName =
-            mangleExternalName(deconstructedName, appendUnderscoreOpt);
-        auto newAttr = mlir::StringAttr::get(context, newName);
-        mlir::SymbolTable::setSymbolName(&funcOrGlobal, newAttr);
-        auto newSymRef = mlir::FlatSymbolRefAttr::get(newAttr);
-        remappings.try_emplace(symName, newSymRef);
-        if (llvm::isa<mlir::func::FuncOp>(funcOrGlobal))
-          funcOrGlobal.setAttr(fir::getInternalFuncNameAttrName(), symName);
+
+  auto renameFuncOrGlobalInModule = [&](mlir::Operation *module) {
+    for (auto &funcOrGlobal : module->getRegion(0).front()) {
+      if (llvm::isa<mlir::func::FuncOp>(funcOrGlobal) ||
+          llvm::isa<fir::GlobalOp>(funcOrGlobal)) {
+        auto symName = funcOrGlobal.getAttrOfType<mlir::StringAttr>(
+            mlir::SymbolTable::getSymbolAttrName());
+        auto deconstructedName = fir::NameUniquer::deconstruct(symName);
+        if (fir::NameUniquer::isExternalFacingUniquedName(deconstructedName)) {
+          auto newName =
+              mangleExternalName(deconstructedName, appendUnderscoreOpt);
+          auto newAttr = mlir::StringAttr::get(context, newName);
+          mlir::SymbolTable::setSymbolName(&funcOrGlobal, newAttr);
+          auto newSymRef = mlir::FlatSymbolRefAttr::get(newAttr);
+          remappings.try_emplace(symName, newSymRef);
+          if (llvm::isa<mlir::func::FuncOp>(funcOrGlobal))
+            funcOrGlobal.setAttr(fir::getInternalFuncNameAttrName(), symName);
+        }
       }
     }
-  }
+  };
+
+  // Update names of external Fortran functions and names of Common Block
+  // globals.
+  renameFuncOrGlobalInModule(op);
+
+  // Do the same in GPU modules.
+  if (auto mod = mlir::dyn_cast_or_null<mlir::ModuleOp>(*op))
+    for (auto gpuMod : mod.getOps<mlir::gpu::GPUModuleOp>())
+      renameFuncOrGlobalInModule(gpuMod);
 
   if (remappings.empty())
     return;

flang/test/Fir/CUDA/cuda-extranal-mangling.mlir

@@ -0,0 +1,13 @@
+// RUN: fir-opt --split-input-file --external-name-interop %s | FileCheck %s
+
+gpu.module @cuda_device_mod {
+  gpu.func @_QPfoo() {
+    fir.call @_QPthreadfence() fastmath<contract> : () -> ()
+    gpu.return
+  }
+  func.func private @_QPthreadfence() attributes {cuf.proc_attr = #cuf.cuda_proc<device>}
+}
+
+// CHECK-LABEL: gpu.func @_QPfoo
+// CHECK: fir.call @threadfence_()
+// CHECK: func.func private @threadfence_()