tenferro_gpu/webgpu/apple.rs
1use std::fmt;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::sync::Arc;
4
5use cubecl::client::ComputeClient;
6use cubecl::Runtime;
7use cubecl_wgpu::{
8 init_device_for_graphics_api, MemoryConfiguration, Metal, PrimaryMemoryMode, RuntimeOptions,
9 WgpuDevice, WgpuRuntime as CubeWgpuRuntime,
10};
11
12use super::{
13 alloc_tensor_in_runtime, download_webgpu_tensor, upload_webgpu_tensor, WebGpuBackend,
14 WebGpuRuntime,
15};
16use tenferro_cpu::CpuBackend;
17use tenferro_tensor::{AllocationDomainId, DType, SharedTensorAllocationDomain, Tensor};
18
19/// Explicit host/device transfer counters for one [`AppleContext`].
20///
21/// Guarded host mappings do not change these counters.
22///
23/// # Examples
24///
25/// ```rust
26/// use tenferro_gpu::apple::AppleTransferStats;
27///
28/// assert_eq!(AppleTransferStats::default().uploaded_bytes, 0);
29/// ```
30#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
31pub struct AppleTransferStats {
32 /// Bytes copied by explicit host-to-managed creation.
33 pub uploaded_bytes: usize,
34 /// Bytes copied by explicit managed-to-host download.
35 pub downloaded_bytes: usize,
36}
37
38pub(super) struct AppleDomainState {
39 pub(super) id: AllocationDomainId,
40 client: ComputeClient<CubeWgpuRuntime>,
41 device_ordinal: usize,
42 uploaded_bytes: AtomicUsize,
43 downloaded_bytes: AtomicUsize,
44}
45
46impl fmt::Debug for AppleDomainState {
47 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
48 formatter
49 .debug_struct("AppleDomainState")
50 .field("id", &self.id)
51 .field("device_ordinal", &self.device_ordinal)
52 .field("transfers", &self.snapshot())
53 .finish_non_exhaustive()
54 }
55}
56
57impl AppleDomainState {
58 pub(super) fn record_upload(&self, bytes: usize) {
59 self.uploaded_bytes.fetch_add(bytes, Ordering::Relaxed);
60 }
61
62 pub(super) fn record_download(&self, bytes: usize) {
63 self.downloaded_bytes.fetch_add(bytes, Ordering::Relaxed);
64 }
65
66 fn snapshot(&self) -> AppleTransferStats {
67 AppleTransferStats {
68 uploaded_bytes: self.uploaded_bytes.load(Ordering::Relaxed),
69 downloaded_bytes: self.downloaded_bytes.load(Ordering::Relaxed),
70 }
71 }
72
73 fn runtime(self: &Arc<Self>) -> WebGpuRuntime {
74 WebGpuRuntime::from_apple_client(self.client.clone(), self.device_ordinal, Arc::clone(self))
75 }
76}
77
78#[derive(Debug)]
79struct AppleAllocationDomain {
80 state: Arc<AppleDomainState>,
81}
82
83impl SharedTensorAllocationDomain for AppleAllocationDomain {
84 fn id(&self) -> AllocationDomainId {
85 self.state.id
86 }
87
88 fn allocate(&self, dtype: DType, shape: &[usize]) -> crate::Result<Tensor> {
89 alloc_tensor_in_runtime(&self.state.runtime(), dtype, shape)
90 }
91}
92
93/// Explicit Apple shared-allocation context for CPU and Metal backends.
94///
95/// Every context owns a fresh host-visible Metal client and allocation domain.
96/// Tensors created through [`Self::upload_tensor`] can be mapped by the paired
97/// CPU backend and launched by the paired Metal backend without implicit
98/// transfers. The initial mapped CPU operations are RustFFT and rank-2
99/// Cholesky; this context does not turn other CPU operations into automatic
100/// fallbacks. Clone [`Self::cpu_backend`] or [`Self::metal_backend`] to obtain
101/// the mutable handle required by an explicitly selected operation.
102///
103/// # Examples
104///
105/// ```rust
106/// use tenferro_gpu::apple::{AppleContext, AppleTransferStats};
107///
108/// match AppleContext::new() {
109/// Ok(context) => {
110/// assert_eq!(
111/// context.cpu_backend().allocation_domain(),
112/// Some(context.domain_id())
113/// );
114/// assert_eq!(context.transfer_stats(), AppleTransferStats::default());
115/// }
116/// Err(error) => assert!(matches!(
117/// error,
118/// tenferro_tensor::Error::RuntimeState { .. }
119/// | tenferro_tensor::Error::BackendSource { .. }
120/// )),
121/// }
122/// ```
123#[derive(Clone)]
124pub struct AppleContext {
125 state: Arc<AppleDomainState>,
126 cpu: CpuBackend,
127 metal: WebGpuBackend,
128}
129
130impl fmt::Debug for AppleContext {
131 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
132 formatter
133 .debug_struct("AppleContext")
134 .field("domain_id", &self.domain_id())
135 .field("transfers", &self.transfer_stats())
136 .finish_non_exhaustive()
137 }
138}
139
140impl AppleContext {
141 /// Create an independent host-visible Metal context and paired CPU backend.
142 ///
143 /// # Examples
144 ///
145 /// ```rust
146 /// use tenferro_gpu::apple::{AppleContext, AppleTransferStats};
147 ///
148 /// match AppleContext::new() {
149 /// Ok(context) => {
150 /// assert_eq!(
151 /// context.cpu_backend().allocation_domain(),
152 /// Some(context.domain_id())
153 /// );
154 /// assert_eq!(context.transfer_stats(), AppleTransferStats::default());
155 /// }
156 /// Err(error) => assert!(matches!(
157 /// error,
158 /// tenferro_tensor::Error::RuntimeState { .. }
159 /// | tenferro_tensor::Error::BackendSource { .. }
160 /// )),
161 /// }
162 /// ```
163 ///
164 /// # Errors
165 ///
166 /// Returns a typed runtime or backend error when Metal or host-visible
167 /// primary buffers are unavailable.
168 pub fn new() -> crate::Result<Self> {
169 static NEXT_ORDINAL: AtomicUsize = AtomicUsize::new(1_000_000);
170 let options = RuntimeOptions {
171 tasks_max: 32,
172 memory_config: MemoryConfiguration::ExclusivePages,
173 primary_memory: PrimaryMemoryMode::HostVisible,
174 };
175 let device = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
176 init_device_for_graphics_api::<Metal>(&WgpuDevice::DefaultDevice, options)
177 }))
178 .map_err(|payload| {
179 crate::Error::runtime_state(
180 "AppleContext::new",
181 format!("failed to initialize host-visible Metal runtime: {payload:?}"),
182 )
183 })?;
184 let client = CubeWgpuRuntime::client(&device);
185 let device_ordinal = NEXT_ORDINAL.fetch_add(1, Ordering::Relaxed);
186 let state = Arc::new(AppleDomainState {
187 id: AllocationDomainId::fresh(),
188 client,
189 device_ordinal,
190 uploaded_bytes: AtomicUsize::new(0),
191 downloaded_bytes: AtomicUsize::new(0),
192 });
193 let domain: Arc<dyn SharedTensorAllocationDomain> = Arc::new(AppleAllocationDomain {
194 state: Arc::clone(&state),
195 });
196 let cpu = CpuBackend::new().with_allocation_domain(domain);
197 let metal = WebGpuBackend::from_runtime(state.runtime());
198 Ok(Self { state, cpu, metal })
199 }
200
201 /// Return this context's allocation-domain identity.
202 ///
203 /// # Examples
204 ///
205 /// ```rust
206 /// use tenferro_gpu::apple::AppleContext;
207 ///
208 /// match AppleContext::new() {
209 /// Ok(context) => assert_eq!(
210 /// context.cpu_backend().allocation_domain(),
211 /// Some(context.domain_id())
212 /// ),
213 /// Err(error) => assert!(matches!(
214 /// error,
215 /// tenferro_tensor::Error::RuntimeState { .. }
216 /// | tenferro_tensor::Error::BackendSource { .. }
217 /// )),
218 /// }
219 /// ```
220 pub fn domain_id(&self) -> AllocationDomainId {
221 self.state.id
222 }
223
224 /// Return the paired CPU backend.
225 ///
226 /// # Examples
227 ///
228 /// ```rust
229 /// use tenferro_gpu::apple::AppleContext;
230 ///
231 /// match AppleContext::new() {
232 /// Ok(context) => assert_eq!(
233 /// context.cpu_backend().allocation_domain(),
234 /// Some(context.domain_id())
235 /// ),
236 /// Err(error) => assert!(matches!(
237 /// error,
238 /// tenferro_tensor::Error::RuntimeState { .. }
239 /// | tenferro_tensor::Error::BackendSource { .. }
240 /// )),
241 /// }
242 /// ```
243 pub fn cpu_backend(&self) -> &CpuBackend {
244 &self.cpu
245 }
246
247 /// Return the paired Metal WebGPU backend.
248 ///
249 /// # Examples
250 ///
251 /// ```rust
252 /// use tenferro_gpu::apple::AppleContext;
253 ///
254 /// match AppleContext::new() {
255 /// Ok(context) => {
256 /// let backend = context.metal_backend();
257 /// assert_eq!(
258 /// backend.runtime_identity(),
259 /// backend.runtime().runtime_identity()
260 /// );
261 /// }
262 /// Err(error) => assert!(matches!(
263 /// error,
264 /// tenferro_tensor::Error::RuntimeState { .. }
265 /// | tenferro_tensor::Error::BackendSource { .. }
266 /// )),
267 /// }
268 /// ```
269 pub fn metal_backend(&self) -> &WebGpuBackend {
270 &self.metal
271 }
272
273 /// Return an atomic snapshot of explicit transfer counters.
274 ///
275 /// # Examples
276 ///
277 /// ```rust
278 /// use tenferro_gpu::apple::{AppleContext, AppleTransferStats};
279 ///
280 /// match AppleContext::new() {
281 /// Ok(context) => assert_eq!(
282 /// context.transfer_stats(),
283 /// AppleTransferStats::default()
284 /// ),
285 /// Err(error) => assert!(matches!(
286 /// error,
287 /// tenferro_tensor::Error::RuntimeState { .. }
288 /// | tenferro_tensor::Error::BackendSource { .. }
289 /// )),
290 /// }
291 /// ```
292 pub fn transfer_stats(&self) -> AppleTransferStats {
293 self.state.snapshot()
294 }
295
296 /// Create a managed tensor by explicitly copying a host tensor once.
297 ///
298 /// # Examples
299 ///
300 /// ```rust
301 /// use tenferro_gpu::apple::{AppleContext, AppleTransferStats};
302 /// use tenferro_tensor::{Tensor, TypedTensor};
303 ///
304 /// match AppleContext::new() {
305 /// Ok(context) => {
306 /// let host = Tensor::from_vec_col_major([2], vec![1.0_f32, 2.0])?;
307 /// let managed: TypedTensor<f32> = match context.upload_tensor(&host)? {
308 /// Tensor::F32(typed) => typed,
309 /// _ => unreachable!("expected f32 tensor"),
310 /// };
311 /// assert_eq!(managed.allocation_domain(), Some(context.domain_id()));
312 /// assert_eq!(managed.with_host_read(|data| data.to_vec())?, [1.0, 2.0]);
313 /// assert_eq!(
314 /// context.transfer_stats(),
315 /// AppleTransferStats {
316 /// uploaded_bytes: 8,
317 /// downloaded_bytes: 0,
318 /// }
319 /// );
320 /// }
321 /// Err(error) => assert!(matches!(
322 /// error,
323 /// tenferro_tensor::Error::RuntimeState { .. }
324 /// | tenferro_tensor::Error::BackendSource { .. }
325 /// )),
326 /// }
327 /// # Ok::<(), tenferro_tensor::Error>(())
328 /// ```
329 ///
330 /// # Errors
331 ///
332 /// Returns [`crate::Error::RuntimeState`] when the input is not host
333 /// resident, [`crate::Error::Validation`] for invalid tensor metadata, or
334 /// [`crate::Error::BackendSource`] when managed allocation or queue
335 /// synchronization fails.
336 pub fn upload_tensor(&self, tensor: &Tensor) -> crate::Result<Tensor> {
337 let output = upload_webgpu_tensor(self.metal.runtime(), tensor)?;
338 self.metal.synchronize()?;
339 Ok(output)
340 }
341
342 /// Explicitly copy a matching managed tensor back to host memory.
343 ///
344 /// # Examples
345 ///
346 /// ```rust
347 /// use tenferro_gpu::apple::{AppleContext, AppleTransferStats};
348 /// use tenferro_tensor::Tensor;
349 ///
350 /// match AppleContext::new() {
351 /// Ok(context) => {
352 /// let host = Tensor::from_vec_col_major([2], vec![1.0_f32, 2.0])?;
353 /// let managed = context.upload_tensor(&host)?;
354 /// let downloaded = context.download_tensor(&managed)?;
355 /// assert_eq!(downloaded.as_slice::<f32>()?, &[1.0, 2.0]);
356 /// assert_eq!(
357 /// context.transfer_stats(),
358 /// AppleTransferStats {
359 /// uploaded_bytes: 8,
360 /// downloaded_bytes: 8,
361 /// }
362 /// );
363 /// }
364 /// Err(error) => assert!(matches!(
365 /// error,
366 /// tenferro_tensor::Error::RuntimeState { .. }
367 /// | tenferro_tensor::Error::BackendSource { .. }
368 /// )),
369 /// }
370 /// # Ok::<(), tenferro_tensor::Error>(())
371 /// ```
372 ///
373 /// # Errors
374 ///
375 /// Returns [`crate::Error::HostAccess`] for a foreign allocation domain or
376 /// a typed backend error when readback fails.
377 pub fn download_tensor(&self, tensor: &Tensor) -> crate::Result<Tensor> {
378 download_webgpu_tensor(self.metal.runtime(), tensor)
379 }
380}