[OpenMP] Introduce kernel environment

[tianshilei1992]

Summary

This patch introduces per kernel environment. Previously, flags such as execution mode are set through global variables with name like __kernel_name_exec_mode. They are accessible on the host by reading the corresponding global variable, but not from the device. Besides, some assumptions, such as no nested parallelism, are not per kernel basis, preventing us applying per kernel optimization in the device runtime.

This is a combination and refinement of patch series D116908, D116909, and D116910.

Depend on D155886.

[arsenm]

Description not clear

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3912–3914

ConstantExpr::getAddrSpaceCast

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3511–3518

Optimizations should be a separate change

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3511–3518

That's gonna break the tests because the interfaces between compiler and runtime are changed.

[jdoerfert]

Link the old review too. Pre-commit, w/o review, the file rename and include changes.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

319

Do we need the offset parts? I think we should only work with field numbers not offsets.

openmp/libomptarget/DeviceRTL/include/Configuration.h

19 ↗

(On Diff #492197)

Precommit.

24 ↗

(On Diff #492197)

Unused for now.

openmp/libomptarget/DeviceRTL/src/Configuration.cpp

26

Keep it for now. Makes the transition easier. Use an || below.

[tianshilei1992]

rebase, add more code, fix comments

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2641

This will be removed properly later.

[jdoerfert]

openmp/libomptarget/DeviceRTL/src/Configuration.cpp

56

This looks wrong. We want either flag to allow us to remove nested parallelism handling. So
__omp_rtl_assume_no_nested_parallelism is good enough, and
!state::getKernelEnvironment().Configuration.MayUseNestedParallelism is good enough.

[tianshilei1992]

openmp/libomptarget/DeviceRTL/src/Configuration.cpp

56

Yeah, since I already updated OpenMPOpt.cpp, __omp_rtl_assume_no_nested_parallelism will not be generated.

[tianshilei1992]

rebase, fix comments, added more code.

features are finished now. time to move to tests.

[tianshilei1992]

fix bugs and add plugin support

[tianshilei1992]

rebase and fix problems

[tianshilei1992]

rebase and fix bugs

[tianshilei1992]

rebase and fix bugs

openmp/libomptarget/plugins/cuda/src/rtl.cpp

26

These changes in old plugins will be gone because the old plugins will be removed.

[tianshilei1992]

rebase

[jdoerfert]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3888

That is not the kernel, at least not when clang emits a debug wrapper as part of -g, I think.
We have one workaround for this problem already but we need to provide a generic way. I'll add something.

3907

We can't have it constant as a whole. We need a constant part and a non-constant part. Let's add them like that from the very beginning.
I'd suggest we have a third struct type that is references from the kernelenv. It contains for now only the Level value, but we can/will add more things. It's an independent global such that we can make the kernelenv constant.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

212

Maybe close the namespace here.

3717

Shouldn't we always have one here?

3724

This should happen above in the initializer. We should assume no-nested-parallelism and go back on that assumption if need be.

4452

And here you need to check nested-parallelism then too.

openmp/libomptarget/DeviceRTL/src/Configuration.cpp

56

__omp_rtl_assume_no_nested_parallelism is generated by clang, the value depends on -fopenmp-assume-no-nested-parallelism

[tianshilei1992]

rebase and fix comments

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3892

This doesn't work as expected. Need to figure it out.

[tianshilei1992]

fix the WA, rebase, and update some tests

[jdoerfert]

LG

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3893

Typo: tt

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

583

We used to default to SPMD w/o a _exec_mode. Unsure if this will cause problems. I am inclined to allow missing _kernel_env as well and use SPMD then.

[tianshilei1992]

rebase, update tests, fix comments

[tianshilei1992]

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

583

Now we need the global variable to call target_init at runtime, so missing it sounds like a bug. I'm not sure if we would optimize the symbol out, but I kinda doubt it.

[tianshilei1992]

Pass two benchmarks: XSBench and RSBench.
The patch will break the old plugins as well. I'm wondering we might want to land it after D142820.

[jdoerfert]

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

583

Missing used to mean, and does, that it's not a "target region" kernel. So it might be a con/de-structor or a kernel coming from elsewhere, e.g., CUDA. We should continue that behavior I think.

[tianshilei1992]

rebase and fix comments

[tianshilei1992]

update tests in LLVM

[tianshilei1992]

Only llvm/test/Transforms/OpenMP/spmdization_constant_prop.ll is left.

[tianshilei1992]

rebase

[tianshilei1992]

update more tests

[tianshilei1992]

add the last test

[ronlieb]

@tianshilei1992 this seems to have broken the amdgpu buildbot, could you please address quickly ?
or revert and fix ?

[tianshilei1992]

In D142569#4290109, @ronlieb wrote:

@tianshilei1992 this seems to have broken the amdgpu buildbot, could you please address quickly ?
or revert and fix ?

Reverted.

[tianshilei1992]

Fix two test failres. It looks like not working on AMDGPU because of the global
read on the host is not correct. Needs to investigate more.

[tianshilei1992]

llvm/test/Transforms/OpenMP/always_inline_device.ll

11 ↗

(On Diff #516110)

@arsenm Does AMDGPU support a global struct with constant initializer? I currently encountered an error that, the generated IR has correct initializer, but when I read from the global struct from the image (ELF), the corresponding memory looks pretty random. I just need one of the first three i8 values.

[tianshilei1992]

ping @arsenm

[JonChesterfield]

How does a function find the corresponding kernel environment at runtime?

[tianshilei1992]

For AMDGPU it is something in the linker. If I directly read the struct from .o file, the value is correct, but if I read from .img which is generated by lld, then the value is messy/random.

@__omp_offloading_32_71401f4c_main_l12_dynamic_environment = internal addrspace(1) global %struct.DynamicEnvironmentTy zeroinitializer
@__omp_offloading_32_71401f4c_main_l12_kernel_environment = local_unnamed_addr addrspace(1) constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 1 }, ptr addrspacecast (ptr addrspace(1) @1 to ptr), ptr addrspacecast (ptr addrspace(1) @__omp_offloading_32_71401f4c_main_l12_dynamic_environment to ptr) }

Does the linkage of the global __omp_offloading_32_71401f4c_main_l12_kernel_environment look suspicious?

[tianshilei1992]

rebase

[tianshilei1992]

rebase.

I'll land it and see if AMD buildbot will be happy.

[tianshilei1992]

rebase and fix test issues

[tianshilei1992]

rebase and prepare for landing

[dhruvachak]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3920

Is there a reason this has to be ExternalLinkage? Can we use GlobalValue::WeakAnyLinkage here? The external linkage leads to multiply defined linker errors downstream on test cases that have a target region in a header file. For some reason, the problem does not repro on the main branch.
@tianshilei1992 @jdoerfert

[tianshilei1992]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3920

It has been fixed.

[dhruvachak]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3920

It has been fixed.

Thanks. I hadn't pulled, so did not see it earlier.

[doru1004]

The changes in this patch do not actually work correctly with kernels that really need to be in Generic mode. The ExecMode value recovered in the kmpc_kernel_init function is 3 in the case where it needs to be 1.

The problem lies with the OpenMPOpt changes since disabling them makes everything work again.

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

 
 void CGOpenMPRuntimeGPU::emitKernelDeinit(CodeGenFunction &CGF,
                                           EntryFunctionState &EST,
                                           bool IsSPMD) {
   if (!IsSPMD)
     emitGenericVarsEpilog(CGF);
 
   CGBuilderTy &Bld = CGF.Builder;
-  OMPBuilder.createTargetDeinit(Bld, IsSPMD);
+  OMPBuilder.createTargetDeinit(Bld);
 }
 
 void CGOpenMPRuntimeGPU::emitSPMDKernel(const OMPExecutableDirective &D,
                                           StringRef ParentName,
                                           llvm::Function *&OutlinedFn,
                                           llvm::Constant *&OutlinedFnID,
                                           bool IsOffloadEntry,
                                           const RegionCodeGenTy &CodeGen) {
void CGOpenMPRuntimeGPU::emitSPMDKernel(const OMPExecutableDirective &D,
   } Action(*this, EST);
   CodeGen.setAction(Action);
   IsInTTDRegion = true;
   emitTargetOutlinedFunctionHelper(D, ParentName, OutlinedFn, OutlinedFnID,
                                    IsOffloadEntry, CodeGen);
   IsInTTDRegion = false;
 }
 
-// Create a unique global variable to indicate the execution mode of this target
-// region. The execution mode is either 'generic', or 'spmd' depending on the
-// target directive. This variable is picked up by the offload library to setup
-// the device appropriately before kernel launch. If the execution mode is
-// 'generic', the runtime reserves one warp for the master, otherwise, all
-// warps participate in parallel work.
-static void setPropertyExecutionMode(CodeGenModule &CGM, StringRef Name,
-                                     bool Mode) {
-  auto *GVMode = new llvm::GlobalVariable(
-      CGM.getModule(), CGM.Int8Ty, /*isConstant=*/true,
-      llvm::GlobalValue::WeakAnyLinkage,
-      llvm::ConstantInt::get(CGM.Int8Ty, Mode ? OMP_TGT_EXEC_MODE_SPMD
-                                              : OMP_TGT_EXEC_MODE_GENERIC),
-      Twine(Name, "_exec_mode"));
-  GVMode->setVisibility(llvm::GlobalVariable::ProtectedVisibility);
-  CGM.addCompilerUsedGlobal(GVMode);
-}
-
 void CGOpenMPRuntimeGPU::emitTargetOutlinedFunction(
     const OMPExecutableDirective &D, StringRef ParentName,
     llvm::Function *&OutlinedFn, llvm::Constant *&OutlinedFnID,
     bool IsOffloadEntry, const RegionCodeGenTy &CodeGen) {
   if (!IsOffloadEntry) // Nothing to do.
     return;
 
   assert(!ParentName.empty() && "Invalid target region parent name!");
 
   bool Mode = supportsSPMDExecutionMode(CGM.getContext(), D);
   if (Mode)
     emitSPMDKernel(D, ParentName, OutlinedFn, OutlinedFnID, IsOffloadEntry,
                    CodeGen);
   else
     emitNonSPMDKernel(D, ParentName, OutlinedFn, OutlinedFnID, IsOffloadEntry,
                       CodeGen);
-
-  setPropertyExecutionMode(CGM, OutlinedFn->getName(), Mode);
 }
 
 CGOpenMPRuntimeGPU::CGOpenMPRuntimeGPU(CodeGenModule &CGM)
     : CGOpenMPRuntime(CGM) {
   llvm::OpenMPIRBuilderConfig Config(CGM.getLangOpts().OpenMPIsDevice, true,
                                      hasRequiresUnifiedSharedMemory(),
                                      CGM.getLangOpts().OpenMPOffloadMandatory);
   OMPBuilder.setConfig(Config);

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

public:
   ///
   /// \param Loc The insert and source location description.
   /// \param IsSPMD Flag to indicate if the kernel is an SPMD kernel or not.
   InsertPointTy createTargetInit(const LocationDescription &Loc, bool IsSPMD);
 
   /// 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.
-  void createTargetDeinit(const LocationDescription &Loc, bool IsSPMD);
+  void createTargetDeinit(const LocationDescription &Loc);
 
   ///}
 
 private:
   // Sets the function attributes expected for the outlined function
   void setOutlinedTargetRegionFunctionAttributes(Function *OutlinedFn,
                                                  int32_t NumTeams,
                                                  int32_t NumThreads);

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

 __OMP_STRUCT_TYPE(Ident, ident_t, false, Int32, Int32, Int32, Int32, Int8Ptr)
 __OMP_STRUCT_TYPE(OffloadEntry, __tgt_offload_entry, false, Int8Ptr, Int8Ptr, SizeTy,
                   Int32, Int32)
 __OMP_STRUCT_TYPE(KernelArgs, __tgt_kernel_arguments, false, Int32, Int32, VoidPtrPtr,
       VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr,
       Int64, Int64, Int32Arr3Ty, Int32Arr3Ty, Int32)
 __OMP_STRUCT_TYPE(AsyncInfo, __tgt_async_info, false, Int8Ptr)
 __OMP_STRUCT_TYPE(DependInfo, kmp_dep_info, false, SizeTy, SizeTy, Int8)
+__OMP_STRUCT_TYPE(ConfigurationEnvironmentTy, ConfigurationEnvironmentTy, false,
+                  Int8, Int8, Int8)
+__OMP_STRUCT_TYPE(KernelEnvironmentTy, KernelEnvironmentTy, false,
+                  ConfigurationEnvironmentTy, IdentPtr, Int16)
 
 #undef __OMP_STRUCT_TYPE
 #undef OMP_STRUCT_TYPE
 
 #ifndef OMP_FUNCTION_TYPE
 #define OMP_FUNCTION_TYPE(VarName, IsVarArg, ReturnType, ...)
 #endif
 
__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, Int8, Int1)
-__OMP_RTL(__kmpc_target_deinit, false, Void, IdentPtr, Int8)
+__OMP_RTL(__kmpc_target_init, false, Int32, KernelEnvironmentTyPtr)
+__OMP_RTL(__kmpc_target_deinit, false, Void,)
 __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_shuffle_int32, false, Int32, Int32, Int16, Int16)

 __OMP_RTL_ATTRS(__tgt_mapper_num_components, ForkAttrs, AttributeSet(),
                 ParamAttrs())
 __OMP_RTL_ATTRS(__tgt_push_mapper_component, ForkAttrs, AttributeSet(),
                 ParamAttrs())
 __OMP_RTL_ATTRS(__kmpc_task_allow_completion_event, DefaultAttrs,
                 ReturnPtrAttrs, ParamAttrs(ReadOnlyPtrAttrs, SExt))
 
 __OMP_RTL_ATTRS(__kmpc_target_init, AttributeSet(), SExt,
-                ParamAttrs(AttributeSet(), SExt, SExt))
+                ParamAttrs(AttributeSet()))
 __OMP_RTL_ATTRS(__kmpc_target_deinit, AttributeSet(), AttributeSet(),
-                ParamAttrs(AttributeSet(), SExt))
+                ParamAttrs())
 __OMP_RTL_ATTRS(__kmpc_parallel_51, AlwaysInlineAttrs, AttributeSet(),
                 ParamAttrs(AttributeSet(), SExt, SExt, SExt, SExt,
                            AttributeSet(), AttributeSet(), AttributeSet(),
                            SizeTyExt))
 __OMP_RTL_ATTRS(__kmpc_serialized_parallel, InaccessibleArgOnlyAttrs,
                 AttributeSet(), ParamAttrs(ReadOnlyPtrAttrs, SExt))
 __OMP_RTL_ATTRS(__kmpc_end_serialized_parallel, InaccessibleArgOnlyAttrs,
                 AttributeSet(), ParamAttrs(ReadOnlyPtrAttrs, SExt))

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

if (!updateToLocation(Loc))
     return Loc.IP;
 
   uint32_t SrcLocStrSize;
   Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
   Constant *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
   ConstantInt *IsSPMDVal = ConstantInt::getSigned(
       IntegerType::getInt8Ty(Int8->getContext()),
       IsSPMD ? OMP_TGT_EXEC_MODE_SPMD : OMP_TGT_EXEC_MODE_GENERIC);
-  ConstantInt *UseGenericStateMachine =
-      ConstantInt::getBool(Int32->getContext(), !IsSPMD);
+  ConstantInt *UseGenericStateMachineVal = ConstantInt::getSigned(
+      IntegerType::getInt8Ty(Int8->getContext()), !IsSPMD);
+  ConstantInt *MayUseNestedParallelismVal =
+      ConstantInt::getSigned(IntegerType::getInt8Ty(Int8->getContext()), true);
+  ConstantInt *DebugIndentionLevelVal =
+      ConstantInt::getSigned(IntegerType::getInt16Ty(Int8->getContext()), 0);
 
+  Function *Kernel = Builder.GetInsertBlock()->getParent();
   Function *Fn = getOrCreateRuntimeFunctionPtr(
       omp::RuntimeFunction::OMPRTL___kmpc_target_init);
+  const DataLayout &DL = Fn->getParent()->getDataLayout();
 
-  CallInst *ThreadKind = Builder.CreateCall(
-      Fn, {Ident, IsSPMDVal, UseGenericStateMachine});
+  Constant *ConfigurationEnvironmentInitializer = ConstantStruct::get(
+      ConfigurationEnvironmentTy, {
+                                      UseGenericStateMachineVal,
+                                      MayUseNestedParallelismVal,
+                                      IsSPMDVal,
+                                  });
+  Constant *KernelEnvironmentInitializer = ConstantStruct::get(
+      KernelEnvironmentTy, {
+                               ConfigurationEnvironmentInitializer,
+                               Ident,
+                               DebugIndentionLevelVal,
+                           });
+  Twine KernelEnvironmentName = Kernel->getName() + "_kernel_info";
+  GlobalVariable *KernelEnvironment = new GlobalVariable(
+      M, KernelEnvironmentTy, /* IsConstant */ true,
+      GlobalValue::ExternalLinkage, KernelEnvironmentInitializer,
+      KernelEnvironmentName,
+      /* InsertBefore */ nullptr, llvm::GlobalValue::NotThreadLocal,
+      DL.getDefaultGlobalsAddressSpace());
+
+  CallInst *ThreadKind = Builder.CreateCall(Fn, {KernelEnvironment});
 
   Value *ExecUserCode = Builder.CreateICmpEQ(
       ThreadKind, ConstantInt::get(ThreadKind->getType(), -1),
       "exec_user_code");
 
   // ThreadKind = __kmpc_target_init(...)
   // if (ThreadKind == -1)
   //   user_code
   // else
   //   return;
 
   auto *UI = Builder.CreateUnreachable();
   BasicBlock *CheckBB = UI->getParent();
   BasicBlock *UserCodeEntryBB = CheckBB->splitBasicBlock(UI, "user_code.entry");
 
OpenMPIRBuilder::createTargetInit(const LocationDescription &Loc, bool IsSPMD) {
   CheckBBTI->eraseFromParent();
   UI->eraseFromParent();
 
   // Continue in the "user_code" block, see diagram above and in
   // openmp/libomptarget/deviceRTLs/common/include/target.h .
   return InsertPointTy(UserCodeEntryBB, UserCodeEntryBB->getFirstInsertionPt());
 }
 
-void OpenMPIRBuilder::createTargetDeinit(const LocationDescription &Loc,
-                                         bool IsSPMD) {
+void OpenMPIRBuilder::createTargetDeinit(const LocationDescription &Loc) {
   if (!updateToLocation(Loc))
     return;
 
-  uint32_t SrcLocStrSize;
-  Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
-  Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
-  ConstantInt *IsSPMDVal = ConstantInt::getSigned(
-      IntegerType::getInt8Ty(Int8->getContext()),
-      IsSPMD ? OMP_TGT_EXEC_MODE_SPMD : OMP_TGT_EXEC_MODE_GENERIC);
-
   Function *Fn = getOrCreateRuntimeFunctionPtr(
       omp::RuntimeFunction::OMPRTL___kmpc_target_deinit);
 
-  Builder.CreateCall(Fn, {Ident, IsSPMDVal});
+  Builder.CreateCall(Fn, {});
 }
 
 void OpenMPIRBuilder::setOutlinedTargetRegionFunctionAttributes(
     Function *OutlinedFn, int32_t NumTeams, int32_t NumThreads) {
   if (Config.isEmbedded()) {
     OutlinedFn->setLinkage(GlobalValue::WeakODRLinkage);
     // TODO: Determine if DSO local can be set to true.
     OutlinedFn->setDSOLocal(false);

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
+#include "llvm/Frontend/OpenMP/OMPDeviceConstants.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/Assumptions.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Value.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Transforms/IPO/Attributor.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 
 #include <algorithm>

 STATISTIC(NumBytesMovedToSharedMemory,
           "Amount of memory pushed to shared memory");
 STATISTIC(NumBarriersEliminated, "Number of redundant barriers eliminated");
 
 #if !defined(NDEBUG)
 static constexpr auto TAG = "[" DEBUG_TYPE "]";
 #endif
 
+namespace KernelInfo {
+
+// struct ConfigurationEnvironmentTy {
+//   uint8_t UseGenericStateMachine;
+//   uint8_t MayUseNestedParallelism;
+//   llvm::omp::OMPTgtExecModeFlags ExecMode;
+// };
+
+// struct KernelEnvironmentTy {
+//   ConfigurationEnvironmentTy Configuration;
+//   IdentTy *Ident;
+//   uint16_t DebugIndentionLevel;
+// };
+
+#define KERNEL_ENVIRONMENT_IDX(MEMBER, IDX)                                    \
+  constexpr const unsigned MEMBER##Idx = IDX;
+
+KERNEL_ENVIRONMENT_IDX(Configuration, 0)
+KERNEL_ENVIRONMENT_IDX(Ident, 1)
+// TODO: Debug level is not used yet.
+// KERNEL_ENVIRONMENT_IDX(DebugIndentionLevel, 2)
+
+#undef KERNEL_ENVIRONMENT_IDX
+
+#define KERNEL_ENVIRONMENT_CONFIGURATION_IDX(MEMBER, IDX)                      \
+  constexpr const unsigned MEMBER##Idx = IDX;
+
+KERNEL_ENVIRONMENT_CONFIGURATION_IDX(UseGenericStateMachine, 0)
+KERNEL_ENVIRONMENT_CONFIGURATION_IDX(MayUseNestedParallelism, 1)
+KERNEL_ENVIRONMENT_CONFIGURATION_IDX(ExecMode, 2)
+
+#undef KERNEL_ENVIRONMENT_CONFIGURATION_IDX
+
+#define KERNEL_ENVIRONMENT_GETTER(MEMBER, RETURNTYPE)                          \
+  RETURNTYPE *get##MEMBER##FromKernelEnvironment(ConstantStruct *KernelEnvC) { \
+    return cast<ConstantStruct>(KernelEnvC->getAggregateElement(MEMBER##Idx)); \
+  }
+
+KERNEL_ENVIRONMENT_GETTER(Ident, Constant)
+KERNEL_ENVIRONMENT_GETTER(Configuration, ConstantStruct)
+
+#undef KERNEL_ENVIRONMENT_GETTER
+
+#define KERNEL_ENVIRONMENT_CONFIGURATION_GETTER(MEMBER)                        \
+  ConstantInt *get##MEMBER##FromKernelEnvironment(                             \
+      ConstantStruct *KernelEnvC) {                                            \
+    ConstantStruct *ConfigC =                                                  \
+        getConfigurationFromKernelEnvironment(KernelEnvC);                     \
+    return dyn_cast<ConstantInt>(ConfigC->getAggregateElement(MEMBER##Idx));   \
+  }
+
+KERNEL_ENVIRONMENT_CONFIGURATION_GETTER(UseGenericStateMachine)
+KERNEL_ENVIRONMENT_CONFIGURATION_GETTER(MayUseNestedParallelism)
+KERNEL_ENVIRONMENT_CONFIGURATION_GETTER(ExecMode)
+
+#undef KERNEL_ENVIRONMENT_CONFIGURATION_GETTER
+
+GlobalVariable *
+getKernelEnvironementGVFromKernelInitCB(CallBase *KernelInitCB) {
+  constexpr const int InitKernelEnvironmentArgNo = 0;
+  return cast<GlobalVariable>(
+      KernelInitCB->getArgOperand(InitKernelEnvironmentArgNo)
+          ->stripPointerCasts());
+}
+
+ConstantStruct *getKernelEnvironementFromKernelInitCB(CallBase *KernelInitCB) {
+  GlobalVariable *KernelEnvGV =
+      getKernelEnvironementGVFromKernelInitCB(KernelInitCB);
+  return cast<ConstantStruct>(KernelEnvGV->getInitializer());
+}
+} // namespace KernelInfo
+
 namespace {
 
 struct AAHeapToShared;
 
 struct AAICVTracker;
 
 /// OpenMP specific information. For now, stores RFIs and ICVs also needed for
 /// Attributor runs.
struct RuntimeFunctionInfo {
     bool IsVarArg;
 
     /// The return type of the function.
     Type *ReturnType;
 
     /// The argument types of the function.
     SmallVector<Type *, 8> ArgumentTypes;
 
     /// The declaration if available.
     Function *Declaration = nullptr;
 
     /// Uses of this runtime function per function containing the use.
     using UseVector = SmallVector<Use *, 16>;
 
     /// Clear UsesMap for runtime function.
     void clearUsesMap() { UsesMap.clear(); }
 
struct KernelInfoState : AbstractState {
   /// we cannot be. If it is assumed, then RequiresFullRuntime should also be
   /// false.
   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 constant kernel environement as taken from and passed to
+  /// __kmpc_target_init.
+  ConstantStruct *KernelEnvC = 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.
if (KIS.KernelInitCB) {
       KernelInitCB = KIS.KernelInitCB;
     }
     if (KIS.KernelDeinitCB) {
       if (KernelDeinitCB && KernelDeinitCB != KIS.KernelDeinitCB)
         llvm_unreachable("Kernel that calls another kernel violates OpenMP-Opt "
                          "assumptions.");
       KernelDeinitCB = KIS.KernelDeinitCB;
     }
+    if (KIS.KernelEnvC) {
+      if (KernelEnvC && KernelEnvC != KIS.KernelEnvC)
+        llvm_unreachable("Kernel that calls another kernel violates OpenMP-Opt "
+                         "assumptions.");
+      KernelEnvC = KIS.KernelEnvC;
+    }
     SPMDCompatibilityTracker ^= KIS.SPMDCompatibilityTracker;
     ReachedKnownParallelRegions ^= KIS.ReachedKnownParallelRegions;
     ReachedUnknownParallelRegions ^= KIS.ReachedUnknownParallelRegions;
     NestedParallelism |= KIS.NestedParallelism;
     return *this;
   }
 
   KernelInfoState operator&=(const KernelInfoState &KIS) {
ChangeStatus updateImpl(Attributor &A) override {
     return Changed;
   }
 };
 
 struct AAExecutionDomainFunction : public AAExecutionDomain {
   AAExecutionDomainFunction(const IRPosition &IRP, Attributor &A)
       : AAExecutionDomain(IRP, A) {}
 
   ~AAExecutionDomainFunction() { delete RPOT; }
 
   void initialize(Attributor &A) override {
     if (getAnchorScope()->isDeclaration()) {
       indicatePessimisticFixpoint();
       return;
     }
     RPOT = new ReversePostOrderTraversal<Function *>(getAnchorScope());
   }
static bool isInitialThreadOnlyEdge(Attributor &A, BranchInst *Edge,
     // Match: -1 == __kmpc_target_init (for non-SPMD kernels only!)
     if (C->isAllOnesValue()) {
       auto *CB = dyn_cast<CallBase>(Cmp->getOperand(0));
       auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
       auto &RFI = OMPInfoCache.RFIs[OMPRTL___kmpc_target_init];
       CB = CB ? OpenMPOpt::getCallIfRegularCall(*CB, &RFI) : nullptr;
       if (!CB)
         return false;
-      const int InitModeArgNo = 1;
-      auto *ModeCI = dyn_cast<ConstantInt>(CB->getOperand(InitModeArgNo));
-      return ModeCI && (ModeCI->getSExtValue() & OMP_TGT_EXEC_MODE_GENERIC);
+      ConstantStruct *KernelEnvC =
+          KernelInfo::getKernelEnvironementFromKernelInitCB(CB);
+      ConstantInt *ExecModeC =
+          KernelInfo::getExecModeFromKernelEnvironment(KernelEnvC);
+      return ExecModeC->getSExtValue() & OMP_TGT_EXEC_MODE_GENERIC;
     }
 
     if (C->isZero()) {
       // Match: 0 == llvm.nvvm.read.ptx.sreg.tid.x()
       if (auto *II = dyn_cast<IntrinsicInst>(Cmp->getOperand(0)))
         if (II->getIntrinsicID() == Intrinsic::nvvm_read_ptx_sreg_tid_x)
           return true;
 
AAKernelInfoFunction(const IRPosition &IRP, Attributor &A)
       : AAKernelInfo(IRP, A) {}
 
   SmallPtrSet<Instruction *, 4> GuardedInstructions;
 
   SmallPtrSetImpl<Instruction *> &getGuardedInstructions() {
     return GuardedInstructions;
   }
 
+  void setConfigurationOfKernelEnvironment(ConstantStruct *ConfigC) {
+    Constant *NewKernelEnvC = ConstantFoldInsertValueInstruction(
+        KernelEnvC, ConfigC, {KernelInfo::ConfigurationIdx});
+    assert(NewKernelEnvC && "Failed to create new kernel environment");
+    KernelEnvC = cast<ConstantStruct>(NewKernelEnvC);
+  }
+
+#define KERNEL_ENVIRONMENT_CONFIGURATION_SETTER(MEMBER)                        \
+  void set##MEMBER##OfKernelEnvironment(ConstantInt *NewVal) {                 \
+    ConstantStruct *ConfigC =                                                  \
+        KernelInfo::getConfigurationFromKernelEnvironment(KernelEnvC);         \
+    Constant *NewConfigC = ConstantFoldInsertValueInstruction(                 \
+        ConfigC, NewVal, {KernelInfo::MEMBER##Idx});                           \
+    assert(NewConfigC && "Failed to create new configuration environment");    \
+    setConfigurationOfKernelEnvironment(cast<ConstantStruct>(NewConfigC));     \
+  }
+
+  KERNEL_ENVIRONMENT_CONFIGURATION_SETTER(UseGenericStateMachine)
+  KERNEL_ENVIRONMENT_CONFIGURATION_SETTER(MayUseNestedParallelism)
+  KERNEL_ENVIRONMENT_CONFIGURATION_SETTER(ExecMode)
+
+#undef KERNEL_ENVIRONMENT_CONFIGURATION_SETTER
+
   /// 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();
void initialize(Attributor &A) override {
     // Ignore kernels without initializers such as global constructors.
     if (!KernelInitCB || !KernelDeinitCB)
       return;
 
     // Add itself to the reaching kernel and set IsKernelEntry.
     ReachingKernelEntries.insert(Fn);
     IsKernelEntry = true;
 
-    // For kernels we might need to initialize/finalize the IsSPMD state and
-    // we need to register a simplification callback so that the Attributor
-    // knows the constant arguments to __kmpc_target_init and
-    // __kmpc_target_deinit might actually change.
+    KernelEnvC =
+        KernelInfo::getKernelEnvironementFromKernelInitCB(KernelInitCB);
 
-    Attributor::SimplifictionCallbackTy StateMachineSimplifyCB =
+    Attributor::SimplifictionCallbackTy KernelConfigurationSimplifyCB =
         [&](const IRPosition &IRP, const AbstractAttribute *AA,
             bool &UsedAssumedInformation) -> std::optional<Value *> {
-      // IRP represents the "use generic state machine" argument of an
-      // __kmpc_target_init call. We will answer this one with the internal
-      // state. As long as we are not in an invalid state, we will create a
-      // custom state machine so the value should be a `i1 false`. If we are
-      // in an invalid state, we won't change the value that is in the IR.
-      if (!ReachedKnownParallelRegions.isValidState())
-        return nullptr;
-      // If we have disabled state machine rewrites, don't make a custom one.
-      if (DisableOpenMPOptStateMachineRewrite)
-        return nullptr;
-      if (AA)
-        A.recordDependence(*this, *AA, DepClassTy::OPTIONAL);
-      UsedAssumedInformation = !isAtFixpoint();
-      auto *FalseVal =
-          ConstantInt::getBool(IRP.getAnchorValue().getContext(), false);
-      return FalseVal;
+      return KernelEnvC;
     };
 
-    Attributor::SimplifictionCallbackTy ModeSimplifyCB =
-        [&](const IRPosition &IRP, const AbstractAttribute *AA,
-            bool &UsedAssumedInformation) -> std::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::getSigned(
-          IntegerType::getInt8Ty(IRP.getAnchorValue().getContext()),
-          SPMDCompatibilityTracker.isAssumed() ? OMP_TGT_EXEC_MODE_SPMD
-                                               : OMP_TGT_EXEC_MODE_GENERIC);
-      return Val;
-    };
-
-    constexpr const int InitModeArgNo = 1;
-    constexpr const int DeinitModeArgNo = 1;
-    constexpr const int InitUseStateMachineArgNo = 2;
-    A.registerSimplificationCallback(
-        IRPosition::callsite_argument(*KernelInitCB, InitUseStateMachineArgNo),
-        StateMachineSimplifyCB);
-    A.registerSimplificationCallback(
-        IRPosition::callsite_argument(*KernelInitCB, InitModeArgNo),
-        ModeSimplifyCB);
-    A.registerSimplificationCallback(
-        IRPosition::callsite_argument(*KernelDeinitCB, DeinitModeArgNo),
-        ModeSimplifyCB);
+    A.registerSimplificationCallback(IRPosition::value(*KernelEnvC),
+                                     KernelConfigurationSimplifyCB);
 
     // Check if we know we are in SPMD-mode already.
-    ConstantInt *ModeArg =
-        dyn_cast<ConstantInt>(KernelInitCB->getArgOperand(InitModeArgNo));
-    if (ModeArg && (ModeArg->getSExtValue() & OMP_TGT_EXEC_MODE_SPMD))
+    ConstantInt *ExecModeC =
+        KernelInfo::getExecModeFromKernelEnvironment(KernelEnvC);
+    ConstantInt *AssumedExecModeC = ConstantInt::get(
+        ExecModeC->getType(),
+        ExecModeC->getSExtValue() | OMP_TGT_EXEC_MODE_GENERIC_SPMD);
+    if (ExecModeC->getSExtValue() & OMP_TGT_EXEC_MODE_SPMD)
       SPMDCompatibilityTracker.indicateOptimisticFixpoint();
-    // This is a generic region but SPMDization is disabled so stop tracking.
     else if (DisableOpenMPOptSPMDization)
+      // This is a generic region but SPMDization is disabled so stop
+      // tracking.
       SPMDCompatibilityTracker.indicatePessimisticFixpoint();
+    else
+      setExecModeOfKernelEnvironment(AssumedExecModeC);
 
     // Register virtual uses of functions we might need to preserve.
     auto RegisterVirtualUse = [&](RuntimeFunction RFKind,
                                   Attributor::VirtualUseCallbackTy &CB) {
       if (!OMPInfoCache.RFIs[RFKind].Declaration)
         return;
       A.registerVirtualUseCallback(*OMPInfoCache.RFIs[RFKind].Declaration, CB);
     };
#undef KERNEL_ENVIRONMENT_CONFIGURATION_SETTER
   /// Modify the IR based on the KernelInfoState as the fixpoint iteration is
   /// finished now.
   ChangeStatus manifest(Attributor &A) override {
     // If we are not looking at a kernel with __kmpc_target_init and
     // __kmpc_target_deinit call we cannot actually manifest the information.
     if (!KernelInitCB || !KernelDeinitCB)
       return ChangeStatus::UNCHANGED;
 
-    /// Insert nested Parallelism global variable
-    Function *Kernel = getAnchorScope();
-    Module &M = *Kernel->getParent();
-    Type *Int8Ty = Type::getInt8Ty(M.getContext());
-    new GlobalVariable(M, Int8Ty, /* isConstant */ true,
-                       GlobalValue::WeakAnyLinkage,
-                       ConstantInt::get(Int8Ty, NestedParallelism ? 1 : 0),
-                       Kernel->getName() + "_nested_parallelism");
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
 
     // If we can we change the execution mode to SPMD-mode otherwise we build a
     // custom state machine.
-    ChangeStatus Changed = ChangeStatus::UNCHANGED;
     if (!changeToSPMDMode(A, Changed)) {
       if (!KernelInitCB->getCalledFunction()->isDeclaration())
-        return buildCustomStateMachine(A);
+        Changed |= buildCustomStateMachine(A);
+    }
+
+    if (KernelEnvC) {
+      ConstantInt *MayUseNestedParallelismC =
+          KernelInfo::getMayUseNestedParallelismFromKernelEnvironment(
+              KernelEnvC);
+      ConstantInt *NewMayUseNestedParallelismC = ConstantInt::get(
+          MayUseNestedParallelismC->getType(), NestedParallelism);
+      setMayUseNestedParallelismOfKernelEnvironment(
+          NewMayUseNestedParallelismC);
+
+      // At last, update the KernelEnvc
+      GlobalVariable *KernelEnvGV =
+          KernelInfo::getKernelEnvironementGVFromKernelInitCB(KernelInitCB);
+      if (KernelEnvGV->getInitializer() != KernelEnvC) {
+        KernelEnvGV->setInitializer(KernelEnvC);
+        Changed = ChangeStatus::CHANGED;
+      }
     }
 
     return Changed;
   }
 
   void insertInstructionGuardsHelper(Attributor &A) {
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
 
bool changeToSPMDMode(Attributor &A, ChangeStatus &Changed) {
     if (Kernel->hasLocalLinkage()) {
       assert(Kernel->hasOneUse() && "Unexpected use of debug kernel wrapper.");
       auto *CB = cast<CallBase>(Kernel->user_back());
       Kernel = CB->getCaller();
     }
     assert(OMPInfoCache.Kernels.count(Kernel) && "Expected kernel function!");
 
     // Check if the kernel is already in SPMD mode, if so, return success.
-    GlobalVariable *ExecMode = Kernel->getParent()->getGlobalVariable(
-        (Kernel->getName() + "_exec_mode").str());
-    assert(ExecMode && "Kernel without exec mode?");
-    assert(ExecMode->getInitializer() && "ExecMode doesn't have initializer!");
-
-    // Set the global exec mode flag to indicate SPMD-Generic mode.
-    assert(isa<ConstantInt>(ExecMode->getInitializer()) &&
-           "ExecMode is not an integer!");
-    const int8_t ExecModeVal =
-        cast<ConstantInt>(ExecMode->getInitializer())->getSExtValue();
+    ConstantStruct *ExistingKernelEnvC =
+        KernelInfo::getKernelEnvironementFromKernelInitCB(KernelInitCB);
+    auto *ExecModeC =
+        KernelInfo::getExecModeFromKernelEnvironment(ExistingKernelEnvC);
+    const int8_t ExecModeVal = ExecModeC->getSExtValue();
     if (ExecModeVal != OMP_TGT_EXEC_MODE_GENERIC)
       return true;
 
     // We will now unconditionally modify the IR, indicate a change.
     Changed = ChangeStatus::CHANGED;
 
     // Do not use instruction guards when no parallel is present inside
     // the target region.
     if (mayContainParallelRegion())
       insertInstructionGuardsHelper(A);
     else
       forceSingleThreadPerWorkgroupHelper(A);
 
     // Adjust the global exec mode flag that tells the runtime what mode this
     // kernel is executed in.
     assert(ExecModeVal == OMP_TGT_EXEC_MODE_GENERIC &&
            "Initially non-SPMD kernel has SPMD exec mode!");
-    ExecMode->setInitializer(
-        ConstantInt::get(ExecMode->getInitializer()->getType(),
-                         ExecModeVal | OMP_TGT_EXEC_MODE_GENERIC_SPMD));
-
-    // Next rewrite the init and deinit calls to indicate we use SPMD-mode now.
-    const int InitModeArgNo = 1;
-    const int DeinitModeArgNo = 1;
-    const int InitUseStateMachineArgNo = 2;
-
-    auto &Ctx = getAnchorValue().getContext();
-    A.changeUseAfterManifest(
-        KernelInitCB->getArgOperandUse(InitModeArgNo),
-        *ConstantInt::getSigned(IntegerType::getInt8Ty(Ctx),
-                                OMP_TGT_EXEC_MODE_SPMD));
-    A.changeUseAfterManifest(
-        KernelInitCB->getArgOperandUse(InitUseStateMachineArgNo),
-        *ConstantInt::getBool(Ctx, false));
-    A.changeUseAfterManifest(
-        KernelDeinitCB->getArgOperandUse(DeinitModeArgNo),
-        *ConstantInt::getSigned(IntegerType::getInt8Ty(Ctx),
-                                OMP_TGT_EXEC_MODE_SPMD));
+    setExecModeOfKernelEnvironment(ConstantInt::get(
+        ExecModeC->getType(), ExecModeVal | OMP_TGT_EXEC_MODE_GENERIC_SPMD));
 
     ++NumOpenMPTargetRegionKernelsSPMD;
 
     auto Remark = [&](OptimizationRemark OR) {
       return OR << "Transformed generic-mode kernel to SPMD-mode.";
     };
     A.emitRemark<OptimizationRemark>(KernelInitCB, "OMP120", Remark);
     return true;
ChangeStatus buildCustomStateMachine(Attributor &A) {
 
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
     if (!OMPInfoCache.runtimeFnsAvailable(
             {OMPRTL___kmpc_get_hardware_num_threads_in_block,
              OMPRTL___kmpc_get_warp_size, OMPRTL___kmpc_barrier_simple_generic,
              OMPRTL___kmpc_kernel_parallel, OMPRTL___kmpc_kernel_end_parallel}))
       return ChangeStatus::UNCHANGED;
 
-    const int InitModeArgNo = 1;
-    const int InitUseStateMachineArgNo = 2;
+    ConstantStruct *ExistingKernelEnvC =
+        KernelInfo::getKernelEnvironementFromKernelInitCB(KernelInitCB);
 
     // 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 *Mode =
-        dyn_cast<ConstantInt>(KernelInitCB->getArgOperand(InitModeArgNo));
+    ConstantInt *UseStateMachineC =
+        KernelInfo::getUseGenericStateMachineFromKernelEnvironment(
+            ExistingKernelEnvC);
+    ConstantInt *ModeC =
+        KernelInfo::getExecModeFromKernelEnvironment(ExistingKernelEnvC);
 
     // 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() || !Mode ||
-        (Mode->getSExtValue() & OMP_TGT_EXEC_MODE_SPMD))
+    if (UseStateMachineC->isZero() ||
+        (ModeC->getSExtValue() & OMP_TGT_EXEC_MODE_SPMD))
       return ChangeStatus::UNCHANGED;
 
     // If not SPMD mode, indicate we use a custom state machine now.
-    auto &Ctx = getAnchorValue().getContext();
-    auto *FalseVal = ConstantInt::getBool(Ctx, false);
-    A.changeUseAfterManifest(
-        KernelInitCB->getArgOperandUse(InitUseStateMachineArgNo), *FalseVal);
+    setUseGenericStateMachineOfKernelEnvironment(
+        ConstantInt::get(UseStateMachineC->getType(), false));
 
     // If we don't actually need a state machine we are done here. This can
     // happen if there simply are no parallel regions. In the resulting kernel
     // all worker threads will simply exit right away, leaving the main thread
     // to do the work alone.
     if (!mayContainParallelRegion()) {
       ++NumOpenMPTargetRegionKernelsWithoutStateMachine;
 
ChangeStatus buildCustomStateMachine(Attributor &A) {
     // SMEndParallelBB:           __kmpc_kernel_end_parallel(...);
     //                          }
     // SMDoneBB:                __kmpc_barrier_simple_generic(...);
     //                          goto SMBeginBB;
     //                       }
     // UserCodeEntryBB:      // user code
     //                       __kmpc_target_deinit(...)
     //
+    auto &Ctx = getAnchorValue().getContext();
     Function *Kernel = getAssociatedFunction();
     assert(Kernel && "Expected an associated function!");
 
     BasicBlock *InitBB = KernelInitCB->getParent();
     BasicBlock *UserCodeEntryBB = InitBB->splitBasicBlock(
         KernelInitCB->getNextNode(), "thread.user_code.check");
     BasicBlock *IsWorkerCheckBB =
         BasicBlock::Create(Ctx, "is_worker_check", Kernel, UserCodeEntryBB);
ChangeStatus buildCustomStateMachine(Attributor &A) {
     WorkFnAI->setDebugLoc(DLoc);
 
     OMPInfoCache.OMPBuilder.updateToLocation(
         OpenMPIRBuilder::LocationDescription(
             IRBuilder<>::InsertPoint(StateMachineBeginBB,
                                      StateMachineBeginBB->end()),
             DLoc));
 
-    Value *Ident = KernelInitCB->getArgOperand(0);
+    Value *Ident = KernelInfo::getIdentFromKernelEnvironment(KernelEnvC);
     Value *GTid = KernelInitCB;
 
     FunctionCallee BarrierFn =
         OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
             M, OMPRTL___kmpc_barrier_simple_generic);
     CallInst *Barrier =
         CallInst::Create(BarrierFn, {Ident, GTid}, "", StateMachineBeginBB);
     OMPInfoCache.setCallingConvention(BarrierFn, Barrier);
ChangeStatus buildCustomStateMachine(Attributor &A) {
     return ChangeStatus::CHANGED;
   }
 
   /// Fixpoint iteration update function. Will be called every time a dependence
   /// changed its state (and in the beginning).
   ChangeStatus updateImpl(Attributor &A) override {
     KernelInfoState StateBefore = getState();
 
+    // When we leave this function this RAII will make sure the member
+    // KernelEnvC is updated properly depending on the state. That member is
+    // used for simplification of values and needs to be up to date at all
+    // times.
+    struct UpdateKernelEnvCRAII {
+      AAKernelInfoFunction &AA;
+
+      UpdateKernelEnvCRAII(AAKernelInfoFunction &AA) : AA(AA) {}
+
+      ~UpdateKernelEnvCRAII() {
+        if (!AA.KernelEnvC)
+          return;
+
+        ConstantStruct *ExistingKernelEnvC =
+            KernelInfo::getKernelEnvironementFromKernelInitCB(AA.KernelInitCB);
+
+        if (!AA.isValidState()) {
+          AA.KernelEnvC = ExistingKernelEnvC;
+          return;
+        }
+
+        if (!AA.ReachedKnownParallelRegions.isValidState())
+          AA.setUseGenericStateMachineOfKernelEnvironment(
+              KernelInfo::getUseGenericStateMachineFromKernelEnvironment(
+                  ExistingKernelEnvC));
+
+        if (!AA.SPMDCompatibilityTracker.isValidState())
+          AA.setExecModeOfKernelEnvironment(
+              KernelInfo::getExecModeFromKernelEnvironment(ExistingKernelEnvC));
+      }
+    } RAII(*this);
+
     // Callback to check a read/write instruction.
     auto CheckRWInst = [&](Instruction &I) {
       // We handle calls later.
       if (isa<CallBase>(I))
         return true;
       // We only care about write effects.
       if (!I.mayWriteToMemory())
         return true;

openmp/libomptarget/DeviceRTL/CMakeLists.txt

 list(TRANSFORM LIBOMPTARGET_LLVM_INCLUDE_DIRS_DEVICERTL PREPEND "-I")
 
 # Set flags for LLVM Bitcode compilation.
 set(bc_flags -c -emit-llvm -std=c++17 -fvisibility=hidden
               ${clang_opt_flags} --offload-device-only
              -nocudalib -nogpulib -nostdinc
              -fopenmp -fopenmp-cuda-mode
              -Wno-unknown-cuda-version
+             -DOMPTARGET_DEVICE_RUNTIME
              -I${include_directory}
              -I${devicertl_base_directory}/../include
              ${LIBOMPTARGET_LLVM_INCLUDE_DIRS_DEVICERTL}
 )
 
 # first create an object target
 add_library(omptarget.devicertl.all_objs OBJECT IMPORTED)
 function(compileDeviceRTLLibrary target_cpu target_name target_triple)

openmp/libomptarget/DeviceRTL/include/Debug.h

#define FunctionTracingRAII()                                                  \
   DebugEntryRAII Entry(__FILE__, __LINE__, __PRETTY_FUNCTION__);
 
 /// An RAII class for handling entries to debug locations. The current location
 /// and function will be printed on entry. Nested levels increase the
 /// indentation shown in the debugging output.
 struct DebugEntryRAII {
   DebugEntryRAII(const char *File, const unsigned Line, const char *Function);
   ~DebugEntryRAII();
-
-  static void init();
 };
 
 #endif

openmp/libomptarget/DeviceRTL/include/Interface.h

 uint32_t __kmpc_get_hardware_num_threads_in_block();
 
 /// External interface to get the warp size.
 uint32_t __kmpc_get_warp_size();
 
 /// Kernel
 ///
 ///{
+// Forward declaration
+struct KernelEnvironmentTy;
+
 int8_t __kmpc_is_spmd_exec_mode();
 
-int32_t __kmpc_target_init(IdentTy *Ident, int8_t Mode,
-                           bool UseGenericStateMachine);
+int32_t __kmpc_target_init(KernelEnvironmentTy &KernelEnvironment);
 
-void __kmpc_target_deinit(IdentTy *Ident, int8_t Mode);
+void __kmpc_target_deinit();
 
 ///}
 
 /// Reduction
 ///
 ///{
 void __kmpc_nvptx_end_reduce(int32_t TId);
 

openmp/libomptarget/DeviceRTL/include/State.h

 #ifndef OMPTARGET_STATE_H
 #define OMPTARGET_STATE_H
 
 #include "Debug.h"
 #include "Mapping.h"
 #include "Types.h"
 #include "Utils.h"
 
+// Forward declaration.
+struct KernelEnvironmentTy;
+
 #pragma omp begin declare target device_type(nohost)
 
 namespace ompx {
 
 namespace memory {
 
 /// Alloca \p Size bytes in shared memory, if possible, for \p Reason.
 ///
void init(ThreadStateTy *PreviousTS) {
     PreviousThreadState = PreviousTS;
   }
 };
 
 extern ThreadStateTy **ThreadStates;
 #pragma omp allocate(ThreadStates) allocator(omp_pteam_mem_alloc)
 
 /// Initialize the state machinery. Must be called by all threads.
-void init(bool IsSPMD);
+void init(bool IsSPMD, KernelEnvironmentTy &KernelEnvironment);
+
+/// Return the kernel environment associated with the current kernel.
+KernelEnvironmentTy &getKernelEnvironment();
 
 /// TODO
 enum ValueKind {
   VK_NThreads,
   VK_Level,
   VK_ActiveLevel,
   VK_MaxActiveLevels,
   VK_RunSched,

openmp/libomptarget/DeviceRTL/src/Configuration.cpp

 //===- Configuration.cpp - OpenMP device configuration interface -- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the data object of the constant device environment and the
 // query API.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Configuration.h"
-#include "DeviceEnvironment.h"
+#include "Environment.h"
 #include "State.h"
 #include "Types.h"
 
 using namespace ompx;
 
 #pragma omp begin declare target device_type(nohost)
 
 // defined by CGOpenMPRuntimeGPU
 extern uint32_t __omp_rtl_debug_kind;
 extern uint32_t __omp_rtl_assume_no_thread_state;
-extern uint32_t __omp_rtl_assume_no_nested_parallelism;
 
 // This variable should be visibile to the plugin so we override the default
 // hidden visibility.
 DeviceEnvironmentTy CONSTANT(__omp_rtl_device_environment)
     __attribute__((used, retain, weak, visibility("protected")));
 
 uint32_t config::getDebugKind() {
   return __omp_rtl_debug_kind & __omp_rtl_device_environment.DebugKind;

 
 bool config::isDebugMode(config::DebugKind Kind) {
   return config::getDebugKind() & Kind;
 }
 
 bool config::mayUseThreadStates() { return !__omp_rtl_assume_no_thread_state; }
 
 bool config::mayUseNestedParallelism() {
-  return !__omp_rtl_assume_no_nested_parallelism;
+  return state::getKernelEnvironment().Configuration.MayUseNestedParallelism;
 }
 
 #pragma omp end declare target

openmp/libomptarget/DeviceRTL/src/Debug.cpp

 //===--- Debug.cpp -------- Debug utilities ----------------------- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains debug utilities
 //
 //===----------------------------------------------------------------------===//
 
 #include "Debug.h"
 #include "Configuration.h"
+#include "Environment.h"
 #include "Interface.h"
 #include "Mapping.h"
+#include "State.h"
 #include "Types.h"
 
 using namespace ompx;
 
 #pragma omp begin declare target device_type(nohost)
 
 extern "C" {
 void __assert_assume(bool condition) { __builtin_assume(condition); }
 
 void __assert_fail(const char *assertion, const char *file, unsigned line,
                    const char *function) {
   PRINTF("%s:%u: %s: Assertion `%s' failed.\n", file, line, function,
          assertion);
   __builtin_trap();
 }
 }
 
-/// Current indentation level for the function trace. Only accessed by thread 0.
-__attribute__((loader_uninitialized)) static uint32_t Level;
-#pragma omp allocate(Level) allocator(omp_pteam_mem_alloc)
-
 DebugEntryRAII::DebugEntryRAII(const char *File, const unsigned Line,
                                const char *Function) {
   if (config::isDebugMode(config::DebugKind::FunctionTracing) &&
       mapping::getThreadIdInBlock() == 0 && mapping::getBlockId() == 0) {
 
+    uint16_t &Level = state::getKernelEnvironment().DebugIndentionLevel;
+
     for (int I = 0; I < Level; ++I)
       PRINTF("%s", "  ");
 
     PRINTF("%s:%u: Thread %u Entering %s\n", File, Line,
            mapping::getThreadIdInBlock(), Function);
     Level++;
   }
 }
 
 DebugEntryRAII::~DebugEntryRAII() {
   if (config::isDebugMode(config::DebugKind::FunctionTracing) &&
-      mapping::getThreadIdInBlock() == 0 && mapping::getBlockId() == 0)
+      mapping::getThreadIdInBlock() == 0 && mapping::getBlockId() == 0) {
+    uint16_t &Level = state::getKernelEnvironment().DebugIndentionLevel;
     Level--;
   }
-
-void DebugEntryRAII::init() { Level = 0; }
+}
 
 #pragma omp end declare target

openmp/libomptarget/DeviceRTL/src/Kernel.cpp

 //===--- Kernel.cpp - OpenMP device kernel interface -------------- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the kernel entry points for the device.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Debug.h"
+#include "Environment.h"
 #include "Interface.h"
 #include "Mapping.h"
 #include "State.h"
 #include "Synchronization.h"
 #include "Types.h"
 
 #include "llvm/Frontend/OpenMP/OMPDeviceConstants.h"
 
 using namespace ompx;
 
 #pragma omp begin declare target device_type(nohost)
 
-static void inititializeRuntime(bool IsSPMD) {
+static void inititializeRuntime(bool IsSPMD,
+                                KernelEnvironmentTy &KernelEnvironment) {
   // Order is important here.
   synchronize::init(IsSPMD);
   mapping::init(IsSPMD);
-  state::init(IsSPMD);
+  state::init(IsSPMD, KernelEnvironment);
 }
 
 /// Simple generic state machine for worker threads.
 static void genericStateMachine(IdentTy *Ident) {
   FunctionTracingRAII();
 
   uint32_t TId = mapping::getThreadIdInBlock();
 

 }
 
 extern "C" {
 
 /// Initialization
 ///
 /// \param Ident               Source location identification, can be NULL.
 ///
-int32_t __kmpc_target_init(IdentTy *Ident, int8_t Mode,
-                           bool UseGenericStateMachine) {
+int32_t __kmpc_target_init(KernelEnvironmentTy &KernelEnvironment) {
   FunctionTracingRAII();
-  const bool IsSPMD =
-      Mode & llvm::omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_SPMD;
+  ConfigurationEnvironmentTy &Configuration = KernelEnvironment.Configuration;
+  bool IsSPMD = Configuration.ExecMode &
+                llvm::omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_SPMD;
+  bool UseGenericStateMachine = Configuration.UseGenericStateMachine;
   if (IsSPMD) {
-    inititializeRuntime(/* IsSPMD */ true);
+    inititializeRuntime(/* IsSPMD */ true, KernelEnvironment);
     synchronize::threadsAligned();
   } else {
-    inititializeRuntime(/* IsSPMD */ false);
+    inititializeRuntime(/* IsSPMD */ false, KernelEnvironment);
     // No need to wait since only the main threads will execute user
     // code and workers will run into a barrier right away.
   }
 
   if (IsSPMD) {
     state::assumeInitialState(IsSPMD);
     return -1;
   }
int32_t __kmpc_target_init(KernelEnvironmentTy &KernelEnvironment) {
   // synchronize::threads call in genericStateMachine before the main thread
   // reaches its corresponding synchronize::threads call: that would permit all
   // active worker threads to proceed before the main thread has actually set
   // state::ParallelRegionFn, and then they would immediately quit without
   // doing any work.  mapping::getBlockSize() does not include any of the main
   // thread's warp, so none of its threads can ever be active worker threads.
   if (UseGenericStateMachine &&
       mapping::getThreadIdInBlock() < mapping::getBlockSize(IsSPMD)) {
-    genericStateMachine(Ident);
+    genericStateMachine(KernelEnvironment.Ident);
   } else {
     // Retrieve the work function just to ensure we always call
     // __kmpc_kernel_parallel even if a custom state machine is used.
     // TODO: this is not super pretty. The problem is we create the call to
     // __kmpc_kernel_parallel in the openmp-opt pass but while we optimize it is
     // not there yet. Thus, we assume we never reach it from
     // __kmpc_target_deinit. That allows us to remove the store in there to
     // ParallelRegionFn, which leads to bad results later on.
     ParallelRegionFnTy WorkFn = nullptr;
     __kmpc_kernel_parallel(&WorkFn);
     ASSERT(WorkFn == nullptr);
   }
 
   return mapping::getThreadIdInBlock();
 }
 
 /// De-Initialization
 ///
 /// In non-SPMD, this function releases the workers trapped in a state machine
 /// and also any memory dynamically allocated by the runtime.
 ///
 /// \param Ident Source location identification, can be NULL.
 ///
-void __kmpc_target_deinit(IdentTy *Ident, int8_t Mode) {
+void __kmpc_target_deinit() {
   FunctionTracingRAII();
-  const bool IsSPMD =
-      Mode & llvm::omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_SPMD;
+  bool IsSPMD = mapping::isSPMDMode();
   state::assumeInitialState(IsSPMD);
   if (IsSPMD)
     return;
 
   // Signal the workers to exit the state machine and exit the kernel.
   state::ParallelRegionFn = nullptr;
 }
 
 int8_t __kmpc_is_spmd_exec_mode() {
   FunctionTracingRAII();
   return mapping::isSPMDMode();
 }
 }
 
 #pragma omp end declare target

openmp/libomptarget/DeviceRTL/src/State.cpp

 //===------ State.cpp - OpenMP State & ICV interface ------------- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 //===----------------------------------------------------------------------===//
 
 #include "State.h"
-#include "Configuration.h"
 #include "Debug.h"
+#include "Environment.h"
 #include "Interface.h"
 #include "Mapping.h"
 #include "Synchronization.h"
 #include "Types.h"
 #include "Utils.h"
 
 using namespace ompx;
 
 #pragma omp begin declare target device_type(nohost)
 
 /// Memory implementation
 ///
 ///{
 
 /// Add worst-case padding so that future allocations are properly aligned.
 /// FIXME: The stack shouldn't require worst-case padding. Alignment needs to be
 /// passed in as an argument and the stack rewritten to support it.
 constexpr const uint32_t Alignment = 16;
 
 /// External symbol to access dynamic shared memory.
 extern unsigned char DynamicSharedBuffer[] __attribute__((aligned(Alignment)));
 #pragma omp allocate(DynamicSharedBuffer) allocator(omp_pteam_mem_alloc)
 
+/// The kernel environment passed to the init method by the compiler.
+static KernelEnvironmentTy *SHARED(KernelEnvironmentPtr);
+
 namespace {
 
 /// Fallback implementations are missing to trigger a link time error.
 /// Implementations for new devices, including the host, should go into a
 /// dedicated begin/end declare variant.
 ///
 ///{
 
int returnValIfLevelIsActive(int Level, int Val, int DefaultVal,
   int ActiveLevel = icv::ActiveLevel;
   if (OMP_UNLIKELY(Level != ActiveLevel))
     return DefaultVal;
   return Val;
 }
 
 } // namespace
 
-void state::init(bool IsSPMD) {
+void state::init(bool IsSPMD, KernelEnvironmentTy &KernelEnvironment) {
   SharedMemorySmartStack.init(IsSPMD);
   if (mapping::isInitialThreadInLevel0(IsSPMD)) {
     TeamState.init(IsSPMD);
-    DebugEntryRAII::init();
     ThreadStates = nullptr;
-  }
+    KernelEnvironmentPtr = &KernelEnvironment;
+  }
+
+KernelEnvironmentTy &state::getKernelEnvironment() {
+  return *KernelEnvironmentPtr;
 }
 
 void state::enterDataEnvironment(IdentTy *Ident) {
   ASSERT(config::mayUseThreadStates() &&
          "Thread state modified while explicitly disabled!");
   if (!config::mayUseThreadStates())
     return;
 

openmp/libomptarget/include/DeviceEnvironment.h

-//===---- device_environment.h - OpenMP GPU device environment ---- C++ -*-===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-//
-// Global device environment
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef _OMPTARGET_DEVICE_ENVIRONMENT_H_
-#define _OMPTARGET_DEVICE_ENVIRONMENT_H_
-
-// deviceRTL uses <stdint> and DeviceRTL uses explicit definitions
-
-struct DeviceEnvironmentTy {
-  uint32_t DebugKind;
-  uint32_t NumDevices;
-  uint32_t DeviceNum;
-  uint32_t DynamicMemSize;
-};
-
-#endif

openmp/libomptarget/include/Environment.h

+//===------------ Environment.h - OpenMP GPU environments --------- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Environments shared between host and device.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _OMPTARGET_ENVIRONMENT_H_
+#define _OMPTARGET_ENVIRONMENT_H_
+
+#ifdef OMPTARGET_DEVICE_RUNTIME
+#include "Types.h"
+#else
+#include "SourceInfo.h"
+
+#include <cstdint>
+
+using IdentTy = ident_t;
+#endif
+
+#include "llvm/Frontend/OpenMP/OMPDeviceConstants.h"
+
+struct DeviceEnvironmentTy {
+  uint32_t DebugKind;
+  uint32_t NumDevices;
+  uint32_t DeviceNum;
+  uint32_t DynamicMemSize;
+};
+
+// NOTE: Please don't change the order of those members as their indices are
+// used in the middle end. Always add the new data member at the end.
+struct ConfigurationEnvironmentTy {
+  uint8_t UseGenericStateMachine;
+  uint8_t MayUseNestedParallelism;
+  llvm::omp::OMPTgtExecModeFlags ExecMode;
+};
+
+// NOTE: Please don't change the order of those members as their indices are
+// used in the middle end. Always add the new data member at the end.
+struct KernelEnvironmentTy {
+  ConfigurationEnvironmentTy Configuration;
+  IdentTy *Ident;
+  /// Current indentation level for the function trace. Only accessed by thread
+  /// 0.
+  uint16_t DebugIndentionLevel;
+};
+
+#endif // _OMPTARGET_ENVIRONMENT_H_

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

 #include <deque>
 #include <mutex>
 #include <string>
 #include <system_error>
 #include <unistd.h>
 #include <unordered_map>
 
 #include "Debug.h"
-#include "DeviceEnvironment.h"
+#include "Environment.h"
 #include "GlobalHandler.h"
 #include "PluginInterface.h"
 #include "Utilities.h"
 #include "UtilitiesRTL.h"
 #include "omptarget.h"
 
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
Error initImpl(GenericDeviceTy &Device, DeviceImageTy &Image) override {
     // TODO: Read the kernel descriptor for the max threads per block. May be
     // read from the image.
 
     return Plugin::success();
   }
 
   /// Launch the AMDGPU kernel function.
   Error launchImpl(GenericDeviceTy &GenericDevice, uint32_t NumThreads,
-                   uint64_t NumBlocks, 
+                   uint64_t NumBlocks,
                    KernelArgsTy &KernelArgs, void *Args,
                    AsyncInfoWrapperTy &AsyncInfoWrapper) const override;
 
   /// The default number of blocks is common to the whole device.
   uint32_t getDefaultNumBlocks(GenericDeviceTy &GenericDevice) const override {
     return GenericDevice.getDefaultNumBlocks();
   }
 

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

 #include <cstddef>
 #include <cstdint>
 #include <list>
 #include <map>
 #include <shared_mutex>
 #include <vector>
 
 #include "Debug.h"
-#include "DeviceEnvironment.h"
+#include "Environment.h"
 #include "GlobalHandler.h"
 #include "JIT.h"
 #include "MemoryManager.h"
 #include "Utilities.h"
 #include "omptarget.h"
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Triple.h"
protected:
   /// the device I has a matrix PeerAccesses with PeerAccesses[J] == AVAILABLE,
   /// the device I can access device J's memory directly. However, notice this
   /// does not mean that device J can access device I's memory directly.
   llvm::SmallVector<PeerAccessState> PeerAccesses;
   std::mutex PeerAccessesLock;
 
   /// Map of host pinned allocations used for optimize device transfers.
   PinnedAllocationMapTy PinnedAllocs;
+
+private:
+  /// Return the kernel environment object for kernel \p Name.
+  Expected<KernelEnvironmentTy>
+  getKernelEnvironmentForKernel(StringRef Name, DeviceImageTy &Image);
 };
 
 /// Class implementing common functionalities of offload plugins. Each plugin
 /// should define the specific plugin class, derive from this generic one, and
 /// implement the necessary virtual function members.
 struct GenericPluginTy {
 
   /// Construct a plugin instance.

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

Error GenericDeviceTy::registerKernelOffloadEntry(
   // Set the device entry address to the kernel address and store the entry on
   // the entry table.
   DeviceEntry.addr = (void *)Kernel;
   Image.getOffloadEntryTable().addEntry(DeviceEntry);
 
   return Plugin::success();
 }
 
-Expected<OMPTgtExecModeFlags>
-GenericDeviceTy::getExecutionModeForKernel(StringRef Name,
+Expected<KernelEnvironmentTy>
+GenericDeviceTy::getKernelEnvironmentForKernel(StringRef Name,
                                                DeviceImageTy &Image) {
-  // Create a metadata object for the exec mode global (auto-generated).
-  StaticGlobalTy<llvm::omp::OMPTgtExecModeFlags> ExecModeGlobal(Name.data(),
-                                                                "_exec_mode");
+  // Create a metadata object for the kernel environment object.
+  StaticGlobalTy<KernelEnvironmentTy> KernelEnv(Name.data(), "_kernel_info");
 
-  // Retrieve execution mode for the kernel. This may fail since some kernels
-  // may not have an execution mode.
+  // Retrieve kernel environment object for the kernel.
   GenericGlobalHandlerTy &GHandler = Plugin::get().getGlobalHandler();
-  if (auto Err = GHandler.readGlobalFromImage(*this, Image, ExecModeGlobal)) {
+  if (auto Err = GHandler.readGlobalFromImage(*this, Image, KernelEnv)) {
     // Consume the error since it is acceptable to fail.
     [[maybe_unused]] std::string ErrStr = toString(std::move(Err));
-    DP("Failed to read execution mode for '%s': %s\n"
-       "Using default SPMD (2) execution mode\n",
-       Name.data(), ErrStr.data());
+    DP("Failed to read kernel environment object for '%s': %s\n", Name.data(),
+       ErrStr.data());
 
-    return OMP_TGT_EXEC_MODE_SPMD;
+    return createStringError(inconvertibleErrorCode(), ErrStr);
   }
 
+  return KernelEnv.getValue();
+}
+
+Expected<OMPTgtExecModeFlags>
+GenericDeviceTy::getExecutionModeForKernel(StringRef Name,
+                                           DeviceImageTy &Image) {
+  auto KernelEnvOrError = getKernelEnvironmentForKernel(Name, Image);
+  // We error out directly if we can't read the kernel environment object.
+  if (!KernelEnvOrError)
+    return KernelEnvOrError.takeError();
+
+  auto &KernelEnv = *KernelEnvOrError;
+  auto ExecMode = KernelEnv.Configuration.ExecMode;
+
   // Check that the retrieved execution mode is valid.
-  if (!GenericKernelTy::isValidExecutionMode(ExecModeGlobal.getValue()))
-    return Plugin::error("Invalid execution mode %d for '%s'",
-                         ExecModeGlobal.getValue(), Name.data());
+  if (!GenericKernelTy::isValidExecutionMode(ExecMode))
+    return Plugin::error("Invalid execution mode %d for '%s'", ExecMode,
+                         Name.data());
 
-  return ExecModeGlobal.getValue();
+  return ExecMode;
 }
 
 Error PinnedAllocationMapTy::registerHostBuffer(void *HstPtr,
                                                 void *DevAccessiblePtr,
                                                 size_t Size) {
   assert(HstPtr && "Invalid pointer");
   assert(DevAccessiblePtr && "Invalid pointer");
 

openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp

 
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <string>
 #include <unordered_map>
 
 #include "Debug.h"
-#include "DeviceEnvironment.h"
+#include "Environment.h"
 #include "GlobalHandler.h"
 #include "PluginInterface.h"
 
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
 #include "llvm/Support/Error.h"
 

openmp/libomptarget/plugins-nextgen/generic-elf-64bit/src/rtl.cpp

 
 #include <cassert>
 #include <cstddef>
 #include <ffi.h>
 #include <string>
 #include <unordered_map>
 
 #include "Debug.h"
-#include "DeviceEnvironment.h"
+#include "Environment.h"
 #include "GlobalHandler.h"
 #include "PluginInterface.h"
 #include "omptarget.h"
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
 #include "llvm/Support/DynamicLibrary.h"

openmp/libomptarget/plugins/amdgpu/src/rtl.cpp

 #include "ELFSymbols.h"
 #include "impl_runtime.h"
 #include "interop_hsa.h"
 
 #include "UtilitiesRTL.h"
 #include "internal.h"
 #include "rt.h"
 
-#include "DeviceEnvironment.h"
+#include "Environment.h"
 #include "get_elf_mach_gfx_name.h"
 #include "omptargetplugin.h"
 #include "print_tracing.h"
 
 using namespace llvm;
 using namespace llvm::object;
 using namespace llvm::ELF;
 using namespace llvm::omp::target::plugin::utils;

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 #include <cuda.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <string>
 #include <vector>
 
 #include "Debug.h"
-#include "DeviceEnvironment.h"
+#include "Environment.h"
 #include "omptarget.h"
 #include "omptargetplugin.h"
 
 #ifndef TARGET_NAME
 #define TARGET_NAME CUDA
 #endif
 #ifndef DEBUG_PREFIX
 #define DEBUG_PREFIX "Target " GETNAME(TARGET_NAME) " RTL"