Skip to content

Guard NEO's scratch allocation with a spill-triggered drain - #609

Draft
michel2323 wants to merge 1 commit into
mainfrom
scratch-hedge
Draft

Guard NEO's scratch allocation with a spill-triggered drain#609
michel2323 wants to merge 1 commit into
mainfrom
scratch-hedge

Conversation

@michel2323

@michel2323 michel2323 commented Aug 11, 2026

Copy link
Copy Markdown
Member

NEO allocates a queue's scratch buffer at the first submission of a kernel whose spill exceeds what is already allocated, and that allocation has no error path: prepareScratchAllocation() calls allocateGraphicsMemoryWithProperties() without a null check, and programSurfaceState() aborts the process (UNRECOVERABLE_IF, scratch_space_controller_xehp_and_later.cpp:78, byte-identical from the 25.18 LTS driver through current master) when it failed. Whether it fails depends on how much dead-but-unfinalized driver memory happens to be live at that one instant — a GC lottery. We have logged repeated aborts of exactly this shape on Aurora (PVC, LTS NEO 25.18) under launch storms from real workloads (ExaModels), where the process dies inside zeCommandQueueExecuteCommandLists with no Julia-side recourse.

This PR removes the lottery instead of the (driver-side) bug:

  • ZeKernel caches its spillMemSize (seeded by properties(), one Int load on the launch path).
  • Each queue tracks a high-water mark of per-thread spill among submitted kernels.
  • Before the first submission that crosses the mark, the launch path retires in-flight work, flushes deferred releases (_run_reclaim_callbacks), and runs GC.gc(false) — so the one allocation the driver cannot recover from happens at the cleanest reachable moment.

The hedge fires once per (queue, scratch tier): one synchronize per workload in practice, and a cached-Int compare on the non-spilling fast path. Opt out with ONEAPI_SCRATCH_HEDGE=0 (the high-water mark is maintained regardless).

Tested with a deliberately spilling kernel (256 live accumulators): the hedge fires exactly once, never for no-spill kernels, and the knob-off path still tracks the mark. Suite green on Aurora LTS (PVC) and the kernel spills on both PVC and DG2.

This is deliberately independent of the immediate-command-list rework (#610): it guards a different resource (scratch) against a different garbage source (USM churn), and remains useful on current drivers — the missing null check is in NEO's shared code, not LTS-specific.

🤖 Generated with Claude Code

https://claude.ai/code/session_01WSFxSBXtckG3BVAT12wYkf

NEO allocates a queue's scratch buffer at the first submission of a kernel
whose spill exceeds what is already allocated, and that allocation has no
error path: allocateGraphicsMemoryWithProperties is called without a null
check, and programSurfaceState aborts the process (UNRECOVERABLE_IF) when
it failed. Whether it fails depends on how much dead-but-unfinalized driver
memory happens to be live at that one instant — a GC lottery.

Remove the lottery: track a per-queue high-water mark of kernel spill sizes,
and before the first submission that crosses it, retire in-flight work,
flush deferred releases, and run finalizers, so the fatal allocation happens
at the cleanest reachable moment. Fires once per (queue, scratch tier);
costs one cached-Int compare on the non-spilling fast path. Opt out with
ONEAPI_SCRATCH_HEDGE=0.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WSFxSBXtckG3BVAT12wYkf
@codecov

codecov Bot commented Aug 11, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.80%. Comparing base (315b088) to head (1ee7863).

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #609      +/-   ##
==========================================
+ Coverage   78.51%   78.80%   +0.29%     
==========================================
  Files          50       50              
  Lines        3490     3505      +15     
==========================================
+ Hits         2740     2762      +22     
+ Misses        750      743       -7     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

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