[OpenMPOpt] Expand SPMDization with guarding

[ggeorgakoudis]

Summary

This patch expands SPMDization by guarding code regions that should be executed only by the main thread. It defines an SPMD guarded execution mode where the main thread executes guarded regions and synchronizes with worker threads using simple barriers to ensure that values are updated. Given it is possible that the same code region executes both from SPMD or SPMDized guarded execution mode, the patch creates a non-guarded version of the region and a flag to guard or not that is managed by the runtime.

TODO: tests, depends on simplification callback fixes

[ggeorgakoudis]

Update for sensible access to __kmpc_is_spmd_guarded_exec_mode

[ggeorgakoudis]

Code cleanup

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

void CGOpenMPRuntimeGPU::emitNonSPMDKernel(const OMPExecutableDirective &D,
   emitTargetOutlinedFunctionHelper(D, ParentName, OutlinedFn, OutlinedFnID,
                                    IsOffloadEntry, CodeGen);
   IsInTTDRegion = false;
 }
 
 void CGOpenMPRuntimeGPU::emitKernelInit(CodeGenFunction &CGF,
                                         EntryFunctionState &EST, bool IsSPMD) {
   CGBuilderTy &Bld = CGF.Builder;
-  Bld.restoreIP(OMPBuilder.createTargetInit(Bld, IsSPMD, requiresFullRuntime()));
+  Bld.restoreIP(OMPBuilder.createTargetInit(
+      Bld, IsSPMD, /* IsSPMDGuarded */ false, requiresFullRuntime()));
   IsInTargetMasterThreadRegion = IsSPMD;
   if (!IsSPMD)
     emitGenericVarsProlog(CGF, EST.Loc);
 }
 
 void CGOpenMPRuntimeGPU::emitKernelDeinit(CodeGenFunction &CGF,
                                           EntryFunctionState &EST,
                                           bool IsSPMD) {

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

public:
   ///
   ///{
 
   /// Create a runtime call for kmpc_target_init
   ///
   /// \param Loc The insert and source location description.
   /// \param IsSPMD Flag to indicate if the kernel is an SPMD kernel or not.
   /// \param RequiresFullRuntime Indicate if a full device runtime is necessary.
-  InsertPointTy createTargetInit(const LocationDescription &Loc, bool IsSPMD, bool RequiresFullRuntime);
+  InsertPointTy createTargetInit(const LocationDescription &Loc, bool IsSPMD,
+                                 bool IsSPMDGuarded, bool RequiresFullRuntime);
 
   /// Create a runtime call for kmpc_target_deinit
   ///
   /// \param Loc The insert and source location description.
   /// \param IsSPMD Flag to indicate if the kernel is an SPMD kernel or not.
   /// \param RequiresFullRuntime Indicate if a full device runtime is necessary.
   void createTargetDeinit(const LocationDescription &Loc, bool IsSPMD, bool RequiresFullRuntime);
 

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

__OMP_RTL(__tgt_target_data_update_nowait_mapper, false, Void, IdentPtr, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr)
 __OMP_RTL(__tgt_mapper_num_components, false, Int64, VoidPtr)
 __OMP_RTL(__tgt_push_mapper_component, false, Void, VoidPtr, VoidPtr, VoidPtr,
           Int64, Int64, VoidPtr)
 __OMP_RTL(__kmpc_task_allow_completion_event, false, VoidPtr, IdentPtr,
           /* Int */ Int32, /* kmp_task_t */ VoidPtr)
 
 /// OpenMP Device runtime functions
-__OMP_RTL(__kmpc_target_init, false, Int32, IdentPtr, Int1, Int1, Int1)
+__OMP_RTL(__kmpc_target_init, false, Int32, IdentPtr, Int1, Int1, Int1, Int1)
 __OMP_RTL(__kmpc_target_deinit, false, Void, IdentPtr, Int1, Int1)
 __OMP_RTL(__kmpc_kernel_prepare_parallel, false, Void, VoidPtr)
 __OMP_RTL(__kmpc_parallel_51, false, Void, IdentPtr, Int32, Int32, Int32, Int32,
           VoidPtr, VoidPtr, VoidPtrPtr, SizeTy)
 __OMP_RTL(__kmpc_kernel_parallel, false, Int1, VoidPtrPtr)
 __OMP_RTL(__kmpc_kernel_end_parallel, false, Void, )
 __OMP_RTL(__kmpc_serialized_parallel, false, Void, IdentPtr, Int32)
 __OMP_RTL(__kmpc_end_serialized_parallel, false, Void, IdentPtr, Int32)

 
 __OMP_RTL(__kmpc_alloc_shared, false, VoidPtr, SizeTy)
 __OMP_RTL(__kmpc_free_shared, false, Void, VoidPtr, SizeTy)
 __OMP_RTL(__kmpc_begin_sharing_variables, false, Void, VoidPtrPtrPtr, SizeTy)
 __OMP_RTL(__kmpc_end_sharing_variables, false, Void, )
 __OMP_RTL(__kmpc_get_shared_variables, false, Void, VoidPtrPtrPtr)
 __OMP_RTL(__kmpc_parallel_level, false, Int8, )
 __OMP_RTL(__kmpc_is_spmd_exec_mode, false, Int8, )
+__OMP_RTL(__kmpc_is_spmd_guarded_exec_mode, false, Int8, )
 __OMP_RTL(__kmpc_barrier_simple_spmd, false, Void, IdentPtr, Int32)
 
 __OMP_RTL(__kmpc_warp_active_thread_mask, false, LanemaskTy,)
 __OMP_RTL(__kmpc_syncwarp, false, Void, LanemaskTy)
 
 __OMP_RTL(__kmpc_is_generic_main_thread_id, false, Int8, Int32)
 
 __OMP_RTL(__last, false, Void, )

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

CallInst *OpenMPIRBuilder::createCachedThreadPrivate(
 
   Function *Fn =
       getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_threadprivate_cached);
 
   return Builder.CreateCall(Fn, Args);
 }
 
 OpenMPIRBuilder::InsertPointTy
-OpenMPIRBuilder::createTargetInit(const LocationDescription &Loc, bool IsSPMD, bool RequiresFullRuntime) {
+OpenMPIRBuilder::createTargetInit(const LocationDescription &Loc, bool IsSPMD,
+                                  bool IsSPMDGuarded,
+                                  bool RequiresFullRuntime) {
   if (!updateToLocation(Loc))
     return Loc.IP;
 
   Constant *SrcLocStr = getOrCreateSrcLocStr(Loc);
   Value *Ident = getOrCreateIdent(SrcLocStr);
   ConstantInt *IsSPMDVal = ConstantInt::getBool(Int32->getContext(), IsSPMD);
+  ConstantInt *IsSPMDGuardedVal =
+      ConstantInt::getBool(Int32->getContext(), IsSPMDGuarded);
   ConstantInt *UseGenericStateMachine =
       ConstantInt::getBool(Int32->getContext(), !IsSPMD);
   ConstantInt *RequiresFullRuntimeVal = ConstantInt::getBool(Int32->getContext(), RequiresFullRuntime);
 
   Function *Fn = getOrCreateRuntimeFunctionPtr(
       omp::RuntimeFunction::OMPRTL___kmpc_target_init);
 
   CallInst *ThreadKind =
-      Builder.CreateCall(Fn, {Ident, IsSPMDVal, UseGenericStateMachine, RequiresFullRuntimeVal});
+      Builder.CreateCall(Fn, {Ident, IsSPMDVal, IsSPMDGuardedVal,
+                              UseGenericStateMachine, RequiresFullRuntimeVal});
 
   Value *ExecUserCode = Builder.CreateICmpEQ(
       ThreadKind, ConstantInt::get(ThreadKind->getType(), -1), "exec_user_code");
 
   // ThreadKind = __kmpc_target_init(...)
   // if (ThreadKind == -1)
   //   user_code
   // else

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/Attributor.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 
 using namespace llvm;
 using namespace omp;
 
 #define DEBUG_TYPE "openmp-opt"
 
 static cl::opt<bool> DisableOpenMPOptimizations(
BooleanStateWithPtrSetVector<Function, /* InsertInvalidates */ false>
       ReachedKnownParallelRegions;
 
   /// State to track what parallel region we might reach.
   BooleanStateWithPtrSetVector<CallBase> ReachedUnknownParallelRegions;
 
   /// State to track if we are in SPMD-mode, assumed or know, and why we decided
   /// we cannot be. If it is assumed, then RequiresFullRuntime should also be
   /// false.
-  BooleanStateWithPtrSetVector<Instruction> SPMDCompatibilityTracker;
+  BooleanStateWithPtrSetVector<Instruction, false> SPMDCompatibilityTracker;
 
   /// 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;

 };
 
 /// The function kernel info abstract attribute, basically, what can we say
 /// about a function with regards to the KernelInfoState.
 struct AAKernelInfoFunction : AAKernelInfo {
   AAKernelInfoFunction(const IRPosition &IRP, Attributor &A)
       : AAKernelInfo(IRP, A) {}
 
+  SmallPtrSet<Instruction *, 4> GuardedInstructions;
+
+  SmallPtrSetImpl<Instruction *> &getGuardedInstructions() {
+    return GuardedInstructions;
+  }
+
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     // This is a high-level transform that might change the constant arguments
     // 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();
Attributor::SimplifictionCallbackTy IsSPMDModeSimplifyCB =
       } else {
         UsedAssumedInformation = false;
       }
       auto *Val = ConstantInt::getBool(IRP.getAnchorValue().getContext(),
                                        SPMDCompatibilityTracker.isAssumed());
       return Val;
     };
 
+    Attributor::SimplifictionCallbackTy IsSPMDGuardedModeSimplifyCB =
+        [&](const IRPosition &IRP, const AbstractAttribute *AA,
+            bool &UsedAssumedInformation) -> Optional<Value *> {
+      // IRP represents the "SPMDCompatibilityTracker" argument of an
+      // __kmpc_target_init or
+      // __kmpc_target_deinit call. We will answer this one with the internal
+      // state.
+      if (!SPMDCompatibilityTracker.isValidState())
+        return nullptr;
+      if (!SPMDCompatibilityTracker.isAtFixpoint()) {
+        if (AA)
+          A.recordDependence(*this, *AA, DepClassTy::OPTIONAL);
+        UsedAssumedInformation = true;
+      } else {
+        UsedAssumedInformation = false;
+      }
+
+      auto *Val = ConstantInt::getBool(IRP.getAnchorValue().getContext(),
+                                       (SPMDCompatibilityTracker.isAssumed() &&
+                                        !SPMDCompatibilityTracker.empty()));
+      return Val;
+    };
+
     Attributor::SimplifictionCallbackTy IsGenericModeSimplifyCB =
         [&](const IRPosition &IRP, const AbstractAttribute *AA,
             bool &UsedAssumedInformation) -> Optional<Value *> {
       // IRP represents the "RequiresFullRuntime" argument of an
       // __kmpc_target_init or __kmpc_target_deinit call. We will answer this
       // one with the internal state of the SPMDCompatibilityTracker, so if
       // generic then true, if SPMD then false.
       if (!SPMDCompatibilityTracker.isValidState())
         return nullptr;
       if (!SPMDCompatibilityTracker.isAtFixpoint()) {
         if (AA)
           A.recordDependence(*this, *AA, DepClassTy::OPTIONAL);
         UsedAssumedInformation = true;
       } else {
         UsedAssumedInformation = false;
       }
       auto *Val = ConstantInt::getBool(IRP.getAnchorValue().getContext(),
                                        !SPMDCompatibilityTracker.isAssumed());
       return Val;
     };
 
     constexpr const int InitIsSPMDArgNo = 1;
+    constexpr const int InitIsSPMDGuardedArgNo = 2;
+    constexpr const int InitUseStateMachineArgNo = 3;
+    constexpr const int InitRequiresFullRuntimeArgNo = 4;
     constexpr const int DeinitIsSPMDArgNo = 1;
-    constexpr const int InitUseStateMachineArgNo = 2;
-    constexpr const int InitRequiresFullRuntimeArgNo = 3;
     constexpr const int DeinitRequiresFullRuntimeArgNo = 2;
     A.registerSimplificationCallback(
         IRPosition::callsite_argument(*KernelInitCB, InitUseStateMachineArgNo),
         StateMachineSimplifyCB);
     A.registerSimplificationCallback(
         IRPosition::callsite_argument(*KernelInitCB, InitIsSPMDArgNo),
         IsSPMDModeSimplifyCB);
     A.registerSimplificationCallback(
+        IRPosition::callsite_argument(*KernelInitCB, InitIsSPMDGuardedArgNo),
+        IsSPMDGuardedModeSimplifyCB);
+    A.registerSimplificationCallback(
         IRPosition::callsite_argument(*KernelDeinitCB, DeinitIsSPMDArgNo),
         IsSPMDModeSimplifyCB);
     A.registerSimplificationCallback(
         IRPosition::callsite_argument(*KernelInitCB,
                                       InitRequiresFullRuntimeArgNo),
         IsGenericModeSimplifyCB);
     A.registerSimplificationCallback(
         IRPosition::callsite_argument(*KernelDeinitCB,
if (!SPMDCompatibilityTracker.isAssumed()) {
 
         LLVM_DEBUG(dbgs() << TAG << "SPMD-incompatible side-effect: "
                           << *NonCompatibleI << "\n");
       }
 
       return false;
     }
 
+    auto CreateGuardedRegion = [&](Instruction *RegionStartI,
+                                   Instruction *RegionEndI) {
+      LoopInfo *LI = nullptr;
+      DominatorTree *DT = nullptr;
+      MemorySSAUpdater *MSU = nullptr;
+      using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+
+      BasicBlock *ParentBB = RegionStartI->getParent();
+      Function *Fn = ParentBB->getParent();
+      Module &M = *Fn->getParent();
+
+      // Create all the blocks and logic.
+      // ParentBB:
+      //    IsSPMDGuarded = __kmpc_is_spmd_guarded_mode()
+      //    if (IsSPMDGuarded)
+      //        goto RegionCheckTidBB
+      // RegionNotguardedBB:
+      //    <execute instructions not guarded>
+      //    goto RegionExitBB
+      // RegionCheckTidBB:
+      //    Tid = __kmpc_hardware_thread_id()
+      //    if (Tid != 0)
+      //        goto RegionBarrierBB
+      // RegionStartBB:
+      //    <execute instructions guarded>
+      //    goto RegionEndBB
+      // RegionEndBB:
+      //    <store escaping values to shared mem>
+      //    goto RegionBarrierBB
+      //  RegionBarrierBB:
+      //    __kmpc_simple_barrier_spmd()
+      //    // second barrier is omitted if lacking escaping values.
+      //    <load escaping values from shared mem>
+      //    __kmpc_simple_barrier_spmd()
+      //    goto RegionExitBB
+      // RegionExitBB:
+      //    <execute rest of instructions>
+
+      BasicBlock *RegionEndBB = SplitBlock(ParentBB, RegionEndI->getNextNode(),
+                                           DT, LI, MSU, "region.guarded.end");
+      BasicBlock *RegionBarrierBB =
+          SplitBlock(RegionEndBB, &*RegionEndBB->getFirstInsertionPt(), DT, LI,
+                     MSU, "region.barrier");
+      BasicBlock *RegionExitBB =
+          SplitBlock(RegionBarrierBB, &*RegionBarrierBB->getFirstInsertionPt(),
+                     DT, LI, MSU, "region.exit");
+      BasicBlock *RegionStartBB =
+          SplitBlock(ParentBB, RegionStartI, DT, LI, MSU, "region.guarded");
+
+      // Create a clone that contains an non-guarded version for parallel
+      // execution.
+      ValueToValueMapTy VMap;
+      BasicBlock *RegionNotguardedBB =
+          CloneBasicBlock(RegionStartBB, VMap, ".not");
+      RegionNotguardedBB->insertInto(Fn, RegionStartBB);
+      RegionNotguardedBB->getTerminator()->setSuccessor(0, RegionExitBB);
+
+      assert(ParentBB->getUniqueSuccessor() == RegionStartBB &&
+             "Expected a different CFG");
+
+      BasicBlock *RegionCheckTidBB = SplitBlock(
+          ParentBB, ParentBB->getTerminator(), DT, LI, MSU, "region.check.tid");
+
+      // Register basic blocks with the Attributor.
+      A.registerManifestAddedBasicBlock(*RegionEndBB);
+      A.registerManifestAddedBasicBlock(*RegionBarrierBB);
+      A.registerManifestAddedBasicBlock(*RegionExitBB);
+      A.registerManifestAddedBasicBlock(*RegionStartBB);
+      A.registerManifestAddedBasicBlock(*RegionCheckTidBB);
+      A.registerManifestAddedBasicBlock(*RegionNotguardedBB);
+
+      bool HasBroadcastValues = false;
+      // Find escaping outputs from the guarded region to outside users and
+      // broadcast their values to them.
+      for (Instruction &I : *RegionStartBB) {
+        SmallPtrSet<Instruction *, 4> OutsideUsers;
+        for (User *Usr : I.users()) {
+          Instruction &UsrI = *cast<Instruction>(Usr);
+          if (UsrI.getParent() != RegionStartBB) {
+            dbgs() << "For I " << I << " in BB " << I.getParent()->getName()
+                   << " found outside user UsrI " << UsrI << " in BB "
+                   << UsrI.getParent()->getName() << "\n";
+            OutsideUsers.insert(&UsrI);
+          }
+        }
+
+        if (OutsideUsers.empty())
+          continue;
+
+        HasBroadcastValues = true;
+
+        // Emit a global variable in shared memory to store the broadcasted
+        // value.
+        auto *SharedMem = new GlobalVariable(
+            M, I.getType(), /* IsConstant */ false,
+            GlobalValue::InternalLinkage, UndefValue::get(I.getType()),
+            I.getName() + ".guarded.output.alloc", nullptr,
+            GlobalValue::NotThreadLocal,
+            static_cast<unsigned>(AddressSpace::Shared));
+
+        // Emit a store instruction to update the value.
+        new StoreInst(&I, SharedMem, RegionEndBB->getTerminator());
+
+        LoadInst *LoadI = new LoadInst(I.getType(), SharedMem,
+                                       I.getName() + ".guarded.output.load",
+                                       RegionBarrierBB->getTerminator());
+
+        PHINode *PN = PHINode::Create(I.getType(), 2, ".phi.guarded",
+                                      &*RegionExitBB->getFirstInsertionPt());
+        PN->addIncoming(VMap[&I], RegionNotguardedBB);
+        PN->addIncoming(LoadI, RegionBarrierBB);
+        // Emit a load instruction and replace uses of the output value.
+        for (Instruction *UsrI : OutsideUsers) {
+          dbgs() << "PN " << *PN << " in UsrI " << *UsrI << " in BB "
+                 << UsrI->getParent()->getName() << "\n";
+          assert(UsrI->getParent() == RegionExitBB &&
+                 "Expected escaping users in exit region");
+          UsrI->replaceUsesOfWith(&I, PN);
+        }
+      }
+
+      auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+
+      // Add check for parallel level in ParentBB.
+      const DebugLoc DL = ParentBB->getTerminator()->getDebugLoc();
+      ParentBB->getTerminator()->eraseFromParent();
+      OpenMPIRBuilder::LocationDescription Loc(
+          InsertPointTy(ParentBB, ParentBB->end()), DL);
+      OMPInfoCache.OMPBuilder.updateToLocation(Loc);
+      auto *SrcLocStr = OMPInfoCache.OMPBuilder.getOrCreateSrcLocStr(Loc);
+      Value *Ident = OMPInfoCache.OMPBuilder.getOrCreateIdent(SrcLocStr);
+      FunctionCallee HardwareTidFn =
+          OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
+              M, OMPRTL___kmpc_get_hardware_thread_id_in_block);
+      // There's got to be a better way to do that...
+      FunctionCallee IsSPMDGuardedFn =
+          M.getFunction("__kmpc_is_spmd_guarded_exec_mode.internalized");
+      // OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
+      //    M, OMPRTL___kmpc_is_spmd_guarded_exec_mode);
+      Value *Tid =
+          OMPInfoCache.OMPBuilder.Builder.CreateCall(HardwareTidFn, {});
+      Value *IsSPMDGuarded =
+          OMPInfoCache.OMPBuilder.Builder.CreateCall(IsSPMDGuardedFn, {});
+      Value *IsSPMDGuardedCheck =
+          OMPInfoCache.OMPBuilder.Builder.CreateIsNull(IsSPMDGuarded);
+      OMPInfoCache.OMPBuilder.Builder
+          .CreateCondBr(IsSPMDGuardedCheck, RegionNotguardedBB,
+                        RegionCheckTidBB)
+          ->setDebugLoc(DL);
+
+      // Add check for Tid in RegionCheckTidBB
+      RegionCheckTidBB->getTerminator()->eraseFromParent();
+      OpenMPIRBuilder::LocationDescription LocRegionCheckTid(
+          InsertPointTy(RegionCheckTidBB, RegionCheckTidBB->end()), DL);
+      OMPInfoCache.OMPBuilder.updateToLocation(LocRegionCheckTid);
+      Value *TidCheck = OMPInfoCache.OMPBuilder.Builder.CreateIsNull(Tid);
+      OMPInfoCache.OMPBuilder.Builder
+          .CreateCondBr(TidCheck, RegionStartBB, RegionBarrierBB)
+          ->setDebugLoc(DL);
+
+      // First barrier for synchronization, ensures main thread has updated
+      // values.
+      FunctionCallee BarrierFn =
+          OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
+              M, OMPRTL___kmpc_barrier_simple_spmd);
+      OMPInfoCache.OMPBuilder.updateToLocation(InsertPointTy(
+          RegionBarrierBB, RegionBarrierBB->getFirstInsertionPt()));
+      OMPInfoCache.OMPBuilder.Builder.CreateCall(BarrierFn, {Ident, Tid})
+          ->setDebugLoc(DL);
+
+      // Second barrier ensures workers have read broadcast values.
+      if (HasBroadcastValues)
+        CallInst::Create(BarrierFn, {Ident, Tid}, "",
+                         RegionBarrierBB->getTerminator())
+            ->setDebugLoc(DL);
+    };
+
+    // SmallPtrSet<BasicBlock *, 4> GuardedBasicBlocks;
+    SmallVector<std::pair<Instruction *, Instruction *>, 4> GuardedRegions;
+
+    for (Instruction *GuardedI : SPMDCompatibilityTracker) {
+      BasicBlock *BB = GuardedI->getParent();
+      auto *CalleeAA = A.lookupAAFor<AAKernelInfo>(
+          IRPosition::function(*GuardedI->getFunction()), nullptr,
+          DepClassTy::NONE);
+      assert(CalleeAA != nullptr && "Expected Callee AAKernelInfo");
+      auto &CalleeAAFunction = *cast<AAKernelInfoFunction>(CalleeAA);
+      // Continue if instruction is already guarded.
+      if (CalleeAAFunction.getGuardedInstructions().contains(GuardedI))
+        continue;
+
+      Instruction *GuardedRegionStart = nullptr, *GuardedRegionEnd = nullptr;
+      for (Instruction &I : *BB) {
+        // If instruction I needs to be guarded update the guarded region
+        // bounds.
+        if (SPMDCompatibilityTracker.contains(&I)) {
+          CalleeAAFunction.getGuardedInstructions().insert(&I);
+          if (GuardedRegionStart)
+            GuardedRegionEnd = &I;
+          else
+            GuardedRegionStart = GuardedRegionEnd = &I;
+
+          continue;
+        }
+
+        // Instruction I does not need guarding, store
+        // any region found and reset bounds.
+        if (GuardedRegionStart) {
+          GuardedRegions.push_back(
+              std::make_pair(GuardedRegionStart, GuardedRegionEnd));
+          GuardedRegionStart = nullptr;
+          GuardedRegionEnd = nullptr;
+        }
+      }
+    }
+
+    for (auto &GR : GuardedRegions)
+      CreateGuardedRegion(GR.first, GR.second);
+
     // Adjust the global exec mode flag that tells the runtime what mode this
     // kernel is executed in.
     Function *Kernel = getAnchorScope();
     GlobalVariable *ExecMode = Kernel->getParent()->getGlobalVariable(
         (Kernel->getName() + "_exec_mode").str());
     assert(ExecMode && "Kernel without exec mode?");
     assert(ExecMode->getInitializer() &&
            ExecMode->getInitializer()->isOneValue() &&
            "Initially non-SPMD kernel has SPMD exec mode!");
 
     // Set the global exec mode flag to indicate SPMD-Generic mode.
     constexpr int SPMDGeneric = 2;
     if (!ExecMode->getInitializer()->isZeroValue())
       ExecMode->setInitializer(
           ConstantInt::get(ExecMode->getInitializer()->getType(), SPMDGeneric));
 
     // Next rewrite the init and deinit calls to indicate we use SPMD-mode now.
     const int InitIsSPMDArgNo = 1;
+    const int InitIsSPMDGuardedArgNo = 2;
+    const int InitUseStateMachineArgNo = 3;
+    const int InitRequiresFullRuntimeArgNo = 4;
     const int DeinitIsSPMDArgNo = 1;
-    const int InitUseStateMachineArgNo = 2;
-    const int InitRequiresFullRuntimeArgNo = 3;
     const int DeinitRequiresFullRuntimeArgNo = 2;
 
     auto &Ctx = getAnchorValue().getContext();
     A.changeUseAfterManifest(KernelInitCB->getArgOperandUse(InitIsSPMDArgNo),
                              *ConstantInt::getBool(Ctx, 1));
     A.changeUseAfterManifest(
+        KernelInitCB->getArgOperandUse(InitIsSPMDGuardedArgNo),
+        *ConstantInt::getBool(Ctx, 1));
+    A.changeUseAfterManifest(
         KernelInitCB->getArgOperandUse(InitUseStateMachineArgNo),
         *ConstantInt::getBool(Ctx, 0));
     A.changeUseAfterManifest(
         KernelDeinitCB->getArgOperandUse(DeinitIsSPMDArgNo),
         *ConstantInt::getBool(Ctx, 1));
     A.changeUseAfterManifest(
         KernelInitCB->getArgOperandUse(InitRequiresFullRuntimeArgNo),
         *ConstantInt::getBool(Ctx, 0));
bool changeToSPMDMode(Attributor &A) {
     return true;
   };
 
   ChangeStatus buildCustomStateMachine(Attributor &A) {
     assert(ReachedKnownParallelRegions.isValidState() &&
            "Custom state machine with invalid parallel region states?");
 
     const int InitIsSPMDArgNo = 1;
-    const int InitUseStateMachineArgNo = 2;
+    const int InitUseStateMachineArgNo = 3;
 
     // Check if the current configuration is non-SPMD and generic state machine.
     // If we already have SPMD mode or a custom state machine we do not need to
     // go any further. If it is anything but a constant something is weird and
     // we give up.
     ConstantInt *UseStateMachine = dyn_cast<ConstantInt>(
         KernelInitCB->getArgOperand(InitUseStateMachineArgNo));
     ConstantInt *IsSPMD =
auto CheckRWInst = [&](Instruction &I) {
       if (!I.mayWriteToMemory())
         return true;
       if (auto *SI = dyn_cast<StoreInst>(&I)) {
         SmallVector<const Value *> Objects;
         getUnderlyingObjects(SI->getPointerOperand(), Objects);
         if (llvm::all_of(Objects,
                          [](const Value *Obj) { return isa<AllocaInst>(Obj); }))
           return true;
+        // Check for AAHeapToStack moved objects to avoid guarding.
+        auto *HS = A.lookupAAFor<AAHeapToStack>(
+            IRPosition::function(*I.getFunction()), this, DepClassTy::OPTIONAL);
+        if (HS)
+          if (llvm::all_of(Objects, [HS](const Value *Obj) {
+                auto *CB = dyn_cast<CallBase>(Obj);
+                if (!CB)
+                  return false;
+                return HS->isAssumedHeapToStack(*CB);
+              })) {
+            return true;
           }
-      // For now we give up on everything but stores.
+      }
+
+      // Insert instruction that needs guarding.
       SPMDCompatibilityTracker.insert(&I);
       return true;
     };
 
     bool UsedAssumedInformationInCheckRWInst = false;
     if (!SPMDCompatibilityTracker.isAtFixpoint())
       if (!A.checkForAllReadWriteInstructions(
               CheckRWInst, *this, UsedAssumedInformationInCheckRWInst))
case OMPRTL___kmpc_parallel_51:
       // pointer and record it. In the off chance it doesn't we assume the
       // worst.
       ReachedUnknownParallelRegions.insert(&CB);
       break;
     case OMPRTL___kmpc_omp_task:
       // We do not look into tasks right now, just give up.
       SPMDCompatibilityTracker.insert(&CB);
       ReachedUnknownParallelRegions.insert(&CB);
-      break;
+      indicatePessimisticFixpoint();
+      return;
     case OMPRTL___kmpc_alloc_shared:
     case OMPRTL___kmpc_free_shared:
       // Return without setting a fixpoint, to be resolved in updateImpl.
       return;
     default:
       // Unknown OpenMP runtime calls cannot be executed in SPMD-mode,
       // generally.
       SPMDCompatibilityTracker.insert(&CB);
-      break;
+      indicatePessimisticFixpoint();
+      return;
     }
     // All other OpenMP runtime calls will not reach parallel regions so they
     // can be safely ignored for now. Since it is a known OpenMP runtime call we
     // have now modeled all effects and there is no need for any update.
     indicateOptimisticFixpoint();
   }
 
   ChangeStatus updateImpl(Attributor &A) override {

openmp/libomptarget/deviceRTLs/common/include/target.h

 ///{
 
 /// Initialization
 ///
 /// Must be called by all threads.
 ///
 /// \param Ident               Source location identification, can be NULL.
 ///
-int32_t __kmpc_target_init(ident_t *Ident, bool IsSPMD,
+int32_t __kmpc_target_init(ident_t *Ident, bool IsSPMD, bool IsSPMDGuarded,
                            bool UseGenericStateMachine,
                            bool RequiresFullRuntime);
 
 /// De-Initialization
 ///
 /// Must be called by the main thread in generic mode, can be called by all
 /// threads. Must be called by all threads in SPMD mode.
 ///

openmp/libomptarget/deviceRTLs/common/src/omptarget.cu

static void __kmpc_generic_kernel_deinit() {
   // Enqueue omp state object for use by another team.
   int slot = usedSlotIdx;
   omptarget_nvptx_device_State[slot].Enqueue(
       omptarget_nvptx_threadPrivateContext);
   // Done with work.  Kill the workers.
   omptarget_nvptx_workFn = 0;
 }
 
-static void __kmpc_spmd_kernel_init(bool RequiresFullRuntime) {
+static void __kmpc_spmd_kernel_init(bool IsSPMDGuarded, bool RequiresFullRuntime) {
   PRINT0(LD_IO, "call to __kmpc_spmd_kernel_init\n");
 
-  setExecutionParameters(Spmd, RequiresFullRuntime ? RuntimeInitialized
+  setExecutionParameters(IsSPMDGuarded ? SpmdGuarded : Spmd,
+                         RequiresFullRuntime ? RuntimeInitialized
                                              : RuntimeUninitialized);
   int threadId = __kmpc_get_hardware_thread_id_in_block();
   if (threadId == 0) {
     usedSlotIdx = __kmpc_impl_smid() % MAX_SM;
   }
 
   if (GetLaneId() == 0) {
     parallelLevel[GetWarpId()] =
         1 + (__kmpc_get_hardware_num_threads_in_block() > 1
if (threadId == 0) {
     int slot = usedSlotIdx;
     omptarget_nvptx_device_State[slot].Enqueue(
         omptarget_nvptx_threadPrivateContext);
   }
 }
 
 // Return true if the current target region is executed in SPMD mode.
 EXTERN int8_t __kmpc_is_spmd_exec_mode() {
-  return (execution_param & ModeMask) == Spmd;
+  return ((execution_param & ModeMask) == Spmd ||
+          (execution_param & ModeMask) == SpmdGuarded);
+}
+
+EXTERN __attribute__((used,retain)) int8_t __kmpc_is_spmd_guarded_exec_mode() {
+   return ((execution_param & ModeMask) == SpmdGuarded);
 }
 
 EXTERN int8_t __kmpc_is_generic_main_thread(kmp_int32 Tid) {
   return !__kmpc_is_spmd_exec_mode() && __kmpc_is_generic_main_thread_id(Tid);
 }
 
 NOINLINE EXTERN int8_t __kmpc_is_generic_main_thread_id(kmp_int32 Tid) {
   return GetMasterThreadID() == Tid;
do {
     }
 
     __kmpc_barrier_simple_spmd(Ident, TId);
 
   } while (true);
 }
 
 EXTERN
-int32_t __kmpc_target_init(ident_t *Ident, bool IsSPMD,
+int32_t __kmpc_target_init(ident_t *Ident, bool IsSPMD, bool IsSPMDGuarded,
                            bool UseGenericStateMachine,
                            bool RequiresFullRuntime) {
   int TId = __kmpc_get_hardware_thread_id_in_block();
   if (IsSPMD)
-    __kmpc_spmd_kernel_init(RequiresFullRuntime);
+    __kmpc_spmd_kernel_init(IsSPMDGuarded, RequiresFullRuntime);
   else
     __kmpc_generic_kernel_init();
 
    if (IsSPMD) {
     __kmpc_barrier_simple_spmd(Ident, TId);
      return -1;
    }
 

openmp/libomptarget/deviceRTLs/common/src/parallel.cu

EXTERN void __kmpc_parallel_51(kmp_Ident *ident, kmp_int32 global_tid,
   if (!if_expr || InParallelRegion) {
     __kmpc_serialized_parallel(ident, global_tid);
     __kmp_invoke_microtask(global_tid, 0, fn, args, nargs);
     __kmpc_end_serialized_parallel(ident, global_tid);
     return;
   }
 
   if (__kmpc_is_spmd_exec_mode()) {
+    // Store spmd_guarded status to check after the parallel region executes.
+    int is_spmd_guarded = __kmpc_is_spmd_guarded_exec_mode();
+    if (is_spmd_guarded) {
+      // No barrier is need on entry since this will be called only from non-guarded
+      // SPMD execution.
+
+      // Disable SPMD guarding for the parallel region. Runtime suport is not needed
+      // by construction of SPMD guarded regions, so simple assignment to Spmd is
+      // enough. Also, a preceding barrier is unnecessary since all threads must be
+      // in non-guarded context when reaching this point.
+      if (__kmpc_get_hardware_thread_id_in_block() == 0)
+        execution_param = Spmd;
+
+      // Barrier to ensure all threads are updated to Spmd.
+      __kmpc_barrier_simple_spmd(ident, 0);
+    }
+
     __kmp_invoke_microtask(global_tid, 0, fn, args, nargs);
+
+    if (is_spmd_guarded) {
+      // Re-enable SPMD guarding. Runtime support is not needed by construction.
+      // Barrier to ensure all threads have finished Spmd execution before
+      // re-enabling guarding.
+      __kmpc_barrier_simple_spmd(ident, 0);
+      if (__kmpc_get_hardware_thread_id_in_block() == 0)
+        execution_param = SpmdGuarded;
+
+      // Barrier to ensure all threads are updated to SpmdGuarded.
+      __kmpc_barrier_simple_spmd(ident, 0);
+    }
     return;
   }
 
   // Handle the num_threads clause.
   if (num_threads != -1)
     __kmpc_push_num_threads(ident, global_tid, num_threads);
 
   __kmpc_kernel_prepare_parallel((void *)wrapper_fn);

openmp/libomptarget/deviceRTLs/common/support.h

 //===--------- support.h - OpenMP GPU support functions ---------- CUDA -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Wrapper to some functions natively supported by the GPU.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef OMPTARGET_SUPPORT_H
 #define OMPTARGET_SUPPORT_H
 
 #include "interface.h"
 #include "target_impl.h"
 
 ////////////////////////////////////////////////////////////////////////////////
 // Execution Parameters
 ////////////////////////////////////////////////////////////////////////////////
 enum ExecutionMode {
   Spmd = 0x00u,
   Generic = 0x01u,
-  ModeMask = 0x01u,
+  SpmdGuarded = 0x02u,
+  ModeMask = 0x03u,
 };
 
 enum RuntimeMode {
   RuntimeInitialized = 0x00u,
-  RuntimeUninitialized = 0x02u,
-  RuntimeMask = 0x02u,
+  RuntimeUninitialized = 0x04u,
+  RuntimeMask = 0x04u,
 };
 
 void setExecutionParameters(ExecutionMode EMode, RuntimeMode RMode);
 bool isGenericMode();
 bool isRuntimeUninitialized();
 bool isRuntimeInitialized();
 
 ////////////////////////////////////////////////////////////////////////////////

openmp/libomptarget/deviceRTLs/interface.h

 typedef enum omp_proc_bind_t {
   omp_proc_bind_false = 0,
   omp_proc_bind_true = 1,
   omp_proc_bind_master = 2,
   omp_proc_bind_close = 3,
   omp_proc_bind_spread = 4
 } omp_proc_bind_t;
 
 EXTERN double omp_get_wtick(void);
 EXTERN double omp_get_wtime(void);
 
 EXTERN void omp_set_num_threads(int num);
 EXTERN int omp_get_num_threads(void);
 EXTERN int omp_get_max_threads(void);
 EXTERN int omp_get_thread_limit(void);
 EXTERN int omp_get_thread_num(void);
 EXTERN int omp_get_num_procs(void);
 EXTERN int omp_in_parallel(void);
 EXTERN int omp_in_final(void);
 EXTERN void omp_set_dynamic(int flag);
 EXTERN int omp_get_dynamic(void);
 EXTERN void omp_set_nested(int flag);
 EXTERN int omp_get_nested(void);
 EXTERN void omp_set_max_active_levels(int level);
 EXTERN int omp_get_max_active_levels(void);
 EXTERN int omp_get_level(void);
 EXTERN int omp_get_active_level(void);
 EXTERN int omp_get_ancestor_thread_num(int level);
 EXTERN int omp_get_team_size(int level);
 
 EXTERN void omp_init_lock(omp_lock_t *lock);
 EXTERN void omp_init_nest_lock(omp_nest_lock_t *lock);
 EXTERN void omp_destroy_lock(omp_lock_t *lock);
 EXTERN void omp_destroy_nest_lock(omp_nest_lock_t *lock);
 EXTERN void omp_set_lock(omp_lock_t *lock);
 EXTERN void omp_set_nest_lock(omp_nest_lock_t *lock);

 // cancel
 EXTERN int32_t __kmpc_cancellationpoint(kmp_Ident *loc, int32_t global_tid,
                                         int32_t cancelVal);
 EXTERN int32_t __kmpc_cancel(kmp_Ident *loc, int32_t global_tid,
                              int32_t cancelVal);
 
 // non standard
 EXTERN int32_t __kmpc_target_init(ident_t *Ident, bool IsSPMD,
+                                  bool IsSPMDGuarded,
                                   bool UseGenericStateMachine,
                                   bool RequiresFullRuntime);
 EXTERN void __kmpc_target_deinit(ident_t *Ident, bool IsSPMD,
                            bool RequiresFullRuntime);
 EXTERN void __kmpc_kernel_prepare_parallel(void *WorkFn);
 EXTERN bool __kmpc_kernel_parallel(void **WorkFn);
 EXTERN void __kmpc_kernel_end_parallel();
 
 EXTERN void __kmpc_data_sharing_init_stack();
 EXTERN void __kmpc_begin_sharing_variables(void ***GlobalArgs, size_t nArgs);

 EXTERN void __kmpc_parallel_51(ident_t *ident, kmp_int32 global_tid,
                                kmp_int32 if_expr, kmp_int32 num_threads,
                                int proc_bind, void *fn, void *wrapper_fn,
                                void **args, size_t nargs);
 
 // SPMD execution mode interrogation function.
 EXTERN int8_t __kmpc_is_spmd_exec_mode();
 
+EXTERN int8_t __kmpc_is_spmd_guarded_exec_mode();
+
 /// Return true if the hardware thread id \p Tid represents the OpenMP main
 /// thread in generic mode outside of a parallel region.
 EXTERN int8_t __kmpc_is_generic_main_thread(kmp_int32 Tid);
 
 /// Return true if the hardware thread id \p Tid represents the OpenMP main
 /// thread in generic mode.
 EXTERN int8_t __kmpc_is_generic_main_thread_id(kmp_int32 Tid);