[Flang][MLIR][OpenMP] Add support for target-cpu and target-features

This patch implements a complete flow to propagate information on the target
CPU and target features, as obtained from Flang's command line arguments, to
LLVM IR. The proposed approach can be summarized as follows:
  - A reference to the `TargetMachine` object is stored in the
  `LoweringBridge`. This enables the addition of the "target_cpu" and
  "llvm.target_features" MLIR discardable attributes to all generated
  `func.func` operations.
  - The Func to LLVM dialect conversion pass is updated to translate instances
  of "llvm.target_features" string attributes into the expected structured
  attribute representation for the corresponding "target_features" argument of
  `llvm.func` operations. The "target_cpu" is passed through by default to a
  new string attribute of that same name added to the `llvm.func` operation.
  - The new "target_cpu" attribute is translated to a "target-cpu" LLVM IR
  function attribute.
  - The existing `omp.target` OpenMP attribute class, which is used to
  represent these two attributes, is removed, together with all uses.
  - The OpenMP to LLVMIR translation is updated to propagate these attributes
  to outlined target regions, so they match the attributes from the function
  they are originally extracted.

This patch is probably best to be split up into different parts to simplify
reviews, but I think it's good to present the whole approach early to see if
there are any conceptual problems with it, as I will only have bandwidth to
split this patch in January.

Author: skatrak

flang/include/flang/Lower/Bridge.h

@@ -21,10 +21,7 @@
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Dialect/Support/KindMapping.h"
 #include "mlir/IR/BuiltinOps.h"
-
-namespace llvm {
-class DataLayout;
-} // namespace llvm
+#include "llvm/Target/TargetMachine.h"
 
 namespace Fortran {
 namespace common {
@@ -64,11 +61,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);
   }
 
   //===--------------------------------------------------------------------===//
@@ -110,6 +107,8 @@ class LoweringBridge {
     return languageFeatures;
   }
 
+  const llvm::TargetMachine &getTargetMachine() const { return targetMachine; }
+
   /// Create a folding context. Careful: this is very expensive.
   Fortran::evaluate::FoldingContext createFoldingContext() const;
 
@@ -147,7 +146,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;
 
@@ -164,6 +163,7 @@ class LoweringBridge {
   const Fortran::lower::LoweringOptions &loweringOptions;
   const std::vector<Fortran::lower::EnvironmentDefault> &envDefaults;
   const Fortran::common::LanguageFeatureControl &languageFeatures;
+  const llvm::TargetMachine &targetMachine;
 };
 
 } // namespace lower

flang/include/flang/Tools/CrossToolHelpers.h

@@ -103,17 +103,6 @@ void setOffloadModuleInterfaceAttributes(
   }
 }
 
-//  Shares assinging of the OpenMP OffloadModuleInterface and its TargetCPU
-//  attribute accross Flang tools (bbc/flang)
-void setOffloadModuleInterfaceTargetAttribute(mlir::ModuleOp &module,
-    llvm::StringRef targetCPU, llvm::StringRef targetFeatures) {
-  // Should be registered by the OpenMPDialect
-  if (auto offloadMod = llvm::dyn_cast<mlir::omp::OffloadModuleInterface>(
-          module.getOperation())) {
-    offloadMod.setTarget(targetCPU, targetFeatures);
-  }
-}
-
 void setOpenMPVersionAttribute(mlir::ModuleOp &module, int64_t version) {
   module.getOperation()->setAttr(
       mlir::StringAttr::get(module.getContext(), llvm::Twine{"omp.version"}),

flang/lib/Frontend/FrontendActions.cpp

@@ -283,16 +283,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());
@@ -301,9 +299,6 @@ bool CodeGenAction::beginSourceFileAction() {
           Fortran::common::LanguageFeature::OpenMP)) {
     setOffloadModuleInterfaceAttributes(*mlirModule,
                                         ci.getInvocation().getLangOpts());
-    setOffloadModuleInterfaceTargetAttribute(
-        *mlirModule, targetMachine.getTargetCPU(),
-        targetMachine.getTargetFeatureString());
     setOpenMPVersionAttribute(*mlirModule,
                               ci.getInvocation().getLangOpts().OpenMPVersion);
   }

flang/lib/Lower/Bridge.cpp

@@ -4326,6 +4326,17 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     assert(blockId == 0 && "invalid blockId");
     assert(activeConstructStack.empty() && "invalid construct stack state");
 
+    // Set target_cpu and target_features attributes to be passed through to the
+    // llvm.func operation during lowering.
+    const llvm::TargetMachine &targetMachine = bridge.getTargetMachine();
+    if (auto targetCPU = targetMachine.getTargetCPU(); !targetCPU.empty())
+      func->setAttr("target_cpu",
+                    mlir::StringAttr::get(func.getContext(), targetCPU));
+    if (auto targetFeatures = targetMachine.getTargetFeatureString();
+        !targetFeatures.empty())
+      func->setAttr("llvm.target_features",
+                    mlir::StringAttr::get(func.getContext(), targetFeatures));
+
     // Manage floating point exception, halting mode, and rounding mode
     // settings at function entry and exit.
     if (!funit.isMainProgram())
@@ -5091,12 +5102,12 @@ 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},
       loweringOptions{loweringOptions}, envDefaults{envDefaults},
-      languageFeatures{languageFeatures} {
+      languageFeatures{languageFeatures}, targetMachine{targetMachine} {
   // Register the diagnostic handler.
   context.getDiagEngine().registerHandler([](mlir::Diagnostic &diag) {
     llvm::raw_ostream &os = llvm::errs();
@@ -5147,6 +5158,6 @@ 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::support::setMLIRDataLayout(*module.get(),
+                                  targetMachine.createDataLayout());
 }

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

flang/test/Lower/target-features.f90

@@ -0,0 +1,20 @@
+! RUN: %flang_fc1 -emit-fir %s -o - | FileCheck %s --check-prefixes=ALL,NONE
+! RUN: %flang_fc1 -emit-fir -triple amdgcn-amd-amdhsa %s -o - | FileCheck %s --check-prefixes=ALL,TRIPLE
+! RUN: %flang_fc1 -emit-fir -target-cpu gfx90a %s -o - | FileCheck %s --check-prefixes=ALL,CPU
+! RUN: %flang_fc1 -emit-fir -triple amdgcn-amd-amdhsa -target-cpu gfx90a %s -o - | FileCheck %s --check-prefixes=ALL,BOTH
+
+! ALL-LABEL: func.func @_QPfoo()
+
+! NONE-NOT: llvm.target_features
+! NONE-NOT: target_cpu
+
+! TRIPLE-NOT: llvm.target_features
+! TRIPLE-SAME: target_cpu = "generic-hsa"
+
+! CPU-NOT: llvm.target_features
+! CPU-SAME: target_cpu = "gfx90a"
+
+! BOTH-SAME: llvm.target_features = "{{[^"]*}}+gfx90a-insts{{[^"]*}}"
+! BOTH-SAME: target_cpu = "gfx90a"
+subroutine foo
+end subroutine

flang/tools/bbc/bbc.cpp

@@ -325,7 +325,6 @@ static mlir::LogicalResult convertFortranSourceToMLIR(
   auto &defKinds = semanticsContext.defaultKinds();
   fir::KindMapping kindMap(
       &ctx, llvm::ArrayRef<fir::KindTy>{fir::fromDefaultKinds(defKinds)});
-  const llvm::DataLayout &dataLayout = targetMachine.createDataLayout();
   std::string targetTriple = targetMachine.getTargetTriple().normalize();
   // Use default lowering options for bbc.
   Fortran::lower::LoweringOptions loweringOptions{};
@@ -336,7 +335,7 @@ static mlir::LogicalResult convertFortranSourceToMLIR(
       ctx, semanticsContext, defKinds, semanticsContext.intrinsics(),
       semanticsContext.targetCharacteristics(), parsing.allCooked(),
       targetTriple, kindMap, loweringOptions, {},
-      semanticsContext.languageFeatures(), &dataLayout);
+      semanticsContext.languageFeatures(), targetMachine);
   burnside.lower(parseTree, semanticsContext);
   mlir::ModuleOp mlirModule = burnside.getModule();
   if (enableOpenMP) {

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -1418,6 +1418,7 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<I64Attr>:$alignment,
     OptionalAttr<LLVM_VScaleRangeAttr>:$vscale_range,
     OptionalAttr<FramePointerKindAttr>:$frame_pointer,
+    OptionalAttr<StrAttr>:$target_cpu,
     OptionalAttr<LLVM_TargetFeaturesAttr>:$target_features
   );
 

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

@@ -71,16 +71,6 @@ def FlagsAttr : OpenMP_Attr<"Flags", "flags"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
-def TargetAttr : OpenMP_Attr<"Target", "target"> {
-  let parameters = (ins
-    StringRefParameter<>:$target_cpu,
-    StringRefParameter<>:$target_features
-  );
-
-  let assemblyFormat = "`<` struct(params) `>`";
-}
-
-
 class OpenMP_Op<string mnemonic, list<Trait> traits = []> :
       Op<OpenMP_Dialect, mnemonic, traits>;
 

mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td

@@ -204,34 +204,6 @@ def OffloadModuleInterface : OpInterface<"OffloadModuleInterface"> {
                       assumeTeamsOversubscription, assumeThreadsOversubscription,
                       assumeNoThreadState, assumeNoNestedParallelism, openmpDeviceVersion));
       }]>,
-    InterfaceMethod<
-      /*description=*/[{
-        Get the Target attribute on the current module if it exists
-        and return the attribute, if it doesn't exist it returns a nullptr.
-      }],
-      /*retTy=*/"mlir::omp::TargetAttr",
-      /*methodName=*/"getTarget",
-      (ins), [{}], [{
-        if (Attribute flags = $_op->getAttr("omp.target"))
-          return ::llvm::dyn_cast_or_null<mlir::omp::TargetAttr>(flags);
-        return nullptr;
-      }]>,
-    InterfaceMethod<
-      /*description=*/[{
-        Set the attribute target on the current module with the
-        specified string arguments - name of cpu and corresponding features.
-      }],
-      /*retTy=*/"void",
-      /*methodName=*/"setTarget",
-      (ins "llvm::StringRef":$targetCPU,
-           "llvm::StringRef":$targetFeatures), [{}], [{
-        if (targetCPU.empty())
-          return;
-        $_op->setAttr(("omp." + mlir::omp::TargetAttr::getMnemonic()).str(),
-                  mlir::omp::TargetAttr::get($_op->getContext(),
-                                             targetCPU.str(),
-                                             targetFeatures.str()));
-      }]>,
     InterfaceMethod<
       /*description=*/[{
         Set a StringAttr on the current module containing the host IR file path. This

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

@@ -64,6 +64,7 @@ using namespace mlir;
 static constexpr StringRef varargsAttrName = "func.varargs";
 static constexpr StringRef linkageAttrName = "llvm.linkage";
 static constexpr StringRef barePtrAttrName = "llvm.bareptr";
+static constexpr StringRef targetFeaturesAttrName = "llvm.target_features";
 
 /// Return `true` if the `op` should use bare pointer calling convention.
 static bool shouldUseBarePtrCallConv(Operation *op,
@@ -79,6 +80,7 @@ static void filterFuncAttributes(FunctionOpInterface func,
   for (const NamedAttribute &attr : func->getDiscardableAttrs()) {
     if (attr.getName() == linkageAttrName ||
         attr.getName() == varargsAttrName ||
+        attr.getName() == targetFeaturesAttrName ||
         attr.getName() == LLVM::LLVMDialect::getReadnoneAttrName())
       continue;
     result.push_back(attr);
@@ -379,6 +381,19 @@ mlir::convertFuncOpToLLVMFuncOp(FunctionOpInterface funcOp,
     newFuncOp.setMemoryAttr(memoryAttr);
   }
 
+  // Create target_features attribute.
+  if (funcOp->hasAttr(targetFeaturesAttrName)) {
+    auto attr = funcOp->getAttrOfType<StringAttr>(targetFeaturesAttrName);
+    if (!attr) {
+      funcOp->emitError() << "Contains " << targetFeaturesAttrName
+                          << " attribute not of type StringAttr";
+      return rewriter.notifyMatchFailure(
+          funcOp, "Contains target features attribute not of type StringAttr");
+    }
+    newFuncOp.setTargetFeaturesAttr(
+        LLVM::TargetFeaturesAttr::get(rewriter.getContext(), attr.strref()));
+  }
+
   // Propagate argument/result attributes to all converted arguments/result
   // obtained after converting a given original argument/result.
   if (ArrayAttr resAttrDicts = funcOp.getAllResultAttrs()) {

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -2312,6 +2312,7 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   if (!targetOpSupported(opInst))
     return failure();
 
+  auto parentFn = opInst.getParentOfType<LLVM::LLVMFuncOp>();
   auto targetOp = cast<omp::TargetOp>(opInst);
   auto &targetRegion = targetOp.getRegion();
   DataLayout dl = DataLayout(opInst.getParentOfType<ModuleOp>());
@@ -2322,6 +2323,23 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   auto bodyCB = [&](InsertPointTy allocaIP,
                     InsertPointTy codeGenIP) -> InsertPointTy {
     builder.restoreIP(codeGenIP);
+
+    // Forward target-cpu and target-features function attributes from the
+    // original function to the new outlined function.
+    llvm::Function *llvmParentFn =
+        moduleTranslation.lookupFunction(parentFn.getName());
+    llvm::Function *llvmOutlinedFn = codeGenIP.getBlock()->getParent();
+    assert(llvmParentFn && llvmOutlinedFn &&
+           "Both parent and outlined functions must exist at this point");
+
+    if (auto attr = llvmParentFn->getFnAttribute("target-cpu");
+        attr.isStringAttribute())
+      llvmOutlinedFn->addFnAttr(attr);
+
+    if (auto attr = llvmParentFn->getFnAttribute("target-features");
+        attr.isStringAttribute())
+      llvmOutlinedFn->addFnAttr(attr);
+
     unsigned argIndex = 0;
     for (auto &mapOp : mapOperands) {
       auto mapInfoOp =
@@ -2339,11 +2357,11 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   };
 
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
-  StringRef parentName = opInst.getParentOfType<LLVM::LLVMFuncOp>().getName();
+  StringRef parentName = parentFn.getName();
 
   // Override parent name if early outlining function
-  if (auto earlyOutlineOp = llvm::dyn_cast<mlir::omp::EarlyOutliningInterface>(
-          opInst.getParentOfType<LLVM::LLVMFuncOp>().getOperation())) {
+  if (auto earlyOutlineOp =
+          llvm::dyn_cast<mlir::omp::EarlyOutliningInterface>(*parentFn)) {
     llvm::StringRef outlineParentName = earlyOutlineOp.getParentName();
     parentName = outlineParentName.empty() ? parentName : outlineParentName;
   }

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -1735,6 +1735,11 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
                                  .value()));
   }
 
+  if (llvm::Attribute attr = func->getFnAttribute("target-cpu");
+      attr.isStringAttribute()) {
+    funcOp.setTargetCpuAttr(StringAttr::get(context, attr.getValueAsString()));
+  }
+
   if (llvm::Attribute attr = func->getFnAttribute("target-features");
       attr.isStringAttribute()) {
     funcOp.setTargetFeaturesAttr(

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -971,6 +971,9 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   if (func.getArmNewZa())
     llvmFunc->addFnAttr("aarch64_pstate_za_new");
 
+  if (auto targetCpu = func.getTargetCpu())
+    llvmFunc->addFnAttr("target-cpu", *targetCpu);
+
   if (auto targetFeatures = func.getTargetFeatures())
     llvmFunc->addFnAttr("target-features", targetFeatures->getFeaturesString());
 

mlir/test/Conversion/FuncToLLVM/convert-funcs.mlir

@@ -61,6 +61,14 @@ func.func @variadic_func(%arg0: i32) attributes { "func.varargs" = true } {
   return
 }
 
+// CHECK-LABEL: llvm.func @target_features()
+// CHECK-SAME: target_features = #llvm.target_features<["+sme", "+sve"]>
+func.func private @target_features() attributes { "llvm.target_features" = "+sme,+sve" }
+
+// CHECK-LABEL: llvm.func @target_cpu()
+// CHECK-SAME: target_cpu = "gfx90a"
+func.func private @target_cpu() attributes { "target_cpu" = "gfx90a" }
+
 // -----
 
 // CHECK-LABEL: llvm.func @private_callee
@@ -86,3 +94,8 @@ func.func private @badllvmlinkage(i32) attributes { "llvm.linkage" = 3 : i64 } /
 func.func @variadic_func(%arg0: i32) attributes { "func.varargs" = true, "llvm.emit_c_interface" } {
   return
 }
+
+// -----
+
+// expected-error@+1{{Contains llvm.target_features attribute not of type StringAttr}}
+func.func private @badllvmtargetfeatures() attributes { "llvm.target_features" = 1 : i64 }

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir

@@ -772,7 +772,7 @@ func.func @insert_slice_chain(
 // CHECK-SAME: bufferization.access = "none"
     %arg2: tensor<62x90xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
 // CHECK-SAME: bufferization.access = "write"
-  -> tensor<62x90xf32> attributes {passthrough = [["target-cpu", "skylake-avx512"], ["prefer-vector-width", "512"]]}
+  -> tensor<62x90xf32> attributes {passthrough = [["prefer-vector-width", "512"]], target_cpu = "skylake-avx512"}
 {
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f32

mlir/test/Target/LLVMIR/Import/target-cpu.ll

@@ -0,0 +1,9 @@
+; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
+
+; CHECK-LABEL: llvm.func @target_cpu()
+; CHECK-SAME: target_cpu = "gfx90a"
+define void @target_cpu() #0 {
+  ret void
+}
+
+attributes #0 = { "target-cpu"="gfx90a" }