[Libomptarget] Add basic support for dynamic shared memory on AMDGPU

[jhuber6]

Summary

This patchs adds the arguments necessary to allocate the size of the
dynamic shared memory via the LIBOMPTARGET_SHARED_MEMORY_SIZE
environment variable. This patch only allocates the memory, AMDGPU has a
limitation that shared memory can only be accessed from the kernel
directly. So this will currently only work with optimizations to inline
the accessor function.

[jdoerfert]

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

1189

+ KernelInfoEntry.group_segment_size;

[JonChesterfield]

Looks about as expected, modulo Johannes' comment above. Could we do the getenv call once, instead of per-launch? Somewhere in the massive class constructor perhaps.

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

1189

Yep, need the sum. Possibly worth checking the stoi result is positive and < 64k or so.

[jhuber6]

Updating to use environment structure and append instead of write.

[JonChesterfield]

LGTM, thanks!

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

 } // namespace
 } // namespace core
 
 struct EnvironmentVariables {
   int NumTeams;
   int TeamLimit;
   int TeamThreadLimit;
   int MaxTeamsDefault;
+  int DynamicMemSize;
 };
 
 template <uint32_t wavesize>
 static constexpr const llvm::omp::GV &getGridValue() {
   return llvm::omp::getAMDGPUGridValues<wavesize>();
 }
 
 struct HSALifetime {
assert(DeviceId >= 0 && DeviceId < DeviceCoarseGrainedMemoryPools.size() &&
            "Invalid device Id");
     return DeviceCoarseGrainedMemoryPools[DeviceId];
   }
 
   hsa_amd_memory_pool_t getHostMemoryPool() {
     return HostFineGrainedMemoryPool;
   }
 
-  static int readEnvElseMinusOne(const char *Env) {
+  static int readEnv(const char *Env, int Default = -1) {
     const char *envStr = getenv(Env);
-    int res = -1;
+    int res = Default;
     if (envStr) {
       res = std::stoi(envStr);
       DP("Parsed %s=%d\n", Env, res);
     }
     return res;
   }
 
   RTLDeviceInfoTy() {
for (int i = 0; i < NumberOfDevices; i++) {
       ThreadsPerGroup[i] = RTLDeviceInfoTy::Default_WG_Size;
       GroupsPerDevice[i] = RTLDeviceInfoTy::DefaultNumTeams;
       ComputeUnits[i] = 1;
       DP("Device %d: Initial groupsPerDevice %d & threadsPerGroup %d\n", i,
          GroupsPerDevice[i], ThreadsPerGroup[i]);
     }
 
     // Get environment variables regarding teams
-    Env.TeamLimit = readEnvElseMinusOne("OMP_TEAM_LIMIT");
-    Env.NumTeams = readEnvElseMinusOne("OMP_NUM_TEAMS");
-    Env.MaxTeamsDefault = readEnvElseMinusOne("OMP_MAX_TEAMS_DEFAULT");
-    Env.TeamThreadLimit = readEnvElseMinusOne("OMP_TEAMS_THREAD_LIMIT");
+    Env.TeamLimit = readEnv("OMP_TEAM_LIMIT");
+    Env.NumTeams = readEnv("OMP_NUM_TEAMS");
+    Env.MaxTeamsDefault = readEnv("OMP_MAX_TEAMS_DEFAULT");
+    Env.TeamThreadLimit = readEnv("OMP_TEAMS_THREAD_LIMIT");
+    Env.DynamicMemSize = readEnv("LIBOMPTARGET_SHARED_MEMORY_SIZE", 0);
 
     // Default state.
     RequiresFlags = OMP_REQ_UNDEFINED;
 
     ConstructionSucceeded = true;
   }
 
   ~RTLDeviceInfoTy() {
int32_t runRegionLocked(int32_t device_id, void *tgt_entry_ptr, void **tgt_args,
   if (KernelInfoTable[device_id].find(kernel_name) ==
       KernelInfoTable[device_id].end()) {
     DP("Kernel %s not found\n", kernel_name.c_str());
     return OFFLOAD_FAIL;
   }
 
   const atl_kernel_info_t KernelInfoEntry =
       KernelInfoTable[device_id][kernel_name];
-  const uint32_t group_segment_size = KernelInfoEntry.group_segment_size;
+  const uint32_t group_segment_size =
+      KernelInfoEntry.group_segment_size + DeviceInfo.Env.DynamicMemSize;
   const uint32_t sgpr_count = KernelInfoEntry.sgpr_count;
   const uint32_t vgpr_count = KernelInfoEntry.vgpr_count;
   const uint32_t sgpr_spill_count = KernelInfoEntry.sgpr_spill_count;
   const uint32_t vgpr_spill_count = KernelInfoEntry.vgpr_spill_count;
 
   assert(arg_num == (int)KernelInfoEntry.explicit_argument_count);
 
   /*
// Run on the device.
     packet->workgroup_size_x = WorkgroupSize;
     packet->workgroup_size_y = 1;
     packet->workgroup_size_z = 1;
     packet->reserved0 = 0;
     packet->grid_size_x = GridSize;
     packet->grid_size_y = 1;
     packet->grid_size_z = 1;
     packet->private_segment_size = KernelInfoEntry.private_segment_size;
-    packet->group_segment_size = KernelInfoEntry.group_segment_size;
+    packet->group_segment_size = group_segment_size;
     packet->kernel_object = KernelInfoEntry.kernel_object;
     packet->kernarg_address = 0;     // use the block allocator
     packet->reserved2 = 0;           // impl writes id_ here
     packet->completion_signal = {0}; // may want a pool of signals
 
     KernelArgPool *ArgPool = nullptr;
     void *kernarg = nullptr;
     {
       auto it = KernelArgPoolMap.find(std::string(KernelInfo->Name));
struct device_environment {
 
   DeviceEnvironmentTy host_device_env;
   symbol_info si;
   bool valid = false;
 
   __tgt_device_image *image;
   const size_t img_size;
 
-  device_environment(int device_id, int number_devices,
+  device_environment(int device_id, int number_devices, int dynamic_mem_size,
                      __tgt_device_image *image, const size_t img_size)
       : image(image), img_size(img_size) {
 
     host_device_env.NumDevices = number_devices;
     host_device_env.DeviceNum = device_id;
     host_device_env.DebugKind = 0;
-    host_device_env.DynamicMemSize = 0;
+    host_device_env.DynamicMemSize = dynamic_mem_size;
     if (char *envStr = getenv("LIBOMPTARGET_DEVICE_RTL_DEBUG")) {
       host_device_env.DebugKind = std::stoi(envStr);
     }
 
     int rc = get_symbol_info_without_loading((char *)image->ImageStart,
                                              img_size, sym(), &si);
     if (rc != 0) {
       DP("Finding global device environment '%s' - symbol missing.\n", sym());
__tgt_target_table *__tgt_rtl_load_binary_locked(int32_t device_id,
   // We do not need to set the ELF version because the caller of this function
   // had to do that to decide the right runtime to use
 
   if (!elf_machine_id_is_amdgcn(image)) {
     return NULL;
   }
 
   {
-    auto env = device_environment(device_id, DeviceInfo.NumberOfDevices, image,
-                                  img_size);
+    auto env =
+        device_environment(device_id, DeviceInfo.NumberOfDevices,
+                           DeviceInfo.Env.DynamicMemSize, image, img_size);
 
     auto &KernelInfo = DeviceInfo.KernelInfoTable[device_id];
     auto &SymbolInfo = DeviceInfo.SymbolInfoTable[device_id];
     hsa_status_t err = module_register_from_memory_to_place(
         KernelInfo, SymbolInfo, (void *)image->ImageStart, img_size, device_id,
         [&](void *data, size_t size) {
           if (image_contains_symbol(data, size, "needs_hostcall_buffer")) {
             __atomic_store_n(&DeviceInfo.hostcall_required, true,

openmp/libomptarget/test/api/omp_dynamic_shared_memory_amdgpu.c

+// RUN: %libomptarget-compile-amdgcn-amd-amdhsa -O1 -mllvm -openmp-opt-inline-device
+// RUN: env LIBOMPTARGET_SHARED_MEMORY_SIZE=256 \
+// RUN:   %libomptarget-run-amdgcn-amd-amdhsa | %fcheck-amdgcn-amd-amdhsa
+// REQUIRES: amdgcn-amd-amdhsa
+
+#include <omp.h>
+#include <stdio.h>
+
+int main() {
+  int x;
+#pragma omp target parallel map(from : x)
+  {
+    int *buf = llvm_omp_target_dynamic_shared_alloc() + 252;
+#pragma omp barrier
+    if (omp_get_thread_num() == 0)
+      *buf = 1;
+#pragma omp barrier
+    if (omp_get_thread_num() == 1)
+      x = *buf;
+  }
+
+  // CHECK: PASS
+  if (x == 1 && llvm_omp_target_dynamic_shared_alloc() == NULL)
+    printf("PASS\n");
+}