diff --git a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp index 84cb8cc5b7d..f876b095a22 100644 --- a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp +++ b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp @@ -9,7 +9,9 @@ #include #include #include +#include #include +#include #include @@ -55,6 +57,191 @@ int64_t first_int(const std::vector& v) { return v.empty() ? 0 : v[0]; } +struct Conv1dPwParams { + uint32_t in_channels; + uint32_t out_channels; + uint32_t length; + uint32_t numel; + uint32_t has_bias; + uint32_t pad0; + uint32_t pad1; + uint32_t pad2; +}; +static_assert( + sizeof(Conv1dPwParams) == 32, + "Conv1dPwParams must match the WGSL Params struct (32 bytes)"); + +// Pointwise conv1d (K=1, groups=1): a per-position matmul over channels. +void add_conv1d_pw_node( + WebGPUGraph& graph, + int in_id, + int weight_id, + int bias_id, + int out_id) { + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& out_tensor = graph.get_tensor(out_id); + + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + const uint32_t in_channels = static_cast(in_tensor.dims.at(1)); + const uint32_t out_channels = static_cast(out_tensor.dims.at(1)); + const uint32_t length = static_cast(out_tensor.dims.at(2)); + if (in_tensor.dims.at(2) != out_tensor.dims.at(2)) { + throw std::runtime_error("conv1d_pw: in/out length (dim 2) mismatch"); + } + + uint64_t out_numel = 1; + for (int64_t d : out_tensor.dims) { + out_numel *= static_cast(d); + } + if (in_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes != out_numel * sizeof(float) || + weight_tensor.nbytes != + static_cast(out_channels) * in_channels * sizeof(float)) { + throw std::runtime_error("conv1d_pw: fp32-only (byte-size mismatch)"); + } + if (out_numel > UINT32_MAX) { + throw std::runtime_error("conv1d_pw: output numel exceeds u32"); + } + + Conv1dPwParams params = {}; + params.in_channels = in_channels; + params.out_channels = out_channels; + params.length = length; + params.numel = static_cast(out_numel); + params.has_bias = has_bias ? 1u : 0u; + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kConv1dPwWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(out_numel), wg_size, "conv1d_pw"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Conv1dPwParams)); + graph.add_uniform_buffer_bytes(sizeof(Conv1dPwParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kConv1dPwWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[5] = {}; + for (int i = 0; i < 4; i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = (i == 1) ? WGPUBufferBindingType_Storage + : WGPUBufferBindingType_ReadOnlyStorage; + } + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 5; + 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); + + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : weight_tensor.buffer; + uint64_t bias_sz = + has_bias ? graph.get_tensor(bias_id).nbytes : weight_tensor.nbytes; + + WGPUBindGroupEntry bg_entries[5] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = weight_tensor.buffer; + bg_entries[2].size = weight_tensor.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = bias_buf; + bg_entries[3].size = bias_sz; + bg_entries[4].binding = 4; + bg_entries[4].buffer = params_buf; + bg_entries[4].size = sizeof(Conv1dPwParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 5; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "conv1d_pw", + workgroup_count.y}); + + // Dynamic shapes: only the length varies; recompute params + dispatch. + WGPUBuffer p_buf = params_buf; + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + in_channels, + out_channels, + has_bias, + wg_size, + dispatch_idx, + p_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.size() != 3) { + throw std::runtime_error("conv1d_pw(resize): input is not 3D"); + } + if (d[1] != static_cast(in_channels)) { + throw std::runtime_error( + "conv1d_pw(resize): in channel count changed"); + } + Conv1dPwParams p = {}; + p.in_channels = static_cast(d[1]); + p.out_channels = out_channels; + p.length = static_cast(d[2]); + p.numel = static_cast( + static_cast(d[0]) * out_channels * d[2]); + p.has_bias = has_bias ? 1u : 0u; + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "conv1d_pw(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + const std::vector out_d = {d[0], out_channels, d[2]}; + g.set_cur_dims(out_id, out_d); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(params_buf); +} + // depthwise-conv1d (groups==C); mirrors Vulkan conv1d_dw (Convolution.cpp:755). void convolution_impl(WebGPUGraph& graph, const std::vector& args) { // args mirror Vulkan conv1d_dw; bias (arg 2) may be Null; out=args.back(). @@ -91,14 +278,23 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { const uint32_t out_len = static_cast(out_tensor.dims.at(2)); const uint32_t kernel_size = static_cast(weight_tensor.dims.at(2)); - // Only the depthwise config (groups==C, weight [C,1,K], not transposed). const bool transposed = graph.get_bool(transposed_id); const int64_t groups = graph.get_int(groups_id); + + // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. + if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && + first_int(graph.get_int_list(stride_id)) == 1 && + first_int(graph.get_int_list(padding_id)) == 0) { + add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); + return; + } + + // Otherwise only the depthwise config (groups==C, weight [C,1,K]). if (transposed || groups != static_cast(channels) || weight_tensor.dims.at(0) != static_cast(channels) || weight_tensor.dims.at(1) != 1) { throw std::runtime_error( - "convolution: only depthwise conv1d (groups==C, weight [C,1,K])"); + "convolution: only depthwise or pointwise conv1d supported"); } const int64_t stride_i = first_int(graph.get_int_list(stride_id)); @@ -297,8 +493,10 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { } // namespace -WEBGPU_REGISTER_OPERATORS { - WEBGPU_REGISTER_OP(aten.convolution.default, convolution_impl); +// Single aten.convolution.default registration lives in Conv2d.cpp, which +// dispatches on input rank; this is the 3D (conv1d) entry it calls. +void conv1d_dispatch(WebGPUGraph& graph, const std::vector& args) { + convolution_impl(graph, args); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.h b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.h new file mode 100644 index 00000000000..bcf7a0c8109 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.h @@ -0,0 +1,23 @@ +/* + * 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 + +#include + +namespace executorch::backends::webgpu { + +// conv1d (3D) entry: pointwise (K=1, groups=1) or depthwise (groups==C). Called +// by the single aten.convolution.default handler in Conv2d.cpp, which +// dispatches on input rank (3D -> conv1d, 4D -> conv2d) so one registration +// serves both. +void conv1d_dispatch(WebGPUGraph& graph, const std::vector& args); + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw.wgsl b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw.wgsl new file mode 100644 index 00000000000..66260ec78ad --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw.wgsl @@ -0,0 +1,45 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + in_channels: u32, + out_channels: u32, + length: u32, + numel: u32, + has_bias: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // Pointwise (K=1): out[n,oc,l] = sum_ic weight[oc,ic] * input[n,ic,l]. + let l = oi % params.length; + let oc = (oi / params.length) % params.out_channels; + let n = oi / (params.length * params.out_channels); + + let w_base = oc * params.in_channels; + let in_base = n * params.in_channels * params.length + l; + var s = 0.0; + for (var ic: u32 = 0u; ic < params.in_channels; ic = ic + 1u) { + s = s + weight[w_base + ic] * input[in_base + ic * params.length]; + } + if (params.has_bias != 0u) { + s = s + bias[oc]; + } + output[oi] = s; +} diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw_wgsl.h b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw_wgsl.h new file mode 100644 index 00000000000..854a7160e95 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw_wgsl.h @@ -0,0 +1,69 @@ +/* + * 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 conv1d_pw.wgsl - DO NOT EDIT. +// wgsl-sha256: d47893e191276cf99ee6fd30002ff70f4a0656bcc8e8f90eaca95d9916b54448 +inline constexpr const char* kConv1dPwWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + in_channels: u32, + out_channels: u32, + length: u32, + numel: u32, + has_bias: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // Pointwise (K=1): out[n,oc,l] = sum_ic weight[oc,ic] * input[n,ic,l]. + let l = oi % params.length; + let oc = (oi / params.length) % params.out_channels; + let n = oi / (params.length * params.out_channels); + + let w_base = oc * params.in_channels; + let in_base = n * params.in_channels * params.length + l; + var s = 0.0; + for (var ic: u32 = 0u; ic < params.in_channels; ic = ic + 1u) { + s = s + weight[w_base + ic] * input[in_base + ic * params.length]; + } + if (params.has_bias != 0u) { + s = s + bias[oc]; + } + output[oi] = s; +} +)"; + +inline constexpr uint32_t kConv1dPwWorkgroupSizeX = 64; +inline constexpr uint32_t kConv1dPwWorkgroupSizeY = 1; +inline constexpr uint32_t kConv1dPwWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp index 5f127f6d5f2..3c9d49903bd 100644 --- a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -236,6 +237,13 @@ void conv2d_impl(WebGPUGraph& graph, const std::vector& args) { const int groups_id = args.at(8); const int out_id = args.at(9); + // aten.convolution.default covers conv1d (3D) and conv2d (4D); the registry + // is first-wins, so one handler must serve both ranks. Route 3D to conv1d. + if (graph.get_tensor(in_id).dims.size() == 3) { + conv1d_dispatch(graph, args); + return; + } + WGPUDevice device = graph.device(); if (graph.get_bool(transposed_id)) {