[Flang][OpenMP][MLIR] Lower OpenMP target attributes

[domada]

Summary

Added functionality which allows to attach information about target features to OpenMP MLIR code.

Scope of changes:

1. Add attribute to OpenMP MLIR dialect which stores target cpu and target features
2. Store target information in MLIR module

[domada]

Added OpenMP target attributes from parent review D145579 .

This review contains complete set of changes required to add target attributes for OpenMP dialect.

[domada]

Scope of changes:

1. Applied @kiranchandramohan suggestion ( https://reviews.llvm.org/D146721 ) for target interface.
2. Patch rebased.
3. Limited scope of patch. Lowering from MLIR to LLVM IR will be present in different patch

[domada]

Updated mlir definition of target attributes. Use MLIR class StringRefParameter instead std::string

[kiranchandramohan]

LGTM.

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

116

128

flang/include/flang/Tools/CrossToolHelpers.h

void setOffloadModuleInterfaceAttributes(
     mlir::ModuleOp &module, bool isDevice) {
   // Should be registered by the OpenMPDialect
   if (auto offloadMod = llvm::dyn_cast<mlir::omp::OffloadModuleInterface>(
           module.getOperation())) {
     offloadMod.setIsDevice(isDevice);
   }
 }
 
+//  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);
+  }
+}
+
 #endif // FORTRAN_TOOLS_CROSS_TOOL_HELPERS_H

flang/lib/Frontend/FrontendActions.cpp

lower::LoweringBridge lb = Fortran::lower::LoweringBridge::create(
       ci.getInvocation().getSemanticsContext().targetCharacteristics(),
       ci.getParsing().allCooked(), ci.getInvocation().getTargetOpts().triple,
       kindMap, ci.getInvocation().getLoweringOpts(),
       ci.getInvocation().getFrontendOpts().envDefaults);
 
   // Fetch module from lb, so we can set
   mlirModule = std::make_unique<mlir::ModuleOp>(lb.getModule());
 
+  if (!setUpTargetMachine())
+    return false;
+
   if (ci.getInvocation().getFrontendOpts().features.IsEnabled(
           Fortran::common::LanguageFeature::OpenMP)) {
     setOffloadModuleInterfaceAttributes(
         *mlirModule, ci.getInvocation().getLangOpts().OpenMPIsDevice);
+    setOffloadModuleInterfaceTargetAttribute(*mlirModule, tm->getTargetCPU(),
+                                             tm->getTargetFeatureString());
   }
 
-  if (!setUpTargetMachine())
-    return false;
   const llvm::DataLayout &dl = tm->createDataLayout();
   setMLIRDataLayout(*mlirModule, dl);
 
   // Create a parse tree and lower it to FIR
   Fortran::parser::Program &parseTree{*ci.getParsing().parseTree()};
   lb.lower(parseTree, ci.getInvocation().getSemanticsContext());
 
   // run the default passes.

flang/test/Lower/OpenMP/target_cpu_features.f90

+!REQUIRES: amdgpu-registered-target
+!RUN: %flang_fc1 -emit-fir -triple amdgcn-amd-amdhsa -target-cpu gfx908 -fopenmp %s -o - | FileCheck %s
+
+!===============================================================================
+! Target_Enter Simple
+!===============================================================================
+
+!CHECK: omp.target = #omp.target<target_cpu = "gfx908",
+!CHECK-SAME: target_features = "+dot3-insts,+dot4-insts,+s-memtime-inst,
+!CHECK-SAME: +16-bit-insts,+s-memrealtime,+dot6-insts,+dl-insts,+wavefrontsize64,
+!CHECK-SAME: +gfx9-insts,+gfx8-insts,+ci-insts,+dot10-insts,+dot7-insts,
+!CHECK-SAME: +dot1-insts,+dot5-insts,+mai-insts,+dpp,+dot2-insts">
+!CHECK-LABEL: func.func @_QPomp_target_simple() {
+subroutine omp_target_simple
+end subroutine omp_target_simple
+

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

let parameters = (ins
     DefaultValuedParameter<"bool", "false">:$assume_threads_oversubscription,
     DefaultValuedParameter<"bool", "false">:$assume_no_thread_state,
     DefaultValuedParameter<"bool", "false">:$assume_no_nested_parallelism
   );
   
   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>;
 
 // Type which can be constraint accepting standard integers and indices.
 def IntLikeType : AnyTypeOf<[AnyInteger, Index]>;
 
 def OpenMP_PointerLikeType : TypeAlias<OpenMP_PointerLikeTypeInterface,
   "OpenMP-compatible variable type">;

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

def OffloadModuleInterface : OpInterface<"OffloadModuleInterface"> {
   }];
 
   let cppNamespace = "::mlir::omp";
 
   let methods = [
     InterfaceMethod<
       /*description=*/[{
         Set the attribute IsDeviceAttr on the current module with the 
-        specified boolean argument. 
+        specified boolean argument.
       }],
       /*retTy=*/"void",
-      /*methodName=*/"setIsDevice", 
+      /*methodName=*/"setIsDevice",
       (ins "bool":$isDevice), [{}], [{
         $_op->setAttr(
           mlir::StringAttr::get($_op->getContext(), llvm::Twine{"omp.is_device"}),
             mlir::omp::IsDeviceAttr::get($_op->getContext(), isDevice));
       }]>,
       InterfaceMethod<
       /*description=*/[{
         Get the IsDeviceAttr attribute on the current module if it exists and return
         its value, if it doesn't exit it returns false by default.
       }],
       /*retTy=*/"bool",
-      /*methodName=*/"getIsDevice", 
+      /*methodName=*/"getIsDevice",
       (ins), [{}], [{
         if (Attribute isDevice = $_op->getAttr("omp.is_device"))
           if (isDevice.isa<mlir::omp::IsDeviceAttr>())
             return isDevice.dyn_cast<IsDeviceAttr>().getIsDevice();
         return false;
       }]>,
       InterfaceMethod<
       /*description=*/[{
InterfaceMethod<
             "bool":$assumeThreadsOversubscription,
             "bool":$assumeNoThreadState,
             "bool":$assumeNoNestedParallelism), [{}], [{
         $_op->setAttr(("omp." + mlir::omp::FlagsAttr::getMnemonic()).str(),
                   mlir::omp::FlagsAttr::get($_op->getContext(), debugKind,
                       assumeTeamsOversubscription, assumeThreadsOversubscription, 
                       assumeNoThreadState, assumeNoNestedParallelism));
       }]>,
+    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 flags.dyn_cast_or_null<mlir::omp::TargetAttr>();
+        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()));
+      }]>,
   ];
 }
 
 #endif // OpenMP_OPS_INTERFACES