[openmp][amdgpu] Make plugin robust to presence of explicit implicit arguments

[JonChesterfield]

Summary

OpenMP (compiler) does not currently request any implicit kernel
arguments. OpenMP (runtime) allocates and initialises a reasonable guess at
the implicit kernel arguments anyway.

This change makes the plugin check the number of explicit arguments, instead
of all arguments, and puts the pointer to hostcall buffer in both the current
location and at the offset expected when implicit arguments are added to the
metadata by D113538.

This is intended to keep things running while fixing the oversight in the
compiler (in D113538). Once that patch lands, and a following one marks
openmp kernels that use printf such that the backend emits an args element
with the right type (instead of hidden_node), the over-allocation can be
removed and the hardcoded 8*e+3 offset replaced with one read from the
.offset of the corresponding metadata element.

[arsenm]

openmp/libomptarget/plugins/amdgpu/impl/system.cpp

495

probably shouldn't touch this line until it's deleted

[JonChesterfield]

openmp/libomptarget/plugins/amdgpu/impl/system.cpp

495

i missed clang-format on a previous patch, but not on this one. Doesn't seem worth going through and putting the whitespace back in

[JonChesterfield]

openmp/libomptarget/plugins/amdgpu/impl/system.cpp

475

There's some evidence suggesting clang won't emit this metadata (or the printf one), it may actually emit

{
".address_space" : "global",
".offset" : 24,
".size" : 8,
".value_kind" : "hidden_none",
},

as a placeholder to indicate it doesn't know the print/hostcall thing is needed

[JonChesterfield]

• use positional encoding, clang isn't setting print or hostcall metadata

[arsenm]

This makes my assert problem go away

[JonChesterfield]

Win for shotgun debugging. The hostcall init changes are inert in trunk so if this got through internal testing we can land it, then your fix, then delete the handling for the current clang output.

[estewart08]

LGTM, passed CI testing.

openmp/libomptarget/plugins/amdgpu/impl/internal.h

typedef struct atl_kernel_info_s {
   uint64_t kernel_object;
   uint32_t group_segment_size;
   uint32_t private_segment_size;
   uint32_t sgpr_count;
   uint32_t vgpr_count;
   uint32_t sgpr_spill_count;
   uint32_t vgpr_spill_count;
   uint32_t kernel_segment_size;
-  uint32_t num_args;
+  uint32_t explicit_argument_count; // excluding implicit arguments
+  uint32_t hostcall_pointer_offset; // UINT32_MAX if not used
 } atl_kernel_info_t;
 
 typedef struct atl_symbol_info_s {
   uint64_t addr;
   uint32_t size;
 } atl_symbol_info_t;
 
 // ---------------------- Kernel End -------------

openmp/libomptarget/plugins/amdgpu/impl/system.cpp

enum class ValueKind {
     HiddenDefaultQueue,
     HiddenCompletionAction,
     HiddenMultiGridSyncArg,
     HiddenHostcallBuffer,
     Unknown
   };
 
   KernelArgMD()
-      : name_(std::string()),  size_(0), offset_(0),
+      : name_(std::string()), size_(0), offset_(0),
         valueKind_(ValueKind::Unknown) {}
 
   // fields
   std::string name_;
   uint32_t size_;
   uint32_t offset_;
   ValueKind valueKind_;
 };
for (size_t i = 0; i < kernelsSize; i++) {
 
     // Make sure that kernelName + ".kd" == symbolName
     if ((kernelName + ".kd") != symbolName) {
       printf("[%s:%d] Kernel name mismatching symbol: %s != %s + .kd\n",
              __FILE__, __LINE__, symbolName.c_str(), kernelName.c_str());
       return HSA_STATUS_ERROR_INVALID_CODE_OBJECT;
     }
 
-    atl_kernel_info_t info = {0, 0, 0, 0, 0, 0, 0, 0, 0};
+    atl_kernel_info_t info = {0, 0, 0, 0, 0, 0, 0, 0, 0, UINT32_MAX};
 
     uint64_t sgpr_count, vgpr_count, sgpr_spill_count, vgpr_spill_count;
     msgpack_errors += map_lookup_uint64_t(element, ".sgpr_count", &sgpr_count);
     if (msgpack_errors != 0) {
       printf("[%s:%d] %s failed\n", __FILE__, __LINE__,
              "sgpr count metadata lookup in kernel metadata");
       return HSA_STATUS_ERROR_INVALID_CODE_OBJECT;
     }
if (kernel_segment_size > 0) {
       msgpack_errors +=
           map_lookup_array(element, ".args", &args_array, &argsSize);
       if (msgpack_errors != 0) {
         printf("[%s:%d] %s failed\n", __FILE__, __LINE__,
                "kernel args metadata lookup in kernel metadata");
         return HSA_STATUS_ERROR_INVALID_CODE_OBJECT;
       }
 
-      info.num_args = argsSize;
-
+      uint32_t explicit_argument_count = 0;
       for (size_t i = 0; i < argsSize; ++i) {
         KernelArgMD lcArg;
 
         msgpack::byte_range args_element;
         msgpack_errors += array_lookup_element(args_array, i, &args_element);
         if (msgpack_errors != 0) {
           printf("[%s:%d] %s failed\n", __FILE__, __LINE__,
                  "iterate args map in kernel args metadata");
if (kernel_segment_size > 0) {
         // v3 has offset field and not align field
         size_t new_offset = lcArg.offset_;
         size_t padding = new_offset - offset;
         offset = new_offset;
         DP("Arg[%lu] \"%s\" (%u, %u)\n", i, lcArg.name_.c_str(), lcArg.size_,
            lcArg.offset_);
         offset += lcArg.size_;
 
+        if (lcArg.valueKind_ == KernelArgMD::ValueKind::HiddenHostcallBuffer) {
+          if (info.hostcall_pointer_offset != UINT32_MAX) {
+            DP("[%s: Multiple hostcall buffers detected\n", kernelName.c_str());
+          }
+          info.hostcall_pointer_offset = lcArg.offset_;
+        }
+
         // check if the arg is a hidden/implicit arg
         // this logic assumes that all hidden args are 8-byte aligned
         if (!isImplicit(lcArg.valueKind_)) {
+          explicit_argument_count++;
           kernel_explicit_args_size += lcArg.size_;
         } else {
           hasHiddenArgs = true;
         }
         kernel_explicit_args_size += padding;
       }
+      info.explicit_argument_count = explicit_argument_count;
     }
 
-    // TODO: Probably don't want this arithmetic 
+    // TODO: Probably don't want this arithmetic
     info.kernel_segment_size =
         (hasHiddenArgs ? kernel_explicit_args_size : kernel_segment_size);
     DP("[%s: kernarg seg size] (%lu --> %u)\n", kernelName.c_str(),
        kernel_segment_size, info.kernel_segment_size);
 
     // kernel received, now add it to the kernel info table
     KernelInfoTable[kernelName] = info;
   }

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

int32_t __tgt_rtl_run_target_team_region_locked(
   const atl_kernel_info_t KernelInfoEntry =
       KernelInfoTable[device_id][kernel_name];
   const uint32_t group_segment_size = KernelInfoEntry.group_segment_size;
   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.num_args);
+  assert(arg_num == (int)KernelInfoEntry.explicit_argument_count);
 
   /*
    * Set limit based on ThreadsPerGroup and GroupsPerDevice
    */
   launchVals LV =
       getLaunchVals(DeviceInfo.WarpSize[device_id], DeviceInfo.Env,
                     KernelInfo->ConstWGSize, KernelInfo->ExecutionMode,
                     num_teams,      // From run_region arg
if (!ArgPool) {
       impl_args->offset_z = 0;
 
       // assign a hostcall buffer for the selected Q
       if (__atomic_load_n(&DeviceInfo.hostcall_required, __ATOMIC_ACQUIRE)) {
         // hostrpc_assign_buffer is not thread safe, and this function is
         // under a multiple reader lock, not a writer lock.
         static pthread_mutex_t hostcall_init_lock = PTHREAD_MUTEX_INITIALIZER;
         pthread_mutex_lock(&hostcall_init_lock);
-        impl_args->hostcall_ptr = hostrpc_assign_buffer(
+        unsigned long buffer = hostrpc_assign_buffer(
             DeviceInfo.HSAAgents[device_id], queue, device_id);
         pthread_mutex_unlock(&hostcall_init_lock);
-        if (!impl_args->hostcall_ptr) {
+        if (!buffer) {
           DP("hostrpc_assign_buffer failed, gpu would dereference null and "
              "error\n");
           return OFFLOAD_FAIL;
         }
+
+        uint32_t HostcallOffset = KernelInfoEntry.hostcall_pointer_offset;
+        if (HostcallOffset != UINT32_MAX) {
+          if (HostcallOffset > (ArgPool->kernarg_segment_size - 8)) {
+            DP("Bad offset of hostcall, exceeds kernarg segment size\n");
+          } else {
+            memcpy(static_cast<char *>(kernarg) + HostcallOffset, &buffer, 8);
+          }
+        }
+        impl_args->hostcall_ptr = buffer;
       }
 
       packet->kernarg_address = kernarg;
     }
 
     hsa_signal_t s = DeviceInfo.FreeSignalPool.pop();
     if (s.handle == 0) {
       DP("Failed to get signal instance\n");