1 //===- OMPContext.cpp ------ Collection of helpers for OpenMP contexts ----===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 /// \file
9 ///
10 /// This file implements helper functions and classes to deal with OpenMP
11 /// contexts as used by `[begin/end] declare variant` and `metadirective`.
12 ///
13 //===----------------------------------------------------------------------===//
14 
15 #include "llvm/Frontend/OpenMP/OMPContext.h"
16 #include "llvm/ADT/StringRef.h"
17 #include "llvm/ADT/StringSwitch.h"
18 #include "llvm/Support/Debug.h"
19 #include "llvm/Support/raw_ostream.h"
20 #include "llvm/TargetParser/Triple.h"
21 
22 #define DEBUG_TYPE "openmp-ir-builder"
23 
24 using namespace llvm;
25 using namespace omp;
26 
27 OMPContext::OMPContext(bool IsDeviceCompilation, Triple TargetTriple) {
28   // Add the appropriate device kind trait based on the triple and the
29   // IsDeviceCompilation flag.
30   ActiveTraits.set(unsigned(IsDeviceCompilation
31                                 ? TraitProperty::device_kind_nohost
32                                 : TraitProperty::device_kind_host));
33   switch (TargetTriple.getArch()) {
34   case Triple::arm:
35   case Triple::armeb:
36   case Triple::aarch64:
37   case Triple::aarch64_be:
38   case Triple::aarch64_32:
39   case Triple::mips:
40   case Triple::mipsel:
41   case Triple::mips64:
42   case Triple::mips64el:
43   case Triple::ppc:
44   case Triple::ppcle:
45   case Triple::ppc64:
46   case Triple::ppc64le:
47   case Triple::x86:
48   case Triple::x86_64:
49     ActiveTraits.set(unsigned(TraitProperty::device_kind_cpu));
50     break;
51   case Triple::amdgcn:
52   case Triple::nvptx:
53   case Triple::nvptx64:
54     ActiveTraits.set(unsigned(TraitProperty::device_kind_gpu));
55     break;
56   default:
57     break;
58   }
59 
60   // Add the appropriate device architecture trait based on the triple.
61 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
62   if (TraitSelector::TraitSelectorEnum == TraitSelector::device_arch) {        \
63     if (TargetTriple.getArch() == TargetTriple.getArchTypeForLLVMName(Str))    \
64       ActiveTraits.set(unsigned(TraitProperty::Enum));                         \
65     if (StringRef(Str) == StringRef("x86_64") &&                               \
66         TargetTriple.getArch() == Triple::x86_64)                              \
67       ActiveTraits.set(unsigned(TraitProperty::Enum));                         \
68   }
69 #include "llvm/Frontend/OpenMP/OMPKinds.def"
70 
71   // TODO: What exactly do we want to see as device ISA trait?
72   //       The discussion on the list did not seem to have come to an agreed
73   //       upon solution.
74 
75   // LLVM is the "OpenMP vendor" but we could also interpret vendor as the
76   // target vendor.
77   ActiveTraits.set(unsigned(TraitProperty::implementation_vendor_llvm));
78 
79   // The user condition true is accepted but not false.
80   ActiveTraits.set(unsigned(TraitProperty::user_condition_true));
81 
82   // This is for sure some device.
83   ActiveTraits.set(unsigned(TraitProperty::device_kind_any));
84 
85   LLVM_DEBUG({
86     dbgs() << "[" << DEBUG_TYPE
87            << "] New OpenMP context with the following properties:\n";
88     for (unsigned Bit : ActiveTraits.set_bits()) {
89       TraitProperty Property = TraitProperty(Bit);
90       dbgs() << "\t " << getOpenMPContextTraitPropertyFullName(Property)
91              << "\n";
92     }
93   });
94 }
95 
96 /// Return true if \p C0 is a subset of \p C1. Note that both arrays are
97 /// expected to be sorted.
98 template <typename T> static bool isSubset(ArrayRef<T> C0, ArrayRef<T> C1) {
99 #ifdef EXPENSIVE_CHECKS
100   assert(llvm::is_sorted(C0) && llvm::is_sorted(C1) &&
101          "Expected sorted arrays!");
102 #endif
103   if (C0.size() > C1.size())
104     return false;
105   auto It0 = C0.begin(), End0 = C0.end();
106   auto It1 = C1.begin(), End1 = C1.end();
107   while (It0 != End0) {
108     if (It1 == End1)
109       return false;
110     if (*It0 == *It1) {
111       ++It0;
112       ++It1;
113       continue;
114     }
115     ++It0;
116   }
117   return true;
118 }
119 
120 /// Return true if \p C0 is a strict subset of \p C1. Note that both arrays are
121 /// expected to be sorted.
122 template <typename T>
123 static bool isStrictSubset(ArrayRef<T> C0, ArrayRef<T> C1) {
124   if (C0.size() >= C1.size())
125     return false;
126   return isSubset<T>(C0, C1);
127 }
128 
129 static bool isStrictSubset(const VariantMatchInfo &VMI0,
130                            const VariantMatchInfo &VMI1) {
131   // If all required traits are a strict subset and the ordered vectors storing
132   // the construct traits, we say it is a strict subset. Note that the latter
133   // relation is not required to be strict.
134   if (VMI0.RequiredTraits.count() >= VMI1.RequiredTraits.count())
135     return false;
136   for (unsigned Bit : VMI0.RequiredTraits.set_bits())
137     if (!VMI1.RequiredTraits.test(Bit))
138       return false;
139   if (!isSubset<TraitProperty>(VMI0.ConstructTraits, VMI1.ConstructTraits))
140     return false;
141   return true;
142 }
143 
144 static int isVariantApplicableInContextHelper(
145     const VariantMatchInfo &VMI, const OMPContext &Ctx,
146     SmallVectorImpl<unsigned> *ConstructMatches, bool DeviceSetOnly) {
147 
148   // The match kind determines if we need to match all traits, any of the
149   // traits, or none of the traits for it to be an applicable context.
150   enum MatchKind { MK_ALL, MK_ANY, MK_NONE };
151 
152   MatchKind MK = MK_ALL;
153   // Determine the match kind the user wants, "all" is the default and provided
154   // to the user only for completeness.
155   if (VMI.RequiredTraits.test(
156           unsigned(TraitProperty::implementation_extension_match_any)))
157     MK = MK_ANY;
158   if (VMI.RequiredTraits.test(
159           unsigned(TraitProperty::implementation_extension_match_none)))
160     MK = MK_NONE;
161 
162   // Helper to deal with a single property that was (not) found in the OpenMP
163   // context based on the match kind selected by the user via
164   // `implementation={extensions(match_[all,any,none])}'
165   auto HandleTrait = [MK](TraitProperty Property,
166                           bool WasFound) -> std::optional<bool> /* Result */ {
167     // For kind "any" a single match is enough but we ignore non-matched
168     // properties.
169     if (MK == MK_ANY) {
170       if (WasFound)
171         return true;
172       return std::nullopt;
173     }
174 
175     // In "all" or "none" mode we accept a matching or non-matching property
176     // respectively and move on. We are not done yet!
177     if ((WasFound && MK == MK_ALL) || (!WasFound && MK == MK_NONE))
178       return std::nullopt;
179 
180     // We missed a property, provide some debug output and indicate failure.
181     LLVM_DEBUG({
182       if (MK == MK_ALL)
183         dbgs() << "[" << DEBUG_TYPE << "] Property "
184                << getOpenMPContextTraitPropertyName(Property, "")
185                << " was not in the OpenMP context but match kind is all.\n";
186       if (MK == MK_NONE)
187         dbgs() << "[" << DEBUG_TYPE << "] Property "
188                << getOpenMPContextTraitPropertyName(Property, "")
189                << " was in the OpenMP context but match kind is none.\n";
190     });
191     return false;
192   };
193 
194   for (unsigned Bit : VMI.RequiredTraits.set_bits()) {
195     TraitProperty Property = TraitProperty(Bit);
196     if (DeviceSetOnly &&
197         getOpenMPContextTraitSetForProperty(Property) != TraitSet::device)
198       continue;
199 
200     // So far all extensions are handled elsewhere, we skip them here as they
201     // are not part of the OpenMP context.
202     if (getOpenMPContextTraitSelectorForProperty(Property) ==
203         TraitSelector::implementation_extension)
204       continue;
205 
206     bool IsActiveTrait = Ctx.ActiveTraits.test(unsigned(Property));
207 
208     // We overwrite the isa trait as it is actually up to the OMPContext hook to
209     // check the raw string(s).
210     if (Property == TraitProperty::device_isa___ANY)
211       IsActiveTrait = llvm::all_of(VMI.ISATraits, [&](StringRef RawString) {
212         return Ctx.matchesISATrait(RawString);
213       });
214 
215     if (std::optional<bool> Result = HandleTrait(Property, IsActiveTrait))
216       return *Result;
217   }
218 
219   if (!DeviceSetOnly) {
220     // We could use isSubset here but we also want to record the match
221     // locations.
222     unsigned ConstructIdx = 0, NoConstructTraits = Ctx.ConstructTraits.size();
223     for (TraitProperty Property : VMI.ConstructTraits) {
224       assert(getOpenMPContextTraitSetForProperty(Property) ==
225                  TraitSet::construct &&
226              "Variant context is ill-formed!");
227 
228       // Verify the nesting.
229       bool FoundInOrder = false;
230       while (!FoundInOrder && ConstructIdx != NoConstructTraits)
231         FoundInOrder = (Ctx.ConstructTraits[ConstructIdx++] == Property);
232       if (ConstructMatches)
233         ConstructMatches->push_back(ConstructIdx - 1);
234 
235       if (std::optional<bool> Result = HandleTrait(Property, FoundInOrder))
236         return *Result;
237 
238       if (!FoundInOrder) {
239         LLVM_DEBUG(dbgs() << "[" << DEBUG_TYPE << "] Construct property "
240                           << getOpenMPContextTraitPropertyName(Property, "")
241                           << " was not nested properly.\n");
242         return false;
243       }
244 
245       // TODO: Verify SIMD
246     }
247 
248     assert(isSubset<TraitProperty>(VMI.ConstructTraits, Ctx.ConstructTraits) &&
249            "Broken invariant!");
250   }
251 
252   if (MK == MK_ANY) {
253     LLVM_DEBUG(dbgs() << "[" << DEBUG_TYPE
254                       << "] None of the properties was in the OpenMP context "
255                          "but match kind is any.\n");
256     return false;
257   }
258 
259   return true;
260 }
261 
262 bool llvm::omp::isVariantApplicableInContext(const VariantMatchInfo &VMI,
263                                              const OMPContext &Ctx,
264                                              bool DeviceSetOnly) {
265   return isVariantApplicableInContextHelper(
266       VMI, Ctx, /* ConstructMatches */ nullptr, DeviceSetOnly);
267 }
268 
269 static APInt getVariantMatchScore(const VariantMatchInfo &VMI,
270                                   const OMPContext &Ctx,
271                                   SmallVectorImpl<unsigned> &ConstructMatches) {
272   APInt Score(64, 1);
273 
274   unsigned NoConstructTraits = VMI.ConstructTraits.size();
275   for (unsigned Bit : VMI.RequiredTraits.set_bits()) {
276     TraitProperty Property = TraitProperty(Bit);
277     // If there is a user score attached, use it.
278     if (VMI.ScoreMap.count(Property)) {
279       const APInt &UserScore = VMI.ScoreMap.lookup(Property);
280       assert(UserScore.uge(0) && "Expect non-negative user scores!");
281       Score += UserScore.getZExtValue();
282       continue;
283     }
284 
285     switch (getOpenMPContextTraitSetForProperty(Property)) {
286     case TraitSet::construct:
287       // We handle the construct traits later via the VMI.ConstructTraits
288       // container.
289       continue;
290     case TraitSet::implementation:
291       // No effect on the score (implementation defined).
292       continue;
293     case TraitSet::user:
294       // No effect on the score.
295       continue;
296     case TraitSet::device:
297       // Handled separately below.
298       break;
299     case TraitSet::invalid:
300       llvm_unreachable("Unknown trait set is not to be used!");
301     }
302 
303     // device={kind(any)} is "as if" no kind selector was specified.
304     if (Property == TraitProperty::device_kind_any)
305       continue;
306 
307     switch (getOpenMPContextTraitSelectorForProperty(Property)) {
308     case TraitSelector::device_kind:
309       Score += (1ULL << (NoConstructTraits + 0));
310       continue;
311     case TraitSelector::device_arch:
312       Score += (1ULL << (NoConstructTraits + 1));
313       continue;
314     case TraitSelector::device_isa:
315       Score += (1ULL << (NoConstructTraits + 2));
316       continue;
317     default:
318       continue;
319     }
320   }
321 
322   unsigned ConstructIdx = 0;
323   assert(NoConstructTraits == ConstructMatches.size() &&
324          "Mismatch in the construct traits!");
325   for (TraitProperty Property : VMI.ConstructTraits) {
326     assert(getOpenMPContextTraitSetForProperty(Property) ==
327                TraitSet::construct &&
328            "Ill-formed variant match info!");
329     (void)Property;
330     // ConstructMatches is the position p - 1 and we need 2^(p-1).
331     Score += (1ULL << ConstructMatches[ConstructIdx++]);
332   }
333 
334   LLVM_DEBUG(dbgs() << "[" << DEBUG_TYPE << "] Variant has a score of " << Score
335                     << "\n");
336   return Score;
337 }
338 
339 int llvm::omp::getBestVariantMatchForContext(
340     const SmallVectorImpl<VariantMatchInfo> &VMIs, const OMPContext &Ctx) {
341 
342   APInt BestScore(64, 0);
343   int BestVMIIdx = -1;
344   const VariantMatchInfo *BestVMI = nullptr;
345 
346   for (unsigned u = 0, e = VMIs.size(); u < e; ++u) {
347     const VariantMatchInfo &VMI = VMIs[u];
348 
349     SmallVector<unsigned, 8> ConstructMatches;
350     // If the variant is not applicable its not the best.
351     if (!isVariantApplicableInContextHelper(VMI, Ctx, &ConstructMatches,
352                                             /* DeviceSetOnly */ false))
353       continue;
354     // Check if its clearly not the best.
355     APInt Score = getVariantMatchScore(VMI, Ctx, ConstructMatches);
356     if (Score.ult(BestScore))
357       continue;
358     // Equal score need subset checks.
359     if (Score.eq(BestScore)) {
360       // Strict subset are never best.
361       if (isStrictSubset(VMI, *BestVMI))
362         continue;
363       // Same score and the current best is no strict subset so we keep it.
364       if (!isStrictSubset(*BestVMI, VMI))
365         continue;
366     }
367     // New best found.
368     BestVMI = &VMI;
369     BestVMIIdx = u;
370     BestScore = Score;
371   }
372 
373   return BestVMIIdx;
374 }
375 
376 TraitSet llvm::omp::getOpenMPContextTraitSetKind(StringRef S) {
377   return StringSwitch<TraitSet>(S)
378 #define OMP_TRAIT_SET(Enum, Str) .Case(Str, TraitSet::Enum)
379 #include "llvm/Frontend/OpenMP/OMPKinds.def"
380       .Default(TraitSet::invalid);
381 }
382 
383 TraitSet
384 llvm::omp::getOpenMPContextTraitSetForSelector(TraitSelector Selector) {
385   switch (Selector) {
386 #define OMP_TRAIT_SELECTOR(Enum, TraitSetEnum, Str, ReqProp)                   \
387   case TraitSelector::Enum:                                                    \
388     return TraitSet::TraitSetEnum;
389 #include "llvm/Frontend/OpenMP/OMPKinds.def"
390   }
391   llvm_unreachable("Unknown trait selector!");
392 }
393 TraitSet
394 llvm::omp::getOpenMPContextTraitSetForProperty(TraitProperty Property) {
395   switch (Property) {
396 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
397   case TraitProperty::Enum:                                                    \
398     return TraitSet::TraitSetEnum;
399 #include "llvm/Frontend/OpenMP/OMPKinds.def"
400   }
401   llvm_unreachable("Unknown trait set!");
402 }
403 StringRef llvm::omp::getOpenMPContextTraitSetName(TraitSet Kind) {
404   switch (Kind) {
405 #define OMP_TRAIT_SET(Enum, Str)                                               \
406   case TraitSet::Enum:                                                         \
407     return Str;
408 #include "llvm/Frontend/OpenMP/OMPKinds.def"
409   }
410   llvm_unreachable("Unknown trait set!");
411 }
412 
413 TraitSelector llvm::omp::getOpenMPContextTraitSelectorKind(StringRef S) {
414   return StringSwitch<TraitSelector>(S)
415 #define OMP_TRAIT_SELECTOR(Enum, TraitSetEnum, Str, ReqProp)                   \
416   .Case(Str, TraitSelector::Enum)
417 #include "llvm/Frontend/OpenMP/OMPKinds.def"
418       .Default(TraitSelector::invalid);
419 }
420 TraitSelector
421 llvm::omp::getOpenMPContextTraitSelectorForProperty(TraitProperty Property) {
422   switch (Property) {
423 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
424   case TraitProperty::Enum:                                                    \
425     return TraitSelector::TraitSelectorEnum;
426 #include "llvm/Frontend/OpenMP/OMPKinds.def"
427   }
428   llvm_unreachable("Unknown trait set!");
429 }
430 StringRef llvm::omp::getOpenMPContextTraitSelectorName(TraitSelector Kind) {
431   switch (Kind) {
432 #define OMP_TRAIT_SELECTOR(Enum, TraitSetEnum, Str, ReqProp)                   \
433   case TraitSelector::Enum:                                                    \
434     return Str;
435 #include "llvm/Frontend/OpenMP/OMPKinds.def"
436   }
437   llvm_unreachable("Unknown trait selector!");
438 }
439 
440 TraitProperty llvm::omp::getOpenMPContextTraitPropertyKind(
441     TraitSet Set, TraitSelector Selector, StringRef S) {
442   // Special handling for `device={isa(...)}` as we accept anything here. It is
443   // up to the target to decide if the feature is available.
444   if (Set == TraitSet::device && Selector == TraitSelector::device_isa)
445     return TraitProperty::device_isa___ANY;
446 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
447   if (Set == TraitSet::TraitSetEnum && Str == S)                               \
448     return TraitProperty::Enum;
449 #include "llvm/Frontend/OpenMP/OMPKinds.def"
450   return TraitProperty::invalid;
451 }
452 TraitProperty
453 llvm::omp::getOpenMPContextTraitPropertyForSelector(TraitSelector Selector) {
454   return StringSwitch<TraitProperty>(
455              getOpenMPContextTraitSelectorName(Selector))
456 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
457   .Case(Str, Selector == TraitSelector::TraitSelectorEnum                      \
458                  ? TraitProperty::Enum                                         \
459                  : TraitProperty::invalid)
460 #include "llvm/Frontend/OpenMP/OMPKinds.def"
461       .Default(TraitProperty::invalid);
462 }
463 StringRef llvm::omp::getOpenMPContextTraitPropertyName(TraitProperty Kind,
464                                                        StringRef RawString) {
465   if (Kind == TraitProperty::device_isa___ANY)
466     return RawString;
467   switch (Kind) {
468 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
469   case TraitProperty::Enum:                                                    \
470     return Str;
471 #include "llvm/Frontend/OpenMP/OMPKinds.def"
472   }
473   llvm_unreachable("Unknown trait property!");
474 }
475 StringRef llvm::omp::getOpenMPContextTraitPropertyFullName(TraitProperty Kind) {
476   switch (Kind) {
477 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
478   case TraitProperty::Enum:                                                    \
479     return "(" #TraitSetEnum "," #TraitSelectorEnum "," Str ")";
480 #include "llvm/Frontend/OpenMP/OMPKinds.def"
481   }
482   llvm_unreachable("Unknown trait property!");
483 }
484 
485 bool llvm::omp::isValidTraitSelectorForTraitSet(TraitSelector Selector,
486                                                 TraitSet Set,
487                                                 bool &AllowsTraitScore,
488                                                 bool &RequiresProperty) {
489   AllowsTraitScore = Set != TraitSet::construct && Set != TraitSet::device;
490   switch (Selector) {
491 #define OMP_TRAIT_SELECTOR(Enum, TraitSetEnum, Str, ReqProp)                   \
492   case TraitSelector::Enum:                                                    \
493     RequiresProperty = ReqProp;                                                \
494     return Set == TraitSet::TraitSetEnum;
495 #include "llvm/Frontend/OpenMP/OMPKinds.def"
496   }
497   llvm_unreachable("Unknown trait selector!");
498 }
499 
500 bool llvm::omp::isValidTraitPropertyForTraitSetAndSelector(
501     TraitProperty Property, TraitSelector Selector, TraitSet Set) {
502   switch (Property) {
503 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
504   case TraitProperty::Enum:                                                    \
505     return Set == TraitSet::TraitSetEnum &&                                    \
506            Selector == TraitSelector::TraitSelectorEnum;
507 #include "llvm/Frontend/OpenMP/OMPKinds.def"
508   }
509   llvm_unreachable("Unknown trait property!");
510 }
511 
512 std::string llvm::omp::listOpenMPContextTraitSets() {
513   std::string S;
514 #define OMP_TRAIT_SET(Enum, Str)                                               \
515   if (StringRef(Str) != "invalid")                                             \
516     S.append("'").append(Str).append("'").append(" ");
517 #include "llvm/Frontend/OpenMP/OMPKinds.def"
518   S.pop_back();
519   return S;
520 }
521 
522 std::string llvm::omp::listOpenMPContextTraitSelectors(TraitSet Set) {
523   std::string S;
524 #define OMP_TRAIT_SELECTOR(Enum, TraitSetEnum, Str, ReqProp)                   \
525   if (TraitSet::TraitSetEnum == Set && StringRef(Str) != "Invalid")            \
526     S.append("'").append(Str).append("'").append(" ");
527 #include "llvm/Frontend/OpenMP/OMPKinds.def"
528   S.pop_back();
529   return S;
530 }
531 
532 std::string
533 llvm::omp::listOpenMPContextTraitProperties(TraitSet Set,
534                                             TraitSelector Selector) {
535   std::string S;
536 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
537   if (TraitSet::TraitSetEnum == Set &&                                         \
538       TraitSelector::TraitSelectorEnum == Selector &&                          \
539       StringRef(Str) != "invalid")                                             \
540     S.append("'").append(Str).append("'").append(" ");
541 #include "llvm/Frontend/OpenMP/OMPKinds.def"
542   if (S.empty())
543     return "<none>";
544   S.pop_back();
545   return S;
546 }
547