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/ADT/Triple.h"
19#include "llvm/Support/Debug.h"
20#include "llvm/Support/raw_ostream.h"
21
22#define DEBUG_TYPE "openmp-ir-builder"
23
24using namespace llvm;
25using namespace omp;
26
27OMPContext::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.
98template <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.
122template <typename T>
123static 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
129static 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
144static 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
262bool llvm::omp::isVariantApplicableInContext(const VariantMatchInfo &VMI,
263                                             const OMPContext &Ctx,
264                                             bool DeviceSetOnly) {
265  return isVariantApplicableInContextHelper(
266      VMI, Ctx, /* ConstructMatches */ nullptr, DeviceSetOnly);
267}
268
269static 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
339int 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
376TraitSet 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
383TraitSet
384llvm::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}
393TraitSet
394llvm::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}
403StringRef 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
413TraitSelector 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}
420TraitSelector
421llvm::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}
430StringRef 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
440TraitProperty 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}
452TraitProperty
453llvm::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}
463StringRef 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}
475StringRef 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
485bool 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
500bool 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
512std::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
522std::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
532std::string
533llvm::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