[ExecuTorch][WebGPU] Port conv1d pointwise (aten.convolution)

backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp

@@ -10,6 +10,7 @@
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 #include <executorch/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw_wgsl.h>
 
 #include <webgpu/webgpu.h>
 
@@ -55,6 +56,191 @@ int64_t first_int(const std::vector<int64_t>& 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<uint32_t>(in_tensor.dims.at(1));
+  const uint32_t out_channels = static_cast<uint32_t>(out_tensor.dims.at(1));
+  const uint32_t length = static_cast<uint32_t>(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<uint64_t>(d);
+  }
+  if (in_tensor.nbytes % sizeof(float) != 0 ||
+      out_tensor.nbytes != out_numel * sizeof(float) ||
+      weight_tensor.nbytes !=
+          static_cast<size_t>(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<uint32_t>(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<uint32_t>(out_numel), wg_size, "conv1d_pw");
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUBuffer params_buf =
+      utils::make_uniform(device, &params, 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<uint32_t>(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<int64_t>(in_channels)) {
+          throw std::runtime_error(
+              "conv1d_pw(resize): in channel count changed");
+        }
+        Conv1dPwParams p = {};
+        p.in_channels = static_cast<uint32_t>(d[1]);
+        p.out_channels = out_channels;
+        p.length = static_cast<uint32_t>(d[2]);
+        p.numel = static_cast<uint32_t>(
+            static_cast<uint64_t>(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<int64_t> 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<int>& args) {
   // args mirror Vulkan conv1d_dw; bias (arg 2) may be Null; out=args.back().
@@ -91,14 +277,23 @@ void convolution_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const uint32_t out_len = static_cast<uint32_t>(out_tensor.dims.at(2));
   const uint32_t kernel_size = static_cast<uint32_t>(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<int64_t>(channels) ||
       weight_tensor.dims.at(0) != static_cast<int64_t>(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));

backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw.wgsl

@@ -0,0 +1,45 @@
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+@group(0) @binding(2) var<storage, read> weight: array<f32>;
+@group(0) @binding(3) var<storage, read> bias: array<f32>;
+
+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<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 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;
+}

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 <cstdint>
+
+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<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+@group(0) @binding(2) var<storage, read> weight: array<f32>;
+@group(0) @binding(3) var<storage, read> bias: array<f32>;
+
+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<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 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