[WIP][MLIR][OMPIRBuilder] Lower OpenMP MLIR target attributes to LLVM IR

[domada]

Summary

OpenMP MLIR attributes are converted to LLVM IR at the end of OpenMP MLIR lowering.

OpenMP Task operation requires additional wrapper which is generated by the callback function as the part of finalization process. Finalization process is done after lowering of OpenMP MLIR target attributes. In consequence wrapper function inherits target attributes from its parent .

Added tests for checking the correctness of the lowering.

[domada]

Patch rebased

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

public:
   /// Gets (if variable with the given name already exist) or creates
   /// internal global variable with the specified Name. The created variable has
   /// linkage CommonLinkage by default and is initialized by null value.
   /// \param Ty Type of the global variable. If it is exist already the type
   /// must be the same.
   /// \param Name Name of the variable.
   GlobalVariable *getOrCreateInternalVariable(Type *Ty, const StringRef &Name,
                                               unsigned AddressSpace = 0);
+
+  /// Add function attribute to all functions which are defined inside module
+  /// but they don't have given attribute
+  /// \param AttributeName Name of the attribute
+  /// \param AttributeValue Value of the attribute
+  void addAttributeToModuleFunctions(StringRef AttributeName,
+                                     StringRef AttributeValue);
 };
 
 /// Class to represented the control flow structure of an OpenMP canonical loop.
 ///
 /// The control-flow structure is standardized for easy consumption by
 /// directives associated with loops. For instance, the worksharing-loop
 /// construct may change this control flow such that each loop iteration is
 /// executed on only one thread. The constraints of a canonical loop in brief

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

public:
 
     /// Verify that assumption cache isn't stale after a region is extracted.
     /// Returns true when verifier finds errors. AssumptionCache is passed as
     /// parameter to make this function stateless.
     static bool verifyAssumptionCache(const Function &OldFunc,
                                       const Function &NewFunc,
                                       AssumptionCache *AC);
 
+    /// Inherit all of the target dependent attributes and white-listed target
+    /// independent attributes. (e.g. If the extracted region contains a call to
+    /// an x86.sse instruction we need to make sure that the extracted region
+    /// has the "target-features" attribute allowing it to be lowered.
+
+    static void inheritTargetDependentAttributes(const Function *oldFunction,
+                                                 Function *newFunction);
     /// Test whether this code extractor is eligible.
     ///
     /// Based on the blocks used when constructing the code extractor,
     /// determine whether it is eligible for extraction.
     ///
     /// Checks that varargs handling (with vastart and vaend) is only done in
     /// the outlined blocks.
     bool isEligible() const;

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

OI.PostOutlineCB = [this, Ident, Tied, Final, IfCondition,
     // be present.
     SmallVector<Type *> WrapperArgTys{Builder.getInt32Ty()};
     if (HasTaskData)
       WrapperArgTys.push_back(OutlinedFn.getArg(0)->getType());
     FunctionCallee WrapperFuncVal = M.getOrInsertFunction(
         (Twine(OutlinedFn.getName()) + ".wrapper").str(),
         FunctionType::get(Builder.getInt32Ty(), WrapperArgTys, false));
     Function *WrapperFunc = dyn_cast<Function>(WrapperFuncVal.getCallee());
+    CodeExtractor::inheritTargetDependentAttributes(&OutlinedFn, WrapperFunc);
     PointerType *WrapperFuncBitcastType =
         FunctionType::get(Builder.getInt32Ty(),
                           {Builder.getInt32Ty(), Builder.getInt8PtrTy()}, false)
             ->getPointerTo();
     Value *WrapperFuncBitcast =
         ConstantExpr::getBitCast(WrapperFunc, WrapperFuncBitcastType);
 
     // Emit the @__kmpc_omp_task_alloc runtime call
if (const auto *CE =
                          Flags, CE->getLinkage());
 
     } else {
       llvm_unreachable("Unsupported entry kind.");
     }
   }
 }
 
+void OpenMPIRBuilder::addAttributeToModuleFunctions(StringRef AttributeName,
+                                                    StringRef AttributeValue) {
+  if (AttributeName.empty() || AttributeValue.empty())
+    return;
+  for (Function &f : M.functions()) {
+    if (f.isDeclaration())
+      continue;
+    if (f.hasFnAttribute(AttributeName))
+      continue;
+    f.addFnAttr(AttributeName, AttributeValue);
+  }
+}
+
 void TargetRegionEntryInfo::getTargetRegionEntryFnName(
     SmallVectorImpl<char> &Name, StringRef ParentName, unsigned DeviceID,
     unsigned FileID, unsigned Line, unsigned Count) {
   raw_svector_ostream OS(Name);
   OS << "__omp_offloading" << llvm::format("_%x", DeviceID)
      << llvm::format("_%x_", FileID) << ParentName << "_l" << Line;
   if (Count)
     OS << "_" << Count;

llvm/lib/Transforms/Utils/CodeExtractor.cpp

if (ReturnInst *RI = dyn_cast<ReturnInst>(Block->getTerminator())) {
         DomTreeNode *NewNode = DT->addNewBlock(New, Block);
 
         for (DomTreeNode *I : Children)
           DT->changeImmediateDominator(I, NewNode);
       }
     }
 }
 
-/// constructFunction - make a function based on inputs and outputs, as follows:
-/// f(in0, ..., inN, out0, ..., outN)
-Function *CodeExtractor::constructFunction(const ValueSet &inputs,
-                                           const ValueSet &outputs,
-                                           BasicBlock *header,
-                                           BasicBlock *newRootNode,
-                                           BasicBlock *newHeader,
-                                           Function *oldFunction,
-                                           Module *M) {
-  LLVM_DEBUG(dbgs() << "inputs: " << inputs.size() << "\n");
-  LLVM_DEBUG(dbgs() << "outputs: " << outputs.size() << "\n");
-
-  // This function returns unsigned, outputs will go back by reference.
-  switch (NumExitBlocks) {
-  case 0:
-  case 1: RetTy = Type::getVoidTy(header->getContext()); break;
-  case 2: RetTy = Type::getInt1Ty(header->getContext()); break;
-  default: RetTy = Type::getInt16Ty(header->getContext()); break;
-  }
-
-  std::vector<Type *> ParamTy;
-  std::vector<Type *> AggParamTy;
-  ValueSet StructValues;
-  const DataLayout &DL = M->getDataLayout();
-
-  // Add the types of the input values to the function's argument list
-  for (Value *value : inputs) {
-    LLVM_DEBUG(dbgs() << "value used in func: " << *value << "\n");
-    if (AggregateArgs && !ExcludeArgsFromAggregate.contains(value)) {
-      AggParamTy.push_back(value->getType());
-      StructValues.insert(value);
-    } else
-      ParamTy.push_back(value->getType());
-  }
-
-  // Add the types of the output values to the function's argument list.
-  for (Value *output : outputs) {
-    LLVM_DEBUG(dbgs() << "instr used in func: " << *output << "\n");
-    if (AggregateArgs && !ExcludeArgsFromAggregate.contains(output)) {
-      AggParamTy.push_back(output->getType());
-      StructValues.insert(output);
-    } else
-      ParamTy.push_back(
-          PointerType::get(output->getType(), DL.getAllocaAddrSpace()));
-  }
-
-  assert(
-      (ParamTy.size() + AggParamTy.size()) ==
-          (inputs.size() + outputs.size()) &&
-      "Number of scalar and aggregate params does not match inputs, outputs");
-  assert((StructValues.empty() || AggregateArgs) &&
-         "Expeced StructValues only with AggregateArgs set");
-
-  // Concatenate scalar and aggregate params in ParamTy.
-  size_t NumScalarParams = ParamTy.size();
-  StructType *StructTy = nullptr;
-  if (AggregateArgs && !AggParamTy.empty()) {
-    StructTy = StructType::get(M->getContext(), AggParamTy);
-    ParamTy.push_back(PointerType::get(StructTy, DL.getAllocaAddrSpace()));
-  }
-
-  LLVM_DEBUG({
-    dbgs() << "Function type: " << *RetTy << " f(";
-    for (Type *i : ParamTy)
-      dbgs() << *i << ", ";
-    dbgs() << ")\n";
-  });
-
-  FunctionType *funcType = FunctionType::get(
-      RetTy, ParamTy, AllowVarArgs && oldFunction->isVarArg());
-
-  std::string SuffixToUse =
-      Suffix.empty()
-          ? (header->getName().empty() ? "extracted" : header->getName().str())
-          : Suffix;
-  // Create the new function
-  Function *newFunction = Function::Create(
-      funcType, GlobalValue::InternalLinkage, oldFunction->getAddressSpace(),
-      oldFunction->getName() + "." + SuffixToUse, M);
-
-  // Inherit all of the target dependent attributes and white-listed
-  // target independent attributes.
-  //  (e.g. If the extracted region contains a call to an x86.sse
-  //  instruction we need to make sure that the extracted region has the
-  //  "target-features" attribute allowing it to be lowered.
+void CodeExtractor::inheritTargetDependentAttributes(
+    const Function *oldFunction, Function *newFunction) {
   // FIXME: This should be changed to check to see if a specific
   //           attribute can not be inherited.
   for (const auto &Attr : oldFunction->getAttributes().getFnAttrs()) {
     if (Attr.isStringAttribute()) {
       if (Attr.getKindAsString() == "thunk")
         continue;
     } else
       switch (Attr.getKindAsEnum()) {
if (Attr.isStringAttribute()) {
       case Attribute::WriteOnly:
       //  These are not really attributes.
       case Attribute::None:
       case Attribute::EndAttrKinds:
       case Attribute::EmptyKey:
       case Attribute::TombstoneKey:
         llvm_unreachable("Not a function attribute");
       }
-
     newFunction->addFnAttr(Attr);
   }
+}
+/// constructFunction - make a function based on inputs and outputs, as follows:
+/// f(in0, ..., inN, out0, ..., outN)
+Function *CodeExtractor::constructFunction(const ValueSet &inputs,
+                                           const ValueSet &outputs,
+                                           BasicBlock *header,
+                                           BasicBlock *newRootNode,
+                                           BasicBlock *newHeader,
+                                           Function *oldFunction, Module *M) {
+  LLVM_DEBUG(dbgs() << "inputs: " << inputs.size() << "\n");
+  LLVM_DEBUG(dbgs() << "outputs: " << outputs.size() << "\n");
+  // This function returns unsigned, outputs will go back by reference.
+  switch (NumExitBlocks) {
+  case 0:
+  case 1:
+    RetTy = Type::getVoidTy(header->getContext());
+    break;
+  case 2:
+    RetTy = Type::getInt1Ty(header->getContext());
+    break;
+  default:
+    RetTy = Type::getInt16Ty(header->getContext());
+    break;
+  }
+
+  std::vector<Type *> ParamTy;
+  std::vector<Type *> AggParamTy;
+  ValueSet StructValues;
+  const DataLayout &DL = M->getDataLayout();
+
+  // Add the types of the input values to the function's argument list
+  for (Value *value : inputs) {
+    LLVM_DEBUG(dbgs() << "value used in func: " << *value << "\n");
+    if (AggregateArgs && !ExcludeArgsFromAggregate.contains(value)) {
+      AggParamTy.push_back(value->getType());
+      StructValues.insert(value);
+    } else
+      ParamTy.push_back(value->getType());
+  }
+
+  // Add the types of the output values to the function's argument list.
+  for (Value *output : outputs) {
+    LLVM_DEBUG(dbgs() << "instr used in func: " << *output << "\n");
+    if (AggregateArgs && !ExcludeArgsFromAggregate.contains(output)) {
+      AggParamTy.push_back(output->getType());
+      StructValues.insert(output);
+    } else
+      ParamTy.push_back(
+          PointerType::get(output->getType(), DL.getAllocaAddrSpace()));
+  }
+
+  assert(
+      (ParamTy.size() + AggParamTy.size()) ==
+          (inputs.size() + outputs.size()) &&
+      "Number of scalar and aggregate params does not match inputs, outputs");
+  assert((StructValues.empty() || AggregateArgs) &&
+         "Expeced StructValues only with AggregateArgs set");
+
+  // Concatenate scalar and aggregate params in ParamTy.
+  size_t NumScalarParams = ParamTy.size();
+  StructType *StructTy = nullptr;
+  if (AggregateArgs && !AggParamTy.empty()) {
+    StructTy = StructType::get(M->getContext(), AggParamTy);
+    ParamTy.push_back(PointerType::get(StructTy, DL.getAllocaAddrSpace()));
+  }
+
+  LLVM_DEBUG({
+    dbgs() << "Function type: " << *RetTy << " f(";
+    for (Type *i : ParamTy)
+      dbgs() << *i << ", ";
+    dbgs() << ")\n";
+  });
+
+  FunctionType *funcType = FunctionType::get(
+      RetTy, ParamTy, AllowVarArgs && oldFunction->isVarArg());
+
+  std::string SuffixToUse =
+      Suffix.empty()
+          ? (header->getName().empty() ? "extracted" : header->getName().str())
+          : Suffix;
+  // Create the new function
+  Function *newFunction = Function::Create(
+      funcType, GlobalValue::InternalLinkage, oldFunction->getAddressSpace(),
+      oldFunction->getName() + "." + SuffixToUse, M);
+
+  inheritTargetDependentAttributes(oldFunction, newFunction);
   newFunction->insert(newFunction->end(), newRootNode);
 
   // Create scalar and aggregate iterators to name all of the arguments we
   // inserted.
   Function::arg_iterator ScalarAI = newFunction->arg_begin();
   Function::arg_iterator AggAI = std::next(ScalarAI, NumScalarParams);
 
   // Rewrite all users of the inputs in the extracted region to use the

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

ompBuilder->createGlobalFlag(
       attribute
           .getAssumeNoNestedParallelism() /*LangOpts().OpenMPNoNestedParallelism*/
       ,
       "__omp_rtl_assume_no_nested_parallelism");
 
   return success();
 }
 
+/// Convert OpenMP MLIR target attributes to LLVM function attributes
+static LogicalResult
+setLLVMFunctionTargetAttr(Operation *Op, omp::TargetAttr targetAttr,
+                          LLVM::ModuleTranslation &moduleTranslation) {
+
+  llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+  ompBuilder->addAttributeToModuleFunctions("target-cpu",
+                                            targetAttr.getTargetCpu());
+  ompBuilder->addAttributeToModuleFunctions("target-features",
+                                            targetAttr.getTargetFeatures());
+  return success();
+}
+
 namespace {
 
 /// Implementation of the dialect interface that converts operations belonging
 /// to the OpenMP dialect to LLVM IR.
 class OpenMPDialectLLVMIRTranslationInterface
     : public LLVMTranslationDialectInterface {
 public:
   using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;

 LogicalResult OpenMPDialectLLVMIRTranslationInterface::amendOperation(
     Operation *op, NamedAttribute attribute,
     LLVM::ModuleTranslation &moduleTranslation) const {
 
   return llvm::TypeSwitch<Attribute, LogicalResult>(attribute.getValue())
       .Case([&](mlir::omp::FlagsAttr rtlAttr) {
         return convertFlagsAttr(op, rtlAttr, moduleTranslation);
       })
+      .Case([&](omp::TargetAttr targetAttr) {
+        return setLLVMFunctionTargetAttr(op, targetAttr, moduleTranslation);
+      })
       .Default([&](Attribute attr) {
         // fall through for omp attributes that do not require lowering and/or
         // have no concrete definition and thus no type to define a case on
         return success();
       });
 
   return failure();
 }

mlir/test/Target/LLVMIR/openmp-llvm.mlir

 
 // CHECK: @__omp_rtl_debug_kind = weak_odr hidden constant i32 0
 // CHECK: @__omp_rtl_assume_teams_oversubscription = weak_odr hidden constant i32 1
 // CHECK: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 0
 // CHECK: @__omp_rtl_assume_no_thread_state = weak_odr hidden constant i32 1
 // CHECK: @__omp_rtl_assume_no_nested_parallelism = weak_odr hidden constant i32 0
 module attributes {omp.flags = #omp.flags<assume_teams_oversubscription = true, assume_no_thread_state = true>, 
                    omp.is_device = #omp.isdevice<is_device = true>} {}
+
+// -----
+
+// CHECK: @omp_target_features_test() #0 {
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.target = #omp.target<target_cpu = "gfx908", target_features = "+dot3-insts,+dot4-insts,+s-memtime-inst,+16-bit-insts,+s-memrealtime,+dot6-insts,+dl-insts,+wavefrontsize64,+gfx9-insts,+gfx8-insts,+ci-insts,+dot10-insts,+dot7-insts,+dot1-insts,+dot5-insts,+mai-insts,+dpp,+dot2-insts">} {
+  llvm.func @omp_target_features_test() {
+    llvm.return
+  }
+}
+
+// CHECK: attributes #0 = { "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,
+// CHECK-SAME: +wavefrontsize64,+gfx9-insts,+gfx8-insts,+ci-insts,+dot10-insts,
+// CHECK-SAME: +dot7-insts,+dot1-insts,+dot5-insts,+mai-insts,+dpp,+dot2-insts"
+
+// -----
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa",  omp.target = #omp.target<target_cpu = "gfx908", target_features = "+dot3-insts,+dot4-insts,+s-memtime-inst,+16-bit-insts,+s-memrealtime,+dot6-insts,+dl-insts,+wavefrontsize64,+gfx9-insts,+gfx8-insts,+ci-insts,+dot10-insts,+dot7-insts,+dot1-insts,+dot5-insts,+mai-insts,+dpp,+dot2-insts">} {
+// CHECK: @test_omp_attr() #0 {
+// CHECK: @test_omp_attr..omp_par() #0 {
+// CHECK: @test_omp_attr..omp_par.wrapper(i32 %{{.*}}) #0 {
+
+  llvm.func @test_omp_attr() {
+    omp.task {
+      omp.terminator
+    }
+  llvm.return
+  }
+}
+// CHECK: attributes #0 = { "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,
+// CHECK-SAME: +wavefrontsize64,+gfx9-insts,+gfx8-insts,+ci-insts,+dot10-insts,
+// CHECK-SAME: +dot7-insts,+dot1-insts,+dot5-insts,+mai-insts,+dpp,+dot2-insts"