tenferro_gpu/webgpu/
mod.rs

1//! CubeCL WebGPU provider runtime and backend skeleton.
2
3use cubecl::prelude::{CubeCount, CubeDim, CubeElement, CubeType, Sequence, TensorBinding};
4use cubecl_wgpu::WgpuRuntime;
5use std::fmt;
6use std::sync::Arc;
7
8use crate::{
9    AccessError, AllocationDomainId, AllocationId, AllocationKey, BackendAllocation,
10    BackendCachedDot, BackendId, BackendRuntimeCache, BackendSession, CompareDir, DType,
11    DeviceAccessError, DeviceAccessRequest, DeviceId, DeviceKind, DotGeneralConfig, Error,
12    GatherConfig, GpuBackendKind, HostAccessError, MemoryKind, PadConfig, Placement,
13    PreparedDeviceAccess, ProviderCapabilities, ProviderReadMapping, ProviderWriteMapping,
14    RootBoundSpan, RootResourceExtent, ScatterConfig, SliceConfig, Tensor, TensorAnalytic,
15    TensorBackend, TensorBuffer, TensorDeviceTransfer, TensorDot, TensorElementwise, TensorFusion,
16    TensorIndexing, TensorRank, TensorRead, TensorReduction, TensorScalar, TensorStructural,
17    TensorViewCanonicalization, TensorWrite, TypedTensor, TypedTensorView, TypedTensorViewMut,
18};
19
20const DEFAULT_CUBE_DIM_X: u32 = 256;
21
22mod apple;
23mod error;
24#[cfg(not(target_family = "wasm"))]
25mod event_domain;
26mod exec_session;
27mod gemm;
28#[doc(hidden)]
29pub mod interop;
30mod kernels;
31mod memory;
32mod runtime;
33mod runtime_adapter;
34mod structural;
35
36pub use apple::{AppleContext, AppleTransferStats};
37pub(crate) use error::{unsupported_dtype, unsupported_operation};
38#[doc(hidden)]
39pub use exec_session::{with_webgpu_exec_session, WebGpuExecSession};
40pub use memory::{download_webgpu_tensor, upload_webgpu_tensor};
41pub use runtime::{webgpu_available, WebGpuRuntime, WebGpuRuntimeIdentity};
42pub use runtime_adapter::{
43    webgpu_runtime_engine_id, webgpu_runtime_engine_registration,
44    webgpu_runtime_engine_registration_with_id, webgpu_runtime_hardware_class,
45};
46
47/// Scalar-independent WebGPU allocation stored behind tensor backend-buffer
48/// trait objects; dtype is carried by the borrowed tensor descriptor.
49pub(crate) struct WebGpuBuffer {
50    handle: cubecl_runtime::server::Handle,
51    byte_len: usize,
52    device_ordinal: usize,
53    managed: Option<Arc<cubecl_runtime::storage::ManagedResource<cubecl_wgpu::WgpuResource>>>,
54    allocation_domain: AllocationDomainId,
55    allocation_id: AllocationId,
56}
57
58static NEXT_WEBGPU_ALLOCATION_ID: std::sync::atomic::AtomicU64 =
59    std::sync::atomic::AtomicU64::new(1);
60
61impl std::fmt::Debug for WebGpuBuffer {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        f.debug_struct("WebGpuBuffer")
64            .field("byte_len", &self.byte_len)
65            .field("device_ordinal", &self.device_ordinal)
66            .field("allocation_domain", &self.allocation_domain)
67            .field("allocation_id", &self.allocation_id)
68            .finish()
69    }
70}
71
72impl WebGpuBuffer {
73    fn new(
74        handle: cubecl_runtime::server::Handle,
75        byte_len: usize,
76        device_ordinal: usize,
77        allocation_domain: AllocationDomainId,
78    ) -> Self {
79        Self {
80            handle,
81            byte_len,
82            device_ordinal,
83            managed: None,
84            allocation_domain,
85            allocation_id: AllocationId::from_backend_id(
86                NEXT_WEBGPU_ALLOCATION_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
87            ),
88        }
89    }
90
91    fn element_len<T: 'static>(&self) -> usize {
92        let element_size = std::mem::size_of::<T>();
93        debug_assert!(element_size != 0 && self.byte_len.is_multiple_of(element_size));
94        self.byte_len / element_size
95    }
96
97    fn new_for_runtime(
98        rt: &WebGpuRuntime,
99        handle: cubecl_runtime::server::Handle,
100        byte_len: usize,
101        op: &'static str,
102    ) -> crate::Result<Self> {
103        let Some(_domain) = rt.allocation_domain() else {
104            return Ok(Self::new(
105                handle,
106                byte_len,
107                rt.device_ordinal(),
108                rt.allocation_domain_id(),
109            ));
110        };
111        let managed = rt
112            .client()
113            .get_resource(handle.clone())
114            .map_err(|error| crate::Error::backend_source(op, error))?;
115        let allocation_id = AllocationId::from_backend_id(managed.resource().allocation_id());
116        Ok(Self {
117            handle,
118            byte_len,
119            device_ordinal: rt.device_ordinal(),
120            managed: Some(Arc::new(managed)),
121            allocation_domain: rt.allocation_domain_id(),
122            allocation_id,
123        })
124    }
125}
126
127/// Opaque provider state produced by the shared storage root.
128#[derive(Debug)]
129pub(crate) struct WebGpuPreparedAccess {
130    handle: cubecl_runtime::server::Handle,
131    byte_len: usize,
132    device_ordinal: usize,
133}
134
135impl PreparedDeviceAccess for WebGpuPreparedAccess {
136    fn as_any(&self) -> &dyn std::any::Any {
137        self
138    }
139
140    fn into_any(self: Box<Self>) -> Box<dyn std::any::Any> {
141        self
142    }
143}
144
145struct WebGpuReadMapping {
146    guard: cubecl_wgpu::WgpuMappedReadGuard,
147    range: std::ops::Range<usize>,
148}
149
150impl std::ops::Deref for WebGpuReadMapping {
151    type Target = [u8];
152
153    fn deref(&self) -> &Self::Target {
154        &self.guard[self.range.clone()]
155    }
156}
157
158impl AsRef<[u8]> for WebGpuReadMapping {
159    fn as_ref(&self) -> &[u8] {
160        self
161    }
162}
163
164struct WebGpuWriteMapping {
165    guard: cubecl_wgpu::WgpuMappedWriteGuard,
166    bytes: Vec<u8>,
167}
168
169impl std::ops::Deref for WebGpuWriteMapping {
170    type Target = [u8];
171
172    fn deref(&self) -> &Self::Target {
173        &self.bytes
174    }
175}
176
177impl std::ops::DerefMut for WebGpuWriteMapping {
178    fn deref_mut(&mut self) -> &mut Self::Target {
179        &mut self.bytes
180    }
181}
182
183impl AsRef<[u8]> for WebGpuWriteMapping {
184    fn as_ref(&self) -> &[u8] {
185        self
186    }
187}
188
189impl AsMut<[u8]> for WebGpuWriteMapping {
190    fn as_mut(&mut self) -> &mut [u8] {
191        self
192    }
193}
194
195impl Drop for WebGpuWriteMapping {
196    fn drop(&mut self) {
197        self.guard.copy_from_slice(&self.bytes);
198    }
199}
200
201fn provider_dtype_size(dtype: DType) -> usize {
202    match dtype {
203        DType::F32 | DType::I32 => core::mem::size_of::<f32>(),
204        DType::F64 | DType::I64 => core::mem::size_of::<f64>(),
205        DType::Bool => core::mem::size_of::<bool>(),
206        DType::C32 => core::mem::size_of::<num_complex::Complex32>(),
207        DType::C64 => core::mem::size_of::<num_complex::Complex64>(),
208    }
209}
210
211fn provider_mapping_range(
212    buffer: &WebGpuBuffer,
213    span: RootBoundSpan,
214    dtype: DType,
215) -> Result<std::ops::Range<usize>, AccessError> {
216    let start = span.byte_offset();
217    let end = start
218        .checked_add(span.byte_len())
219        .ok_or_else(|| AccessError::Provider {
220            message: "WebGPU mapping span overflows".to_owned(),
221        })?;
222    if end > buffer.byte_len {
223        return Err(AccessError::Provider {
224            message: "WebGPU mapping span exceeds the allocation".to_owned(),
225        });
226    }
227    let element_size = provider_dtype_size(dtype);
228    if !start.is_multiple_of(element_size) || !span.byte_len().is_multiple_of(element_size) {
229        return Err(AccessError::Provider {
230            message: "WebGPU mapping span is not element-aligned".to_owned(),
231        });
232    }
233    Ok(start..end)
234}
235
236// SAFETY: WebGpuBuffer owns exactly one CubeCL allocation handle. Its provider
237// guards retain the underlying managed resource for every borrowed mapping;
238// the root importer consumes the buffer exactly once and the provider handle
239// is never used as a public ownership authority.
240unsafe impl BackendAllocation for WebGpuBuffer {
241    fn root_extent(&self) -> RootResourceExtent {
242        RootResourceExtent::try_new(
243            AllocationKey::new(self.allocation_domain, self.allocation_id),
244            0,
245            self.byte_len,
246            8,
247        )
248        .expect("WebGPU allocation metadata is constructed with a valid extent")
249    }
250
251    fn provider_kind(&self) -> BackendId {
252        BackendId::WebGpu
253    }
254
255    fn capabilities(&self) -> ProviderCapabilities {
256        if self.managed.is_some() {
257            ProviderCapabilities::host()
258        } else {
259            ProviderCapabilities::none()
260        }
261    }
262
263    fn prepare_device_access(
264        &self,
265        request: DeviceAccessRequest<'_>,
266    ) -> Result<Box<dyn PreparedDeviceAccess>, DeviceAccessError> {
267        if request.allocation_domain() != self.allocation_domain
268            || request.allocation_id() != self.allocation_id
269        {
270            return Err(DeviceAccessError::InvalidRequest {
271                message: "prepared request does not match the WebGPU allocation identity"
272                    .to_owned(),
273            });
274        }
275        if request.byte_len() > self.byte_len {
276            return Err(DeviceAccessError::InvalidRequest {
277                message: "prepared request exceeds the WebGPU allocation extent".to_owned(),
278            });
279        }
280        Ok(Box::new(WebGpuPreparedAccess {
281            handle: self.handle.clone(),
282            byte_len: self.byte_len,
283            device_ordinal: self.device_ordinal,
284        }))
285    }
286
287    fn map_read(
288        &self,
289        span: RootBoundSpan,
290        dtype: DType,
291    ) -> Result<ProviderReadMapping<'_>, AccessError> {
292        let managed = self.managed.as_ref().ok_or(AccessError::Unsupported {
293            backend: "cubecl-webgpu",
294        })?;
295        let range = provider_mapping_range(self, span, dtype)?;
296        let guard = managed
297            .resource()
298            .map_read()
299            .map_err(|error| AccessError::Provider {
300                message: error.to_string(),
301            })?;
302        if range.end > guard.len() {
303            return Err(AccessError::Provider {
304                message: "WebGPU host mapping is shorter than the checked root extent".to_owned(),
305            });
306        }
307        Ok(ProviderReadMapping::from_guard(WebGpuReadMapping {
308            guard,
309            range,
310        }))
311    }
312
313    fn map_write(
314        &self,
315        span: RootBoundSpan,
316        dtype: DType,
317    ) -> Result<ProviderWriteMapping<'_>, AccessError> {
318        let managed = self.managed.as_ref().ok_or(AccessError::Unsupported {
319            backend: "cubecl-webgpu",
320        })?;
321        let range = provider_mapping_range(self, span, dtype)?;
322        let guard = managed
323            .resource()
324            .map_write()
325            .map_err(|error| AccessError::Provider {
326                message: error.to_string(),
327            })?;
328        if range.end > guard.len() {
329            return Err(AccessError::Provider {
330                message: "WebGPU host mapping is shorter than the checked root extent".to_owned(),
331            });
332        }
333        let bytes = vec![0_u8; range.len()];
334        Ok(ProviderWriteMapping::from_guard(WebGpuWriteMapping {
335            guard,
336            bytes,
337        }))
338    }
339
340    fn as_any(&self) -> &dyn std::any::Any {
341        self
342    }
343
344    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
345        self
346    }
347}
348
349pub(super) fn prepared_webgpu_tensor<T: TensorScalar + 'static>(
350    tensor: &TypedTensor<T>,
351    op: &'static str,
352) -> crate::Result<WebGpuPreparedAccess> {
353    let prepared = tensor.prepare_device_read(op)?;
354    prepared
355        .into_any()
356        .downcast::<WebGpuPreparedAccess>()
357        .map(|prepared| *prepared)
358        .map_err(|_| crate::Error::runtime_state(op, "expected a WebGPU prepared allocation"))
359}
360
361impl WebGpuPreparedAccess {
362    pub(crate) const fn device_ordinal(&self) -> usize {
363        self.device_ordinal
364    }
365}
366
367pub(super) fn prepared_webgpu_view<T: TensorScalar + 'static, R: TensorRank>(
368    view: &TypedTensorView<'_, T, R>,
369    op: &'static str,
370) -> crate::Result<WebGpuPreparedAccess> {
371    let prepared = view.prepare_device_read(op)?;
372    prepared
373        .into_any()
374        .downcast::<WebGpuPreparedAccess>()
375        .map(|prepared| *prepared)
376        .map_err(|_| crate::Error::runtime_state(op, "expected a WebGPU prepared allocation"))
377}
378
379fn checked_shape_product(op: &'static str, shape: &[usize]) -> crate::Result<usize> {
380    shape
381        .iter()
382        .try_fold(1usize, |acc, &dim| acc.checked_mul(dim))
383        .ok_or_else(|| {
384            Error::invalid_argument(
385                op,
386                "shape",
387                format!("shape product overflow for shape {shape:?}"),
388            )
389        })
390}
391
392fn cube_count_for_len(len: usize) -> crate::Result<CubeCount> {
393    let cubes = len.div_ceil(DEFAULT_CUBE_DIM_X as usize);
394    let cubes = u32::try_from(cubes).map_err(|_| {
395        Error::invalid_argument(
396            "cube_count_for_len",
397            "length",
398            format!(
399                "1D WebGPU launch for {len} elements requires {cubes} cubes, \
400                 which exceeds u32::MAX"
401            ),
402        )
403    })?;
404    Ok(CubeCount::Static(cubes.max(1), 1, 1))
405}
406
407fn cube_dim_1d() -> CubeDim {
408    CubeDim::new_1d(DEFAULT_CUBE_DIM_X)
409}
410
411fn comptime_sequence<T: CubeType + Clone>(values: &[T]) -> Sequence<T> {
412    let mut out = Sequence::new();
413    for value in values {
414        out.push(value.clone());
415    }
416    out
417}
418
419fn typed_tensor_binding_with_layout<T: CubeElement + TensorScalar + Clone>(
420    tensor: &TypedTensor<T>,
421    shape: &[usize],
422    strides: &[usize],
423    op: &'static str,
424) -> crate::Result<TensorBinding<WgpuRuntime>> {
425    if shape.len() != strides.len() {
426        return Err(Error::rank_mismatch(op, shape.len(), strides.len()));
427    }
428    let prepared = prepared_webgpu_tensor(tensor, op)?;
429    let layout_len = checked_shape_product(op, shape)?;
430    if layout_len != tensor.n_elements() {
431        return Err(Error::runtime_state(
432            op,
433            format!(
434                "WebGPU tensor binding layout covers {layout_len} elements, tensor has {}",
435                tensor.n_elements()
436            ),
437        ));
438    }
439
440    let (shape, strides) = if shape.is_empty() {
441        (vec![1], vec![1])
442    } else {
443        (shape.to_vec(), strides.to_vec())
444    };
445
446    // SAFETY: The tensor root prepared the provider allocation for this exact
447    // descriptor before the binding is constructed. The caller-provided
448    // shape/stride metadata covers the validated logical tensor extent.
449    Ok(unsafe { TensorBinding::from_raw_parts(prepared.handle, strides.into(), shape.into()) })
450}
451
452pub(super) fn ensure_resident_on_runtime<T: TensorScalar + 'static>(
453    rt: &WebGpuRuntime,
454    tensor: &TypedTensor<T>,
455    op: &'static str,
456) -> crate::Result<()> {
457    let view = tensor.as_view();
458    let expected_allocation_domain = rt.allocation_domain_id();
459    let Some(actual_allocation_domain) = tensor.allocation_domain() else {
460        return Err(Error::runtime_state(
461            op,
462            "expected a WebGPU backend tensor, got host storage",
463        ));
464    };
465    if actual_allocation_domain != expected_allocation_domain {
466        return Err(Error::host_access(
467            op,
468            HostAccessError::ForeignDomain {
469                expected: expected_allocation_domain,
470                actual: actual_allocation_domain,
471            },
472        ));
473    }
474    if !matches!(view.backend_family(), Some("webgpu" | "cubecl-webgpu")) {
475        return Err(Error::runtime_state(
476            op,
477            "expected a WebGPU allocation from the selected provider",
478        ));
479    }
480    ensure_placement_resident_on_runtime(rt, tensor.placement(), op)
481}
482
483fn ensure_placement_resident_on_runtime(
484    rt: &WebGpuRuntime,
485    placement: &Placement,
486    op: &'static str,
487) -> crate::Result<()> {
488    let expected_memory = if rt.allocation_domain().is_some() {
489        MemoryKind::Managed
490    } else {
491        MemoryKind::Device
492    };
493    if placement.memory_kind != expected_memory {
494        return Err(Error::runtime_state(
495            op,
496            format!(
497                "expected WebGPU tensor placement, got {:?}",
498                placement.memory_kind
499            ),
500        ));
501    }
502    match &placement.device {
503        Some(device)
504            if device.kind == DeviceKind::Gpu(GpuBackendKind::WebGpu)
505                && device.ordinal == rt.device_ordinal() =>
506        {
507            Ok(())
508        }
509        Some(device) => Err(Error::runtime_state(
510            op,
511            format!(
512                "expected WebGPU tensor resident on webgpu:{}, got {:?}:{}",
513                rt.device_ordinal(),
514                device.kind,
515                device.ordinal
516            ),
517        )),
518        None => Err(Error::runtime_state(
519            op,
520            format!(
521                "expected WebGPU tensor resident on webgpu:{}, got missing device metadata",
522                rt.device_ordinal()
523            ),
524        )),
525    }
526}
527
528pub(super) fn typed_from_webgpu<T: TensorScalar + Send + Sync + 'static>(
529    shape: Vec<usize>,
530    buffer: WebGpuBuffer,
531    rt: &WebGpuRuntime,
532) -> crate::Result<TypedTensor<T>> {
533    let expected_len = checked_shape_product("typed_from_webgpu", &shape)?;
534    if expected_len != buffer.element_len::<T>() {
535        return Err(Error::runtime_state(
536            "typed_from_webgpu",
537            format!(
538                "WebGPU allocation has {} elements, shape requires {expected_len}",
539                buffer.element_len::<T>()
540            ),
541        ));
542    }
543    TypedTensor::from_backend_allocation(shape, Box::new(buffer), webgpu_placement(rt))
544}
545
546fn alloc_output<T: CubeElement + TensorScalar + Clone + Send + Sync + 'static>(
547    rt: &WebGpuRuntime,
548    shape: &[usize],
549    op: &'static str,
550) -> crate::Result<TypedTensor<T>> {
551    let len = checked_shape_product(op, shape)?;
552    let bytes = len.checked_mul(core::mem::size_of::<T>()).ok_or_else(|| {
553        Error::invalid_argument(
554            op,
555            "shape",
556            format!("WebGPU output byte length overflow for shape {shape:?}"),
557        )
558    })?;
559    let handle = rt.client().empty(bytes);
560    let buffer = WebGpuBuffer::new_for_runtime(rt, handle, bytes, op)?;
561    typed_from_webgpu(shape.to_vec(), buffer, rt)
562}
563
564pub(super) fn alloc_tensor_in_runtime(
565    rt: &WebGpuRuntime,
566    dtype: DType,
567    shape: &[usize],
568) -> crate::Result<Tensor> {
569    match dtype {
570        DType::F32 => alloc_output::<f32>(rt, shape, "apple_alloc").map(Tensor::F32),
571        DType::F64 => alloc_output::<f64>(rt, shape, "apple_alloc").map(Tensor::F64),
572        DType::I32 => alloc_output::<i32>(rt, shape, "apple_alloc").map(Tensor::I32),
573        DType::I64 => alloc_output::<i64>(rt, shape, "apple_alloc").map(Tensor::I64),
574        DType::C32 => {
575            alloc_output::<num_complex::Complex32>(rt, shape, "apple_alloc").map(Tensor::C32)
576        }
577        DType::C64 => {
578            alloc_output::<num_complex::Complex64>(rt, shape, "apple_alloc").map(Tensor::C64)
579        }
580        DType::Bool => {
581            let len = checked_shape_product("apple_alloc", shape)?;
582            let handle = rt.client().empty(len);
583            let buffer = WebGpuBuffer::new_for_runtime(rt, handle, len, "apple_alloc")?;
584            Ok(Tensor::Bool(TypedTensor::from_backend_allocation(
585                shape.to_vec(),
586                Box::new(buffer),
587                webgpu_placement(rt),
588            )?))
589        }
590    }
591}
592
593fn webgpu_placement(rt: &WebGpuRuntime) -> Placement {
594    Placement {
595        memory_kind: if rt.allocation_domain().is_some() {
596            MemoryKind::Managed
597        } else {
598            MemoryKind::Device
599        },
600        device: Some(DeviceId {
601            kind: DeviceKind::Gpu(GpuBackendKind::WebGpu),
602            ordinal: rt.device_ordinal(),
603        }),
604        cpu_affinity: None,
605    }
606}
607
608/// CubeCL WebGPU tensor backend.
609///
610/// # Examples
611///
612/// ```
613/// use tenferro_gpu::webgpu::WebGpuBackend;
614///
615/// let _ctor: fn(usize) -> tenferro_tensor::Result<WebGpuBackend> = WebGpuBackend::new;
616/// ```
617#[doc(hidden)]
618struct WebGpuBackendSessionMarker;
619
620#[derive(Clone)]
621pub struct WebGpuBackend {
622    runtime: WebGpuRuntime,
623}
624
625impl fmt::Debug for WebGpuBackend {
626    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
627        f.debug_struct("WebGpuBackend")
628            .field("runtime", &self.runtime)
629            .finish_non_exhaustive()
630    }
631}
632
633impl WebGpuBackend {
634    /// Initialize a WebGPU backend for a discrete GPU ordinal.
635    ///
636    /// # Examples
637    ///
638    /// ```
639    /// use tenferro_gpu::webgpu::WebGpuBackend;
640    ///
641    /// let _ctor: fn(usize) -> tenferro_tensor::Result<WebGpuBackend> = WebGpuBackend::new;
642    /// ```
643    ///
644    /// # Errors
645    ///
646    /// Returns [`crate::Error::RuntimeState`] when no adapter/device is
647    /// available, or [`crate::Error::BackendSource`] when CubeCL initialization
648    /// fails.
649    pub fn new(device_ordinal: usize) -> crate::Result<Self> {
650        WebGpuRuntime::new(device_ordinal).map(Self::from_runtime)
651    }
652
653    /// Initialize a WebGPU backend using CubeCL's default adapter selection.
654    ///
655    /// # Examples
656    ///
657    /// ```
658    /// use tenferro_gpu::webgpu::WebGpuBackend;
659    ///
660    /// let _ctor: fn() -> tenferro_tensor::Result<WebGpuBackend> = WebGpuBackend::new_default;
661    /// ```
662    ///
663    /// # Errors
664    ///
665    /// Returns [`crate::Error::RuntimeState`] when default adapter selection
666    /// is unavailable, or [`crate::Error::BackendSource`] when initialization
667    /// fails.
668    pub fn new_default() -> crate::Result<Self> {
669        WebGpuRuntime::new_default().map(Self::from_runtime)
670    }
671
672    /// Build a WebGPU backend from an initialized runtime.
673    ///
674    /// # Examples
675    ///
676    /// ```
677    /// use tenferro_gpu::{webgpu::WebGpuBackend, webgpu::WebGpuRuntime};
678    ///
679    /// let _from_runtime: fn(WebGpuRuntime) -> WebGpuBackend = WebGpuBackend::from_runtime;
680    /// ```
681    pub fn from_runtime(runtime: WebGpuRuntime) -> Self {
682        Self { runtime }
683    }
684
685    /// Return this backend's WebGPU runtime.
686    ///
687    /// # Examples
688    ///
689    /// ```
690    /// use tenferro_gpu::{webgpu::WebGpuBackend, webgpu::WebGpuRuntime};
691    ///
692    /// let _runtime: fn(&WebGpuBackend) -> &WebGpuRuntime = WebGpuBackend::runtime;
693    /// ```
694    pub fn runtime(&self) -> &WebGpuRuntime {
695        &self.runtime
696    }
697
698    /// Return the opaque identity of this exact executable backend instance.
699    ///
700    /// Clones of a backend return the same identity. Independently initialized
701    /// backends return different identities even when they target the same
702    /// WebGPU device ordinal. This also covers Apple-backed WebGPU runtimes.
703    ///
704    /// # Examples
705    ///
706    /// ```
707    /// use tenferro_gpu::webgpu::WebGpuBackend;
708    ///
709    /// let _identity = WebGpuBackend::runtime_identity;
710    /// ```
711    pub fn runtime_identity(&self) -> WebGpuRuntimeIdentity {
712        self.runtime.runtime_identity()
713    }
714
715    /// Block until queued WebGPU work completes.
716    ///
717    /// # Examples
718    ///
719    /// ```
720    /// use tenferro_gpu::webgpu::WebGpuBackend;
721    ///
722    /// let _sync: fn(&WebGpuBackend) -> tenferro_tensor::Result<()> = WebGpuBackend::synchronize;
723    /// ```
724    ///
725    /// # Errors
726    ///
727    /// Returns [`crate::Error::BackendSource`] when queue flush or
728    /// synchronization fails, or [`crate::Error::RuntimeState`] when the
729    /// runtime has lost its device state.
730    pub fn synchronize(&self) -> crate::Result<()> {
731        self.runtime.synchronize()
732    }
733}
734
735fn unsupported_op(op: &'static str) -> crate::Error {
736    crate::Error::unsupported(
737        op,
738        "WebGPU backend does not support this operation yet; upload/download explicitly and use a supported backend operation",
739    )
740}
741
742macro_rules! unsupported {
743    ($op:literal) => {
744        Err(unsupported_op($op))
745    };
746}
747
748impl TensorElementwise for WebGpuBackend {
749    fn add(&mut self, _lhs: &Tensor, _rhs: &Tensor) -> crate::Result<Tensor> {
750        unsupported!("webgpu_add")
751    }
752
753    fn sub(&mut self, _lhs: &Tensor, _rhs: &Tensor) -> crate::Result<Tensor> {
754        unsupported!("webgpu_sub")
755    }
756
757    fn mul(&mut self, _lhs: &Tensor, _rhs: &Tensor) -> crate::Result<Tensor> {
758        unsupported!("webgpu_mul")
759    }
760
761    fn neg(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
762        unsupported!("webgpu_neg")
763    }
764
765    fn conj(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
766        unsupported!("webgpu_conj")
767    }
768
769    fn div(&mut self, _lhs: &Tensor, _rhs: &Tensor) -> crate::Result<Tensor> {
770        unsupported!("webgpu_div")
771    }
772
773    fn abs(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
774        unsupported!("webgpu_abs")
775    }
776
777    fn sign(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
778        unsupported!("webgpu_sign")
779    }
780
781    fn maximum(&mut self, _lhs: &Tensor, _rhs: &Tensor) -> crate::Result<Tensor> {
782        unsupported!("webgpu_maximum")
783    }
784
785    fn minimum(&mut self, _lhs: &Tensor, _rhs: &Tensor) -> crate::Result<Tensor> {
786        unsupported!("webgpu_minimum")
787    }
788
789    fn compare(
790        &mut self,
791        _lhs: &Tensor,
792        _rhs: &Tensor,
793        _dir: &CompareDir,
794    ) -> crate::Result<Tensor> {
795        unsupported!("webgpu_compare")
796    }
797
798    fn select(
799        &mut self,
800        _pred: &Tensor,
801        _on_true: &Tensor,
802        _on_false: &Tensor,
803    ) -> crate::Result<Tensor> {
804        unsupported!("webgpu_select")
805    }
806
807    fn clamp(
808        &mut self,
809        _input: &Tensor,
810        _lower: &Tensor,
811        _upper: &Tensor,
812    ) -> crate::Result<Tensor> {
813        unsupported!("webgpu_clamp")
814    }
815}
816
817impl TensorAnalytic for WebGpuBackend {
818    fn exp(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
819        unsupported!("webgpu_exp")
820    }
821
822    fn log(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
823        unsupported!("webgpu_log")
824    }
825
826    fn sin(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
827        unsupported!("webgpu_sin")
828    }
829
830    fn cos(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
831        unsupported!("webgpu_cos")
832    }
833
834    fn tanh(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
835        unsupported!("webgpu_tanh")
836    }
837
838    fn sqrt(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
839        unsupported!("webgpu_sqrt")
840    }
841
842    fn rsqrt(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
843        unsupported!("webgpu_rsqrt")
844    }
845
846    fn pow(&mut self, _lhs: &Tensor, _rhs: &Tensor) -> crate::Result<Tensor> {
847        unsupported!("webgpu_pow")
848    }
849
850    fn expm1(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
851        unsupported!("webgpu_expm1")
852    }
853
854    fn log1p(&mut self, _input: &Tensor) -> crate::Result<Tensor> {
855        unsupported!("webgpu_log1p")
856    }
857}
858
859impl TensorStructural for WebGpuBackend {
860    fn to_contiguous_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
861        structural::to_contiguous_read(self, input)
862    }
863
864    fn copy_read_into(&mut self, _src: TensorRead<'_>, _dst: TensorWrite<'_>) -> crate::Result<()> {
865        unsupported!("WebGpuBackend::copy_read_into")
866    }
867
868    fn transpose(&mut self, input: &Tensor, perm: &[usize]) -> crate::Result<Tensor> {
869        structural::transpose(self, input, perm)
870    }
871
872    fn reshape(&mut self, _input: &Tensor, _shape: &[usize]) -> crate::Result<Tensor> {
873        unsupported!("webgpu_reshape")
874    }
875
876    fn broadcast_in_dim(
877        &mut self,
878        _input: &Tensor,
879        _shape: &[usize],
880        _dims: &[usize],
881    ) -> crate::Result<Tensor> {
882        unsupported!("webgpu_broadcast_in_dim")
883    }
884
885    fn cast(&mut self, _input: &Tensor, _to: DType) -> crate::Result<Tensor> {
886        unsupported!("webgpu_cast")
887    }
888
889    fn extract_diagonal(
890        &mut self,
891        _input: &Tensor,
892        _axis_a: usize,
893        _axis_b: usize,
894    ) -> crate::Result<Tensor> {
895        unsupported!("webgpu_extract_diagonal")
896    }
897
898    fn embed_diagonal(
899        &mut self,
900        _input: &Tensor,
901        _axis_a: usize,
902        _axis_b: usize,
903    ) -> crate::Result<Tensor> {
904        unsupported!("webgpu_embed_diagonal")
905    }
906
907    fn tril(&mut self, _input: &Tensor, _k: i64) -> crate::Result<Tensor> {
908        unsupported!("webgpu_tril")
909    }
910
911    fn triu(&mut self, _input: &Tensor, _k: i64) -> crate::Result<Tensor> {
912        unsupported!("webgpu_triu")
913    }
914}
915
916impl TensorViewCanonicalization<f32, tenferro_tensor::DynRank> for WebGpuBackend {
917    fn to_contiguous(
918        &mut self,
919        view: &TypedTensorView<'_, f32>,
920    ) -> crate::Result<TypedTensor<f32>> {
921        structural::to_contiguous_f32(self, view)
922    }
923
924    fn copy_into(
925        &mut self,
926        _src: &TypedTensorView<'_, f32>,
927        _dst: &mut TypedTensorViewMut<'_, f32>,
928    ) -> crate::Result<()> {
929        unsupported!("WebGpuBackend::copy_into")
930    }
931}
932
933impl TensorReduction for WebGpuBackend {
934    fn reduce_sum(&mut self, _input: &Tensor, _axes: &[usize]) -> crate::Result<Tensor> {
935        unsupported!("webgpu_reduce_sum")
936    }
937
938    fn reduce_prod(&mut self, _input: &Tensor, _axes: &[usize]) -> crate::Result<Tensor> {
939        unsupported!("webgpu_reduce_prod")
940    }
941
942    fn reduce_max(&mut self, _input: &Tensor, _axes: &[usize]) -> crate::Result<Tensor> {
943        unsupported!("webgpu_reduce_max")
944    }
945
946    fn reduce_min(&mut self, _input: &Tensor, _axes: &[usize]) -> crate::Result<Tensor> {
947        unsupported!("webgpu_reduce_min")
948    }
949}
950
951impl TensorDot for WebGpuBackend {
952    fn dot_general(
953        &mut self,
954        lhs: &Tensor,
955        rhs: &Tensor,
956        config: &DotGeneralConfig,
957    ) -> crate::Result<Tensor> {
958        gemm::dot_general(self, lhs, rhs, config)
959    }
960
961    fn dot_general_with_conj(
962        &mut self,
963        lhs: &Tensor,
964        rhs: &Tensor,
965        config: &DotGeneralConfig,
966        lhs_conj: bool,
967        rhs_conj: bool,
968    ) -> crate::Result<Tensor> {
969        gemm::dot_general_with_conj(self, lhs, rhs, config, lhs_conj, rhs_conj)
970    }
971}
972
973impl TensorIndexing for WebGpuBackend {
974    fn gather(
975        &mut self,
976        _operand: &Tensor,
977        _start_indices: &Tensor,
978        _config: &GatherConfig,
979    ) -> crate::Result<Tensor> {
980        unsupported!("webgpu_gather")
981    }
982
983    fn scatter(
984        &mut self,
985        _operand: &Tensor,
986        _scatter_indices: &Tensor,
987        _updates: &Tensor,
988        _config: &ScatterConfig,
989    ) -> crate::Result<Tensor> {
990        unsupported!("webgpu_scatter")
991    }
992
993    fn slice(&mut self, _input: &Tensor, _config: &SliceConfig) -> crate::Result<Tensor> {
994        unsupported!("webgpu_slice")
995    }
996
997    fn dynamic_slice(
998        &mut self,
999        _input: &Tensor,
1000        _starts: &Tensor,
1001        _slice_sizes: &[usize],
1002    ) -> crate::Result<Tensor> {
1003        unsupported!("webgpu_dynamic_slice")
1004    }
1005
1006    fn dynamic_update_slice(
1007        &mut self,
1008        _operand: &Tensor,
1009        _update: &Tensor,
1010        _starts: &Tensor,
1011    ) -> crate::Result<Tensor> {
1012        unsupported!("webgpu_dynamic_update_slice")
1013    }
1014
1015    fn pad(&mut self, _input: &Tensor, _config: &PadConfig) -> crate::Result<Tensor> {
1016        unsupported!("webgpu_pad")
1017    }
1018
1019    fn concatenate(&mut self, _inputs: &[&Tensor], _axis: usize) -> crate::Result<Tensor> {
1020        unsupported!("webgpu_concatenate")
1021    }
1022
1023    fn reverse(&mut self, _input: &Tensor, _axes: &[usize]) -> crate::Result<Tensor> {
1024        unsupported!("webgpu_reverse")
1025    }
1026}
1027
1028impl TensorFusion for WebGpuBackend {}
1029
1030impl TensorBuffer for WebGpuBackend {}
1031
1032impl TensorDeviceTransfer for WebGpuBackend {
1033    fn download_to_host(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor> {
1034        let tensor = tensor.as_tensor().ok_or_else(|| {
1035            crate::Error::unsupported(
1036                "WebGpuBackend::download_to_host",
1037                "WebGPU transfer currently requires an owned tensor; materialize a view explicitly first",
1038            )
1039        })?;
1040        download_webgpu_tensor(self.runtime(), tensor)
1041    }
1042
1043    fn upload_host_tensor(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor> {
1044        let tensor = tensor.as_tensor().ok_or_else(|| {
1045            crate::Error::unsupported(
1046                "WebGpuBackend::upload_host_tensor",
1047                "WebGPU transfer currently requires an owned tensor; materialize a view explicitly first",
1048            )
1049        })?;
1050        upload_webgpu_tensor(self.runtime(), tensor)
1051    }
1052}
1053
1054impl BackendRuntimeCache for WebGpuBackend {
1055    type RuntimeCache = ();
1056}
1057
1058impl BackendSession for WebGpuBackend {
1059    fn session_type_id(&self) -> std::any::TypeId {
1060        std::any::TypeId::of::<WebGpuBackendSessionMarker>()
1061    }
1062
1063    unsafe fn session_data_mut(&mut self) -> *mut () {
1064        self as *mut Self as *mut ()
1065    }
1066}
1067
1068impl BackendCachedDot for WebGpuBackend {}
1069
1070impl TensorBackend for WebGpuBackend {}
tenferro_gpu/webgpu/mod.rs

tenferro_gpu/webgpu/
mod.rs
