[libomptarget] Add register usage info to kernel metadata

[dhruvachak]

Summary

Add register usage information to the runtime metadata so that it can be used during kernel launch (that change will be in a different commit). Add this information to the kernel trace.

[JonChesterfield]

Looks good to me, thanks.

We should probably use uint64_t everywhere, instead of sometimes truncating to uint32_t, but that pattern and the 0,0,0,0, one are pre-existing.

Let's go with this and report errors on implausible (e.g. > 4 billion) register counts as a separate patch, along with sanity checking requested LDS etc.

[dhruvachak]

In D98829#2634094, @JonChesterfield wrote:

Looks good to me, thanks.

We should probably use uint64_t everywhere, instead of sometimes truncating to uint32_t, but that pattern and the 0,0,0,0, one are pre-existing.

Let's go with this and report errors on implausible (e.g. > 4 billion) register counts as a separate patch, along with sanity checking requested LDS etc.

Thanks, Jon. I don't have commit access. Can you please land this change?

[JonChesterfield]

Sure. The process for requesting access has changed since I did it, can't remember what it is now.

[JonChesterfield]

Ah, apologies. My 'land it' script has clobbered the author field. Will need to look up how to not do that in future

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

 
 #define ATMI_WAIT_STATE HSA_WAIT_STATE_BLOCKED
 
 // ---------------------- Kernel Start -------------
 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;
   std::vector<uint64_t> arg_alignments;
   std::vector<uint64_t> arg_offsets;
   std::vector<uint64_t> arg_sizes;
 } atl_kernel_info_t;
 
 typedef struct atl_symbol_info_s {

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

for (size_t i = 0; i < kernelsSize; i++) {
     msgpack::byte_range element;
     msgpack_errors += array_lookup_element(kernel_array, i, &element);
     msgpackErrorCheck(element lookup in kernel metadata, msgpack_errors);
 
     msgpack_errors += map_lookup_string(element, ".name", &kernelName);
     msgpack_errors += map_lookup_string(element, ".symbol", &symbolName);
     msgpackErrorCheck(strings lookup in kernel metadata, msgpack_errors);
 
-    atl_kernel_info_t info = {0, 0, 0, 0, 0, {}, {}, {}};
+    atl_kernel_info_t info = {0, 0, 0, 0, 0, 0, 0, 0, 0, {}, {}, {}};
+
+    uint64_t sgpr_count, vgpr_count, sgpr_spill_count, vgpr_spill_count;
+    msgpack_errors += map_lookup_uint64_t(element, ".sgpr_count", &sgpr_count);
+    msgpackErrorCheck(sgpr count metadata lookup in kernel metadata,
+                      msgpack_errors);
+    info.sgpr_count = sgpr_count;
+
+    msgpack_errors += map_lookup_uint64_t(element, ".vgpr_count", &vgpr_count);
+    msgpackErrorCheck(vgpr count metadata lookup in kernel metadata,
+                      msgpack_errors);
+    info.vgpr_count = vgpr_count;
+
+    msgpack_errors +=
+        map_lookup_uint64_t(element, ".sgpr_spill_count", &sgpr_spill_count);
+    msgpackErrorCheck(sgpr spill count metadata lookup in kernel metadata,
+                      msgpack_errors);
+    info.sgpr_spill_count = sgpr_spill_count;
+
+    msgpack_errors +=
+        map_lookup_uint64_t(element, ".vgpr_spill_count", &vgpr_spill_count);
+    msgpackErrorCheck(vgpr spill count metadata lookup in kernel metadata,
+                      msgpack_errors);
+    info.vgpr_spill_count = vgpr_spill_count;
+
     size_t kernel_explicit_args_size = 0;
     uint64_t kernel_segment_size;
     msgpack_errors += map_lookup_uint64_t(element, ".kernarg_segment_size",
                                           &kernel_segment_size);
     msgpackErrorCheck(kernarg segment size metadata lookup in kernel metadata,
                       msgpack_errors);
 
     // create a map from symbol to name

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

int32_t __tgt_rtl_run_target_team_region_locked(
   for (int32_t i = 0; i < arg_num; ++i) {
     ptrs[i] = (void *)((intptr_t)tgt_args[i] + tgt_offsets[i]);
     args[i] = &ptrs[i];
     DP("Offseted base: arg[%d]:" DPxMOD "\n", i, DPxPTR(ptrs[i]));
   }
 
   KernelTy *KernelInfo = (KernelTy *)tgt_entry_ptr;
 
+  std::string kernel_name = std::string(KernelInfo->Name);
+  uint32_t sgpr_count, vgpr_count, sgpr_spill_count, vgpr_spill_count;
+
+  {
+    assert(KernelInfoTable[device_id].find(kernel_name) !=
+           KernelInfoTable[device_id].end());
+    auto it = KernelInfoTable[device_id][kernel_name];
+    sgpr_count = it.sgpr_count;
+    vgpr_count = it.vgpr_count;
+    sgpr_spill_count = it.sgpr_spill_count;
+    vgpr_spill_count = it.vgpr_spill_count;
+  }
+
   /*
    * Set limit based on ThreadsPerGroup and GroupsPerDevice
    */
   int num_groups = 0;
 
   int threadsPerGroup = RTLDeviceInfoTy::Default_WG_Size;
 
   getLaunchVals(threadsPerGroup, num_groups, KernelInfo->ConstWGSize,
                 KernelInfo->ExecutionMode, DeviceInfo.EnvTeamLimit,
                 DeviceInfo.EnvNumTeams,
                 num_teams,      // From run_region arg
                 thread_limit,   // From run_region arg
                 loop_tripcount, // From run_region arg
                 KernelInfo->device_id);
 
   if (print_kernel_trace >= LAUNCH) {
     // enum modes are SPMD, GENERIC, NONE 0,1,2
     // if doing rtl timing, print to stderr, unless stdout requested.
     bool traceToStdout = print_kernel_trace & (RTL_TO_STDOUT | RTL_TIMING);
     fprintf(traceToStdout ? stdout : stderr,
             "DEVID:%2d SGN:%1d ConstWGSize:%-4d args:%2d teamsXthrds:(%4dX%4d) "
-            "reqd:(%4dX%4d) n:%s\n",
+            "reqd:(%4dX%4d) sgpr_count:%u vgpr_count:%u sgpr_spill_count:%u "
+            "vgpr_spill_count:%u tripcount:%lu n:%s\n",
             device_id, KernelInfo->ExecutionMode, KernelInfo->ConstWGSize,
             arg_num, num_groups, threadsPerGroup, num_teams, thread_limit,
-            KernelInfo->Name);
+            sgpr_count, vgpr_count, sgpr_spill_count, vgpr_spill_count,
+            loop_tripcount, KernelInfo->Name);
   }
 
   // Run on the device.
   {
     hsa_queue_t *queue = DeviceInfo.HSAQueues[device_id];
     uint64_t packet_id = acquire_available_packet_id(queue);
 
     const uint32_t mask = queue->size - 1; // size is a power of 2
// Run on the device.
     packet->grid_size_z = 1;
     packet->private_segment_size = 0;
     packet->group_segment_size = 0;
     packet->kernel_object = 0;
     packet->kernarg_address = 0;     // use the block allocator
     packet->reserved2 = 0;           // atmi writes id_ here
     packet->completion_signal = {0}; // may want a pool of signals
 
-    std::string kernel_name = std::string(KernelInfo->Name);
     {
       assert(KernelInfoTable[device_id].find(kernel_name) !=
              KernelInfoTable[device_id].end());
       auto it = KernelInfoTable[device_id][kernel_name];
       packet->kernel_object = it.kernel_object;
       packet->private_segment_size = it.private_segment_size;
       packet->group_segment_size = it.group_segment_size;
       assert(arg_num == (int)it.num_args);