[OpenMPOpt][WIP][NOT-FOR-REVIEW] Track an ICV that determines if we are in a SPMD GPU region

[jdoerfert]

Summary

NOTE: This is a draft.

We can use the new ICV tacking to do fancy optimiztions w/o much effort.
This adds all the machinery to track if we are in an "SPMD region" or
not. The corresponding getter is replaced accordingly, which can easily
result in significant performance boosts. As the ICV tacking is still a
bit in flux, we'll wait till there are a few more pieces in place. Then
we'll add tests to this and review it properly.

llvm/include/llvm/Frontend/OpenMP/OMPConstants.h

 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 };
 
 #define ICV_DATA_ENV(Enum, ...)                                                \
   constexpr auto Enum = omp::InternalControlVar::Enum;
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 
 enum class ICVInitValue {
-#define ICV_DATA_ENV(Enum, Name, EnvVar, Init) Init,
+#define ICV_INIT_VALUE(Enum, Name) Enum,
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 };
 
-#define ICV_DATA_ENV(Enum, Name, EnvVar, Init)                                 \
-  constexpr auto Init = omp::ICVInitValue::Init;
+#define ICV_INIT_VALUE(Enum, Name)                                             \
+  constexpr auto Enum = omp::ICVInitValue::Enum;
+#include "llvm/Frontend/OpenMP/OMPKinds.def"
+
+enum class ICVInterferenceKind {
+#define ICV_INTERFERENCE_KIND(Enum, Name) Enum,
+#include "llvm/Frontend/OpenMP/OMPKinds.def"
+};
+
+#define ICV_INTERFERENCE_KIND(Enum, Name)                                      \
+  constexpr auto Enum = omp::ICVInterferenceKind::Enum;
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 
 /// IDs for all omp runtime library (RTL) functions.
 enum class RuntimeFunction {
 #define OMP_RTL(Enum, ...) Enum,
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 };
 

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

 #define OMP_TYPE(VarName, InitValue)
 #endif
 
 #define __OMP_TYPE(VarName) OMP_TYPE(VarName, Type::get##VarName##Ty(Ctx))
 
 __OMP_TYPE(Void)
 __OMP_TYPE(Int1)
 __OMP_TYPE(Int8)
+__OMP_TYPE(Int16)
 __OMP_TYPE(Int32)
 __OMP_TYPE(Int64)
 __OMP_TYPE(Int8Ptr)
 __OMP_TYPE(Int32Ptr)
 __OMP_TYPE(Int64Ptr)
 
 OMP_TYPE(SizeTy, M.getDataLayout().getIntPtrType(Ctx))
 

 #undef OMP_FUNCTION_TYPE
 
 ///}
 
 /// Internal Control Variables information
 ///
 ///{
 
+#ifndef ICV_INIT_VALUE
+#define ICV_INIT_VALUE(Enum, Name)
+#endif
+
+#define __ICV_INIT_VALUE(Name) ICV_INIT_VALUE(ICV_##Name, #Name)
+
+__ICV_INIT_VALUE(ZERO)
+__ICV_INIT_VALUE(FALSE)
+__ICV_INIT_VALUE(IMPLEMENTATION_DEFINED)
+__ICV_INIT_VALUE(LAST)
+
+#undef __ICV_INIT_VALUE
+#undef ICV_INIT_VALUE
+
+#ifndef ICV_INTERFERENCE_KIND
+#define ICV_INTERFERENCE_KIND(Enum, Name)
+#endif
+
+#define __ICV_INTERFERENCE_KIND(Name)                                          \
+  ICV_INTERFERENCE_KIND(ICV_INTERFERENCE_KIND_##Name, Name)
+
+__ICV_INTERFERENCE_KIND(anything)
+__ICV_INTERFERENCE_KIND(nothing)
+
+#undef __ICV_INTERFERENCE_KIND
+#undef ICV_INTERFERENCE_KIND
+
 #ifndef ICV_DATA_ENV
 #define ICV_DATA_ENV(Enum, Name, EnvVarName, Init)
 #endif
 
 #define __ICV_DATA_ENV(Name, EnvVarName, Init)                                 \
   ICV_DATA_ENV(ICV_##Name, #Name, #EnvVarName, Init)
 
 __ICV_DATA_ENV(nthreads, OMP_NUM_THREADS, ICV_IMPLEMENTATION_DEFINED)
 __ICV_DATA_ENV(active_levels, NONE, ICV_ZERO)
 __ICV_DATA_ENV(cancel, OMP_CANCELLATION, ICV_FALSE)
+__ICV_DATA_ENV(GPU_SPMD, NONE, ICV_FALSE)
 __ICV_DATA_ENV(__last, last, ICV_LAST)
 
 #undef __ICV_DATA_ENV
 #undef ICV_DATA_ENV
 
 #ifndef ICV_RT_SET
-#define ICV_RT_SET(Name, RTL)
+#define ICV_RT_SET(Name, RTL, ARGUMENT_NO, INTERFERNCE_KIND)
 #endif
 
-#define __ICV_RT_SET(Name, RTL) ICV_RT_SET(ICV_##Name, OMPRTL_##RTL)
+#define __ICV_RT_SET(Name, RTL, ARGUMENT_NO, INTERFERNCE_KIND)                 \
+  ICV_RT_SET(ICV_##Name, OMPRTL_##RTL, ARGUMENT_NO,                            \
+             ICV_INTERFERENCE_KIND_##INTERFERNCE_KIND)
 
-__ICV_RT_SET(nthreads, omp_set_num_threads)
+__ICV_RT_SET(nthreads, omp_set_num_threads, /* ArgumentNo */ 1, anything)
+__ICV_RT_SET(GPU_SPMD, __kmpc_spmd_kernel_init, /* ArgumentNo */ -1, nothing)
 
 #undef __ICV_RT_SET
 #undef ICV_RT_SET
 
 #ifndef ICV_RT_GET
 #define ICV_RT_GET(Name, RTL)
 #endif
 
 #define __ICV_RT_GET(Name, RTL) ICV_RT_GET(ICV_##Name, OMPRTL_##RTL)
 
 __ICV_RT_GET(nthreads, omp_get_max_threads)
 __ICV_RT_GET(active_levels, omp_get_active_level)
 __ICV_RT_GET(cancel, omp_get_cancellation)
+__ICV_RT_GET(GPU_SPMD, __kmpc_is_spmd_exec_mode)
 
 #undef __ICV_RT_GET
 #undef ICV_RT_GET
 
 ///}
 
 /// Runtime library function (and their attributes)
 ///

 __OMP_RTL(omp_get_partition_place_nums, false, Void, Int32Ptr)
 
 __OMP_RTL(omp_set_num_threads, false, Void, Int32)
 __OMP_RTL(omp_set_dynamic, false, Void, Int32)
 __OMP_RTL(omp_set_nested, false, Void, Int32)
 __OMP_RTL(omp_set_schedule, false, Void, Int32, Int32)
 __OMP_RTL(omp_set_max_active_levels, false, Void, Int32)
 
+__OMP_RTL(__kmpc_is_spmd_exec_mode, false, Int8, )
+__OMP_RTL(__kmpc_spmd_kernel_init, false, Void, Int32, Int16, Int16)
+
 __OMP_RTL(__kmpc_master, false, Int32, IdentPtr, Int32)
 __OMP_RTL(__kmpc_end_master, false, Void, IdentPtr, Int32)
 __OMP_RTL(__kmpc_critical, false, Void, IdentPtr, Int32, KmpCriticalNamePtrTy)
 __OMP_RTL(__kmpc_critical_with_hint, false, Void, IdentPtr, Int32,
           KmpCriticalNamePtrTy, Int32)
 __OMP_RTL(__kmpc_end_critical, false, Void, IdentPtr, Int32,
           KmpCriticalNamePtrTy)
 

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

struct InternalControlVarInfo {
     ICVInitValue InitKind;
 
     /// Initial value.
     ConstantInt *InitValue;
 
     /// Setter RTL function associated with this ICV.
     RuntimeFunction Setter;
 
+    /// Which argument of the setter defines the value.
+    int SetterArgNo;
+
+    /// Kind to describe how this ICV is impacted by calls, e.g., if the value
+    /// can change or not by an unknown call.
+    ICVInterferenceKind InterferenceKind;
+
     /// Getter RTL function associated with this ICV.
     RuntimeFunction Getter;
 
     /// RTL Function corresponding to the override clause of this ICV
     RuntimeFunction Clause;
   };
 
   /// Generic information that describes a runtime function
struct OMPInformationCache : public InformationCache {
   /// Map from ICV kind to the ICV description.
   EnumeratedArray<InternalControlVarInfo, InternalControlVar,
                   InternalControlVar::ICV___last>
       ICVs;
 
   /// Helper to initialize all internal control variable information for those
   /// defined in OMPKinds.def.
   void initializeInternalControlVars() {
-#define ICV_RT_SET(_Name, RTL)                                                 \
+#define ICV_RT_SET(_Name, RTL, ARGUMENT_NO, INTERFERENCE_KIND)                 \
   {                                                                            \
     auto &ICV = ICVs[_Name];                                                   \
     ICV.Setter = RTL;                                                          \
+    ICV.SetterArgNo = ARGUMENT_NO;                                             \
+    ICV.InterferenceKind = INTERFERENCE_KIND;                                  \
   }
 #define ICV_RT_GET(Name, RTL)                                                  \
   {                                                                            \
     auto &ICV = ICVs[Name];                                                    \
     ICV.Getter = RTL;                                                          \
   }
 #define ICV_DATA_ENV(Enum, _Name, _EnvVarName, Init)                           \
   {                                                                            \
bool deduplicateICVGetters(InternalControlVar &ICV, Attributor &A) {
     return Changed;
   }
 
   // Map of ICV to their values at specific program point.
   EnumeratedArray<SmallSetVector<ICVValue, 4>, InternalControlVar,
                   InternalControlVar::ICV___last>
       ICVValuesMap;
 
-  // Currently only nthreads is being tracked.
-  // this array will only grow with time.
-  InternalControlVar TrackableICVs[1] = {ICV_nthreads};
+  InternalControlVar TrackableICVs[2] = {ICV_nthreads, ICV_GPU_SPMD};
 
   ChangeStatus updateImpl(Attributor &A) override {
     ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
 
     Function *F = getAnchorScope();
 
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
 
     for (InternalControlVar ICV : TrackableICVs) {
-      auto &SetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Setter];
+      auto &ICVInfo = OMPInfoCache.ICVs[ICV];
+      auto &SetterRFI = OMPInfoCache.RFIs[ICVInfo.Setter];
 
       auto TrackValues = [&](Use &U, Function &) {
         CallInst *CI = OpenMPOpt::getCallIfRegularCall(U);
         if (!CI)
           return false;
 
-        // FIXME: handle setters with more that 1 arguments.
-        /// Track new value.
-        if (ICVValuesMap[ICV].insert(ICVValue(CI, CI->getArgOperand(0))))
+        // Track new value.
+        Value *ICVIRValue;
+        if (ICVInfo.SetterArgNo >= 0)
+          ICVIRValue = CI->getArgOperand(ICVInfo.SetterArgNo);
+        else
+          ICVIRValue = ConstantInt::get(Type::getInt8Ty(CI->getContext()), 1);
+
+        if (ICVValuesMap[ICV].insert(ICVValue(CI, ICVIRValue)))
           HasChanged = ChangeStatus::CHANGED;
 
         return false;
       };
 
       SetterRFI.foreachUse(TrackValues, F);
     }
 
     return HasChanged;
   }
 
   /// Return the value with which \p I can be replaced for specific \p ICV.
   Value *getReplacementValue(InternalControlVar ICV, const Instruction *I,
                              Attributor &A) override {
     const BasicBlock *CurrBB = I->getParent();
 
     auto &ValuesSet = ICVValuesMap[ICV];
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
-    auto &GetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Getter];
+    auto &ICVInfo = OMPInfoCache.ICVs[ICV];
+    auto &GetterRFI = OMPInfoCache.RFIs[ICVInfo.Getter];
 
     for (const auto &ICVVal : ValuesSet) {
       if (CurrBB == ICVVal.Inst->getParent()) {
         if (!ICVVal.Inst->comesBefore(I))
           continue;
+        // If nothing can interfere with this ICV, we don't need to check.
+        if (ICVInfo.InterferenceKind == ICV_INTERFERENCE_KIND_nothing)
+          return ICVVal.TrackedValue;
 
         // both instructions are in the same BB and at \p I we know the ICV
         // value.
         while (I != ICVVal.Inst) {
           // we don't yet know if a call might update an ICV.
           // TODO: check callsite AA for value.
           if (const auto *CB = dyn_cast<CallBase>(I))
             if (CB->getCalledFunction() != GetterRFI.Declaration)