diff --git a/crates/cubecl-wgpu/Cargo.toml b/crates/cubecl-wgpu/Cargo.toml index 0c5ddca1d9..568bbfb470 100644 --- a/crates/cubecl-wgpu/Cargo.toml +++ b/crates/cubecl-wgpu/Cargo.toml @@ -66,7 +66,6 @@ cubecl-cpp = { package = "t4a-cubecl-cpp", path = "../cubecl-cpp", version = "0. bytemuck = { workspace = true } async-channel = { workspace = true } -derive-new = { workspace = true } hashbrown = { workspace = true } cfg-if = { workspace = true } diff --git a/crates/cubecl-wgpu/src/backend/base.rs b/crates/cubecl-wgpu/src/backend/base.rs index b5f1b14503..01cf7c0057 100644 --- a/crates/cubecl-wgpu/src/backend/base.rs +++ b/crates/cubecl-wgpu/src/backend/base.rs @@ -1,6 +1,5 @@ use super::wgsl; -use crate::AutoRepresentationRef; -use crate::WgpuServer; +use crate::{AutoRepresentationRef, PrimaryMemoryMode, WgpuServer}; use cubecl_core::MemoryConfiguration; use cubecl_core::{ ExecutionMode, WgpuCompilationOptions, hash::StableHash, server::KernelArguments, @@ -234,41 +233,77 @@ impl WgpuServer { } } -pub async fn request_device(adapter: &Adapter) -> (Device, Queue) { - if let Some(result) = request_vulkan_device(adapter).await { +pub(crate) const HOST_VISIBLE_PRIMARY_UNSUPPORTED: &str = "Host-visible primary memory requires wgpu::Features::MAPPABLE_PRIMARY_BUFFERS, but the selected adapter or device does not support it."; + +pub(crate) fn requested_features( + adapter: &Adapter, + primary_memory: PrimaryMemoryMode, +) -> wgpu::Features { + let features = adapter.features(); + match primary_memory { + PrimaryMemoryMode::DeviceLocal => { + features.difference(wgpu::Features::MAPPABLE_PRIMARY_BUFFERS) + } + PrimaryMemoryMode::HostVisible => { + if !features.contains(wgpu::Features::MAPPABLE_PRIMARY_BUFFERS) { + panic!("{HOST_VISIBLE_PRIMARY_UNSUPPORTED}"); + } + features + } + } +} + +pub async fn request_device( + adapter: &Adapter, + primary_memory: PrimaryMemoryMode, +) -> (Device, Queue) { + let _ = requested_features(adapter, primary_memory); + if let Some(result) = request_vulkan_device(adapter, primary_memory).await { return result; } - if let Some(result) = request_metal_device(adapter).await { + if let Some(result) = request_metal_device(adapter, primary_memory).await { return result; } - wgsl::request_device(adapter).await + wgsl::request_device(adapter, primary_memory).await } #[cfg(feature = "spirv")] -async fn request_vulkan_device(adapter: &Adapter) -> Option<(Device, Queue)> { +async fn request_vulkan_device( + adapter: &Adapter, + primary_memory: PrimaryMemoryMode, +) -> Option<(Device, Queue)> { if is_vulkan(adapter) { - vulkan::request_vulkan_device(adapter).await + vulkan::request_vulkan_device(adapter, primary_memory).await } else { None } } #[cfg(not(feature = "spirv"))] -async fn request_vulkan_device(_adapter: &Adapter) -> Option<(Device, Queue)> { +async fn request_vulkan_device( + _adapter: &Adapter, + _primary_memory: PrimaryMemoryMode, +) -> Option<(Device, Queue)> { None } #[cfg(all(feature = "msl", target_os = "macos"))] -async fn request_metal_device(adapter: &Adapter) -> Option<(Device, Queue)> { +async fn request_metal_device( + adapter: &Adapter, + primary_memory: PrimaryMemoryMode, +) -> Option<(Device, Queue)> { if is_metal(adapter) { - Some(metal::request_metal_device(adapter).await) + Some(metal::request_metal_device(adapter, primary_memory).await) } else { None } } #[cfg(not(all(feature = "msl", target_os = "macos")))] -async fn request_metal_device(_adapter: &Adapter) -> Option<(Device, Queue)> { +async fn request_metal_device( + _adapter: &Adapter, + _primary_memory: PrimaryMemoryMode, +) -> Option<(Device, Queue)> { None } diff --git a/crates/cubecl-wgpu/src/backend/metal.rs b/crates/cubecl-wgpu/src/backend/metal.rs index 244e7965c2..fc82209b0a 100644 --- a/crates/cubecl-wgpu/src/backend/metal.rs +++ b/crates/cubecl-wgpu/src/backend/metal.rs @@ -16,11 +16,14 @@ use wgpu::{ hal::{self, Adapter, metal}, }; -pub async fn request_metal_device(adapter: &wgpu::Adapter) -> (wgpu::Device, wgpu::Queue) { +use crate::{PrimaryMemoryMode, backend::requested_features}; + +pub async fn request_metal_device( + adapter: &wgpu::Adapter, + primary_memory: PrimaryMemoryMode, +) -> (wgpu::Device, wgpu::Queue) { let limits = adapter.limits(); - let features = adapter - .features() - .difference(Features::MAPPABLE_PRIMARY_BUFFERS); + let features = requested_features(adapter, primary_memory); unsafe { let hal_adapter = adapter.as_hal::().unwrap(); request_device(adapter, &hal_adapter, features, limits) diff --git a/crates/cubecl-wgpu/src/backend/vulkan.rs b/crates/cubecl-wgpu/src/backend/vulkan.rs index 4f1908cd79..a136c54be4 100644 --- a/crates/cubecl-wgpu/src/backend/vulkan.rs +++ b/crates/cubecl-wgpu/src/backend/vulkan.rs @@ -21,7 +21,7 @@ use wgpu::{ }, }; -use crate::{AutoCompiler, WgpuServer}; +use crate::{AutoCompiler, PrimaryMemoryMode, WgpuServer, backend::requested_features}; mod features; @@ -36,11 +36,12 @@ pub fn bindings( (buffers, meta, repr.uniform_info) } -pub async fn request_vulkan_device(adapter: &wgpu::Adapter) -> Option<(wgpu::Device, wgpu::Queue)> { +pub async fn request_vulkan_device( + adapter: &wgpu::Adapter, + primary_memory: PrimaryMemoryMode, +) -> Option<(wgpu::Device, wgpu::Queue)> { let limits = adapter.limits(); - let features = adapter - .features() - .difference(Features::MAPPABLE_PRIMARY_BUFFERS); + let features = requested_features(adapter, primary_memory); unsafe { let hal_adapter = adapter.as_hal::().unwrap(); request_device(adapter, &hal_adapter, features, limits) diff --git a/crates/cubecl-wgpu/src/backend/wgsl.rs b/crates/cubecl-wgpu/src/backend/wgsl.rs index cef99135c5..0c3aa3e600 100644 --- a/crates/cubecl-wgpu/src/backend/wgsl.rs +++ b/crates/cubecl-wgpu/src/backend/wgsl.rs @@ -1,12 +1,10 @@ +use crate::{PrimaryMemoryMode, WgslCompiler, backend::requested_features}; use cubecl_core::{Compiler, prelude::Visibility, server::KernelArguments}; use cubecl_core::{ WgpuCompilationOptions, ir::{ElemType, UIntKind}, }; use cubecl_ir::{DeviceProperties, Type}; -use wgpu::Features; - -use crate::WgslCompiler; pub fn bindings( repr: &::Representation, @@ -27,14 +25,15 @@ pub fn bindings( (bindings, meta, false) } -pub async fn request_device(adapter: &wgpu::Adapter) -> (wgpu::Device, wgpu::Queue) { +pub async fn request_device( + adapter: &wgpu::Adapter, + primary_memory: PrimaryMemoryMode, +) -> (wgpu::Device, wgpu::Queue) { let limits = adapter.limits(); adapter .request_device(&wgpu::DeviceDescriptor { label: None, - required_features: adapter - .features() - .difference(Features::MAPPABLE_PRIMARY_BUFFERS), + required_features: requested_features(adapter, primary_memory), required_limits: limits, // The default is MemoryHints::Performance, which tries to do some bigger // block allocations. However, we already batch allocations, so we diff --git a/crates/cubecl-wgpu/src/compute/mem_manager.rs b/crates/cubecl-wgpu/src/compute/mem_manager.rs index d03f795270..ebf2ec9202 100644 --- a/crates/cubecl-wgpu/src/compute/mem_manager.rs +++ b/crates/cubecl-wgpu/src/compute/mem_manager.rs @@ -1,4 +1,4 @@ -use crate::{WgpuResource, WgpuStorage}; +use crate::{PrimaryMemoryMode, WgpuResource, WgpuStorage}; use cubecl_common::stub::Arc; use cubecl_core::{ MemoryConfiguration, @@ -28,22 +28,42 @@ impl WgpuMemManager { device: wgpu::Device, memory_properties: MemoryDeviceProperties, memory_config: MemoryConfiguration, + primary_memory: PrimaryMemoryMode, logger: Arc, ) -> Self { + let (memory_config_main, main_usages, host_visible) = match primary_memory { + PrimaryMemoryMode::DeviceLocal => ( + memory_config, + BufferUsages::STORAGE + | BufferUsages::COPY_SRC + | BufferUsages::COPY_DST + | BufferUsages::INDIRECT, + false, + ), + PrimaryMemoryMode::HostVisible => ( + MemoryConfiguration::ExclusivePages, + BufferUsages::STORAGE + | BufferUsages::COPY_SRC + | BufferUsages::COPY_DST + | BufferUsages::INDIRECT + | BufferUsages::MAP_READ + | BufferUsages::MAP_WRITE, + true, + ), + }; + // Allocate storage & memory management for the main memory buffers. Any calls // to empty() or create() with a small enough size will be allocated from this // main memory pool. let memory_main = MemoryManagement::from_configuration( - WgpuStorage::new( + WgpuStorage::new_with_host_visibility( memory_properties.alignment as usize, device.clone(), - BufferUsages::STORAGE - | BufferUsages::COPY_SRC - | BufferUsages::COPY_DST - | BufferUsages::INDIRECT, + main_usages, + host_visible, ), &memory_properties, - memory_config, + memory_config_main, logger.clone(), MemoryManagementOptions::new("Main GPU Memory"), ); @@ -104,14 +124,17 @@ impl WgpuMemManager { let resource = self .memory_pool_staging .get_resource(binding.clone(), None, None) - .unwrap(); + .unwrap() + .with_lease(binding.clone()); Ok((resource, binding)) } pub(crate) fn get_resource(&mut self, binding: Binding) -> Result { + let lease = binding.memory.clone(); self.memory_pool .get_resource(binding.memory, binding.offset_start, binding.offset_end) + .map(|resource| resource.with_lease(lease)) } pub(crate) fn reserve_uniform(&mut self, size: u64) -> WgpuResource { @@ -121,11 +144,15 @@ impl WgpuMemManager { .expect("Must have enough memory for a uniform"); // Keep track of this uniform until it is released. self.uniforms.push(slice.clone()); + let binding = slice.binding(); let handle = self .memory_uniforms - .get_storage(slice.binding()) + .get_storage(binding.clone()) .expect("Failed to find storage!"); - self.memory_uniforms.storage().get(&handle) + self.memory_uniforms + .storage() + .get(&handle) + .with_lease(binding) } pub(crate) fn memory_usage(&self) -> cubecl_runtime::memory_management::MemoryUsage { diff --git a/crates/cubecl-wgpu/src/compute/schedule.rs b/crates/cubecl-wgpu/src/compute/schedule.rs index 55897489d7..0962cb0c88 100644 --- a/crates/cubecl-wgpu/src/compute/schedule.rs +++ b/crates/cubecl-wgpu/src/compute/schedule.rs @@ -1,4 +1,4 @@ -use crate::{WgpuResource, stream::WgpuStream}; +use crate::{GpuAccessToken, HostAccessError, PrimaryMemoryMode, WgpuResource, stream::WgpuStream}; use alloc::sync::Arc; use cubecl_common::{bytes::Bytes, profile::TimingMethod}; use cubecl_core::{ @@ -19,6 +19,8 @@ pub enum ScheduleTask { data: Bytes, /// The target buffer resource. buffer: WgpuResource, + /// Reservation held until the queue write is submitted. + reservation: GpuAccessToken, }, /// Represents a task to execute a compute pipeline. Execute { @@ -50,6 +52,7 @@ impl core::fmt::Debug for ScheduleTask { pub struct BindingsResource { /// List of WGPU resources used in the task. pub resources: Vec, + pub(crate) reservations: Vec, /// Metadata for uniform bindings. pub info: MetadataBindingInfo, } @@ -68,6 +71,7 @@ pub struct WgpuStreamFactory { queue: wgpu::Queue, memory_properties: MemoryDeviceProperties, memory_config: MemoryConfiguration, + primary_memory: PrimaryMemoryMode, timing_method: TimingMethod, tasks_max: usize, logger: Arc, @@ -85,6 +89,7 @@ impl StreamFactory for WgpuStreamFactory { self.queue.clone(), self.memory_properties.clone(), self.memory_config.clone(), + self.primary_memory, self.timing_method, self.tasks_max, self.logger.clone(), @@ -94,11 +99,13 @@ impl StreamFactory for WgpuStreamFactory { impl ScheduledWgpuBackend { /// Creates a new `ScheduledWgpuBackend` with the given WGPU device, queue, and configurations. + #[allow(clippy::too_many_arguments)] pub fn new( device: wgpu::Device, queue: wgpu::Queue, memory_properties: MemoryDeviceProperties, memory_config: MemoryConfiguration, + primary_memory: PrimaryMemoryMode, timing_method: TimingMethod, tasks_max: usize, logger: Arc, @@ -109,6 +116,7 @@ impl ScheduledWgpuBackend { queue, memory_properties, memory_config, + primary_memory, timing_method, tasks_max, logger, @@ -120,15 +128,19 @@ impl ScheduledWgpuBackend { impl BindingsResource { /// Converts metadata and scalar bindings into WGPU resources for a stream. - pub fn into_resources(mut self, stream: &mut WgpuStream) -> Vec { + pub fn into_resources( + mut self, + stream: &mut WgpuStream, + ) -> Result<(Vec, Vec), HostAccessError> { // If metadata contains data, create a uniform buffer for it. if !self.info.data.is_empty() { let info = stream.create_uniform(bytemuck::cast_slice(&self.info.data)); + self.reservations.push(info.acquire_gpu()?); self.resources.push(info); } // Return the complete list of resources. - self.resources + Ok((self.resources, self.reservations)) } } diff --git a/crates/cubecl-wgpu/src/compute/server.rs b/crates/cubecl-wgpu/src/compute/server.rs index 36199a934f..04cdc97e22 100644 --- a/crates/cubecl-wgpu/src/compute/server.rs +++ b/crates/cubecl-wgpu/src/compute/server.rs @@ -1,6 +1,6 @@ use super::storage::{WgpuResource, WgpuStorage}; use crate::schedule::{BindingsResource, ScheduleTask, ScheduledWgpuBackend}; -use crate::{AutoCompiler, AutoRepresentation}; +use crate::{AutoCompiler, AutoRepresentation, PrimaryMemoryMode}; use alloc::sync::Arc; use cubecl_common::{ backtrace::BackTrace, @@ -68,6 +68,7 @@ impl WgpuServer { device: wgpu::Device, queue: wgpu::Queue, tasks_max: usize, + primary_memory: PrimaryMemoryMode, backend: wgpu::Backend, timing_method: TimingMethod, utilities: ServerUtilities, @@ -77,6 +78,7 @@ impl WgpuServer { queue.clone(), memory_properties, memory_config, + primary_memory, timing_method, tasks_max, utilities.logger.clone(), @@ -117,19 +119,25 @@ impl WgpuServer { } } - fn prepare_bindings(&mut self, bindings: KernelArguments) -> Result { + fn prepare_bindings( + &mut self, + bindings: KernelArguments, + ) -> Result { // Store all the resources we'll be using. This could be eliminated if // there was a way to tie the lifetime of the resource to the memory handle. let mut resources = Vec::with_capacity(bindings.buffers.len()); + let mut reservations = Vec::with_capacity(bindings.buffers.len()); for b in bindings.buffers.into_iter() { let stream = self.scheduler.stream(&b.stream); let resource = stream.mem_manage.get_resource(b)?; + reservations.push(resource.acquire_gpu()?); resources.push(resource); } Ok(BindingsResource { resources, + reservations, info: bindings.info, }) } @@ -295,6 +303,9 @@ impl ComputeServer for WgpuServer { Ok(val) => val, Err(err) => return Box::pin(async move { Err(err.into()) }), }; + if let Err(err) = resource.ensure_gpu_access() { + return Box::pin(async move { Err(err.into()) }); + } resources.push((resource, desc.shape, desc.elem_size)); } @@ -322,9 +333,17 @@ impl ComputeServer for WgpuServer { return; } }; + let reservation = match resource.acquire_gpu() { + Ok(reservation) => reservation, + Err(err) => { + stream.error(err.into()); + return; + } + }; let task = ScheduleTask::Write { data, buffer: resource, + reservation, }; self.scheduler.register(stream_id, task, &[]); @@ -376,7 +395,7 @@ impl ComputeServer for WgpuServer { Err(err) => { // We make the stream that would execute the kernel in error. let stream = self.scheduler.stream(&stream_id); - stream.errors.push(ServerError::Io(err)); + stream.errors.push(err); return; } }; @@ -463,3 +482,60 @@ pub(crate) fn contiguous_strides(shape: &Shape) -> Strides { } strides } + +#[cfg(all(test, target_os = "macos"))] +mod tests { + use super::*; + use crate::{HostAccessError, RuntimeOptions, WgpuDevice, runtime}; + use cubecl_common::future; + use cubecl_runtime::server::Handle; + + #[test] + fn scheduled_write_reserves_gpu_access_until_queue_submission() { + let instance = wgpu::Instance::new(wgpu::InstanceDescriptor { + backends: wgpu::Backends::METAL, + ..wgpu::InstanceDescriptor::new_without_display_handle() + }); + if future::block_on(instance.enumerate_adapters(wgpu::Backends::METAL)).is_empty() { + eprintln!("skipping scheduled-write reservation test: no Metal adapter is available"); + return; + } + + let options = RuntimeOptions { + tasks_max: 32, + memory_config: MemoryConfiguration::ExclusivePages, + primary_memory: PrimaryMemoryMode::HostVisible, + }; + let setup = future::block_on(runtime::create_setup_for_device( + &WgpuDevice::DefaultDevice, + wgpu::Backend::Metal, + options.primary_memory, + )); + let mut server = runtime::create_server(setup, options); + let stream_id = StreamId::current(); + let handle = Handle::new(stream_id, 16); + server.initialize_memory(handle.memory.clone(), 16, stream_id); + let managed = server + .get_resource(handle.clone().binding(), stream_id) + .unwrap(); + let resource = managed.resource().clone(); + + server.write( + vec![( + CopyDescriptor::new(handle.binding(), [16].into(), [1].into(), 1), + Bytes::from_bytes_vec(vec![7; 16]), + )], + stream_id, + ); + + let host_attempt = std::thread::spawn(move || resource.map_read().unwrap_err()); + assert_eq!( + host_attempt.join().unwrap(), + HostAccessError::GpuAccessInProgress + ); + + server.flush(stream_id).unwrap(); + let guard = managed.resource().map_read().unwrap(); + assert_eq!(&*guard, &[7; 16]); + } +} diff --git a/crates/cubecl-wgpu/src/compute/storage.rs b/crates/cubecl-wgpu/src/compute/storage.rs index 479d51dd20..32cc56c4f3 100644 --- a/crates/cubecl-wgpu/src/compute/storage.rs +++ b/crates/cubecl-wgpu/src/compute/storage.rs @@ -1,7 +1,18 @@ -use cubecl_core::server::IoError; -use cubecl_runtime::storage::{ComputeStorage, StorageHandle, StorageId, StorageUtilization}; +use cubecl_common::backtrace::BackTrace; +use cubecl_core::server::{IoError, ServerError}; +use cubecl_runtime::{ + memory_management::ManagedMemoryBinding, + storage::{ComputeStorage, StorageHandle, StorageId, StorageUtilization}, +}; use hashbrown::HashMap; -use std::num::NonZeroU64; +use std::{ + num::NonZeroU64, + ops::Deref, + sync::{ + Arc, + atomic::{AtomicU64, AtomicUsize, Ordering}, + }, +}; use wgpu::BufferUsages; /// Minimum buffer size in bytes. The WebGPU spec requires buffer sizes > 0, and shaders @@ -9,12 +20,203 @@ use wgpu::BufferUsages; /// 32 bytes covers the largest possible binding type (`vec4`). const MIN_BUFFER_SIZE: u64 = 32; +const ACCESS_GPU_IDLE: usize = 0; +const ACCESS_HOST_MAPPED: usize = 1 << (usize::BITS - 1); +const ACCESS_GPU_COUNT_MASK: usize = ACCESS_HOST_MAPPED - 1; + +static NEXT_ALLOCATION_ID: AtomicU64 = AtomicU64::new(1); + +/// Error returned when host and GPU access to a WGPU allocation cannot be coordinated. +#[non_exhaustive] +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum HostAccessError { + /// The allocation was created without host-visible primary memory enabled. + DeviceLocalAllocation, + /// Another host mapping is already active for the physical allocation. + OverlappingHostMapping, + /// GPU work has reserved the allocation and has not yet been submitted. + GpuAccessInProgress, + /// GPU use was requested while a host mapping is active. + MappedForHost, + /// WGPU failed to complete the asynchronous map callback. + MapCallbackFailure(String), + /// WGPU failed while polling the device for mapping completion. + DevicePollFailure(String), +} + +impl core::fmt::Display for HostAccessError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Self::DeviceLocalAllocation => { + f.write_str("the WGPU allocation is device-local and cannot be mapped by the host") + } + Self::OverlappingHostMapping => { + f.write_str("the WGPU allocation already has an active host mapping") + } + Self::GpuAccessInProgress => f.write_str( + "the WGPU allocation has GPU work pending submission and cannot be mapped by the host", + ), + Self::MappedForHost => f.write_str( + "the WGPU allocation is mapped for host access and cannot be used by the GPU", + ), + Self::MapCallbackFailure(reason) => { + write!(f, "the WGPU map callback failed: {reason}") + } + Self::DevicePollFailure(reason) => { + write!( + f, + "polling the WGPU device for host mapping failed: {reason}" + ) + } + } + } +} + +impl std::error::Error for HostAccessError {} + +impl From for ServerError { + fn from(error: HostAccessError) -> Self { + Self::Generic { + reason: error.to_string(), + backtrace: BackTrace::capture(), + } + } +} + +#[derive(Debug)] +struct AllocationAccessInner { + allocation_id: u64, + state: AtomicUsize, +} + +#[derive(Clone, Debug)] +struct AllocationAccess { + inner: Arc, +} + +impl AllocationAccess { + fn new() -> Self { + let allocation_id = NEXT_ALLOCATION_ID + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + current.checked_add(1) + }) + .expect("WGPU allocation ID overflowed"); + + Self { + inner: Arc::new(AllocationAccessInner { + allocation_id, + state: AtomicUsize::new(ACCESS_GPU_IDLE), + }), + } + } + + fn allocation_id(&self) -> u64 { + self.inner.allocation_id + } + + fn acquire_host(&self) -> Result { + self.inner + .state + .compare_exchange( + ACCESS_GPU_IDLE, + ACCESS_HOST_MAPPED, + Ordering::Acquire, + Ordering::Relaxed, + ) + .map_err(|state| { + if state == ACCESS_HOST_MAPPED { + HostAccessError::OverlappingHostMapping + } else { + HostAccessError::GpuAccessInProgress + } + })?; + + Ok(HostAccessToken { + access: Some(self.inner.clone()), + }) + } + + fn ensure_gpu_access(&self) -> Result<(), HostAccessError> { + if self.inner.state.load(Ordering::Acquire) == ACCESS_HOST_MAPPED { + Err(HostAccessError::MappedForHost) + } else { + Ok(()) + } + } + + fn acquire_gpu(&self) -> Result { + let mut state = self.inner.state.load(Ordering::Acquire); + loop { + if state == ACCESS_HOST_MAPPED { + return Err(HostAccessError::MappedForHost); + } + assert_ne!( + state, ACCESS_GPU_COUNT_MASK, + "too many simultaneous WGPU GPU reservations" + ); + match self.inner.state.compare_exchange_weak( + state, + state + 1, + Ordering::Acquire, + Ordering::Relaxed, + ) { + Ok(_) => { + return Ok(GpuAccessToken { + access: Some(self.inner.clone()), + lease: None, + }); + } + Err(current) => state = current, + } + } + } +} + +/// RAII reservation preventing host mapping until queued GPU use has been submitted. +#[doc(hidden)] +#[derive(Debug)] +pub struct GpuAccessToken { + access: Option>, + lease: Option, +} + +impl Drop for GpuAccessToken { + fn drop(&mut self) { + if let Some(access) = self.access.take() { + let previous = access.state.fetch_sub(1, Ordering::Release); + debug_assert!(previous > 0 && previous < ACCESS_HOST_MAPPED); + } + } +} + +#[derive(Debug)] +struct HostAccessToken { + access: Option>, +} + +impl Drop for HostAccessToken { + fn drop(&mut self) { + if let Some(access) = self.access.take() { + access.state.store(ACCESS_GPU_IDLE, Ordering::Release); + } + } +} + +#[derive(Debug)] +struct WgpuAllocation { + buffer: wgpu::Buffer, + device: Option, + allocation_access: AllocationAccess, + host_visible: bool, +} + /// Buffer storage for wgpu. pub struct WgpuStorage { - memory: HashMap, + memory: HashMap>, device: wgpu::Device, buffer_usages: BufferUsages, mem_alignment: usize, + host_visible: bool, } impl core::fmt::Debug for WgpuStorage { @@ -24,7 +226,14 @@ impl core::fmt::Debug for WgpuStorage { } /// The memory resource that can be allocated for wgpu. -#[derive(new, Debug)] +/// +/// For [`crate::WgpuRuntime`], callers can resolve a public `CubeCL` allocation handle with +/// [`cubecl_runtime::client::ComputeClient::get_resource`]. Call the returned +/// [`cubecl_runtime::storage::ManagedResource`]'s `resource` method to borrow this value, and then +/// use [`Self::allocation_id`], [`Self::map_read`], or [`Self::map_write`]. Managed resources, +/// cloned `WgpuResource` values, mapped guards, and queued GPU reservations each retain the +/// memory-manager lease for as long as they can access the allocation. +#[derive(Clone, Debug)] pub struct WgpuResource { /// The wgpu buffer. pub buffer: wgpu::Buffer, @@ -36,9 +245,200 @@ pub struct WgpuResource { /// /// The result considers the offset. pub size: u64, + allocation: Arc, + lease: Option, } impl WgpuResource { + /// Creates a device-local resource from an existing WGPU buffer. + /// + /// Resources allocated by [`WgpuStorage`] additionally retain their device so they can opt in + /// to synchronous host mapping. + pub fn new(buffer: wgpu::Buffer, offset: u64, size: u64) -> Self { + let allocation = Arc::new(WgpuAllocation { + buffer: buffer.clone(), + device: None, + allocation_access: AllocationAccess::new(), + host_visible: false, + }); + + Self { + buffer, + offset, + size, + allocation, + lease: None, + } + } + + fn from_allocation(allocation: Arc, offset: u64, size: u64) -> Self { + Self { + buffer: allocation.buffer.clone(), + offset, + size, + allocation, + lease: None, + } + } + + pub(crate) fn with_lease(mut self, lease: ManagedMemoryBinding) -> Self { + self.lease = Some(lease); + self + } + + /// Returns the process-unique identity of the underlying physical allocation. + pub fn allocation_id(&self) -> u64 { + self.allocation.allocation_access.allocation_id() + } + + /// Validates that the resource can currently be used by the GPU. + pub fn ensure_gpu_access(&self) -> Result<(), HostAccessError> { + self.allocation.allocation_access.ensure_gpu_access() + } + + pub(crate) fn acquire_gpu(&self) -> Result { + let mut token = self.allocation.allocation_access.acquire_gpu()?; + token.lease = self.lease.clone(); + Ok(token) + } + + fn acquire_host(&self) -> Result { + if !self.allocation.host_visible { + return Err(HostAccessError::DeviceLocalAllocation); + } + self.allocation.allocation_access.acquire_host() + } + + #[cfg(not(target_family = "wasm"))] + fn map_range(&self) -> core::ops::Range { + let end = self.offset + self.size; + self.offset..end.next_multiple_of(wgpu::COPY_BUFFER_ALIGNMENT) + } + + /// Maps the exact logical resource range for synchronous host reads. + #[cfg(not(target_family = "wasm"))] + pub fn map_read(&self) -> Result { + let token = self.acquire_host()?; + + if self.size == 0 { + return Ok(WgpuMappedReadGuard { + allocation: self.allocation.clone(), + view: None, + logical_len: 0, + token: Some(token), + _lease: self.lease.clone(), + }); + } + + let range = self.map_range(); + let (sender, receiver) = std::sync::mpsc::sync_channel(1); + self.buffer + .map_async(wgpu::MapMode::Read, range.clone(), move |result| { + let _ = sender.send(result); + }); + + let device = self + .allocation + .device + .as_ref() + .expect("host-visible WGPU allocations always retain their device"); + if let Err(err) = device.poll(wgpu::PollType::wait_indefinitely()) { + self.buffer.unmap(); + return Err(HostAccessError::DevicePollFailure(err.to_string())); + } + + let callback_result = match receiver.recv() { + Ok(result) => result, + Err(err) => { + self.buffer.unmap(); + return Err(HostAccessError::MapCallbackFailure(err.to_string())); + } + }; + if let Err(err) = callback_result { + self.buffer.unmap(); + return Err(HostAccessError::MapCallbackFailure(err.to_string())); + } + + let view = self.buffer.get_mapped_range(range); + Ok(WgpuMappedReadGuard { + allocation: self.allocation.clone(), + view: Some(view), + logical_len: self.size as usize, + token: Some(token), + _lease: self.lease.clone(), + }) + } + + /// Maps the exact logical resource range for synchronous host writes. + #[cfg(not(target_family = "wasm"))] + pub fn map_write(&self) -> Result { + let token = self.acquire_host()?; + + if self.size == 0 { + return Ok(WgpuMappedWriteGuard { + allocation: self.allocation.clone(), + view: None, + logical_len: 0, + token: Some(token), + _lease: self.lease.clone(), + }); + } + + let range = self.map_range(); + let (sender, receiver) = std::sync::mpsc::sync_channel(1); + self.buffer + .map_async(wgpu::MapMode::Write, range.clone(), move |result| { + let _ = sender.send(result); + }); + + let device = self + .allocation + .device + .as_ref() + .expect("host-visible WGPU allocations always retain their device"); + if let Err(err) = device.poll(wgpu::PollType::wait_indefinitely()) { + self.buffer.unmap(); + return Err(HostAccessError::DevicePollFailure(err.to_string())); + } + + let callback_result = match receiver.recv() { + Ok(result) => result, + Err(err) => { + self.buffer.unmap(); + return Err(HostAccessError::MapCallbackFailure(err.to_string())); + } + }; + if let Err(err) = callback_result { + self.buffer.unmap(); + return Err(HostAccessError::MapCallbackFailure(err.to_string())); + } + + let view = self.buffer.get_mapped_range_mut(range); + Ok(WgpuMappedWriteGuard { + allocation: self.allocation.clone(), + view: Some(view), + logical_len: self.size as usize, + token: Some(token), + _lease: self.lease.clone(), + }) + } + + /// Synchronous mapping is unavailable on WebAssembly because `Device::poll` cannot block. + #[cfg(target_family = "wasm")] + pub fn map_read(&self) -> Result { + Err(HostAccessError::DevicePollFailure( + "synchronous host mapping is unsupported on WebAssembly".to_string(), + )) + } + + /// Synchronous mapping is unavailable on WebAssembly because `Device::poll` cannot block. + #[cfg(target_family = "wasm")] + pub fn map_write(&self) -> Result { + Err(HostAccessError::DevicePollFailure( + "synchronous host mapping is unsupported on WebAssembly".to_string(), + )) + } + /// Return the binding view of the buffer. pub fn as_wgpu_bind_resource(&self) -> wgpu::BindingResource<'_> { // wgpu enforces 4-byte alignment for buffer binding sizes per the WebGPU spec. @@ -59,17 +459,114 @@ impl WgpuResource { }; wgpu::BindingResource::Buffer(binding) } + + /// Returns the binding view after checking that no host mapping is active. + pub fn try_as_wgpu_bind_resource(&self) -> Result, HostAccessError> { + self.ensure_gpu_access()?; + Ok(self.as_wgpu_bind_resource()) + } +} + +/// Owning RAII guard for a synchronously mapped host-readable WGPU resource. +#[derive(Debug)] +pub struct WgpuMappedReadGuard { + allocation: Arc, + view: Option, + logical_len: usize, + token: Option, + _lease: Option, +} + +impl Deref for WgpuMappedReadGuard { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + match &self.view { + Some(view) => &view[..self.logical_len], + None => &[], + } + } +} + +impl Drop for WgpuMappedReadGuard { + fn drop(&mut self) { + drop(self.view.take()); + if self.logical_len != 0 { + self.allocation.buffer.unmap(); + } + drop(self.token.take()); + } +} + +/// Owning RAII guard for a synchronously mapped host-writable WGPU resource. +/// +/// WGPU write mappings may be write-combining memory, so wgpu 29 intentionally exposes them as +/// [`wgpu::WriteOnly`] rather than `&mut [u8]`. Use [`Self::as_write_only`] for zero-copy writes. +#[derive(Debug)] +pub struct WgpuMappedWriteGuard { + allocation: Arc, + view: Option, + logical_len: usize, + token: Option, + _lease: Option, +} + +impl WgpuMappedWriteGuard { + /// Returns a zero-copy write-only view of the exact logical resource range. + pub fn as_write_only(&mut self) -> Option> { + self.view + .as_mut() + .map(|view| view.slice(..self.logical_len)) + } + + /// Copies bytes directly into the mapped logical resource range. + pub fn copy_from_slice(&mut self, bytes: &[u8]) { + assert_eq!(bytes.len(), self.logical_len); + if let Some(mut view) = self.as_write_only() { + view.copy_from_slice(bytes); + } + } + + /// Returns the logical resource length in bytes. + pub fn len(&self) -> usize { + self.logical_len + } + + /// Returns whether the logical resource is empty. + pub fn is_empty(&self) -> bool { + self.logical_len == 0 + } +} + +impl Drop for WgpuMappedWriteGuard { + fn drop(&mut self) { + drop(self.view.take()); + if self.logical_len != 0 { + self.allocation.buffer.unmap(); + } + drop(self.token.take()); + } } /// Keeps actual wgpu buffer references in a hashmap with ids as key. impl WgpuStorage { /// Create a new storage on the given [device](wgpu::Device). pub fn new(mem_alignment: usize, device: wgpu::Device, usages: BufferUsages) -> Self { + Self::new_with_host_visibility(mem_alignment, device, usages, false) + } + + pub(crate) fn new_with_host_visibility( + mem_alignment: usize, + device: wgpu::Device, + usages: BufferUsages, + host_visible: bool, + ) -> Self { Self { memory: HashMap::new(), device, buffer_usages: usages, mem_alignment, + host_visible, } } } @@ -82,8 +579,8 @@ impl ComputeStorage for WgpuStorage { } fn get(&mut self, handle: &StorageHandle) -> Self::Resource { - let buffer = self.memory.get(&handle.id).unwrap(); - WgpuResource::new(buffer.clone(), handle.offset(), handle.size()) + let allocation = self.memory.get(&handle.id).unwrap(); + WgpuResource::from_allocation(allocation.clone(), handle.offset(), handle.size()) } #[cfg_attr( @@ -102,7 +599,15 @@ impl ComputeStorage for WgpuStorage { mapped_at_creation: false, }); - self.memory.insert(id, buffer); + self.memory.insert( + id, + Arc::new(WgpuAllocation { + buffer, + device: Some(self.device.clone()), + allocation_access: AllocationAccess::new(), + host_visible: self.host_visible, + }), + ); Ok(StorageHandle::new( id, StorageUtilization { offset: 0, size }, @@ -118,3 +623,62 @@ impl ComputeStorage for WgpuStorage { // We don't wait for dealloc } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cloned_allocation_access_retains_nonzero_identity() { + let access = Arc::new(AllocationAccess::new()); + let cloned = access.clone(); + + assert_ne!(access.allocation_id(), 0); + assert_eq!(access.allocation_id(), cloned.allocation_id()); + } + + #[test] + fn host_access_excludes_other_host_and_gpu_access_until_drop() { + let access = AllocationAccess::new(); + let token = access.acquire_host().unwrap(); + + assert_eq!( + access.acquire_host().unwrap_err(), + HostAccessError::OverlappingHostMapping + ); + assert_eq!( + access.ensure_gpu_access(), + Err(HostAccessError::MappedForHost) + ); + + drop(token); + + assert_eq!(access.ensure_gpu_access(), Ok(())); + assert!(access.acquire_host().is_ok()); + } + + #[test] + fn gpu_reservation_excludes_host_access_until_submission() { + let access = AllocationAccess::new(); + let worker_access = access.clone(); + let (reserved_tx, reserved_rx) = std::sync::mpsc::sync_channel(1); + let (submitted_tx, submitted_rx) = std::sync::mpsc::sync_channel(1); + + let worker = std::thread::spawn(move || { + let reservation = worker_access.acquire_gpu().unwrap(); + reserved_tx.send(()).unwrap(); + submitted_rx.recv().unwrap(); + drop(reservation); + }); + + reserved_rx.recv().unwrap(); + assert_eq!( + access.acquire_host().unwrap_err(), + HostAccessError::GpuAccessInProgress + ); + + submitted_tx.send(()).unwrap(); + worker.join().unwrap(); + assert!(access.acquire_host().is_ok()); + } +} diff --git a/crates/cubecl-wgpu/src/compute/stream.rs b/crates/cubecl-wgpu/src/compute/stream.rs index e723a240a1..6d28169a0f 100644 --- a/crates/cubecl-wgpu/src/compute/stream.rs +++ b/crates/cubecl-wgpu/src/compute/stream.rs @@ -1,5 +1,7 @@ use super::{mem_manager::WgpuMemManager, poll::WgpuPoll, timings::QueryProfiler}; -use crate::{WgpuResource, controller::WgpuAllocController, schedule::ScheduleTask}; +use crate::{ + GpuAccessToken, WgpuResource, controller::WgpuAllocController, schedule::ScheduleTask, +}; use cubecl_common::{ backtrace::BackTrace, bytes::Bytes, @@ -42,15 +44,21 @@ pub struct WgpuStream { /// Used to prevent wgpu staging buffer pool exhaustion during bulk writes /// (e.g. model loading with hundreds of tensors). pending_write_count: usize, + /// GPU access reservations held until the associated commands are submitted. + pending_gpu_reservations: Vec, + /// GPU access reservations held until direct queue writes are submitted. + pending_write_gpu_reservations: Vec, } impl WgpuStream { /// Creates a new WGPU stream. + #[allow(clippy::too_many_arguments)] pub fn new( device: wgpu::Device, queue: wgpu::Queue, memory_properties: MemoryDeviceProperties, memory_config: MemoryConfiguration, + primary_memory: crate::PrimaryMemoryMode, timing_method: TimingMethod, tasks_max: usize, logger: Arc, @@ -71,8 +79,13 @@ impl WgpuStream { let poll = WgpuPoll::new(device.clone()); #[allow(unused_mut)] - let mut mem_manage = - WgpuMemManager::new(device.clone(), memory_properties, memory_config, logger); + let mut mem_manage = WgpuMemManager::new( + device.clone(), + memory_properties, + memory_config, + primary_memory, + logger, + ); Self { mem_manage, @@ -91,6 +104,8 @@ impl WgpuStream { poll, submission_load: SubmissionLoad::default(), pending_write_count: 0, + pending_gpu_reservations: Vec::new(), + pending_write_gpu_reservations: Vec::new(), } } @@ -101,7 +116,11 @@ impl WgpuStream { /// * `task` - The task to execute. pub fn enqueue_task(&mut self, task: ScheduleTask) { match task { - ScheduleTask::Write { data, buffer } => { + ScheduleTask::Write { + data, + buffer, + reservation, + } => { // It is important to flush before writing, as the write operation is inserted // into the QUEUE not the encoder. We want to make sure all outstanding work // happens _before_ the write operation. @@ -111,15 +130,25 @@ impl WgpuStream { flush: false, }) .ok(); - self.write_to_buffer(&buffer, &data); + self.write_to_buffer(&buffer, &data, reservation); } ScheduleTask::Execute { pipeline, count, resources, } => { - let resources = resources.into_resources(self); - self.register_pipeline(pipeline, resources.iter(), &count); + let (resources, reservations) = match resources.into_resources(self) { + Ok(resources) => resources, + Err(err) => { + self.error(err.into()); + return; + } + }; + if let Err(err) = + self.register_pipeline(pipeline, resources.iter(), reservations, &count) + { + self.error(err.into()); + } } } } @@ -143,6 +172,10 @@ impl WgpuStream { let mut callbacks = Vec::with_capacity(descriptors.len()); for (resource, shape, elem_size) in descriptors { + let reservation = match resource.acquire_gpu() { + Ok(reservation) => reservation, + Err(err) => return Box::pin(async move { Err(err.into()) }), + }; let size = shape.iter().product::() * elem_size; // Zero-sized resources don't need a GPU copy. @@ -165,6 +198,7 @@ impl WgpuStream { 0, aligned_len, ); + self.pending_gpu_reservations.push(reservation); staging_info.push(Some((staging, binding, size))); } @@ -375,7 +409,10 @@ impl WgpuStream { pub(crate) fn create_uniform(&mut self, data: &[u8]) -> WgpuResource { let resource = self.mem_manage.reserve_uniform(data.len() as u64); - self.write_to_buffer(&resource, data); + let reservation = resource + .acquire_gpu() + .expect("new uniform allocation cannot be host-mapped"); + self.write_to_buffer(&resource, data, reservation); resource } @@ -383,7 +420,12 @@ impl WgpuStream { // so you have to be really careful about the ordering of operations here. // Any buffer which has outstanding (not yet flushed) compute work should // NOT be copied to. - fn write_to_buffer(&mut self, resource: &WgpuResource, data: &[u8]) { + fn write_to_buffer( + &mut self, + resource: &WgpuResource, + data: &[u8], + reservation: GpuAccessToken, + ) { // Nothing to write for zero-sized resources. if resource.size == 0 { return; @@ -416,6 +458,7 @@ impl WgpuStream { } self.pending_write_count += 1; + self.pending_write_gpu_reservations.push(reservation); // Prevent wgpu staging buffer pool exhaustion during bulk writes (e.g. model // loading with hundreds of tensors). queue.write_buffer() is async — wgpu @@ -436,6 +479,7 @@ impl WgpuStream { label: Some("CubeCL Write Flush Encoder"), }); let index = self.queue.submit([write_flush_encoder.finish()]); + self.pending_write_gpu_reservations.clear(); // Wait for the GPU to finish processing these writes before continuing. #[cfg(not(target_family = "wasm"))] @@ -465,7 +509,7 @@ impl WgpuStream { } pub fn flush(&mut self, mode: StreamErrorMode) -> Result<(), ServerError> { - if self.tasks_count == 0 { + if self.tasks_count == 0 && self.pending_write_count == 0 { return self.flush_errors(mode); } @@ -485,6 +529,8 @@ impl WgpuStream { // This will _first_ fire off all pending write_buffer work. let index = self.queue.submit([tasks_encoder.finish()]); + self.pending_gpu_reservations.clear(); + self.pending_write_gpu_reservations.clear(); self.submission_load .regulate(&self.device, self.tasks_count, index); @@ -525,19 +571,31 @@ impl WgpuStream { &mut self, pipeline: Arc, resources: impl Iterator, + mut reservations: Vec, dispatch: &CubeCount, - ) { + ) -> Result<(), crate::HostAccessError> { if dispatch.is_empty() { - return; + return Ok(()); } + let indirect = match dispatch.clone() { + CubeCount::Dynamic(binding) => { + let resource = self.mem_manage.get_resource(binding).unwrap(); + reservations.push(resource.acquire_gpu()?); + Some(resource) + } + CubeCount::Static(..) => None, + }; + let entries = resources .enumerate() - .map(|(i, r)| wgpu::BindGroupEntry { - binding: i as u32, - resource: r.as_wgpu_bind_resource(), + .map(|(i, r)| { + Ok(wgpu::BindGroupEntry { + binding: i as u32, + resource: r.as_wgpu_bind_resource(), + }) }) - .collect::>(); + .collect::, crate::HostAccessError>>()?; // Start a new compute pass if needed. The forget_lifetime allows // to store this with a 'static lifetime, but the compute pass must @@ -578,12 +636,14 @@ impl WgpuStream { CubeCount::Static(x, y, z) => { pass.dispatch_workgroups(x, y, z); } - CubeCount::Dynamic(binding) => { - let res = self.mem_manage.get_resource(binding).unwrap(); + CubeCount::Dynamic(_) => { + let res = indirect.expect("dynamic dispatch resource was prepared"); pass.dispatch_workgroups_indirect(&res.buffer, res.offset); } } + self.pending_gpu_reservations.extend(reservations); self.flush_if_needed(); + Ok(()) } pub(crate) fn flush_errors_queue(&mut self) -> Vec { diff --git a/crates/cubecl-wgpu/src/lib.rs b/crates/cubecl-wgpu/src/lib.rs index fe1c7bad21..c76e56df02 100644 --- a/crates/cubecl-wgpu/src/lib.rs +++ b/crates/cubecl-wgpu/src/lib.rs @@ -3,9 +3,6 @@ // publish closure. #![allow(unexpected_cfgs)] -#[macro_use] -extern crate derive_new; - extern crate alloc; mod backend; diff --git a/crates/cubecl-wgpu/src/runtime.rs b/crates/cubecl-wgpu/src/runtime.rs index 0e5ae4f2e8..8985851fb1 100644 --- a/crates/cubecl-wgpu/src/runtime.rs +++ b/crates/cubecl-wgpu/src/runtime.rs @@ -25,8 +25,13 @@ pub struct WgpuRuntime; impl DeviceService for WgpuServer { fn init(device_id: cubecl_common::device::DeviceId) -> Self { let device = WgpuDevice::from_id(device_id); - let setup = future::block_on(create_setup_for_device(&device, AutoGraphicsApi::backend())); - create_server(setup, RuntimeOptions::default()) + let options = RuntimeOptions::default(); + let setup = future::block_on(create_setup_for_device( + &device, + AutoGraphicsApi::backend(), + options.primary_memory, + )); + create_server(setup, options) } fn utilities(&self) -> ServerUtilitiesHandle { @@ -174,12 +179,24 @@ fn enumerate_all_adapters(instance: wgpu::Instance, backend: wgpu::Backend) -> V cubecl_common::future::block_on(instance.enumerate_adapters(backend.into())) } +/// Controls where WGPU primary-memory allocations are placed. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum PrimaryMemoryMode { + /// Preserve the standard device-local WGPU allocation behavior. + #[default] + DeviceLocal, + /// Allocate exclusive primary buffers that can be mapped directly by the host. + HostVisible, +} + /// The values that control how a WGPU Runtime will perform its calculations. pub struct RuntimeOptions { /// Control the amount of compute tasks to be aggregated into a single GPU command. pub tasks_max: usize, /// Configures the memory management. pub memory_config: MemoryConfiguration, + /// Configures whether primary allocations are device-local or host-visible. + pub primary_memory: PrimaryMemoryMode, } impl Default for RuntimeOptions { @@ -199,6 +216,7 @@ impl Default for RuntimeOptions { Self { tasks_max, memory_config: MemoryConfiguration::default(), + primary_memory: PrimaryMemoryMode::DeviceLocal, } } } @@ -245,6 +263,38 @@ pub fn init_device(setup: WgpuSetup, options: RuntimeOptions) -> WgpuDevice { device_id } +/// Select a setup for graphics API `G` and register it under a fresh [`WgpuDevice`] ID. +/// +/// Unlike [`init_setup`], this function does not register a client under `selector`, so it can be +/// called repeatedly with the same selector to create independent runtime clients. +/// +/// On WebAssembly, use [`init_device_for_graphics_api_async`] instead. +pub fn init_device_for_graphics_api( + selector: &WgpuDevice, + options: RuntimeOptions, +) -> WgpuDevice { + cfg_if::cfg_if! { + if #[cfg(target_family = "wasm")] { + let _ = (selector, options); + panic!("Creating a wgpu device synchronously is unsupported on wasm. Use init_device_for_graphics_api_async instead"); + } else { + future::block_on(init_device_for_graphics_api_async::(selector, options)) + } + } +} + +/// Async version of [`init_device_for_graphics_api`]. +/// +/// The selector is used only to choose an adapter. The returned [`WgpuDevice::Existing`] ID is +/// globally unique and owns an independently registered runtime client. +pub async fn init_device_for_graphics_api_async( + selector: &WgpuDevice, + options: RuntimeOptions, +) -> WgpuDevice { + let setup = create_setup_for_device(selector, G::backend(), options.primary_memory).await; + init_device(setup, options) +} + /// Like [`init_setup_async`], but synchronous. /// On wasm, it is necessary to use [`init_setup_async`] instead. pub fn init_setup(device: &WgpuDevice, options: RuntimeOptions) -> WgpuSetup { @@ -265,7 +315,7 @@ pub async fn init_setup_async( device: &WgpuDevice, options: RuntimeOptions, ) -> WgpuSetup { - let setup = create_setup_for_device(device, G::backend()).await; + let setup = create_setup_for_device(device, G::backend(), options.primary_memory).await; let return_setup = setup.clone(); let server = create_server(setup, options); let _ = ComputeClient::::init(device, server); @@ -273,6 +323,15 @@ pub async fn init_setup_async( } pub(crate) fn create_server(setup: WgpuSetup, options: RuntimeOptions) -> WgpuServer { + if options.primary_memory == PrimaryMemoryMode::HostVisible + && !setup + .device + .features() + .contains(wgpu::Features::MAPPABLE_PRIMARY_BUFFERS) + { + panic!("{}", backend::HOST_VISIBLE_PRIMARY_UNSUPPORTED); + } + let limits = setup.device.limits(); let adapter_limits = setup.adapter.limits(); let mut adapter_info = setup.adapter.get_info(); @@ -377,6 +436,7 @@ pub(crate) fn create_server(setup: WgpuSetup, options: RuntimeOptions) -> WgpuSe setup.device.clone(), setup.queue, options.tasks_max, + options.primary_memory, setup.backend, time_measurement, ServerUtilities::new(device_props, logger, setup.backend, allocator), @@ -388,9 +448,10 @@ pub(crate) fn create_server(setup: WgpuSetup, options: RuntimeOptions) -> WgpuSe pub(crate) async fn create_setup_for_device( device: &WgpuDevice, backend: wgpu::Backend, + primary_memory: PrimaryMemoryMode, ) -> WgpuSetup { let (instance, adapter) = request_adapter(device, backend).await; - let (device, queue) = backend::request_device(&adapter).await; + let (device, queue) = backend::request_device(&adapter, primary_memory).await; log::info!( "Created wgpu compute server on device {:?} => {:?}", @@ -605,3 +666,16 @@ fn get_device_override() -> Option { override_device }) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn runtime_options_default_to_device_local_primary_memory() { + assert_eq!( + RuntimeOptions::default().primary_memory, + PrimaryMemoryMode::DeviceLocal + ); + } +} diff --git a/crates/cubecl-wgpu/tests/host_visible_primary.rs b/crates/cubecl-wgpu/tests/host_visible_primary.rs new file mode 100644 index 0000000000..4dcb3dfc97 --- /dev/null +++ b/crates/cubecl-wgpu/tests/host_visible_primary.rs @@ -0,0 +1,162 @@ +#![cfg(target_os = "macos")] + +use cubecl::prelude::*; +use cubecl_core as cubecl; +use cubecl_runtime::server::ServerError; +use cubecl_wgpu::{ + HostAccessError, MemoryConfiguration, Metal, PrimaryMemoryMode, RuntimeOptions, WgpuDevice, + WgpuRuntime, init_device_for_graphics_api, init_setup, +}; + +#[cube(launch)] +fn increment(values: &mut Array) { + if ABSOLUTE_POS < values.len() { + values[ABSOLUTE_POS] += 1; + } +} + +#[test] +fn compute_client_handle_resolves_to_guarded_host_visible_resource() { + let instance = wgpu::Instance::new(wgpu::InstanceDescriptor { + backends: wgpu::Backends::METAL, + ..wgpu::InstanceDescriptor::new_without_display_handle() + }); + let adapters = + cubecl_common::future::block_on(instance.enumerate_adapters(wgpu::Backends::METAL)); + if adapters.is_empty() { + eprintln!("skipping host-visible primary test: no Metal adapter is available"); + return; + } + + let device = WgpuDevice::DefaultDevice; + let _setup = init_setup::( + &device, + RuntimeOptions { + tasks_max: 32, + memory_config: MemoryConfiguration::ExclusivePages, + primary_memory: PrimaryMemoryMode::HostVisible, + }, + ); + let client = WgpuRuntime::client(&device); + let handle = client.empty(4 * core::mem::size_of::()); + + // This is the public downstream path used by tenferro: Handle -> ComputeClient -> + // ManagedResource -> &WgpuResource. + let managed = client.get_resource(handle.clone()).unwrap(); + let resource = managed.resource(); + let allocation_id = resource.allocation_id(); + let cloned_resource = resource.clone(); + assert_ne!(allocation_id, 0); + assert_eq!(cloned_resource.allocation_id(), allocation_id); + + { + let mut guard = resource.map_write().unwrap(); + guard.copy_from_slice(u32::as_bytes(&[1, 2, 3, 4])); + + assert_eq!( + resource.map_read().unwrap_err(), + HostAccessError::OverlappingHostMapping + ); + + assert_eq!( + resource.ensure_gpu_access(), + Err(HostAccessError::MappedForHost) + ); + + increment::launch( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(4), + unsafe { ArrayArg::from_raw_parts(handle.clone(), 4) }, + ); + let error = client.flush().unwrap_err(); + assert!(matches!(error, ServerError::ServerUnhealthy { .. })); + assert!(format!("{error:?}").contains("mapped for host access")); + } + + increment::launch( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(4), + unsafe { ArrayArg::from_raw_parts(handle.clone(), 4) }, + ); + // Barrier on the server channel: binding preparation and scheduler registration have completed, + // but the queued encoder has not been submitted yet. + client.exclusive(|| ()).unwrap(); + let pending_resource = resource.clone(); + let host_attempt = std::thread::spawn(move || pending_resource.map_read().unwrap_err()); + assert_eq!( + host_attempt.join().unwrap(), + HostAccessError::GpuAccessInProgress + ); + client.flush().unwrap(); + + let guard = resource.map_read().unwrap(); + assert_eq!(u32::from_bytes(&guard), &[2, 3, 4, 5]); + assert_eq!(resource.allocation_id(), allocation_id); + drop(guard); + drop(managed); + drop(handle); + + // A cloned resource can escape the ManagedResource borrow. It must retain the managed-memory + // binding so cleanup cannot recycle its slice for a new allocation while the clone is alive. + let replacement_handle = client.empty(4 * core::mem::size_of::()); + let replacement_managed = client.get_resource(replacement_handle).unwrap(); + let replacement = replacement_managed.resource(); + assert_ne!(replacement.allocation_id(), cloned_resource.allocation_id()); + + { + let mut replacement_guard = replacement.map_write().unwrap(); + replacement_guard.copy_from_slice(u32::as_bytes(&[10, 20, 30, 40])); + } + let stale_guard = cloned_resource.map_read().unwrap(); + drop(cloned_resource); + + // The mapped guard is now the only remaining lease for the original allocation. + let guard_replacement_handle = client.empty(4 * core::mem::size_of::()); + let guard_replacement_managed = client.get_resource(guard_replacement_handle).unwrap(); + assert_ne!( + guard_replacement_managed.resource().allocation_id(), + allocation_id + ); + assert_eq!(u32::from_bytes(&stale_guard), &[2, 3, 4, 5]); +} + +#[test] +fn independent_host_visible_metal_devices_use_fresh_context_ids() { + let instance = wgpu::Instance::new(wgpu::InstanceDescriptor { + backends: wgpu::Backends::METAL, + ..wgpu::InstanceDescriptor::new_without_display_handle() + }); + let adapters = + cubecl_common::future::block_on(instance.enumerate_adapters(wgpu::Backends::METAL)); + if adapters.is_empty() { + eprintln!("skipping host-visible primary test: no Metal adapter is available"); + return; + } + + let options = || RuntimeOptions { + tasks_max: 32, + memory_config: MemoryConfiguration::ExclusivePages, + primary_memory: PrimaryMemoryMode::HostVisible, + }; + let selector = WgpuDevice::DefaultDevice; + let first = init_device_for_graphics_api::(&selector, options()); + let second = init_device_for_graphics_api::(&selector, options()); + + assert_ne!(first, second); + assert!(matches!(first, WgpuDevice::Existing(_))); + assert!(matches!(second, WgpuDevice::Existing(_))); + + let first_client = WgpuRuntime::client(&first); + let second_client = WgpuRuntime::client(&second); + let first_handle = first_client.empty(core::mem::size_of::()); + let second_handle = second_client.empty(core::mem::size_of::()); + let first_resource = first_client.get_resource(first_handle).unwrap(); + let second_resource = second_client.get_resource(second_handle).unwrap(); + + assert_ne!( + first_resource.resource().allocation_id(), + second_resource.resource().allocation_id() + ); +} diff --git a/docs/superpowers/plans/2026-07-19-fresh-graphics-api-device.md b/docs/superpowers/plans/2026-07-19-fresh-graphics-api-device.md new file mode 100644 index 0000000000..d2df21d5b1 --- /dev/null +++ b/docs/superpowers/plans/2026-07-19-fresh-graphics-api-device.md @@ -0,0 +1,66 @@ +# Fresh Graphics-API Device Initialization Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add sync and async WGPU initialization helpers that create independently registered clients under fresh device IDs for a selected graphics API. + +**Architecture:** Select the adapter and create `WgpuSetup` without registering the selector, then reuse `init_device` for fresh-ID generation and client registration. Preserve all existing initialization entry points. + +**Tech Stack:** Rust, CubeCL WGPU runtime, wgpu Metal backend, cargo test. + +## Global Constraints + +- Preserve the behavior and signatures of `init_setup`, `init_setup_async`, and `init_device`. +- Return a fresh `WgpuDevice::Existing` ID on every successful call. +- Support `PrimaryMemoryMode::HostVisible` and both synchronous native and asynchronous initialization. +- Keep the fork minimal and upstream-close. + +--- + +### Task 1: Fresh Graphics-API Initialization + +**Files:** +- Modify: `crates/cubecl-wgpu/src/runtime.rs` +- Test: `crates/cubecl-wgpu/tests/host_visible_primary.rs` + +**Interfaces:** +- Consumes: `create_setup_for_device`, `init_device`, `GraphicsApi`, `RuntimeOptions`, `WgpuDevice`. +- Produces: `init_device_for_graphics_api(&WgpuDevice, RuntimeOptions) -> WgpuDevice` and `init_device_for_graphics_api_async(&WgpuDevice, RuntimeOptions) -> WgpuDevice`. + +- [ ] **Step 1: Write the failing integration test** + +Add a macOS Metal test that calls `init_device_for_graphics_api::` twice with host-visible options, asserts the IDs differ, resolves `WgpuRuntime::client` for both IDs, allocates one handle per client, and resolves both through `ComputeClient::get_resource`. + +- [ ] **Step 2: Run the test to verify it fails** + +Run: `cargo test -p t4a-cubecl-wgpu --test host_visible_primary independent_host_visible_metal_devices_use_fresh_context_ids --features msl -- --exact --nocapture` + +Expected: compilation fails because `init_device_for_graphics_api` does not exist. + +- [ ] **Step 3: Implement the async helper** + +Add: + +```rust +pub async fn init_device_for_graphics_api_async( + selector: &WgpuDevice, + options: RuntimeOptions, +) -> WgpuDevice { + let setup = create_setup_for_device(selector, G::backend(), options.primary_memory).await; + init_device(setup, options) +} +``` + +- [ ] **Step 4: Implement the synchronous wrapper** + +Add a native wrapper using `future::block_on(init_device_for_graphics_api_async::(selector, options))`; on WebAssembly, panic with guidance to use the async function, matching `init_setup`. + +- [ ] **Step 5: Run focused and API validation** + +Run the focused integration command from Step 2, then `cargo fmt --all -- --check`, `cargo clippy -p t4a-cubecl-wgpu --all-targets --features std,msl -- -D warnings`, and `RUSTDOCFLAGS='-D warnings' cargo doc -p t4a-cubecl-wgpu --features std,msl --no-deps`. + +Expected: all commands pass; on a host without Metal the integration test reports its existing skip path. + +- [ ] **Step 6: Commit with the related review fixes** + +Stage the runtime and integration changes together with the guarded host-visible primary-memory review fixes and commit with an intentional fix message. diff --git a/docs/superpowers/plans/2026-07-19-host-visible-primary.md b/docs/superpowers/plans/2026-07-19-host-visible-primary.md new file mode 100644 index 0000000000..36ede370e8 --- /dev/null +++ b/docs/superpowers/plans/2026-07-19-host-visible-primary.md @@ -0,0 +1,93 @@ +# Host-Visible WGPU Primary Memory Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add an opt-in WGPU primary-memory mode whose whole-resource allocations can be mapped by the host and used by Metal-backed WGPU without copying tensor bytes. + +**Architecture:** Keep the default device-local WGPU path unchanged. An opt-in `PrimaryMemoryMode::HostVisible` requests `MAPPABLE_PRIMARY_BUFFERS`, forces the main pool to `ExclusivePages`, allocates main buffers with `MAP_READ | MAP_WRITE` in addition to their GPU usages, and returns guarded mapped ranges from `WgpuResource`. Every allocation carries shared access state so a CPU mapping and GPU binding are mutually exclusive; dropping a mapping unmaps the resource and restores GPU access. + +**Tech Stack:** Rust, CubeCL 0.10 fork, wgpu 26, Metal-backed WGPU, Cargo tests. + +## Global Constraints + +- Default behavior and non-Metal adapters remain unchanged. +- Host-visible primary memory is explicit; no runtime fallback or hidden copy is allowed. +- Host-visible allocations use whole-resource pooling (`MemoryConfiguration::ExclusivePages`) and stable shared allocation identity. +- A host mapping waits for submitted GPU work through wgpu mapping completion. +- GPU binding while a host guard is live is rejected with a typed error; overlapping host guards are rejected with a typed error. +- Dropping a host guard always unmaps and makes the resource GPU-usable again. +- Do not publish or release crates in this task; a git revision is sufficient. +- Keep fork-only changes small and compatible with the upstream 0.10 API shape. + +--- + +### Task 1: Opt-in host-visible WGPU allocations and guarded mapping + +**Files:** +- Modify: `crates/cubecl-wgpu/src/runtime.rs` +- Modify: `crates/cubecl-wgpu/src/backend/metal.rs` +- Modify: `crates/cubecl-wgpu/src/backend/base.rs` +- Modify: `crates/cubecl-wgpu/src/compute/mem_manager.rs` +- Modify: `crates/cubecl-wgpu/src/compute/storage.rs` +- Modify: `crates/cubecl-wgpu/src/compute/server.rs` +- Modify: `crates/cubecl-wgpu/src/compute/schedule.rs` +- Test: `crates/cubecl-wgpu/src/compute/storage.rs` +- Test: `crates/cubecl-wgpu/tests/host_visible_primary.rs` + +**Interfaces:** +- Produces: public `PrimaryMemoryMode::{DeviceLocal, HostVisible}`. +- Produces: `RuntimeOptions { tasks_max, memory_config, primary_memory }`, with `DeviceLocal` as the default. +- Produces: `WgpuResource::allocation_id() -> u64` returning an identity shared by clones/views of one physical allocation. +- Produces: `WgpuResource::map_read() -> Result` and `WgpuResource::map_write() -> Result`. +- Produces: read and write guards implementing `Deref`; the write guard also implements `DerefMut`. +- Produces: public `HostAccessError` variants that distinguish a device-local allocation, an overlapping host mapping, a mapped resource requested for GPU use, a map callback failure, and device polling failure. +- Consumes later: tenferro obtains a `WgpuResource` through `ComputeClient::get_resource`, maps it synchronously, and records `allocation_id()` for zero-copy assertions. + +- [ ] **Step 1: Add failing unit tests for allocation identity and access-state transitions** + + Add tests around the access-state object used by `WgpuResource`. Assert that resource clones retain the same nonzero allocation ID, one host access excludes another, GPU access is rejected while a host token exists, and dropping the token restores GPU access. Keep these tests independent of physical GPU availability. + +- [ ] **Step 2: Run the focused unit test and confirm it fails** + + Run: `cargo test -p t4a-cubecl-wgpu compute::storage::tests --features std` + + Expected: compilation fails because the identity/access-state API does not exist. + +- [ ] **Step 3: Implement stable identity, typed access errors, and RAII mapped guards** + + Store each physical buffer in a cloneable allocation object containing the `wgpu::Buffer`, `wgpu::Device`, a process-unique nonzero `u64` allocation ID, a host-visible flag, and an atomic access state. `WgpuStorage::get` must create offset/size views that share this object. Mapping must compare-and-swap GPU-idle to host-mapped before calling `map_async`; every error path must restore GPU-idle. After callback completion, poll with `wgpu::PollType::Wait`, obtain the exact logical `[offset, offset + size)` mapped range, and return an owning guard. Guard drop must release the mapped view before calling `unmap`, then restore GPU-idle. The GPU binding path must call a fallible `ensure_gpu_access()` before producing a binding resource. + +- [ ] **Step 4: Add opt-in runtime configuration and Metal feature negotiation** + + Add `PrimaryMemoryMode` to `RuntimeOptions`, defaulting to `DeviceLocal`. Thread the choice through WGPU device/setup creation and `WgpuServer`/`WgpuMemManager`. When `HostVisible` is selected, require `wgpu::Features::MAPPABLE_PRIMARY_BUFFERS`; return or panic with a precise unsupported-feature message at setup rather than silently falling back. The Metal backend must request the feature only in this mode. Other backends must either request it when supported or reject setup explicitly. + +- [ ] **Step 5: Allocate host-visible main buffers as exclusive whole resources** + + In host-visible mode, force only the main memory pool to `MemoryConfiguration::ExclusivePages` and add `MAP_READ | MAP_WRITE` to its existing `STORAGE | COPY_SRC | COPY_DST | INDIRECT` usages. Mark those resources host-visible. Staging and uniform pools remain unchanged. Default mode must retain its caller-provided memory configuration and usages. + +- [ ] **Step 6: Reject GPU scheduling while a host guard is live** + + Make WGPU binding preparation validate every `WgpuResource` with `ensure_gpu_access()`. Convert `HostAccessError::MappedForHost` into the existing asynchronous server error path without panicking. Cover direct kernel buffers, dynamic dispatch buffers, writes, and other scheduled tasks that use primary resources. + +- [ ] **Step 7: Add a macOS Metal integration test** + + Under `#[cfg(target_os = "macos")]`, initialize a Metal WGPU device with host-visible primary memory and `ExclusivePages`. Create one `u32` buffer, map-write values, drop the guard, launch a kernel that increments them, map-read the result, and assert values and an unchanged allocation ID. While a write guard is live, attempt GPU use and assert a typed/server error; after dropping it, assert the same resource remains usable. Skip with a clear message only when no Metal adapter is available. + +- [ ] **Step 8: Verify formatting, focused tests, and the WGPU crate** + + Run: + + ```bash + cargo fmt --all -- --check + cargo test -p t4a-cubecl-wgpu --features std + cargo clippy -p t4a-cubecl-wgpu --all-targets --features std -- -D warnings + ``` + + Expected: all commands pass, including the macOS host-visible transition test on an available Metal adapter. + +- [ ] **Step 9: Commit** + + ```bash + git add crates/cubecl-wgpu docs/superpowers/plans/2026-07-19-host-visible-primary.md + git commit -m "feat(wgpu): add guarded host-visible primary memory" + ``` diff --git a/docs/superpowers/specs/2026-07-19-fresh-graphics-api-device-design.md b/docs/superpowers/specs/2026-07-19-fresh-graphics-api-device-design.md new file mode 100644 index 0000000000..6e5ce7bd0a --- /dev/null +++ b/docs/superpowers/specs/2026-07-19-fresh-graphics-api-device-design.md @@ -0,0 +1,37 @@ +# Fresh Graphics-API Device Initialization Design + +## Goal + +Allow callers to select a WGPU graphics API repeatedly without registering the selector device globally, returning a fresh `WgpuDevice::Existing` ID for each independent runtime client. + +## API + +Add these public functions beside the existing initialization helpers: + +```rust +pub fn init_device_for_graphics_api( + selector: &WgpuDevice, + options: RuntimeOptions, +) -> WgpuDevice; + +pub async fn init_device_for_graphics_api_async( + selector: &WgpuDevice, + options: RuntimeOptions, +) -> WgpuDevice; +``` + +The synchronous function is unavailable on WebAssembly in the same way as `init_setup`; callers there use the async function. + +## Data Flow + +The async function calls `create_setup_for_device(selector, G::backend(), options.primary_memory)` directly, so the selector is used only for adapter selection. It then passes the setup and options to the existing `init_device`, which generates a globally unique `WgpuDevice::Existing` ID, creates the server, and registers the client under that ID. The synchronous function blocks on the async function outside WebAssembly. + +Existing `init_setup`, `init_setup_async`, and `init_device` behavior remains unchanged. + +## Errors and Compatibility + +Adapter/device setup keeps the existing panic behavior because the lower-level setup helpers are unchanged. Host-visible feature validation remains in `create_server`. No existing public signature or default changes. + +## Testing + +On macOS with an available Metal adapter, initialize two host-visible Metal devices through the new API. Assert the returned IDs are distinct, obtain a client for each ID, allocate resources independently, and verify both resource paths resolve. The existing host-visible integration test remains the downstream-path and mapping behavior test.