feat(rocjitsu): Lazy vgpr allocation system

emulation/rocjitsu/CMakeLists.txt

@@ -405,7 +405,6 @@ add_subdirectory(fuzz/decode)
 if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
     add_subdirectory(tools/rocjitsu)
 endif()
-# Tests and scaling tools.
 if(BUILD_TESTING)
     add_subdirectory(tests)
 endif()

emulation/rocjitsu/docs/simdojo.md

@@ -209,6 +209,95 @@ controllers.
 indexed read/write access. **VectorReg** (`vector_reg.h`) models a
 SIMD-width vector register with per-element access.
 
+### Register-file storage
+
+The functional GPU model provisions vector-register storage for every fixed
+wavefront slot in each simulated compute unit. Slots are independent of
+hardware occupancy: a wave that uses many registers still occupies one
+simulator slot, and register pressure does not reduce the number of admitted
+waves. For the shipped gfx942 configuration, the vector backing has the
+following dimensions:
+
+| Quantity                  | Value                                      |
+|---------------------------|--------------------------------------------|
+| Simulated compute units   | 320                                        |
+| Wavefront slots per CU    | 32                                         |
+| Registers per slot        | 512 combined VGPR and AGPR entries         |
+| Register width            | 64 lanes × 4 bytes                         |
+| Complete vector backing   | 1,280 MiB                                  |
+
+`RegisterFile` divides the physical VGPR index space into fixed 4 KiB chunks.
+An untouched chunk is represented by a null entry and reads as zero; its
+storage is allocated and zero-initialized on the first mutable access. Retiring
+a wave releases wholly covered chunks and clears any shared boundary chunk.
+The implementation uses only portable C++ allocation and does not require
+virtual-memory APIs. Mutable handles and handles into materialized registers
+remain stable until their allocation retires, but the complete register file is
+deliberately not contiguous. Const access to an unmaterialized register observes
+its logical zero value through shared immutable backing without allocating a
+chunk. That read-only handle is an ephemeral value observation, not a persistent
+storage identity, and need not observe a later write through another handle.
+
+This is a functional storage model, not a model of a physical register file.
+In particular, filling all 32 simulator slots with high-pressure waves can
+represent an occupancy that hardware could not admit. That distinction is
+important when interpreting memory measurements or using register pressure to
+approximate hardware throughput.
+
+#### Representative workloads
+
+The fixed-topology and IREE measurements were collected on 2026-08-10 from
+rocjitsu revision `49e6ddafeb9fd8cadda1c78b891466fd498ec676`. Fixed topology
+reports the median of six processes; IREE reports the median of five processes,
+all of which matched the checked-in output exactly. The CTS rows report the
+median of three source-identical processes from rocjitsu revision
+`4d5779c15beb0d9495b319dce5d906e966275e23` and test-corpus revision
+`ce5da512188dd40de0bc7da298ec11b587d8fdd3`; all tests passed. Every build was
+Release, pinned to CPU 8, and used an AMD EPYC 9554 host running Linux 6.8.
+
+| Workload                             | Eager runtime | Software runtime | Eager peak RSS | Software peak RSS | Effect                              |
+|--------------------------------------|--------------:|-----------------:|---------------:|------------------:|-------------------------------------|
+| Fixed topology, no kernel            |        0.87 s |          0.125 s |  1,511,424 KiB |       201,728 KiB | 85.6% faster, 86.7% less RSS        |
+| IREE softmax, exact output            |        2.50 s |           0.99 s |  1,630,208 KiB |       326,656 KiB | 60.4% faster, 80.0% less RSS        |
+| CTS math                             |       28.37 s |          27.94 s |  2,235,924 KiB |       926,788 KiB | 1.5% faster, 58.5% less RSS         |
+| CTS MFMA                             |      151.40 s |         150.78 s |  2,235,392 KiB |       926,720 KiB | 0.4% faster, 58.5% less RSS         |
+
+Unlike fixed topology and the shorter IREE workload, the CTS suites are
+dominated by simulated instruction execution and therefore amortize the fixed
+initialization work that lazy storage eliminates. Across three runs, the
+startup proxy (whole-process wall time minus test-body time) fell from about
+1.29 seconds to 0.54 seconds. Test-body time increased by 1.2% for math and was
+unchanged to measurement precision for MFMA (+0.1%). The result is a smaller
+end-to-end runtime improvement even as peak RSS falls by 58.5%.
+
+#### Worst-case saturation validation
+
+The qualified gfx942 saturation workload was measured on 2026-08-10 from
+rocjitsu revision `49e6ddafeb9fd8cadda1c78b891466fd498ec676`. The
+representative case places eight waves on each of one XCD's 40 CUs and reports
+three-launch medians. The all-slot case deliberately fills all 32 functional
+simulator slots on all 320 CUs with 253-VGPR waves—an occupancy that hardware
+could not admit—and reports one checksum-validated launch.
+
+| Saturation runtime                          |       Eager | Software lazy | Change      |
+|---------------------------------------------|------------:|--------------:|-------------|
+| 253-VGPR initialization-only median launch  |  156.201 ms |    157.479 ms | 0.8% slower |
+| 253-VGPR eight-update median launch         | 1,332.955 ms |  1,335.671 ms | 0.2% slower |
+| All-slot 253-VGPR eight-update timed launch |    48.127 s |      48.780 s | 1.4% slower |
+
+Peak RSS is a process-wide high-water measurement. The representative sweep
+therefore has one result covering all 9--253-VGPR pressure points rather than
+separate values for its individual launch timings.
+
+| Saturation peak RSS                |          Eager | Software lazy | Change         |
+|------------------------------------|---------------:|--------------:|----------------|
+| Complete 9--253-VGPR sweep         | 1,609,728 KiB  |  318,280 KiB  | 80.2% less RSS |
+| All slots, 253 VGPR, eight updates | 1,622,016 KiB  |  967,064 KiB  | 40.4% less RSS |
+
+The qualified workload, complete methodology, and reproduction commands are
+pinned at rocjitsu test corpus revision
+[`b3f5fa8a263ee2cd83666a62bc18731b09b1166a`](https://github.com/ROCm/rocjitsu-test-corpus/tree/b3f5fa8a263ee2cd83666a62bc18731b09b1166a/corpus/kernels/benchmarks/vgpr_saturation).
+
 ---
 
 ## Messages and Events

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/config/checkpoint.cpp

@@ -8,7 +8,8 @@
 #include "checkpoint_generated.h"
 #include "flatbuffers/flatbuffers.h"
 
-#include <cstring>
+#include <algorithm>
+#include <cstddef>
 #include <fstream>
 #include <stdexcept>
 #include <vector>
@@ -18,6 +19,32 @@ namespace config {
 
 namespace {
 
+flatbuffers::Offset<flatbuffers::Vector<uint8_t>>
+serialize_vgpr_block(flatbuffers::FlatBufferBuilder &builder, const amdgpu::ComputeUnitCore &cu,
+                     uint32_t base) {
+  const size_t register_bytes = static_cast<size_t>(cu.wf_size()) * sizeof(uint32_t);
+  const size_t block_bytes = static_cast<size_t>(cu.vgpr_allocation_block_size()) * register_bytes;
+  uint8_t *serialized = nullptr;
+  const auto offset = builder.CreateUninitializedVector<uint8_t>(block_bytes, &serialized);
+  cu.copy_raw_vgprs_to(base, cu.vgpr_allocation_block_size(),
+                       {reinterpret_cast<std::byte *>(serialized), block_bytes});
+  return offset;
+}
+
+/// Restore sparse checkpoint data into a freshly allocated, zeroed VGPR block.
+/// Zero source registers are skipped to preserve lazy backing; this is a full
+/// restore, rather than a merge, only while the destination begins entirely
+/// zeroed.
+void restore_vgpr_block_into_zeroed_storage(amdgpu::ComputeUnitCore &cu, uint32_t base,
+                                            const flatbuffers::Vector<uint8_t> &stored) {
+  const size_t register_bytes = static_cast<size_t>(cu.wf_size()) * sizeof(uint32_t);
+  const size_t block_bytes = static_cast<size_t>(cu.vgpr_allocation_block_size()) * register_bytes;
+  const size_t copy_size = std::min<size_t>(stored.size(), block_bytes);
+  cu.restore_raw_vgprs_into_zeroed_storage(
+      base, cu.vgpr_allocation_block_size(),
+      {reinterpret_cast<const std::byte *>(stored.data()), copy_size});
+}
+
 /// @brief Serialize the SoC configuration into a FlatBuffer SimulationConfig.
 flatbuffers::Offset<fb::SimulationConfig>
 serialize_config(flatbuffers::FlatBufferBuilder &builder, const SoC &soc,
@@ -121,10 +148,7 @@ void save_checkpoint(const std::string &path, const SoC &soc, uint64_t tick,
 
           auto sgprs_vec =
               builder.CreateVector(cu->sgpr_data(w->sgpr_alloc().base), w->num_sgprs());
-          size_t vgpr_bytes = static_cast<size_t>(cu->vgpr_allocation_block_size()) *
-                              static_cast<size_t>(w->wf_size()) * sizeof(uint32_t);
-          auto vgprs_vec =
-              builder.CreateVector(cu->raw_vgpr_data(w->vgpr_alloc().base), vgpr_bytes);
+          auto vgprs_vec = serialize_vgpr_block(builder, *cu, w->vgpr_alloc().base);
 
           auto wfs = fb::CreateWavefrontState(builder, w->wf_id(), w->wg_id(), w->pc, w->exec_raw(),
                                               w->vcc(), w->m0(), w->is_halted(), w->status_raw(),
@@ -261,12 +285,10 @@ LoadedConfig restore_checkpoint(const std::string &path) {
             }
           }
 
-          if (auto *vgprs = wf_state->vgprs()) {
-            size_t vgpr_bytes = static_cast<size_t>(cu->vgpr_allocation_block_size()) *
-                                static_cast<size_t>(wf->wf_size()) * sizeof(uint32_t);
-            size_t copy_size = std::min<size_t>(vgprs->size(), vgpr_bytes);
-            std::memcpy(cu->raw_vgpr_data(wf->vgpr_alloc().base), vgprs->data(), copy_size);
-          }
+          // dispatch_wf_at() has just allocated this block, so RegisterFile::allocate()'s
+          // zero-state postcondition satisfies the sparse restore helper's precondition.
+          if (auto *vgprs = wf_state->vgprs())
+            restore_vgpr_block_into_zeroed_storage(*cu, wf->vgpr_alloc().base, *vgprs);
         }
       }
     }

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/cdna5/isa.h

@@ -70,6 +70,7 @@ struct Isa : amdgpu::RdnaIsaBase {
   using StatusReg = rocjitsu::cdna5::StatusReg;
 
   static constexpr uint32_t WF_SIZE_MAX = 32; ///< gfx1250 is Wave32-only.
+  static constexpr uint32_t MAX_WF_SLOTS = 64;
   // gfx1250 uses the RDNA base decoder shape, but its MODE[27] enables VGPR
   // high-bank indexing rather than RDNA4 DISABLE_PERF.
   static constexpr bool MODE_HAS_GPR_IDX_EN = true;

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/cdna_isa_base.h

@@ -91,10 +91,13 @@ class CdnaStatusReg : public ::util::Bitfield<32> {
 /// `Isa` struct adds `Decoder`, `MachineInst`, `OperandType`, and `StatusReg`
 /// type aliases.
 struct CdnaIsaBase {
-  static constexpr uint32_t WF_SIZE = 64;               ///< Lanes per wavefront (Wave64).
-  static constexpr uint32_t WF_SIZE_MAX = 64;           ///< CDNA is Wave64-only.
-  static constexpr uint32_t MAX_SGPRS_PER_WF = 102;     ///< SGPRs per wavefront.
-  static constexpr uint32_t MAX_VGPRS_PER_WF = 256;     ///< VGPRs per wavefront.
+  static constexpr uint32_t WF_SIZE = 64;           ///< Lanes per wavefront (Wave64).
+  static constexpr uint32_t WF_SIZE_MAX = 64;       ///< CDNA is Wave64-only.
+  static constexpr uint32_t MAX_WF_SLOTS = 32;      ///< Maximum simulated CU wave slots.
+  static constexpr uint32_t MAX_SGPRS_PER_WF = 102; ///< SGPRs per wavefront.
+  static constexpr uint32_t MAX_VGPRS_PER_WF = 256; ///< VGPRs per wavefront.
+  static constexpr uint32_t MAX_ADDRESSABLE_VGPRS_PER_WF =
+      MAX_VGPRS_PER_WF;                                 ///< Maximum ordinary VGPR address span.
   static constexpr uint32_t MAX_ACC_VGPRS_PER_WF = 0;   ///< AccVGPRs (0 = absent; CDNA1 default).
   static constexpr uint8_t WAITCNT_LGKMCNT_MASK = 0x0F; ///< lgkmcnt mask in S_WAITCNT [11:8].
   static constexpr bool SRAM_ECC = false;               ///< CDNA1 default: no SRAM ECC.