tenferro_gpu/webgpu/
runtime_adapter.rs

1use std::mem::{size_of, size_of_val};
2use std::sync::Arc;
3
4use tenferro_runtime::program::{CoreSemanticOp, SemanticOpRef, SemanticOperationView};
5use tenferro_runtime::{
6    assemble_executable_engine_registration, CoreCapabilityBundle, CoreCapabilityKind,
7    CorePrepareContext, DotGeneralPreparation, DotGeneralPrepareRequest, EngineId,
8    EngineRegistration, EngineRegistrationMetadata, ExecutableEngineRegistrationConfig,
9    ExecutionContextIdentity, HardwareClassId, InputIngressContract, InputPlacementContract,
10    InputSignature, InputSignatureContract, InputSpecializationProjection,
11    InputSpecializationRequirements, LayoutProjection, LayoutSpecialization, PrepareCapability,
12    PrepareError, PreparedOperation, PreparedOperationBinding, PreparedOperationPlan,
13    ProviderContractError, ProviderDeviceIdentity, ProviderId, ResidentOutputContract,
14    RuntimeConfigError, RuntimeInputContract, SpecializationError, SpecializationProjection,
15    SpecializationRequirements, StorageClass, UnsupportedReason,
16};
17use tenferro_tensor::{
18    AllocationDomainId, DeviceKind, GpuBackendKind, MemoryKind, Placement, TensorRead, TensorView,
19};
20
21#[cfg(not(target_family = "wasm"))]
22use super::event_domain::WebGpuEventDomainDriver;
23use super::{prepared_webgpu_view, WebGpuBackend};
24#[cfg(target_family = "wasm")]
25use tenferro_runtime::{
26    assemble_preparation_only_engine_registration, PreparationOnlyEngineRegistrationConfig,
27};
28
29const WEBGPU_ENGINE_ID: &str = "tenferro-webgpu.default.v1";
30const WEBGPU_HARDWARE_CLASS_ID: &str = "tenferro-webgpu.device.v1";
31const WEBGPU_STORAGE_CLASS_ID: &str = "tenferro.storage.device.v1";
32
33/// Return the canonical WebGPU runtime engine identifier.
34///
35/// # Examples
36///
37/// ```
38/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
39/// # #[cfg(feature = "webgpu")]
40/// # {
41/// let engine = tenferro_gpu::webgpu::webgpu_runtime_engine_id()?;
42/// assert_eq!(engine.as_str(), "tenferro-webgpu.default.v1");
43/// # }
44/// # Ok(())
45/// # }
46/// ```
47///
48/// # Errors
49///
50/// Returns [`RuntimeConfigError`] if the built-in WebGPU engine identifier
51/// violates runtime identifier validation.
52pub fn webgpu_runtime_engine_id() -> Result<EngineId, RuntimeConfigError> {
53    EngineId::new(WEBGPU_ENGINE_ID).map_err(RuntimeConfigError::from)
54}
55
56/// Return the canonical WebGPU runtime hardware class.
57///
58/// # Examples
59///
60/// ```
61/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
62/// # #[cfg(feature = "webgpu")]
63/// # {
64/// let hardware = tenferro_gpu::webgpu::webgpu_runtime_hardware_class()?;
65/// assert_eq!(hardware.as_str(), "tenferro-webgpu.device.v1");
66/// # }
67/// # Ok(())
68/// # }
69/// ```
70///
71/// # Errors
72///
73/// Returns [`RuntimeConfigError`] if the built-in WebGPU hardware class
74/// violates runtime identifier validation.
75pub fn webgpu_runtime_hardware_class() -> Result<HardwareClassId, RuntimeConfigError> {
76    HardwareClassId::new(WEBGPU_HARDWARE_CLASS_ID).map_err(RuntimeConfigError::from)
77}
78
79/// Build a runtime engine registration for a [`WebGpuBackend`].
80///
81/// The registration exposes WebGPU direct preparation for `dot_general` and the
82/// runtime-owned tensor backend execution bridge. Other core operation families
83/// remain unregistered until their WebGPU runtime preparation contracts are
84/// implemented.
85///
86/// # Examples
87///
88/// ```
89/// use tenferro_gpu::webgpu::WebGpuBackend;
90///
91/// let _register: fn(
92///     &WebGpuBackend,
93/// ) -> Result<tenferro_runtime::EngineRegistration, tenferro_runtime::RuntimeConfigError> =
94///     tenferro_gpu::webgpu::webgpu_runtime_engine_registration;
95/// ```
96///
97/// # Errors
98///
99/// Returns [`RuntimeConfigError`] if one of the built-in WebGPU runtime
100/// identifiers violates runtime validation or if the registration is internally
101/// invalid.
102pub fn webgpu_runtime_engine_registration(
103    backend: &WebGpuBackend,
104) -> Result<EngineRegistration, RuntimeConfigError> {
105    webgpu_runtime_engine_registration_with_id(backend, webgpu_runtime_engine_id()?)
106}
107
108/// Build a WebGPU runtime engine registration with a caller-selected engine ID.
109///
110/// The provider/device identity remains tied to the selected WebGPU device;
111/// only the runtime placement identity is supplied by the caller.
112///
113/// # Errors
114///
115/// Returns [`RuntimeConfigError`] if the caller-selected engine identifier,
116/// the selected backend's provider/device identity, or the assembled hardware
117/// and storage metadata fails validation.
118pub fn webgpu_runtime_engine_registration_with_id(
119    backend: &WebGpuBackend,
120    engine_id: EngineId,
121) -> Result<EngineRegistration, RuntimeConfigError> {
122    let backend = Arc::new(backend.clone());
123    let dot_general: Arc<dyn DotGeneralPreparation> = backend.clone();
124    let execution_backend = backend.as_ref().clone();
125
126    let mut capabilities = CoreCapabilityBundle::builder();
127    capabilities.dot_general(dot_general);
128
129    let storage = webgpu_runtime_storage_class()?;
130    let default_storage = storage.clone();
131    let placement_storage = storage.clone();
132    let signature_storage = storage.clone();
133    let runtime_storage = storage.clone();
134    let resident_storage = storage.clone();
135    let runtime = backend.runtime();
136    let device_ordinal = runtime.device_ordinal();
137    let allocation_domain = runtime.allocation_domain_id();
138    let managed_domain = runtime.allocation_domain().map(|domain| domain.id);
139    let provider_device_identity = ProviderDeviceIdentity::new(
140        ProviderId::new("tenferro.webgpu")?,
141        format!("device:{device_ordinal}"),
142    )?;
143    let ingress = InputIngressContract::new(
144        InputPlacementContract::new(move |placement, candidate| {
145            candidate == &placement_storage
146                && webgpu_input_placement(placement, device_ordinal, managed_domain)
147        }),
148        InputSignatureContract::new(move |placement, family, domain, candidate| {
149            candidate == &signature_storage
150                && webgpu_input_signature(
151                    placement,
152                    family,
153                    domain,
154                    device_ordinal,
155                    managed_domain,
156                    allocation_domain,
157                )
158        }),
159        RuntimeInputContract::new(move |input: &TensorRead<'_>, candidate| {
160            candidate == &runtime_storage
161                && webgpu_input_tensor(input, device_ordinal, managed_domain, allocation_domain)
162        }),
163        ResidentOutputContract::new(move |input: &TensorRead<'_>, candidate| {
164            candidate == &resident_storage
165                && webgpu_input_tensor(input, device_ordinal, managed_domain, allocation_domain)
166        }),
167    );
168    let capabilities = capabilities.build();
169    #[cfg(not(target_family = "wasm"))]
170    {
171        let metadata = EngineRegistrationMetadata::new(
172            engine_id,
173            provider_device_identity,
174            webgpu_runtime_hardware_class()?,
175            Arc::from(vec![storage]),
176            default_storage,
177            capabilities,
178        );
179        assemble_executable_engine_registration(ExecutableEngineRegistrationConfig::new(
180            metadata,
181            execution_backend,
182            Arc::new(WebGpuEventDomainDriver::new(backend.runtime().clone())),
183            ingress,
184            None,
185        ))
186    }
187    #[cfg(target_family = "wasm")]
188    {
189        let metadata = EngineRegistrationMetadata::new(
190            engine_id,
191            provider_device_identity,
192            webgpu_runtime_hardware_class()?,
193            Arc::from(vec![storage]),
194            default_storage,
195            capabilities,
196        );
197        assemble_preparation_only_engine_registration(PreparationOnlyEngineRegistrationConfig::new(
198            metadata,
199            ExecutionContextIdentity::of::<WebGpuBackend>(),
200        ))
201    }
202}
203
204fn webgpu_input_signature(
205    placement: &Placement,
206    backend_family: Option<&'static str>,
207    input_domain: Option<AllocationDomainId>,
208    device_ordinal: usize,
209    managed_domain: Option<AllocationDomainId>,
210    allocation_domain: AllocationDomainId,
211) -> bool {
212    webgpu_input_placement(placement, device_ordinal, managed_domain)
213        && matches!(backend_family, Some("webgpu" | "cubecl-webgpu"))
214        && input_domain == Some(allocation_domain)
215}
216
217fn webgpu_input_placement(
218    placement: &Placement,
219    device_ordinal: usize,
220    allocation_domain: Option<AllocationDomainId>,
221) -> bool {
222    placement.memory_kind
223        == if allocation_domain.is_some() {
224            MemoryKind::Managed
225        } else {
226            MemoryKind::Device
227        }
228        && matches!(
229            &placement.device,
230            Some(device)
231                if device.kind == DeviceKind::Gpu(GpuBackendKind::WebGpu)
232                    && device.ordinal == device_ordinal
233        )
234}
235
236fn webgpu_input_tensor(
237    input: &TensorRead<'_>,
238    device_ordinal: usize,
239    managed_domain: Option<AllocationDomainId>,
240    allocation_domain: AllocationDomainId,
241) -> bool {
242    webgpu_input_placement(input.placement(), device_ordinal, managed_domain)
243        && matches!(input.backend_family(), Some("webgpu" | "cubecl-webgpu"))
244        && input.allocation_domain() == Some(allocation_domain)
245        && webgpu_input_has_owned_buffer(input, device_ordinal, allocation_domain)
246}
247
248fn webgpu_input_has_owned_buffer(
249    input: &TensorRead<'_>,
250    device_ordinal: usize,
251    allocation_domain: AllocationDomainId,
252) -> bool {
253    match input.clone().tensor_view() {
254        TensorView::F32(view) => {
255            webgpu_view_has_owner::<f32>(&view, device_ordinal, allocation_domain)
256        }
257        TensorView::F64(view) => {
258            webgpu_view_has_owner::<f64>(&view, device_ordinal, allocation_domain)
259        }
260        TensorView::I32(view) => {
261            webgpu_view_has_owner::<i32>(&view, device_ordinal, allocation_domain)
262        }
263        TensorView::I64(view) => {
264            webgpu_view_has_owner::<i64>(&view, device_ordinal, allocation_domain)
265        }
266        TensorView::Bool(view) => {
267            webgpu_view_has_owner::<bool>(&view, device_ordinal, allocation_domain)
268        }
269        TensorView::C32(view) => webgpu_view_has_owner::<num_complex::Complex32>(
270            &view,
271            device_ordinal,
272            allocation_domain,
273        ),
274        TensorView::C64(view) => webgpu_view_has_owner::<num_complex::Complex64>(
275            &view,
276            device_ordinal,
277            allocation_domain,
278        ),
279    }
280}
281
282fn webgpu_view_has_owner<T: tenferro_tensor::TensorScalar + 'static>(
283    view: &tenferro_tensor::TypedTensorView<'_, T>,
284    device_ordinal: usize,
285    allocation_domain: AllocationDomainId,
286) -> bool {
287    view.backend_family()
288        .is_some_and(|family| family == "webgpu" || family == "cubecl-webgpu")
289        && view.allocation_domain() == Some(allocation_domain)
290        && matches!(
291            &view.placement().device,
292            Some(device)
293                if device.kind == DeviceKind::Gpu(GpuBackendKind::WebGpu)
294                    && device.ordinal == device_ordinal
295        )
296        && prepared_webgpu_view(view, "webgpu_input_tensor")
297            .is_ok_and(|prepared| prepared.device_ordinal() == device_ordinal)
298}
299
300#[cfg(test)]
301#[path = "tests/runtime_adapter.rs"]
302mod tests;
303
304fn webgpu_runtime_storage_class() -> Result<StorageClass, RuntimeConfigError> {
305    StorageClass::new(WEBGPU_STORAGE_CLASS_ID).map_err(RuntimeConfigError::from)
306}
307
308#[derive(Debug)]
309struct WebGpuPreparedOperation {
310    binding: PreparedOperationBinding,
311    specialization: SpecializationProjection,
312}
313
314impl PreparedOperation for WebGpuPreparedOperation {
315    fn binding(&self) -> &PreparedOperationBinding {
316        &self.binding
317    }
318
319    fn specialization(&self) -> &SpecializationProjection {
320        &self.specialization
321    }
322
323    fn retained_bytes(&self) -> usize {
324        checked_specialization_heap_retained_bytes(&self.specialization).unwrap_or(usize::MAX)
325    }
326}
327
328impl DotGeneralPreparation for WebGpuBackend {
329    fn prepare(
330        &self,
331        request: DotGeneralPrepareRequest<'_>,
332    ) -> Result<PrepareCapability, PrepareError> {
333        prepare_webgpu_dot_general(request.operation(), request.context())
334    }
335}
336
337fn prepare_webgpu_dot_general(
338    operation: SemanticOperationView<'_>,
339    context: &CorePrepareContext<'_>,
340) -> Result<PrepareCapability, PrepareError> {
341    validate_webgpu_runtime_context(context)?;
342    let SemanticOpRef::Core(op) = operation.op() else {
343        return Err(wrong_family_error("extension"));
344    };
345    if !matches!(op, CoreSemanticOp::DotGeneral { .. }) {
346        return Err(wrong_family_error(core_operation_name(op)));
347    }
348
349    let minimum = dot_general_specialization_requirements(context.inputs())?;
350    let merged =
351        merge_specialization_requirements(context.specialization().requirements(), &minimum)?;
352    if &merged != context.specialization().requirements() {
353        return Ok(PrepareCapability::NeedsSpecialization(merged));
354    }
355
356    Ok(PrepareCapability::Prepared(
357        PreparedOperationPlan::metadata(Arc::new(WebGpuPreparedOperation {
358            binding: context.binding().clone(),
359            specialization: context.specialization().clone(),
360        })),
361    ))
362}
363
364fn validate_webgpu_runtime_context(context: &CorePrepareContext<'_>) -> Result<(), PrepareError> {
365    let expected_context = ExecutionContextIdentity::of::<WebGpuBackend>();
366    if context.binding().context_identity() != expected_context {
367        return Err(PrepareError::ProviderContract {
368            source: ProviderContractError::WrongOperationFamily {
369                expected: CoreCapabilityKind::DotGeneral,
370                operation: "webgpu-context-mismatch",
371            },
372        });
373    }
374    if context.binding().hardware_class().as_str() != WEBGPU_HARDWARE_CLASS_ID {
375        return Err(PrepareError::ProviderContract {
376            source: ProviderContractError::WrongOperationFamily {
377                expected: CoreCapabilityKind::DotGeneral,
378                operation: "webgpu-hardware-mismatch",
379            },
380        });
381    }
382    if context.resolved_placement().storage_class().as_str() != WEBGPU_STORAGE_CLASS_ID {
383        return Err(PrepareError::Unsupported {
384            reason: UnsupportedReason::StorageClass {
385                storage_class: context.resolved_placement().storage_class().clone(),
386            },
387        });
388    }
389    Ok(())
390}
391
392fn dot_general_specialization_requirements(
393    inputs: &InputSignature,
394) -> Result<SpecializationRequirements, PrepareError> {
395    let mut requirements = Vec::with_capacity(inputs.entries().len());
396    for (input, entry) in inputs.entries().iter().enumerate() {
397        let mut builder = InputSpecializationRequirements::builder();
398        builder
399            .dtype(true)
400            .rank(true)
401            .concrete_dimensions(concrete_axes_for_rank(input, entry.shape().len())?)
402            .layout(LayoutSpecialization::Class);
403        requirements.push(builder.build().map_err(specialization_requirements_error)?);
404    }
405    Ok(SpecializationRequirements::new(requirements))
406}
407
408fn concrete_axes_for_rank(input: usize, rank: usize) -> Result<Vec<u32>, PrepareError> {
409    if u32::try_from(rank).is_err() {
410        return Err(PrepareError::Specialization {
411            source: SpecializationError::ProjectionOverflow { input, rank },
412        });
413    }
414    let mut axes = Vec::with_capacity(rank);
415    for axis in 0..rank {
416        axes.push(
417            u32::try_from(axis).map_err(|_| PrepareError::Specialization {
418                source: SpecializationError::ProjectionOverflow { input, rank },
419            })?,
420        );
421    }
422    Ok(axes)
423}
424
425fn merge_specialization_requirements(
426    current: &SpecializationRequirements,
427    minimum: &SpecializationRequirements,
428) -> Result<SpecializationRequirements, PrepareError> {
429    debug_assert_eq!(current.inputs().len(), minimum.inputs().len());
430    let inputs = current
431        .inputs()
432        .iter()
433        .zip(minimum.inputs())
434        .map(|(current, minimum)| merge_input_requirements(current, minimum))
435        .collect::<Result<Vec<_>, _>>()?;
436    Ok(SpecializationRequirements::new(inputs))
437}
438
439fn merge_input_requirements(
440    current: &InputSpecializationRequirements,
441    minimum: &InputSpecializationRequirements,
442) -> Result<InputSpecializationRequirements, PrepareError> {
443    let mut axes = current.concrete_dimensions().to_vec();
444    for axis in minimum.concrete_dimensions() {
445        if !axes.contains(axis) {
446            axes.push(*axis);
447        }
448    }
449    let layout = current.layout().max(minimum.layout());
450    let rank = current.specializes_rank()
451        || minimum.specializes_rank()
452        || !axes.is_empty()
453        || layout == LayoutSpecialization::ExactStrides;
454    let alignment = match (current.alignment_log2(), minimum.alignment_log2()) {
455        (Some(left), Some(right)) => Some(left.max(right)),
456        (Some(value), None) | (None, Some(value)) => Some(value),
457        (None, None) => None,
458    };
459    let mut builder = InputSpecializationRequirements::builder();
460    builder
461        .dtype(current.specializes_dtype() || minimum.specializes_dtype())
462        .rank(rank)
463        .concrete_dimensions(axes)
464        .placement(current.placement().max(minimum.placement()))
465        .layout(layout)
466        .alignment_log2(alignment);
467    builder.build().map_err(specialization_requirements_error)
468}
469
470fn specialization_requirements_error(
471    source: tenferro_runtime::InputSpecializationRequirementsError,
472) -> PrepareError {
473    PrepareError::Engine {
474        source: Arc::new(source),
475    }
476}
477
478fn wrong_family_error(operation: &'static str) -> PrepareError {
479    PrepareError::ProviderContract {
480        source: ProviderContractError::WrongOperationFamily {
481            expected: CoreCapabilityKind::DotGeneral,
482            operation,
483        },
484    }
485}
486
487fn core_operation_name(op: &CoreSemanticOp) -> &'static str {
488    match op {
489        CoreSemanticOp::DotGeneral { .. } => "dot_general",
490        _ => "non-dot-general-core-operation",
491    }
492}
493
494fn checked_specialization_heap_retained_bytes(
495    specialization: &SpecializationProjection,
496) -> Option<usize> {
497    let requirements = specialization.requirements();
498    checked_sum([
499        requirements
500            .inputs()
501            .len()
502            .checked_mul(size_of::<InputSpecializationRequirements>())?,
503        checked_sum(
504            requirements
505                .inputs()
506                .iter()
507                .map(|input| size_of_val(input.concrete_dimensions())),
508        )?,
509        specialization
510            .inputs()
511            .len()
512            .checked_mul(size_of::<InputSpecializationProjection>())?,
513        checked_sum_options(
514            specialization
515                .inputs()
516                .iter()
517                .map(input_projection_retained_bytes),
518        )?,
519    ])
520}
521
522fn input_projection_retained_bytes(projection: &InputSpecializationProjection) -> Option<usize> {
523    size_of_val(projection.concrete_dimensions()).checked_add(match projection.layout() {
524        Some(LayoutProjection::ExactStrides(strides)) if strides.spilled() => {
525            size_of_val(strides.as_slice())
526        }
527        _ => 0,
528    })
529}
530
531fn checked_sum(values: impl IntoIterator<Item = usize>) -> Option<usize> {
532    values
533        .into_iter()
534        .try_fold(0usize, |sum, value| sum.checked_add(value))
535}
536
537fn checked_sum_options(values: impl IntoIterator<Item = Option<usize>>) -> Option<usize> {
538    values
539        .into_iter()
540        .try_fold(0usize, |sum, value| sum.checked_add(value?))
541}
tenferro_gpu/webgpu/runtime_adapter.rs

tenferro_gpu/webgpu/
runtime_adapter.rs
