Skip to content

fix(core): pre-allocate trackedRefs in pass encoders (ML OOM fix)#196

Merged
kolkov merged 1 commit into
mainfrom
fix/tracked-refs-memory
May 26, 2026
Merged

fix(core): pre-allocate trackedRefs in pass encoders (ML OOM fix)#196
kolkov merged 1 commit into
mainfrom
fix/tracked-refs-memory

Conversation

@kolkov
Copy link
Copy Markdown
Contributor

@kolkov kolkov commented May 26, 2026

Summary

  • Pre-allocate trackedRefs with capacity 64 in BeginComputePass, BeginRenderPass, CreateCommandEncoder
  • Eliminates O(log N) intermediate backing arrays from append() starting at nil
  • For Born ML workloads (10K dispatches × 4 buffers): 17 doublings → 1 allocation per pass

Investigation

All transfer points already nil the source correctly:

  • ComputePassEncoder.End()p.trackedRefs = nil
  • CommandEncoder.Finish()e.trackedRefs = nil
  • postSubmit()cb.trackedRefs = nil
  • DestroyQueue.Triage() → zeros TrackedSubmission struct ✅

Rust uses mem::take() (swap with empty Vec) — Vec capacity reused via EncoderInFlight ownership. Our pre-alloc achieves same effect.

Test plan

  • Build + Tests: 15/15
  • Lint: 0 issues
  • CI
  • Born ML pprof: heap < 1 GB (was 3.8 GB)

BeginComputePass, BeginRenderPass, and CreateCommandEncoder now pre-allocate
trackedRefs with capacity 64. Previously Phase 2 tracking (v0.28.8) used
bare append() from nil — O(log N) reallocations with abandoned backing
arrays. For 10K dispatches × 4 buffers, each encoder did 17 doublings
(1→2→4→...→65536). Pre-allocation eliminates all intermediate arrays.

Investigation confirmed all transfer points already nil the source:
- ComputePassEncoder.End() → p.trackedRefs = nil
- CommandEncoder.Finish() → e.trackedRefs = nil
- postSubmit() → cb.trackedRefs = nil
- DestroyQueue.Triage() → zeros TrackedSubmission struct

Rust uses mem::take() (swap with empty) at transfer, Vec capacity is
reused via EncoderInFlight ownership chain. Our pre-alloc achieves the
same effect — single allocation per pass, no intermediate arrays.
@kolkov kolkov merged commit 1e82cd4 into main May 26, 2026
10 checks passed
@kolkov kolkov deleted the fix/tracked-refs-memory branch May 26, 2026 16:37
@codecov
Copy link
Copy Markdown

codecov Bot commented May 26, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant