[OpenMP][NFCI] Introduce the kernel environment for target regions

[jdoerfert]

Summary

This introduces the kernel environment which contains information passed
by the compiler to a GPU kernel. For now it just encapsulated the
ident_t object that we passed explicitly before but we will add more
content, including mutable content, later on.

[jdoerfert]

Move more values into the kernel environment (and configuration environment), update tests

[ye-luo]

the kernel environment which contains information passed by the compiler to a GPU kernel.

DId you mean this environment is baked into kernel at compile time? So there is no additional H2D cost at each call, right?

[jdoerfert]

In D113140#3123162, @ye-luo wrote:

the kernel environment which contains information passed by the compiler to a GPU kernel.

DId you mean this environment is baked into kernel at compile time? So there is no additional H2D cost at each call, right?

Yes, this environment is a constant fed into the kernel at compile time.
It will allow us to pass more information, e.g., if you want to support sequentialiazed parallel regions, in a way that caused the "feature" to be removed during compilation.
It also allows us to communicate dynamic information in a way that will make us pay for it only if we "enable" it, e.g., if you don't want printf support on AMDGPU you don't have to pay for a buffer.

[jdoerfert]

Split up, main content starts with D116908.

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

 #define OMP_STRUCT_TYPE(VarName, StructName, ...)
 #endif
 
 #define __OMP_STRUCT_TYPE(VarName, Name, ...)                                  \
   OMP_STRUCT_TYPE(VarName, "struct." #Name, __VA_ARGS__)
 
 __OMP_STRUCT_TYPE(Ident, ident_t, Int32, Int32, Int32, Int32, Int8Ptr)
 __OMP_STRUCT_TYPE(AsyncInfo, __tgt_async_info, Int8Ptr)
+__OMP_STRUCT_TYPE(KernelEnvironmentTy, KernelEnvironmentTy, Ident)
 
 #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, Int1)
-__OMP_RTL(__kmpc_target_deinit, false, Void, IdentPtr, Int8, Int1)
+__OMP_RTL(__kmpc_target_init, false, Int32, KernelEnvironmentTyPtr, Int8, Int1, Int1)
+__OMP_RTL(__kmpc_target_deinit, false, Void, Int8, 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_shuffle_int32, false, Int32, Int32, Int16, Int16)

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

OpenMPIRBuilder::createTargetInit(const LocationDescription &Loc, bool IsSPMD,
   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});
+  Function *Kernel = Builder.GetInsertBlock()->getParent();
+  const DataLayout &DL = Fn->getParent()->getDataLayout();
+  Constant *KernelEnvironmentInitializer = ConstantStruct::get(
+      KernelEnvironmentTy, {cast<GlobalVariable>(Ident)->getInitializer()});
+  std::string KernelEnvironmentName =
+      (Kernel->getName() + "_kernel_info").str();
+  GlobalVariable *KernelEnvironment = new GlobalVariable(
+      M, KernelEnvironmentTy, /*IsConstant*/ false,
+      llvm::GlobalValue::ExternalLinkage, KernelEnvironmentInitializer,
+      KernelEnvironmentName,
+      /*InsertBefore=*/nullptr, llvm::GlobalValue::NotThreadLocal,
+      DL.getDefaultGlobalsAddressSpace());
+  auto *KernelEnvironmentCasted =
+      ConstantExpr::getPointerBitCastOrAddrSpaceCast(KernelEnvironment,
+                                                     KernelEnvironmentTyPtr);
+
+  CallInst *ThreadKind =
+      Builder.CreateCall(Fn, {KernelEnvironmentCasted, IsSPMDVal,
+                              UseGenericStateMachine, RequiresFullRuntimeVal});
 
   Value *ExecUserCode = Builder.CreateICmpEQ(
       ThreadKind, ConstantInt::get(ThreadKind->getType(), -1),
       "exec_user_code");
 
   // ThreadKind = __kmpc_target_init(...)
   // if (ThreadKind == -1)
   //   user_code

 }
 
 void OpenMPIRBuilder::createTargetDeinit(const LocationDescription &Loc,
                                          bool IsSPMD,
                                          bool RequiresFullRuntime) {
   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);
   ConstantInt *RequiresFullRuntimeVal =
       ConstantInt::getBool(Int32->getContext(), RequiresFullRuntime);
 
   Function *Fn = getOrCreateRuntimeFunctionPtr(
       omp::RuntimeFunction::OMPRTL___kmpc_target_deinit);
 
-  Builder.CreateCall(Fn, {Ident, IsSPMDVal, RequiresFullRuntimeVal});
+  Builder.CreateCall(Fn, {IsSPMDVal, RequiresFullRuntimeVal});
 }
 
 std::string OpenMPIRBuilder::getNameWithSeparators(ArrayRef<StringRef> Parts,
                                                    StringRef FirstSeparator,
                                                    StringRef Separator) {
   SmallString<128> Buffer;
   llvm::raw_svector_ostream OS(Buffer);
   StringRef Sep = FirstSeparator;

openmp/libomptarget/DeviceRTL/CMakeLists.txt

set(src_files
   ${source_directory}/Reduction.cpp
   ${source_directory}/State.cpp
   ${source_directory}/Synchronization.cpp
   ${source_directory}/Tasking.cpp
   ${source_directory}/Utils.cpp
   ${source_directory}/Workshare.cpp
 )
 
-set(clang_opt_flags -O1 -mllvm -openmp-opt-disable -DSHARED_SCRATCHPAD_SIZE=2048)
+set(clang_opt_flags -O1 -mllvm -openmp-opt-disable -DSHARED_SCRATCHPAD_SIZE=2048 -DOMPTARGET_DEVICE_RUNTIME)
 set(link_opt_flags  -O1        -openmp-opt-disable)
 
 # Prepend -I to each list element
 set (LIBOMPTARGET_LLVM_INCLUDE_DIRS_DEVICERTL "${LIBOMPTARGET_LLVM_INCLUDE_DIRS}")
 list(TRANSFORM LIBOMPTARGET_LLVM_INCLUDE_DIRS_DEVICERTL PREPEND "-I")
 
 # Set flags for LLVM Bitcode compilation.
 set(bc_flags -S -x c++ -std=c++17

openmp/libomptarget/DeviceRTL/include/Interface.h

 //===-------- Interface.h - OpenMP 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
 //
 //===----------------------------------------------------------------------===//
 //
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef OMPTARGET_DEVICERTL_INTERFACE_H
 #define OMPTARGET_DEVICERTL_INTERFACE_H
 
+#include "KernelEnvironment.h"
 #include "Types.h"
 
 /// External API
 ///
 ///{
 
 extern "C" {
 

 /// External interface to get the number of threads.
 uint32_t __kmpc_get_hardware_num_threads_in_block();
 
 /// Kernel
 ///
 ///{
 int8_t __kmpc_is_spmd_exec_mode();
 
-int32_t __kmpc_target_init(IdentTy *Ident, int8_t Mode,
+int32_t __kmpc_target_init(_OMP::KernelEnvironmentTy KernelEnv, int8_t Mode,
                            bool UseGenericStateMachine, bool);
 
-void __kmpc_target_deinit(IdentTy *Ident, int8_t Mode, bool);
+void __kmpc_target_deinit(int8_t Mode, bool);
 
 ///}
 
 /// Reduction
 ///
 ///{
 void __kmpc_nvptx_end_reduce(int32_t TId);
 

openmp/libomptarget/DeviceRTL/include/State.h

 //===-------- State.h - 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
 //
 //===----------------------------------------------------------------------===//
 //
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef OMPTARGET_STATE_H
 #define OMPTARGET_STATE_H
 
 #include "Debug.h"
+#include "KernelEnvironment.h"
 #include "Types.h"
 
 #pragma omp declare target
 
 namespace _OMP {
 
 namespace state {
 
 inline constexpr uint32_t SharedScratchpadSize = SHARED_SCRATCHPAD_SIZE;
 
 /// Initialize the state machinery. Must be called by all threads.
-void init(bool IsSPMD);
+void init(bool IsSPMD, KernelEnvironmentTy &KernelEnv);
+
+/// 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/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 "Interface.h"
+#include "KernelEnvironment.h"
 #include "Mapping.h"
 #include "State.h"
 #include "Synchronization.h"
 #include "Types.h"
 
 using namespace _OMP;
 
 #pragma omp declare target
 
-static void inititializeRuntime(bool IsSPMD) {
+static void inititializeRuntime(bool IsSPMD, KernelEnvironmentTy &KernelEnv) {
   // Order is important here.
   synchronize::init(IsSPMD);
   mapping::init(IsSPMD);
-  state::init(IsSPMD);
+  state::init(IsSPMD, KernelEnv);
 }
 
 /// Simple generic state machine for worker threads.
-static void genericStateMachine(IdentTy *Ident) {
+static void genericStateMachine(IdentTy &Ident) {
   FunctionTracingRAII();
 
   uint32_t TId = mapping::getThreadIdInBlock();
 
   do {
     ParallelRegionFnTy WorkFn = 0;
 
     // Wait for the signal that we have a new work function.
do {
 
     synchronize::threads();
 
   } while (true);
 }
 
 extern "C" {
 
-/// Initialization
-///
-/// \param Ident               Source location identification, can be NULL.
-///
-int32_t __kmpc_target_init(IdentTy *Ident, int8_t Mode,
+int32_t __kmpc_target_init(KernelEnvironmentTy KernelEnv, int8_t Mode,
                            bool UseGenericStateMachine, bool) {
   FunctionTracingRAII();
   const bool IsSPMD = Mode & OMP_TGT_EXEC_MODE_SPMD;
   if (IsSPMD) {
-    inititializeRuntime(/* IsSPMD */ true);
+    inititializeRuntime(/* IsSPMD */ true, KernelEnv);
     synchronize::threadsAligned();
   } else {
-    inititializeRuntime(/* IsSPMD */ false);
+    inititializeRuntime(/* IsSPMD */ false, KernelEnv);
     // 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;
   }
 
   if (mapping::isInitialThreadInLevel0(IsSPMD))
     return -1;
 
   if (UseGenericStateMachine)
-    genericStateMachine(Ident);
+    genericStateMachine(KernelEnv.Ident);
 
   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, bool) {
+void __kmpc_target_deinit(int8_t Mode, bool) {
   FunctionTracingRAII();
   const bool IsSPMD = Mode & OMP_TGT_EXEC_MODE_SPMD;
   state::assumeInitialState(IsSPMD);
   if (IsSPMD)
     return;
 
   // Signal the workers to exit the state machine and exit the kernel.
   state::ParallelRegionFn = nullptr;

openmp/libomptarget/DeviceRTL/src/State.cpp

 
 /// Add worst-case padding so that future allocations are properly aligned.
 constexpr const uint32_t Alignment = 8;
 
 /// 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.
 ///
 ///{
 
void *&state::lookupPtr(ValueKind Kind, bool IsReadonly) {
   case state::VK_ParallelRegionFn:
     return TeamState.ParallelRegionFnVar;
   default:
     break;
   }
   __builtin_unreachable();
 }
 
-void state::init(bool IsSPMD) {
+void state::init(bool IsSPMD, KernelEnvironmentTy &KernelEnv) {
   SharedMemorySmartStack.init(IsSPMD);
-  if (mapping::isInitialThreadInLevel0(IsSPMD))
+  if (mapping::isInitialThreadInLevel0(IsSPMD)) {
     TeamState.init(IsSPMD);
+    KernelEnvironmentPtr = &KernelEnv;
+  }
 
   ThreadStates[mapping::getThreadIdInBlock()] = nullptr;
 }
 
+KernelEnvironmentTy &state::getKernelEnvironment() {
+  return *KernelEnvironmentPtr;
+}
+
 void state::enterDataEnvironment() {
   unsigned TId = mapping::getThreadIdInBlock();
   ThreadStateTy *NewThreadState =
       static_cast<ThreadStateTy *>(__kmpc_alloc_shared(sizeof(ThreadStateTy)));
   NewThreadState->init(ThreadStates[TId]);
   ThreadStates[TId] = NewThreadState;
 }
 

openmp/libomptarget/deviceRTLs/amdgcn/CMakeLists.txt

set(cu_cmd ${CLANG_TOOL}
     -emit-llvm
     -Xclang -aux-triple -Xclang ${aux_triple}
     -fopenmp -fopenmp-cuda-mode -Xclang -fopenmp-is-device
     -D__AMDGCN__
     -Xclang -target-cpu -Xclang ${mcpu}
     -fvisibility=default
     -Wno-unused-value
     -nogpulib
+    -DOMPTARGET_OLD_DEVICE_RUNTIME
     -O${optimization_level}
     ${CUDA_DEBUG}
     -I${CMAKE_CURRENT_SOURCE_DIR}/src
     -I${devicertl_base_directory}/common/include
     -I${devicertl_base_directory}
     -I${devicertl_base_directory}/../include
     ${LIBOMPTARGET_LLVM_INCLUDE_DIRS_AMDGCN})
 

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

do {
     }
 
     __kmpc_barrier_simple_spmd(Ident, TId);
 
   } while (true);
 }
 
 EXTERN
-int32_t __kmpc_target_init(ident_t *Ident, int8_t Mode,
+int32_t __kmpc_target_init(KernelEnvironmentTy KernelEnv, int8_t Mode,
                            bool UseGenericStateMachine,
                            bool RequiresFullRuntime) {
   const bool IsSPMD = Mode & OMP_TGT_EXEC_MODE_SPMD;
   int TId = __kmpc_get_hardware_thread_id_in_block();
   if (IsSPMD)
     __kmpc_spmd_kernel_init(RequiresFullRuntime);
   else
     __kmpc_generic_kernel_init();
 
    if (IsSPMD) {
-    __kmpc_barrier_simple_spmd(Ident, TId);
+    __kmpc_barrier_simple_spmd(&KernelEnv.Ident, TId);
      return -1;
    }
 
    if (TId == GetMasterThreadID())
      return -1;
 
   if (UseGenericStateMachine)
-    __kmpc_target_region_state_machine(Ident);
+    __kmpc_target_region_state_machine(&KernelEnv.Ident);
 
   return TId;
 }
 
 EXTERN
-void __kmpc_target_deinit(ident_t *Ident, int8_t Mode,
+void __kmpc_target_deinit(int8_t Mode,
                           bool RequiresFullRuntime) {
   const bool IsSPMD = Mode & OMP_TGT_EXEC_MODE_SPMD;
   if (IsSPMD)
     __kmpc_spmd_kernel_deinit(RequiresFullRuntime);
   else
     __kmpc_generic_kernel_deinit();
 }
 
 #pragma omp end declare target

openmp/libomptarget/deviceRTLs/interface.h

EXTERN void __kmpc_taskloop(kmp_Ident *loc, uint32_t global_tid,
                             int32_t sched, uint64_t grainsize, void *task_dup);
 
 // 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);
 
+#include "KernelEnvironment.h"
 // non standard
-EXTERN int32_t __kmpc_target_init(ident_t *Ident, int8_t Mode,
+EXTERN int32_t __kmpc_target_init(KernelEnvironmentTy, int8_t Mode,
                                   bool UseGenericStateMachine,
                                   bool RequiresFullRuntime);
-EXTERN void __kmpc_target_deinit(ident_t *Ident, int8_t Mode,
+EXTERN void __kmpc_target_deinit(int8_t Mode,
                                  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_end_sharing_variables();

openmp/libomptarget/deviceRTLs/nvptx/CMakeLists.txt

else()
   list(APPEND bc_flags -DOMPTARGET_NVPTX_DEBUG=0)
 endif()
 
 # Create target to build all Bitcode libraries.
 add_custom_target(omptarget-nvptx-bc)
 
 # Generate a Bitcode library for all the compute capabilities the user requested
 foreach(sm ${nvptx_sm_list})
-  set(cuda_flags -Xclang -target-cpu -Xclang sm_${sm} "-D__CUDA_ARCH__=${sm}0")
+  set(cuda_flags -Xclang -target-cpu -Xclang sm_${sm} "-D__CUDA_ARCH__=${sm}0" -DOMPTARGET_OLD_DEVICE_RUNTIME)
   set(bc_files "")
   foreach(src ${cuda_src_files})
     get_filename_component(infile ${src} ABSOLUTE)
     get_filename_component(outfile ${src} NAME)
     set(outfile "${outfile}-sm_${sm}.bc")
 
     add_custom_command(OUTPUT ${outfile}
       COMMAND ${cuda_compiler} ${bc_flags}

openmp/libomptarget/include/KernelEnvironment.h

+//===---- KernelEnvironment.h - OpenMP GPU kernel 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
+//
+//===----------------------------------------------------------------------===//
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef OMPTARGET_KERNEL_ENVIRONMENT_H
+#define OMPTARGET_KERNEL_ENVIRONMENT_H
+
+// deviceRTL uses <stdint> and DeviceRTL uses explicit definitions
+
+#ifdef OMPTARGET_DEVICE_RUNTIME
+#include "Types.h"
+#else
+#ifdef OMPTARGET_OLD_DEVICE_RUNTIME
+#include "interface.h"
+using IdentTy = ident_t;
+#else
+#include <SourceInfo.h>
+using IdentTy = ident_t;
+#endif
+#endif
+
+#ifdef OMPTARGET_DEVICE_RUNTIME
+namespace _OMP {
+#endif
+
+struct KernelEnvironmentTy {
+  IdentTy Ident;
+};
+
+#ifdef OMPTARGET_DEVICE_RUNTIME
+} // namespace _OMP
+#endif
+
+#endif