flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp

   1 //===- ExternalNameConversion.cpp -- convert name with external convention ===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "flang/Common/Fortran.h"
  10 #include "flang/Optimizer/Dialect/FIRDialect.h"
  11 #include "flang/Optimizer/Dialect/FIROps.h"
  12 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
  13 #include "flang/Optimizer/Support/InternalNames.h"
  14 #include "flang/Optimizer/Transforms/Passes.h"
  15 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
  16 #include "mlir/IR/Attributes.h"
  17 #include "mlir/IR/SymbolTable.h"
  18 #include "mlir/Pass/Pass.h"
  19
  20 namespace fir {
  21 #define GEN_PASS_DEF_EXTERNALNAMECONVERSION
  22 #include "flang/Optimizer/Transforms/Passes.h.inc"
  23 } // namespace fir
  24
  25 using namespace mlir;
  26
  27 //===----------------------------------------------------------------------===//
  28 // Helper functions
  29 //===----------------------------------------------------------------------===//
  30
  31 /// Mangle the name with gfortran convention.
  32 std::string
  33 mangleExternalName(const std::pair<fir::NameUniquer::NameKind,
  34                                    fir::NameUniquer::DeconstructedName>
  35                        result,
  36                    bool appendUnderscore) {
  37   if (result.first == fir::NameUniquer::NameKind::COMMON &&
  38       result.second.name.empty())
  39     return Fortran::common::blankCommonObjectName;
  40   return Fortran::common::GetExternalAssemblyName(result.second.name,
  41                                                   appendUnderscore);
  42 }
  43
  44 namespace {
  45
  46 class ExternalNameConversionPass
  47     : public fir::impl::ExternalNameConversionBase<ExternalNameConversionPass> {
  48 public:
  49   using ExternalNameConversionBase<
  50       ExternalNameConversionPass>::ExternalNameConversionBase;
  51
  52   mlir::ModuleOp getModule() { return getOperation(); }
  53   void runOnOperation() override;
  54 };
  55 } // namespace
  56
  57 void ExternalNameConversionPass::runOnOperation() {
  58   auto op = getOperation();
  59   auto *context = &getContext();
  60
  61   llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr> remappings;
  62
  63   auto processFctOrGlobal = [&](mlir::Operation &funcOrGlobal) {
  64     auto symName = funcOrGlobal.getAttrOfType<mlir::StringAttr>(
  65         mlir::SymbolTable::getSymbolAttrName());
  66     auto deconstructedName = fir::NameUniquer::deconstruct(symName);
  67     if (fir::NameUniquer::isExternalFacingUniquedName(deconstructedName)) {
  68       auto newName = mangleExternalName(deconstructedName, appendUnderscoreOpt);
  69       auto newAttr = mlir::StringAttr::get(context, newName);
  70       mlir::SymbolTable::setSymbolName(&funcOrGlobal, newAttr);
  71       auto newSymRef = mlir::FlatSymbolRefAttr::get(newAttr);
  72       remappings.try_emplace(symName, newSymRef);
  73       if (llvm::isa<mlir::func::FuncOp>(funcOrGlobal))
  74         funcOrGlobal.setAttr(fir::getInternalFuncNameAttrName(), symName);
  75     }
  76   };
  77
  78   auto renameFuncOrGlobalInModule = [&](mlir::Operation *module) {
  79     for (auto &op : module->getRegion(0).front()) {
  80       if (mlir::isa<mlir::func::FuncOp, fir::GlobalOp>(op)) {
  81         processFctOrGlobal(op);
  82       } else if (auto gpuMod = mlir::dyn_cast<mlir::gpu::GPUModuleOp>(op)) {
  83         for (auto &gpuOp : gpuMod.getBodyRegion().front())
  84           if (mlir::isa<mlir::func::FuncOp, fir::GlobalOp,
  85                         mlir::gpu::GPUFuncOp>(gpuOp))
  86             processFctOrGlobal(gpuOp);
  87       }
  88     }
  89   };
  90
  91   // Update names of external Fortran functions and names of Common Block
  92   // globals.
  93   renameFuncOrGlobalInModule(op);
  94
  95   if (remappings.empty())
  96     return;
  97
  98   // Update all uses of the functions and globals that have been renamed.
  99   op.walk([&remappings](mlir::Operation *nestedOp) {
 100     llvm::SmallVector<std::pair<mlir::StringAttr, mlir::SymbolRefAttr>> updates;
 101     for (const mlir::NamedAttribute &attr : nestedOp->getAttrDictionary())
 102       if (auto symRef = llvm::dyn_cast<mlir::SymbolRefAttr>(attr.getValue())) {
 103         if (auto remap = remappings.find(symRef.getLeafReference());
 104             remap != remappings.end()) {
 105           mlir::SymbolRefAttr symAttr = mlir::FlatSymbolRefAttr(remap->second);
 106           if (mlir::isa<mlir::gpu::LaunchFuncOp>(nestedOp))
 107             symAttr = mlir::SymbolRefAttr::get(
 108                 symRef.getRootReference(),
 109                 {mlir::FlatSymbolRefAttr(remap->second)});
 110           updates.emplace_back(std::pair<mlir::StringAttr, mlir::SymbolRefAttr>{
 111               attr.getName(), symAttr});
 112         }
 113       }
 114     for (auto update : updates)
 115       nestedOp->setAttr(update.first, update.second);
 116   });
 117 }