fix(hrr): record the code object HIP loads, not the whole fat binary

projects/clr/hipamd/src/hip_code_object.cpp

@@ -703,10 +703,10 @@ Var* StatCO::FindDeferredManagedVar(const void* ptr) {
 }
 
 // ================================================================================================
-void StatCO::ForEachFatBinaryBlob(void (*cb)(const void*)) const {
+void StatCO::ForEachFatBinary(void (*cb)(const void*, FatBinaryInfo*)) const {
   std::scoped_lock lock(sclock_);
-  for (const auto& [data, _] : modules_) {
-    cb(data);
+  for (const auto& [data, fat_binary_info] : modules_) {
+    cb(data, fat_binary_info);
   }
 }
 

projects/clr/hipamd/src/hip_code_object.hpp

@@ -195,8 +195,9 @@ class StatCO : public CodeObject {
   // Resize device-specific data structures for all registered functions and variables
   void ResizeForDevices(size_t device_count);
 
-  // Iterate all registered fat binary data pointers — for HRR capture post-registration sweep.
-  void ForEachFatBinaryBlob(void (*cb)(const void*)) const;
+  // Iterate all registered fat binaries as (data pointer, FatBinaryInfo, null until HIP
+  // digests it) — for HRR capture post-registration sweep.
+  void ForEachFatBinary(void (*cb)(const void*, FatBinaryInfo*)) const;
 
   // Iterate all registered __device__ globals as (host shadow address, symbol
   // name, size, device address) — the same post-registration sweep for HRR,

projects/clr/hipamd/src/hrr/hip_capture.cpp

@@ -47,6 +47,7 @@
 #include "platform/kernel.hpp"     // amd::KernelSignature
 #include "opencl/amdocl/cl_kernel.h"  // T_POINTER enum
 #include "os/os.hpp"               // amd::Os::installExceptionHandlers()
+#include "elf/elf.hpp"             // amd::Elf::getElfSize()
 
 // Fat binary format structs (ClangOffloadBundleUncompressedHeader, etc.)
 #include "../hip_code_object.hpp"
@@ -951,6 +952,21 @@ static hrr_cap::Hash128 hash_program_image(const amd::Program* prog, const uint8
                      (const void*)prog, (unsigned long long)size, (const void*)data);
     return {0, 0};
   }
+  // The span's length is not a readable extent: a code object HIP took from a
+  // registered fat binary carries the size its bundle entry declares, and
+  // nothing bounds that. An ELF's own headers do, so hash no further than its
+  // sections and segments reach, and refuse one whose headers reach past the
+  // declared length.
+  if (amd::Elf::isElfMagic(reinterpret_cast<const char*>(data))) {
+    const uint64_t elf_size = amd::Elf::getElfSize(data, size);
+    if (elf_size == 0) {
+      LogPrintfWarning("[HRR capture] program %p: ELF at %p does not fit the %llu bytes"
+                       " HIP holds for it — not recording a code object for it",
+                       (const void*)prog, (const void*)data, (unsigned long long)size);
+      return {0, 0};
+    }
+    size = static_cast<size_t>(elf_size);
+  }
   return hrr_cap::writer::write_code_object(data, size);
 }
 
@@ -1363,112 +1379,68 @@ hipError_t capture_hipLaunchByPtr(const void* func) {
 }
 
 // ---------------------------------------------------------------------------
-// Fat binary registration — record the binary blob
+// Fat binary registration — record the code object HIP took from it
 //
-// We capture the clang offload bundle blob so the replay can load it via
-// hipModuleLoadData, making all kernel names resolvable.
+// The bundle itself is not stored. HIP is handed a bare pointer and reads only
+// the header and the code object it picks for each GPU, with a size COMGR
+// gives it; a whole-bundle size could only come from the header, and nothing
+// bounds that. Replay resolves kernels and globals in the code objects the
+// archive holds, so the register event names the code object, not a blob.
 // ---------------------------------------------------------------------------
 
-// Byte size of a clang offload bundle, or 0 when the header is not one or fails
-// the checks HIP applies to it. HIP receives a registered fat binary as a bare
-// pointer, so it knows no image bound (image_bound in hip_fatbin.cpp) and these
-// checks are all that stands between the header and the read.
-static size_t compute_bundle_size(const void* blob) {
-  if (!blob) return 0;
-  const char* p = static_cast<const char*>(blob);
-
-  // Compressed format: magic "CCOB", header contains totalSize at byte 8.
-  if (std::memcmp(p, hip::symbols::kOffloadBundleCompressedMagicStr,
-                  hip::symbols::kOffloadBundleCompressedMagicStrSize - 1) == 0) {
-    constexpr size_t kHeaderSize =
-        offsetof(hip::symbols::ClangOffloadBundleCompressedHeader, compressedBinarydesc);
-    const auto* hdr = static_cast<const hip::symbols::ClangOffloadBundleCompressedHeader*>(blob);
-    if (hdr->totalSize < kHeaderSize) {
-      LogPrintfWarning(
-          "[HRR capture] Fat binary at %p not recorded: compressed bundle "
-          "totalSize %u is smaller than its %zu-byte header",
-          blob, static_cast<unsigned>(hdr->totalSize), kHeaderSize);
-      return 0;
-    }
-    return static_cast<size_t>(hdr->totalSize);
-  }
-
-  // Uncompressed format: magic "__CLANG_OFFLOAD_BUNDLE__" + numOfCodeObjects + entries.
-  if (std::memcmp(p, hip::symbols::kOffloadBundleUncompressedMagicStr,
-                  hip::symbols::kOffloadBundleUncompressedMagicStrSize - 1) != 0) {
-    return 0;  // Unknown format
-  }
-  const auto* hdr = static_cast<const hip::symbols::ClangOffloadBundleUncompressedHeader*>(blob);
-  uint64_t n = hdr->numOfCodeObjects;
-  if (n == 0) return 0;
-
-  // HIP hands COMGR only the first 4096 bytes, and COMGR reads the entry table from
-  // that slice and fails if a read runs off its end. __hipRegisterFatBinary passes
-  // no length, so this bound is also what keeps the walk below in range.
-  constexpr size_t kEntryTableLimit = 4096;
-  constexpr size_t kEntryHeaderSize = offsetof(hip::symbols::ClangOffloadBundleInfo, bundleEntryId);
-  size_t pos = offsetof(hip::symbols::ClangOffloadBundleUncompressedHeader, desc);
-  uint64_t end = 0;
-  for (uint64_t i = 0; i < n; i++) {
-    const auto* entry = reinterpret_cast<const hip::symbols::ClangOffloadBundleInfo*>(p + pos);
-    if (kEntryTableLimit - pos < kEntryHeaderSize ||
-        entry->bundleEntryIdSize > kEntryTableLimit - pos - kEntryHeaderSize) {
-      LogPrintfWarning(
-          "[HRR capture] Fat binary at %p not recorded: bundle entry %llu ends "
-          "past the first %zu bytes, where HIP looks for it",
-          blob, static_cast<unsigned long long>(i), kEntryTableLimit);
-      return 0;
-    }
-    pos += kEntryHeaderSize + static_cast<size_t>(entry->bundleEntryIdSize);
-    // HIP makes this check only for the code objects it picks for this host's GPUs.
-    // Every entry counts towards the end here, so capture checks them all.
-    const uint64_t entry_end = entry->offset + entry->size;
-    if (entry_end < entry->offset) {
-      LogPrintfWarning(
-          "[HRR capture] Fat binary at %p not recorded: bundle entry %llu has "
-          "an offset and size that overflow",
-          blob, static_cast<unsigned long long>(i));
-      return 0;
-    }
-    if (entry_end > end) end = entry_end;
+// Record the code object HIP extracted from a digested fat binary for the
+// current device, or device 0 before one is current, the device the global
+// variable sweep resolves on. A fat binary HIP has not digested yet has
+// nothing loaded; a launch from it records its code object then.
+//
+// Returns the code object's hash, which the register event carries so replay
+// loads it on the device that is current at that event, as HIP did here.
+// The size HIP holds is the bundle entry's declared one, so only an ELF, whose
+// headers bound it (see hash_program_image), is recorded.
+static hrr_cap::Hash128 record_fat_binary_code_object(hip::FatBinaryInfo* fbi) {
+  if (!fbi) return {0, 0};
+  int dev = hip::ihipGetDevice();
+  if (dev < 0) dev = 0;
+  if (static_cast<size_t>(dev) >= hip::g_devices.size()) return {0, 0};
+  amd::Program* prog = fbi->GetProgram(dev);
+  if (!prog) return {0, 0};
+  const uint8_t* image = nullptr;
+  size_t size = 0;
+  {
+    std::lock_guard<std::mutex> lk(g_prog_hash_mu);
+    const amd::Device* device = hip::g_devices[dev]->devices()[0];
+    const auto& bin = prog->binary(*device);
+    image = std::get<0>(bin);
+    size = std::get<1>(bin).first;
   }
-  static_assert(sizeof(size_t) >= sizeof(uint64_t), "a 64-bit bundle size must fit in size_t");
-  return static_cast<size_t>(end);
+  if (!image || size < sizeof(Elf64_Ehdr) ||
+      !amd::Elf::isElfMagic(reinterpret_cast<const char*>(image)))
+    return {0, 0};
+  const hrr_cap::Hash128 h = hash_program_image(prog, image, size);
+  remember_program_hash(prog, image, size, h);
+  return h;
+}
+
+// The register event names the code object in its blob hash, with a blob
+// size of 0; a bundle blob (archives from before) always has a size.
+static void write_register_fat_binary_event(uint64_t ret, hrr_cap::Hash128 co) {
+  hrr_args___hipRegisterFatBinary a{};
+  a.ret = ret;
+  a.blob_hash_lo = co.lo;
+  a.blob_hash_hi = co.hi;
+  hrr_cap::writer::write_event_raw(HRR_API_HIPREGISTERFATBINARY, &a.hdr, sizeof(a));
 }
 
 void** capture___hipRegisterFatBinary(const void* data) {
   void** r = g_real_compiler_table.__hipRegisterFatBinary_fn(data);
   // Shim is only installed when capture is active — no hip_capture_enabled() check needed.
 
-  // data is a __CudaFatBinaryWrapper* { magic, version, binary, dummy }.
-  // Capture the fat binary blob (binary field) so replay can load it via hipModuleLoadData.
-  struct __HRRFatBinaryWrapper { uint32_t magic; uint32_t version; const void* binary; const void* dummy; };
-  const auto* wrapper = static_cast<const __HRRFatBinaryWrapper*>(data);
-  const void* blob = (wrapper && (wrapper->magic == 0x48495046u /*HIPF*/ ||
-                                   wrapper->magic == 0x4B504948u /*HIPK*/))
-                     ? wrapper->binary : nullptr;
   // PlatformState is initialized before this shim is installed, so HIP has already
-  // digested the bundle, and no handle means it could not load it on this host: a bad
-  // header or, more often, no code object for any GPU here. Then no launch from it can
-  // succeed and be recorded, so replay never needs the blob.
-  if (blob && !r) {
-    LogPrintfWarning(
-        "[HRR capture] Fat binary at %p not recorded: HIP did not register it "
-        "(no code object for any GPU here, or a bad header)",
-        blob);
-    blob = nullptr;
-  }
-  size_t blob_size = blob ? compute_bundle_size(blob) : 0;
-
-  hrr_args___hipRegisterFatBinary a{};
-  a.ret      = reinterpret_cast<uint64_t>(r);
-  a.blob_size = static_cast<uint64_t>(blob_size);
-  if (blob && blob_size > 0) {
-    auto h = hrr_cap::writer::write_blob(blob, blob_size);
-    a.blob_hash_lo = h.lo;
-    a.blob_hash_hi = h.hi;
-  }
-  hrr_cap::writer::write_event_raw(HRR_API_HIPREGISTERFATBINARY, &a.hdr, sizeof(a));
+  // digested the bundle and r, when set, is its FatBinaryInfo**.
+  const hrr_cap::Hash128 co =
+      r ? record_fat_binary_code_object(*reinterpret_cast<hip::FatBinaryInfo**>(r))
+        : hrr_cap::Hash128{0, 0};
+  write_register_fat_binary_event(reinterpret_cast<uint64_t>(r), co);
   return r;
 }
 
@@ -2609,23 +2581,13 @@ void hip_capture_install_compiler_table(const HipCompilerDispatchTable& shims) {
 // Init / Shutdown
 // ---------------------------------------------------------------------------
 
-// Record a single fat binary blob as a HRR_API_HIPREGISTERFATBINARY event.
-// blob_ptr is the fbwrapper->binary pointer (the actual clang offload bundle).
-// These were registered before HIP initialized and are not digested yet, so
-// unlike the live shim there is no HIP verdict to check, only the header, and a
-// bundle with no code object for this host is still recorded.
-static void record_fat_binary_blob(const void* blob_ptr) {
-  if (!blob_ptr) return;
-  size_t blob_size = compute_bundle_size(blob_ptr);
-  if (blob_size == 0) return;
-
-  hrr_args___hipRegisterFatBinary a{};
-  a.ret      = 0;  // handle not meaningful at init time
-  a.blob_size = static_cast<uint64_t>(blob_size);
-  auto h = hrr_cap::writer::write_blob(blob_ptr, blob_size);
-  a.blob_hash_lo = h.lo;
-  a.blob_hash_hi = h.hi;
-  hrr_cap::writer::write_event_raw(HRR_API_HIPREGISTERFATBINARY, &a.hdr, sizeof(a));
+// Record a fat binary registered before capture started as a
+// HRR_API_HIPREGISTERFATBINARY event, with the code object HIP has loaded from
+// it so far (fbi is null until HIP digests it).
+static void record_fat_binary(const void* data, hip::FatBinaryInfo* fbi) {
+  if (!data) return;
+  // ret 0: the handle is not meaningful at init time.
+  write_register_fat_binary_event(0, record_fat_binary_code_object(fbi));
 }
 
 // Record one registered __device__ global as a HRR_API_HIPREGISTERVAR event.
@@ -2744,14 +2706,13 @@ void hip_capture_init() {
     // the compiler dispatch table is fully populated.
     hip_capture_build_compiler_table();
 
-    // Retroactively record fat binaries that fired before our shims were live.
-    // __hipRegisterFatBinary fires at app static-init, before hip_capture_init().
-    hip::PlatformState::Instance().StatCO().ForEachFatBinaryBlob(record_fat_binary_blob);
-
-    // And the __device__ globals registered against them, for the same reason
-    // and in this order: the fat binaries have to be in the archive before the
-    // symbols that live inside them.
+    // Retroactively record the __device__ globals and fat binaries registered
+    // before our shims were live: both fire at app static-init, before
+    // hip_capture_init(). Globals first, because resolving one makes HIP load
+    // the code object holding it, which the fat binary sweep then records.
+    // Replay loads every code object before it replays any event.
     hip::PlatformState::Instance().StatCO().ForEachGlobalVar(record_registered_var);
+    hip::PlatformState::Instance().StatCO().ForEachFatBinary(record_fat_binary);
 
     std::call_once(g_hrr_atexit_once, [] { std::atexit(hip_capture_shutdown); });
 
@@ -2767,7 +2728,7 @@ void hip_capture_init() {
   #else
     // HRR capture is disabled
     // below to avoid -Wunused-function / -Wunused-variable
-    (void)&record_fat_binary_blob;
+    (void)&record_fat_binary;
     (void)&record_registered_var;
     (void)&hrr_install_clr_exception_handler;
     (void)&g_hrr_atexit_once;

projects/hrr/DESIGN.md

@@ -52,7 +52,7 @@ through one of these shims, which records a fixed-size `hrr_event_header` (32 by
 type, sequence ID, timestamp, thread ID) followed by a variable-length payload of
 serialised arguments into a streaming `events.bin` file.
 
-A subset of APIs requires hand-written shims (`MANUAL_CAPTURE_APIS`): kernel launches (4 variants), host-to-device memcpy (4 variants), module load (3 variants), `__hipRegisterFatBinary`, host memory registration (`hipHostRegister`, `hipHostUnregister`), all `hipMemcpy3D` variants (4), array creation, VMM, stream/memory attribute APIs, and others where struct pointer serialisation requires custom handling. All others are code-generated by `gen_hrr_api_args.py` from `hip_api_trace.hpp`. Fat binaries registered before shim install (C++ static-init) are recovered retroactively via `ForEachFatBinaryBlob` at init.
+A subset of APIs requires hand-written shims (`MANUAL_CAPTURE_APIS`): kernel launches (4 variants), host-to-device memcpy (4 variants), module load (3 variants), `__hipRegisterFatBinary`, host memory registration (`hipHostRegister`, `hipHostUnregister`), all `hipMemcpy3D` variants (4), array creation, VMM, stream/memory attribute APIs, and others where struct pointer serialisation requires custom handling. All others are code-generated by `gen_hrr_api_args.py` from `hip_api_trace.hpp`. Fat binaries registered before shim install (C++ static-init) are recovered retroactively via `ForEachFatBinary` at init.
 
 H2D memcpy captures the host source buffer as a content-addressed blob (FNV-1a-128 hash). D2H memcpy captures the expected output blob after a forced stream sync. Kernel argument pointers are detected via `kernel->signature()` (in-tree access to `amd::KernelSignature`) — no ELF parsing needed.
 
@@ -487,7 +487,7 @@ manually: `const uint8_t* p = pl + sizeof(hrr_event_header);`
 
 | API | Extra fields |
 |-----|-------------|
-| `__hipRegisterFatBinary` | `blob_hash_lo u64`, `blob_hash_hi u64`, `blob_size u64` |
+| `__hipRegisterFatBinary` | `blob_hash_lo u64`, `blob_hash_hi u64`, `blob_size u64` (zero since capture stopped storing bundles) |
 | `hipModuleLoadData`, `hipModuleLoadDataEx`, `hipModuleLoad` | `co_hash_lo u64`, `co_hash_hi u64`, `module_id u32` |
 | `hipMemcpy`, `hipMemcpyAsync`, `hipMemcpyHtoD`, `hipMemcpyHtoDAsync` | `blob_hash_lo u64`, `blob_hash_hi u64` |
 | `hipHostRegister` | `blob_hash_lo u64`, `blob_hash_hi u64` |
@@ -698,9 +698,17 @@ before suspecting the `<<<>>>` / kernarg capture path.
 
 `__hipRegisterFatBinary` fires at C++ static-init time before shims are installed.
 At init, all already-registered fat binaries are recovered retroactively via
-`ForEachFatBinaryBlob`. At playback, each captured blob is loaded with
-`hipModuleLoadData`; the resulting `hipModule_t` is stored in `module_map` keyed by
-hash and scanned for kernel name lookups.
+`ForEachFatBinary`. Capture does not store the bundle: its size could only come
+from the header, which nothing bounds. It records the code object HIP extracted
+from the bundle for the current device, bounded by the ELF's own headers since the
+size HIP knows is the one the bundle entry declares, and the register event names
+that code object (blob hash set, blob size 0). Replay pre-loads every code object,
+so kernels and `__device__` globals resolve as they do for module loads. HIP builds
+a loaded code object for one device, so replay keeps a module per code object and
+device, and pre-loads each code object on every device a registration naming it, a
+module load or a kernel launch used it on, the device current on that thread at
+that event; on a host with GPUs of more than one kind it may load on no other. Archives from before this
+change still carry a blob, which playback loads with `hipModuleLoadData`.
 
 ## Multi-Threaded Capture
 
@@ -1109,7 +1117,7 @@ kernel.
 `__hipRegisterFatBinary` fires at C++ static-initialisation time, before
 `hip_capture_init()` installs shims. All fat binaries registered before the shim is
 active are missed by the normal capture path and recovered retroactively via
-`PlatformState::Instance().StatCO().ForEachFatBinaryBlob(...)` at init time.
+`PlatformState::Instance().StatCO().ForEachFatBinary(...)` at init time.
 
 If an application loads additional shared libraries _after_ HIP init that register
 their own fat binaries, those will be captured normally through the shim. The edge

projects/hrr/include/hrr/hrr_api_args.h

@@ -195,7 +195,7 @@ typedef struct {
 typedef struct {
     hrr_event_header hdr;
     uint64_t ret;
-    uint64_t blob_hash_lo;  /* FNV-1a-128 lo of fat binary */
+    uint64_t blob_hash_lo;  /* FNV-1a-128 lo of fat binary, or of its code object if blob_size is 0 */
     uint64_t blob_hash_hi;  /* FNV-1a-128 hi */
     uint64_t blob_size;  /* fat binary byte count */
 } hrr_args___hipRegisterFatBinary;