1 //===-- IPO/OpenMPOpt.cpp - Collection of OpenMP specific optimizations ---===//
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 //
9 // OpenMP specific optimizations:
10 //
11 // - Deduplication of runtime calls, e.g., omp_get_thread_num.
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #include "llvm/Transforms/IPO/OpenMPOpt.h"
16 
17 #include "llvm/ADT/EnumeratedArray.h"
18 #include "llvm/ADT/Statistic.h"
19 #include "llvm/Analysis/CallGraph.h"
20 #include "llvm/Analysis/CallGraphSCCPass.h"
21 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
22 #include "llvm/Frontend/OpenMP/OMPConstants.h"
23 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
24 #include "llvm/InitializePasses.h"
25 #include "llvm/Support/CommandLine.h"
26 #include "llvm/Transforms/IPO.h"
27 #include "llvm/Transforms/IPO/Attributor.h"
28 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
29 #include "llvm/Analysis/ValueTracking.h"
30 
31 using namespace llvm;
32 using namespace omp;
33 
34 #define DEBUG_TYPE "openmp-opt"
35 
36 static cl::opt<bool> DisableOpenMPOptimizations(
37     "openmp-opt-disable", cl::ZeroOrMore,
38     cl::desc("Disable OpenMP specific optimizations."), cl::Hidden,
39     cl::init(false));
40 
41 static cl::opt<bool> PrintICVValues("openmp-print-icv-values", cl::init(false),
42                                     cl::Hidden);
43 static cl::opt<bool> PrintOpenMPKernels("openmp-print-gpu-kernels",
44                                         cl::init(false), cl::Hidden);
45 
46 static cl::opt<bool> HideMemoryTransferLatency(
47     "openmp-hide-memory-transfer-latency",
48     cl::desc("[WIP] Tries to hide the latency of host to device memory"
49              " transfers"),
50     cl::Hidden, cl::init(false));
51 
52 
53 STATISTIC(NumOpenMPRuntimeCallsDeduplicated,
54           "Number of OpenMP runtime calls deduplicated");
55 STATISTIC(NumOpenMPParallelRegionsDeleted,
56           "Number of OpenMP parallel regions deleted");
57 STATISTIC(NumOpenMPRuntimeFunctionsIdentified,
58           "Number of OpenMP runtime functions identified");
59 STATISTIC(NumOpenMPRuntimeFunctionUsesIdentified,
60           "Number of OpenMP runtime function uses identified");
61 STATISTIC(NumOpenMPTargetRegionKernels,
62           "Number of OpenMP target region entry points (=kernels) identified");
63 STATISTIC(
64     NumOpenMPParallelRegionsReplacedInGPUStateMachine,
65     "Number of OpenMP parallel regions replaced with ID in GPU state machines");
66 
67 #if !defined(NDEBUG)
68 static constexpr auto TAG = "[" DEBUG_TYPE "]";
69 #endif
70 
71 namespace {
72 
73 struct AAICVTracker;
74 
75 /// OpenMP specific information. For now, stores RFIs and ICVs also needed for
76 /// Attributor runs.
77 struct OMPInformationCache : public InformationCache {
78   OMPInformationCache(Module &M, AnalysisGetter &AG,
79                       BumpPtrAllocator &Allocator, SetVector<Function *> &CGSCC,
80                       SmallPtrSetImpl<Kernel> &Kernels)
81       : InformationCache(M, AG, Allocator, &CGSCC), OMPBuilder(M),
82         Kernels(Kernels) {
83 
84     OMPBuilder.initialize();
85     initializeRuntimeFunctions();
86     initializeInternalControlVars();
87   }
88 
89   /// Generic information that describes an internal control variable.
90   struct InternalControlVarInfo {
91     /// The kind, as described by InternalControlVar enum.
92     InternalControlVar Kind;
93 
94     /// The name of the ICV.
95     StringRef Name;
96 
97     /// Environment variable associated with this ICV.
98     StringRef EnvVarName;
99 
100     /// Initial value kind.
101     ICVInitValue InitKind;
102 
103     /// Initial value.
104     ConstantInt *InitValue;
105 
106     /// Setter RTL function associated with this ICV.
107     RuntimeFunction Setter;
108 
109     /// Getter RTL function associated with this ICV.
110     RuntimeFunction Getter;
111 
112     /// RTL Function corresponding to the override clause of this ICV
113     RuntimeFunction Clause;
114   };
115 
116   /// Generic information that describes a runtime function
117   struct RuntimeFunctionInfo {
118 
119     /// The kind, as described by the RuntimeFunction enum.
120     RuntimeFunction Kind;
121 
122     /// The name of the function.
123     StringRef Name;
124 
125     /// Flag to indicate a variadic function.
126     bool IsVarArg;
127 
128     /// The return type of the function.
129     Type *ReturnType;
130 
131     /// The argument types of the function.
132     SmallVector<Type *, 8> ArgumentTypes;
133 
134     /// The declaration if available.
135     Function *Declaration = nullptr;
136 
137     /// Uses of this runtime function per function containing the use.
138     using UseVector = SmallVector<Use *, 16>;
139 
140     /// Clear UsesMap for runtime function.
141     void clearUsesMap() { UsesMap.clear(); }
142 
143     /// Boolean conversion that is true if the runtime function was found.
144     operator bool() const { return Declaration; }
145 
146     /// Return the vector of uses in function \p F.
147     UseVector &getOrCreateUseVector(Function *F) {
148       std::shared_ptr<UseVector> &UV = UsesMap[F];
149       if (!UV)
150         UV = std::make_shared<UseVector>();
151       return *UV;
152     }
153 
154     /// Return the vector of uses in function \p F or `nullptr` if there are
155     /// none.
156     const UseVector *getUseVector(Function &F) const {
157       auto I = UsesMap.find(&F);
158       if (I != UsesMap.end())
159         return I->second.get();
160       return nullptr;
161     }
162 
163     /// Return how many functions contain uses of this runtime function.
164     size_t getNumFunctionsWithUses() const { return UsesMap.size(); }
165 
166     /// Return the number of arguments (or the minimal number for variadic
167     /// functions).
168     size_t getNumArgs() const { return ArgumentTypes.size(); }
169 
170     /// Run the callback \p CB on each use and forget the use if the result is
171     /// true. The callback will be fed the function in which the use was
172     /// encountered as second argument.
173     void foreachUse(SmallVectorImpl<Function *> &SCC,
174                     function_ref<bool(Use &, Function &)> CB) {
175       for (Function *F : SCC)
176         foreachUse(CB, F);
177     }
178 
179     /// Run the callback \p CB on each use within the function \p F and forget
180     /// the use if the result is true.
181     void foreachUse(function_ref<bool(Use &, Function &)> CB, Function *F) {
182       SmallVector<unsigned, 8> ToBeDeleted;
183       ToBeDeleted.clear();
184 
185       unsigned Idx = 0;
186       UseVector &UV = getOrCreateUseVector(F);
187 
188       for (Use *U : UV) {
189         if (CB(*U, *F))
190           ToBeDeleted.push_back(Idx);
191         ++Idx;
192       }
193 
194       // Remove the to-be-deleted indices in reverse order as prior
195       // modifications will not modify the smaller indices.
196       while (!ToBeDeleted.empty()) {
197         unsigned Idx = ToBeDeleted.pop_back_val();
198         UV[Idx] = UV.back();
199         UV.pop_back();
200       }
201     }
202 
203   private:
204     /// Map from functions to all uses of this runtime function contained in
205     /// them.
206     DenseMap<Function *, std::shared_ptr<UseVector>> UsesMap;
207   };
208 
209   /// An OpenMP-IR-Builder instance
210   OpenMPIRBuilder OMPBuilder;
211 
212   /// Map from runtime function kind to the runtime function description.
213   EnumeratedArray<RuntimeFunctionInfo, RuntimeFunction,
214                   RuntimeFunction::OMPRTL___last>
215       RFIs;
216 
217   /// Map from ICV kind to the ICV description.
218   EnumeratedArray<InternalControlVarInfo, InternalControlVar,
219                   InternalControlVar::ICV___last>
220       ICVs;
221 
222   /// Helper to initialize all internal control variable information for those
223   /// defined in OMPKinds.def.
224   void initializeInternalControlVars() {
225 #define ICV_RT_SET(_Name, RTL)                                                 \
226   {                                                                            \
227     auto &ICV = ICVs[_Name];                                                   \
228     ICV.Setter = RTL;                                                          \
229   }
230 #define ICV_RT_GET(Name, RTL)                                                  \
231   {                                                                            \
232     auto &ICV = ICVs[Name];                                                    \
233     ICV.Getter = RTL;                                                          \
234   }
235 #define ICV_DATA_ENV(Enum, _Name, _EnvVarName, Init)                           \
236   {                                                                            \
237     auto &ICV = ICVs[Enum];                                                    \
238     ICV.Name = _Name;                                                          \
239     ICV.Kind = Enum;                                                           \
240     ICV.InitKind = Init;                                                       \
241     ICV.EnvVarName = _EnvVarName;                                              \
242     switch (ICV.InitKind) {                                                    \
243     case ICV_IMPLEMENTATION_DEFINED:                                           \
244       ICV.InitValue = nullptr;                                                 \
245       break;                                                                   \
246     case ICV_ZERO:                                                             \
247       ICV.InitValue = ConstantInt::get(                                        \
248           Type::getInt32Ty(OMPBuilder.Int32->getContext()), 0);                \
249       break;                                                                   \
250     case ICV_FALSE:                                                            \
251       ICV.InitValue = ConstantInt::getFalse(OMPBuilder.Int1->getContext());    \
252       break;                                                                   \
253     case ICV_LAST:                                                             \
254       break;                                                                   \
255     }                                                                          \
256   }
257 #include "llvm/Frontend/OpenMP/OMPKinds.def"
258   }
259 
260   /// Returns true if the function declaration \p F matches the runtime
261   /// function types, that is, return type \p RTFRetType, and argument types
262   /// \p RTFArgTypes.
263   static bool declMatchesRTFTypes(Function *F, Type *RTFRetType,
264                                   SmallVector<Type *, 8> &RTFArgTypes) {
265     // TODO: We should output information to the user (under debug output
266     //       and via remarks).
267 
268     if (!F)
269       return false;
270     if (F->getReturnType() != RTFRetType)
271       return false;
272     if (F->arg_size() != RTFArgTypes.size())
273       return false;
274 
275     auto RTFTyIt = RTFArgTypes.begin();
276     for (Argument &Arg : F->args()) {
277       if (Arg.getType() != *RTFTyIt)
278         return false;
279 
280       ++RTFTyIt;
281     }
282 
283     return true;
284   }
285 
286   // Helper to collect all uses of the declaration in the UsesMap.
287   unsigned collectUses(RuntimeFunctionInfo &RFI, bool CollectStats = true) {
288     unsigned NumUses = 0;
289     if (!RFI.Declaration)
290       return NumUses;
291     OMPBuilder.addAttributes(RFI.Kind, *RFI.Declaration);
292 
293     if (CollectStats) {
294       NumOpenMPRuntimeFunctionsIdentified += 1;
295       NumOpenMPRuntimeFunctionUsesIdentified += RFI.Declaration->getNumUses();
296     }
297 
298     // TODO: We directly convert uses into proper calls and unknown uses.
299     for (Use &U : RFI.Declaration->uses()) {
300       if (Instruction *UserI = dyn_cast<Instruction>(U.getUser())) {
301         if (ModuleSlice.count(UserI->getFunction())) {
302           RFI.getOrCreateUseVector(UserI->getFunction()).push_back(&U);
303           ++NumUses;
304         }
305       } else {
306         RFI.getOrCreateUseVector(nullptr).push_back(&U);
307         ++NumUses;
308       }
309     }
310     return NumUses;
311   }
312 
313   // Helper function to recollect uses of all runtime functions.
314   void recollectUses() {
315     for (int Idx = 0; Idx < RFIs.size(); ++Idx) {
316       auto &RFI = RFIs[static_cast<RuntimeFunction>(Idx)];
317       RFI.clearUsesMap();
318       collectUses(RFI, /*CollectStats*/ false);
319     }
320   }
321 
322   /// Helper to initialize all runtime function information for those defined
323   /// in OpenMPKinds.def.
324   void initializeRuntimeFunctions() {
325     Module &M = *((*ModuleSlice.begin())->getParent());
326 
327     // Helper macros for handling __VA_ARGS__ in OMP_RTL
328 #define OMP_TYPE(VarName, ...)                                                 \
329   Type *VarName = OMPBuilder.VarName;                                          \
330   (void)VarName;
331 
332 #define OMP_ARRAY_TYPE(VarName, ...)                                           \
333   ArrayType *VarName##Ty = OMPBuilder.VarName##Ty;                             \
334   (void)VarName##Ty;                                                           \
335   PointerType *VarName##PtrTy = OMPBuilder.VarName##PtrTy;                     \
336   (void)VarName##PtrTy;
337 
338 #define OMP_FUNCTION_TYPE(VarName, ...)                                        \
339   FunctionType *VarName = OMPBuilder.VarName;                                  \
340   (void)VarName;                                                               \
341   PointerType *VarName##Ptr = OMPBuilder.VarName##Ptr;                         \
342   (void)VarName##Ptr;
343 
344 #define OMP_STRUCT_TYPE(VarName, ...)                                          \
345   StructType *VarName = OMPBuilder.VarName;                                    \
346   (void)VarName;                                                               \
347   PointerType *VarName##Ptr = OMPBuilder.VarName##Ptr;                         \
348   (void)VarName##Ptr;
349 
350 #define OMP_RTL(_Enum, _Name, _IsVarArg, _ReturnType, ...)                     \
351   {                                                                            \
352     SmallVector<Type *, 8> ArgsTypes({__VA_ARGS__});                           \
353     Function *F = M.getFunction(_Name);                                        \
354     if (declMatchesRTFTypes(F, OMPBuilder._ReturnType, ArgsTypes)) {           \
355       auto &RFI = RFIs[_Enum];                                                 \
356       RFI.Kind = _Enum;                                                        \
357       RFI.Name = _Name;                                                        \
358       RFI.IsVarArg = _IsVarArg;                                                \
359       RFI.ReturnType = OMPBuilder._ReturnType;                                 \
360       RFI.ArgumentTypes = std::move(ArgsTypes);                                \
361       RFI.Declaration = F;                                                     \
362       unsigned NumUses = collectUses(RFI);                                     \
363       (void)NumUses;                                                           \
364       LLVM_DEBUG({                                                             \
365         dbgs() << TAG << RFI.Name << (RFI.Declaration ? "" : " not")           \
366                << " found\n";                                                  \
367         if (RFI.Declaration)                                                   \
368           dbgs() << TAG << "-> got " << NumUses << " uses in "                 \
369                  << RFI.getNumFunctionsWithUses()                              \
370                  << " different functions.\n";                                 \
371       });                                                                      \
372     }                                                                          \
373   }
374 #include "llvm/Frontend/OpenMP/OMPKinds.def"
375 
376     // TODO: We should attach the attributes defined in OMPKinds.def.
377   }
378 
379   /// Collection of known kernels (\see Kernel) in the module.
380   SmallPtrSetImpl<Kernel> &Kernels;
381 };
382 
383 /// Used to map the values physically (in the IR) stored in an offload
384 /// array, to a vector in memory.
385 struct OffloadArray {
386   /// Physical array (in the IR).
387   AllocaInst *Array = nullptr;
388   /// Mapped values.
389   SmallVector<Value *, 8> StoredValues;
390   /// Last stores made in the offload array.
391   SmallVector<StoreInst *, 8> LastAccesses;
392 
393   OffloadArray() = default;
394 
395   /// Initializes the OffloadArray with the values stored in \p Array before
396   /// instruction \p Before is reached. Returns false if the initialization
397   /// fails.
398   /// This MUST be used immediately after the construction of the object.
399   bool initialize(AllocaInst &Array, Instruction &Before) {
400     if (!Array.getAllocatedType()->isArrayTy())
401       return false;
402 
403     if (!getValues(Array, Before))
404       return false;
405 
406     this->Array = &Array;
407     return true;
408   }
409 
410   static const unsigned BasePtrsArgNum = 2;
411   static const unsigned PtrsArgNum = 3;
412   static const unsigned SizesArgNum = 4;
413 
414 private:
415   /// Traverses the BasicBlock where \p Array is, collecting the stores made to
416   /// \p Array, leaving StoredValues with the values stored before the
417   /// instruction \p Before is reached.
418   bool getValues(AllocaInst &Array, Instruction &Before) {
419     // Initialize container.
420     const uint64_t NumValues =
421         Array.getAllocatedType()->getArrayNumElements();
422     StoredValues.assign(NumValues, nullptr);
423     LastAccesses.assign(NumValues, nullptr);
424 
425     // TODO: This assumes the instruction \p Before is in the same
426     //  BasicBlock as Array. Make it general, for any control flow graph.
427     BasicBlock *BB = Array.getParent();
428     if (BB != Before.getParent())
429       return false;
430 
431     const DataLayout &DL = Array.getModule()->getDataLayout();
432     const unsigned int PointerSize = DL.getPointerSize();
433 
434     for (Instruction &I : *BB) {
435       if (&I == &Before)
436         break;
437 
438       if (!isa<StoreInst>(&I))
439         continue;
440 
441       auto *S = cast<StoreInst>(&I);
442       int64_t Offset = -1;
443       auto *Dst = GetPointerBaseWithConstantOffset(S->getPointerOperand(),
444                                                    Offset, DL);
445       if (Dst == &Array) {
446         int64_t Idx = Offset / PointerSize;
447         StoredValues[Idx] = getUnderlyingObject(S->getValueOperand());
448         LastAccesses[Idx] = S;
449       }
450     }
451 
452     return isFilled();
453   }
454 
455   /// Returns true if all values in StoredValues and
456   /// LastAccesses are not nullptrs.
457   bool isFilled() {
458     const unsigned NumValues = StoredValues.size();
459     for (unsigned I = 0; I < NumValues; ++I) {
460       if (!StoredValues[I] || !LastAccesses[I])
461         return false;
462     }
463 
464     return true;
465   }
466 };
467 
468 struct OpenMPOpt {
469 
470   using OptimizationRemarkGetter =
471       function_ref<OptimizationRemarkEmitter &(Function *)>;
472 
473   OpenMPOpt(SmallVectorImpl<Function *> &SCC, CallGraphUpdater &CGUpdater,
474             OptimizationRemarkGetter OREGetter,
475             OMPInformationCache &OMPInfoCache, Attributor &A)
476       : M(*(*SCC.begin())->getParent()), SCC(SCC), CGUpdater(CGUpdater),
477         OREGetter(OREGetter), OMPInfoCache(OMPInfoCache), A(A) {}
478 
479   /// Check if any remarks are enabled for openmp-opt
480   bool remarksEnabled() {
481     auto &Ctx = M.getContext();
482     return Ctx.getDiagHandlerPtr()->isAnyRemarkEnabled(DEBUG_TYPE);
483   }
484 
485   /// Run all OpenMP optimizations on the underlying SCC/ModuleSlice.
486   bool run() {
487     if (SCC.empty())
488       return false;
489 
490     bool Changed = false;
491 
492     LLVM_DEBUG(dbgs() << TAG << "Run on SCC with " << SCC.size()
493                       << " functions in a slice with "
494                       << OMPInfoCache.ModuleSlice.size() << " functions\n");
495 
496     if (PrintICVValues)
497       printICVs();
498     if (PrintOpenMPKernels)
499       printKernels();
500 
501     Changed |= rewriteDeviceCodeStateMachine();
502 
503     Changed |= runAttributor();
504 
505     // Recollect uses, in case Attributor deleted any.
506     OMPInfoCache.recollectUses();
507 
508     Changed |= deduplicateRuntimeCalls();
509     Changed |= deleteParallelRegions();
510     if (HideMemoryTransferLatency)
511       Changed |= hideMemTransfersLatency();
512     if (remarksEnabled())
513       analysisGlobalization();
514 
515     return Changed;
516   }
517 
518   /// Print initial ICV values for testing.
519   /// FIXME: This should be done from the Attributor once it is added.
520   void printICVs() const {
521     InternalControlVar ICVs[] = {ICV_nthreads, ICV_active_levels, ICV_cancel};
522 
523     for (Function *F : OMPInfoCache.ModuleSlice) {
524       for (auto ICV : ICVs) {
525         auto ICVInfo = OMPInfoCache.ICVs[ICV];
526         auto Remark = [&](OptimizationRemark OR) {
527           return OR << "OpenMP ICV " << ore::NV("OpenMPICV", ICVInfo.Name)
528                     << " Value: "
529                     << (ICVInfo.InitValue
530                             ? ICVInfo.InitValue->getValue().toString(10, true)
531                             : "IMPLEMENTATION_DEFINED");
532         };
533 
534         emitRemarkOnFunction(F, "OpenMPICVTracker", Remark);
535       }
536     }
537   }
538 
539   /// Print OpenMP GPU kernels for testing.
540   void printKernels() const {
541     for (Function *F : SCC) {
542       if (!OMPInfoCache.Kernels.count(F))
543         continue;
544 
545       auto Remark = [&](OptimizationRemark OR) {
546         return OR << "OpenMP GPU kernel "
547                   << ore::NV("OpenMPGPUKernel", F->getName()) << "\n";
548       };
549 
550       emitRemarkOnFunction(F, "OpenMPGPU", Remark);
551     }
552   }
553 
554   /// Return the call if \p U is a callee use in a regular call. If \p RFI is
555   /// given it has to be the callee or a nullptr is returned.
556   static CallInst *getCallIfRegularCall(
557       Use &U, OMPInformationCache::RuntimeFunctionInfo *RFI = nullptr) {
558     CallInst *CI = dyn_cast<CallInst>(U.getUser());
559     if (CI && CI->isCallee(&U) && !CI->hasOperandBundles() &&
560         (!RFI || CI->getCalledFunction() == RFI->Declaration))
561       return CI;
562     return nullptr;
563   }
564 
565   /// Return the call if \p V is a regular call. If \p RFI is given it has to be
566   /// the callee or a nullptr is returned.
567   static CallInst *getCallIfRegularCall(
568       Value &V, OMPInformationCache::RuntimeFunctionInfo *RFI = nullptr) {
569     CallInst *CI = dyn_cast<CallInst>(&V);
570     if (CI && !CI->hasOperandBundles() &&
571         (!RFI || CI->getCalledFunction() == RFI->Declaration))
572       return CI;
573     return nullptr;
574   }
575 
576 private:
577   /// Try to delete parallel regions if possible.
578   bool deleteParallelRegions() {
579     const unsigned CallbackCalleeOperand = 2;
580 
581     OMPInformationCache::RuntimeFunctionInfo &RFI =
582         OMPInfoCache.RFIs[OMPRTL___kmpc_fork_call];
583 
584     if (!RFI.Declaration)
585       return false;
586 
587     bool Changed = false;
588     auto DeleteCallCB = [&](Use &U, Function &) {
589       CallInst *CI = getCallIfRegularCall(U);
590       if (!CI)
591         return false;
592       auto *Fn = dyn_cast<Function>(
593           CI->getArgOperand(CallbackCalleeOperand)->stripPointerCasts());
594       if (!Fn)
595         return false;
596       if (!Fn->onlyReadsMemory())
597         return false;
598       if (!Fn->hasFnAttribute(Attribute::WillReturn))
599         return false;
600 
601       LLVM_DEBUG(dbgs() << TAG << "Delete read-only parallel region in "
602                         << CI->getCaller()->getName() << "\n");
603 
604       auto Remark = [&](OptimizationRemark OR) {
605         return OR << "Parallel region in "
606                   << ore::NV("OpenMPParallelDelete", CI->getCaller()->getName())
607                   << " deleted";
608       };
609       emitRemark<OptimizationRemark>(CI, "OpenMPParallelRegionDeletion",
610                                      Remark);
611 
612       CGUpdater.removeCallSite(*CI);
613       CI->eraseFromParent();
614       Changed = true;
615       ++NumOpenMPParallelRegionsDeleted;
616       return true;
617     };
618 
619     RFI.foreachUse(SCC, DeleteCallCB);
620 
621     return Changed;
622   }
623 
624   /// Try to eliminate runtime calls by reusing existing ones.
625   bool deduplicateRuntimeCalls() {
626     bool Changed = false;
627 
628     RuntimeFunction DeduplicableRuntimeCallIDs[] = {
629         OMPRTL_omp_get_num_threads,
630         OMPRTL_omp_in_parallel,
631         OMPRTL_omp_get_cancellation,
632         OMPRTL_omp_get_thread_limit,
633         OMPRTL_omp_get_supported_active_levels,
634         OMPRTL_omp_get_level,
635         OMPRTL_omp_get_ancestor_thread_num,
636         OMPRTL_omp_get_team_size,
637         OMPRTL_omp_get_active_level,
638         OMPRTL_omp_in_final,
639         OMPRTL_omp_get_proc_bind,
640         OMPRTL_omp_get_num_places,
641         OMPRTL_omp_get_num_procs,
642         OMPRTL_omp_get_place_num,
643         OMPRTL_omp_get_partition_num_places,
644         OMPRTL_omp_get_partition_place_nums};
645 
646     // Global-tid is handled separately.
647     SmallSetVector<Value *, 16> GTIdArgs;
648     collectGlobalThreadIdArguments(GTIdArgs);
649     LLVM_DEBUG(dbgs() << TAG << "Found " << GTIdArgs.size()
650                       << " global thread ID arguments\n");
651 
652     for (Function *F : SCC) {
653       for (auto DeduplicableRuntimeCallID : DeduplicableRuntimeCallIDs)
654         Changed |= deduplicateRuntimeCalls(
655             *F, OMPInfoCache.RFIs[DeduplicableRuntimeCallID]);
656 
657       // __kmpc_global_thread_num is special as we can replace it with an
658       // argument in enough cases to make it worth trying.
659       Value *GTIdArg = nullptr;
660       for (Argument &Arg : F->args())
661         if (GTIdArgs.count(&Arg)) {
662           GTIdArg = &Arg;
663           break;
664         }
665       Changed |= deduplicateRuntimeCalls(
666           *F, OMPInfoCache.RFIs[OMPRTL___kmpc_global_thread_num], GTIdArg);
667     }
668 
669     return Changed;
670   }
671 
672   /// Tries to hide the latency of runtime calls that involve host to
673   /// device memory transfers by splitting them into their "issue" and "wait"
674   /// versions. The "issue" is moved upwards as much as possible. The "wait" is
675   /// moved downards as much as possible. The "issue" issues the memory transfer
676   /// asynchronously, returning a handle. The "wait" waits in the returned
677   /// handle for the memory transfer to finish.
678   bool hideMemTransfersLatency() {
679     auto &RFI = OMPInfoCache.RFIs[OMPRTL___tgt_target_data_begin_mapper];
680     bool Changed = false;
681     auto SplitMemTransfers = [&](Use &U, Function &Decl) {
682       auto *RTCall = getCallIfRegularCall(U, &RFI);
683       if (!RTCall)
684         return false;
685 
686       OffloadArray OffloadArrays[3];
687       if (!getValuesInOffloadArrays(*RTCall, OffloadArrays))
688         return false;
689 
690       LLVM_DEBUG(dumpValuesInOffloadArrays(OffloadArrays));
691 
692       // TODO: Check if can be moved upwards.
693       bool WasSplit = false;
694       Instruction *WaitMovementPoint = canBeMovedDownwards(*RTCall);
695       if (WaitMovementPoint)
696         WasSplit = splitTargetDataBeginRTC(*RTCall, *WaitMovementPoint);
697 
698       Changed |= WasSplit;
699       return WasSplit;
700     };
701     RFI.foreachUse(SCC, SplitMemTransfers);
702 
703     return Changed;
704   }
705 
706   void analysisGlobalization() {
707     auto &RFI =
708         OMPInfoCache.RFIs[OMPRTL___kmpc_data_sharing_coalesced_push_stack];
709 
710     auto checkGlobalization = [&](Use &U, Function &Decl) {
711       if (CallInst *CI = getCallIfRegularCall(U, &RFI)) {
712         auto Remark = [&](OptimizationRemarkAnalysis ORA) {
713           return ORA
714                  << "Found thread data sharing on the GPU. "
715                  << "Expect degraded performance due to data globalization.";
716         };
717         emitRemark<OptimizationRemarkAnalysis>(CI, "OpenMPGlobalization",
718                                                Remark);
719       }
720 
721       return false;
722     };
723 
724     RFI.foreachUse(SCC, checkGlobalization);
725     return;
726   }
727 
728   /// Maps the values stored in the offload arrays passed as arguments to
729   /// \p RuntimeCall into the offload arrays in \p OAs.
730   bool getValuesInOffloadArrays(CallInst &RuntimeCall,
731                                 MutableArrayRef<OffloadArray> OAs) {
732     assert(OAs.size() == 3 && "Need space for three offload arrays!");
733 
734     // A runtime call that involves memory offloading looks something like:
735     // call void @__tgt_target_data_begin_mapper(arg0, arg1,
736     //   i8** %offload_baseptrs, i8** %offload_ptrs, i64* %offload_sizes,
737     // ...)
738     // So, the idea is to access the allocas that allocate space for these
739     // offload arrays, offload_baseptrs, offload_ptrs, offload_sizes.
740     // Therefore:
741     // i8** %offload_baseptrs.
742     Value *BasePtrsArg =
743         RuntimeCall.getArgOperand(OffloadArray::BasePtrsArgNum);
744     // i8** %offload_ptrs.
745     Value *PtrsArg = RuntimeCall.getArgOperand(OffloadArray::PtrsArgNum);
746     // i8** %offload_sizes.
747     Value *SizesArg = RuntimeCall.getArgOperand(OffloadArray::SizesArgNum);
748 
749     // Get values stored in **offload_baseptrs.
750     auto *V = getUnderlyingObject(BasePtrsArg);
751     if (!isa<AllocaInst>(V))
752       return false;
753     auto *BasePtrsArray = cast<AllocaInst>(V);
754     if (!OAs[0].initialize(*BasePtrsArray, RuntimeCall))
755       return false;
756 
757     // Get values stored in **offload_baseptrs.
758     V = getUnderlyingObject(PtrsArg);
759     if (!isa<AllocaInst>(V))
760       return false;
761     auto *PtrsArray = cast<AllocaInst>(V);
762     if (!OAs[1].initialize(*PtrsArray, RuntimeCall))
763       return false;
764 
765     // Get values stored in **offload_sizes.
766     V = getUnderlyingObject(SizesArg);
767     // If it's a [constant] global array don't analyze it.
768     if (isa<GlobalValue>(V))
769       return isa<Constant>(V);
770     if (!isa<AllocaInst>(V))
771       return false;
772 
773     auto *SizesArray = cast<AllocaInst>(V);
774     if (!OAs[2].initialize(*SizesArray, RuntimeCall))
775       return false;
776 
777     return true;
778   }
779 
780   /// Prints the values in the OffloadArrays \p OAs using LLVM_DEBUG.
781   /// For now this is a way to test that the function getValuesInOffloadArrays
782   /// is working properly.
783   /// TODO: Move this to a unittest when unittests are available for OpenMPOpt.
784   void dumpValuesInOffloadArrays(ArrayRef<OffloadArray> OAs) {
785     assert(OAs.size() == 3 && "There are three offload arrays to debug!");
786 
787     LLVM_DEBUG(dbgs() << TAG << " Successfully got offload values:\n");
788     std::string ValuesStr;
789     raw_string_ostream Printer(ValuesStr);
790     std::string Separator = " --- ";
791 
792     for (auto *BP : OAs[0].StoredValues) {
793       BP->print(Printer);
794       Printer << Separator;
795     }
796     LLVM_DEBUG(dbgs() << "\t\toffload_baseptrs: " << Printer.str() << "\n");
797     ValuesStr.clear();
798 
799     for (auto *P : OAs[1].StoredValues) {
800       P->print(Printer);
801       Printer << Separator;
802     }
803     LLVM_DEBUG(dbgs() << "\t\toffload_ptrs: " << Printer.str() << "\n");
804     ValuesStr.clear();
805 
806     for (auto *S : OAs[2].StoredValues) {
807       S->print(Printer);
808       Printer << Separator;
809     }
810     LLVM_DEBUG(dbgs() << "\t\toffload_sizes: " << Printer.str() << "\n");
811   }
812 
813   /// Returns the instruction where the "wait" counterpart \p RuntimeCall can be
814   /// moved. Returns nullptr if the movement is not possible, or not worth it.
815   Instruction *canBeMovedDownwards(CallInst &RuntimeCall) {
816     // FIXME: This traverses only the BasicBlock where RuntimeCall is.
817     //  Make it traverse the CFG.
818 
819     Instruction *CurrentI = &RuntimeCall;
820     bool IsWorthIt = false;
821     while ((CurrentI = CurrentI->getNextNode())) {
822 
823       // TODO: Once we detect the regions to be offloaded we should use the
824       //  alias analysis manager to check if CurrentI may modify one of
825       //  the offloaded regions.
826       if (CurrentI->mayHaveSideEffects() || CurrentI->mayReadFromMemory()) {
827         if (IsWorthIt)
828           return CurrentI;
829 
830         return nullptr;
831       }
832 
833       // FIXME: For now if we move it over anything without side effect
834       //  is worth it.
835       IsWorthIt = true;
836     }
837 
838     // Return end of BasicBlock.
839     return RuntimeCall.getParent()->getTerminator();
840   }
841 
842   /// Splits \p RuntimeCall into its "issue" and "wait" counterparts.
843   bool splitTargetDataBeginRTC(CallInst &RuntimeCall,
844                                Instruction &WaitMovementPoint) {
845     // Create stack allocated handle (__tgt_async_info) at the beginning of the
846     // function. Used for storing information of the async transfer, allowing to
847     // wait on it later.
848     auto &IRBuilder = OMPInfoCache.OMPBuilder;
849     auto *F = RuntimeCall.getCaller();
850     Instruction *FirstInst = &(F->getEntryBlock().front());
851     AllocaInst *Handle = new AllocaInst(
852         IRBuilder.AsyncInfo, F->getAddressSpace(), "handle", FirstInst);
853 
854     // Add "issue" runtime call declaration:
855     // declare %struct.tgt_async_info @__tgt_target_data_begin_issue(i64, i32,
856     //   i8**, i8**, i64*, i64*)
857     FunctionCallee IssueDecl = IRBuilder.getOrCreateRuntimeFunction(
858         M, OMPRTL___tgt_target_data_begin_mapper_issue);
859 
860     // Change RuntimeCall call site for its asynchronous version.
861     SmallVector<Value *, 8> Args;
862     for (auto &Arg : RuntimeCall.args())
863       Args.push_back(Arg.get());
864     Args.push_back(Handle);
865 
866     CallInst *IssueCallsite =
867         CallInst::Create(IssueDecl, Args, /*NameStr=*/"", &RuntimeCall);
868     RuntimeCall.eraseFromParent();
869 
870     // Add "wait" runtime call declaration:
871     // declare void @__tgt_target_data_begin_wait(i64, %struct.__tgt_async_info)
872     FunctionCallee WaitDecl = IRBuilder.getOrCreateRuntimeFunction(
873         M, OMPRTL___tgt_target_data_begin_mapper_wait);
874 
875     // Add call site to WaitDecl.
876     const unsigned DeviceIDArgNum = 0;
877     Value *WaitParams[2] = {
878         IssueCallsite->getArgOperand(DeviceIDArgNum), // device_id.
879         Handle                                        // handle to wait on.
880     };
881     CallInst::Create(WaitDecl, WaitParams, /*NameStr=*/"", &WaitMovementPoint);
882 
883     return true;
884   }
885 
886   static Value *combinedIdentStruct(Value *CurrentIdent, Value *NextIdent,
887                                     bool GlobalOnly, bool &SingleChoice) {
888     if (CurrentIdent == NextIdent)
889       return CurrentIdent;
890 
891     // TODO: Figure out how to actually combine multiple debug locations. For
892     //       now we just keep an existing one if there is a single choice.
893     if (!GlobalOnly || isa<GlobalValue>(NextIdent)) {
894       SingleChoice = !CurrentIdent;
895       return NextIdent;
896     }
897     return nullptr;
898   }
899 
900   /// Return an `struct ident_t*` value that represents the ones used in the
901   /// calls of \p RFI inside of \p F. If \p GlobalOnly is true, we will not
902   /// return a local `struct ident_t*`. For now, if we cannot find a suitable
903   /// return value we create one from scratch. We also do not yet combine
904   /// information, e.g., the source locations, see combinedIdentStruct.
905   Value *
906   getCombinedIdentFromCallUsesIn(OMPInformationCache::RuntimeFunctionInfo &RFI,
907                                  Function &F, bool GlobalOnly) {
908     bool SingleChoice = true;
909     Value *Ident = nullptr;
910     auto CombineIdentStruct = [&](Use &U, Function &Caller) {
911       CallInst *CI = getCallIfRegularCall(U, &RFI);
912       if (!CI || &F != &Caller)
913         return false;
914       Ident = combinedIdentStruct(Ident, CI->getArgOperand(0),
915                                   /* GlobalOnly */ true, SingleChoice);
916       return false;
917     };
918     RFI.foreachUse(SCC, CombineIdentStruct);
919 
920     if (!Ident || !SingleChoice) {
921       // The IRBuilder uses the insertion block to get to the module, this is
922       // unfortunate but we work around it for now.
923       if (!OMPInfoCache.OMPBuilder.getInsertionPoint().getBlock())
924         OMPInfoCache.OMPBuilder.updateToLocation(OpenMPIRBuilder::InsertPointTy(
925             &F.getEntryBlock(), F.getEntryBlock().begin()));
926       // Create a fallback location if non was found.
927       // TODO: Use the debug locations of the calls instead.
928       Constant *Loc = OMPInfoCache.OMPBuilder.getOrCreateDefaultSrcLocStr();
929       Ident = OMPInfoCache.OMPBuilder.getOrCreateIdent(Loc);
930     }
931     return Ident;
932   }
933 
934   /// Try to eliminate calls of \p RFI in \p F by reusing an existing one or
935   /// \p ReplVal if given.
936   bool deduplicateRuntimeCalls(Function &F,
937                                OMPInformationCache::RuntimeFunctionInfo &RFI,
938                                Value *ReplVal = nullptr) {
939     auto *UV = RFI.getUseVector(F);
940     if (!UV || UV->size() + (ReplVal != nullptr) < 2)
941       return false;
942 
943     LLVM_DEBUG(
944         dbgs() << TAG << "Deduplicate " << UV->size() << " uses of " << RFI.Name
945                << (ReplVal ? " with an existing value\n" : "\n") << "\n");
946 
947     assert((!ReplVal || (isa<Argument>(ReplVal) &&
948                          cast<Argument>(ReplVal)->getParent() == &F)) &&
949            "Unexpected replacement value!");
950 
951     // TODO: Use dominance to find a good position instead.
952     auto CanBeMoved = [this](CallBase &CB) {
953       unsigned NumArgs = CB.getNumArgOperands();
954       if (NumArgs == 0)
955         return true;
956       if (CB.getArgOperand(0)->getType() != OMPInfoCache.OMPBuilder.IdentPtr)
957         return false;
958       for (unsigned u = 1; u < NumArgs; ++u)
959         if (isa<Instruction>(CB.getArgOperand(u)))
960           return false;
961       return true;
962     };
963 
964     if (!ReplVal) {
965       for (Use *U : *UV)
966         if (CallInst *CI = getCallIfRegularCall(*U, &RFI)) {
967           if (!CanBeMoved(*CI))
968             continue;
969 
970           auto Remark = [&](OptimizationRemark OR) {
971             auto newLoc = &*F.getEntryBlock().getFirstInsertionPt();
972             return OR << "OpenMP runtime call "
973                       << ore::NV("OpenMPOptRuntime", RFI.Name) << " moved to "
974                       << ore::NV("OpenMPRuntimeMoves", newLoc->getDebugLoc());
975           };
976           emitRemark<OptimizationRemark>(CI, "OpenMPRuntimeCodeMotion", Remark);
977 
978           CI->moveBefore(&*F.getEntryBlock().getFirstInsertionPt());
979           ReplVal = CI;
980           break;
981         }
982       if (!ReplVal)
983         return false;
984     }
985 
986     // If we use a call as a replacement value we need to make sure the ident is
987     // valid at the new location. For now we just pick a global one, either
988     // existing and used by one of the calls, or created from scratch.
989     if (CallBase *CI = dyn_cast<CallBase>(ReplVal)) {
990       if (CI->getNumArgOperands() > 0 &&
991           CI->getArgOperand(0)->getType() == OMPInfoCache.OMPBuilder.IdentPtr) {
992         Value *Ident = getCombinedIdentFromCallUsesIn(RFI, F,
993                                                       /* GlobalOnly */ true);
994         CI->setArgOperand(0, Ident);
995       }
996     }
997 
998     bool Changed = false;
999     auto ReplaceAndDeleteCB = [&](Use &U, Function &Caller) {
1000       CallInst *CI = getCallIfRegularCall(U, &RFI);
1001       if (!CI || CI == ReplVal || &F != &Caller)
1002         return false;
1003       assert(CI->getCaller() == &F && "Unexpected call!");
1004 
1005       auto Remark = [&](OptimizationRemark OR) {
1006         return OR << "OpenMP runtime call "
1007                   << ore::NV("OpenMPOptRuntime", RFI.Name) << " deduplicated";
1008       };
1009       emitRemark<OptimizationRemark>(CI, "OpenMPRuntimeDeduplicated", Remark);
1010 
1011       CGUpdater.removeCallSite(*CI);
1012       CI->replaceAllUsesWith(ReplVal);
1013       CI->eraseFromParent();
1014       ++NumOpenMPRuntimeCallsDeduplicated;
1015       Changed = true;
1016       return true;
1017     };
1018     RFI.foreachUse(SCC, ReplaceAndDeleteCB);
1019 
1020     return Changed;
1021   }
1022 
1023   /// Collect arguments that represent the global thread id in \p GTIdArgs.
1024   void collectGlobalThreadIdArguments(SmallSetVector<Value *, 16> &GTIdArgs) {
1025     // TODO: Below we basically perform a fixpoint iteration with a pessimistic
1026     //       initialization. We could define an AbstractAttribute instead and
1027     //       run the Attributor here once it can be run as an SCC pass.
1028 
1029     // Helper to check the argument \p ArgNo at all call sites of \p F for
1030     // a GTId.
1031     auto CallArgOpIsGTId = [&](Function &F, unsigned ArgNo, CallInst &RefCI) {
1032       if (!F.hasLocalLinkage())
1033         return false;
1034       for (Use &U : F.uses()) {
1035         if (CallInst *CI = getCallIfRegularCall(U)) {
1036           Value *ArgOp = CI->getArgOperand(ArgNo);
1037           if (CI == &RefCI || GTIdArgs.count(ArgOp) ||
1038               getCallIfRegularCall(
1039                   *ArgOp, &OMPInfoCache.RFIs[OMPRTL___kmpc_global_thread_num]))
1040             continue;
1041         }
1042         return false;
1043       }
1044       return true;
1045     };
1046 
1047     // Helper to identify uses of a GTId as GTId arguments.
1048     auto AddUserArgs = [&](Value &GTId) {
1049       for (Use &U : GTId.uses())
1050         if (CallInst *CI = dyn_cast<CallInst>(U.getUser()))
1051           if (CI->isArgOperand(&U))
1052             if (Function *Callee = CI->getCalledFunction())
1053               if (CallArgOpIsGTId(*Callee, U.getOperandNo(), *CI))
1054                 GTIdArgs.insert(Callee->getArg(U.getOperandNo()));
1055     };
1056 
1057     // The argument users of __kmpc_global_thread_num calls are GTIds.
1058     OMPInformationCache::RuntimeFunctionInfo &GlobThreadNumRFI =
1059         OMPInfoCache.RFIs[OMPRTL___kmpc_global_thread_num];
1060 
1061     GlobThreadNumRFI.foreachUse(SCC, [&](Use &U, Function &F) {
1062       if (CallInst *CI = getCallIfRegularCall(U, &GlobThreadNumRFI))
1063         AddUserArgs(*CI);
1064       return false;
1065     });
1066 
1067     // Transitively search for more arguments by looking at the users of the
1068     // ones we know already. During the search the GTIdArgs vector is extended
1069     // so we cannot cache the size nor can we use a range based for.
1070     for (unsigned u = 0; u < GTIdArgs.size(); ++u)
1071       AddUserArgs(*GTIdArgs[u]);
1072   }
1073 
1074   /// Kernel (=GPU) optimizations and utility functions
1075   ///
1076   ///{{
1077 
1078   /// Check if \p F is a kernel, hence entry point for target offloading.
1079   bool isKernel(Function &F) { return OMPInfoCache.Kernels.count(&F); }
1080 
1081   /// Cache to remember the unique kernel for a function.
1082   DenseMap<Function *, Optional<Kernel>> UniqueKernelMap;
1083 
1084   /// Find the unique kernel that will execute \p F, if any.
1085   Kernel getUniqueKernelFor(Function &F);
1086 
1087   /// Find the unique kernel that will execute \p I, if any.
1088   Kernel getUniqueKernelFor(Instruction &I) {
1089     return getUniqueKernelFor(*I.getFunction());
1090   }
1091 
1092   /// Rewrite the device (=GPU) code state machine create in non-SPMD mode in
1093   /// the cases we can avoid taking the address of a function.
1094   bool rewriteDeviceCodeStateMachine();
1095 
1096   ///
1097   ///}}
1098 
1099   /// Emit a remark generically
1100   ///
1101   /// This template function can be used to generically emit a remark. The
1102   /// RemarkKind should be one of the following:
1103   ///   - OptimizationRemark to indicate a successful optimization attempt
1104   ///   - OptimizationRemarkMissed to report a failed optimization attempt
1105   ///   - OptimizationRemarkAnalysis to provide additional information about an
1106   ///     optimization attempt
1107   ///
1108   /// The remark is built using a callback function provided by the caller that
1109   /// takes a RemarkKind as input and returns a RemarkKind.
1110   template <typename RemarkKind,
1111             typename RemarkCallBack = function_ref<RemarkKind(RemarkKind &&)>>
1112   void emitRemark(Instruction *Inst, StringRef RemarkName,
1113                   RemarkCallBack &&RemarkCB) const {
1114     Function *F = Inst->getParent()->getParent();
1115     auto &ORE = OREGetter(F);
1116 
1117     ORE.emit(
1118         [&]() { return RemarkCB(RemarkKind(DEBUG_TYPE, RemarkName, Inst)); });
1119   }
1120 
1121   /// Emit a remark on a function. Since only OptimizationRemark is supporting
1122   /// this, it can't be made generic.
1123   void
1124   emitRemarkOnFunction(Function *F, StringRef RemarkName,
1125                        function_ref<OptimizationRemark(OptimizationRemark &&)>
1126                            &&RemarkCB) const {
1127     auto &ORE = OREGetter(F);
1128 
1129     ORE.emit([&]() {
1130       return RemarkCB(OptimizationRemark(DEBUG_TYPE, RemarkName, F));
1131     });
1132   }
1133 
1134   /// The underlying module.
1135   Module &M;
1136 
1137   /// The SCC we are operating on.
1138   SmallVectorImpl<Function *> &SCC;
1139 
1140   /// Callback to update the call graph, the first argument is a removed call,
1141   /// the second an optional replacement call.
1142   CallGraphUpdater &CGUpdater;
1143 
1144   /// Callback to get an OptimizationRemarkEmitter from a Function *
1145   OptimizationRemarkGetter OREGetter;
1146 
1147   /// OpenMP-specific information cache. Also Used for Attributor runs.
1148   OMPInformationCache &OMPInfoCache;
1149 
1150   /// Attributor instance.
1151   Attributor &A;
1152 
1153   /// Helper function to run Attributor on SCC.
1154   bool runAttributor() {
1155     if (SCC.empty())
1156       return false;
1157 
1158     registerAAs();
1159 
1160     ChangeStatus Changed = A.run();
1161 
1162     LLVM_DEBUG(dbgs() << "[Attributor] Done with " << SCC.size()
1163                       << " functions, result: " << Changed << ".\n");
1164 
1165     return Changed == ChangeStatus::CHANGED;
1166   }
1167 
1168   /// Populate the Attributor with abstract attribute opportunities in the
1169   /// function.
1170   void registerAAs() {
1171     if (SCC.empty())
1172       return;
1173 
1174     // Create CallSite AA for all Getters.
1175     for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
1176       auto ICVInfo = OMPInfoCache.ICVs[static_cast<InternalControlVar>(Idx)];
1177 
1178       auto &GetterRFI = OMPInfoCache.RFIs[ICVInfo.Getter];
1179 
1180       auto CreateAA = [&](Use &U, Function &Caller) {
1181         CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &GetterRFI);
1182         if (!CI)
1183           return false;
1184 
1185         auto &CB = cast<CallBase>(*CI);
1186 
1187         IRPosition CBPos = IRPosition::callsite_function(CB);
1188         A.getOrCreateAAFor<AAICVTracker>(CBPos);
1189         return false;
1190       };
1191 
1192       GetterRFI.foreachUse(SCC, CreateAA);
1193     }
1194   }
1195 };
1196 
1197 Kernel OpenMPOpt::getUniqueKernelFor(Function &F) {
1198   if (!OMPInfoCache.ModuleSlice.count(&F))
1199     return nullptr;
1200 
1201   // Use a scope to keep the lifetime of the CachedKernel short.
1202   {
1203     Optional<Kernel> &CachedKernel = UniqueKernelMap[&F];
1204     if (CachedKernel)
1205       return *CachedKernel;
1206 
1207     // TODO: We should use an AA to create an (optimistic and callback
1208     //       call-aware) call graph. For now we stick to simple patterns that
1209     //       are less powerful, basically the worst fixpoint.
1210     if (isKernel(F)) {
1211       CachedKernel = Kernel(&F);
1212       return *CachedKernel;
1213     }
1214 
1215     CachedKernel = nullptr;
1216     if (!F.hasLocalLinkage())
1217       return nullptr;
1218   }
1219 
1220   auto GetUniqueKernelForUse = [&](const Use &U) -> Kernel {
1221     if (auto *Cmp = dyn_cast<ICmpInst>(U.getUser())) {
1222       // Allow use in equality comparisons.
1223       if (Cmp->isEquality())
1224         return getUniqueKernelFor(*Cmp);
1225       return nullptr;
1226     }
1227     if (auto *CB = dyn_cast<CallBase>(U.getUser())) {
1228       // Allow direct calls.
1229       if (CB->isCallee(&U))
1230         return getUniqueKernelFor(*CB);
1231       // Allow the use in __kmpc_kernel_prepare_parallel calls.
1232       if (Function *Callee = CB->getCalledFunction())
1233         if (Callee->getName() == "__kmpc_kernel_prepare_parallel")
1234           return getUniqueKernelFor(*CB);
1235       return nullptr;
1236     }
1237     // Disallow every other use.
1238     return nullptr;
1239   };
1240 
1241   // TODO: In the future we want to track more than just a unique kernel.
1242   SmallPtrSet<Kernel, 2> PotentialKernels;
1243   OMPInformationCache::foreachUse(F, [&](const Use &U) {
1244     PotentialKernels.insert(GetUniqueKernelForUse(U));
1245   });
1246 
1247   Kernel K = nullptr;
1248   if (PotentialKernels.size() == 1)
1249     K = *PotentialKernels.begin();
1250 
1251   // Cache the result.
1252   UniqueKernelMap[&F] = K;
1253 
1254   return K;
1255 }
1256 
1257 bool OpenMPOpt::rewriteDeviceCodeStateMachine() {
1258   OMPInformationCache::RuntimeFunctionInfo &KernelPrepareParallelRFI =
1259       OMPInfoCache.RFIs[OMPRTL___kmpc_kernel_prepare_parallel];
1260 
1261   bool Changed = false;
1262   if (!KernelPrepareParallelRFI)
1263     return Changed;
1264 
1265   for (Function *F : SCC) {
1266 
1267     // Check if the function is uses in a __kmpc_kernel_prepare_parallel call at
1268     // all.
1269     bool UnknownUse = false;
1270     bool KernelPrepareUse = false;
1271     unsigned NumDirectCalls = 0;
1272 
1273     SmallVector<Use *, 2> ToBeReplacedStateMachineUses;
1274     OMPInformationCache::foreachUse(*F, [&](Use &U) {
1275       if (auto *CB = dyn_cast<CallBase>(U.getUser()))
1276         if (CB->isCallee(&U)) {
1277           ++NumDirectCalls;
1278           return;
1279         }
1280 
1281       if (isa<ICmpInst>(U.getUser())) {
1282         ToBeReplacedStateMachineUses.push_back(&U);
1283         return;
1284       }
1285       if (!KernelPrepareUse && OpenMPOpt::getCallIfRegularCall(
1286                                    *U.getUser(), &KernelPrepareParallelRFI)) {
1287         KernelPrepareUse = true;
1288         ToBeReplacedStateMachineUses.push_back(&U);
1289         return;
1290       }
1291       UnknownUse = true;
1292     });
1293 
1294     // Do not emit a remark if we haven't seen a __kmpc_kernel_prepare_parallel
1295     // use.
1296     if (!KernelPrepareUse)
1297       continue;
1298 
1299     {
1300       auto Remark = [&](OptimizationRemark OR) {
1301         return OR << "Found a parallel region that is called in a target "
1302                      "region but not part of a combined target construct nor "
1303                      "nesed inside a target construct without intermediate "
1304                      "code. This can lead to excessive register usage for "
1305                      "unrelated target regions in the same translation unit "
1306                      "due to spurious call edges assumed by ptxas.";
1307       };
1308       emitRemarkOnFunction(F, "OpenMPParallelRegionInNonSPMD", Remark);
1309     }
1310 
1311     // If this ever hits, we should investigate.
1312     // TODO: Checking the number of uses is not a necessary restriction and
1313     // should be lifted.
1314     if (UnknownUse || NumDirectCalls != 1 ||
1315         ToBeReplacedStateMachineUses.size() != 2) {
1316       {
1317         auto Remark = [&](OptimizationRemark OR) {
1318           return OR << "Parallel region is used in "
1319                     << (UnknownUse ? "unknown" : "unexpected")
1320                     << " ways; will not attempt to rewrite the state machine.";
1321         };
1322         emitRemarkOnFunction(F, "OpenMPParallelRegionInNonSPMD", Remark);
1323       }
1324       continue;
1325     }
1326 
1327     // Even if we have __kmpc_kernel_prepare_parallel calls, we (for now) give
1328     // up if the function is not called from a unique kernel.
1329     Kernel K = getUniqueKernelFor(*F);
1330     if (!K) {
1331       {
1332         auto Remark = [&](OptimizationRemark OR) {
1333           return OR << "Parallel region is not known to be called from a "
1334                        "unique single target region, maybe the surrounding "
1335                        "function has external linkage?; will not attempt to "
1336                        "rewrite the state machine use.";
1337         };
1338         emitRemarkOnFunction(F, "OpenMPParallelRegionInMultipleKernesl",
1339                              Remark);
1340       }
1341       continue;
1342     }
1343 
1344     // We now know F is a parallel body function called only from the kernel K.
1345     // We also identified the state machine uses in which we replace the
1346     // function pointer by a new global symbol for identification purposes. This
1347     // ensures only direct calls to the function are left.
1348 
1349     {
1350       auto RemarkParalleRegion = [&](OptimizationRemark OR) {
1351         return OR << "Specialize parallel region that is only reached from a "
1352                      "single target region to avoid spurious call edges and "
1353                      "excessive register usage in other target regions. "
1354                      "(parallel region ID: "
1355                   << ore::NV("OpenMPParallelRegion", F->getName())
1356                   << ", kernel ID: "
1357                   << ore::NV("OpenMPTargetRegion", K->getName()) << ")";
1358       };
1359       emitRemarkOnFunction(F, "OpenMPParallelRegionInNonSPMD",
1360                            RemarkParalleRegion);
1361       auto RemarkKernel = [&](OptimizationRemark OR) {
1362         return OR << "Target region containing the parallel region that is "
1363                      "specialized. (parallel region ID: "
1364                   << ore::NV("OpenMPParallelRegion", F->getName())
1365                   << ", kernel ID: "
1366                   << ore::NV("OpenMPTargetRegion", K->getName()) << ")";
1367       };
1368       emitRemarkOnFunction(K, "OpenMPParallelRegionInNonSPMD", RemarkKernel);
1369     }
1370 
1371     Module &M = *F->getParent();
1372     Type *Int8Ty = Type::getInt8Ty(M.getContext());
1373 
1374     auto *ID = new GlobalVariable(
1375         M, Int8Ty, /* isConstant */ true, GlobalValue::PrivateLinkage,
1376         UndefValue::get(Int8Ty), F->getName() + ".ID");
1377 
1378     for (Use *U : ToBeReplacedStateMachineUses)
1379       U->set(ConstantExpr::getBitCast(ID, U->get()->getType()));
1380 
1381     ++NumOpenMPParallelRegionsReplacedInGPUStateMachine;
1382 
1383     Changed = true;
1384   }
1385 
1386   return Changed;
1387 }
1388 
1389 /// Abstract Attribute for tracking ICV values.
1390 struct AAICVTracker : public StateWrapper<BooleanState, AbstractAttribute> {
1391   using Base = StateWrapper<BooleanState, AbstractAttribute>;
1392   AAICVTracker(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
1393 
1394   void initialize(Attributor &A) override {
1395     Function *F = getAnchorScope();
1396     if (!F || !A.isFunctionIPOAmendable(*F))
1397       indicatePessimisticFixpoint();
1398   }
1399 
1400   /// Returns true if value is assumed to be tracked.
1401   bool isAssumedTracked() const { return getAssumed(); }
1402 
1403   /// Returns true if value is known to be tracked.
1404   bool isKnownTracked() const { return getAssumed(); }
1405 
1406   /// Create an abstract attribute biew for the position \p IRP.
1407   static AAICVTracker &createForPosition(const IRPosition &IRP, Attributor &A);
1408 
1409   /// Return the value with which \p I can be replaced for specific \p ICV.
1410   virtual Optional<Value *> getReplacementValue(InternalControlVar ICV,
1411                                                 const Instruction *I,
1412                                                 Attributor &A) const {
1413     return None;
1414   }
1415 
1416   /// Return an assumed unique ICV value if a single candidate is found. If
1417   /// there cannot be one, return a nullptr. If it is not clear yet, return the
1418   /// Optional::NoneType.
1419   virtual Optional<Value *>
1420   getUniqueReplacementValue(InternalControlVar ICV) const = 0;
1421 
1422   // Currently only nthreads is being tracked.
1423   // this array will only grow with time.
1424   InternalControlVar TrackableICVs[1] = {ICV_nthreads};
1425 
1426   /// See AbstractAttribute::getName()
1427   const std::string getName() const override { return "AAICVTracker"; }
1428 
1429   /// See AbstractAttribute::getIdAddr()
1430   const char *getIdAddr() const override { return &ID; }
1431 
1432   /// This function should return true if the type of the \p AA is AAICVTracker
1433   static bool classof(const AbstractAttribute *AA) {
1434     return (AA->getIdAddr() == &ID);
1435   }
1436 
1437   static const char ID;
1438 };
1439 
1440 struct AAICVTrackerFunction : public AAICVTracker {
1441   AAICVTrackerFunction(const IRPosition &IRP, Attributor &A)
1442       : AAICVTracker(IRP, A) {}
1443 
1444   // FIXME: come up with better string.
1445   const std::string getAsStr() const override { return "ICVTrackerFunction"; }
1446 
1447   // FIXME: come up with some stats.
1448   void trackStatistics() const override {}
1449 
1450   /// We don't manifest anything for this AA.
1451   ChangeStatus manifest(Attributor &A) override {
1452     return ChangeStatus::UNCHANGED;
1453   }
1454 
1455   // Map of ICV to their values at specific program point.
1456   EnumeratedArray<DenseMap<Instruction *, Value *>, InternalControlVar,
1457                   InternalControlVar::ICV___last>
1458       ICVReplacementValuesMap;
1459 
1460   ChangeStatus updateImpl(Attributor &A) override {
1461     ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
1462 
1463     Function *F = getAnchorScope();
1464 
1465     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
1466 
1467     for (InternalControlVar ICV : TrackableICVs) {
1468       auto &SetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Setter];
1469 
1470       auto &ValuesMap = ICVReplacementValuesMap[ICV];
1471       auto TrackValues = [&](Use &U, Function &) {
1472         CallInst *CI = OpenMPOpt::getCallIfRegularCall(U);
1473         if (!CI)
1474           return false;
1475 
1476         // FIXME: handle setters with more that 1 arguments.
1477         /// Track new value.
1478         if (ValuesMap.insert(std::make_pair(CI, CI->getArgOperand(0))).second)
1479           HasChanged = ChangeStatus::CHANGED;
1480 
1481         return false;
1482       };
1483 
1484       auto CallCheck = [&](Instruction &I) {
1485         Optional<Value *> ReplVal = getValueForCall(A, &I, ICV);
1486         if (ReplVal.hasValue() &&
1487             ValuesMap.insert(std::make_pair(&I, *ReplVal)).second)
1488           HasChanged = ChangeStatus::CHANGED;
1489 
1490         return true;
1491       };
1492 
1493       // Track all changes of an ICV.
1494       SetterRFI.foreachUse(TrackValues, F);
1495 
1496       A.checkForAllInstructions(CallCheck, *this, {Instruction::Call},
1497                                 /* CheckBBLivenessOnly */ true);
1498 
1499       /// TODO: Figure out a way to avoid adding entry in
1500       /// ICVReplacementValuesMap
1501       Instruction *Entry = &F->getEntryBlock().front();
1502       if (HasChanged == ChangeStatus::CHANGED && !ValuesMap.count(Entry))
1503         ValuesMap.insert(std::make_pair(Entry, nullptr));
1504     }
1505 
1506     return HasChanged;
1507   }
1508 
1509   /// Hepler to check if \p I is a call and get the value for it if it is
1510   /// unique.
1511   Optional<Value *> getValueForCall(Attributor &A, const Instruction *I,
1512                                     InternalControlVar &ICV) const {
1513 
1514     const auto *CB = dyn_cast<CallBase>(I);
1515     if (!CB)
1516       return None;
1517 
1518     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
1519     auto &GetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Getter];
1520     auto &SetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Setter];
1521     Function *CalledFunction = CB->getCalledFunction();
1522 
1523     // Indirect call, assume ICV changes.
1524     if (CalledFunction == nullptr)
1525       return nullptr;
1526     if (CalledFunction == GetterRFI.Declaration)
1527       return None;
1528     if (CalledFunction == SetterRFI.Declaration) {
1529       if (ICVReplacementValuesMap[ICV].count(I))
1530         return ICVReplacementValuesMap[ICV].lookup(I);
1531 
1532       return nullptr;
1533     }
1534 
1535     // Since we don't know, assume it changes the ICV.
1536     if (CalledFunction->isDeclaration())
1537       return nullptr;
1538 
1539     const auto &ICVTrackingAA =
1540         A.getAAFor<AAICVTracker>(*this, IRPosition::callsite_returned(*CB));
1541 
1542     if (ICVTrackingAA.isAssumedTracked())
1543       return ICVTrackingAA.getUniqueReplacementValue(ICV);
1544 
1545     // If we don't know, assume it changes.
1546     return nullptr;
1547   }
1548 
1549   // We don't check unique value for a function, so return None.
1550   Optional<Value *>
1551   getUniqueReplacementValue(InternalControlVar ICV) const override {
1552     return None;
1553   }
1554 
1555   /// Return the value with which \p I can be replaced for specific \p ICV.
1556   Optional<Value *> getReplacementValue(InternalControlVar ICV,
1557                                         const Instruction *I,
1558                                         Attributor &A) const override {
1559     const auto &ValuesMap = ICVReplacementValuesMap[ICV];
1560     if (ValuesMap.count(I))
1561       return ValuesMap.lookup(I);
1562 
1563     SmallVector<const Instruction *, 16> Worklist;
1564     SmallPtrSet<const Instruction *, 16> Visited;
1565     Worklist.push_back(I);
1566 
1567     Optional<Value *> ReplVal;
1568 
1569     while (!Worklist.empty()) {
1570       const Instruction *CurrInst = Worklist.pop_back_val();
1571       if (!Visited.insert(CurrInst).second)
1572         continue;
1573 
1574       const BasicBlock *CurrBB = CurrInst->getParent();
1575 
1576       // Go up and look for all potential setters/calls that might change the
1577       // ICV.
1578       while ((CurrInst = CurrInst->getPrevNode())) {
1579         if (ValuesMap.count(CurrInst)) {
1580           Optional<Value *> NewReplVal = ValuesMap.lookup(CurrInst);
1581           // Unknown value, track new.
1582           if (!ReplVal.hasValue()) {
1583             ReplVal = NewReplVal;
1584             break;
1585           }
1586 
1587           // If we found a new value, we can't know the icv value anymore.
1588           if (NewReplVal.hasValue())
1589             if (ReplVal != NewReplVal)
1590               return nullptr;
1591 
1592           break;
1593         }
1594 
1595         Optional<Value *> NewReplVal = getValueForCall(A, CurrInst, ICV);
1596         if (!NewReplVal.hasValue())
1597           continue;
1598 
1599         // Unknown value, track new.
1600         if (!ReplVal.hasValue()) {
1601           ReplVal = NewReplVal;
1602           break;
1603         }
1604 
1605         // if (NewReplVal.hasValue())
1606         // We found a new value, we can't know the icv value anymore.
1607         if (ReplVal != NewReplVal)
1608           return nullptr;
1609       }
1610 
1611       // If we are in the same BB and we have a value, we are done.
1612       if (CurrBB == I->getParent() && ReplVal.hasValue())
1613         return ReplVal;
1614 
1615       // Go through all predecessors and add terminators for analysis.
1616       for (const BasicBlock *Pred : predecessors(CurrBB))
1617         if (const Instruction *Terminator = Pred->getTerminator())
1618           Worklist.push_back(Terminator);
1619     }
1620 
1621     return ReplVal;
1622   }
1623 };
1624 
1625 struct AAICVTrackerFunctionReturned : AAICVTracker {
1626   AAICVTrackerFunctionReturned(const IRPosition &IRP, Attributor &A)
1627       : AAICVTracker(IRP, A) {}
1628 
1629   // FIXME: come up with better string.
1630   const std::string getAsStr() const override {
1631     return "ICVTrackerFunctionReturned";
1632   }
1633 
1634   // FIXME: come up with some stats.
1635   void trackStatistics() const override {}
1636 
1637   /// We don't manifest anything for this AA.
1638   ChangeStatus manifest(Attributor &A) override {
1639     return ChangeStatus::UNCHANGED;
1640   }
1641 
1642   // Map of ICV to their values at specific program point.
1643   EnumeratedArray<Optional<Value *>, InternalControlVar,
1644                   InternalControlVar::ICV___last>
1645       ICVReplacementValuesMap;
1646 
1647   /// Return the value with which \p I can be replaced for specific \p ICV.
1648   Optional<Value *>
1649   getUniqueReplacementValue(InternalControlVar ICV) const override {
1650     return ICVReplacementValuesMap[ICV];
1651   }
1652 
1653   ChangeStatus updateImpl(Attributor &A) override {
1654     ChangeStatus Changed = ChangeStatus::UNCHANGED;
1655     const auto &ICVTrackingAA = A.getAAFor<AAICVTracker>(
1656         *this, IRPosition::function(*getAnchorScope()));
1657 
1658     if (!ICVTrackingAA.isAssumedTracked())
1659       return indicatePessimisticFixpoint();
1660 
1661     for (InternalControlVar ICV : TrackableICVs) {
1662       Optional<Value *> &ReplVal = ICVReplacementValuesMap[ICV];
1663       Optional<Value *> UniqueICVValue;
1664 
1665       auto CheckReturnInst = [&](Instruction &I) {
1666         Optional<Value *> NewReplVal =
1667             ICVTrackingAA.getReplacementValue(ICV, &I, A);
1668 
1669         // If we found a second ICV value there is no unique returned value.
1670         if (UniqueICVValue.hasValue() && UniqueICVValue != NewReplVal)
1671           return false;
1672 
1673         UniqueICVValue = NewReplVal;
1674 
1675         return true;
1676       };
1677 
1678       if (!A.checkForAllInstructions(CheckReturnInst, *this, {Instruction::Ret},
1679                                      /* CheckBBLivenessOnly */ true))
1680         UniqueICVValue = nullptr;
1681 
1682       if (UniqueICVValue == ReplVal)
1683         continue;
1684 
1685       ReplVal = UniqueICVValue;
1686       Changed = ChangeStatus::CHANGED;
1687     }
1688 
1689     return Changed;
1690   }
1691 };
1692 
1693 struct AAICVTrackerCallSite : AAICVTracker {
1694   AAICVTrackerCallSite(const IRPosition &IRP, Attributor &A)
1695       : AAICVTracker(IRP, A) {}
1696 
1697   void initialize(Attributor &A) override {
1698     Function *F = getAnchorScope();
1699     if (!F || !A.isFunctionIPOAmendable(*F))
1700       indicatePessimisticFixpoint();
1701 
1702     // We only initialize this AA for getters, so we need to know which ICV it
1703     // gets.
1704     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
1705     for (InternalControlVar ICV : TrackableICVs) {
1706       auto ICVInfo = OMPInfoCache.ICVs[ICV];
1707       auto &Getter = OMPInfoCache.RFIs[ICVInfo.Getter];
1708       if (Getter.Declaration == getAssociatedFunction()) {
1709         AssociatedICV = ICVInfo.Kind;
1710         return;
1711       }
1712     }
1713 
1714     /// Unknown ICV.
1715     indicatePessimisticFixpoint();
1716   }
1717 
1718   ChangeStatus manifest(Attributor &A) override {
1719     if (!ReplVal.hasValue() || !ReplVal.getValue())
1720       return ChangeStatus::UNCHANGED;
1721 
1722     A.changeValueAfterManifest(*getCtxI(), **ReplVal);
1723     A.deleteAfterManifest(*getCtxI());
1724 
1725     return ChangeStatus::CHANGED;
1726   }
1727 
1728   // FIXME: come up with better string.
1729   const std::string getAsStr() const override { return "ICVTrackerCallSite"; }
1730 
1731   // FIXME: come up with some stats.
1732   void trackStatistics() const override {}
1733 
1734   InternalControlVar AssociatedICV;
1735   Optional<Value *> ReplVal;
1736 
1737   ChangeStatus updateImpl(Attributor &A) override {
1738     const auto &ICVTrackingAA = A.getAAFor<AAICVTracker>(
1739         *this, IRPosition::function(*getAnchorScope()));
1740 
1741     // We don't have any information, so we assume it changes the ICV.
1742     if (!ICVTrackingAA.isAssumedTracked())
1743       return indicatePessimisticFixpoint();
1744 
1745     Optional<Value *> NewReplVal =
1746         ICVTrackingAA.getReplacementValue(AssociatedICV, getCtxI(), A);
1747 
1748     if (ReplVal == NewReplVal)
1749       return ChangeStatus::UNCHANGED;
1750 
1751     ReplVal = NewReplVal;
1752     return ChangeStatus::CHANGED;
1753   }
1754 
1755   // Return the value with which associated value can be replaced for specific
1756   // \p ICV.
1757   Optional<Value *>
1758   getUniqueReplacementValue(InternalControlVar ICV) const override {
1759     return ReplVal;
1760   }
1761 };
1762 
1763 struct AAICVTrackerCallSiteReturned : AAICVTracker {
1764   AAICVTrackerCallSiteReturned(const IRPosition &IRP, Attributor &A)
1765       : AAICVTracker(IRP, A) {}
1766 
1767   // FIXME: come up with better string.
1768   const std::string getAsStr() const override {
1769     return "ICVTrackerCallSiteReturned";
1770   }
1771 
1772   // FIXME: come up with some stats.
1773   void trackStatistics() const override {}
1774 
1775   /// We don't manifest anything for this AA.
1776   ChangeStatus manifest(Attributor &A) override {
1777     return ChangeStatus::UNCHANGED;
1778   }
1779 
1780   // Map of ICV to their values at specific program point.
1781   EnumeratedArray<Optional<Value *>, InternalControlVar,
1782                   InternalControlVar::ICV___last>
1783       ICVReplacementValuesMap;
1784 
1785   /// Return the value with which associated value can be replaced for specific
1786   /// \p ICV.
1787   Optional<Value *>
1788   getUniqueReplacementValue(InternalControlVar ICV) const override {
1789     return ICVReplacementValuesMap[ICV];
1790   }
1791 
1792   ChangeStatus updateImpl(Attributor &A) override {
1793     ChangeStatus Changed = ChangeStatus::UNCHANGED;
1794     const auto &ICVTrackingAA = A.getAAFor<AAICVTracker>(
1795         *this, IRPosition::returned(*getAssociatedFunction()));
1796 
1797     // We don't have any information, so we assume it changes the ICV.
1798     if (!ICVTrackingAA.isAssumedTracked())
1799       return indicatePessimisticFixpoint();
1800 
1801     for (InternalControlVar ICV : TrackableICVs) {
1802       Optional<Value *> &ReplVal = ICVReplacementValuesMap[ICV];
1803       Optional<Value *> NewReplVal =
1804           ICVTrackingAA.getUniqueReplacementValue(ICV);
1805 
1806       if (ReplVal == NewReplVal)
1807         continue;
1808 
1809       ReplVal = NewReplVal;
1810       Changed = ChangeStatus::CHANGED;
1811     }
1812     return Changed;
1813   }
1814 };
1815 } // namespace
1816 
1817 const char AAICVTracker::ID = 0;
1818 
1819 AAICVTracker &AAICVTracker::createForPosition(const IRPosition &IRP,
1820                                               Attributor &A) {
1821   AAICVTracker *AA = nullptr;
1822   switch (IRP.getPositionKind()) {
1823   case IRPosition::IRP_INVALID:
1824   case IRPosition::IRP_FLOAT:
1825   case IRPosition::IRP_ARGUMENT:
1826   case IRPosition::IRP_CALL_SITE_ARGUMENT:
1827     llvm_unreachable("ICVTracker can only be created for function position!");
1828   case IRPosition::IRP_RETURNED:
1829     AA = new (A.Allocator) AAICVTrackerFunctionReturned(IRP, A);
1830     break;
1831   case IRPosition::IRP_CALL_SITE_RETURNED:
1832     AA = new (A.Allocator) AAICVTrackerCallSiteReturned(IRP, A);
1833     break;
1834   case IRPosition::IRP_CALL_SITE:
1835     AA = new (A.Allocator) AAICVTrackerCallSite(IRP, A);
1836     break;
1837   case IRPosition::IRP_FUNCTION:
1838     AA = new (A.Allocator) AAICVTrackerFunction(IRP, A);
1839     break;
1840   }
1841 
1842   return *AA;
1843 }
1844 
1845 PreservedAnalyses OpenMPOptPass::run(LazyCallGraph::SCC &C,
1846                                      CGSCCAnalysisManager &AM,
1847                                      LazyCallGraph &CG, CGSCCUpdateResult &UR) {
1848   if (!containsOpenMP(*C.begin()->getFunction().getParent(), OMPInModule))
1849     return PreservedAnalyses::all();
1850 
1851   if (DisableOpenMPOptimizations)
1852     return PreservedAnalyses::all();
1853 
1854   SmallVector<Function *, 16> SCC;
1855   // If there are kernels in the module, we have to run on all SCC's.
1856   bool SCCIsInteresting = !OMPInModule.getKernels().empty();
1857   for (LazyCallGraph::Node &N : C) {
1858     Function *Fn = &N.getFunction();
1859     SCC.push_back(Fn);
1860 
1861     // Do we already know that the SCC contains kernels,
1862     // or that OpenMP functions are called from this SCC?
1863     if (SCCIsInteresting)
1864       continue;
1865     // If not, let's check that.
1866     SCCIsInteresting |= OMPInModule.containsOMPRuntimeCalls(Fn);
1867   }
1868 
1869   if (!SCCIsInteresting || SCC.empty())
1870     return PreservedAnalyses::all();
1871 
1872   FunctionAnalysisManager &FAM =
1873       AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
1874 
1875   AnalysisGetter AG(FAM);
1876 
1877   auto OREGetter = [&FAM](Function *F) -> OptimizationRemarkEmitter & {
1878     return FAM.getResult<OptimizationRemarkEmitterAnalysis>(*F);
1879   };
1880 
1881   CallGraphUpdater CGUpdater;
1882   CGUpdater.initialize(CG, C, AM, UR);
1883 
1884   SetVector<Function *> Functions(SCC.begin(), SCC.end());
1885   BumpPtrAllocator Allocator;
1886   OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG, Allocator,
1887                                 /*CGSCC*/ Functions, OMPInModule.getKernels());
1888 
1889   Attributor A(Functions, InfoCache, CGUpdater);
1890 
1891   OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
1892   bool Changed = OMPOpt.run();
1893   if (Changed)
1894     return PreservedAnalyses::none();
1895 
1896   return PreservedAnalyses::all();
1897 }
1898 
1899 namespace {
1900 
1901 struct OpenMPOptLegacyPass : public CallGraphSCCPass {
1902   CallGraphUpdater CGUpdater;
1903   OpenMPInModule OMPInModule;
1904   static char ID;
1905 
1906   OpenMPOptLegacyPass() : CallGraphSCCPass(ID) {
1907     initializeOpenMPOptLegacyPassPass(*PassRegistry::getPassRegistry());
1908   }
1909 
1910   void getAnalysisUsage(AnalysisUsage &AU) const override {
1911     CallGraphSCCPass::getAnalysisUsage(AU);
1912   }
1913 
1914   bool doInitialization(CallGraph &CG) override {
1915     // Disable the pass if there is no OpenMP (runtime call) in the module.
1916     containsOpenMP(CG.getModule(), OMPInModule);
1917     return false;
1918   }
1919 
1920   bool runOnSCC(CallGraphSCC &CGSCC) override {
1921     if (!containsOpenMP(CGSCC.getCallGraph().getModule(), OMPInModule))
1922       return false;
1923     if (DisableOpenMPOptimizations || skipSCC(CGSCC))
1924       return false;
1925 
1926     SmallVector<Function *, 16> SCC;
1927     // If there are kernels in the module, we have to run on all SCC's.
1928     bool SCCIsInteresting = !OMPInModule.getKernels().empty();
1929     for (CallGraphNode *CGN : CGSCC) {
1930       Function *Fn = CGN->getFunction();
1931       if (!Fn || Fn->isDeclaration())
1932         continue;
1933       SCC.push_back(Fn);
1934 
1935       // Do we already know that the SCC contains kernels,
1936       // or that OpenMP functions are called from this SCC?
1937       if (SCCIsInteresting)
1938         continue;
1939       // If not, let's check that.
1940       SCCIsInteresting |= OMPInModule.containsOMPRuntimeCalls(Fn);
1941     }
1942 
1943     if (!SCCIsInteresting || SCC.empty())
1944       return false;
1945 
1946     CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
1947     CGUpdater.initialize(CG, CGSCC);
1948 
1949     // Maintain a map of functions to avoid rebuilding the ORE
1950     DenseMap<Function *, std::unique_ptr<OptimizationRemarkEmitter>> OREMap;
1951     auto OREGetter = [&OREMap](Function *F) -> OptimizationRemarkEmitter & {
1952       std::unique_ptr<OptimizationRemarkEmitter> &ORE = OREMap[F];
1953       if (!ORE)
1954         ORE = std::make_unique<OptimizationRemarkEmitter>(F);
1955       return *ORE;
1956     };
1957 
1958     AnalysisGetter AG;
1959     SetVector<Function *> Functions(SCC.begin(), SCC.end());
1960     BumpPtrAllocator Allocator;
1961     OMPInformationCache InfoCache(
1962         *(Functions.back()->getParent()), AG, Allocator,
1963         /*CGSCC*/ Functions, OMPInModule.getKernels());
1964 
1965     Attributor A(Functions, InfoCache, CGUpdater);
1966 
1967     OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
1968     return OMPOpt.run();
1969   }
1970 
1971   bool doFinalization(CallGraph &CG) override { return CGUpdater.finalize(); }
1972 };
1973 
1974 } // end anonymous namespace
1975 
1976 void OpenMPInModule::identifyKernels(Module &M) {
1977 
1978   NamedMDNode *MD = M.getOrInsertNamedMetadata("nvvm.annotations");
1979   if (!MD)
1980     return;
1981 
1982   for (auto *Op : MD->operands()) {
1983     if (Op->getNumOperands() < 2)
1984       continue;
1985     MDString *KindID = dyn_cast<MDString>(Op->getOperand(1));
1986     if (!KindID || KindID->getString() != "kernel")
1987       continue;
1988 
1989     Function *KernelFn =
1990         mdconst::dyn_extract_or_null<Function>(Op->getOperand(0));
1991     if (!KernelFn)
1992       continue;
1993 
1994     ++NumOpenMPTargetRegionKernels;
1995 
1996     Kernels.insert(KernelFn);
1997   }
1998 }
1999 
2000 bool llvm::omp::containsOpenMP(Module &M, OpenMPInModule &OMPInModule) {
2001   if (OMPInModule.isKnown())
2002     return OMPInModule;
2003 
2004   auto RecordFunctionsContainingUsesOf = [&](Function *F) {
2005     for (User *U : F->users())
2006       if (auto *I = dyn_cast<Instruction>(U))
2007         OMPInModule.FuncsWithOMPRuntimeCalls.insert(I->getFunction());
2008   };
2009 
2010   // MSVC doesn't like long if-else chains for some reason and instead just
2011   // issues an error. Work around it..
2012   do {
2013 #define OMP_RTL(_Enum, _Name, ...)                                             \
2014   if (Function *F = M.getFunction(_Name)) {                                    \
2015     RecordFunctionsContainingUsesOf(F);                                        \
2016     OMPInModule = true;                                                        \
2017   }
2018 #include "llvm/Frontend/OpenMP/OMPKinds.def"
2019   } while (false);
2020 
2021   // Identify kernels once. TODO: We should split the OMPInformationCache into a
2022   // module and an SCC part. The kernel information, among other things, could
2023   // go into the module part.
2024   if (OMPInModule.isKnown() && OMPInModule) {
2025     OMPInModule.identifyKernels(M);
2026     return true;
2027   }
2028 
2029   return OMPInModule = false;
2030 }
2031 
2032 char OpenMPOptLegacyPass::ID = 0;
2033 
2034 INITIALIZE_PASS_BEGIN(OpenMPOptLegacyPass, "openmpopt",
2035                       "OpenMP specific optimizations", false, false)
2036 INITIALIZE_PASS_DEPENDENCY(CallGraphWrapperPass)
2037 INITIALIZE_PASS_END(OpenMPOptLegacyPass, "openmpopt",
2038                     "OpenMP specific optimizations", false, false)
2039 
2040 Pass *llvm::createOpenMPOptLegacyPass() { return new OpenMPOptLegacyPass(); }
2041