SWDEV-346018 - Support unbundling of code object in runtime for hiprtc.

Change-Id: Idda30def999b28500738c7b48560d36ca6a8e730


[ROCm/clr commit: 4bba4faa0c]

projects/clr/hipamd/src/hiprtc/hiprtcComgrHelper.cpp

@@ -28,6 +28,90 @@ THE SOFTWARE.
 namespace hiprtc {
 
 namespace helpers {
+
+size_t constexpr strLiteralLength(char const* str) {
+  return *str ? 1 + strLiteralLength(str + 1) : 0;
+}
+constexpr char const* AMDGCN_TARGET_TRIPLE = "amdgcn-amd-amdhsa-";
+constexpr char const* CLANG_OFFLOAD_BUNDLER_MAGIC_STR = "__CLANG_OFFLOAD_BUNDLE__";
+static constexpr size_t bundle_magic_string_size = strLiteralLength(CLANG_OFFLOAD_BUNDLER_MAGIC_STR);
+
+struct __ClangOffloadBundleInfo {
+  uint64_t offset;
+  uint64_t size;
+  uint64_t bundleEntryIdSize;
+  const char bundleEntryId[1];
+};
+
+struct __ClangOffloadBundleHeader {
+  const char magic[bundle_magic_string_size - 1];
+  uint64_t numOfCodeObjects;
+  __ClangOffloadBundleInfo desc[1];
+};
+
+// Consumes the string 'consume_' from the starting of the given input
+// eg: input = amdgcn-amd-amdhsa--gfx908 and consume_ is amdgcn-amd-amdhsa--
+// input will become gfx908.
+static bool consume(std::string& input, std::string consume_) {
+  if (input.substr(0, consume_.size()) != consume_) {
+    return false;
+  }
+  input = input.substr(consume_.size());
+  return true;
+}
+
+bool isCodeObjectCompatibleWithDevice(std::string bundleEntryId, std::string isa) {
+  // If it is a direct match then return true.
+  if (bundleEntryId == isa) {
+    return true;
+  }
+
+  consume(bundleEntryId, std::string("hip") + '-' + std::string(AMDGCN_TARGET_TRIPLE));
+  consume(isa, std::string(AMDGCN_TARGET_TRIPLE) + '-');
+
+  if (bundleEntryId == isa) {
+    return true;
+  }
+
+  return false;
+}
+
+bool UnbundleBitCode(const std::vector<char>& bundled_llvm_bitcode, const std::string& isa,
+                     size_t& co_offset, size_t& co_size) {
+  std::string magic(bundled_llvm_bitcode.begin(),
+                    bundled_llvm_bitcode.begin() + bundle_magic_string_size);
+  if (magic.compare(CLANG_OFFLOAD_BUNDLER_MAGIC_STR)) {
+    // Handle case where the whole file is unbundled
+    return true;
+  }
+
+  std::string bundled_llvm_bitcode_s(bundled_llvm_bitcode.begin(), bundled_llvm_bitcode.begin()
+                                                                   + bundled_llvm_bitcode.size());
+  const void* data = reinterpret_cast<const void*>(bundled_llvm_bitcode_s.c_str());
+  const auto obheader
+    = reinterpret_cast<const __ClangOffloadBundleHeader*>(data);
+  const auto* desc = &obheader->desc[0];
+  for (uint64_t idx=0; idx < obheader->numOfCodeObjects; ++idx,
+                desc = reinterpret_cast<const __ClangOffloadBundleInfo*>(
+                       reinterpret_cast<uintptr_t>(&desc->bundleEntryId[0]) +
+                       desc->bundleEntryIdSize)) {
+    const void* image = reinterpret_cast<const void*>(reinterpret_cast<uintptr_t>(obheader) +
+                                                      desc->offset);
+    const size_t image_size = desc->size;
+    std::string bundleEntryId{desc->bundleEntryId, desc->bundleEntryIdSize};
+
+    // Check if the device id and code object id are compatible
+    if (isCodeObjectCompatibleWithDevice(bundleEntryId, isa)) {
+      co_offset = (reinterpret_cast<uintptr_t>(image) - reinterpret_cast<uintptr_t>(data));
+      co_size = image_size;
+      std::cout<<"bundleEntryId: "<<bundleEntryId<<" Isa:"<<isa<<" Offset: "<<co_offset<<" Size: "
+               << co_size <<std::endl;
+      break;
+    }
+  }
+  return true;
+}
+
 bool addCodeObjData(amd_comgr_data_set_t& input, const std::vector<char>& source,
                     const std::string& name, const amd_comgr_data_kind_t type) {
   amd_comgr_data_t data;

projects/clr/hipamd/src/hiprtc/hiprtcComgrHelper.hpp

@@ -31,6 +31,8 @@ THE SOFTWARE.
 
 namespace hiprtc {
 namespace helpers {
+bool UnbundleBitCode(const std::vector<char>& bundled_bit_code, const std::string& isa,
+                    size_t& co_offset, size_t& co_size);
 bool addCodeObjData(amd_comgr_data_set_t& input, const std::vector<char>& source,
                     const std::string& name, const amd_comgr_data_kind_t type);
 bool extractBuildLog(amd_comgr_data_set_t dataSet, std::string& buildLog);

projects/clr/hipamd/src/hiprtc/hiprtcInternal.cpp

@@ -500,29 +500,51 @@ amd_comgr_data_kind_t RTCLinkProgram::GetCOMGRDataKind(hiprtcJITInputType input_
 
 bool RTCLinkProgram::AddLinkerFile(std::string file_path, hiprtcJITInputType input_type) {
   amd::ScopedLock lock(lock_);
+  std::vector<char> llvm_bitcode;
 
+  // Get the file size.
   struct stat stat_buf;
   if (stat(file_path.c_str(), &stat_buf)) {
     return false;
   }
 
-  std::string link_file_name_("Linker Program");
-  std::vector<char> llvm_bitcode(stat_buf.st_size);
+  // Read the file contents
+  std::string link_file_name("Linker Program");
+  std::vector<char> link_file_info(stat_buf.st_size);
   std::ifstream bc_file(file_path, std::ios_base::in | std::ios_base::binary);
   if (!bc_file.good()) {
     return true;
   }
 
-  bc_file.read(llvm_bitcode.data(), stat_buf.st_size);
+  bc_file.read(link_file_info.data(), stat_buf.st_size);
   bc_file.close();
 
+  // If this is bundled bitcode then unbundle this.
+  if (input_type == HIPRTC_JIT_INPUT_LLVM_BUNDLED_BITCODE) {
+    if (!findIsa()) {
+      return false;
+    }
+
+    size_t co_offset = 0;
+    size_t co_size = 0;
+    if(!UnbundleBitCode(link_file_info, isa_, co_offset, co_size)) {
+      LogError("Error in hiprtc: unable to unbundle the llvm bitcode");
+      return false;
+    }
+
+    llvm_bitcode.assign(link_file_info.begin() + co_offset,
+                        link_file_info.begin() + co_offset + co_size);
+  } else {
+    llvm_bitcode.assign(link_file_name.begin(), link_file_name.end());
+  }
+
   amd_comgr_data_kind_t data_kind;
   if((data_kind = GetCOMGRDataKind(input_type)) == AMD_COMGR_DATA_KIND_UNDEF) {
     LogError("Cannot find the correct COMGR data kind");
     return false;
   }
 
-  if (!addCodeObjData(link_input_, llvm_bitcode, link_file_name_, data_kind)) {
+  if (!addCodeObjData(link_input_, llvm_bitcode, link_file_name, data_kind)) {
     LogError("Error in hiprtc: unable to add linked code object");
     return false;
   }
@@ -533,9 +555,28 @@ bool RTCLinkProgram::AddLinkerFile(std::string file_path, hiprtcJITInputType inp
 bool RTCLinkProgram::AddLinkerData(void* image_ptr, size_t image_size, std::string link_file_name,
                                    hiprtcJITInputType input_type) {
   amd::ScopedLock lock(lock_);
-
   char* image_char_buf = reinterpret_cast<char*>(image_ptr);
-  std::vector<char> llvm_bitcode(image_char_buf, image_char_buf + image_size);
+  std::vector<char> llvm_bitcode;
+
+  if (input_type == HIPRTC_JIT_INPUT_LLVM_BUNDLED_BITCODE) {
+    std::vector<char> bundled_llvm_bitcode(image_char_buf, image_char_buf + image_size);
+
+    if (!findIsa()) {
+      return false;
+    }
+
+    size_t co_offset = 0;
+    size_t co_size = 0;
+    if(!UnbundleBitCode(bundled_llvm_bitcode, isa_, co_offset, co_size)) {
+      LogError("Error in hiprtc: unable to unbundle the llvm bitcode");
+      return false;
+    }
+
+    llvm_bitcode.assign(bundled_llvm_bitcode.begin() + co_offset,
+                        bundled_llvm_bitcode.begin() + co_offset + co_size);
+  } else {
+    llvm_bitcode.assign(image_char_buf, image_char_buf + image_size);
+  }
 
   amd_comgr_data_kind_t data_kind;
   if((data_kind = GetCOMGRDataKind(input_type)) == AMD_COMGR_DATA_KIND_UNDEF) {
@@ -543,7 +584,7 @@ bool RTCLinkProgram::AddLinkerData(void* image_ptr, size_t image_size, std::stri
     return false;
   }
 
-  if(!addCodeObjData(link_input_,llvm_bitcode , link_file_name, data_kind)) {
+  if(!addCodeObjData(link_input_, llvm_bitcode , link_file_name, data_kind)) {
     LogError("Error in hiprtc: unable to add linked code object");
     return false;
   }