[OpenMP] Create custom state machines for generic target regions

[jdoerfert]

Summary

In the spirit of TRegions [0], this patch creates a custom state
machine for a generic target region based on the potentially called
parallel regions.

The code analysis is done interprocedurally via an abstract attribute
(AAKernelInfo). All outermost parallel regions are collected and we
check if there might be unknown outermost parallel regions for which
we need an indirect call. Other AAKernelInfo extensions are expected.

[0] https://link.springer.com/chapter/10.1007/978-3-030-28596-8_11

[jdoerfert]

NOTE: tests are coming soon.

[jdoerfert]

Fix porting mistakes

[jdoerfert]

Add comments and tests

[jdoerfert]

FWIW, here is what clang currently generates: https://godbolt.org/z/ncWq8f437
With these two patches the result is better in basically every case except the last where there is no actual improvement to be made.

[jdoerfert]

Rebase

[jdoerfert]

Update tests with --function-signature to avoid check-label matching calls

[ggeorgakoudis]

LG, some remarks

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

509

Nit, concatenation?

612

Assertion?

686

skeleton

[ggeorgakoudis]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

496

Is equality comparison enough for ParallelRegions or should we change ParallelRegions refer to same function? I guess size equality should be enough since we only add ParallelRegions, but still checking.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

496

I will check the content, it should not matter how we use it but you're right, better to plan for the future.

[jdoerfert]

Addressed comments
Also moved logic from the AAKernelInfoFunction::updateImpl to
AAKernelInfoCallSite::initialize as it looks cleaner to fix the
state of runtime calls right away.

[jdoerfert]

[Wrong review -.- ]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

509

It won't always be described properly by "concatenation". Already the |= part is on the border, once we do SPMD detection there will be a &= as well.

[jdoerfert]

ping?

[ggeorgakoudis]

LGTM, minor comment, fix tests?

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

794

Remove dump()?

[jdoerfert]

Rebase

[jdoerfert]

Update test and check for internalized functions too.

[jdoerfert]

Track unknown parallel calls and add remarks for the state machine

[jdoerfert]

rebase, remarks, and more tests

[jhuber6]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2902

I think it should be clear that the optimization pass changed something. I feel just saying "Is executed" is a little ambiguous to what actually happened.

2926

Should this be an analysis remark? It pairs with the analysis remarks below.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2902

sure.

2926

this is the matching remark to the ones above. Why would we make one transformation and the other one analysis, this is describing the transformation as the other two are.

[jdoerfert]

Use a proper state to track unknown and known reached parallel regions.

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2865

Is it good to return ahead of time? In later patches we also create AAKernelInfoFunction for "regular" functions.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2865

It is good to keep control flow simple. Later, as needed, we can have an initializeKernel and initializeNonKernel, for example.

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

460

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

466

If the state is invalid or in pessimistic state, do we want to clear the set?

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

460

size_t is probably the worst, but sure.

466

no, we want the content to emit remarks.

[JonChesterfield]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2901

I'm seeing this assert fire. The kernel contains both the init and deinit calls, but the InitRFI/DeinitRFI for each use doesn't see any uses, so StoreCallBase isn't called, and nothing is assumed to Storage. Kernel otherwise looks uninteresting, run into while trying to track down another bug.

For my reference, issue_001.c from aomp, debug clang build at trunk today, invoked as:

clang -O2 -fopenmp -fopenmp-targets=amdgcn-amd-amdhsa -Xopenmp-target=amdgcn-amd-amdhsa -march=gfx906 issue_001.c -o issue_001
gets bitcode corruption, invalid record

If passed -openmp-opt-disable=true, no corruption seen. If that is used to extract amdgpu bitcode which is fed back into opt, this assert trips.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2901

Can you emit the IR (module) before openmp-opt?

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 STATISTIC(NumOpenMPParallelRegionsDeleted,
           "Number of OpenMP parallel regions deleted");
 STATISTIC(NumOpenMPRuntimeFunctionsIdentified,
           "Number of OpenMP runtime functions identified");
 STATISTIC(NumOpenMPRuntimeFunctionUsesIdentified,
           "Number of OpenMP runtime function uses identified");
 STATISTIC(NumOpenMPTargetRegionKernels,
           "Number of OpenMP target region entry points (=kernels) identified");
+STATISTIC(NumOpenMPTargetRegionKernelsSPMD,
+          "Number of OpenMP target region entry points (=kernels) executed in "
+          "SPMD-mode instead of generic-mode");
+STATISTIC(NumOpenMPTargetRegionKernelsCustomStateMachine,
+          "Number of OpenMP target region entry points (=kernels) executed in "
+          "generic-mode with customized state machines");
 STATISTIC(
     NumOpenMPParallelRegionsReplacedInGPUStateMachine,
     "Number of OpenMP parallel regions replaced with ID in GPU state machines");
 STATISTIC(NumOpenMPParallelRegionsMerged,
           "Number of OpenMP parallel regions merged");
 
 #if !defined(NDEBUG)
 static constexpr auto TAG = "[" DEBUG_TYPE "]";
void foreachUse(function_ref<bool(Use &, Function &)> CB, Function *F) {
         UV.pop_back();
       }
     }
 
   private:
     /// Map from functions to all uses of this runtime function contained in
     /// them.
     DenseMap<Function *, std::shared_ptr<UseVector>> UsesMap;
+
+  public:
+    /// Iterators for the uses of this runtime function.
+    decltype(UsesMap)::iterator begin() { return UsesMap.begin(); }
+    decltype(UsesMap)::iterator end() { return UsesMap.end(); }
   };
 
   /// An OpenMP-IR-Builder instance
   OpenMPIRBuilder OMPBuilder;
 
   /// Map from runtime function kind to the runtime function description.
   EnumeratedArray<RuntimeFunctionInfo, RuntimeFunction,
                   RuntimeFunction::OMPRTL___last>
       RFIs;
 
+  /// Map from function declarations/definitions to their runtime enum type.
+  DenseMap<Function *, RuntimeFunction> RuntimeFunctionIDMap;
+
   /// 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 OMP_STRUCT_TYPE(VarName, ...)                                          \
   PointerType *VarName##Ptr = OMPBuilder.VarName##Ptr;                         \
   (void)VarName##Ptr;
 
 #define OMP_RTL(_Enum, _Name, _IsVarArg, _ReturnType, ...)                     \
   {                                                                            \
     SmallVector<Type *, 8> ArgsTypes({__VA_ARGS__});                           \
     Function *F = M.getFunction(_Name);                                        \
     if (declMatchesRTFTypes(F, OMPBuilder._ReturnType, ArgsTypes)) {           \
+      RuntimeFunctionIDMap[F] = _Enum;                                         \
       auto &RFI = RFIs[_Enum];                                                 \
       RFI.Kind = _Enum;                                                        \
       RFI.Name = _Name;                                                        \
       RFI.IsVarArg = _IsVarArg;                                                \
       RFI.ReturnType = OMPBuilder._ReturnType;                                 \
       RFI.ArgumentTypes = std::move(ArgsTypes);                                \
       RFI.Declaration = F;                                                     \
       unsigned NumUses = collectUses(RFI);                                     \
#include "llvm/Frontend/OpenMP/OMPKinds.def"
 
     // TODO: We should attach the attributes defined in OMPKinds.def.
   }
 
   /// Collection of known kernels (\see Kernel) in the module.
   SmallPtrSetImpl<Kernel> &Kernels;
 };
 
+struct KernelInfoState : AbstractState {
+  KernelInfoState() {}
+  KernelInfoState(bool BestState) {
+    if (!BestState)
+      indicatePessimisticFixpoint();
+  }
+
+  /// See AbstractState::isValidState(...)
+  bool isValidState() const override { return true; }
+
+  /// See AbstractState::isAtFixpoint(...)
+  bool isAtFixpoint() const override { return IsAtFixpoint; }
+
+  /// See AbstractState::indicatePessimisticFixpoint(...)
+  ChangeStatus indicatePessimisticFixpoint() override {
+    IsAtFixpoint = true;
+    MayReachUnknownParallelRegion = true;
+    return ChangeStatus::CHANGED;
+  }
+
+  /// See AbstractState::indicateOptimisticFixpoint(...)
+  ChangeStatus indicateOptimisticFixpoint() override {
+    IsAtFixpoint = true;
+    return ChangeStatus::UNCHANGED;
+  }
+
+  /// Return the assumed state
+  KernelInfoState &getAssumed() { return *this; }
+  const KernelInfoState &getAssumed() const { return *this; }
+
+  bool operator==(const KernelInfoState &RHS) const {
+    if ((MayReachUnknownParallelRegion != RHS.MayReachUnknownParallelRegion))
+      return false;
+    return
+           ParallelRegions.size() == RHS.ParallelRegions.size();
+  }
+
+  /// Return empty set as the best state of potential values.
+  static KernelInfoState getBestState() { return KernelInfoState(true); }
+
+  static KernelInfoState getBestState(KernelInfoState &KIS) {
+    return getBestState();
+  }
+
+  /// Return full set as the worst state of potential values.
+  static KernelInfoState getWorstState() { return KernelInfoState(false); }
+
+  /// "Clamp" this state with \p KIS.
+  KernelInfoState operator^=(const KernelInfoState &KIS) {
+    MayReachUnknownParallelRegion |= KIS.MayReachUnknownParallelRegion;
+    ParallelRegions.insert(KIS.ParallelRegions.begin(),
+                           KIS.ParallelRegions.end());
+    return *this;
+  }
+
+  KernelInfoState operator&=(const KernelInfoState &KIS) {
+    return (*this ^= KIS);
+  }
+
+public:
+  bool IsAtFixpoint = false;
+
+  SmallSetVector<Function *, 2> ParallelRegions;
+  CallBase *KernelInitCB = nullptr;
+  CallBase *KernelDeinitCB = nullptr;
+  bool MayReachUnknownParallelRegion = false;
+};
+
+struct AAKernelInfo : public StateWrapper<KernelInfoState, AbstractAttribute> {
+  using Base = StateWrapper<KernelInfoState, AbstractAttribute>;
+  AAKernelInfo(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
+
+  ChangeStatus manifest(Attributor &A) override {
+    if (!KernelInitCB || !KernelDeinitCB)
+      return ChangeStatus::UNCHANGED;
+
+    // Callback to check a call instruction.
+    auto CheckCallInst = [&](Instruction &I) {
+      auto &CB = cast<CallBase>(I);
+      CB.addAttribute(AttributeList::FunctionIndex, Attribute::AlwaysInline);
+      return true;
+    };
+    if (!A.checkForAllCallLikeInstructions(CheckCallInst, *this))
+      return indicatePessimisticFixpoint();
+
+    auto &Ctx = getAnchorValue().getContext();
+    const int InitIsSPMDArgNo = 1;
+    const int InitUseStateMachineArgNo = 2;
+    ConstantInt *UseStateMachine = dyn_cast<ConstantInt>(
+        KernelInitCB->getArgOperand(InitUseStateMachineArgNo));
+    ConstantInt *IsSPMD =
+        dyn_cast<ConstantInt>(KernelInitCB->getArgOperand(InitIsSPMDArgNo));
+
+    // If we are stuck with generic mode, try to create a custom device (=GPU)
+    // state machine which is specialized for the parallel regions that are
+    // reachable by the kernel.
+    if (!UseStateMachine || UseStateMachine->isZero() || !IsSPMD ||
+        !IsSPMD->isZero()) {
+      return ChangeStatus::UNCHANGED;
+    }
+
+    // Indicate we use a custom state machine now.
+    A.changeUseAfterManifest(
+        KernelInitCB->getArgOperandUse(InitUseStateMachineArgNo),
+        *ConstantInt::getBool(Ctx, 0));
+
+    ++NumOpenMPTargetRegionKernelsCustomStateMachine;
+
+    // If we don't need a state machine we are done.
+    if (!MayReachUnknownParallelRegion && ParallelRegions.empty())
+      return ChangeStatus::CHANGED;
+
+    // Create all the blocks:
+    //
+    //                       InitCB = __kmpc_target_init(...)
+    //                       bool IsWorker = InitCB >= 0;
+    //                       if (IsWorker) {
+    // SMBeginBB:               __kmpc_barrier_simple_spmd(...);
+    //                         void *WorkFn;
+    //                         bool Active = __kmpc_kernel_parallel(&WorkFn);
+    //                         if (!WorkFn) return;
+    // SMIsActiveCheckBB:       if (Active) {
+    // SMIfCascadeCurrentBB:      if      (WorkFn == <ParFn0>)
+    //                              ParFn0(...);
+    // SMIfCascadeCurrentBB:      else if (WorkFn == <ParFn1>)
+    //                              ParFn1(...);
+    //                            ...
+    // SMIfCascadeCurrentBB:      else
+    //                              ((WorkFnTy*)WorkFn)(...);
+    // SMEndParallelBB:           __kmpc_kernel_end_parallel(...);
+    //                          }
+    // SMDoneBB:                __kmpc_barrier_simple_spmd(...);
+    //                          goto SMBeginBB;
+    //                       }
+    // UserCodeEntryBB:      // user code
+    //                       __kmpc_target_deinit(...)
+    //
+    Function *Kernel = getAssociatedFunction();
+    BasicBlock *InitBB = KernelInitCB->getParent();
+    BasicBlock *UserCodeEntryBB = InitBB->splitBasicBlock(
+        KernelInitCB->getNextNode(), "thread.user_code.check");
+    BasicBlock *StateMachineBeginBB = BasicBlock::Create(
+        Ctx, "worker_state_machine.begin", Kernel, UserCodeEntryBB);
+    BasicBlock *StateMachineFinishedBB = BasicBlock::Create(
+        Ctx, "worker_state_machine.finished", Kernel, UserCodeEntryBB);
+    BasicBlock *StateMachineIsActiveCheckBB = BasicBlock::Create(
+        Ctx, "worker_state_machine.is_active.check", Kernel, UserCodeEntryBB);
+    BasicBlock *StateMachineIfCascadeCurrentBB =
+        BasicBlock::Create(Ctx, "worker_state_machine.parallel_region.check",
+                           Kernel, UserCodeEntryBB);
+    BasicBlock *StateMachineEndParallelBB =
+        BasicBlock::Create(Ctx, "worker_state_machine.parallel_region.end",
+                           Kernel, UserCodeEntryBB);
+    BasicBlock *StateMachineDoneBarrierBB = BasicBlock::Create(
+        Ctx, "worker_state_machine.done.barrier", Kernel, UserCodeEntryBB);
+
+    ReturnInst::Create(Ctx, StateMachineFinishedBB);
+
+    InitBB->getTerminator()->eraseFromParent();
+    Value *IsWorker =
+        ICmpInst::Create(ICmpInst::ICmp, llvm::CmpInst::ICMP_NE, KernelInitCB,
+                         ConstantInt::get(KernelInitCB->getType(), -1),
+                         "thread.is_worker", InitBB);
+    BranchInst::Create(StateMachineBeginBB, UserCodeEntryBB, IsWorker, InitBB);
+
+    // Create local storage for the work function pointer.
+    Type *VoidPtrTy = Type::getInt8PtrTy(Ctx);
+    AllocaInst *WorkFnAI = new AllocaInst(VoidPtrTy, 0, "worker.work_fn.addr",
+                                          &Kernel->getEntryBlock().front());
+
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    OMPInfoCache.OMPBuilder.updateToLocation(IRBuilder<>::InsertPoint(
+        StateMachineBeginBB, StateMachineBeginBB->end()));
+
+    Value *Ident = KernelInitCB->getArgOperand(0);
+    Value *GTid = KernelInitCB;
+
+    Module &M = *Kernel->getParent();
+    FunctionCallee BarrierFn =
+        OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
+            M, OMPRTL___kmpc_barrier_simple_spmd);
+    CallInst::Create(BarrierFn, {Ident, GTid}, "", StateMachineBeginBB);
+
+    FunctionCallee KernelParallelFn =
+        OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
+            M, OMPRTL___kmpc_kernel_parallel);
+    Value *IsActiveWorker = CallInst::Create(
+        KernelParallelFn, {WorkFnAI}, "worker.is_active", StateMachineBeginBB);
+    Value *WorkFn = new LoadInst(VoidPtrTy, WorkFnAI, "worker.work_fn",
+                                 StateMachineBeginBB);
+
+    FunctionType *ParallelRegionFnTy = FunctionType::get(
+        Type::getVoidTy(Ctx), {Type::getInt16Ty(Ctx), Type::getInt32Ty(Ctx)},
+        false);
+    Value *WorkFnCast = BitCastInst::CreatePointerBitCastOrAddrSpaceCast(
+        WorkFn, ParallelRegionFnTy->getPointerTo(), "worker.work_fn.addr_cast",
+        StateMachineBeginBB);
+
+    Value *IsDone = ICmpInst::Create(ICmpInst::ICmp, llvm::CmpInst::ICMP_EQ,
+                                     WorkFn, Constant::getNullValue(VoidPtrTy),
+                                     "worker.is_done", StateMachineBeginBB);
+    BranchInst::Create(StateMachineFinishedBB, StateMachineIsActiveCheckBB,
+                       IsDone, StateMachineBeginBB);
+
+    BranchInst::Create(StateMachineIfCascadeCurrentBB,
+                       StateMachineDoneBarrierBB, IsActiveWorker,
+                       StateMachineIsActiveCheckBB);
+
+    Value *ZeroArg =
+        Constant::getNullValue(ParallelRegionFnTy->getParamType(0));
+
+    if (!ParallelRegions.empty()) {
+
+      for (int i = 0, e = ParallelRegions.size(); i < e; ++i) {
+        auto *ParallelRegion = ParallelRegions[i];
+        BasicBlock *PRExecuteBB = BasicBlock::Create(
+            Ctx, "worker_state_machine.parallel_region.execute", Kernel,
+            StateMachineEndParallelBB);
+        CallInst::Create(ParallelRegion, {ZeroArg, GTid}, "", PRExecuteBB);
+        BranchInst::Create(StateMachineEndParallelBB, PRExecuteBB);
+
+        BasicBlock *PRNextBB = BasicBlock::Create(
+            Ctx, "worker_state_machine.parallel_region.check", Kernel,
+            StateMachineEndParallelBB);
+
+        // Check if we need to compare the pointer at all or if we can just
+        // call the parallel region function.
+        Value *IsPR;
+        if (i + 1 < e || MayReachUnknownParallelRegion)
+          IsPR = ICmpInst::Create(ICmpInst::ICmp, llvm::CmpInst::ICMP_EQ,
+                                  WorkFnCast, ParallelRegion,
+                                  "worker.check_parallel_region",
+                                  StateMachineIfCascadeCurrentBB);
+        else
+          IsPR = ConstantInt::getTrue(Ctx);
+
+        BranchInst::Create(PRExecuteBB, PRNextBB, IsPR,
+                           StateMachineIfCascadeCurrentBB);
+        StateMachineIfCascadeCurrentBB = PRNextBB;
+      }
+    }
+
+    if (MayReachUnknownParallelRegion) {
+      StateMachineIfCascadeCurrentBB->setName(
+          "worker_state_machine.parallel_region.fallback.execute");
+      CallInst::Create(ParallelRegionFnTy, WorkFnCast, {ZeroArg, GTid}, "",
+                       StateMachineIfCascadeCurrentBB);
+    }
+    BranchInst::Create(StateMachineEndParallelBB,
+                       StateMachineIfCascadeCurrentBB);
+
+    CallInst::Create(OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
+                         M, OMPRTL___kmpc_kernel_end_parallel),
+                     {}, "", StateMachineEndParallelBB);
+    BranchInst::Create(StateMachineDoneBarrierBB, StateMachineEndParallelBB);
+
+    CallInst::Create(BarrierFn, {Ident, GTid}, "", StateMachineDoneBarrierBB);
+    BranchInst::Create(StateMachineBeginBB, StateMachineDoneBarrierBB);
+
+    return ChangeStatus::CHANGED;
+  }
+
+  /// Statistics are tracked as part of manifest for now.
+  void trackStatistics() const override {}
+
+  /// See AbstractAttribute::getAsStr()
+  const std::string getAsStr() const override {
+    return std::string("#PR: ") + std::to_string(ParallelRegions.size()) + (MayReachUnknownParallelRegion ? " + Unknown PR" : "");
+  }
+
+  /// Create an abstract attribute biew for the position \p IRP.
+  static AAKernelInfo &createForPosition(const IRPosition &IRP, Attributor &A);
+
+  /// See AbstractAttribute::getName()
+  const std::string getName() const override { return "AAKernelInfo"; }
+
+  /// See AbstractAttribute::getIdAddr()
+  const char *getIdAddr() const override { return &ID; }
+
+  /// This function should return true if the type of the \p AA is AAKernelInfo
+  static bool classof(const AbstractAttribute *AA) {
+    return (AA->getIdAddr() == &ID);
+  }
+
+  static const char ID;
+};
+
+struct AAKernelInfoFunction : AAKernelInfo {
+  AAKernelInfoFunction(const IRPosition &IRP, Attributor &A)
+      : AAKernelInfo(IRP, A) {}
+
+  void initialize(Attributor &A) override {
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    KernelInfoState StateBefore = getState();
+
+    // Callback to check a call instruction.
+    auto CheckCallInst = [&](Instruction &I) {
+      auto &CB = cast<CallBase>(I);
+      if (!CB.mayWriteToMemory())
+        return true;
+      if (isa<IntrinsicInst>(CB))
+        return true;
+      if (CB.hasFnAttr("no_openmp"))
+        return true;
+
+      Function *Callee = CB.getCalledFunction();
+      auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+      const auto &It = OMPInfoCache.RuntimeFunctionIDMap.find(Callee);
+      if (It == OMPInfoCache.RuntimeFunctionIDMap.end()) {
+        auto &CBAA = A.getAAFor<AAKernelInfo>(
+            *this, IRPosition::callsite_function(CB), DepClassTy::REQUIRED);
+        if (CBAA.getState().isValidState()) {
+          getState() ^= CBAA.getState();
+        } else {
+          MayReachUnknownParallelRegion = true;
+        }
+        return true;
+      }
+
+      const unsigned int WrapperFunctionArgNo = 6;
+      RuntimeFunction RF = It->getSecond();
+      switch (RF) {
+      case OMPRTL___kmpc_target_init:
+        if (KernelInitCB && KernelInitCB != &CB)
+          return false;
+        KernelInitCB = &CB;
+        return true;
+      case OMPRTL___kmpc_target_deinit:
+        if (KernelDeinitCB && KernelDeinitCB != &CB)
+          return false;
+        KernelDeinitCB = &CB;
+        return true;
+      case OMPRTL___kmpc_parallel_51:
+        if (auto *ParallelRegion = dyn_cast<Function>(
+                CB.getArgOperand(WrapperFunctionArgNo)->stripPointerCasts())) {
+          ParallelRegions.insert(ParallelRegion);
+          return true;
+        }
+        MayReachUnknownParallelRegion = true;
+        return true;
+      case OMPRTL___kmpc_omp_task:
+        MayReachUnknownParallelRegion = true;
+        return true;
+      default:
+        break;
+      }
+      return true;
+    };
+    if (!A.checkForAllCallLikeInstructions(CheckCallInst, *this))
+      return indicatePessimisticFixpoint();
+
+    return StateBefore == getState() ? ChangeStatus::UNCHANGED
+                                     : ChangeStatus::CHANGED;
+  }
+};
+
+struct AAKernelInfoCallSite : AAKernelInfo {
+  AAKernelInfoCallSite(const IRPosition &IRP, Attributor &A)
+      : AAKernelInfo(IRP, A) {}
+
+  /// See AbstractAttribute::initialize(...).
+  void initialize(Attributor &A) override {
+    AAKernelInfo::initialize(A);
+    Function *F = getAssociatedFunction();
+    if (!F || F->isDeclaration())
+      indicatePessimisticFixpoint();
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    // TODO: Once we have call site specific value information we can provide
+    //       call site specific liveness information and then it makes
+    //       sense to specialize attributes for call sites arguments instead of
+    //       redirecting requests to the callee argument.
+    Function *F = getAssociatedFunction();
+    const IRPosition &FnPos = IRPosition::function(*F);
+    auto &FnAA = A.getAAFor<AAKernelInfo>(*this, FnPos, DepClassTy::REQUIRED);
+    if (getState() == FnAA.getState())
+      return ChangeStatus::UNCHANGED;
+    getState() = FnAA.getState();
+    return ChangeStatus::CHANGED;
+  }
+};
+
 /// Used to map the values physically (in the IR) stored in an offload
 /// array, to a vector in memory.
 struct OffloadArray {
   /// Physical array (in the IR).
   AllocaInst *Array = nullptr;
   /// Mapped values.
   SmallVector<Value *, 8> StoredValues;
   /// Last stores made in the offload array.
bool run(bool IsModulePass) {
 
     bool Changed = false;
 
     LLVM_DEBUG(dbgs() << TAG << "Run on SCC with " << SCC.size()
                       << " functions in a slice with "
                       << OMPInfoCache.ModuleSlice.size() << " functions\n");
 
     if (IsModulePass) {
       Changed |= runAttributor();
 
       if (remarksEnabled())
         analysisGlobalization();
     } else {
       if (PrintICVValues)
         printICVs();
       if (PrintOpenMPKernels)
         printKernels();
 
-      Changed |= rewriteDeviceCodeStateMachine();
-
-      Changed |= runAttributor();
+      Changed |= runAttributor(IsModulePass);
 
       // Recollect uses, in case Attributor deleted any.
       OMPInfoCache.recollectUses();
 
       Changed |= deleteParallelRegions();
+      Changed |= rewriteDeviceCodeStateMachine();
+
       if (HideMemoryTransferLatency)
         Changed |= hideMemTransfersLatency();
       Changed |= deduplicateRuntimeCalls();
       if (EnableParallelRegionMerging) {
         if (mergeParallelRegions()) {
           deduplicateRuntimeCalls();
           Changed = true;
         }
private:
 
   /// OpenMP-specific information cache. Also Used for Attributor runs.
   OMPInformationCache &OMPInfoCache;
 
   /// Attributor instance.
   Attributor &A;
 
   /// Helper function to run Attributor on SCC.
-  bool runAttributor() {
+  bool runAttributor(bool IsModulePass) {
     if (SCC.empty())
       return false;
 
-    registerAAs();
+    registerAAs(IsModulePass);
 
     ChangeStatus Changed = A.run();
 
     LLVM_DEBUG(dbgs() << "[Attributor] Done with " << SCC.size()
                       << " functions, result: " << Changed << ".\n");
 
     return Changed == ChangeStatus::CHANGED;
   }
 
   /// Populate the Attributor with abstract attribute opportunities in the
   /// function.
-  void registerAAs() {
+  void registerAAs(bool IsModulePass) {
     if (SCC.empty())
       return;
 
     // 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];
for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
 
       GetterRFI.foreachUse(SCC, CreateAA);
     }
 
     for (auto &F : M) {
       if (!F.isDeclaration())
         A.getOrCreateAAFor<AAExecutionDomain>(IRPosition::function(F));
     }
+    if (IsModulePass) {
+      for (Function *Kernel : OMPInfoCache.Kernels)
+        A.getOrCreateAAFor<AAKernelInfo>(IRPosition::function(*Kernel));
+    }
   }
 };
 
 Kernel OpenMPOpt::getUniqueKernelFor(Function &F) {
   if (!OMPInfoCache.ModuleSlice.count(&F))
     return nullptr;
 
   // Use a scope to keep the lifetime of the CachedKernel short.
if (!KernelParallelUse)
       };
       emitRemarkOnFunction(F, "OpenMPParallelRegionInNonSPMD", Remark);
     }
 
     // If this ever hits, we should investigate.
     // TODO: Checking the number of uses is not a necessary restriction and
     // should be lifted.
     if (UnknownUse || NumDirectCalls != 1 ||
-        ToBeReplacedStateMachineUses.size() != 2) {
+        ToBeReplacedStateMachineUses.size() > 2) {
       {
         auto Remark = [&](OptimizationRemark OR) {
           return OR << "Parallel region is used in "
                     << (UnknownUse ? "unknown" : "unexpected")
                     << " ways; will not attempt to rewrite the state machine.";
         };
         emitRemarkOnFunction(F, "OpenMPParallelRegionInNonSPMD", Remark);
       }
ChangeStatus AAExecutionDomainFunction::updateImpl(Attributor &A) {
   return (NumSingleThreadedBBs == SingleThreadedBBs.size())
              ? ChangeStatus::UNCHANGED
              : ChangeStatus::CHANGED;
 }
 
 } // namespace
 
 const char AAICVTracker::ID = 0;
+const char AAKernelInfo::ID = 0;
 const char AAExecutionDomain::ID = 0;
 
 AAICVTracker &AAICVTracker::createForPosition(const IRPosition &IRP,
                                               Attributor &A) {
   AAICVTracker *AA = nullptr;
   switch (IRP.getPositionKind()) {
   case IRPosition::IRP_INVALID:
   case IRPosition::IRP_FLOAT:
AAExecutionDomain &AAExecutionDomain::createForPosition(const IRPosition &IRP,
   case IRPosition::IRP_CALL_SITE:
     llvm_unreachable(
         "AAExecutionDomain can only be created for function position!");
   case IRPosition::IRP_FUNCTION:
     AA = new (A.Allocator) AAExecutionDomainFunction(IRP, A);
     break;
   }
 
   return *AA;
 }
 
+AAKernelInfo &AAKernelInfo::createForPosition(const IRPosition &IRP,
+                                              Attributor &A) {
+  AAKernelInfo *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:
+    llvm_unreachable("KernelInfo can only be created for function position!");
+  case IRPosition::IRP_CALL_SITE:
+    AA = new (A.Allocator) AAKernelInfoCallSite(IRP, A);
+    break;
+  case IRPosition::IRP_FUNCTION:
+    AA = new (A.Allocator) AAKernelInfoFunction(IRP, A);
+    break;
+  }
+
+  return *AA;
+}
+
 PreservedAnalyses OpenMPOptPass::run(Module &M, ModuleAnalysisManager &AM) {
   if (!containsOpenMP(M, OMPInModule))
     return PreservedAnalyses::all();
 
   if (DisableOpenMPOptimizations)
     return PreservedAnalyses::all();
 
   // Look at every function definition in the Module.
   SmallVector<Function *, 16> SCC;
   for (Function &Fn : M)
     if (!Fn.isDeclaration())
       SCC.push_back(&Fn);
 
   if (SCC.empty())
     return PreservedAnalyses::all();
 
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
 
   AnalysisGetter AG(FAM);
 
   auto OREGetter = [&FAM](Function *F) -> OptimizationRemarkEmitter & {
     return FAM.getResult<OptimizationRemarkEmitterAnalysis>(*F);
   };
 
   BumpPtrAllocator Allocator;
   CallGraphUpdater CGUpdater;
 
   SetVector<Function *> Functions(SCC.begin(), SCC.end());
   OMPInformationCache InfoCache(M, AG, Allocator, /*CGSCC*/ Functions,
                                 OMPInModule.getKernels());
 
-  Attributor A(Functions, InfoCache, CGUpdater);
+  Attributor A(Functions, InfoCache, CGUpdater, nullptr, false);
 
   OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
   bool Changed = OMPOpt.run(true);
   if (Changed)
     return PreservedAnalyses::none();
 
   return PreservedAnalyses::all();
 }
 
 PreservedAnalyses OpenMPOptCGSCCPass::run(LazyCallGraph::SCC &C,
                                           CGSCCAnalysisManager &AM,
                                           LazyCallGraph &CG,
PreservedAnalyses OpenMPOptCGSCCPass::run(LazyCallGraph::SCC &C,
   BumpPtrAllocator Allocator;
   CallGraphUpdater CGUpdater;
   CGUpdater.initialize(CG, C, AM, UR);
 
   SetVector<Function *> Functions(SCC.begin(), SCC.end());
   OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG, Allocator,
                                 /*CGSCC*/ Functions, OMPInModule.getKernels());
 
-  Attributor A(Functions, InfoCache, CGUpdater);
+  Attributor A(Functions, InfoCache, CGUpdater, nullptr, false);
 
   OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
   bool Changed = OMPOpt.run(false);
   if (Changed)
     return PreservedAnalyses::none();
 
   return PreservedAnalyses::all();
 }
bool runOnSCC(CallGraphSCC &CGSCC) override {
 
     AnalysisGetter AG;
     SetVector<Function *> Functions(SCC.begin(), SCC.end());
     BumpPtrAllocator Allocator;
     OMPInformationCache InfoCache(
         *(Functions.back()->getParent()), AG, Allocator,
         /*CGSCC*/ Functions, OMPInModule.getKernels());
 
-    Attributor A(Functions, InfoCache, CGUpdater);
+    Attributor A(Functions, InfoCache, CGUpdater, nullptr, false);
 
     OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
     return OMPOpt.run(false);
   }
 
   bool doFinalization(CallGraph &CG) override { return CGUpdater.finalize(); }
 };