feat(engine): the single compute dispatch runs on the same machinery as a batch

[tato123]

Summary

The two compute-dispatch escalate ops now reach the GPU through one machinery. run_compute_kernel builds a BatchedComputeKernelDispatch of one and rides GpuContext::dispatch_compute_kernel_batch: barriers and dispatch land in one recording on the engine's shared cached recorder — one submission, one fence wait — replacing a throwaway per-dispatch transition recorder (own command pool, own fence, own submit-and-wait) followed by VulkanComputeKernel::dispatch's drain, submit, and post-wait. A dispatch whose bindings needed a transition paid 2 submissions and 3 stalls; one whose bindings sat right paid 1 and 2. Both are now 1 and 1, counted, not timed.

The kernel fence's fail-then-hang flaw leaves this path with it: a failed vkQueueSubmit leaves that fence unsignaled forever and wedges the next dispatch; the recorder tracks in-flight state so it cannot.

Through the batch method, never the batch handler — a single dispatch's refusals keep naming the binding, never "dispatch 0 of this batch" (the ticket's ruled-out shape). The one message inside the method that carried batch vocabulary now names the binding alone when the recording holds one dispatch.

Review round 1 surfaced a real narrowing and it is fixed in the shared machinery: the batch path published tracked layouts per first-touch registration per image, but a cross-process resolve mints a fresh registration cell per call — one image bound at two same-kind slots (only kind clashes are refused) left the second cell stale, and the surface-share dedup kept exactly that cell. Publication now walks every binding, so every cell naming an image learns the layout the recording left it in; the batch op inherits the fix. Falsified: narrowing the walk fails the new test with UNDEFINED vs GENERAL.

Outside tests the sources are +136/−124 (the extracted shared helpers and the per-binding publish carry docs); the wheel diff is two comments that claimed the single op dispatches on the kernel's own fence.

Closes

Closes #1890

Exit criteria

• The single compute-dispatch escalate op records the same barriers and layout transitions the batched op does — in the same recording as its dispatch, on the shared recorder — and submits and waits once
• Existing engine and rig Python kernel suites pass unchanged, including refusal messages (all single-dispatch refusals fire in the shared planner/resolver, ahead of the recording — verified by the unchanged suite)
• No doc still describes the single path as dispatching on the kernel's own fence (engine rustdoc, wheel comments; the graphics and ray-tracing claims remain true and remain)
• The batching cost test distinguishes the paths exactly: N separate ops cost N submissions and N stalls, the batch costs 1 and 1

Test plan

Every gate green at HEAD: cargo fmt --check, CI's workspace clippy (no warning in any changed hunk), cargo xtask run-local-ci-gates, cargo test -p streamlib-engine --lib (1351 passed), -p streamlib-python-wheel --lib (72), -p xtask (226), cargo doc -p streamlib --no-deps warning-free.

Engine, on the rig (RTX 3090) — compute_kernel_dispatch 30 pass (3 new), graphics_kernel_dispatch 25, ray_tracing_kernel_dispatch 25:

• a_single_dispatch_costs_one_submission_and_one_stall_and_rests_its_layouts — warm-up then one measured op via the device's submission/fence-wait counters, plus the tracked resting layouts (sampled → SHADER_READ_ONLY_OPTIMAL, storage → GENERAL) and the output pixels. Falsified — restoring the kernel-fence path fails it: stalls 2 vs 1.
• every_registration_cell_naming_one_image_learns_the_landed_layout — two hand-minted registration cells over one image (the shape a path-2 resolve produces), both asserted GENERAL after the recording. Falsified — narrowing the publish walk to first-touch leaves cell B at UNDEFINED.
• a_single_dispatch_failing_inside_the_recording_leaves_the_recorder_usable — undeclared push constants refuse at set_push_constants, after the barriers are recorded; the next dispatch succeeds. The single-path analogue of the batch's abort test.
• a_batch_costs_one_submission_and_one_stall_where_separate_dispatches_cost_n — separate arm tightened from >= to exactly N submissions and N stalls.

Python, on the rig — test_compute_kernel.py + test_kernel_dispatch_batch.py: 21 pass after a maturin develop --release rebuild.

Reviewed by review-pr (two rounds) and rust-craftsmanship-reviewer; every should-fix from both is in this diff.

Notes for owner

1. The kernel-fence hang is still reachable from four callers. VulkanComputeKernel::dispatch keeps its fail-then-deadlock fence (verified in review: reset_fences then any failure before/at submit leaves it unsignaled; the next call blocks forever) and the tone mapper, both colour-converter paths, and sdk/vulkan-jpeg still call it. Engine-layer defect; it retires naturally with the already-anticipated follow-up migrating those callers as part of the plan's bindings-at-dispatch convergence. Only reachable on a failed submit, so not filed as blocking.
2. The per-binding publication fix improves the batch op beyond the ticket's letter. It corrects semantics the batch shipped with in feat(python)!: batched kernel dispatch — one submission, one fence #1889 (its PR already noted cross-process layout-publication gaps). Same-process dispatches were never affected — Path-1 resolves share one cell.
3. Residual, pre-existing in class: on a failed wait_for_completion the local registration cells advance (deliberately — the queue owns the transitions once submitted) but the surface-share store is only updated on an Ok escalate, so a failed wait leaves store and cells divergent. Unchanged by this diff; noting for whoever next works on cross-process kernel bindings.
4. The "names a texture with no image" refusal is unreachable through the escalate ops today (buffer-backed surfaces refuse earlier, at resolve), so neither wording arm of binding_location_in_this_recording has test coverage. Left as-is rather than adding a contrived direct-call test.
5. docs/testing-hardware.md has drifted: all five crates in its tier-1 --exclude list are gone from the workspace (gate-runner finding, unrelated to this diff).
6. macOS cross-compile not run: every changed engine line is #[cfg(target_os = "linux")], and --target aarch64-apple-darwin remains broken on this box by the pre-existing iceoryx2 libproc.h failure.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Improved single compute-dispatch reliability by routing it through the same submission and synchronization process as batched dispatches.
  • Ensured surface layouts are correctly updated for all bindings, including multiple slots referencing the same image.
  • Improved recovery after recording failures so subsequent dispatches can continue normally.
• Documentation
  • Clarified compute and ray-tracing completion and error-handling behavior.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 556f67e4-29d3-4e38-8cca-35002d67b6e6

📥 Commits

Reviewing files that changed from the base of the PR and between 9ef5ff8 and 39a9de9.

📒 Files selected for processing (3)

• runtime/streamlib-engine/src/core/compiler/compiler_ops/subprocess_escalate.rs
• runtime/streamlib-engine/src/core/context/gpu_context.rs
• sdk/streamlib-python-wheel/src/python_helper_process_pixel_exchange.rs

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

Changes

Single compute dispatches now use the shared batch recorder as a batch of one. Image layouts are tracked by Vulkan image and published to all matching registrations. Tests cover submission counts, recorder recovery, and layout publication.

Compute dispatch unification

Layer / File(s)	Summary
Shared recording and layout tracking runtime/streamlib-engine/src/core/context/gpu_context.rs	The recorder supports single-dispatch recordings. Layout tracking uses image keys and publishes final layouts to every associated binding registration.
Escalate dispatch integration runtime/streamlib-engine/src/core/compiler/compiler_ops/subprocess_escalate.rs, sdk/streamlib-python-wheel/src/python_helper_process_pixel_exchange.rs	The single-dispatch path creates a recording of one and shares dispatch and layout publication with batch execution. Related comments describe the updated completion behavior.
Dispatch validation and layout assertions runtime/streamlib-engine/src/core/compiler/compiler_ops/subprocess_escalate.rs	Tests verify one submission and one wait, recorder reuse after an in-recording refusal, exact separate-dispatch costs, and layout publication across registrations.

Estimated code review effort: 4 (Complex) | ~45 minutes

Merge Risk: ⚪ Minimal · up to 39a9d

The single-dispatch path now shares the batch recording and publication machinery, with validated submission, layout-tracking, refusal, and cost behavior. No actionable merge-blocking risk remains after normal checks and review.

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant EscalationHandler
  participant SharedRecorder
  participant GPU
  participant SurfaceShare
  Client->>EscalationHandler: run_compute_kernel
  EscalationHandler->>SharedRecorder: record one compute dispatch
  SharedRecorder->>GPU: submit barriers and dispatch
  GPU-->>SharedRecorder: complete submission
  SharedRecorder-->>EscalationHandler: return after one wait
  EscalationHandler->>SurfaceShare: publish landed image layouts
  SurfaceShare-->>Client: complete operation

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly states that single compute dispatch uses the same machinery as batch dispatch.
Linked Issues check	✅ Passed	The changes satisfy the coding objectives in [#1890], including shared recording, layout publication, error preservation, and cost coverage.
Out of Scope Changes check	✅ Passed	All changes support [#1890], including implementation, tests, documentation, and related comments; no unrelated code changes are evident.
Docstring Coverage	✅ Passed	Docstring coverage is 90.91% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 22 functions across 3 files.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/1890-single-dispatch-on-batch-machinery

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}