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

tenferro_gpu/webgpu/
mod.rs
