diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index cdea955fca8..75f89fe17cb 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -57,6 +57,7 @@ set(WEBGPU_SRCS runtime/ops/linear_fp32/LinearFp32.cpp runtime/ops/et_vk_conv2d/Conv2d.cpp runtime/ops/et_vk_sdpa/EtVkSdpa.cpp + runtime/ops/embedding/Embedding.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/embedding/Embedding.cpp b/backends/webgpu/runtime/ops/embedding/Embedding.cpp new file mode 100644 index 00000000000..9b8d670ad66 --- /dev/null +++ b/backends/webgpu/runtime/ops/embedding/Embedding.cpp @@ -0,0 +1,125 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +struct EmbeddingParams { + uint32_t embed_dim; + uint32_t num_elements; + uint32_t _pad[2]; +}; +static_assert( + sizeof(EmbeddingParams) == 16, + "EmbeddingParams must be 16 bytes"); + +// aten.embedding.default args: [weight, indices, padding_idx, +// scale_grad_by_freq, sparse, out]. Forward is a plain row gather +// (padding_idx/scale_grad/sparse only affect the backward pass); out[row, :] = +// weight[indices[row], :]. fp32 weight, int32 indices (the backend's index +// convention; mirrors index.wgsl / embedding_q4gsw). +void embedding_impl(WebGPUGraph& graph, const std::vector& args) { + const int weight_id = args.at(0); + const int indices_id = args.at(1); + const int out_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + const auto& weight = graph.get_tensor(weight_id); + const auto& indices = graph.get_tensor(indices_id); + const auto& out = graph.get_tensor(out_id); + + if (weight.buffer == nullptr || indices.buffer == nullptr || + out.buffer == nullptr) { + throw std::runtime_error("WebGPU embedding: null buffer binding"); + } + if (weight.dims.size() < 2 || out.dims.empty() || indices.dims.empty()) { + throw std::runtime_error("WebGPU embedding: malformed dims"); + } + + const uint32_t embed_dim = static_cast(out.dims.back()); + if (embed_dim == 0) { + throw std::runtime_error("WebGPU embedding: zero embed_dim"); + } + if (static_cast(weight.dims.back()) != embed_dim) { + throw std::runtime_error( + "WebGPU embedding: weight row width != out embed_dim"); + } + + const uint64_t out_numel = utils::numel(out.dims); + const uint32_t num_indices = static_cast(out_numel / embed_dim); + + const uint64_t indices_numel = utils::numel(indices.dims); + // Per-type byte guards (no runtime dtype): indices int32, weight/out fp32. + if (indices_numel != num_indices || + indices.nbytes != indices_numel * sizeof(int32_t)) { + throw std::runtime_error( + "WebGPU embedding: dtype/byte-size mismatch (indices int32, out fp32)"); + } + utils::check_fp32(out, "embedding", "output"); + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kEmbeddingWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, static_cast(out_numel), wg_size, "embedding"); + + EmbeddingParams params = {}; + params.embed_dim = embed_dim; + params.num_elements = static_cast(out_numel); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(EmbeddingParams)); + graph.add_uniform_buffer_bytes(sizeof(EmbeddingParams)); + + WGPUConstantEntry wg_size_constant = utils::make_wg_size_constant(wg_size); + + // out (rw) + indices/weight (ro storage) + uniform. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kEmbeddingWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + indices.buffer, + indices.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(EmbeddingParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.embedding.default, embedding_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/embedding/embedding.wgsl b/backends/webgpu/runtime/ops/embedding/embedding.wgsl new file mode 100644 index 00000000000..6099489a9bf --- /dev/null +++ b/backends/webgpu/runtime/ops/embedding/embedding.wgsl @@ -0,0 +1,28 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_indices: array; +@group(0) @binding(2) var t_weight: array; + +struct Params { + embed_dim: u32, + num_elements: u32, + _pad0: u32, + _pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +// fp32 embedding row-gather: out[row, col] = weight[indices[row], col]. +// One thread per output element. int32 indices (the backend's index convention, +// see index.wgsl / embedding_q4gsw). +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let row = idx / params.embed_dim; + let col = idx % params.embed_dim; + let token = u32(t_indices[row]); + t_out[idx] = t_weight[token * params.embed_dim + col]; +} diff --git a/backends/webgpu/runtime/ops/embedding/embedding_wgsl.h b/backends/webgpu/runtime/ops/embedding/embedding_wgsl.h new file mode 100644 index 00000000000..37ded649581 --- /dev/null +++ b/backends/webgpu/runtime/ops/embedding/embedding_wgsl.h @@ -0,0 +1,52 @@ +/* + * 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 embedding.wgsl - DO NOT EDIT. +// wgsl-sha256: 7ccb183b182070ea17138a4c06ec0e3419893e8312a6769bd945acd7223f9e9c +inline constexpr const char* kEmbeddingWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_indices: array; +@group(0) @binding(2) var t_weight: array; + +struct Params { + embed_dim: u32, + num_elements: u32, + _pad0: u32, + _pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +// fp32 embedding row-gather: out[row, col] = weight[indices[row], col]. +// One thread per output element. int32 indices (the backend's index convention, +// see index.wgsl / embedding_q4gsw). +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let row = idx / params.embed_dim; + let col = idx % params.embed_dim; + let token = u32(t_indices[row]); + t_out[idx] = t_weight[token * params.embed_dim + col]; +} +)"; + +inline constexpr uint32_t kEmbeddingWorkgroupSizeX = 64; +inline constexpr uint32_t kEmbeddingWorkgroupSizeY = 1; +inline constexpr uint32_t kEmbeddingWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu