[OpenMP][AMDGPU] Extract code object version from the ELF

[saiislam]

Summary

Added CodeObjectVersion as a new field in KernelArgPool and
modified relevant constructors to populate it.

[JonChesterfield]

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

128

I don't think the memory pool needs to know what format the elf is using

481

This is per code object, how does a global scalar capture that?

1382

This seems likely to be something already available in the elf library

[saiislam]

Fixed usage of CodeObjectVersion and ImplicitArgs size in the old plugin.
A separate patch will be required to support cov5 for the next gen plugin.

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

128

Fixed in the new diff.

kernargSizeIncludingImplicit() needs the CodeObjectVersion because size of ImplicitArgs in cov4 is 56 bytes whereas it is 256 bytes in cov5.

481

As we do not support mixing of code object versions, the value remains the same for all.

This field gets set during call to __tgt_rtl_load_binary_locked() and is later used as DeviceInfo().CodeObjectVersion.

1382

I took the style that two other functions are using for looking up symbols and getting platform flags in this file. *.getHeader() seem to be the simplest way to get the header.

[saiislam]

Ping.

[saiislam]

Not required after D139730

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

enum IMPLICITARGS : uint16_t {
   COV5_GRID_DIMS_SIZE = 2,
 
   COV5_HOSTCALL_PTR_OFFSET = 80,
 
   COV5_HEAPV1_PTR_OFFSET = 96,
   COV5_HEAPV1_PTR_SIZE = 8
 };
 
+const uint16_t getImplicitArgsSize(uint16_t Version);
+
 #endif

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

const char *get_elf_mach_gfx_name(uint32_t EFlags) {
   case EF_AMDGPU_MACH_AMDGCN_GFX1102:
     return "gfx1102";
   case EF_AMDGPU_MACH_AMDGCN_GFX1103:
     return "gfx1103";
   default:
     return "--unknown gfx";
   }
 }
+
+const uint16_t getImplicitArgsSize(uint16_t Version) {
+  return Version < ELFABIVERSION_AMDGPU_HSA_V5 ? IMPLICITARGS::COV4_SIZE
+                                               : IMPLICITARGS::COV5_SIZE;
+}

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

 struct KernelArgPool {
 private:
   static pthread_mutex_t Mutex;
 
 public:
   uint32_t KernargSegmentSize;
   void *KernargRegion = nullptr;
   std::queue<int> FreeKernargSegments;
+  uint16_t CodeObjectVersion;
 
   uint32_t kernargSizeIncludingImplicit() {
-    return KernargSegmentSize + sizeof(AMDGPUImplicitArgsTy);
+    return KernargSegmentSize + getImplicitArgsSize(CodeObjectVersion);
   }
 
   ~KernelArgPool() {
     if (KernargRegion) {
       auto R = hsa_amd_memory_pool_free(KernargRegion);
       if (R != HSA_STATUS_SUCCESS) {
         DP("hsa_amd_memory_pool_free failed: %s\n", get_error_string(R));
       }
     }
   }
 
   // Can't really copy or move a mutex
   KernelArgPool() = default;
   KernelArgPool(const KernelArgPool &) = delete;
   KernelArgPool(KernelArgPool &&) = delete;
 
-  KernelArgPool(uint32_t KernargSegmentSize, hsa_amd_memory_pool_t &MemoryPool)
-      : KernargSegmentSize(KernargSegmentSize) {
+  KernelArgPool(uint32_t KernargSegmentSize, hsa_amd_memory_pool_t &MemoryPool,
+                uint16_t CodeObjectVersion)
+      : KernargSegmentSize(KernargSegmentSize),
+        CodeObjectVersion(CodeObjectVersion) {
 
     // impl uses one pool per kernel for all gpus, with a fixed upper size
     // preserving that exact scheme here, including the queue<int>
 
     hsa_status_t Err = hsa_amd_memory_pool_allocate(
         MemoryPool, kernargSizeIncludingImplicit() * MAX_NUM_KERNELS, 0,
         &KernargRegion);
 
struct KernelTy {
   int16_t ConstWGSize;
   int32_t DeviceId;
   void *CallStackAddr = nullptr;
   const char *Name;
 
   KernelTy(llvm::omp::OMPTgtExecModeFlags ExecutionMode, int16_t ConstWgSize,
            int32_t DeviceId, void *CallStackAddr, const char *Name,
            uint32_t KernargSegmentSize,
-           hsa_amd_memory_pool_t &KernArgMemoryPool)
+           hsa_amd_memory_pool_t &KernArgMemoryPool, uint16_t CodeObjectVersion)
       : ExecutionMode(ExecutionMode), ConstWGSize(ConstWgSize),
         DeviceId(DeviceId), CallStackAddr(CallStackAddr), Name(Name) {
     DP("Construct kernelinfo: ExecMode %d\n", ExecutionMode);
 
     std::string N(Name);
     if (KernelArgPoolMap.find(N) == KernelArgPoolMap.end()) {
-      KernelArgPoolMap.insert(
-          std::make_pair(N, std::unique_ptr<KernelArgPool>(new KernelArgPool(
-                                KernargSegmentSize, KernArgMemoryPool))));
+      KernelArgPoolMap.insert(std::make_pair(
+          N, std::unique_ptr<KernelArgPool>(new KernelArgPool(
+                 KernargSegmentSize, KernArgMemoryPool, CodeObjectVersion))));
     }
   }
 };
 
 /// List that contains all the kernels.
 /// FIXME: we may need this to be per device and per library.
 std::list<KernelTy> KernelsList;
 
public:
 
   // Device properties
   std::vector<int> ComputeUnits;
   std::vector<int> GroupsPerDevice;
   std::vector<int> ThreadsPerGroup;
   std::vector<int> WarpSize;
   std::vector<std::string> GPUName;
   std::vector<std::string> TargetID;
+  uint16_t CodeObjectVersion;
 
   // OpenMP properties
   std::vector<int> NumTeams;
   std::vector<int> NumThreads;
 
   // OpenMP Environment properties
   EnvironmentVariables Env;
 
static uint64_t acquireAvailablePacketId(hsa_queue_t *Queue) {
   bool Full = true;
   while (Full) {
     Full =
         PacketId >= (Queue->size + hsa_queue_load_read_index_scacquire(Queue));
   }
   return PacketId;
 }
 
+const uint16_t getCodeObjectVersionFromELF(__tgt_device_image *Image) {
+  char *ImageBegin = (char *)Image->ImageStart;
+  size_t ImageSize = (char *)Image->ImageEnd - ImageBegin;
+
+  StringRef Buffer = StringRef(ImageBegin, ImageSize);
+  auto ElfOrErr = ObjectFile::createELFObjectFile(MemoryBufferRef(Buffer, ""),
+                                                  /*InitContent=*/false);
+  if (!ElfOrErr) {
+    REPORT("Failed to load ELF: %s\n", toString(ElfOrErr.takeError()).c_str());
+    return 1;
+  }
+
+  if (const auto *ELFObj = dyn_cast<ELF64LEObjectFile>(ElfOrErr->get())) {
+    auto Header = ELFObj->getELFFile().getHeader();
+    uint16_t Version = (uint8_t)(Header.e_ident[EI_ABIVERSION]);
+    DP("ELFABIVERSION Version: %u\n", Version);
+    return Version;
+  }
+  return 0;
+}
+
 int32_t runRegionLocked(int32_t DeviceId, void *TgtEntryPtr, void **TgtArgs,
                         ptrdiff_t *TgtOffsets, int32_t ArgNum, int32_t NumTeams,
                         int32_t ThreadLimit, uint64_t LoopTripcount) {
   // Set the context we are using
   // update thread limit content in gpu memory if un-initialized or specified
   // from host
 
   DP("Run target team region thread_limit %d\n", ThreadLimit);
int32_t runRegionLocked(int32_t DeviceId, void *TgtEntryPtr, void **TgtArgs,
   // Run on the device.
   {
     hsa_queue_t *Queue = DeviceInfo().HSAQueueSchedulers[DeviceId].next();
     if (!Queue) {
       return OFFLOAD_FAIL;
     }
     uint64_t PacketId = acquireAvailablePacketId(Queue);
 
+    uint16_t CodeObjectVersion = DeviceInfo().CodeObjectVersion;
     const uint32_t Mask = Queue->size - 1; // size is a power of 2
     hsa_kernel_dispatch_packet_t *Packet =
         (hsa_kernel_dispatch_packet_t *)Queue->base_address + (PacketId & Mask);
 
     // packet->header is written last
     Packet->setup = UINT16_C(1) << HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS;
     Packet->workgroup_size_x = WorkgroupSize;
     Packet->workgroup_size_y = 1;
if (!ArgPool) {
       }
 
       // Copy explicit arguments
       for (int I = 0; I < ArgNum; I++) {
         memcpy((char *)KernArg + sizeof(void *) * I, Args[I], sizeof(void *));
       }
 
       // Initialize implicit arguments. TODO: Which of these can be dropped
-      AMDGPUImplicitArgsTy *ImplArgs = reinterpret_cast<AMDGPUImplicitArgsTy *>(
-          static_cast<char *>(KernArg) + ArgPool->KernargSegmentSize);
-      memset(ImplArgs, 0,
-             sizeof(AMDGPUImplicitArgsTy)); // may not be necessary
-      ImplArgs->OffsetX = 0;
-      ImplArgs->OffsetY = 0;
-      ImplArgs->OffsetZ = 0;
+      uint8_t *ImplArgs =
+          static_cast<uint8_t *>(KernArg) + sizeof(void *) * ArgNum;
+      memset(ImplArgs, 0, getImplicitArgsSize(CodeObjectVersion));
 
       uint64_t Buffer = 0;
       // assign a hostcall buffer for the selected Q
       if (__atomic_load_n(&DeviceInfo().HostcallRequired, __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 HostcallInitLock = PTHREAD_MUTEX_INITIALIZER;
         pthread_mutex_lock(&HostcallInitLock);
         Buffer = hostrpc_assign_buffer(
             DeviceInfo().HSAAgents[DeviceId], Queue, DeviceId);
         pthread_mutex_unlock(&HostcallInitLock);
         if (!Buffer) {
           DP("hostrpc_assign_buffer failed, gpu would dereference null and "
              "error\n");
           return OFFLOAD_FAIL;
         }
       }
 
       DP("Implicit argument count: %d\n",
           KernelInfoEntry.implicit_argument_count);
+
+      if (CodeObjectVersion < llvm::ELF::ELFABIVERSION_AMDGPU_HSA_V5) {
         DP("Setting Hostcall buffer for COV4\n");
         memcpy(&ImplArgs[IMPLICITARGS::COV4_HOSTCALL_PTR_OFFSET], &Buffer,
                IMPLICITARGS::HOSTCALL_PTR_SIZE);
+      } else {
+        DP("Code object version 5 is not yet supported\n");
+        return OFFLOAD_FAIL;
+      }
 
       Packet->kernarg_address = KernArg;
     }
 
     hsa_signal_t S = DeviceInfo().FreeSignalPool.pop();
     if (S.handle == 0) {
       DP("Failed to get signal instance\n");
       return OFFLOAD_FAIL;
__tgt_target_table *__tgt_rtl_load_binary_locked(int32_t DeviceId,
   DeviceInfo().clearOffloadEntriesTable(DeviceId);
 
   // 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 (!elfMachineIdIsAmdgcn(Image))
     return NULL;
 
+  DeviceInfo().CodeObjectVersion = getCodeObjectVersionFromELF(Image);
+
   {
     auto Env =
         DeviceEnvironment(DeviceId, DeviceInfo().NumberOfDevices,
                           DeviceInfo().Env.DynamicMemSize, Image, ImgSize);
 
     auto &KernelInfo = DeviceInfo().KernelInfoTable[DeviceId];
     auto &SymbolInfo = DeviceInfo().SymbolInfoTable[DeviceId];
     hsa_status_t Err = moduleRegisterFromMemoryToPlace(
if (Err == HSA_STATUS_SUCCESS) {
          "value "
          "GENERIC (1)\n",
          ExecModeName);
     }
     check("Loading computation property", Err);
 
     KernelsList.push_back(KernelTy(ExecModeVal, WGSizeVal, DeviceId,
                                    CallStackAddr, E->name, KernargSegmentSize,
-                                   DeviceInfo().KernArgPool));
+                                   DeviceInfo().KernArgPool,
+                                   DeviceInfo().CodeObjectVersion));
     __tgt_offload_entry Entry = *E;
     Entry.addr = (void *)&KernelsList.back();
     DeviceInfo().addOffloadEntry(DeviceId, Entry);
     DP("Entry point %ld maps to %s\n", E - HostBegin, E->name);
   }
 
   return DeviceInfo().getOffloadEntriesTable(DeviceId);
 }