pack.cpp

#include <hipFatbin.h>

#include <cassert>
#include <cstring>
#include <string>
#include <string_view>
#include <concepts>
#include <mutex>
#include <sstream>

#include "handle.hpp"

namespace {
static void finalize(fatbin::Handle *h) {
  constexpr std::string_view hipv4 = "hipv4-";
  constexpr std::string_view gcn_triple = "amdgcn-amd-amdhsa--";
  constexpr std::string_view spirv_triple = "spirv64-amd-amdhsa--amdgcnspirv";

  std::vector<std::string> arch_names;
  arch_names.reserve(h->inputs_.size());
  size_t total_size = 0;

  for (const auto &i : h->inputs_) {
    total_size += i.bin_.size();
    std::string arch(hipv4);
    if (i.type() == fatbin::BinaryType::HOST) {
      arch = i.arch();
    } else if (i.type() == fatbin::BinaryType::BINARY) {
      arch += gcn_triple;
      arch += i.arch();
    } else if (i.type() == fatbin::BinaryType::SPIRV) {
      arch += spirv_triple;
    } else {
      // TODO add support of more types
      assert(false && "unsupported binary type");
    }
    arch_names.push_back(arch);
  }

  size_t header_offset = 0;
  for (size_t i = 0; i < arch_names.size(); i++) {
    header_offset += 3 * 8;
    header_offset += arch_names[i].size();
  }

  constexpr std::string_view uncompressed_bundle_header =
      "__CLANG_OFFLOAD_BUNDLE__";

  std::vector<unsigned char> out(total_size + 2000, 0);
  size_t copy_index = 0;

  // Write bundle identifier
  std::memcpy(out.data() + copy_index, uncompressed_bundle_header.data(),
              uncompressed_bundle_header.size());
  copy_index += uncompressed_bundle_header.size();

  // Write number of bundles
  uint64_t bundles_count = arch_names.size();
  std::memcpy(out.data() + copy_index, &bundles_count, sizeof(bundles_count));
  copy_index += sizeof(bundles_count);

  size_t write_offset = header_offset + copy_index;

  assert(h->inputs_.size() == arch_names.size());
  for (size_t i = 0; i < h->inputs_.size(); i++) {
    uint64_t input_offset = write_offset;
    uint64_t input_size = h->inputs_[i].bin().size();
    uint64_t bundle_entry_id_size = arch_names[i].size();

    // Copy offsets etc
    std::memcpy(out.data() + copy_index, &input_offset, sizeof(input_offset));
    copy_index += sizeof(input_offset);
    std::memcpy(out.data() + copy_index, &input_size, sizeof(input_size));
    copy_index += sizeof(input_size);
    std::memcpy(out.data() + copy_index, &bundle_entry_id_size,
                sizeof(bundle_entry_id_size));
    copy_index += sizeof(bundle_entry_id_size);
    std::memcpy(out.data() + copy_index, arch_names[i].data(),
                arch_names[i].size());
    copy_index += arch_names[i].size();

    // Write the binary
    std::memcpy(out.data() + write_offset, h->inputs_[i].bin().data(),
                h->inputs_[i].bin().size());
    write_offset += h->inputs_[i].bin().size();
  }

  out.resize(write_offset);
  h->output_ = std::move(out);
  h->built_ = true;
}
}  // namespace

extern "C" {
hipFatbinResult hipFatbinGet(hipFatbinHandle handle, void *buffer) {
  if (handle == nullptr || buffer == nullptr) {
    return HIPFATBIN_ERROR_NULL_POINTER;
  }
  auto h = reinterpret_cast<fatbin::Handle *>(handle);
  if (!h->built()) finalize(h);

  std::memcpy(buffer, h->output_.data(), h->output_.size());
  return HIPFATBIN_SUCCESS;
}

hipFatbinResult hipFatbinSize(hipFatbinHandle handle, size_t *size) {
  if (handle == nullptr || size == nullptr) {
    return HIPFATBIN_ERROR_NULL_POINTER;
  }
  auto h = reinterpret_cast<fatbin::Handle *>(handle);
  if (!h->built()) finalize(h);

  *size = h->output_.size();
  return HIPFATBIN_SUCCESS;
}
}