From a136c793e316bcd00f6ff34243b585f93573cdec Mon Sep 17 00:00:00 2001 From: Valentin Clement Date: Mon, 9 Dec 2024 13:58:03 -0800 Subject: [PATCH 1/2] [flang][cuda] Additional update to ExternalNameConversion --- .../Transforms/ExternalNameConversion.cpp | 57 +++++++++++-------- .../test/Fir/CUDA/cuda-extranal-mangling.mlir | 16 +++++- 2 files changed, 46 insertions(+), 27 deletions(-) diff --git a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp index cfd90ff723793..eaa40a35e3860 100644 --- a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp +++ b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp @@ -60,23 +60,30 @@ void ExternalNameConversionPass::runOnOperation() { llvm::DenseMap remappings; + auto processFctOrGlobal = [&](mlir::Operation &funcOrGlobal) { + auto symName = funcOrGlobal.getAttrOfType( + 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(funcOrGlobal)) + funcOrGlobal.setAttr(fir::getInternalFuncNameAttrName(), symName); + } + }; + auto renameFuncOrGlobalInModule = [&](mlir::Operation *module) { - for (auto &funcOrGlobal : module->getRegion(0).front()) { - if (llvm::isa(funcOrGlobal) || - llvm::isa(funcOrGlobal)) { - auto symName = funcOrGlobal.getAttrOfType( - 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(funcOrGlobal)) - funcOrGlobal.setAttr(fir::getInternalFuncNameAttrName(), symName); - } + for (auto &op : module->getRegion(0).front()) { + if (mlir::isa(op)) { + processFctOrGlobal(op); + } else if (auto gpuMod = mlir::dyn_cast(op)) { + for (auto &gpuOp : gpuMod.getBodyRegion().front()) + if (mlir::isa(gpuOp)) + processFctOrGlobal(gpuOp); } } }; @@ -85,11 +92,6 @@ void ExternalNameConversionPass::runOnOperation() { // globals. renameFuncOrGlobalInModule(op); - // Do the same in GPU modules. - if (auto mod = mlir::dyn_cast_or_null(*op)) - for (auto gpuMod : mod.getOps()) - renameFuncOrGlobalInModule(gpuMod); - if (remappings.empty()) return; @@ -97,11 +99,16 @@ void ExternalNameConversionPass::runOnOperation() { op.walk([&remappings](mlir::Operation *nestedOp) { llvm::SmallVector> updates; for (const mlir::NamedAttribute &attr : nestedOp->getAttrDictionary()) - if (auto symRef = llvm::dyn_cast(attr.getValue())) - if (auto remap = remappings.find(symRef.getRootReference()); - remap != remappings.end()) + if (auto symRef = llvm::dyn_cast(attr.getValue())) { + if (auto remap = remappings.find(symRef.getLeafReference()); + remap != remappings.end()) { + mlir::SymbolRefAttr symAttr = mlir::FlatSymbolRefAttr(remap->second); + if (mlir::isa(nestedOp)) + symAttr = mlir::SymbolRefAttr::get(symRef.getRootReference(), {mlir::FlatSymbolRefAttr(remap->second)}); updates.emplace_back(std::pair{ - attr.getName(), mlir::SymbolRefAttr(remap->second)}); + attr.getName(), symAttr}); + } + } for (auto update : updates) nestedOp->setAttr(update.first, update.second); }); diff --git a/flang/test/Fir/CUDA/cuda-extranal-mangling.mlir b/flang/test/Fir/CUDA/cuda-extranal-mangling.mlir index 551a89a7018c2..cd028a201e6fa 100644 --- a/flang/test/Fir/CUDA/cuda-extranal-mangling.mlir +++ b/flang/test/Fir/CUDA/cuda-extranal-mangling.mlir @@ -1,13 +1,25 @@ // RUN: fir-opt --split-input-file --external-name-interop %s | FileCheck %s +module @mod attributes {gpu.container_module} { + gpu.module @cuda_device_mod { - gpu.func @_QPfoo() { + gpu.func @_QPfoo() kernel { fir.call @_QPthreadfence() fastmath : () -> () gpu.return } func.func private @_QPthreadfence() attributes {cuf.proc_attr = #cuf.cuda_proc} } -// CHECK-LABEL: gpu.func @_QPfoo +func.func @test() -> () { + %0 = llvm.mlir.constant(0 : i64) : i64 + %1 = llvm.mlir.constant(0 : i32) : i32 + gpu.launch_func @cuda_device_mod::@_QPfoo blocks in (%0, %0, %0) threads in (%0, %0, %0) : i64 dynamic_shared_memory_size %1 + return +} + +// CHECK-LABEL: gpu.func @foo_() // CHECK: fir.call @threadfence_() // CHECK: func.func private @threadfence_() +// CHECK: gpu.launch_func @cuda_device_mod::@foo_ + +} From 3c336e34b692e1d98dadb903b1950ffbabaa89a4 Mon Sep 17 00:00:00 2001 From: Valentin Clement Date: Mon, 9 Dec 2024 14:16:59 -0800 Subject: [PATCH 2/2] clang-format --- .../Optimizer/Transforms/ExternalNameConversion.cpp | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp index eaa40a35e3860..4f6974ee52695 100644 --- a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp +++ b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp @@ -62,11 +62,10 @@ void ExternalNameConversionPass::runOnOperation() { auto processFctOrGlobal = [&](mlir::Operation &funcOrGlobal) { auto symName = funcOrGlobal.getAttrOfType( - mlir::SymbolTable::getSymbolAttrName()); + mlir::SymbolTable::getSymbolAttrName()); auto deconstructedName = fir::NameUniquer::deconstruct(symName); if (fir::NameUniquer::isExternalFacingUniquedName(deconstructedName)) { - auto newName = - mangleExternalName(deconstructedName, appendUnderscoreOpt); + auto newName = mangleExternalName(deconstructedName, appendUnderscoreOpt); auto newAttr = mlir::StringAttr::get(context, newName); mlir::SymbolTable::setSymbolName(&funcOrGlobal, newAttr); auto newSymRef = mlir::FlatSymbolRefAttr::get(newAttr); @@ -82,7 +81,8 @@ void ExternalNameConversionPass::runOnOperation() { processFctOrGlobal(op); } else if (auto gpuMod = mlir::dyn_cast(op)) { for (auto &gpuOp : gpuMod.getBodyRegion().front()) - if (mlir::isa(gpuOp)) + if (mlir::isa(gpuOp)) processFctOrGlobal(gpuOp); } } @@ -104,7 +104,9 @@ void ExternalNameConversionPass::runOnOperation() { remap != remappings.end()) { mlir::SymbolRefAttr symAttr = mlir::FlatSymbolRefAttr(remap->second); if (mlir::isa(nestedOp)) - symAttr = mlir::SymbolRefAttr::get(symRef.getRootReference(), {mlir::FlatSymbolRefAttr(remap->second)}); + symAttr = mlir::SymbolRefAttr::get( + symRef.getRootReference(), + {mlir::FlatSymbolRefAttr(remap->second)}); updates.emplace_back(std::pair{ attr.getName(), symAttr}); }