[OMPIRBuilder][MLIR] Fix target-features and target-cpu handling

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -103,14 +103,10 @@ class OpenMPIRBuilderConfig {
   /// Flag for specifying if offloading is mandatory.
   std::optional<bool> OpenMPOffloadMandatory;
 
-  /// Name of the target processor.
-  StringRef TargetCPU;
-  /// String representation of the target processor's features.
-  StringRef TargetFeatures;
-
   /// First separator used between the initial two parts of a name.
   std::optional<StringRef> FirstSeparator;
-  /// Separator used between all of the rest consecutive parts of s name
+
+  /// Separator used between all of the rest consecutive parts of a name
   std::optional<StringRef> Separator;
 
   // Grid Value for the GPU target

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -156,22 +156,14 @@ static Value *getTypeSizeInBytesValue(IRBuilder<> &Builder, Module &M,
   return Builder.getInt64(getTypeSizeInBytes(M, Type));
 }
 
-static const omp::GV &getGridValue(const Triple &T, StringRef Features) {
-  if (T.isAMDGPU()) {
-    if (Features.count("+wavefrontsize64"))
-      return omp::getAMDGPUGridValues<64>();
-    return omp::getAMDGPUGridValues<32>();
-  }
-  if (T.isNVPTX())
-    return omp::NVPTXGridValues;
-  llvm_unreachable("No grid value available for this architecture!");
-}
-
 static const omp::GV &getGridValue(const Triple &T, Function *Kernel) {
   if (T.isAMDGPU()) {
     StringRef Features =
         Kernel->getFnAttribute("target-features").getValueAsString();
-    return getGridValue(T, Features);
+
+    if (Features.count("+wavefrontsize64"))
+      return omp::getAMDGPUGridValues<64>();
+    return omp::getAMDGPUGridValues<32>();
   }
   if (T.isNVPTX())
     return omp::NVPTXGridValues;
@@ -5638,7 +5630,7 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTargetInit(
   Function *Kernel = Builder.GetInsertBlock()->getParent();
 
   // Set the grid value in the config needed for lowering later on
-  Config.setGridValue(getGridValue(T, Config.TargetFeatures));
+  Config.setGridValue(getGridValue(T, Kernel));
 
   // Manifest the launch configuration in the metadata matching the kernel
   // environment.
@@ -5896,11 +5888,6 @@ void OpenMPIRBuilder::setOutlinedTargetRegionFunctionAttributes(
     if (T.isAMDGCN())
       OutlinedFn->setCallingConv(CallingConv::AMDGPU_KERNEL);
   }
-
-  if (!Config.TargetCPU.empty())
-    OutlinedFn->addFnAttr("target-cpu", Config.TargetCPU);
-  if (!Config.TargetFeatures.empty())
-    OutlinedFn->addFnAttr("target-features", Config.TargetFeatures);
 }
 
 Constant *OpenMPIRBuilder::createOutlinedFunctionID(Function *OutlinedFn,
@@ -6249,6 +6236,18 @@ static Function *createOutlinedFunction(
   auto Func = Function::Create(FuncType, GlobalValue::InternalLinkage, FuncName,
                                Builder.GetInsertBlock()->getModule());
 
+  // Forward target-cpu and target-features function attributes from the
+  // original function to the new outlined function.
+  Function *ParentFn = Builder.GetInsertBlock()->getParent();
+
+  auto TargetCpuAttr = ParentFn->getFnAttribute("target-cpu");
+  if (TargetCpuAttr.isStringAttribute())
+    Func->addFnAttr(TargetCpuAttr);
+
+  auto TargetFeaturesAttr = ParentFn->getFnAttribute("target-features");
+  if (TargetFeaturesAttr.isStringAttribute())
+    Func->addFnAttr(TargetFeaturesAttr);
+
   if (OMPBuilder.Config.isTargetDevice()) {
     std::vector<llvm::WeakTrackingVH> LLVMCompilerUsed;
     emitExecutionMode(OMPBuilder, Builder, FuncName, false, LLVMCompilerUsed);

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

@@ -2840,22 +2840,6 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
   auto bodyCB = [&](InsertPointTy allocaIP,
                     InsertPointTy codeGenIP) -> InsertPointTy {
-    // 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);
-
     builder.restoreIP(codeGenIP);
     unsigned argIndex = 0;
     for (auto &mapOp : mapOperands) {

mlir/test/Target/LLVMIR/omptarget-parallel-wsloop.mlir

@@ -39,9 +39,9 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
 
 // CHECK:      define internal void @[[LOOP_BODY_FUNC]](i32 %[[CNT:.*]], ptr %[[LOOP_BODY_ARG_PTR:.*]]) #[[ATTRS2:[0-9]+]] {
 
-// CHECK:      attributes #[[ATTRS2:.*]] = {
+// CHECK:      attributes #[[ATTRS1]] = {
 // CHECK-SAME:  "target-cpu"="gfx90a"
 // CHECK-SAME:  "target-features"="+gfx9-insts,+wavefrontsize64"
-// CHECK:      attributes #[[ATTRS1:.*]] = {
+// CHECK:      attributes #[[ATTRS2]] = {
 // CHECK-SAME:  "target-cpu"="gfx90a"
 // CHECK-SAME:  "target-features"="+gfx9-insts,+wavefrontsize64"