[lua1235]

Motivation

Rocjitsu currently eagerly allocates the entire VGPR file for every simulated CU. This completely dominates RSS for most kernels: the gfx942 topology alone commits approximately 1,280 MiB (1.25 GiB) of fixed VGPR storage. This PR replaces eager allocation with demand-paged backing for VGPRs. RSS is significantly reduced on current benchmarks and workloads, while performance is also maintained or improved.

Technical Details

• Implement lazy VGPR allocation using a portable software page table with 4KiB pages, allocated and zero-initialized on first non-const access.
• Contiguous VGPR range guarantee is lost: This was only ever used for checkpointing save and restore: A replacement mechanism to copy pages into the checkpoint buffer is implemented.
• Add a saturation benchmark with full-output checksum validation and compiled resource-profile verification.

Issue Tracking

N/A

Test Plan

Perform benchmarking on workloads available in rocjitsu-test-coprpus/IREE.

Test Result

IREE e2e Softmax Test

The softmax_static_10x256x256xf32 IREE corpus workload processes 655,360 FP32 elements and writes a 2.5 MiB output. It was compiled independently for gfx942 and gfx950, executed through IREE's HIP runtime under the corresponding rocjitsu simulator, and compared bit-for-bit with the checked-in NumPy golden output.

Each exact-output total combines five baseline and five lazy-storage executions. The process measurements retained for the RSS and runtime summary were:

Target	Build	Mean peak RSS (KiB)	Mean runtime
gfx942	Baseline	1,627,477	2.09 s
gfx942	Lazy registers	316,328	0.70 s
gfx950	Baseline	1,387,861	1.84 s
gfx950	Lazy registers	338,881	0.72 s

Corpus Kernel Results

Workload	Register pressure	Baseline median	Lazy median	Baseline peak RSS	Lazy peak RSS	Result
Startup, gfx942	Fixed topology only	0.86 s	0.13 s	1,511,424 KiB	198,656 KiB	84.9% faster, 86.9% less RSS
Corpus matvec set	6–88 VGPR	3.26 s	1.86 s	1,616,700 KiB	305,864 KiB	42.9% faster, 81.1% less RSS
Corpus HipKittens 256^3	248 VGPR	2.03 s	0.60 s	1,626,112 KiB	314,368 KiB	70.4% faster, 80.7% less RSS
Corpus hip-matmul 128^3	Up to 133 VGPR plus AGPR	6.34 s	4.91 s	1,615,000 KiB	302,940 KiB	22.6% faster, 81% less RSS

All three completed corpus workloads produced their expected correctness results. The fixed-RSS reduction is 1,312,768 KiB. That closely matches the original estimate of approximately 1,281 MiB to eagerly committed VGPR files.

Adversarial Kernel Analysis

The storage cost has three distinct sections:

1. Untouched VGPR pages consume no RSS.
2. The first access during a wave dispatch commits pages, with cost (RSS and time) growing approximately linearly with the number of touched pages.
3. Subsequent operations use the same contiguous pointer path as the preexisting eager implementation and show no pressure-dependent penalty. Physical pages are returned when the wave retires.

The worst case is therefore an extremely short kernel that declares and touches nearly every VGPR across many CUs but performs very little work with them. This is not a common real-world workload pattern. The observed runtime penalty is about 8–9% for that deliberately adversarial initialization-only case. Amortization from eight rounds of computation (applying one XOR-multiply-add sequence) over every live VGPR reduces the penalty below 0.7%, while retaining an 81.2% peak-RSS reduction.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/rocm-systems/blob/develop/CONTRIBUTING.md.

[newling]

Nice speed-up! I'll review this, but just trying to understand

gfx942 topology alone commits approximately 1,280 MiB

My back-of-the-envelope: About 300 CUs, each 64 lanes with 512 registers, say. That's 300 * 64 * 512 * 4 bytes = 40 MB... any idea what I'm missing?

I wonder if a similar lazy idea can apply to LDS (if not already?)

[kuhar]

It would be useful to include the measurements that drove this change. benchmarks/vgpr_saturation/README.md explains how to collect data, but the commit and README do not report before/after peak RSS or per-dispatch wall time. Those numbers matter here because the fixed 16-byte offset creates partial-page clears and an all-block reset at idle, so the RSS reduction may come with launch-time overhead. Please record a representative topology and workload alongside both measurements so reviewers and future maintainers have a baseline.

Since this adds about 550 lines of standalone benchmark tooling that no build or CI target references, could we also give it a lightweight build/check path (possibly opt-in) so changes to the HIP source and gfx942 resource assumptions do not go unnoticed?

We should document the zero-on-allocation invariant at emulation/rocjitsu/lib/simdojo/include/simdojo/components/register_file.h:210. Clearing moved out of ComputeUnitCore::dispatch_wf_at, and VgprRedispatchTest.RecycledAllocationStartsZero shows that allocate() must return an all-zero block, regardless of storage policy, to prevent values from a retired wave from reaching the next dispatch. An @post on allocate() would make clear that the reset path is correctness-critical, not only an RSS optimization.

[lua1235]

Nice speed-up! I'll review this, but just trying to understand

gfx942 topology alone commits approximately 1,280 MiB

My back-of-the-envelope: About 300 CUs, each 64 lanes with 512 registers, say. That's 300 * 64 * 512 * 4 bytes = 40 MB... any idea what I'm missing?

I wonder if a similar lazy idea can apply to LDS (if not already?)

Your calculation covers 1 wave per CU, but each CU has 32 wavefront slots, which were all eagerly allocated.
320 CUs × 32 slots × 512 VGPRs × 64 lanes × 4 bytes = 1,280 MiB

IIRC LDS is also eagerly backed by a zero-initialized vector, so the same idea could apply. The potential win is smaller, though: roughly 320 CUs × 64 KiB = 20 MiB on gfx942, versus 1,280 MiB for VGPRs.

[newling]

Your calculation covers 1 wave per CU, but each CU has 32 wavefront slots, which were all eagerly allocated.
320 CUs × 32 slots × 512 VGPRs × 64 lanes × 4 bytes = 1,280 MiB

For the actual physical GPU footprint I think I was off by a factor of 4 (I forgot the x4 SIMDs). My understanding is that there are 32 wavefront slots, but not all slots are filled when there is register pressure. The key bit of info that I was missing though was that rocjitsu is not modelling physical registers. So that allocation of 1,280 MB is larger (by 8x?) than the actual physical register file size. I wonder if we'll eventually want to model occupancy (i.e. when not all wave slots can be filled because a waves all require too many registers) for non-functional (cycle-accurate) mode. In which case the allocation will be 8x smaller.

[lua1235]

Your calculation covers 1 wave per CU, but each CU has 32 wavefront slots, which were all eagerly allocated.
320 CUs × 32 slots × 512 VGPRs × 64 lanes × 4 bytes = 1,280 MiB

For the actual physical GPU footprint I think I was off by a factor of 4 (I forgot the x4 SIMDs). My understanding is that there are 32 wavefront slots, but not all slots are filled when there is register pressure. The key bit of info that I was missing though was that rocjitsu is not modelling physical registers. So that allocation of 1,280 MB is larger (by 8x?) than the actual physical register file size. I wonder if we'll eventually want to model occupancy (i.e. when not all wave slots can be filled because a waves all require too many registers) for non-functional (cycle-accurate) mode. In which case the allocation will be 8x smaller.

Yup, exactly 8x. The 1280 MB is impossible on real hardware. More just a quirk of how functional works.

[atgutier]

Thanks, really nice work!

I have a request to improve the architecture of this.

Could you add an allocator/demand or lazy storage header-only class/lib into util and then hide the Linux-specific bits behind the implementation?

#if checks are probably ok directly in that implementation, but high-level VMEM abstraction APIs could also be useful in an os-specific header lib in util as well.

That latter part can go in as a later refactor if needed.

Something along the lines of this (but don't need to take this verbatim if you have better ideas) would suffice:

 class ReclaimableBuffer {
  public:
    void allocate(size_t bytes, size_t alignment);
    void zero_and_reclaim(size_t offset, size_t bytes);

    std::byte *data();
    const std::byte *data() const;
  };

Its contract should guarantee:

• Newly allocated bytes read as zero.
• zero_and_reclaim() restores the requested range to zero.
• Bytes outside that range are preserved.
• Reclamation is best-effort; correctness never depends on it.
• The buffer address remains stable until destruction.

Implementations:

• Linux: mmap, madvise, munmap, and page-boundary handling.
• Portable fallback: aligned allocation plus memset/std::fill.
• Future Windows implementation: reserved/committed virtual memory and page reset.
• Other platforms can use the portable implementation.

Then DemandPagedRegisterStorage<T> becomes a thin typed adapter:

  template <typename T>
  class DemandPagedRegisterStorage {
  public:
    void init(uint32_t count) {
      buffer_.allocate(count * sizeof(T), alignof(T));
      data_ = reinterpret_cast<T *>(buffer_.data());
    }

    void reset(uint32_t base, uint32_t count) {
      buffer_.zero_and_reclaim(base * sizeof(T), count * sizeof(T));
    }

    T &operator[](uint32_t index) { return data_[index]; }

  private:
    ReclaimableBuffer buffer_;
    T *data_ = nullptr;
  };

[newling]

This speeds up a hipblaslt gemm I'm currently emulating on by a nice 25% 😃

I got an agent to review this https://github.com/newling/agent-public-reviews/blob/main/rocm-systems/review_pr9779.md, it didn't find anything blocking, just a request for documentation. Feel free to take what you like from it...

The one lingering question I have is whether modeling physical VGPRs would get the same speed-up (smaller memory footprint) and how hard/complex that would be, and if it's something we want anyway.

[lua1235]

This speeds up a hipblaslt gemm I'm currently emulating on by a nice 25% 😃

I got an agent to review this https://github.com/newling/agent-public-reviews/blob/main/rocm-systems/review_pr9779.md, it didn't find anything blocking, just a request for documentation. Feel free to take what you like from it...

The one lingering question I have is whether modeling physical VGPRs would get the same speed-up (smaller memory footprint) and how hard/complex that would be, and if it's something we want anyway.

Good question! A physical VGPR model would definitely cut the max RSS down in a comparable way (A hard upper bound, in fact), but speed-up seems harder to answer. There would be extra allocation/scheduling machinery, and simGPU throughput would naturally decrease under register pressure. But it's also not clear what effect that would actually have on simulation speed without an implementation to play with.

I actually think these two things compose though: A physical VGPR model will still need backing for its VGPRs, and that backing would probably want to be lazily allocated :D.

[kuhar]

Original anchor: emulation/rocjitsu/benchmarks/vgpr_saturation/CMakeLists.txt:64

I think we should remove ALL from vgpr_saturation_check and make this verifier an explicit target or labelled CTest. With BUILD_TESTING=ON, every default build currently compiles the benchmark and requires each kernel to use exactly Count + 1 VGPRs. That makes a compiler register-allocation change fail the simulator build even when the source is unchanged. Keeping the check opt-in or in CI would still catch benchmark drift without turning a compiler-version property into a prerequisite for ordinary development.

Original anchor: emulation/rocjitsu/lib/simdojo/include/simdojo/components/register_file.h:104

I think we should page-align the demand-paged register backing so each VGPR block can be reclaimed in full. A Wave64 block is exactly 128 KiB (512 registers x 64 lanes x 4 bytes), but the current 16-byte offset makes every block straddle page boundaries. In the churn probe, that leaves about 128 KiB resident per CU after only one block remains live, or roughly 40 MiB across the 320-CU configuration; page alignment leaves no residual pages from the freed blocks. The cache-index rationale seems weak for one multi-megabyte file per CU, whose accesses already span many cache sets. Passing page-size alignment for this policy would also make the whole-file cleanup unnecessary.

Original anchor: emulation/rocjitsu/lib/simdojo/include/simdojo/components/register_file.h:181

We should define what const reads from a free block mean, because the storage policies currently expose different results. Eager storage retains the retired values until the next allocation, while demand-paged storage clears them during free(). The simplest fix may be to document free-block reads as unspecified (or give the const accessors the same allocated-block precondition); otherwise callers can observe a backing-policy detail that looks like part of RegisterFile's contract.

[therock-pr-bot]

❌ PR Check — Action Required

Check	Status	Details
📝 PR Description	❌ Fail	Error: PR description must reference a JIRA ID, ISSUE ID, or a GitHub closing keyword. Expected: include a JIRA ID / ISSUE ID line (separator : or -, or omitted; value may be a JIRA key, a number with/without #, or a link), OR a closing keyword + issue reference. Accepted examples: • JIRA ID : TESTAUTO-6039 • JIRA ID - #330 • JIRA ID #330 • JIRA ID (on separate line) ROCM-25757 • ISSUE ID : TESTUTO-3334 • ISSUE ID #3334 • ISSUE ID - TESTAUTO-3433 • ISSUE ID (on separate line) AIRUNTIME-2352 • ISSUE ID : https://github.com/<org_name>/<repo_name>/issues/1234 • Closes #10 • Fixes octo-org/octo-repo#100 • Resolves: #123 • #123 • https://github.com/<org_name>/<repo_name>/issues/123 Current: no valid JIRA/ISSUE/closing-keyword reference found
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—

⚠️ 1 policy check(s) failed. Please address the issues above before this PR can be Reviewed.

🚫 Please fix the failed policies

• ❌ PR Description

The Not ready to Review label was added to this PR. Once all policies pass, the label is removed automatically.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

[therock-pr-bot]

🚫 Please fix the failed policies before requesting reviews.

The following policy checks failed:

• ❌ PR Description

The Not ready to Review label has been added to this PR.
Once all policies pass, the label will be removed automatically.

[kuhar]

A few comments

[newling]

I have a bias toward staying near the top of the C++ abstraction stack, maybe just because of what I am familiar with, but I’m uneasy about reaching down to mmap, madvise, page-boundary handling, and platform-specific behavior to avoid eagerly backing unused VGPR storage. Is there any precedent for this? I see that these details are contained in ReclaimableBuffer, but its implementation and #if defined(linux) behavior might still add portability and maintenance complexity.

My understanding is that each CU currently owns one contiguous VGPR backing store containing fixed-size blocks for all its wavefront slots. Is CU-wide contiguity an important requirement, or primarily a property of the current RegisterFile implementation?

Have we considered or prototyped storage owned independently by each wavefront slot—for example, a std::vector grown to the required architectural register span when a wave is dispatched? Each slot could retain and reuse its capacity as a high-water mark, or reclaim it according to a separate policy if that proves useful. This would preserve contiguous storage within each active wave while avoiding OS-specific virtual-memory management.

It would require adapting the current CU-global physical-register indexing, and it would initialize all registers requested by the wave rather than backing only pages that are actually touched, so I’m not claiming it would necessarily be faster. Still, I would be interested in seeing the complexity/performance trade-off against this more straightforward design, particularly using the workloads already measured in this PR.

[atgutier]

Approved with comment:

[MINOR] lib/util/include/util/reclaimable_buffer.h:162 — MAP_NORESERVE
permits deferred fatal allocation failure

Reasoning: Linux explicitly allows a writable MAP_NORESERVE mapping to
succeed and later raise SIGSEGV when a page cannot be committed.
Consequently, ReclaimableBuffer::allocate() can report success, while an
ordinary VGPR write later terminates the process rather than producing
std::bad_alloc. That behavior is currently absent from the class contract.

Evidence: The flag is added unconditionally when available. Anonymous
mappings are already demand-paged without it; MAP_NORESERVE changes commit
guarantees. This is especially relevant because the buffer intentionally
reserves very large sparse address ranges.

Suggested fix: Either remove MAP_NORESERVE, or retain it as an explicit
sparse-allocation policy and document that first-touch failure can be
fatal. The latter is likely appropriate for this simulator, but callers
should not believe allocation success guarantees future commit.

Confidence: High.

[MINOR] benchmarks/vgpr_saturation/CMakeLists.txt:23 — The benchmark
verifier is part of ordinary CTest and adds non-hermetic toolchain
requirements

Reasoning: With BUILD_TESTING=ON and amdclang++ available, configuration
requires Python, and lines 67–75 register the compiler-resource verifier
in the default unfiltered CTest suite. The verifier pins exact compiler-
dependent VGPR counts, so a legitimate compiler upgrade can fail the
normal rocjitsu test suite even though the simulator remains correct. A
CTest label does not exclude a test from ordinary ctest.

Evidence: The benchmark build target is correctly opt-in, but the verifier
test is not. It recompiles HIP source using the installed ROCm compiler
and requires Count + 1 VGPRs for every kernel.

Suggested fix: Gate this behind an explicit option such as
RJ_ENABLE_BENCHMARK_CHECKS, defaulting off, or move it into the runtime/
compiler corpus. Keep the explicit vgpr_saturation_check target for
reproducing the measurements.

Confidence: High on the mechanism; medium on whether this repository
intentionally accepts compiler-resource checks in its default suite.

[lua1235]

I've moved the vgpr saturation benchmark to corpus. #39

[lua1235]

I have a bias toward staying near the top of the C++ abstraction stack, maybe just because of what I am familiar with, but I’m uneasy about reaching down to mmap, madvise, page-boundary handling, and platform-specific behavior to avoid eagerly backing unused VGPR storage. Is there any precedent for this? I see that these details are contained in ReclaimableBuffer, but its implementation and #if defined(linux) behavior might still add portability and maintenance complexity.

I'm testing a CU-wide lazy fallback which I initially meant to replace the current eager fallback, but could also replace the linux specific implementation if it turns out to be comparable.

My understanding is that each CU currently owns one contiguous VGPR backing store containing fixed-size blocks for all its wavefront slots. Is CU-wide contiguity an important requirement, or primarily a property of the current RegisterFile implementation?

Have we considered or prototyped storage owned independently by each wavefront slot—for example, a std::vector grown to the required architectural register span when a wave is dispatched? Each slot could retain and reuse its capacity as a high-water mark, or reclaim it according to a separate policy if that proves useful. This would preserve contiguous storage within each active wave while avoiding OS-specific virtual-memory management.

It would require adapting the current CU-global physical-register indexing, and it would initialize all registers requested by the wave rather than backing only pages that are actually touched, so I’m not claiming it would necessarily be faster. Still, I would be interested in seeing the complexity/performance trade-off against this more straightforward design, particularly using the workloads already measured in this PR.

As to the per-wave storage, I'd like to defer that as a separate refactor if we want it. I prefer it as an abstraction, but my intuition tells me a centralized design would still be easier to optimize. (Just spit balling though).

As you mentioned, it'd require significant changes to how vgprs are addressed, and also to several interfaces which rely on the current design. For this, I'd want a stable optimized implementation to bench against first before changing anything. The preexisting eager allocation is very easy to beat, so comparing against that is not very valuable.

[newling]

Sounds like a good plan, thanks @lua1235.
One option, if you don't have time to explore this further, is to leave a small comment on the design decision with the new class, referencing back to this github PR.

[lua1235]

@newling @kuhar @atgutier Could I get a rereview when you guys have time? I replaced the Linux-specific allocator with a portable software-only page table. It achieves essentially the same performance gains as the Linux specific allocator.

The tradeoff is that VGPR access goes from direct indexing into contiguous storage to a pointer-table load, predictable null check, and some inexpensive index arithmetic. Empirically, this cost is measurable in isolated register-access microbenchmarks, but e2e workloads remain mostly unchanged: shorter workloads improve because of substantially faster simulator startup, while long workloads largely amortize both effects.

e2e measurements are documented in docs/simdojo.md.

[kuhar]

Added some comments

[newling]

I replaced the Linux-specific allocator with a portable software-only page table. It achieves essentially the same performance gains as the Linux specific allocator.

I like this new direction, good news that it was possible. I haven't been through any details myself, I've just confirmed that it looks like (nice) vanilla C++ now. I created an agent review, there are 2 suggested actionable items (might be duplicates of Jakub's) : https://github.com/newling/agent-public-reviews/blob/main/rocm-systems/review_pr9779_followup.md

[atgutier]

One more agent comment:

[BLOCKER] emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/compute_unit.h:492 and emulation/
rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/mma_exec.h:2802 — optimized MFMA/WMMA
paths read beyond lazy VGPR chunks

Reasoning: The new contract explicitly guarantees that raw_vgpr_data() spans exactly one VGPR.
VgprReadRegion::reg_data() at register_access.h:663-668 forwards that single-register pointer.
However, the optimized matrix paths still treat the result as a contiguous multi-register array—for
example:

const uint32_t *c_words = reads.acc->reg_data();
uint32_t raw = c_words[out.reg * wf + out.lane];

The new storage places registers in independent 4 KiB allocations at register_file.h:116-119. For
Wave32, each VGPR is 128 bytes, so a chunk contains 32 registers. The existing
Gfx1250WmmaTest.F16Fp8K64MatchesReferenceLayout places its four-register accumulator at VGPR 30;
accessing VGPRs 32–33 walks beyond the first chunk.

Evidence:

• Exact-head GCC 13 ASan/AVX-512 build succeeded.

• Running Gfx1250WmmaTest.F16Fp8K64MatchesReferenceLayout produced a 4-byte heap-buffer-overflow at
  mma_exec.h:2802.

• ASan reported the read exactly one byte-range past the 4,096-byte
  SoftwareLazyRegisterStorage<VectorReg<32>>::Chunk allocated at register_file.h:118.

• The same binary passes the test with RJ_FORCE_SCALAR=1, isolating the defect to the optimized
  multi-register pointer path.

• Thirteen fast-path groups retain this assumption: mma_exec.h:1842, 1920, 2036, 2116, 2557, 2683,
  2785, 3338, 4057, 4132, 4221, 4310, and 4398.

• Current sanitizer CI does not enable the AVX-512 path, so its green status does not cover this
  failure.

Non-sanitized execution can read unrelated allocator memory, crash, or silently produce incorrect
matrix results.

Suggested fix: Remove the dense-span assumption from all optimized matrix paths. Gather each multi-
VGPR region into contiguous scratch using reg_data(relative_reg)/lanes(relative_reg), or introduce a
region API that explicitly produces a contiguous gathered view. Restoring an allocation-block-level
contiguity guarantee is another option, but its RSS and performance impact needs remeasurement.

Add an AVX-512 sanitizer regression that proves the optimized path was selected; otherwise this test
silently exercises only the scalar fallback.

Confidence: high—reproduced against exact PR head 8e7a5d7.

[kuhar]

Just a few comments