diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 8efff290e3d..cdea955fca8 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -56,6 +56,7 @@ set(WEBGPU_SRCS runtime/ops/layer_norm/NativeLayerNorm.cpp runtime/ops/linear_fp32/LinearFp32.cpp runtime/ops/et_vk_conv2d/Conv2d.cpp + runtime/ops/et_vk_sdpa/EtVkSdpa.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp new file mode 100644 index 00000000000..dbe6b03f30c --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp @@ -0,0 +1,303 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct QkParams { + uint32_t B; + uint32_t H; + uint32_t S_q; + uint32_t S_kv; + uint32_t D; + uint32_t has_mask; + uint32_t _pad0; + float scale; +}; +static_assert(sizeof(QkParams) == 32, "QkParams must be 32 bytes"); + +struct AvParams { + uint32_t B; + uint32_t H; + uint32_t S_q; + uint32_t S_kv; + uint32_t D; + uint32_t _pad0; + uint32_t _pad1; + uint32_t _pad2; +}; +static_assert(sizeof(AvParams) == 32, "AvParams must be 32 bytes"); + +// Mirrors the Params struct in sdpa_softmax.wgsl (file-local in Sdpa.cpp, so +// re-declared here for the reuse). +struct SoftmaxParams { + uint32_t num_rows; + uint32_t row_width; + uint32_t _pad0; + uint32_t _pad1; +}; +static_assert(sizeof(SoftmaxParams) == 16, "SoftmaxParams must be 16 bytes"); + +// aten op: et_vk.sdpa.default. Args: [q, k, v, attn_mask, scale, out] (mirrors +// Vulkan fused_sdpa_impl, SDPA.cpp). Non-causal, no KV-cache; all tensors are +// DSHB [B, H, S, D], row-major. Three dispatches: QK (scaled, optional additive +// mask) -> softmax (reused sdpa_softmax.wgsl) -> AV. +void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() != 6) { + throw std::runtime_error("WebGPU et_vk.sdpa: expected 6 args"); + } + const int q_id = args.at(0); + const int k_id = args.at(1); + const int v_id = args.at(2); + const int mask_id = args.at(3); + const int scale_id = args.at(4); + const int out_id = args.at(5); + + WGPUDevice device = graph.device(); + + const auto& q = graph.get_tensor(q_id); + const auto& k = graph.get_tensor(k_id); + const auto& v = graph.get_tensor(v_id); + const auto& out = graph.get_tensor(out_id); + + if (q.dims.size() < 3) { + throw std::runtime_error("WebGPU et_vk.sdpa: q rank < 3"); + } + const int rank = static_cast(q.dims.size()); + const uint32_t D = static_cast(q.dims[rank - 1]); + const uint32_t S_q = static_cast(q.dims[rank - 2]); + const uint32_t H = static_cast(q.dims[rank - 3]); + if (D == 0 || S_q == 0 || H == 0) { + throw std::runtime_error("WebGPU et_vk.sdpa: zero D/S_q/H"); + } + // QK/AV kernels view q/k/v/out as vec4 over D; every model in scope + // (Whisper/Voxtral/DaViT/BART/Hiera) uses D=64 or 128, always %4==0. + utils::check_vec4_aligned(D, "et_vk.sdpa", "D"); + const uint64_t q_numel = utils::check_fp32(q, "et_vk.sdpa", "q"); + const uint32_t B = static_cast(q_numel / (uint64_t(H) * S_q * D)); + + // Asymmetric seq supported (S_q != S_kv, e.g. Hiera pooled query): k/v carry + // S_kv, q carries S_q; B/H/D must match across q/k/v. out is [B, H, S_q, D]. + // When S_q == S_kv this is plain self-attention (bit-identical to before). + if (k.dims != v.dims) { + throw std::runtime_error("WebGPU et_vk.sdpa: k/v shape mismatch"); + } + if (k.dims.size() != q.dims.size()) { + throw std::runtime_error("WebGPU et_vk.sdpa: q/k rank mismatch"); + } + if (static_cast(k.dims[rank - 1]) != D || + static_cast(k.dims[rank - 3]) != H) { + throw std::runtime_error("WebGPU et_vk.sdpa: q/k/v must share H and D"); + } + const uint32_t S_kv = static_cast(k.dims[rank - 2]); + if (S_kv == 0) { + throw std::runtime_error("WebGPU et_vk.sdpa: zero S_kv"); + } + // Leading (batch) dims must agree across q/k/v. + for (int d = 0; d < rank - 3; ++d) { + if (k.dims[d] != q.dims[d]) { + throw std::runtime_error("WebGPU et_vk.sdpa: q/k batch dims mismatch"); + } + } + // out must be [B, H, S_q, D] (same as q's shape). + if (out.dims != q.dims) { + throw std::runtime_error( + "WebGPU et_vk.sdpa: out shape must match q [B, H, S_q, D]"); + } + + const bool has_mask = + graph.get_value_type(mask_id) == WebGPUGraph::ValueType::Tensor; + if (has_mask) { + // The QK shader indexes mask as [B, H, S_q, S_kv] row-major; require it. + const auto& mask = graph.get_tensor(mask_id); + if (mask.nbytes != uint64_t(B) * H * S_q * S_kv * sizeof(float)) { + throw std::runtime_error( + "WebGPU et_vk.sdpa: attn_mask must be [B, H, S_q, S_kv] fp32"); + } + } + + float scale = 1.0f / std::sqrt(static_cast(D)); + const auto scale_type = graph.get_value_type(scale_id); + if (scale_type == WebGPUGraph::ValueType::Double) { + scale = static_cast(graph.get_double(scale_id)); + } else if (scale_type != WebGPUGraph::ValueType::Null) { + throw std::runtime_error("WebGPU et_vk.sdpa: scale must be Double or None"); + } + + const uint64_t num_rows = uint64_t(B) * H * S_q; // attn_weights rows + const uint64_t aw_numel = num_rows * S_kv; // [B, H, S_q, S_kv] + const uint64_t out_numel = uint64_t(B) * H * S_q * D; + const size_t aw_bytes = static_cast(aw_numel) * sizeof(float); + + // Up-front dispatch-limit checks (throw BEFORE any buffer alloc → no leak). + // QK = one thread per (b,h,s) row; AV = one per (b,h,s,d4) vec4 (4 output + // elements/thread); softmax = one workgroup per row. compute_1d_workgroup_ + // count throws past the limit. + const uint32_t qk_wg_size = + utils::clamp_workgroup_size(device, kEtVkSdpaQkWorkgroupSizeX); + const uint32_t qk_wg_count = utils::compute_1d_workgroup_count( + device, static_cast(num_rows), qk_wg_size, "et_vk_sdpa_qk"); + const uint32_t av_wg_size = + utils::clamp_workgroup_size(device, kEtVkSdpaAvWorkgroupSizeX); + const uint64_t out_numel4 = + out_numel / 4; // exact: D % 4 == 0 (checked above) + const uint32_t av_wg_count = utils::compute_1d_workgroup_count( + device, static_cast(out_numel4), av_wg_size, "et_vk_sdpa_av"); + // Near-square 2D grid of workgroups (1 workgroup = 1 row) past the 65535 + // per-dimension ceiling; sdpa_softmax.wgsl recovers the flat row index from + // @builtin(num_workgroups), so no override constant is needed here. + utils::WgCount softmax_grid = utils::compute_2d_workgroup_count( + device, + static_cast(num_rows), + /*workgroup_size=*/1, + "et_vk_sdpa_softmax"); + utils::check_fp32(out, "et_vk.sdpa", "out"); + + // NOTE: graph.create_scratch_buffer allocates a fresh buffer per call (no + // pooling), so a multi-layer graph holds 2 × num_layers × aw_bytes live, not + // 2 × max(aw_bytes) (e.g. S=1024 × 12 layers ≈ 1.1 GB) — a memory-headroom + // gap in the shared allocator, not this op; not a correctness issue. + WGPUBuffer attn_buf = graph.create_scratch_buffer(aw_bytes); + WGPUBuffer softmax_buf = graph.create_scratch_buffer(aw_bytes); + + // ---- Dispatch 1: QK (one thread per (b,h,s) row) ---- + { + QkParams p = {}; + p.B = B; + p.H = H; + p.S_q = S_q; + p.S_kv = S_kv; + p.D = D; + p.has_mask = has_mask ? 1u : 0u; + p.scale = scale; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, &p, sizeof(QkParams)); + graph.add_uniform_buffer_bytes(sizeof(QkParams)); + + // 4-byte dummy storage to satisfy the mask binding when absent (shader + // never reads it under has_mask == 0). + utils::OptionalBinding mask = utils::make_optional_binding( + device, + has_mask, + has_mask ? graph.get_tensor(mask_id).buffer : nullptr, + has_mask ? graph.get_tensor(mask_id).nbytes : 0); + + WGPUConstantEntry wg_const = utils::make_wg_size_constant(qk_wg_size); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kEtVkSdpaQkWGSL, + { + {0, WGPUBufferBindingType_Storage, attn_buf, aw_bytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, q.buffer, q.nbytes}, + {2, WGPUBufferBindingType_ReadOnlyStorage, k.buffer, k.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + mask.buffer, + mask.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(QkParams)}, + }, + &wg_const, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, qk_wg_count}); + + wgpuBufferRelease(uniform_buffer); + if (mask.owned_dummy != nullptr) { + wgpuBufferRelease(mask.owned_dummy); + } + } + + // ---- Dispatch 2: softmax over the last dim (reuse sdpa_softmax.wgsl) ---- + { + SoftmaxParams p = {}; + p.num_rows = static_cast(num_rows); + p.row_width = S_kv; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, &p, sizeof(SoftmaxParams)); + graph.add_uniform_buffer_bytes(sizeof(SoftmaxParams)); + + // sdpa_softmax.wgsl hardcodes @workgroup_size(64,1,1); no override + // constant is needed to decode the near-square 2D grid. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kSdpaSoftmaxWGSL, + { + {0, WGPUBufferBindingType_Storage, softmax_buf, aw_bytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, attn_buf, aw_bytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(SoftmaxParams)}, + }); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, softmax_grid.x, softmax_grid.y); + + wgpuBufferRelease(uniform_buffer); + } + + // ---- Dispatch 3: AV (one thread per (b,h,s,d4) vec4 output element) ---- + { + AvParams p = {}; + p.B = B; + p.H = H; + p.S_q = S_q; + p.S_kv = S_kv; + p.D = D; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, &p, sizeof(AvParams)); + graph.add_uniform_buffer_bytes(sizeof(AvParams)); + + WGPUConstantEntry wg_const = utils::make_wg_size_constant(av_wg_size); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kEtVkSdpaAvWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, softmax_buf, aw_bytes}, + {2, WGPUBufferBindingType_ReadOnlyStorage, v.buffer, v.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(AvParams)}, + }, + &wg_const, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, av_wg_count}); + + wgpuBufferRelease(uniform_buffer); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.sdpa.default, et_vk_sdpa_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl new file mode 100644 index 00000000000..3ccd1d200a2 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av.wgsl @@ -0,0 +1,46 @@ +@group(0) @binding(0) var out: array>; +@group(0) @binding(1) var sm: array; +@group(0) @binding(2) var v: array>; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, AV phase. out[b,h,s,d]=sum_c sm[b,h,s,c]*v[b,h,c,d]. +// DSHB layout, row-major: out/sm rows over S_q, v rows over S_kv; v/out viewed +// as vec4 over D (caller guarantees D % 4 == 0 for every model in scope). +// Supports asymmetric seq (S_q != S_kv); reduces to self-attention when +// S_q == S_kv. ONE thread per (b, h, s, d4) computing 4 output elements; the +// thread contracts over c (0..S_kv, scalar — S_kv isn't guaranteed % 4 == 0). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let d4_count = params.D / 4u; + let total = params.B * params.H * params.S_q * d4_count; + let i = gid.x; + if (i >= total) { + return; + } + let d4 = i % d4_count; + let s = (i / d4_count) % params.S_q; + let h = (i / (d4_count * params.S_q)) % params.H; + let b = i / (d4_count * params.S_q * params.H); + + let smbase = ((b * params.H + h) * params.S_q + s) * params.S_kv; + let vblock4 = (b * params.H + h) * params.S_kv; // first V row of this (b, h) + + var acc: vec4 = vec4(0.0); + for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) { + acc = acc + sm[smbase + c] * v[(vblock4 + c) * d4_count + d4]; + } + out[i] = acc; +} diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h new file mode 100644 index 00000000000..6f625437589 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h @@ -0,0 +1,70 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from et_vk_sdpa_av.wgsl - DO NOT EDIT. +// wgsl-sha256: 0f91462226b666d2fc05381866caed709a6ef3ed2f2562addfe9a23ac81f173b +inline constexpr const char* kEtVkSdpaAvWGSL = R"( +@group(0) @binding(0) var out: array>; +@group(0) @binding(1) var sm: array; +@group(0) @binding(2) var v: array>; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, AV phase. out[b,h,s,d]=sum_c sm[b,h,s,c]*v[b,h,c,d]. +// DSHB layout, row-major: out/sm rows over S_q, v rows over S_kv; v/out viewed +// as vec4 over D (caller guarantees D % 4 == 0 for every model in scope). +// Supports asymmetric seq (S_q != S_kv); reduces to self-attention when +// S_q == S_kv. ONE thread per (b, h, s, d4) computing 4 output elements; the +// thread contracts over c (0..S_kv, scalar — S_kv isn't guaranteed % 4 == 0). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let d4_count = params.D / 4u; + let total = params.B * params.H * params.S_q * d4_count; + let i = gid.x; + if (i >= total) { + return; + } + let d4 = i % d4_count; + let s = (i / d4_count) % params.S_q; + let h = (i / (d4_count * params.S_q)) % params.H; + let b = i / (d4_count * params.S_q * params.H); + + let smbase = ((b * params.H + h) * params.S_q + s) * params.S_kv; + let vblock4 = (b * params.H + h) * params.S_kv; // first V row of this (b, h) + + var acc: vec4 = vec4(0.0); + for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) { + acc = acc + sm[smbase + c] * v[(vblock4 + c) * d4_count + d4]; + } + out[i] = acc; +} +)"; + +inline constexpr uint32_t kEtVkSdpaAvWorkgroupSizeX = 64; +inline constexpr uint32_t kEtVkSdpaAvWorkgroupSizeY = 1; +inline constexpr uint32_t kEtVkSdpaAvWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl new file mode 100644 index 00000000000..f1c4f29ed8c --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk.wgsl @@ -0,0 +1,55 @@ +@group(0) @binding(0) var attn: array; +@group(0) @binding(1) var q: array>; +@group(0) @binding(2) var k: array>; +@group(0) @binding(3) var mask: array; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + has_mask: u32, + _pad0: u32, + scale: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D], +// k [B, H, S_kv, D]; q/k viewed as vec4 over D (caller guarantees +// D % 4 == 0 for every model in scope). Supports asymmetric seq (S_q != S_kv, +// e.g. Hiera pooled query); when S_q == S_kv this reduces to plain +// self-attention. ONE thread per (b, h, s) ROW of attn_weights +// [B, H, S_q, S_kv]; the thread loops over c (0..S_kv) and d4 (0..D/4). Row +// count = B*H*S_q stays well under the 65535 1D dispatch limit for any ViT. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let num_rows = params.B * params.H * params.S_q; + let row = gid.x; + if (row >= num_rows) { + return; + } + let s = row % params.S_q; + let h = (row / params.S_q) % params.H; + let b = row / (params.S_q * params.H); + + let d4_count = params.D / 4u; + let qbase4 = ((b * params.H + h) * params.S_q + s) * d4_count; + let kblock4 = (b * params.H + h) * params.S_kv; // first K row of this (b, h) + let arow = ((b * params.H + h) * params.S_q + s) * params.S_kv; + + for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) { + let kbase4 = (kblock4 + c) * d4_count; + var acc: f32 = 0.0; + for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) { + acc = acc + dot(q[qbase4 + d4], k[kbase4 + d4]); + } + acc = acc * params.scale; + if (params.has_mask != 0u) { + acc = acc + mask[arow + c]; + } + attn[arow + c] = acc; + } +} diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h new file mode 100644 index 00000000000..030e9dec24e --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h @@ -0,0 +1,79 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from et_vk_sdpa_qk.wgsl - DO NOT EDIT. +// wgsl-sha256: cdaa33d5b652395821e1d74222b113e72391b9d618358199abd0772cb821ef40 +inline constexpr const char* kEtVkSdpaQkWGSL = R"( +@group(0) @binding(0) var attn: array; +@group(0) @binding(1) var q: array>; +@group(0) @binding(2) var k: array>; +@group(0) @binding(3) var mask: array; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + has_mask: u32, + _pad0: u32, + scale: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D], +// k [B, H, S_kv, D]; q/k viewed as vec4 over D (caller guarantees +// D % 4 == 0 for every model in scope). Supports asymmetric seq (S_q != S_kv, +// e.g. Hiera pooled query); when S_q == S_kv this reduces to plain +// self-attention. ONE thread per (b, h, s) ROW of attn_weights +// [B, H, S_q, S_kv]; the thread loops over c (0..S_kv) and d4 (0..D/4). Row +// count = B*H*S_q stays well under the 65535 1D dispatch limit for any ViT. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let num_rows = params.B * params.H * params.S_q; + let row = gid.x; + if (row >= num_rows) { + return; + } + let s = row % params.S_q; + let h = (row / params.S_q) % params.H; + let b = row / (params.S_q * params.H); + + let d4_count = params.D / 4u; + let qbase4 = ((b * params.H + h) * params.S_q + s) * d4_count; + let kblock4 = (b * params.H + h) * params.S_kv; // first K row of this (b, h) + let arow = ((b * params.H + h) * params.S_q + s) * params.S_kv; + + for (var c: u32 = 0u; c < params.S_kv; c = c + 1u) { + let kbase4 = (kblock4 + c) * d4_count; + var acc: f32 = 0.0; + for (var d4: u32 = 0u; d4 < d4_count; d4 = d4 + 1u) { + acc = acc + dot(q[qbase4 + d4], k[kbase4 + d4]); + } + acc = acc * params.scale; + if (params.has_mask != 0u) { + acc = acc + mask[arow + c]; + } + attn[arow + c] = acc; + } +} +)"; + +inline constexpr uint32_t kEtVkSdpaQkWorkgroupSizeX = 64; +inline constexpr uint32_t kEtVkSdpaQkWorkgroupSizeY = 1; +inline constexpr uint32_t kEtVkSdpaQkWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu