[AbstractAttributor] Fold function calls to `__kmpc_is_spmd_exec_mode` if possible

[tianshilei1992]

Summary

In the device runtime there are many function calls to __kmpc_is_spmd_exec_mode
to query the execution mode of current kernels. In many cases, user programs
only contain target region executing in one mode. As a consequence, those runtime
function calls will only return one value. If we can get rid of these function
calls during compliation, it can potentially improve performance.

In this patch, we use AAKernelInfo to analyze kernel execution. Basically, for
each kernel (device) function F, we collect all kernel entries K that can
reach F. A new AA, AAFoldRuntimeCall, is created for each call site. In each
iteration, it will check all reaching kernel entries, and update the folded value
accordingly.

In the future we will support more function.

[tianshilei1992]

Tests will come soon

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3573

This might not be always true. What if later there is a kernel entry in another mode?

[tianshilei1992]

add test case and update some code. note: test case is not right currently. keep debugging.

[jdoerfert]

We need to properly tack when SPMDCompatibility is known.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3359

I don't think we need this anymore, do we?

3439

Do we need/want a KernelInfoState here? Wouldn't a BooleanState suffice to do the tacking of valid/invalid and fixpoint. Plus the Optional<Constant *> SimplifiedValue;.

3488

3573

I don't think that can happen

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3573

Function call in is_mixed_helper in the test case is folded by mistake.

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3439

Nope because we access CallerKernelInfoAA.ReachingKernelEntries directly in foldIsSPMDExecMode

[tianshilei1992]

fixed issue that assumed information is not being used

[tianshilei1992]

fix format

[tianshilei1992]

fix comments

[tianshilei1992]

llvm/test/Transforms/OpenMP/is_spmd_exec_mode_fold.ll

95

The function call here can only be folded if we use assumed information as SPMDization doesn't change the state to known.

[tianshilei1992]

removed value simplification call back registration as it can affect SPMDization

[tianshilei1992]

not use assumed information

[tianshilei1992]

refined getAsStr and changed to BooleanState

[tianshilei1992]

rebase and reuse assumed information

[jdoerfert]

Minor nits, basically good to go.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3540–3543

Same below please.

3571

Technically we could even try to stay with nullptr but not pessimistic fixppoint here.

3593

Remove the empty conditionals above.

llvm/test/Transforms/OpenMP/is_spmd_exec_mode_fold.ll

130

Can we have a second is_generic helper like this:

define internal void @is_non_spmd_helper2() {
;
  %isSPMD = call i8 @__kmpc_is_spmd_exec_mode()
  %c = icmp eq i8 %isSPMD, 0
  br i1 %c, label %t, label %f
t:
  call void @foo()
  ret void
f:
  call void bar()
  ret void
}

[tianshilei1992]

rebase and fix comments

[tianshilei1992]

leftover

[tianshilei1992]

reorg some code

[jdoerfert]

Change the title of the commit, clang format. LGTM

[tianshilei1992]

rebase, fix tests, clang-formats

[tianshilei1992]

clang-format

[tianshilei1992]

Reopen this to fix hangs.

[tianshilei1992]

fixed an issue that insertion of ReachingKernelEntries invalidates it

[jdoerfert]

LG, one nit.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3579–3583

below we might get to all generic but here we can give up

[tianshilei1992]

fixed comments

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3250

I don't think we need this anymore, do we?

[tianshilei1992]

fixed a series of critical issues

[tianshilei1992]

update

[Meinersbur]

Thank you for the revert, it has fixed many failures in the buildbot openmp-offload-cuda-project (https://lab.llvm.org/staging/#/builders/155/builds/1885) with only 5 remaining (4 by PR49940, PR50738 and PR50739; and 1 caused by either D101976 or D101977)

Strangely, openmp-offload-cuda-runtime only followed much later when the next clean build was pending (https://lab.llvm.org/staging/#/builders/154/builds/1509). There might be a dependency problem in the LLVM_ENABLE_RUNTIMES build.

[tianshilei1992]

In D105787#2882375, @Meinersbur wrote:

Thank you for the revert, it has fixed many failures in the buildbot openmp-offload-cuda-project (https://lab.llvm.org/staging/#/builders/155/builds/1885) with only 5 remaining (4 by PR49940, PR50738 and PR50739; and 1 caused by either D101976 or D101977)

Strangely, openmp-offload-cuda-runtime only followed much later when the next clean build was pending (https://lab.llvm.org/staging/#/builders/154/builds/1509). There might be a dependency problem in the LLVM_ENABLE_RUNTIMES build.

but I committed it again after all known issues have been fixed. Does it look good now?

[ye-luo]

@tianshilei1992
There are still issues remaining after you reapply the patch.
One of them seems being fixed by 97c8f60bbaf0986320fdfd03b11328b91c730a96
But more remaining
https://bugs.llvm.org/show_bug.cgi?id=51124

[tianshilei1992]

In D105787#2885399, @ye-luo wrote:

@tianshilei1992
There are still issues remaining after you reapply the patch.
One of them seems being fixed by 97c8f60bbaf0986320fdfd03b11328b91c730a96
But more remaining
https://bugs.llvm.org/show_bug.cgi?id=51124

Thanks for the report. Would you mind giving https://reviews.llvm.org/D106209 a shot to see if it can fix the issue?

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

struct KernelInfoState : AbstractState {
   /// The __kmpc_target_init call in this kernel, if any. If we find more than
   /// one we abort as the kernel is malformed.
   CallBase *KernelInitCB = nullptr;
 
   /// The __kmpc_target_deinit call in this kernel, if any. If we find more than
   /// one we abort as the kernel is malformed.
   CallBase *KernelDeinitCB = nullptr;
 
+  /// Flag to indicate if the associated function is a kernel entry.
+  bool IsKernelEntry = false;
+
+  /// State to track what kernel entries can reach the associated function.
+  BooleanStateWithPtrSetVector<Function> ReachingKernelEntries;
+
   /// Abstract State interface
   ///{
 
   KernelInfoState() {}
   KernelInfoState(bool BestState) {
     if (!BestState)
       indicatePessimisticFixpoint();
   }
struct KernelInfoState : AbstractState {
 
   bool operator==(const KernelInfoState &RHS) const {
     if (SPMDCompatibilityTracker != RHS.SPMDCompatibilityTracker)
       return false;
     if (ReachedKnownParallelRegions != RHS.ReachedKnownParallelRegions)
       return false;
     if (ReachedUnknownParallelRegions != RHS.ReachedUnknownParallelRegions)
       return false;
+    if (ReachingKernelEntries != RHS.ReachingKernelEntries)
+      return false;
     return true;
   }
 
   /// Return empty set as the best state of potential values.
   static KernelInfoState getBestState() { return KernelInfoState(true); }
 
   static KernelInfoState getBestState(KernelInfoState &KIS) {
     return getBestState();
void initialize(Attributor &A) override {
     // of the init and dinit calls. We need to tell the Attributor about this
     // to avoid other parts using the current constant value for simpliication.
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
 
     Function *Fn = getAnchorScope();
     if (!OMPInfoCache.Kernels.count(Fn))
       return;
 
+    // Add itself to the reaching kernel and set IsKernelEntry.
+    ReachingKernelEntries.insert(Fn);
+    IsKernelEntry = true;
+
     OMPInformationCache::RuntimeFunctionInfo &InitRFI =
         OMPInfoCache.RFIs[OMPRTL___kmpc_target_init];
     OMPInformationCache::RuntimeFunctionInfo &DeinitRFI =
         OMPInfoCache.RFIs[OMPRTL___kmpc_target_deinit];
 
     // For kernels we perform more initialization work, first we find the init
     // and deinit calls.
     auto StoreCallBase = [](Use &U,
auto CheckRWInst = [&](Instruction &I) {
       return true;
     };
 
     bool UsedAssumedInformationInCheckRWInst = false;
     if (!A.checkForAllReadWriteInstructions(
             CheckRWInst, *this, UsedAssumedInformationInCheckRWInst))
       SPMDCompatibilityTracker.indicatePessimisticFixpoint();
 
+    if (!IsKernelEntry)
+       updateReachingKernelEntries(A);
+
     // Callback to check a call instruction.
     auto CheckCallInst = [&](Instruction &I) {
       auto &CB = cast<CallBase>(I);
       auto &CBAA = A.getAAFor<AAKernelInfo>(
           *this, IRPosition::callsite_function(CB), DepClassTy::OPTIONAL);
       if (CBAA.getState().isValidState())
         getState() ^= CBAA.getState();
+
+      Function *Callee = CB.getCalledFunction();
+      if (Callee) {
+        // We need to propagate information to the callee, but since the
+        // construction of AA always starts with kernel entries, we have to
+        // create AAKernelInfoFunction for all called functions. However, here
+        // the caller doesn't depend on the callee.
+        // TODO: We might want to change the dependence here later if we need
+        // information from callee to caller.
+        A.getOrCreateAAFor<AAKernelInfo>(IRPosition::function(*Callee), this,
+                                         DepClassTy::NONE);
+      }
+
       return true;
     };
 
     bool UsedAssumedInformationInCheckCallInst = false;
     if (!A.checkForAllCallLikeInstructions(
             CheckCallInst, *this, UsedAssumedInformationInCheckCallInst))
       return indicatePessimisticFixpoint();
 
     return StateBefore == getState() ? ChangeStatus::UNCHANGED
                                      : ChangeStatus::CHANGED;
   }
+
+private:
+  /// Update info regarding reaching kernels.
+  void updateReachingKernelEntries(Attributor &A) {
+    auto PredCallSite = [&](AbstractCallSite ACS) {
+      Function *Caller = ACS.getInstruction()->getFunction();
+
+      assert(Caller && "Caller is nullptr");
+
+      auto &CAA =
+          A.getOrCreateAAFor<AAKernelInfo>(IRPosition::function(*Caller));
+      if (CAA.isValidState()) {
+        ReachingKernelEntries ^= CAA.ReachingKernelEntries;
+        return true;
+      }
+
+      // We lost track of the caller of the associated function, any kernel
+      // could reach now.
+      ReachingKernelEntries.indicatePessimisticFixpoint();
+
+      return true;
+    };
+
+    bool AllCallSitesKnown;
+    if (!A.checkForAllCallSites(PredCallSite, *this,
+                                true /* RequireAllCallSites */,
+                                AllCallSitesKnown))
+      ReachingKernelEntries.indicatePessimisticFixpoint();
+  }
 };
 
 /// The call site kernel info abstract attribute, basically, what can we say
 /// about a call site with regards to the KernelInfoState. For now this simply
 /// forwards the information from the callee.
 struct AAKernelInfoCallSite : AAKernelInfo {
   AAKernelInfoCallSite(const IRPosition &IRP, Attributor &A)
       : AAKernelInfo(IRP, A) {}
if (It == OMPInfoCache.RuntimeFunctionIDMap.end()) {
       return;
     }
 
     const unsigned int WrapperFunctionArgNo = 6;
     RuntimeFunction RF = It->getSecond();
     switch (RF) {
     // All the functions we know are compatible with SPMD mode.
     case OMPRTL___kmpc_is_spmd_exec_mode:
     case OMPRTL___kmpc_for_static_fini:
     case OMPRTL___kmpc_global_thread_num:
     case OMPRTL___kmpc_single:
     case OMPRTL___kmpc_end_single:
     case OMPRTL___kmpc_master:
     case OMPRTL___kmpc_end_master:
     case OMPRTL___kmpc_barrier:
       break;
     case OMPRTL___kmpc_for_static_init_4:
ChangeStatus updateImpl(Attributor &A) override {
     auto &FnAA = A.getAAFor<AAKernelInfo>(*this, FnPos, DepClassTy::REQUIRED);
     if (getState() == FnAA.getState())
       return ChangeStatus::UNCHANGED;
     getState() = FnAA.getState();
     return ChangeStatus::CHANGED;
   }
 };
 
+struct AAFoldRuntimeCall
+    : public StateWrapper<BooleanState, AbstractAttribute> {
+  using Base = StateWrapper<BooleanState, AbstractAttribute>;
+
+  AAFoldRuntimeCall(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
+
+  /// Statistics are tracked as part of manifest for now.
+  void trackStatistics() const override {}
+
+  /// Create an abstract attribute biew for the position \p IRP.
+  static AAFoldRuntimeCall &createForPosition(const IRPosition &IRP,
+                                              Attributor &A);
+
+  /// See AbstractAttribute::getName()
+  const std::string getName() const override { return "AAFoldRuntimeCall"; }
+
+  /// See AbstractAttribute::getIdAddr()
+  const char *getIdAddr() const override { return &ID; }
+
+  /// This function should return true if the type of the \p AA is
+  /// AAFoldRuntimeCall
+  static bool classof(const AbstractAttribute *AA) {
+    return (AA->getIdAddr() == &ID);
+  }
+
+  static const char ID;
+};
+
+struct AAFoldRuntimeCallCallSite : AAFoldRuntimeCall {
+  AAFoldRuntimeCallCallSite(const IRPosition &IRP, Attributor &A)
+      : AAFoldRuntimeCall(IRP, A) {}
+
+  /// See AbstractAttribute::getAsStr()
+  const std::string getAsStr() const override {
+    if (!isValidState())
+      return "<invalid>";
+
+    std::string Str("simplified value: ");
+
+    if (!SimplifiedValue.hasValue())
+      return Str + std::string("none");
+
+    if (!SimplifiedValue.getValue())
+      return Str + std::string("nullptr");
+
+    if (ConstantInt *CI = dyn_cast<ConstantInt>(SimplifiedValue.getValue()))
+      return Str + std::to_string(CI->getSExtValue());
+
+    return Str + std::string("unknown");
+  }
+
+  void initialize(Attributor &A) override {
+    Function *Callee = getAssociatedFunction();
+
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    const auto &It = OMPInfoCache.RuntimeFunctionIDMap.find(Callee);
+    assert(It != OMPInfoCache.RuntimeFunctionIDMap.end() &&
+           "Expected a known OpenMP runtime function");
+
+    RFKind = It->getSecond();
+
+    CallBase &CB = cast<CallBase>(getAssociatedValue());
+    A.registerSimplificationCallback(
+        IRPosition::callsite_returned(CB),
+        [&](const IRPosition &IRP, const AbstractAttribute *AA,
+            bool &UsedAssumedInformation) -> Optional<Value *> {
+          if (!isAtFixpoint()) {
+            UsedAssumedInformation = true;
+            if (AA)
+              A.recordDependence(*this, *AA, DepClassTy::OPTIONAL);
+          }
+          return SimplifiedValue;
+        });
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
+
+    switch (RFKind) {
+    case OMPRTL___kmpc_is_spmd_exec_mode:
+      Changed = Changed | foldIsSPMDExecMode(A);
+      break;
+    default:
+      llvm_unreachable("Unhandled OpenMP runtime function!");
+    }
+
+    return Changed;
+  }
+
+  /// Fold __kmpc_is_spmd_exec_mode into a constant if possible.
+  ChangeStatus foldIsSPMDExecMode(Attributor &A) {
+    BooleanState StateBefore = getState();
+
+    unsigned AssumedSPMDCount = 0, KnownSPMDCount = 0;
+    unsigned AssumedNonSPMDCount = 0, KnownNonSPMDCount = 0;
+    auto &CallerKernelInfoAA = A.getAAFor<AAKernelInfo>(
+        *this, IRPosition::function(*getAnchorScope()), DepClassTy::REQUIRED);
+
+    for (Kernel K : CallerKernelInfoAA.ReachingKernelEntries) {
+      auto &AA = A.getAAFor<AAKernelInfo>(*this, IRPosition::function(*K),
+                                          DepClassTy::REQUIRED);
+
+      if (!AA.isValidState()) {
+        SimplifiedValue = nullptr;
+        return indicatePessimisticFixpoint();
+      }
+
+      if (AA.SPMDCompatibilityTracker.isAssumed()) {
+        if (AA.SPMDCompatibilityTracker.isAtFixpoint())
+          ++KnownSPMDCount;
+        else
+          ++AssumedSPMDCount;
+      } else {
+        if (AA.SPMDCompatibilityTracker.isAtFixpoint())
+          ++KnownNonSPMDCount;
+        else
+          ++AssumedNonSPMDCount;
+      }
+    }
+
+    if (KnownSPMDCount && KnownNonSPMDCount) {
+      SimplifiedValue = nullptr;
+      return getState() == StateBefore ? ChangeStatus::UNCHANGED
+                                       : ChangeStatus::CHANGED;
+    }
+
+    if (AssumedSPMDCount && AssumedNonSPMDCount) {
+      SimplifiedValue = nullptr;
+      return getState() == StateBefore ? ChangeStatus::UNCHANGED
+                                       : ChangeStatus::CHANGED;
+    }
+
+    auto &Ctx = getAnchorValue().getContext();
+    if (KnownSPMDCount || AssumedSPMDCount) {
+      assert(KnownNonSPMDCount == 0 && AssumedNonSPMDCount == 0 &&
+             "Expected only SPMD kernels!");
+      // All reaching kernels are in SPMD mode. Update all function calls to
+      // __kmpc_is_spmd_exec_mode to 1.
+      SimplifiedValue = ConstantInt::get(Type::getInt8Ty(Ctx), true);
+    } else {
+      assert(KnownSPMDCount == 0 && AssumedSPMDCount == 0 &&
+             "Expected only non-SPMD kernels!");
+      // All reaching kernels are in non-SPMD mode. Update all function
+      // calls to __kmpc_is_spmd_exec_mode to 0.
+      SimplifiedValue = ConstantInt::get(Type::getInt8Ty(Ctx), false);
+    }
+
+    return getState() == StateBefore ? ChangeStatus::UNCHANGED
+                                     : ChangeStatus::CHANGED;
+  }
+
+  ChangeStatus manifest(Attributor &A) override {
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
+
+    if (SimplifiedValue.hasValue() && SimplifiedValue.getValue()) {
+      Instruction &CB = *getCtxI();
+      A.changeValueAfterManifest(CB, **SimplifiedValue);
+      A.deleteAfterManifest(CB);
+      Changed = ChangeStatus::CHANGED;
+    }
+
+    return Changed;
+  }
+
+private:
+  /// An optional value the associated value is assumed to fold to. That is, we
+  /// assume the associated value (which is a call) can be replaced by this
+  /// simplified value.
+  Optional<Value *> SimplifiedValue;
+
+  /// The runtime function kind of the callee of the associated call site.
+  RuntimeFunction RFKind;
+};
+
 } // namespace
 
 void OpenMPOpt::registerAAs(bool IsModulePass) {
   if (SCC.empty())
 
     return;
   if (IsModulePass) {
     // Ensure we create the AAKernelInfo AAs first and without triggering an
     // update. This will make sure we register all value simplification
     // callbacks before any other AA has the chance to create an AAValueSimplify
     // or similar.
     for (Function *Kernel : OMPInfoCache.Kernels)
       A.getOrCreateAAFor<AAKernelInfo>(
           IRPosition::function(*Kernel), /* QueryingAA */ nullptr,
           DepClassTy::NONE, /* ForceUpdate */ false,
           /* UpdateAfterInit */ false);
+
+    auto &IsSPMDRFI = OMPInfoCache.RFIs[OMPRTL___kmpc_is_spmd_exec_mode];
+    IsSPMDRFI.foreachUse(SCC, [&](Use &U, Function &) {
+      CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &IsSPMDRFI);
+      if (!CI)
+        return false;
+      A.getOrCreateAAFor<AAFoldRuntimeCall>(
+          IRPosition::callsite_function(*CI), /* QueryingAA */ nullptr,
+          DepClassTy::NONE, /* ForceUpdate */ false,
+          /* UpdateAfterInit */ false);
+      return false;
+    });
   }
 
   // Create CallSite AA for all Getters.
   for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
     auto ICVInfo = OMPInfoCache.ICVs[static_cast<InternalControlVar>(Idx)];
 
     auto &GetterRFI = OMPInfoCache.RFIs[ICVInfo.Getter];
 
if (isOpenMPDevice(M))
       A.getOrCreateAAFor<AAHeapToStack>(IRPosition::function(*F));
   }
 }
 
 const char AAICVTracker::ID = 0;
 const char AAKernelInfo::ID = 0;
 const char AAExecutionDomain::ID = 0;
 const char AAHeapToShared::ID = 0;
+const char AAFoldRuntimeCall::ID = 0;
 
 AAICVTracker &AAICVTracker::createForPosition(const IRPosition &IRP,
                                               Attributor &A) {
   AAICVTracker *AA = nullptr;
   switch (IRP.getPositionKind()) {
   case IRPosition::IRP_INVALID:
   case IRPosition::IRP_FLOAT:
   case IRPosition::IRP_ARGUMENT:
AAKernelInfo &AAKernelInfo::createForPosition(const IRPosition &IRP,
   case IRPosition::IRP_FUNCTION:
     AA = new (A.Allocator) AAKernelInfoFunction(IRP, A);
     break;
   }
 
   return *AA;
 }
 
+AAFoldRuntimeCall &AAFoldRuntimeCall::createForPosition(const IRPosition &IRP,
+                                                        Attributor &A) {
+  AAFoldRuntimeCall *AA = nullptr;
+  switch (IRP.getPositionKind()) {
+  case IRPosition::IRP_INVALID:
+  case IRPosition::IRP_FLOAT:
+  case IRPosition::IRP_ARGUMENT:
+  case IRPosition::IRP_RETURNED:
+  case IRPosition::IRP_CALL_SITE_RETURNED:
+  case IRPosition::IRP_CALL_SITE_ARGUMENT:
+  case IRPosition::IRP_FUNCTION:
+    llvm_unreachable("KernelInfo can only be created for call site position!");
+  case IRPosition::IRP_CALL_SITE:
+    AA = new (A.Allocator) AAFoldRuntimeCallCallSite(IRP, A);
+    break;
+  }
+
+  return *AA;
+}
+
 PreservedAnalyses OpenMPOptPass::run(Module &M, ModuleAnalysisManager &AM) {
   if (!containsOpenMP(M))
     return PreservedAnalyses::all();
   if (DisableOpenMPOptimizations)
     return PreservedAnalyses::all();
 
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();

llvm/test/Transforms/OpenMP/is_spmd_exec_mode_fold.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-globals
+; RUN: opt -S -passes=openmp-opt < %s | FileCheck %s
+target triple = "nvptx64"
+
+%struct.ident_t = type { i32, i32, i32, i32, i8* }
+
+@is_spmd_exec_mode = weak constant i8 0
+@will_be_spmd_exec_mode = weak constant i8 1
+@non_spmd_exec_mode = weak constant i8 1
+@will_not_be_spmd_exec_mode = weak constant i8 1
+@G = external global i8
[e-mail] = appending global [4 x i8*] [i8* @is_spmd_exec_mode, i8* @will_be_spmd_exec_mode, i8* @non_spmd_exec_mode, i8* @will_not_be_spmd_exec_mode ], section "llvm.metadata"
+
+;.
+; CHECK: @[[IS_SPMD_EXEC_MODE:[a-zA-Z0-9_$"\\.-]+]] = weak constant i8 0
+; CHECK: @[[WILL_BE_SPMD_EXEC_MODE:[a-zA-Z0-9_$"\\.-]+]] = weak constant i8 0
+; CHECK: @[[NON_SPMD_EXEC_MODE:[a-zA-Z0-9_$"\\.-]+]] = weak constant i8 1
+; CHECK: @[[WILL_NOT_BE_SPMD_EXEC_MODE:[a-zA-Z0-9_$"\\.-]+]] = weak constant i8 1
+; CHECK: @[[G:[a-zA-Z0-9_$"\\.-]+]] = external global i8
+; CHECK: @[[LLVM_COMPILER_USED:[a-zA-Z0-9_$"\\.-]+]] = appending global [4 x i8*] [i8* @is_spmd_exec_mode, i8* @will_be_spmd_exec_mode, i8* @non_spmd_exec_mode, i8* @will_not_be_spmd_exec_mode], section "llvm.metadata"
+;.
+define weak void @is_spmd() {
+; CHECK-LABEL: define {{[^@]+}}@is_spmd() {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 true, i1 false, i1 false)
+; CHECK-NEXT:    call void @is_spmd_helper1()
+; CHECK-NEXT:    call void @is_spmd_helper2()
+; CHECK-NEXT:    call void @is_mixed_helper()
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 true, i1 false, i1 false)
+  call void @is_spmd_helper1()
+  call void @is_spmd_helper2()
+  call void @is_mixed_helper()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+  ret void
+}
+
+define weak void @will_be_spmd() {
+; CHECK-LABEL: define {{[^@]+}}@will_be_spmd() {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 true, i1 false, i1 false)
+; CHECK-NEXT:    call void @is_spmd_helper2()
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+  call void @is_spmd_helper2()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+  ret void
+}
+
+define weak void @non_spmd() {
+; CHECK-LABEL: define {{[^@]+}}@non_spmd() {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+; CHECK-NEXT:    call void @is_generic_helper1()
+; CHECK-NEXT:    call void @is_generic_helper2()
+; CHECK-NEXT:    call void @is_mixed_helper()
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+  call void @is_generic_helper1()
+  call void @is_generic_helper2()
+  call void @is_mixed_helper()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+  ret void
+}
+
+define weak void @will_not_be_spmd() {
+; CHECK-LABEL: define {{[^@]+}}@will_not_be_spmd() {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+; CHECK-NEXT:    call void @is_generic_helper1()
+; CHECK-NEXT:    call void @is_generic_helper2()
+; CHECK-NEXT:    call void @is_mixed_helper()
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+  call void @is_generic_helper1()
+  call void @is_generic_helper2()
+  call void @is_mixed_helper()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+  ret void
+}
+
+define internal void @is_spmd_helper1() {
+; CHECK-LABEL: define {{[^@]+}}@is_spmd_helper1() {
+; CHECK-NEXT:    store i8 1, i8* @G, align 1
+; CHECK-NEXT:    ret void
+;
+  %isSPMD = call i8 @__kmpc_is_spmd_exec_mode()
+  store i8 %isSPMD, i8* @G
+  ret void
+}
+
+define internal void @is_spmd_helper2() {
+; CHECK-LABEL: define {{[^@]+}}@is_spmd_helper2() {
+; CHECK-NEXT:    br label [[F:%.*]]
+; CHECK:       t:
+; CHECK-NEXT:    unreachable
+; CHECK:       f:
+; CHECK-NEXT:    ret void
+;
+  %isSPMD = call i8 @__kmpc_is_spmd_exec_mode()
+  %c = icmp eq i8 %isSPMD, 0
+  br i1 %c, label %t, label %f
+t:
+  call void @spmd_compatible()
+  ret void
+f:
+  ret void
+}
+
+define internal void @is_generic_helper1() {
+; CHECK-LABEL: define {{[^@]+}}@is_generic_helper1() {
+; CHECK-NEXT:    store i8 0, i8* @G, align 1
+; CHECK-NEXT:    ret void
+;
+  %isSPMD = call i8 @__kmpc_is_spmd_exec_mode()
+  store i8 %isSPMD, i8* @G
+  ret void
+}
+
+define internal void @is_generic_helper2() {
+; CHECK-LABEL: define {{[^@]+}}@is_generic_helper2() {
+; CHECK-NEXT:    br label [[T:%.*]]
+; CHECK:       t:
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    ret void
+; CHECK:       f:
+; CHECK-NEXT:    unreachable
+;
+  %isSPMD = call i8 @__kmpc_is_spmd_exec_mode()
+  %c = icmp eq i8 %isSPMD, 0
+  br i1 %c, label %t, label %f
+t:
+  call void @foo()
+  ret void
+f:
+  call void @bar()
+  ret void
+}
+
+define internal void @is_mixed_helper() {
+; CHECK-LABEL: define {{[^@]+}}@is_mixed_helper() {
+; CHECK-NEXT:    [[ISSPMD:%.*]] = call i8 @__kmpc_is_spmd_exec_mode()
+; CHECK-NEXT:    store i8 undef, i8* @G, align 1
+; CHECK-NEXT:    ret void
+;
+  %isSPMD = call i8 @__kmpc_is_spmd_exec_mode()
+  store i8 %isSPMD, i8* @G
+  ret void
+}
+
+declare void @spmd_compatible() "llvm.assume"="ompx_spmd_amenable"
+declare i8 @__kmpc_is_spmd_exec_mode()
+declare i32 @__kmpc_target_init(%struct.ident_t*, i1 zeroext, i1 zeroext, i1 zeroext) #1
+declare void @__kmpc_target_deinit(%struct.ident_t* nocapture readnone, i1 zeroext, i1 zeroext) #1
+declare void @foo()
+declare void @bar()
+
+!llvm.module.flags = !{!0, !1}
+!nvvm.annotations = !{!2, !3, !4, !5}
+
+!0 = !{i32 7, !"openmp", i32 50}
+!1 = !{i32 7, !"openmp-device", i32 50}
+!2 = !{void ()* @is_spmd, !"kernel", i32 1}
+!3 = !{void ()* @will_be_spmd, !"kernel", i32 1}
+!4 = !{void ()* @non_spmd, !"kernel", i32 1}
+!5 = !{void ()* @will_not_be_spmd, !"kernel", i32 1}
+;.
+; CHECK: attributes #[[ATTR0:[0-9]+]] = { "llvm.assume"="ompx_spmd_amenable" }
+;.
+; CHECK: [[META0:![0-9]+]] = !{i32 7, !"openmp", i32 50}
+; CHECK: [[META1:![0-9]+]] = !{i32 7, !"openmp-device", i32 50}
+; CHECK: [[META2:![0-9]+]] = !{void ()* @is_spmd, !"kernel", i32 1}
+; CHECK: [[META3:![0-9]+]] = !{void ()* @will_be_spmd, !"kernel", i32 1}
+; CHECK: [[META4:![0-9]+]] = !{void ()* @non_spmd, !"kernel", i32 1}
+; CHECK: [[META5:![0-9]+]] = !{void ()* @will_not_be_spmd, !"kernel", i32 1}
+;.