[Flang][Lower] Attach target_cpu and target_features attributes to ML…

…IR functions (#78289)

This patch forwards the target CPU and features information from the
Flang frontend to MLIR func.func operation attributes, which are later
used to populate the target_cpu and target_features llvm.func
attributes.

This is achieved in two stages:

1. Introduce the `fir.target_cpu` and `fir.target_features` module
attributes with information from the target machine immediately after
the initial creation of the MLIR module in the lowering bridge.

2. Update the target rewrite flang pass to get this information from the
module and pass it along to all func.func MLIR operations, respectively
as attributes named `target_cpu` and `target_features`. These attributes
will be automatically picked up during Func to LLVM dialect lowering and
used to initialize the corresponding llvm.func named attributes.

The target rewrite and FIR to LLVM lowering passes are updated with the
ability to override these module attributes, and the `CodeGenSpecifics`
optimizer class is augmented to make this information available to
target-specific MLIR transformations.

This completes a full flow by which target CPU and features make it all
the way from compiler options to LLVM IR function attributes.

flang/include/flang/Lower/Bridge.h

@@ -24,7 +24,7 @@
 #include <set>
 
 namespace llvm {
-class DataLayout;
+class TargetMachine;
 } // namespace llvm
 
 namespace Fortran {
@@ -65,11 +65,11 @@ class LoweringBridge {
          const Fortran::lower::LoweringOptions &loweringOptions,
          const std::vector<Fortran::lower::EnvironmentDefault> &envDefaults,
          const Fortran::common::LanguageFeatureControl &languageFeatures,
-         const llvm::DataLayout *dataLayout = nullptr) {
+         const llvm::TargetMachine &targetMachine) {
     return LoweringBridge(ctx, semanticsContext, defaultKinds, intrinsics,
                           targetCharacteristics, allCooked, triple, kindMap,
                           loweringOptions, envDefaults, languageFeatures,
-                          dataLayout);
+                          targetMachine);
   }
 
   //===--------------------------------------------------------------------===//
@@ -148,7 +148,7 @@ class LoweringBridge {
       const Fortran::lower::LoweringOptions &loweringOptions,
       const std::vector<Fortran::lower::EnvironmentDefault> &envDefaults,
       const Fortran::common::LanguageFeatureControl &languageFeatures,
-      const llvm::DataLayout *dataLayout);
+      const llvm::TargetMachine &targetMachine);
   LoweringBridge() = delete;
   LoweringBridge(const LoweringBridge &) = delete;
 

flang/include/flang/Optimizer/CodeGen/CGPasses.td

@@ -29,6 +29,10 @@ def FIRToLLVMLowering : Pass<"fir-to-llvm-ir", "mlir::ModuleOp"> {
            "Override module's target triple.">,
     Option<"forcedDataLayout", "datalayout", "std::string", /*default=*/"",
            "Override module's data layout.">,
+    Option<"forcedTargetCPU", "target-cpu", "std::string", /*default=*/"",
+           "Override module's target CPU.">,
+    Option<"forcedTargetFeatures", "target-features", "std::string",
+           /*default=*/"", "Override module's target features.">,
     Option<"applyTBAA", "apply-tbaa", "bool", /*default=*/"false",
            "Attach TBAA tags to memory accessing operations.">
   ];
@@ -60,6 +64,10 @@ def TargetRewritePass : Pass<"target-rewrite", "mlir::ModuleOp"> {
   let options = [
     Option<"forcedTargetTriple", "target", "std::string", /*default=*/"",
            "Override module's target triple.">,
+    Option<"forcedTargetCPU", "target-cpu", "std::string", /*default=*/"",
+           "Override module's target CPU.">,
+    Option<"forcedTargetFeatures", "target-features", "std::string",
+           /*default=*/"", "Override module's target features.">,
     Option<"noCharacterConversion", "no-character-conversion",
            "bool", /*default=*/"false",
            "Disable target-specific conversion of CHARACTER.">,

flang/include/flang/Optimizer/CodeGen/Target.h

@@ -15,6 +15,7 @@
 
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/Support/KindMapping.h"
+#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "llvm/TargetParser/Triple.h"
 #include <memory>
@@ -70,17 +71,19 @@ class CodeGenSpecifics {
   using TypeAndAttr = std::tuple<mlir::Type, Attributes>;
   using Marshalling = std::vector<TypeAndAttr>;
 
-  static std::unique_ptr<CodeGenSpecifics> get(mlir::MLIRContext *ctx,
-                                               llvm::Triple &&trp,
-                                               KindMapping &&kindMap,
-                                               const mlir::DataLayout &dl);
+  static std::unique_ptr<CodeGenSpecifics>
+  get(mlir::MLIRContext *ctx, llvm::Triple &&trp, KindMapping &&kindMap,
+      llvm::StringRef targetCPU, mlir::LLVM::TargetFeaturesAttr targetFeatures,
+      const mlir::DataLayout &dl);
 
   static TypeAndAttr getTypeAndAttr(mlir::Type t) { return TypeAndAttr{t, {}}; }
 
   CodeGenSpecifics(mlir::MLIRContext *ctx, llvm::Triple &&trp,
-                   KindMapping &&kindMap, const mlir::DataLayout &dl)
+                   KindMapping &&kindMap, llvm::StringRef targetCPU,
+                   mlir::LLVM::TargetFeaturesAttr targetFeatures,
+                   const mlir::DataLayout &dl)
       : context{*ctx}, triple{std::move(trp)}, kindMap{std::move(kindMap)},
-        dataLayout{&dl} {}
+        targetCPU{targetCPU}, targetFeatures{targetFeatures}, dataLayout{&dl} {}
   CodeGenSpecifics() = delete;
   virtual ~CodeGenSpecifics() {}
 
@@ -161,6 +164,12 @@ class CodeGenSpecifics {
   // Returns width in bits of C/C++ 'int' type size.
   virtual unsigned char getCIntTypeWidth() const = 0;
 
+  llvm::StringRef getTargetCPU() const { return targetCPU; }
+
+  mlir::LLVM::TargetFeaturesAttr getTargetFeatures() const {
+    return targetFeatures;
+  }
+
   const mlir::DataLayout &getDataLayout() const {
     assert(dataLayout && "dataLayout must be set");
     return *dataLayout;
@@ -170,6 +179,8 @@ class CodeGenSpecifics {
   mlir::MLIRContext &context;
   llvm::Triple triple;
   KindMapping kindMap;
+  llvm::StringRef targetCPU;
+  mlir::LLVM::TargetFeaturesAttr targetFeatures;
   const mlir::DataLayout *dataLayout = nullptr;
 };
 

flang/include/flang/Optimizer/Dialect/Support/FIRContext.h

@@ -17,6 +17,7 @@
 #ifndef FORTRAN_OPTIMIZER_SUPPORT_FIRCONTEXT_H
 #define FORTRAN_OPTIMIZER_SUPPORT_FIRCONTEXT_H
 
+#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/TargetParser/Triple.h"
 
@@ -50,6 +51,19 @@ KindMapping getKindMapping(mlir::ModuleOp mod);
 /// If a ModuleOp cannot be reached, the function returns default KindMapping.
 KindMapping getKindMapping(mlir::Operation *op);
 
+/// Set the target CPU for the module. `cpu` must not be deallocated while
+/// module `mod` is still live.
+void setTargetCPU(mlir::ModuleOp mod, llvm::StringRef cpu);
+
+/// Get the target CPU string from the Module or return a null reference.
+llvm::StringRef getTargetCPU(mlir::ModuleOp mod);
+
+/// Set the target features for the module.
+void setTargetFeatures(mlir::ModuleOp mod, llvm::StringRef features);
+
+/// Get the target features from the Module.
+mlir::LLVM::TargetFeaturesAttr getTargetFeatures(mlir::ModuleOp mod);
+
 /// Helper for determining the target from the host, etc. Tools may use this
 /// function to provide a consistent interpretation of the `--target=<string>`
 /// command-line option.

flang/lib/Frontend/FrontendActions.cpp

@@ -285,16 +285,14 @@ bool CodeGenAction::beginSourceFileAction() {
       ci.getSemanticsContext().defaultKinds();
   fir::KindMapping kindMap(mlirCtx.get(), llvm::ArrayRef<fir::KindTy>{
                                               fir::fromDefaultKinds(defKinds)});
-  const llvm::DataLayout &dl = targetMachine.createDataLayout();
-
   lower::LoweringBridge lb = Fortran::lower::LoweringBridge::create(
       *mlirCtx, ci.getSemanticsContext(), defKinds,
       ci.getSemanticsContext().intrinsics(),
       ci.getSemanticsContext().targetCharacteristics(),
       ci.getParsing().allCooked(), ci.getInvocation().getTargetOpts().triple,
       kindMap, ci.getInvocation().getLoweringOpts(),
       ci.getInvocation().getFrontendOpts().envDefaults,
-      ci.getInvocation().getFrontendOpts().features, &dl);
+      ci.getInvocation().getFrontendOpts().features, targetMachine);
 
   // Fetch module from lb, so we can set
   mlirModule = std::make_unique<mlir::ModuleOp>(lb.getModule());

flang/lib/Lower/Bridge.cpp

@@ -64,6 +64,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Target/TargetMachine.h"
 #include <optional>
 
 #define DEBUG_TYPE "flang-lower-bridge"
@@ -5092,7 +5093,7 @@ Fortran::lower::LoweringBridge::LoweringBridge(
     const Fortran::lower::LoweringOptions &loweringOptions,
     const std::vector<Fortran::lower::EnvironmentDefault> &envDefaults,
     const Fortran::common::LanguageFeatureControl &languageFeatures,
-    const llvm::DataLayout *dataLayout)
+    const llvm::TargetMachine &targetMachine)
     : semanticsContext{semanticsContext}, defaultKinds{defaultKinds},
       intrinsics{intrinsics}, targetCharacteristics{targetCharacteristics},
       cooked{&cooked}, context{context}, kindMap{kindMap},
@@ -5148,6 +5149,8 @@ Fortran::lower::LoweringBridge::LoweringBridge(
   assert(module.get() && "module was not created");
   fir::setTargetTriple(*module.get(), triple);
   fir::setKindMapping(*module.get(), kindMap);
-  if (dataLayout)
-    fir::support::setMLIRDataLayout(*module.get(), *dataLayout);
+  fir::setTargetCPU(*module.get(), targetMachine.getTargetCPU());
+  fir::setTargetFeatures(*module.get(), targetMachine.getTargetFeatureString());
+  fir::support::setMLIRDataLayout(*module.get(),
+                                  targetMachine.createDataLayout());
 }

flang/lib/Optimizer/CodeGen/CodeGen.cpp

@@ -3879,6 +3879,12 @@ class FIRToLLVMLowering
       fir::support::setMLIRDataLayout(mod, dl);
     }
 
+    if (!forcedTargetCPU.empty())
+      fir::setTargetCPU(mod, forcedTargetCPU);
+
+    if (!forcedTargetFeatures.empty())
+      fir::setTargetFeatures(mod, forcedTargetFeatures);
+
     // Run dynamic pass pipeline for converting Math dialect
     // operations into other dialects (llvm, func, etc.).
     // Some conversions of Math operations cannot be done

flang/lib/Optimizer/CodeGen/Target.cpp

@@ -1057,51 +1057,57 @@ struct TargetLoongArch64 : public GenericTarget<TargetLoongArch64> {
 // TODO: Add other targets to this file as needed.
 std::unique_ptr<fir::CodeGenSpecifics>
 fir::CodeGenSpecifics::get(mlir::MLIRContext *ctx, llvm::Triple &&trp,
-                           KindMapping &&kindMap, const mlir::DataLayout &dl) {
+                           KindMapping &&kindMap, llvm::StringRef targetCPU,
+                           mlir::LLVM::TargetFeaturesAttr targetFeatures,
+                           const mlir::DataLayout &dl) {
   switch (trp.getArch()) {
   default:
     break;
   case llvm::Triple::ArchType::x86:
     if (trp.isOSWindows())
       return std::make_unique<TargetI386Win>(ctx, std::move(trp),
-                                             std::move(kindMap), dl);
+                                             std::move(kindMap), targetCPU,
+                                             targetFeatures, dl);
     else
       return std::make_unique<TargetI386>(ctx, std::move(trp),
-                                          std::move(kindMap), dl);
+                                          std::move(kindMap), targetCPU,
+                                          targetFeatures, dl);
   case llvm::Triple::ArchType::x86_64:
     if (trp.isOSWindows())
       return std::make_unique<TargetX86_64Win>(ctx, std::move(trp),
-                                               std::move(kindMap), dl);
+                                               std::move(kindMap), targetCPU,
+                                               targetFeatures, dl);
     else
       return std::make_unique<TargetX86_64>(ctx, std::move(trp),
-                                            std::move(kindMap), dl);
+                                            std::move(kindMap), targetCPU,
+                                            targetFeatures, dl);
   case llvm::Triple::ArchType::aarch64:
-    return std::make_unique<TargetAArch64>(ctx, std::move(trp),
-                                           std::move(kindMap), dl);
+    return std::make_unique<TargetAArch64>(
+        ctx, std::move(trp), std::move(kindMap), targetCPU, targetFeatures, dl);
   case llvm::Triple::ArchType::ppc64:
-    return std::make_unique<TargetPPC64>(ctx, std::move(trp),
-                                         std::move(kindMap), dl);
+    return std::make_unique<TargetPPC64>(
+        ctx, std::move(trp), std::move(kindMap), targetCPU, targetFeatures, dl);
   case llvm::Triple::ArchType::ppc64le:
-    return std::make_unique<TargetPPC64le>(ctx, std::move(trp),
-                                           std::move(kindMap), dl);
+    return std::make_unique<TargetPPC64le>(
+        ctx, std::move(trp), std::move(kindMap), targetCPU, targetFeatures, dl);
   case llvm::Triple::ArchType::sparc:
-    return std::make_unique<TargetSparc>(ctx, std::move(trp),
-                                         std::move(kindMap), dl);
+    return std::make_unique<TargetSparc>(
+        ctx, std::move(trp), std::move(kindMap), targetCPU, targetFeatures, dl);
   case llvm::Triple::ArchType::sparcv9:
-    return std::make_unique<TargetSparcV9>(ctx, std::move(trp),
-                                           std::move(kindMap), dl);
+    return std::make_unique<TargetSparcV9>(
+        ctx, std::move(trp), std::move(kindMap), targetCPU, targetFeatures, dl);
   case llvm::Triple::ArchType::riscv64:
-    return std::make_unique<TargetRISCV64>(ctx, std::move(trp),
-                                           std::move(kindMap), dl);
+    return std::make_unique<TargetRISCV64>(
+        ctx, std::move(trp), std::move(kindMap), targetCPU, targetFeatures, dl);
   case llvm::Triple::ArchType::amdgcn:
-    return std::make_unique<TargetAMDGPU>(ctx, std::move(trp),
-                                          std::move(kindMap), dl);
+    return std::make_unique<TargetAMDGPU>(
+        ctx, std::move(trp), std::move(kindMap), targetCPU, targetFeatures, dl);
   case llvm::Triple::ArchType::nvptx64:
-    return std::make_unique<TargetNVPTX>(ctx, std::move(trp),
-                                         std::move(kindMap), dl);
+    return std::make_unique<TargetNVPTX>(
+        ctx, std::move(trp), std::move(kindMap), targetCPU, targetFeatures, dl);
   case llvm::Triple::ArchType::loongarch64:
-    return std::make_unique<TargetLoongArch64>(ctx, std::move(trp),
-                                               std::move(kindMap), dl);
+    return std::make_unique<TargetLoongArch64>(
+        ctx, std::move(trp), std::move(kindMap), targetCPU, targetFeatures, dl);
   }
   TODO(mlir::UnknownLoc::get(ctx), "target not implemented");
 }

flang/lib/Optimizer/CodeGen/TargetRewrite.cpp

@@ -90,6 +90,12 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
     if (!forcedTargetTriple.empty())
       fir::setTargetTriple(mod, forcedTargetTriple);
 
+    if (!forcedTargetCPU.empty())
+      fir::setTargetCPU(mod, forcedTargetCPU);
+
+    if (!forcedTargetFeatures.empty())
+      fir::setTargetFeatures(mod, forcedTargetFeatures);
+
     // TargetRewrite will require querying the type storage sizes, if it was
     // not set already, create a DataLayoutSpec for the ModuleOp now.
     std::optional<mlir::DataLayout> dl =
@@ -102,9 +108,9 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
       return;
     }
 
-    auto specifics =
-        fir::CodeGenSpecifics::get(mod.getContext(), fir::getTargetTriple(mod),
-                                   fir::getKindMapping(mod), *dl);
+    auto specifics = fir::CodeGenSpecifics::get(
+        mod.getContext(), fir::getTargetTriple(mod), fir::getKindMapping(mod),
+        fir::getTargetCPU(mod), fir::getTargetFeatures(mod), *dl);
 
     setMembers(specifics.get(), &rewriter, &*dl);
 
@@ -666,8 +672,21 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
   /// As the type signature is being changed, this must also update the
   /// function itself to use any new arguments, etc.
   mlir::LogicalResult convertTypes(mlir::ModuleOp mod) {
-    for (auto fn : mod.getOps<mlir::func::FuncOp>())
+    mlir::MLIRContext *ctx = mod->getContext();
+    auto targetCPU = specifics->getTargetCPU();
+    mlir::StringAttr targetCPUAttr =
+        targetCPU.empty() ? nullptr : mlir::StringAttr::get(ctx, targetCPU);
+    auto targetFeaturesAttr = specifics->getTargetFeatures();
+
+    for (auto fn : mod.getOps<mlir::func::FuncOp>()) {
+      if (targetCPUAttr)
+        fn->setAttr("target_cpu", targetCPUAttr);
+
+      if (targetFeaturesAttr)
+        fn->setAttr("target_features", targetFeaturesAttr);
+
       convertSignature(fn);
+    }
     return mlir::success();
   }
 

flang/lib/Optimizer/CodeGen/TypeConverter.cpp

@@ -32,9 +32,9 @@ LLVMTypeConverter::LLVMTypeConverter(mlir::ModuleOp module, bool applyTBAA,
                                      const mlir::DataLayout &dl)
     : mlir::LLVMTypeConverter(module.getContext()),
       kindMapping(getKindMapping(module)),
-      specifics(CodeGenSpecifics::get(module.getContext(),
-                                      getTargetTriple(module),
-                                      getKindMapping(module), dl)),
+      specifics(CodeGenSpecifics::get(
+          module.getContext(), getTargetTriple(module), getKindMapping(module),
+          getTargetCPU(module), getTargetFeatures(module), dl)),
       tbaaBuilder(std::make_unique<TBAABuilder>(module->getContext(), applyTBAA,
                                                 forceUnifiedTBAATree)) {
   LLVM_DEBUG(llvm::dbgs() << "FIR type converter\n");

flang/lib/Optimizer/Dialect/Support/FIRContext.cpp

@@ -60,6 +60,42 @@ fir::KindMapping fir::getKindMapping(mlir::Operation *op) {
   return getKindMapping(moduleOp);
 }
 
+static constexpr const char *targetCpuName = "fir.target_cpu";
+
+void fir::setTargetCPU(mlir::ModuleOp mod, llvm::StringRef cpu) {
+  if (cpu.empty())
+    return;
+
+  auto *ctx = mod.getContext();
+  mod->setAttr(targetCpuName, mlir::StringAttr::get(ctx, cpu));
+}
+
+llvm::StringRef fir::getTargetCPU(mlir::ModuleOp mod) {
+  if (auto attr = mod->getAttrOfType<mlir::StringAttr>(targetCpuName))
+    return attr.getValue();
+
+  return {};
+}
+
+static constexpr const char *targetFeaturesName = "fir.target_features";
+
+void fir::setTargetFeatures(mlir::ModuleOp mod, llvm::StringRef features) {
+  if (features.empty())
+    return;
+
+  auto *ctx = mod.getContext();
+  mod->setAttr(targetFeaturesName,
+               mlir::LLVM::TargetFeaturesAttr::get(ctx, features));
+}
+
+mlir::LLVM::TargetFeaturesAttr fir::getTargetFeatures(mlir::ModuleOp mod) {
+  if (auto attr = mod->getAttrOfType<mlir::LLVM::TargetFeaturesAttr>(
+          targetFeaturesName))
+    return attr;
+
+  return {};
+}
+
 std::string fir::determineTargetTriple(llvm::StringRef triple) {
   // Treat "" or "default" as stand-ins for the default machine.
   if (triple.empty() || triple == "default")

flang/test/Driver/save-mlir-temps.f90

@@ -51,9 +51,9 @@
 ! Content to check from the MLIR outputs
 !--------------------------
 ! MLIR-FIR-NOT: llvm.func
-! MLIR-FIR: func.func @{{.*}}main(){{.*}}{
+! MLIR-FIR: func.func @{{.*}}main(){{.*}}
 
-! MLIR-FIR-NOT: func.func
-! MLIR-LLVMIR: llvm.func @{{.*}}main(){{.*}}{
+! MLIR-LLVMIR-NOT: func.func
+! MLIR-LLVMIR: llvm.func @{{.*}}main(){{.*}}
 
 end program