diff --git a/CHANGELOG.md b/CHANGELOG.md index 82eedfa..7621cc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,18 @@ _Disclaimer: this changelog is updated using generative AI, but is still verified manually._ +## Unreleased + +### Added +- **Non-blocking GPU→CPU readback.** A new `GpuReadback` type stages a buffer copy and returns immediately: `request`/`request_copy` start the transfer and `try_take` polls for completion (`is_idle` reports whether a readback is in flight). `GpuTimestamps` gained matching `request_read`/`try_take`/`is_idle` for non-blocking timestamp readback. Both are driven by a new `Backend::poll` maintenance method (a non-blocking counterpart to `synchronize`) implemented across the WebGPU, CUDA, Metal, and CPU backends. The `metal` feature now pulls in `block` for command-buffer completion handlers. +- `WebGpu::from_device`: build a WebGPU backend on top of an already-created wgpu instance/adapter/device/queue instead of creating its own. +- `from_wgpu` constructors on `GpuBufferSlice`, `GpuBufferSliceMut`, and `WebGpuBufferSlice` to wrap a foreign `wgpu::Buffer` (owned by another library sharing the device) as a khal buffer slice. +- `force_cpu_coroutines` option for `#[spirv_bindgen]`: forces the CPU backend to dispatch a kernel through the coroutine scheduler so `barrier_wait()` synchronizes, even for kernels with no `#[spirv(workgroup)]` parameter. Replaces the old dummy-workgroup-parameter hack. +- The CUDA backend's `load_module_bytes` now accepts a pre-linked CUBIN (detected via the ELF magic) in addition to PTX text, for modules referencing symbols the driver JIT cannot resolve on its own (e.g. libdevice `__nv_*` math). + +### Fixed +- WebGPU entry-point lookup on wasm now mirrors naga's identifier sanitization, which appends a trailing `_` to names ending in a digit, so kernels like `reduce_add_f32` resolve correctly. + ## v0.2.0 ### Added diff --git a/crates/khal-derive/src/spirv_bindgen/cpu.rs b/crates/khal-derive/src/spirv_bindgen/cpu.rs index 6275f58..0bdc112 100644 --- a/crates/khal-derive/src/spirv_bindgen/cpu.rs +++ b/crates/khal-derive/src/spirv_bindgen/cpu.rs @@ -222,6 +222,7 @@ pub(super) fn generate_cpu_dispatch_block( original_params: &[OriginalParam], workgroup_size: [u32; 3], func_ident: &syn::Ident, + force_coroutines: bool, ) -> proc_macro2::TokenStream { let wg_x = workgroup_size[0]; let wg_y = workgroup_size[1]; @@ -250,7 +251,10 @@ pub(super) fn generate_cpu_dispatch_block( if cfg_variant_sets.is_empty() { // No cfg-gated params: single variant with all params. let all_params: Vec<&OriginalParam> = original_params.iter().collect(); - let body = generate_variant_body(&all_params, workgroup_size); + let mut body = generate_variant_body(&all_params, workgroup_size); + // `force_cpu_coroutines`: force the coroutine path even with no workgroup param, + // so `barrier_wait()` actually synchronizes. + body.has_workgroup |= force_coroutines; let dispatch_loop = gen_dispatch_loop(&body, workgroup_size, func_ident); quote! { @@ -284,7 +288,8 @@ pub(super) fn generate_cpu_dispatch_block( }) .collect(); - let body = generate_variant_body(&active_params, workgroup_size); + let mut body = generate_variant_body(&active_params, workgroup_size); + body.has_workgroup |= force_coroutines; let dispatch_loop = gen_dispatch_loop(&body, workgroup_size, func_ident); quote! { diff --git a/crates/khal-derive/src/spirv_bindgen/mod.rs b/crates/khal-derive/src/spirv_bindgen/mod.rs index b816dac..48a0854 100644 --- a/crates/khal-derive/src/spirv_bindgen/mod.rs +++ b/crates/khal-derive/src/spirv_bindgen/mod.rs @@ -344,7 +344,17 @@ pub(crate) fn spirv_bindgen(attr: TokenStream, item: TokenStream) -> TokenStream // #[spirv_bindgen(MyName)] // #[spirv_bindgen(spirv_passthrough)] // #[spirv_bindgen(MyName, spirv_passthrough)] + // #[spirv_bindgen(force_cpu_coroutines)] // force CPU coroutine dispatch + // #[spirv_bindgen(MyName, force_cpu_coroutines)] let mut spirv_passthrough = false; + // When set, the CPU backend always dispatches this kernel through the + // coroutine scheduler (the path that makes `barrier_wait()` actually + // synchronize), even when the kernel declares no `#[spirv(workgroup)]` + // parameter. Use for kernels that rely on barriers but exchange data through + // storage buffers rather than shared memory — otherwise their threads run + // sequentially with no-op barriers on the CPU backend. Replaces the old + // "dummy `#[spirv(workgroup)] _cpu_marker` parameter" hack. + let mut force_cpu_coroutines = false; let struct_name: syn::Ident = if attr.is_empty() { // No explicit name provided - derive from function name (snake_case to PascalCase) let func_name = func.sig.ident.to_string(); @@ -358,6 +368,8 @@ pub(crate) fn spirv_bindgen(attr: TokenStream, item: TokenStream) -> TokenStream for ident in &args { if ident == "spirv_passthrough" { spirv_passthrough = true; + } else if ident == "force_cpu_coroutines" { + force_cpu_coroutines = true; } else { if name.is_some() { return syn::Error::new_spanned( @@ -618,8 +630,12 @@ pub(crate) fn spirv_bindgen(attr: TokenStream, item: TokenStream) -> TokenStream // ── CPU dispatch code generation ───────────────────────────────────── let func_ident = &func.sig.ident; - let cpu_dispatch_block = - cpu::generate_cpu_dispatch_block(&original_params, workgroup_size, func_ident); + let cpu_dispatch_block = cpu::generate_cpu_dispatch_block( + &original_params, + workgroup_size, + func_ident, + force_cpu_coroutines, + ); // ── CUDA (nvptx64) kernel entry point generation ───────────────────── let func_name_str = func.sig.ident.to_string(); @@ -743,8 +759,20 @@ pub(crate) fn spirv_bindgen(attr: TokenStream, item: TokenStream) -> TokenStream format!("{}::{}", module, entry_point) }; + // On wasm/WebGPU, wgpu transpiles the SPIR-V to WGSL via naga and matches + // the requested entry point against the generated WGSL function name. naga's + // Namer sanitizes identifiers: `::` -> `_`, and additionally appends a + // trailing `_` when the name ends in a digit (it reserves that suffix for its + // own numeric disambiguation). Mirror both rules so the lookup matches, e.g. + // `linalg::reduce::reduce_add_f32` -> `linalg_reduce_reduce_add_f32_`. #[cfg(target_arch = "wasm32")] - let full_entry = full_entry.replace("::", "_"); + let full_entry = { + let mut full_entry = full_entry.replace("::", "_"); + if matches!(full_entry.bytes().last(), Some(b'0'..=b'9')) { + full_entry.push('_'); + } + full_entry + }; Self::from_bytes(backend, file.contents(), &full_entry) } diff --git a/crates/khal/Cargo.toml b/crates/khal/Cargo.toml index 3e3472d..8b10d64 100644 --- a/crates/khal/Cargo.toml +++ b/crates/khal/Cargo.toml @@ -13,7 +13,7 @@ derive = ["khal-derive"] webgpu = ["khal-derive?/webgpu"] cpu = ["khal-derive?/cpu"] cuda = ["dep:cudarc", "khal-derive?/cuda"] -metal = ["dep:metal", "dep:naga", "khal-derive?/webgpu"] +metal = ["dep:metal", "dep:naga", "dep:block", "khal-derive?/webgpu"] push_constants = [] subgroup_ops = [] @@ -40,6 +40,9 @@ paste = "1" [target.'cfg(target_os = "macos")'.dependencies] metal = { version = "0.32", optional = true } naga = { version = "29", optional = true, features = ["spv-in", "msl-out"] } +# Objective-C block support for Metal command-buffer completion handlers +# (non-blocking readback completion). Matches metal-rs's own `block` dep. +block = { version = "0.1.6", optional = true } [lints.rust] # The `objc` crate's `msg_send!` macro (used by the Metal backend) expands to diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index 7ecc166..d41c1a9 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -241,6 +241,16 @@ pub enum GpuBufferSlice<'a, T: DeviceValue> { } impl<'a, T: DeviceValue> GpuBufferSlice<'a, T> { + /// Wraps a foreign `wgpu::Buffer` (owned by another library that shares this + /// backend's device) as a read-only GPU buffer slice over its whole range. + /// + /// The element type `T` is purely a compile-time tag; callers are responsible + /// for ensuring the buffer's byte layout matches `T`. + #[cfg(feature = "webgpu")] + pub fn from_wgpu(buffer: &'a wgpu::Buffer) -> Self { + Self::WebGpu(crate::backend::webgpu::WebGpuBufferSlice::from_wgpu(buffer)) + } + /// Returns the underlying CPU slice. Panics if this is not a CPU buffer slice. #[cfg(feature = "cpu")] pub fn unwrap_slice(&self) -> &[T] { @@ -424,6 +434,17 @@ impl<'a, T: DeviceValue + bytemuck::Pod> GpuBufferSliceMut<'a, T> { } impl<'a, T: DeviceValue> GpuBufferSliceMut<'a, T> { + /// Wraps a foreign `wgpu::Buffer` (owned by another library that shares this + /// backend's device) as a mutable storage GPU buffer slice over its whole + /// range, suitable for passing to a kernel `.call(...)` as a mutable binding. + /// + /// The element type `T` is purely a compile-time tag; callers are responsible + /// for ensuring the buffer's byte layout matches `T`. + #[cfg(feature = "webgpu")] + pub fn from_wgpu(buffer: &'a wgpu::Buffer) -> Self { + Self::WebGpu(crate::backend::webgpu::WebGpuBufferSlice::from_wgpu(buffer)) + } + /// Returns the underlying mutable CPU slice. Panics if this is not a CPU buffer slice. #[cfg(feature = "cpu")] pub fn unwrap_slice(&mut self) -> &mut [T] { @@ -679,6 +700,21 @@ impl CpuTimestamps { entries: Arc::new(Mutex::new(Vec::new())), } } + + /// No-op: CPU pass timings are collected synchronously on pass drop, so + /// results are available as soon as the step returns. + fn request_read(&mut self) {} + + /// Always ready: returns the timings collected so far this frame. + fn try_take(&mut self) -> Option> { + Some(self.entries.lock().unwrap().clone()) + } + + /// Always idle: CPU readback is synchronous, so there's never an in-flight + /// readback to wait on before recording a new frame. + fn is_idle(&self) -> bool { + true + } } /// Backend-agnostic GPU timestamp manager for profiling compute passes. @@ -723,6 +759,25 @@ impl GpuTimestamps { !matches!(self, GpuTimestamps::Noop) } + /// Whether no non-blocking readback is in flight, so a new frame can be + /// recorded and resolved. Use this to gate recording: while a + /// [`request_read`](Self::request_read) is pending, recording another frame + /// would clobber the readback in progress (and, on wgpu, resolve into a + /// mapped staging buffer). Mirrors [`GpuReadback::is_idle`]. + pub fn is_idle(&self) -> bool { + match self { + #[cfg(feature = "webgpu")] + GpuTimestamps::WebGpu(ts) => ts.is_idle(), + #[cfg(feature = "cuda")] + GpuTimestamps::Cuda(ts) => ts.is_idle(), + #[cfg(feature = "metal")] + GpuTimestamps::Metal(ts) => ts.is_idle(), + #[cfg(feature = "cpu")] + GpuTimestamps::Cpu(ts) => ts.is_idle(), + GpuTimestamps::Noop => true, + } + } + /// Resets the timestamp manager for a new frame. pub fn reset(&mut self) { match self { @@ -750,6 +805,47 @@ impl GpuTimestamps { } } + /// Initiates a non-blocking readback of the timed passes. + /// + /// Call once per frame after the passes have been submitted (and, on the + /// wgpu backend, after [`resolve`](Self::resolve)). Returns immediately + /// without blocking on the GPU. Poll for completion with + /// [`try_take`](Self::try_take). + pub fn request_read(&mut self, backend: &GpuBackend) { + match (self, backend) { + #[cfg(feature = "webgpu")] + (GpuTimestamps::WebGpu(ts), GpuBackend::WebGpu(_)) => ts.request_read(), + #[cfg(feature = "cuda")] + (GpuTimestamps::Cuda(ts), _) => ts.request_read(), + #[cfg(feature = "metal")] + (GpuTimestamps::Metal(ts), GpuBackend::Metal(metal)) => ts.request_read(metal), + #[cfg(feature = "cpu")] + (GpuTimestamps::Cpu(ts), _) => ts.request_read(), + _ => {} + } + } + + /// Non-blocking poll of a readback started by [`request_read`](Self::request_read). + /// + /// Returns `Some(results)` once the GPU work that produced the timestamps + /// has completed; `None` while it is still in flight. This drives the + /// backend with a non-blocking [`poll`](Backend::poll), so completion + /// progresses as long as it is called periodically (e.g. once per frame). + pub fn try_take(&mut self, backend: &GpuBackend) -> Option> { + backend.poll(); + match self { + #[cfg(feature = "webgpu")] + GpuTimestamps::WebGpu(ts) => ts.try_take(), + #[cfg(feature = "cuda")] + GpuTimestamps::Cuda(ts) => ts.try_take(), + #[cfg(feature = "metal")] + GpuTimestamps::Metal(ts) => ts.try_take(), + #[cfg(feature = "cpu")] + GpuTimestamps::Cpu(ts) => ts.try_take(), + GpuTimestamps::Noop => Some(Vec::new()), + } + } + /// Reads back timestamp results after GPU synchronization. /// /// Returns per-pass durations in milliseconds. Call after `backend.synchronize()`. @@ -768,6 +864,225 @@ impl GpuTimestamps { } } +/// Tracks completion of a [`GpuReadback`]'s in-flight copy. +enum ReadbackState { + /// No readback in flight. + Idle, + /// The staging copy is host-visible as soon as the submission is processed + /// (CPU backend — synchronous wrt submit). + Ready, + /// A wgpu buffer map is in flight; the channel fires once the GPU has + /// finished and the staging buffer is mapped. + #[cfg(feature = "webgpu")] + WebGpu(async_channel::Receiver>), + /// A Metal completion flag, flipped by the fence's command-buffer completion + /// handler once the copy has landed in the shared staging buffer. + #[cfg(feature = "metal")] + Metal(std::sync::Arc), +} + +/// Non-blocking GPU→CPU readback of a small buffer region. +/// +/// Makes it easier to stage non-blocking readback to eventually access data +/// transfered from the device. +/// - [`request`](Self::request) copies the sources into an internal staging buffer +/// and starts an asynchronous readback that returns immediately. +/// - [`try_take`](Self::try_take) polls for completion and copies the data out +/// once the GPU has actually finished. +/// +/// Reuse one instance across frames to pipeline the readback. +pub struct GpuReadback { + staging: GpuBuffer, + len: usize, + state: ReadbackState, +} + +// SAFETY: mirrors `MetalTimestamps` — the only non-`Send`/`Sync` payload is a +// Metal command buffer, whose underlying Objective-C object Apple documents as +// thread-safe. The wgpu channel and `GpuBuffer` are already `Send`/`Sync`. +unsafe impl Send for GpuReadback {} +unsafe impl Sync for GpuReadback {} + +impl GpuReadback { + /// Creates a readback with an internal staging buffer of `len` elements. + pub fn new(backend: &GpuBackend, len: usize) -> Result { + let staging = + backend.uninit_buffer::(len, BufferUsages::MAP_READ | BufferUsages::COPY_DST)?; + Ok(Self { + staging, + len, + state: ReadbackState::Idle, + }) + } + + /// Whether no readback is currently in flight (ready to [`request`](Self::request)). + /// + /// Use this to pipeline a single-slot readback: only issue a new request + /// once the previous one has been consumed by [`try_take`](Self::try_take), + /// otherwise a freshly-requested copy is never old enough to have completed. + pub fn is_idle(&self) -> bool { + matches!(self.state, ReadbackState::Idle) + } + + /// Number of elements the staging buffer holds. + pub fn len(&self) -> usize { + self.len + } + + /// Whether the staging buffer is empty. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Copies the first `len` elements of `source` (from `source_offset`) into + /// the staging buffer and starts a non-blocking readback. Use this for bulk + /// readback of a whole buffer; see [`request`](Self::request) for gathering + /// individual scalars from several sources. + pub fn request_copy( + &mut self, + backend: &GpuBackend, + source: &GpuBuffer, + source_offset: usize, + ) -> Result<(), GpuBackendError> { + let mut encoder = backend.begin_encoding(); + encoder.copy_buffer_to_buffer(source, source_offset, &mut self.staging, 0, self.len)?; + backend.submit(encoder)?; + self.state = self.begin_completion(backend); + Ok(()) + } + + /// Gathers several `(source, source_offset, count)` ranges into the staging + /// buffer back-to-back (in the given order) and starts a non-blocking + /// readback. Each range copies `count` elements from `source[source_offset..]`; + /// use `count == 1` to gather individual scalars, or a larger `count` to read + /// a whole sub-buffer (e.g. a per-batch count array). The combined element + /// count must not exceed the readback's [`len`](Self::len). A previously + /// requested-but-untaken readback is discarded. + pub fn request( + &mut self, + backend: &GpuBackend, + sources: &[(&GpuBuffer, usize, usize)], + ) -> Result<(), GpuBackendError> { + let mut encoder = backend.begin_encoding(); + let mut dst_offset = 0; + for (source, source_offset, count) in sources.iter() { + encoder.copy_buffer_to_buffer( + source, + *source_offset, + &mut self.staging, + dst_offset, + *count, + )?; + dst_offset += count; + } + backend.submit(encoder)?; + self.state = self.begin_completion(backend); + Ok(()) + } + + /// Sets up backend-specific completion tracking after the copy is submitted. + fn begin_completion(&self, backend: &GpuBackend) -> ReadbackState { + match (&self.staging, backend) { + #[cfg(feature = "webgpu")] + (GpuBuffer::WebGpu(buffer), _) => { + let (sender, receiver) = async_channel::bounded(1); + buffer.slice(..).map_async(wgpu::MapMode::Read, move |v| { + let _ = sender.force_send(v); + }); + ReadbackState::WebGpu(receiver) + } + #[cfg(feature = "metal")] + (GpuBuffer::Metal(_), GpuBackend::Metal(metal)) => { + ReadbackState::Metal(metal.commit_completion()) + } + #[allow(unreachable_patterns)] + _ => ReadbackState::Ready, + } + } + + /// Non-blocking poll of a readback started by [`request`](Self::request). + /// + /// Returns `true` and fills `out` once the data is ready; returns `false` + /// (leaving `out` untouched) while the GPU is still running. Drives the + /// backend with a non-blocking [`poll`](Backend::poll), so completion + /// progresses as long as this is called periodically. + pub fn try_take(&mut self, backend: &GpuBackend, out: &mut [T]) -> bool { + backend.poll(); + match std::mem::replace(&mut self.state, ReadbackState::Idle) { + ReadbackState::Idle => false, + ReadbackState::Ready => { + self.read_staging_into(out); + true + } + #[cfg(feature = "webgpu")] + ReadbackState::WebGpu(rx) => match rx.try_recv() { + Ok(Ok(())) => { + self.read_staging_into(out); + true + } + // Map failed / channel closed: stop waiting, leave `out` as-is. + Ok(Err(_)) | Err(async_channel::TryRecvError::Closed) => true, + Err(async_channel::TryRecvError::Empty) => { + self.state = ReadbackState::WebGpu(rx); + false + } + }, + #[cfg(feature = "metal")] + ReadbackState::Metal(done) => { + if done.load(std::sync::atomic::Ordering::Acquire) { + self.read_staging_into(out); + true + } else { + self.state = ReadbackState::Metal(done); + false + } + } + } + } + + /// Copies the host-visible staging contents into `out` (up to `out.len()`). + fn read_staging_into(&self, out: &mut [T]) { + let n = out.len().min(self.len); + if n == 0 { + return; + } + match &self.staging { + #[cfg(feature = "cpu")] + GpuBuffer::Cpu(data) => out[..n].copy_from_slice(&data[..n]), + #[cfg(feature = "metal")] + GpuBuffer::Metal(buffer) => { + // SAFETY: shared-storage buffer; contents() is host-visible and + // valid, and `T: AnyBitPattern` makes any bit pattern a valid T. + unsafe { + std::ptr::copy_nonoverlapping( + buffer.raw().contents() as *const T, + out.as_mut_ptr(), + n, + ); + } + } + #[cfg(feature = "webgpu")] + GpuBuffer::WebGpu(buffer) => { + let slice = buffer.slice(..); + let data = slice.get_mapped_range(); + // SAFETY: T: AnyBitPattern; copy at most the smaller of the + // mapped range and the destination. + unsafe { + std::ptr::copy_nonoverlapping( + data.as_ptr(), + out.as_mut_ptr() as *mut u8, + data.len().min(n * std::mem::size_of::()), + ); + } + drop(data); + buffer.unmap(); + } + #[allow(unreachable_patterns)] + _ => {} + } + } +} + impl Backend for GpuBackend { const NAME: &'static str = "any"; const TARGET: super::CompileTarget = super::CompileTarget::Wgsl; @@ -1024,6 +1339,19 @@ impl Backend for GpuBackend { } } + fn poll(&self) { + match self { + #[cfg(feature = "webgpu")] + GpuBackend::WebGpu(backend) => backend.poll(), + #[cfg(feature = "cuda")] + GpuBackend::Cuda(backend) => backend.poll(), + #[cfg(feature = "metal")] + GpuBackend::Metal(backend) => backend.poll(), + #[cfg(feature = "cpu")] + GpuBackend::Cpu => {} + } + } + async fn read_buffer( &self, buffer: &Self::Buffer, diff --git a/crates/khal/src/backend/cuda.rs b/crates/khal/src/backend/cuda.rs index cf715e9..2cba404 100644 --- a/crates/khal/src/backend/cuda.rs +++ b/crates/khal/src/backend/cuda.rs @@ -219,6 +219,26 @@ impl CudaTimestamps { self.pending.lock().unwrap().clear(); } + /// Whether no recorded-but-untaken timestamps remain (safe to record a new + /// frame). False while events from a previous frame await [`try_take`]. + pub fn is_idle(&self) -> bool { + self.pending.lock().unwrap().is_empty() + } + + /// Initiates a non-blocking readback. No-op for CUDA: completion is probed + /// directly via the recorded events in [`try_take`](Self::try_take). + pub fn request_read(&mut self) {} + + /// Non-blocking poll of the recorded timestamps. + /// + /// Returns `Some(results)` once every event has completed, or `None` while + /// any is still pending. `event::elapsed` reports `CUDA_ERROR_NOT_READY` + /// (an `Err`) until both events of a pair have been recorded by the GPU, so + /// a successful read doubles as the readiness probe. + pub fn try_take(&mut self) -> Option> { + self.read().ok() + } + /// Reads timestamp results. Must be called after stream synchronization. pub fn read(&self) -> Result, CudaBackendError> { let pending = self.pending.lock().unwrap(); diff --git a/crates/khal/src/backend/metal.rs b/crates/khal/src/backend/metal.rs index 713d3f0..2f3d352 100644 --- a/crates/khal/src/backend/metal.rs +++ b/crates/khal/src/backend/metal.rs @@ -15,6 +15,7 @@ use crate::backend::{ use crate::shader::{BindGroupLayoutInfo, ShaderArgsError}; use bytemuck::{AnyBitPattern, NoUninit}; // metal re-exports objc; pull in its macros so msg_send! / sel! resolve. +use block::ConcreteBlock; use metal::objc::runtime::Object; use metal::objc::{msg_send, sel, sel_impl}; use metal::{ @@ -26,6 +27,7 @@ use metal::{ use std::collections::{BTreeMap, HashMap}; use std::marker::PhantomData; use std::ops::RangeBounds; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; // ── Core backend ─────────────────────────────────────────────────────── @@ -48,6 +50,29 @@ pub struct Metal { unsafe impl Send for Metal {} unsafe impl Sync for Metal {} +impl Metal { + /// Commits an empty command buffer with a completion handler that flips the + /// returned flag once the GPU finishes it. + /// + /// Because the queue completes command buffers in commit order, the flag + /// being set means every earlier submission has finished. The handler fires + /// asynchronously on the GPU's completion — unlike polling `status()`, this + /// needs no `synchronize()`/drain to make progress, so the non-blocking + /// readback paths ([`MetalTimestamps`], `GpuReadback`) work on their own. + pub(crate) fn commit_completion(&self) -> Arc { + let done = Arc::new(AtomicBool::new(false)); + let done_handler = done.clone(); + let cb = self.queue.new_command_buffer(); + let handler = ConcreteBlock::new(move |_cb: &metal::CommandBufferRef| { + done_handler.store(true, Ordering::Release); + }) + .copy(); + cb.add_completed_handler(&handler); + cb.commit(); + done + } +} + /// Cached info needed to issue GPU timestamp queries on this device. struct MetalTimingCaps { /// The device's "timestamp" common counter set. @@ -327,6 +352,12 @@ pub struct MetalTimestamps { labels: Vec, /// Tick → nanosecond multiplier captured from the backend at creation. period_ns: f64, + /// Completion flag for the fence committed by [`request_read`](Self::request_read), + /// flipped by its command-buffer completion handler. Because the queue + /// completes buffers in commit order, this being set means every earlier + /// buffer — including the ones that sampled these timestamps — has finished + /// writing the shared sample buffer. + done: Option>, } // SAFETY: CounterSampleBuffer wraps an MTLCounterSampleBuffer, thread-safe. @@ -356,6 +387,7 @@ impl MetalTimestamps { next_index: 0, labels: Vec::with_capacity(capacity as usize), period_ns: caps.period_ns, + done: None, }) } @@ -363,6 +395,37 @@ impl MetalTimestamps { pub fn reset(&mut self) { self.next_index = 0; self.labels.clear(); + self.done = None; + } + + /// Whether no non-blocking readback is in flight (safe to record a new frame). + pub fn is_idle(&self) -> bool { + self.done.is_none() + } + + /// Initiates a non-blocking readback of the sampled timestamps. + /// + /// Call once after the frame's passes have been submitted. Commits an empty + /// fence command buffer and returns immediately; poll for completion with + /// [`try_take`](Self::try_take). + pub fn request_read(&mut self, metal: &Metal) { + self.done = Some(metal.commit_completion()); + } + + /// Non-blocking poll of a readback started by [`request_read`](Self::request_read). + /// + /// Returns `Some(results)` once the fence's completion handler has fired (so + /// the sampled values are valid), or `None` while the GPU is still running. + pub fn try_take(&mut self) -> Option> { + let ready = self + .done + .as_ref() + .is_some_and(|d| d.load(Ordering::Acquire)); + if !ready { + return None; + } + self.done = None; + Some(self.read().unwrap_or_default()) } /// Reads back timestamp results after GPU synchronization. diff --git a/crates/khal/src/backend/mod.rs b/crates/khal/src/backend/mod.rs index cee795f..a7e051d 100644 --- a/crates/khal/src/backend/mod.rs +++ b/crates/khal/src/backend/mod.rs @@ -240,6 +240,14 @@ pub trait Backend: 'static + Sized + MaybeSendSync { ) -> Self::Dispatch<'a>; /// Blocks until all submitted GPU work has completed. fn synchronize(&self) -> Result<(), Self::Error>; + /// Non-blocking maintenance poll. + /// + /// Lets the backend make progress on already-submitted work and fire any + /// completion callbacks (e.g. buffer map callbacks on the wgpu backend). + /// Unlike [`synchronize`](Self::synchronize), this never blocks waiting on + /// the GPU. Call it once per frame to drive non-blocking readbacks such as + /// [`crate::backend::GpuTimestamps::try_take`]. + fn poll(&self) {} /// Submits the recorded commands in the encoder for execution. fn submit(&self, encoder: Self::Encoder) -> Result<(), Self::Error>; diff --git a/crates/khal/src/backend/webgpu.rs b/crates/khal/src/backend/webgpu.rs index f074eb1..f4921c5 100644 --- a/crates/khal/src/backend/webgpu.rs +++ b/crates/khal/src/backend/webgpu.rs @@ -40,6 +40,17 @@ pub struct WebGpuBufferSlice<'a> { pub(crate) byte_len: u64, } +impl<'a> WebGpuBufferSlice<'a> { + /// Wraps a foreign `wgpu::Buffer` (e.g. one owned by another library sharing + /// this device) as a buffer slice spanning its whole range. + pub fn from_wgpu(buffer: &'a wgpu::Buffer) -> Self { + Self { + inner: buffer.slice(..), + byte_len: buffer.size(), + } + } +} + impl<'a> From> for wgpu::BindingResource<'a> { fn from(slice: WebGpuBufferSlice<'a>) -> Self { slice.inner.into() @@ -185,6 +196,28 @@ impl WebGpu { }) } + /// Builds a WebGPU backend on top of an already-created wgpu device/queue, + /// instead of creating its own (as [`Self::new`] does). + pub fn from_device(instance: Instance, adapter: Adapter, device: Device, queue: Queue) -> Self { + let timestamp_supported = device.features().contains(wgpu::Features::TIMESTAMP_QUERY); + let is_vulkan_backend = adapter.get_info().backend == wgpu::Backend::Vulkan; + let spirv_passthrough_enabled = is_vulkan_backend + && device + .features() + .contains(wgpu::Features::PASSTHROUGH_SHADERS); + + Self { + _instance: instance, + _adapter: adapter, + device, + queue, + force_buffer_copy_src: false, + hacks: vec![], + spirv_passthrough_enabled, + timestamp_supported, + } + } + /// Adds a regex-based text replacement to apply to WGSL source before compilation. pub fn append_hack(&mut self, regex: Regex, replace_pattern: String) { self.hacks.push((regex, replace_pattern)); @@ -564,6 +597,11 @@ impl Backend for WebGpu { Ok(()) } + fn poll(&self) { + // Non-blocking: process completed work and fire map callbacks. + let _ = self.device.poll(wgpu::PollType::Poll); + } + async fn read_buffer( &self, buffer: &Self::Buffer, @@ -891,6 +929,21 @@ impl crate::backend::Buffer for Buffer { } /// WebGPU timestamp query manager for profiling compute passes. +/// State of a non-blocking timestamp readback initiated by +/// [`WebGpuTimestamps::request_read`]. +enum TimestampReadState { + /// No readback in flight. + Idle, + /// A readback was requested for a frame that recorded no passes. + Empty, + /// A buffer map is in flight; `rx` fires once the GPU has finished and the + /// staging buffer is mapped. + Pending { + rx: async_channel::Receiver>, + query_count: u32, + }, +} + pub struct WebGpuTimestamps { query_set: wgpu::QuerySet, resolve_buffer: wgpu::Buffer, @@ -899,6 +952,7 @@ pub struct WebGpuTimestamps { next_index: u32, labels: Vec, timestamp_period: f32, + read_state: TimestampReadState, } impl WebGpuTimestamps { @@ -936,6 +990,7 @@ impl WebGpuTimestamps { next_index: 0, labels: Vec::with_capacity(capacity as usize), timestamp_period, + read_state: TimestampReadState::Idle, }) } @@ -943,6 +998,14 @@ impl WebGpuTimestamps { pub fn reset(&mut self) { self.next_index = 0; self.labels.clear(); + self.read_state = TimestampReadState::Idle; + } + + /// Whether no non-blocking readback is in flight (safe to record + resolve a + /// new frame). While a [`request_read`](Self::request_read) is pending, the + /// staging buffer is mapped, so resolving into it again would be invalid. + pub fn is_idle(&self) -> bool { + matches!(self.read_state, TimestampReadState::Idle) } /// Resolves timestamp queries to a buffer and copies to staging for readback. @@ -1006,22 +1069,87 @@ impl WebGpuTimestamps { self.staging_buffer.unmap(); } + Ok(self.durations_from_raw(&raw_timestamps)) + } + + /// Initiates a non-blocking readback of the resolved timestamps. + /// + /// Call once after the frame's passes have been submitted (and after + /// [`resolve`](Self::resolve)). This issues an asynchronous buffer map and + /// returns immediately without blocking on the GPU. Poll for completion + /// with [`try_take`](Self::try_take). + pub fn request_read(&mut self) { + if self.next_index == 0 { + self.read_state = TimestampReadState::Empty; + return; + } + let query_count = self.next_index * 2; + let (sender, receiver) = async_channel::bounded(1); + self.staging_buffer + .slice(..) + .map_async(wgpu::MapMode::Read, move |v| { + let _ = sender.force_send(v); + }); + self.read_state = TimestampReadState::Pending { + rx: receiver, + query_count, + }; + } + + /// Non-blocking poll of a readback started by [`request_read`](Self::request_read). + /// + /// Returns `Some(results)` once the GPU work has completed and the staging + /// buffer is mapped; `None` while still in flight. The caller must drive + /// completion by polling the device (see [`WebGpu::poll`]). + pub fn try_take(&mut self) -> Option> { + match std::mem::replace(&mut self.read_state, TimestampReadState::Idle) { + TimestampReadState::Idle => None, + TimestampReadState::Empty => Some(Vec::new()), + TimestampReadState::Pending { rx, query_count } => match rx.try_recv() { + Ok(Ok(())) => { + let mut raw = vec![0u64; query_count as usize]; + let buffer_slice = self.staging_buffer.slice(..); + let data = buffer_slice.get_mapped_range(); + // SAFETY: u64 is AnyBitPattern; we copy at most the smaller + // of the mapped range and our destination. + unsafe { + std::ptr::copy_nonoverlapping( + data.as_ptr(), + raw.as_mut_ptr() as *mut u8, + data.len().min(raw.len() * std::mem::size_of::()), + ); + } + drop(data); + self.staging_buffer.unmap(); + Some(self.durations_from_raw(&raw)) + } + // Map failed: report no results rather than stalling the caller. + Ok(Err(_)) | Err(async_channel::TryRecvError::Closed) => Some(Vec::new()), + // Not ready yet: keep the pending state and report back later. + Err(async_channel::TryRecvError::Empty) => { + self.read_state = TimestampReadState::Pending { rx, query_count }; + None + } + }, + } + } + + /// Converts raw `(begin, end)` tick pairs into labeled millisecond durations. + fn durations_from_raw(&self, raw: &[u64]) -> Vec { let period_ms = self.timestamp_period as f64 / 1_000_000.0; - let results = self - .labels + self.labels .iter() .enumerate() .map(|(i, label)| { - let begin = raw_timestamps[i * 2]; - let end = raw_timestamps[i * 2 + 1]; + let begin = raw[i * 2]; + let end = raw[i * 2 + 1]; let duration_ms = (end.wrapping_sub(begin)) as f64 * period_ms; GpuTimestamp { label: label.clone(), duration_ms, } }) - .collect(); - Ok(results) + .collect() } /// Allocates a begin/end timestamp pair for a labeled pass.