From 6c4def0e887d7b611b4db9b15d74d5db3c6d32e6 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:48:48 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 2 + .../runtime/ops/embedding/Embedding.cpp | 156 +++++++++++++++ .../runtime/ops/embedding/embedding.wgsl | 25 +++ .../runtime/ops/embedding/embedding_wgsl.h | 49 +++++ backends/webgpu/runtime/ops/gather/Gather.cpp | 182 ++++++++++++++++++ .../webgpu/runtime/ops/gather/gather.wgsl | 39 ++++ .../webgpu/runtime/ops/gather/gather_wgsl.h | 63 ++++++ 7 files changed, 516 insertions(+) create mode 100644 backends/webgpu/runtime/ops/embedding/Embedding.cpp create mode 100644 backends/webgpu/runtime/ops/embedding/embedding.wgsl create mode 100644 backends/webgpu/runtime/ops/embedding/embedding_wgsl.h create mode 100644 backends/webgpu/runtime/ops/gather/Gather.cpp create mode 100644 backends/webgpu/runtime/ops/gather/gather.wgsl create mode 100644 backends/webgpu/runtime/ops/gather/gather_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index dd644600fe8..f92ed119b4e 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -61,7 +61,9 @@ set(WEBGPU_SRCS runtime/ops/sub/BinaryOp.cpp runtime/ops/where/Where.cpp runtime/ops/compare/Compare.cpp + runtime/ops/gather/Gather.cpp runtime/ops/linear/Linear.cpp + runtime/ops/embedding/Embedding.cpp runtime/ops/logical_not/LogicalNot.cpp ) diff --git a/backends/webgpu/runtime/ops/embedding/Embedding.cpp b/backends/webgpu/runtime/ops/embedding/Embedding.cpp new file mode 100644 index 00000000000..a8ebee9263c --- /dev/null +++ b/backends/webgpu/runtime/ops/embedding/Embedding.cpp @@ -0,0 +1,156 @@ +/* + * 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 num_elements; + uint32_t dim; + uint32_t _pad[2]; +}; + +// aten.embedding: out[row, :] = weight[indices[row], :] (fp32 weight, i32 idx). +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("embedding: null buffer binding"); + } + if (weight.dims.size() != 2) { + throw std::runtime_error("embedding: weight must be 2D [vocab, dim]"); + } + const uint32_t dim = static_cast(weight.dims[1]); + if (dim == 0) { + throw std::runtime_error("embedding: dim == 0"); + } + const size_t out_numel = out.nbytes / sizeof(float); + // Index is the int32 downcast of the int64 indices (mirror index op). + const size_t index_numel = indices.nbytes / sizeof(int32_t); + if (out.nbytes != out_numel * sizeof(float) || + weight.nbytes % sizeof(float) != 0 || + indices.nbytes != index_numel * sizeof(int32_t)) { + throw std::runtime_error("embedding: fp32 weight/out + i32 indices required"); + } + if (out_numel != index_numel * dim) { + throw std::runtime_error("embedding: out numel != num_indices * dim"); + } + + uint32_t num_elements = static_cast(out_numel); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kEmbeddingWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_elements, wg_size, "embedding"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + EmbeddingParams params = {}; + params.num_elements = num_elements; + params.dim = dim; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(EmbeddingParams)); + graph.add_uniform_buffer_bytes(sizeof(EmbeddingParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kEmbeddingWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Storage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[4] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = weight.buffer; + bg_entries[0].size = weight.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = indices.buffer; + bg_entries[1].size = indices.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = out.buffer; + bg_entries[2].size = out.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = uniform_buffer; + bg_entries[3].size = sizeof(EmbeddingParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + 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..9fece09d293 --- /dev/null +++ b/backends/webgpu/runtime/ops/embedding/embedding.wgsl @@ -0,0 +1,25 @@ +@group(0) @binding(0) var weight: array; +@group(0) @binding(1) var indices: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, + dim: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + let row = idx / params.dim; + let col = idx % params.dim; + let row_id = u32(indices[row]); + output[idx] = weight[row_id * params.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..4a790c74d65 --- /dev/null +++ b/backends/webgpu/runtime/ops/embedding/embedding_wgsl.h @@ -0,0 +1,49 @@ +/* + * 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: e7c94da588a9dd40c0e1be7d3c6ef33310c71390c57f6f330ae23a18866d0ea7 +inline constexpr const char* kEmbeddingWGSL = R"( +@group(0) @binding(0) var weight: array; +@group(0) @binding(1) var indices: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, + dim: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + let row = idx / params.dim; + let col = idx % params.dim; + let row_id = u32(indices[row]); + output[idx] = weight[row_id * params.dim + col]; +} +)"; + +inline constexpr uint32_t kEmbeddingWorkgroupSizeX = 256; +inline constexpr uint32_t kEmbeddingWorkgroupSizeY = 1; +inline constexpr uint32_t kEmbeddingWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/gather/Gather.cpp b/backends/webgpu/runtime/ops/gather/Gather.cpp new file mode 100644 index 00000000000..d7e170129d1 --- /dev/null +++ b/backends/webgpu/runtime/ops/gather/Gather.cpp @@ -0,0 +1,182 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +struct GatherParams { + uint32_t dim; + uint32_t _pad[3]; +}; + +// gather: out[c] = self[c] with c[dim] replaced by index[c]. +void gather_impl(WebGPUGraph& graph, const std::vector& args) { + const int self_id = args.at(0); + const int dim_id = args.at(1); + const int index_id = args.at(2); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int) { + throw std::runtime_error("gather: dim is not an int"); + } + WGPUDevice device = graph.device(); + const auto& self_tensor = graph.get_tensor(self_id); + const auto& index_tensor = graph.get_tensor(index_id); + const auto& out_tensor = graph.get_tensor(out_id); + + if (self_tensor.buffer == nullptr || index_tensor.buffer == nullptr || + out_tensor.buffer == nullptr) { + throw std::runtime_error("gather: null buffer binding"); + } + + const int64_t ndim = static_cast(out_tensor.dims.size()); + int64_t dim = graph.get_int(dim_id); + if (dim < 0) { + dim += ndim; + } + if (ndim == 0 || dim < 0 || dim >= ndim) { + throw std::runtime_error("gather: dim out of range"); + } + + TensorMeta out_meta; + TensorMeta self_meta; + fill_tensor_meta(out_tensor, &out_meta); + fill_tensor_meta(self_tensor, &self_meta); + if (out_meta.ndim != self_meta.ndim) { + throw std::runtime_error("gather: self/out rank mismatch"); + } + + const size_t out_numel = out_tensor.nbytes / sizeof(float); + const size_t index_numel = index_tensor.nbytes / sizeof(int32_t); + if (out_tensor.nbytes != out_numel * sizeof(float) || + self_tensor.nbytes % sizeof(float) != 0 || + index_tensor.nbytes != index_numel * sizeof(int32_t)) { + throw std::runtime_error("gather: fp32 self/out + i32 index required"); + } + if (out_numel != index_numel) { + throw std::runtime_error("gather: out numel != index numel"); + } + + uint32_t wg_size = utils::clamp_workgroup_size(device, kGatherWorkgroupSizeX); + uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, out_meta.numel, wg_size, "gather"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + GatherParams params = {}; + params.dim = static_cast(dim); + + WGPUBuffer out_meta_buf = + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); + WGPUBuffer self_meta_buf = + utils::make_uniform(device, &self_meta, sizeof(TensorMeta)); + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(GatherParams)); + graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(GatherParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kGatherWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[6] = {}; + entries[0].binding = 0; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[2].binding = 2; + entries[2].buffer.type = WGPUBufferBindingType_Storage; + entries[3].binding = 3; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + entries[4].binding = 4; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + entries[5].binding = 5; + entries[5].buffer.type = WGPUBufferBindingType_Uniform; + for (auto& e : entries) { + e.visibility = WGPUShaderStage_Compute; + } + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 6; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[6] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = self_tensor.buffer; + bg_entries[0].size = self_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = index_tensor.buffer; + bg_entries[1].size = index_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = out_tensor.buffer; + bg_entries[2].size = out_tensor.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = out_meta_buf; + bg_entries[3].size = sizeof(TensorMeta); + bg_entries[4].binding = 4; + bg_entries[4].buffer = self_meta_buf; + bg_entries[4].size = sizeof(TensorMeta); + bg_entries[5].binding = 5; + bg_entries[5].buffer = params_buf; + bg_entries[5].size = sizeof(GatherParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 6; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch({pipeline, bind_group, workgroup_count}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + wgpuBufferRelease(out_meta_buf); + wgpuBufferRelease(self_meta_buf); + wgpuBufferRelease(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.gather.default, gather_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/gather/gather.wgsl b/backends/webgpu/runtime/ops/gather/gather.wgsl new file mode 100644 index 00000000000..fb62206be64 --- /dev/null +++ b/backends/webgpu/runtime/ops/gather/gather.wgsl @@ -0,0 +1,39 @@ +@group(0) @binding(0) var self_: array; +@group(0) @binding(1) var indices: array; +@group(0) @binding(2) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var self_meta: TensorMeta; + +struct GatherParams { + dim: u32, +} +@group(0) @binding(5) var params: GatherParams; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let o = gid.x; + if (o >= out_meta.numel) { + return; + } + var rem = o; + var self_idx: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let c = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + var coord = c; + if (d == params.dim) { + coord = u32(indices[o]); + } + self_idx = self_idx + coord * self_meta.strides[d]; + } + output[o] = self_[self_idx]; +} diff --git a/backends/webgpu/runtime/ops/gather/gather_wgsl.h b/backends/webgpu/runtime/ops/gather/gather_wgsl.h new file mode 100644 index 00000000000..5b6922e21ea --- /dev/null +++ b/backends/webgpu/runtime/ops/gather/gather_wgsl.h @@ -0,0 +1,63 @@ +/* + * 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 gather.wgsl - DO NOT EDIT. +// wgsl-sha256: cd19a9ed2753e2a97e96fb90f7a72e8f40d08feb6c84fbed7f2c52e71d77a03c +inline constexpr const char* kGatherWGSL = R"( +@group(0) @binding(0) var self_: array; +@group(0) @binding(1) var indices: array; +@group(0) @binding(2) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var self_meta: TensorMeta; + +struct GatherParams { + dim: u32, +} +@group(0) @binding(5) var params: GatherParams; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let o = gid.x; + if (o >= out_meta.numel) { + return; + } + var rem = o; + var self_idx: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let c = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + var coord = c; + if (d == params.dim) { + coord = u32(indices[o]); + } + self_idx = self_idx + coord * self_meta.strides[d]; + } + output[o] = self_[self_idx]; +} +)"; + +inline constexpr uint32_t kGatherWorkgroupSizeX = 256; +inline constexpr uint32_t kGatherWorkgroupSizeY = 1; +inline constexpr uint32_t kGatherWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu