[OpenMP] omp begin/end declare variant - part 2, sema ("+CG")

This is the second part loosely extracted from D71179 and cleaned up.

This patch provides semantic analysis support for `omp begin/end declare
variant`, mostly as defined in OpenMP technical report 8 (TR8) [0].
The sema handling makes code generation obsolete as we generate "the
right" calls that can just be handled as usual. This handling also
applies to the existing, albeit problematic, `omp declare variant
support`. As a consequence a lot of unneeded code generation and
complexity is removed.

A major purpose of this patch is to provide proper `math.h`/`cmath`
support for OpenMP target offloading. See PR42061, PR42798, PR42799. The
current code was developed with this feature in mind, see [1].

The logic is as follows:

If we have seen a `#pragma omp begin declare variant match(<SELECTOR>)`
but not the corresponding `end declare variant`, and we find a function
definition we will:
  1) Create a function declaration for the definition we were about to generate.
  2) Create a function definition but with a mangled name (according to
     `<SELECTOR>`).
  3) Annotate the declaration with the `OMPDeclareVariantAttr`, the same
     one used already for `omp declare variant`, using and the mangled
     function definition as specialization for the context defined by
     `<SELECTOR>`.

When a call is created we inspect it. If the target has an
`OMPDeclareVariantAttr` attribute we try to specialize the call. To this
end, all variants are checked, the best applicable one is picked and a
new call to the specialization is created. The new call is used instead
of the original one to the base function. To keep the AST printing and
tooling possible we utilize the PseudoObjectExpr. The original call is
the syntactic expression, the specialized call is the semantic
expression.

[0] https://www.openmp.org/wp-content/uploads/openmp-TR8.pdf
[1] https://reviews.llvm.org/D61399#change-496lQkg0mhRN

Reviewers: kiranchandramohan, ABataev, RaviNarayanaswamy, gtbercea, grokos, sdmitriev, JonChesterfield, hfinkel, fghanim, aaron.ballman

Subscribers: bollu, guansong, openmp-commits, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D75779

Author: jdoerfert

clang/include/clang/AST/Decl.h

@@ -4542,6 +4542,13 @@ inline bool IsEnumDeclScoped(EnumDecl *ED) {
   return ED->isScoped();
 }
 
+/// OpenMP variants are mangled early based on their OpenMP context selector.
+/// The new name looks likes this:
+///  <name> + OpenMPVariantManglingSeparatorStr + <mangled OpenMP context>
+static constexpr StringRef getOpenMPVariantManglingSeparatorStr() {
+  return ".ompvariant";
+}
+
 } // namespace clang
 
 #endif // LLVM_CLANG_AST_DECL_H

clang/include/clang/AST/OpenMPClause.h

@@ -7146,6 +7146,9 @@ class OMPTraitInfo {
   friend class ASTContext;
 
 public:
+  /// Reconstruct a (partial) OMPTraitInfo object from a mangled name.
+  OMPTraitInfo(StringRef MangledName);
+
   struct OMPTraitProperty {
     llvm::omp::TraitProperty Kind = llvm::omp::TraitProperty::invalid;
   };
@@ -7184,6 +7187,9 @@ class OMPTraitInfo {
                              llvm::omp::VariantMatchInfo &VMI,
                              bool DeviceSetOnly) const;
 
+  /// Return a string representation identifying this context selector.
+  std::string getMangledName() const;
+
   /// Print a human readable representation into \p OS.
   void print(llvm::raw_ostream &OS, const PrintingPolicy &Policy) const;
 };

clang/include/clang/Basic/DiagnosticParseKinds.td

@@ -1234,6 +1234,9 @@ def err_omp_declare_simd_inbranch_notinbranch : Error<
   "unexpected '%0' clause, '%1' is specified already">;
 def err_expected_end_declare_target_or_variant : Error<
   "expected '#pragma omp end declare %select{target|variant}0'">;
+def err_expected_begin_declare_variant
+    : Error<"'#pragma omp end declare variant' with no matching '#pragma omp "
+            "begin declare variant'">;
 def err_omp_declare_target_unexpected_clause: Error<
   "unexpected '%0' clause, only %select{'to' or 'link'|'to', 'link' or 'device_type'}1 clauses expected">;
 def err_omp_expected_clause: Error<

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -10186,6 +10186,10 @@ def err_omp_non_lvalue_in_map_or_motion_clauses: Error<
   >;
 def err_omp_event_var_expected : Error<
   "expected variable of the 'omp_event_handle_t' type%select{|, not %1}0">;
+def warn_nested_declare_variant
+    : Warning<"nesting `omp begin/end declare variant` is not supported yet; "
+              "nested context ignored">,
+      InGroup<SourceUsesOpenMP>;
 } // end of OpenMP category
 
 let CategoryName = "Related Result Type Issue" in {

clang/include/clang/Basic/IdentifierTable.h

@@ -108,7 +108,10 @@ class alignas(IdentifierInfoAlignment) IdentifierInfo {
   // True if this is the 'import' contextual keyword.
   unsigned IsModulesImport : 1;
 
-  // 29 bits left in a 64-bit word.
+  // True if this is a mangled OpenMP variant name.
+  unsigned IsMangledOpenMPVariantName : 1;
+
+  // 28 bits left in a 64-bit word.
 
   // Managed by the language front-end.
   void *FETokenInfo = nullptr;
@@ -121,7 +124,7 @@ class alignas(IdentifierInfoAlignment) IdentifierInfo {
         IsPoisoned(false), IsCPPOperatorKeyword(false),
         NeedsHandleIdentifier(false), IsFromAST(false), ChangedAfterLoad(false),
         FEChangedAfterLoad(false), RevertedTokenID(false), OutOfDate(false),
-        IsModulesImport(false) {}
+        IsModulesImport(false), IsMangledOpenMPVariantName(false) {}
 
 public:
   IdentifierInfo(const IdentifierInfo &) = delete;
@@ -371,6 +374,12 @@ class alignas(IdentifierInfoAlignment) IdentifierInfo {
       RecomputeNeedsHandleIdentifier();
   }
 
+  /// Determine whether this is the mangled name of an OpenMP variant.
+  bool isMangledOpenMPVariantName() const { return IsMangledOpenMPVariantName; }
+
+  /// Set whether this is the mangled name of an OpenMP variant.
+  void setMangledOpenMPVariantName(bool I) { IsMangledOpenMPVariantName = I; }
+
   /// Return true if this identifier is an editor placeholder.
   ///
   /// Editor placeholders are produced by the code-completion engine and are

clang/include/clang/Sema/Sema.h

@@ -9812,14 +9812,55 @@ class Sema final {
                                      MapT &Map, unsigned Selector = 0,
                                      SourceRange SrcRange = SourceRange());
 
-  /// Marks all the functions that might be required for the currently active
-  /// OpenMP context.
-  void markOpenMPDeclareVariantFuncsReferenced(SourceLocation Loc,
-                                               FunctionDecl *Func,
-                                               bool MightBeOdrUse);
+  /// Helper to keep information about the current `omp begin/end declare
+  /// variant` nesting.
+  struct OMPDeclareVariantScope {
+    /// The associated OpenMP context selector.
+    OMPTraitInfo *TI;
+
+    /// The associated OpenMP context selector mangling.
+    std::string NameSuffix;
+
+    OMPDeclareVariantScope(OMPTraitInfo &TI)
+        : TI(&TI), NameSuffix(TI.getMangledName()) {}
+  };
+
+  /// The current `omp begin/end declare variant` scopes.
+  SmallVector<OMPDeclareVariantScope, 4> OMPDeclareVariantScopes;
+
+  /// The declarator \p D defines a function in the scope \p S which is nested
+  /// in an `omp begin/end declare variant` scope. In this method we create a
+  /// declaration for \p D and rename \p D according to the OpenMP context
+  /// selector of the surrounding scope.
+  FunctionDecl *
+  ActOnStartOfFunctionDefinitionInOpenMPDeclareVariantScope(Scope *S,
+                                                            Declarator &D);
+
+  /// Register \p FD as specialization of \p BaseFD in the current `omp
+  /// begin/end declare variant` scope.
+  void ActOnFinishedFunctionDefinitionInOpenMPDeclareVariantScope(
+      FunctionDecl *FD, FunctionDecl *BaseFD);
 
 public:
-  /// Struct to store the context selectors info for declare variant directive.
+
+  /// Can we exit a scope at the moment.
+  bool isInOpenMPDeclareVariantScope() {
+    return !OMPDeclareVariantScopes.empty();
+  }
+
+  /// Given the potential call expression \p Call, determine if there is a
+  /// specialization via the OpenMP declare variant mechanism available. If
+  /// there is, return the specialized call expression, otherwise return the
+  /// original \p Call.
+  ExprResult ActOnOpenMPCall(Sema &S, ExprResult Call, Scope *Scope,
+                             SourceLocation LParenLoc, MultiExprArg ArgExprs,
+                             SourceLocation RParenLoc, Expr *ExecConfig);
+
+  /// Handle a `omp begin declare variant`.
+  void ActOnOpenMPBeginDeclareVariant(SourceLocation Loc, OMPTraitInfo &TI);
+
+  /// Handle a `omp end declare variant`.
+  void ActOnOpenMPEndDeclareVariant();
 
   /// Checks if the variant/multiversion functions are compatible.
   bool areMultiversionVariantFunctionsCompatible(

clang/lib/AST/DeclarationName.cpp

@@ -17,6 +17,7 @@
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclTemplate.h"
+#include "clang/AST/OpenMPClause.h"
 #include "clang/AST/PrettyPrinter.h"
 #include "clang/AST/Type.h"
 #include "clang/AST/TypeLoc.h"
@@ -138,8 +139,19 @@ void DeclarationName::print(raw_ostream &OS,
                             const PrintingPolicy &Policy) const {
   switch (getNameKind()) {
   case DeclarationName::Identifier:
-    if (const IdentifierInfo *II = getAsIdentifierInfo())
-      OS << II->getName();
+    if (const IdentifierInfo *II = getAsIdentifierInfo()) {
+      StringRef Name = II->getName();
+      // If this is a mangled OpenMP variant name we strip off the mangling for
+      // printing. It should not be visible to the user at all.
+      if (II->isMangledOpenMPVariantName()) {
+        std::pair<StringRef, StringRef> NameContextPair =
+            Name.split(getOpenMPVariantManglingSeparatorStr());
+        OS << NameContextPair.first << "["
+           << OMPTraitInfo(NameContextPair.second) << "]";
+      } else {
+        OS << Name;
+      }
+    }
     return;
 
   case DeclarationName::ObjCZeroArgSelector:

clang/lib/AST/OpenMPClause.cpp

@@ -1967,6 +1967,63 @@ void OMPTraitInfo::print(llvm::raw_ostream &OS,
   }
 }
 
+std::string OMPTraitInfo::getMangledName() const {
+  std::string MangledName;
+  llvm::raw_string_ostream OS(MangledName);
+  for (const OMPTraitInfo::OMPTraitSet &Set : Sets) {
+    OS << '.' << 'S' << unsigned(Set.Kind);
+    for (const OMPTraitInfo::OMPTraitSelector &Selector : Set.Selectors) {
+
+      bool AllowsTraitScore = false;
+      bool RequiresProperty = false;
+      isValidTraitSelectorForTraitSet(
+          Selector.Kind, Set.Kind, AllowsTraitScore, RequiresProperty);
+      OS << '.' << 's' << unsigned(Selector.Kind);
+
+      if (!RequiresProperty ||
+          Selector.Kind == TraitSelector::user_condition)
+        continue;
+
+      for (const OMPTraitInfo::OMPTraitProperty &Property : Selector.Properties)
+        OS << '.' << 'P'
+           << getOpenMPContextTraitPropertyName(Property.Kind);
+    }
+  }
+  return OS.str();
+}
+
+OMPTraitInfo::OMPTraitInfo(StringRef MangledName) {
+  unsigned long U;
+  do {
+    if (!MangledName.consume_front(".S"))
+      break;
+    if (MangledName.consumeInteger(10, U))
+      break;
+    Sets.push_back(OMPTraitSet());
+    OMPTraitSet &Set = Sets.back();
+    Set.Kind = TraitSet(U);
+    do {
+      if (!MangledName.consume_front(".s"))
+        break;
+      if (MangledName.consumeInteger(10, U))
+        break;
+      Set.Selectors.push_back(OMPTraitSelector());
+      OMPTraitSelector &Selector = Set.Selectors.back();
+      Selector.Kind = TraitSelector(U);
+      do {
+        if (!MangledName.consume_front(".P"))
+          break;
+        Selector.Properties.push_back(OMPTraitProperty());
+        OMPTraitProperty &Property = Selector.Properties.back();
+        std::pair<StringRef, StringRef> PropRestPair = MangledName.split('.');
+        Property.Kind =
+            getOpenMPContextTraitPropertyKind(Set.Kind, PropRestPair.first);
+        MangledName = PropRestPair.second;
+      } while (true);
+    } while (true);
+  } while (true);
+}
+
 llvm::raw_ostream &clang::operator<<(llvm::raw_ostream &OS,
                                      const OMPTraitInfo &TI) {
   LangOptions LO;

clang/lib/CodeGen/CGOpenMPRuntime.cpp

@@ -1275,52 +1275,6 @@ CGOpenMPRuntime::CGOpenMPRuntime(CodeGenModule &CGM, StringRef FirstSeparator,
   loadOffloadInfoMetadata();
 }
 
-bool CGOpenMPRuntime::tryEmitDeclareVariant(const GlobalDecl &NewGD,
-                                            const GlobalDecl &OldGD,
-                                            llvm::GlobalValue *OrigAddr,
-                                            bool IsForDefinition) {
-  // Emit at least a definition for the aliasee if the the address of the
-  // original function is requested.
-  if (IsForDefinition || OrigAddr)
-    (void)CGM.GetAddrOfGlobal(NewGD);
-  StringRef NewMangledName = CGM.getMangledName(NewGD);
-  llvm::GlobalValue *Addr = CGM.GetGlobalValue(NewMangledName);
-  if (Addr && !Addr->isDeclaration()) {
-    const auto *D = cast<FunctionDecl>(OldGD.getDecl());
-    const CGFunctionInfo &FI = CGM.getTypes().arrangeGlobalDeclaration(NewGD);
-    llvm::Type *DeclTy = CGM.getTypes().GetFunctionType(FI);
-
-    // Create a reference to the named value.  This ensures that it is emitted
-    // if a deferred decl.
-    llvm::GlobalValue::LinkageTypes LT = CGM.getFunctionLinkage(OldGD);
-
-    // Create the new alias itself, but don't set a name yet.
-    auto *GA =
-        llvm::GlobalAlias::create(DeclTy, 0, LT, "", Addr, &CGM.getModule());
-
-    if (OrigAddr) {
-      assert(OrigAddr->isDeclaration() && "Expected declaration");
-
-      GA->takeName(OrigAddr);
-      OrigAddr->replaceAllUsesWith(
-          llvm::ConstantExpr::getBitCast(GA, OrigAddr->getType()));
-      OrigAddr->eraseFromParent();
-    } else {
-      GA->setName(CGM.getMangledName(OldGD));
-    }
-
-    // Set attributes which are particular to an alias; this is a
-    // specialization of the attributes which may be set on a global function.
-    if (D->hasAttr<WeakAttr>() || D->hasAttr<WeakRefAttr>() ||
-        D->isWeakImported())
-      GA->setLinkage(llvm::Function::WeakAnyLinkage);
-
-    CGM.SetCommonAttributes(OldGD, GA);
-    return true;
-  }
-  return false;
-}
-
 void CGOpenMPRuntime::clear() {
   InternalVars.clear();
   // Clean non-target variable declarations possibly used only in debug info.
@@ -1334,14 +1288,6 @@ void CGOpenMPRuntime::clear() {
       continue;
     GV->eraseFromParent();
   }
-  // Emit aliases for the deferred aliasees.
-  for (const auto &Pair : DeferredVariantFunction) {
-    StringRef MangledName = CGM.getMangledName(Pair.second.second);
-    llvm::GlobalValue *Addr = CGM.GetGlobalValue(MangledName);
-    // If not able to emit alias, just emit original declaration.
-    (void)tryEmitDeclareVariant(Pair.second.first, Pair.second.second, Addr,
-                                /*IsForDefinition=*/false);
-  }
 }
 
 std::string CGOpenMPRuntime::getName(ArrayRef<StringRef> Parts) const {
@@ -11393,57 +11339,6 @@ Address CGOpenMPRuntime::getAddressOfLocalVariable(CodeGenFunction &CGF,
   return Address(Addr, Align);
 }
 
-/// Finds the variant function that matches current context with its context
-/// selector.
-static const FunctionDecl *getDeclareVariantFunction(CodeGenModule &CGM,
-                                                     const FunctionDecl *FD) {
-  if (!FD->hasAttrs() || !FD->hasAttr<OMPDeclareVariantAttr>())
-    return FD;
-
-  SmallVector<Expr *, 8> VariantExprs;
-  SmallVector<VariantMatchInfo, 8> VMIs;
-  for (const auto *A : FD->specific_attrs<OMPDeclareVariantAttr>()) {
-    const OMPTraitInfo &TI = *A->getTraitInfos();
-    VMIs.push_back(VariantMatchInfo());
-    TI.getAsVariantMatchInfo(CGM.getContext(), VMIs.back(),
-                             /* DeviceSetOnly */ false);
-    VariantExprs.push_back(A->getVariantFuncRef());
-  }
-
-  OMPContext Ctx(CGM.getLangOpts().OpenMPIsDevice, CGM.getTriple());
-  // FIXME: Keep the context in the OMPIRBuilder so we can add constructs as we
-  //        build them.
-
-  int BestMatchIdx = getBestVariantMatchForContext(VMIs, Ctx);
-  if (BestMatchIdx < 0)
-    return FD;
-
-  return cast<FunctionDecl>(
-      cast<DeclRefExpr>(VariantExprs[BestMatchIdx]->IgnoreParenImpCasts())
-          ->getDecl());
-}
-
-bool CGOpenMPRuntime::emitDeclareVariant(GlobalDecl GD, bool IsForDefinition) {
-  const auto *D = cast<FunctionDecl>(GD.getDecl());
-  // If the original function is defined already, use its definition.
-  StringRef MangledName = CGM.getMangledName(GD);
-  llvm::GlobalValue *Orig = CGM.GetGlobalValue(MangledName);
-  if (Orig && !Orig->isDeclaration())
-    return false;
-  const FunctionDecl *NewFD = getDeclareVariantFunction(CGM, D);
-  // Emit original function if it does not have declare variant attribute or the
-  // context does not match.
-  if (NewFD == D)
-    return false;
-  GlobalDecl NewGD = GD.getWithDecl(NewFD);
-  if (tryEmitDeclareVariant(NewGD, GD, Orig, IsForDefinition)) {
-    DeferredVariantFunction.erase(D);
-    return true;
-  }
-  DeferredVariantFunction.insert(std::make_pair(D, std::make_pair(NewGD, GD)));
-  return true;
-}
-
 CGOpenMPRuntime::NontemporalDeclsRAII::NontemporalDeclsRAII(
     CodeGenModule &CGM, const OMPLoopDirective &S)
     : CGM(CGM), NeedToPush(S.hasClausesOfKind<OMPNontemporalClause>()) {