[OpenMP] Allow traits for the OpenMP context selector isa

It was unclear what `isa` was supposed to mean so we did not provide any
traits for this context selector. With this patch we will allow *any*
string or identifier. We use the target attribute and target info to
determine if the trait matches. In other words, we will check if the
provided value is a target feature that is available (at the call site).

Fixes PR46338

Reviewed By: ABataev

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

Author: jdoerfert

clang/include/clang/AST/OpenMPClause.h

@@ -7635,6 +7635,10 @@ class OMPClausePrinter final : public OMPClauseVisitor<OMPClausePrinter> {
 
 struct OMPTraitProperty {
   llvm::omp::TraitProperty Kind = llvm::omp::TraitProperty::invalid;
+
+  /// The raw string as we parsed it. This is needed for the `isa` trait set
+  /// (which accepts anything) and (later) extensions.
+  StringRef RawString;
 };
 struct OMPTraitSelector {
   Expr *ScoreOrCondition = nullptr;
@@ -7692,6 +7696,23 @@ class OMPTraitInfo {
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const OMPTraitInfo &TI);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const OMPTraitInfo *TI);
 
+/// Clang specific specialization of the OMPContext to lookup target features.
+struct TargetOMPContext final : public llvm::omp::OMPContext {
+
+  TargetOMPContext(ASTContext &ASTCtx,
+                   std::function<void(StringRef)> &&DiagUnknownTrait,
+                   const FunctionDecl *CurrentFunctionDecl);
+  virtual ~TargetOMPContext() = default;
+
+  /// See llvm::omp::OMPContext::matchesISATrait
+  bool matchesISATrait(StringRef RawString) const override;
+
+private:
+  std::function<bool(StringRef)> FeatureValidityCheck;
+  std::function<void(StringRef)> DiagUnknownTrait;
+  llvm::StringMap<bool> FeatureMap;
+};
+
 } // namespace clang
 
 #endif // LLVM_CLANG_AST_OPENMPCLAUSE_H

clang/include/clang/Basic/DiagnosticParseKinds.td

@@ -1278,6 +1278,11 @@ def warn_omp_declare_variant_string_literal_or_identifier
               "%select{set|selector|property}0; "
               "%select{set|selector|property}0 skipped">,
       InGroup<OpenMPClauses>;
+def warn_unknown_begin_declare_variant_isa_trait
+    : Warning<"isa trait '%0' is not known to the current target; verify the "
+              "spelling or consider restricting the context selector with the "
+              "'arch' selector further">,
+      InGroup<SourceUsesOpenMP>;
 def note_omp_declare_variant_ctx_options
     : Note<"context %select{set|selector|property}0 options are: %1">;
 def warn_omp_declare_variant_expected

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -10320,6 +10320,11 @@ def warn_nested_declare_variant
     : Warning<"nesting `omp begin/end declare variant` is not supported yet; "
               "nested context ignored">,
       InGroup<SourceUsesOpenMP>;
+def warn_unknown_declare_variant_isa_trait
+    : Warning<"isa trait '%0' is not known to the current target; verify the "
+              "spelling or consider restricting the context selector with the "
+              "'arch' selector further">,
+      InGroup<SourceUsesOpenMP>;
 def err_omp_non_pointer_type_array_shaping_base : Error<
   "expected expression with a pointer to a complete type as a base of an array "
   "shaping operation">;

clang/lib/AST/OpenMPClause.cpp

@@ -17,6 +17,7 @@
 #include "clang/AST/DeclOpenMP.h"
 #include "clang/Basic/LLVM.h"
 #include "clang/Basic/OpenMPKinds.h"
+#include "clang/Basic/TargetInfo.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -2131,9 +2132,10 @@ void OMPTraitInfo::getAsVariantMatchInfo(ASTContext &ASTCtx,
                 Selector.ScoreOrCondition->getIntegerConstantExpr(ASTCtx))
           VMI.addTrait(CondVal->isNullValue()
                            ? TraitProperty::user_condition_false
-                           : TraitProperty::user_condition_true);
+                           : TraitProperty::user_condition_true,
+                       "<condition>");
         else
-          VMI.addTrait(TraitProperty::user_condition_false);
+          VMI.addTrait(TraitProperty::user_condition_false, "<condition>");
         continue;
       }
 
@@ -2143,11 +2145,12 @@ void OMPTraitInfo::getAsVariantMatchInfo(ASTContext &ASTCtx,
         if ((Score = Selector.ScoreOrCondition->getIntegerConstantExpr(ASTCtx)))
           ScorePtr = &*Score;
         else
-          VMI.addTrait(TraitProperty::user_condition_false);
+          VMI.addTrait(TraitProperty::user_condition_false,
+                       "<non-constant-score>");
       }
 
       for (const OMPTraitProperty &Property : Selector.Properties)
-        VMI.addTrait(Set.Kind, Property.Kind, ScorePtr);
+        VMI.addTrait(Set.Kind, Property.Kind, Property.RawString, ScorePtr);
 
       if (Set.Kind != TraitSet::construct)
         continue;
@@ -2204,7 +2207,8 @@ void OMPTraitInfo::print(llvm::raw_ostream &OS,
           if (!FirstProperty)
             OS << ", ";
           FirstProperty = false;
-          OS << getOpenMPContextTraitPropertyName(Property.Kind);
+          OS << getOpenMPContextTraitPropertyName(Property.Kind,
+                                                  Property.RawString);
         }
       }
       OS << ")";
@@ -2231,7 +2235,9 @@ std::string OMPTraitInfo::getMangledName() const {
         continue;
 
       for (const OMPTraitProperty &Property : Selector.Properties)
-        OS << '$' << 'P' << getOpenMPContextTraitPropertyName(Property.Kind);
+        OS << '$' << 'P'
+           << getOpenMPContextTraitPropertyName(Property.Kind,
+                                                Property.RawString);
     }
   }
   return OS.str();
@@ -2261,8 +2267,9 @@ OMPTraitInfo::OMPTraitInfo(StringRef MangledName) {
         Selector.Properties.push_back(OMPTraitProperty());
         OMPTraitProperty &Property = Selector.Properties.back();
         std::pair<StringRef, StringRef> PropRestPair = MangledName.split('$');
-        Property.Kind =
-            getOpenMPContextTraitPropertyKind(Set.Kind, PropRestPair.first);
+        Property.RawString = PropRestPair.first;
+        Property.Kind = getOpenMPContextTraitPropertyKind(
+            Set.Kind, Selector.Kind, PropRestPair.first);
         MangledName = PropRestPair.second;
       } while (true);
     } while (true);
@@ -2280,3 +2287,24 @@ llvm::raw_ostream &clang::operator<<(llvm::raw_ostream &OS,
                                      const OMPTraitInfo *TI) {
   return TI ? OS << *TI : OS;
 }
+
+TargetOMPContext::TargetOMPContext(
+    ASTContext &ASTCtx, std::function<void(StringRef)> &&DiagUnknownTrait,
+    const FunctionDecl *CurrentFunctionDecl)
+    : OMPContext(ASTCtx.getLangOpts().OpenMPIsDevice,
+                 ASTCtx.getTargetInfo().getTriple()),
+      FeatureValidityCheck([&](StringRef FeatureName) {
+        return ASTCtx.getTargetInfo().isValidFeatureName(FeatureName);
+      }),
+      DiagUnknownTrait(std::move(DiagUnknownTrait)) {
+  ASTCtx.getFunctionFeatureMap(FeatureMap, CurrentFunctionDecl);
+}
+
+bool TargetOMPContext::matchesISATrait(StringRef RawString) const {
+  auto It = FeatureMap.find(RawString);
+  if (It != FeatureMap.end())
+    return It->second;
+  if (!FeatureValidityCheck(RawString))
+    DiagUnknownTrait(RawString);
+  return false;
+}

clang/lib/Parse/ParseOpenMP.cpp

@@ -869,7 +869,8 @@ void Parser::parseOMPTraitPropertyKind(
     return;
   }
 
-  TIProperty.Kind = getOpenMPContextTraitPropertyKind(Set, Name);
+  TIProperty.RawString = Name;
+  TIProperty.Kind = getOpenMPContextTraitPropertyKind(Set, Selector, Name);
   if (TIProperty.Kind != TraitProperty::invalid) {
     if (checkForDuplicates(*this, Name, NameLoc, Seen, CONTEXT_TRAIT_LVL))
       TIProperty.Kind = TraitProperty::invalid;
@@ -910,7 +911,7 @@ void Parser::parseOMPTraitPropertyKind(
        {TraitSet::construct, TraitSet::user, TraitSet::implementation,
         TraitSet::device}) {
     TraitProperty PropertyForName =
-        getOpenMPContextTraitPropertyKind(PotentialSet, Name);
+        getOpenMPContextTraitPropertyKind(PotentialSet, Selector, Name);
     if (PropertyForName == TraitProperty::invalid)
       continue;
     Diag(NameLoc, diag::note_omp_declare_variant_ctx_try)
@@ -949,8 +950,8 @@ static bool checkExtensionProperty(Parser &P, SourceLocation Loc,
     for (OMPTraitProperty &SeenProp : TISelector.Properties)
       if (IsMatchExtension(SeenProp)) {
         P.Diag(Loc, diag::err_omp_variant_ctx_second_match_extension);
-        StringRef SeenName =
-            llvm::omp::getOpenMPContextTraitPropertyName(SeenProp.Kind);
+        StringRef SeenName = llvm::omp::getOpenMPContextTraitPropertyName(
+            SeenProp.Kind, SeenProp.RawString);
         SourceLocation SeenLoc = Seen[SeenName];
         P.Diag(SeenLoc, diag::note_omp_declare_variant_ctx_used_here)
             << CONTEXT_TRAIT_LVL << SeenName;
@@ -995,11 +996,13 @@ void Parser::parseOMPContextProperty(OMPTraitSelector &TISelector,
   }
 
   Diag(PropertyLoc, diag::warn_omp_ctx_incompatible_property_for_selector)
-      << getOpenMPContextTraitPropertyName(TIProperty.Kind)
+      << getOpenMPContextTraitPropertyName(TIProperty.Kind,
+                                           TIProperty.RawString)
       << getOpenMPContextTraitSelectorName(TISelector.Kind)
       << getOpenMPContextTraitSetName(Set);
   Diag(PropertyLoc, diag::note_omp_ctx_compatible_set_and_selector_for_property)
-      << getOpenMPContextTraitPropertyName(TIProperty.Kind)
+      << getOpenMPContextTraitPropertyName(TIProperty.Kind,
+                                           TIProperty.RawString)
       << getOpenMPContextTraitSelectorName(
              getOpenMPContextTraitSelectorForProperty(TIProperty.Kind))
       << getOpenMPContextTraitSetName(
@@ -1045,8 +1048,8 @@ void Parser::parseOMPTraitSelectorKind(
   for (const auto &PotentialSet :
        {TraitSet::construct, TraitSet::user, TraitSet::implementation,
         TraitSet::device}) {
-    TraitProperty PropertyForName =
-        getOpenMPContextTraitPropertyKind(PotentialSet, Name);
+    TraitProperty PropertyForName = getOpenMPContextTraitPropertyKind(
+        PotentialSet, TraitSelector::invalid, Name);
     if (PropertyForName == TraitProperty::invalid)
       continue;
     Diag(NameLoc, diag::note_omp_declare_variant_ctx_is_a)
@@ -1140,7 +1143,8 @@ void Parser::parseOMPContextSelector(
 
   if (!RequiresProperty) {
     TISelector.Properties.push_back(
-        {getOpenMPContextTraitPropertyForSelector(TISelector.Kind)});
+        {getOpenMPContextTraitPropertyForSelector(TISelector.Kind),
+         getOpenMPContextTraitSelectorName(TISelector.Kind)});
     return;
   }
 
@@ -1157,7 +1161,8 @@ void Parser::parseOMPContextSelector(
     if (!Condition.isUsable())
       return FinishSelector();
     TISelector.ScoreOrCondition = Condition.get();
-    TISelector.Properties.push_back({TraitProperty::user_condition_unknown});
+    TISelector.Properties.push_back(
+        {TraitProperty::user_condition_unknown, "<condition>"});
     return;
   }
 
@@ -1236,8 +1241,8 @@ void Parser::parseOMPTraitSetKind(OMPTraitSet &TISet,
   for (const auto &PotentialSet :
        {TraitSet::construct, TraitSet::user, TraitSet::implementation,
         TraitSet::device}) {
-    TraitProperty PropertyForName =
-        getOpenMPContextTraitPropertyKind(PotentialSet, Name);
+    TraitProperty PropertyForName = getOpenMPContextTraitPropertyKind(
+        PotentialSet, TraitSelector::invalid, Name);
     if (PropertyForName == TraitProperty::invalid)
       continue;
     Diag(NameLoc, diag::note_omp_declare_variant_ctx_is_a)
@@ -1820,8 +1825,15 @@ Parser::DeclGroupPtrTy Parser::ParseOpenMPDeclarativeDirectiveWithExtDecl(
     VariantMatchInfo VMI;
     ASTContext &ASTCtx = Actions.getASTContext();
     TI.getAsVariantMatchInfo(ASTCtx, VMI);
-    OMPContext OMPCtx(ASTCtx.getLangOpts().OpenMPIsDevice,
-                      ASTCtx.getTargetInfo().getTriple());
+
+    std::function<void(StringRef)> DiagUnknownTrait = [this, Loc](
+                                                          StringRef ISATrait) {
+      // TODO Track the selector locations in a way that is accessible here to
+      // improve the diagnostic location.
+      Diag(Loc, diag::warn_unknown_begin_declare_variant_isa_trait) << ISATrait;
+    };
+    TargetOMPContext OMPCtx(ASTCtx, std::move(DiagUnknownTrait),
+                            /* CurrentFunctionDecl */ nullptr);
 
     if (isVariantApplicableInContext(VMI, OMPCtx, /* DeviceSetOnly */ true)) {
       Actions.ActOnOpenMPBeginDeclareVariant(Loc, TI);

clang/lib/Sema/SemaOpenMP.cpp

@@ -5895,8 +5895,15 @@ ExprResult Sema::ActOnOpenMPCall(ExprResult Call, Scope *Scope,
     return Call;
 
   ASTContext &Context = getASTContext();
-  OMPContext OMPCtx(getLangOpts().OpenMPIsDevice,
-                    Context.getTargetInfo().getTriple());
+  std::function<void(StringRef)> DiagUnknownTrait = [this,
+                                                     CE](StringRef ISATrait) {
+    // TODO Track the selector locations in a way that is accessible here to
+    // improve the diagnostic location.
+    Diag(CE->getBeginLoc(), diag::warn_unknown_declare_variant_isa_trait)
+        << ISATrait;
+  };
+  TargetOMPContext OMPCtx(Context, std::move(DiagUnknownTrait),
+                          getCurFunctionDecl());
 
   SmallVector<Expr *, 4> Exprs;
   SmallVector<VariantMatchInfo, 4> VMIs;
@@ -5908,7 +5915,8 @@ ExprResult Sema::ActOnOpenMPCall(ExprResult Call, Scope *Scope,
       VariantMatchInfo VMI;
       OMPTraitInfo &TI = A->getTraitInfo();
       TI.getAsVariantMatchInfo(Context, VMI);
-      if (!isVariantApplicableInContext(VMI, OMPCtx, /* DeviceSetOnly */ false))
+      if (!isVariantApplicableInContext(VMI, OMPCtx,
+                                        /* DeviceSetOnly */ false))
         continue;
 
       VMIs.push_back(VMI);

clang/test/OpenMP/declare_variant_device_isa_codegen_1.c

@@ -0,0 +1,49 @@
+// RUN: %clang_cc1 -verify -fopenmp -x c -triple %itanium_abi_triple -emit-llvm %s -o - -fopenmp-version=50 | FileCheck %s --check-prefix=GENERIC
+// RUN: %clang_cc1 -fopenmp -x c++ -std=c++11 -triple %itanium_abi_triple -fexceptions -fcxx-exceptions -emit-pch -o %t -fopenmp-version=50 %s
+// RUN: %clang_cc1 -fopenmp -x c++ -triple %itanium_abi_triple -fexceptions -fcxx-exceptions -std=c++11 -include-pch %t -verify %s -emit-llvm -o - -fopenmp-version=50 | FileCheck %s --check-prefix=GENERIC
+
+// RUN: %clang_cc1 -target-feature +avx512f -verify -fopenmp -x c -triple %itanium_abi_triple -emit-llvm %s -o - -fopenmp-version=50 | FileCheck %s --check-prefix=WITHFEATURE
+// RUN: %clang_cc1 -target-feature +avx512f -fopenmp -x c++ -std=c++11 -triple %itanium_abi_triple -fexceptions -fcxx-exceptions -emit-pch -o %t -fopenmp-version=50 %s
+// RUN: %clang_cc1 -target-feature +avx512f -fopenmp -x c++ -triple %itanium_abi_triple -fexceptions -fcxx-exceptions -std=c++11 -include-pch %t -verify %s -emit-llvm -o - -fopenmp-version=50 | FileCheck %s --check-prefix=WITHFEATURE
+
+// expected-no-diagnostics
+
+// Test taken from PR46338 (by linna su)
+
+#ifndef HEADER
+#define HEADER
+
+void base_saxpy(int, float, float *, float *);
+void avx512_saxpy(int, float, float *, float *);
+
+#pragma omp declare variant(avx512_saxpy) \
+    match(device = {isa(avx512f)})
+void base_saxpy(int n, float s, float *x, float *y) {
+#pragma omp parallel for
+  for (int i = 0; i < n; i++)
+    y[i] = s * x[i] + y[i];
+}
+
+void avx512_saxpy(int n, float s, float *x, float *y) {
+#pragma omp parallel for simd simdlen(16) aligned(x, y : 64)
+  for (int i = 0; i < n; i++)
+    y[i] = s * x[i] + y[i];
+}
+
+void caller(int n, float s, float *x, float *y) {
+  // GENERIC:     define void @{{.*}}caller
+  // GENERIC:      call void @{{.*}}base_saxpy
+  // WITHFEATURE: define void @{{.*}}caller
+  // WITHFEATURE:  call void @{{.*}}avx512_saxpy
+  base_saxpy(n, s, x, y);
+}
+
+__attribute__((target("avx512f"))) void variant_caller(int n, float s, float *x, float *y) {
+  // GENERIC:     define void @{{.*}}variant_caller
+  // GENERIC:      call void @{{.*}}avx512_saxpy
+  // WITHFEATURE: define void @{{.*}}variant_caller
+  // WITHFEATURE:  call void @{{.*}}avx512_saxpy
+  base_saxpy(n, s, x, y);
+}
+
+#endif

clang/test/OpenMP/declare_variant_messages.c

@@ -137,6 +137,18 @@ void marked_variant(void);
 #pragma omp declare variant(marked_variant) match(xxx={}) // expected-warning {{'xxx' is not a valid context set in a `declare variant`; set ignored}} expected-warning {{variant function in '#pragma omp declare variant' is itself marked as '#pragma omp declare variant'}} expected-note {{context set options are: 'construct' 'device' 'implementation' 'user'}} expected-note {{the ignored set spans until here}}
 void marked(void);
 
+#pragma omp declare variant(foo) match(device = {isa("foo")})
+int unknown_isa_trait(void);
+#pragma omp declare variant(foo) match(device = {isa(foo)})
+int unknown_isa_trait2(void);
+#pragma omp declare variant(foo) match(device = {kind(fpga), isa(bar)})
+int ignored_isa_trait(void);
+
+void caller() {
+  unknown_isa_trait();  // expected-warning {{isa trait 'foo' is not known to the current target; verify the spelling or consider restricting the context selector with the 'arch' selector further}}
+  unknown_isa_trait2(); // expected-warning {{isa trait 'foo' is not known to the current target; verify the spelling or consider restricting the context selector with the 'arch' selector further}}
+  ignored_isa_trait();
+}
 
 #pragma omp declare variant // expected-error {{function declaration is expected after 'declare variant' directive}}