[flang][cuda] Adapt ExternalNameConversion to work in gpu module

[clementval]

No description provided.

[llvmbot]

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

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

---

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

2 Files Affected:

• (modified) flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp (+29-16)
• (added) flang/test/Fir/CUDA/cuda-extranal-mangling.mlir (+13)

diff --git a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
index 648628fd1c9af0..e5919120941fbd 100644
--- a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
+++ b/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,24 +59,36 @@ void ExternalNameConversionPass::runOnOperation() {
   auto *context = &getContext();
 
   llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr> remappings;
+
+  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.
-  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);
-      }
+  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);
     }
   }
 
diff --git a/flang/test/Fir/CUDA/cuda-extranal-mangling.mlir b/flang/test/Fir/CUDA/cuda-extranal-mangling.mlir
new file mode 100644
index 00000000000000..551a89a7018c28
--- /dev/null
+++ b/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_()