[ExecuTorch][WebGPU] Add et_vk.adamw_step (on-GPU AdamW optimizer step) to the WebGPU backend

backends/vulkan/custom_ops_lib.py

@@ -1203,6 +1203,58 @@ def fused_ce_backward(ctx, grad_loss, grad_dlogits):
 fused_ce_op = getattr(getattr(torch.ops, namespace), name)
 
 
+###########################
+## adamw_step (training) ##
+###########################
+
+
+def adamw_step_impl(
+    param: torch.Tensor,
+    m: torch.Tensor,
+    v: torch.Tensor,
+    grad: torch.Tensor,
+    lr: float,
+    beta1: float,
+    beta2: float,
+    eps: float,
+    weight_decay: float,
+    bias_correction1: float,
+    bias_correction2: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    param.mul_(1.0 - lr * weight_decay)
+    m.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+    v.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+    mhat = m / bias_correction1
+    denom = (v / bias_correction2).sqrt() + eps
+    param.addcdiv_(mhat, denom, value=-lr)
+    return param, m, v
+
+
+def adamw_step_meta(
+    param: torch.Tensor,
+    m: torch.Tensor,
+    v: torch.Tensor,
+    grad: torch.Tensor,
+    lr: float,
+    beta1: float,
+    beta2: float,
+    eps: float,
+    weight_decay: float,
+    bias_correction1: float,
+    bias_correction2: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    return param, m, v
+
+
+name = "adamw_step"
+lib.define(
+    f"{name}(Tensor(a!) param, Tensor(b!) m, Tensor(c!) v, Tensor grad, float lr, float beta1, float beta2, float eps, float weight_decay, float bias_correction1, float bias_correction2) -> (Tensor(a!), Tensor(b!), Tensor(c!))"
+)
+lib.impl(name, adamw_step_impl, "CompositeExplicitAutograd")
+lib.impl(name, adamw_step_meta, "Meta")
+adamw_step_op = getattr(getattr(torch.ops, namespace), name)
+
+
 # STE weight gradient d_out^T @ x through the frozen 4-bit linear_q4gsw base.
 def linear_q4gsw_dw_impl(
     d_out: torch.Tensor,

backends/vulkan/op_registry.py

@@ -1797,6 +1797,14 @@ def register_logical_not():
     )
 
 
+@update_features("et_vk::adamw_step")
+def register_adamw_step():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=utils.FP_T,
+    )
+
+
 @update_features(exir_ops.edge.et_vk.linear_q4gsw_dw.default)
 def register_linear_q4gsw_dw():
     return OpFeatures(

backends/webgpu/CMakeLists.txt

@@ -70,6 +70,7 @@ set(WEBGPU_SRCS
     runtime/ops/linear/Linear.cpp
     runtime/ops/embedding/Embedding.cpp
     runtime/ops/logical_not/LogicalNot.cpp
+    runtime/ops/adamw/AdamwStep.cpp
     runtime/ops/quantized_linear/QuantizedLinearDw.cpp
     runtime/ops/quantized_linear/QuantizedLinearRequant.cpp
 )

backends/webgpu/runtime/ops/adamw/AdamwStep.cpp

@@ -0,0 +1,181 @@
+/*
+ * 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 <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct AdamwStepParams {
+  uint32_t numel;
+  uint32_t _pad0;
+  uint32_t _pad1;
+  uint32_t _pad2;
+  float lr;
+  float beta1;
+  float beta2;
+  float eps;
+  float weight_decay;
+  float bias_correction1;
+  float bias_correction2;
+  float _pad3;
+};
+static_assert(sizeof(AdamwStepParams) == 48, "params must be 48 bytes");
+
+// AdamW step over an fp32 latent (elementwise, in place); mirrors torch AdamW.
+void adamw_step_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int param_id = args.at(0);
+  const int m_id = args.at(1);
+  const int v_id = args.at(2);
+  const int grad_id = args.at(3);
+
+  WGPUDevice device = graph.device();
+  const auto& param = graph.get_tensor(param_id);
+  const auto& m = graph.get_tensor(m_id);
+  const auto& v = graph.get_tensor(v_id);
+  const auto& grad = graph.get_tensor(grad_id);
+
+  uint64_t numel = 1;
+  for (int64_t d : param.dims) {
+    numel *= static_cast<uint64_t>(d);
+  }
+  if (param.dims.empty() || numel == 0) {
+    throw std::runtime_error("adamw_step: empty param");
+  }
+  const uint64_t bytes = numel * sizeof(float);
+  if (param.nbytes != bytes || m.nbytes != bytes || v.nbytes != bytes ||
+      grad.nbytes != bytes) {
+    throw std::runtime_error(
+        "adamw_step: param/m/v/grad must be fp32 and same numel");
+  }
+  if (numel > UINT32_MAX) {
+    throw std::runtime_error("adamw_step: numel exceeds u32");
+  }
+
+  auto scalar = [&](int id, const char* name) -> float {
+    if (graph.get_value_type(id) != WebGPUGraph::ValueType::Double) {
+      throw std::runtime_error(
+          std::string("adamw_step: ") + name + " must be a float scalar");
+    }
+    return static_cast<float>(graph.get_double(id));
+  };
+
+  AdamwStepParams params = {};
+  params.numel = static_cast<uint32_t>(numel);
+  params.lr = scalar(args.at(4), "lr");
+  params.beta1 = scalar(args.at(5), "beta1");
+  params.beta2 = scalar(args.at(6), "beta2");
+  params.eps = scalar(args.at(7), "eps");
+  params.weight_decay = scalar(args.at(8), "weight_decay");
+  params.bias_correction1 = scalar(args.at(9), "bias_correction1");
+  params.bias_correction2 = scalar(args.at(10), "bias_correction2");
+  if (params.bias_correction1 == 0.0f || params.bias_correction2 == 0.0f) {
+    throw std::runtime_error("adamw_step: bias corrections must be non-zero");
+  }
+
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kAdamwStepWorkgroupSizeX);
+  const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, params.numel, wg_size, "adamw_step");
+
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, sizeof(params));
+  graph.add_uniform_buffer_bytes(sizeof(params));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kAdamwStepWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[5] = {};
+  for (uint32_t i = 0; i <= 2; i++) {
+    entries[i].binding = i;
+    entries[i].visibility = WGPUShaderStage_Compute;
+    entries[i].buffer.type = WGPUBufferBindingType_Storage;
+  }
+  entries[3].binding = 3;
+  entries[3].visibility = WGPUShaderStage_Compute;
+  entries[3].buffer.type = 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);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  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[5] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = param.buffer;
+  bg_entries[0].size = param.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = m.buffer;
+  bg_entries[1].size = m.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = v.buffer;
+  bg_entries[2].size = v.nbytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = grad.buffer;
+  bg_entries[3].size = grad.nbytes;
+  bg_entries[4].binding = 4;
+  bg_entries[4].buffer = uniform_buffer;
+  bg_entries[4].size = sizeof(params);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 5;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  graph.add_dispatch({pipeline, bind_group, workgroup_count, "adamw_step"});
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  graph.own_uniform_buffer(uniform_buffer);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(et_vk.adamw_step.default, adamw_step_impl);
+}
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/ops/adamw/adamw_step.wgsl

@@ -0,0 +1,41 @@
+
+@group(0) @binding(0) var<storage, read_write> t_param: array<f32>;
+@group(0) @binding(1) var<storage, read_write> t_m: array<f32>;
+@group(0) @binding(2) var<storage, read_write> t_v: array<f32>;
+@group(0) @binding(3) var<storage, read> t_grad: array<f32>;
+
+struct Params {
+  numel: u32,
+  _pad0: u32,
+  _pad1: u32,
+  _pad2: u32,
+  lr: f32,
+  beta1: f32,
+  beta2: f32,
+  eps: f32,
+  weight_decay: f32,
+  bias_correction1: f32,
+  bias_correction2: f32,
+  _pad3: f32,
+}
+@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>) {
+  let i = gid.x;
+  if (i >= params.numel) {
+    return;
+  }
+  let g = t_grad[i];
+  var p = t_param[i];
+  p = p - params.lr * params.weight_decay * p;
+  let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g;
+  let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g;
+  t_m[i] = m;
+  t_v[i] = v;
+  let mhat = m / params.bias_correction1;
+  let vhat = v / params.bias_correction2;
+  t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps);
+}

backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h

@@ -0,0 +1,65 @@
+/*
+ * 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 adamw_step.wgsl - DO NOT EDIT.
+// wgsl-sha256: 0957c04168872db5e2b39cf5f26beefba27f3b5514ec69cece4de5145e97156f
+inline constexpr const char* kAdamwStepWGSL = R"(
+
+@group(0) @binding(0) var<storage, read_write> t_param: array<f32>;
+@group(0) @binding(1) var<storage, read_write> t_m: array<f32>;
+@group(0) @binding(2) var<storage, read_write> t_v: array<f32>;
+@group(0) @binding(3) var<storage, read> t_grad: array<f32>;
+
+struct Params {
+  numel: u32,
+  _pad0: u32,
+  _pad1: u32,
+  _pad2: u32,
+  lr: f32,
+  beta1: f32,
+  beta2: f32,
+  eps: f32,
+  weight_decay: f32,
+  bias_correction1: f32,
+  bias_correction2: f32,
+  _pad3: f32,
+}
+@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>) {
+  let i = gid.x;
+  if (i >= params.numel) {
+    return;
+  }
+  let g = t_grad[i];
+  var p = t_param[i];
+  p = p - params.lr * params.weight_decay * p;
+  let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g;
+  let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g;
+  t_m[i] = m;
+  t_v[i] = v;
+  let mhat = m / params.bias_correction1;
+  let vhat = v / params.bias_correction2;
+  t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps);
+}
+)";
+
+inline constexpr uint32_t kAdamwStepWorkgroupSizeX = 64;
+inline constexpr uint32_t kAdamwStepWorkgroupSizeY = 1;
+inline constexpr uint32_t kAdamwStepWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu