[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 ↗

(On Diff #509956)

128 ↗

(On Diff #509956)

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);
 };
 
 /// Data structure to contain the information needed to uniquely identify
 /// a target entry.
 struct TargetRegionEntryInfo {
   std::string ParentName;
   unsigned DeviceID;
   unsigned FileID;

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

 public:
   using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;
 
   /// Translates the given operation to LLVM IR using the provided IR builder
   /// and saving the state in `moduleTranslation`.
   LogicalResult
   convertOperation(Operation *op, llvm::IRBuilderBase &builder,
                    LLVM::ModuleTranslation &moduleTranslation) const final;
+  LogicalResult
+  amendOperation(Operation *op, NamedAttribute attribute,
+                 LLVM::ModuleTranslation &moduleTranslation) const final;
 };
 
 } // namespace
 
+/// Convert OpenMP MLIR target attributes to LLVM function attributes
+static void
+trySetLLVMFunctionTargetAttr(Operation *Op, NamedAttribute namedAttr,
+                             StringAttr valueAttr,
+                             LLVM::ModuleTranslation &moduleTranslation) {
+  std::string LLVMFuncAttrName;
+  if (namedAttr.getName() == "omp.target_cpu")
+    LLVMFuncAttrName = "target-cpu";
+  else if (namedAttr.getName() == "omp.target_cpu_features")
+    LLVMFuncAttrName = "target-features";
+  else
+    return;
+
+  llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+  ompBuilder->addAttributeToModuleFunctions(LLVMFuncAttrName,
+                                            valueAttr.getValue().str());
+}
+
+/// Given an OpenMP MLIR attribute, create the corresponding LLVM-IR, runtime
+/// calls, or operation amendments
+LogicalResult OpenMPDialectLLVMIRTranslationInterface::amendOperation(
+    Operation *op, NamedAttribute attribute,
+    LLVM::ModuleTranslation &moduleTranslation) const {
+
+  return llvm::TypeSwitch<Attribute, LogicalResult>(attribute.getValue())
+      .Case([&](mlir::StringAttr attr) {
+        // check if given string attributes relate to omp.target_cpu or
+        // omp.target_cpu_features. If yes, try to add LLVM function
+        // attributes
+        trySetLLVMFunctionTargetAttr(op, attribute, attr, moduleTranslation);
+        return success();
+      })
+      .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
+        // e.g. omp.is_device
+        return success();
+      });
+
+  return failure();
+}
+
 /// Given an OpenMP MLIR operation, create the corresponding LLVM IR
 /// (including OpenMP runtime calls).
 LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
     Operation *op, llvm::IRBuilderBase &builder,
     LLVM::ModuleTranslation &moduleTranslation) const {
 
   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
 

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

omp.atomic.capture memory_order(acq_rel) {
     omp.atomic.update %arg1 : !llvm.ptr {
     ^bb0(%xval: i32):
       %newval = llvm.add %xval, %expr : i32
       omp.yield(%newval : i32)
     }
   }
   llvm.return
 }
+
+// -----
+
+// CHECK: @omp_target_features_test() #0 {
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.target_cpu = "gfx908", omp.target_cpu_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_cpu = "gfx908", omp.target_cpu_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"