[ExecuTorch][WebGPU] Add softmax ops (_softmax, _log_softmax) to the WebGPU backend

Pull Request resolved: #20919

Add `aten._softmax` + `aten._log_softmax` handlers (max-stable online softmax) for the training tail.

Key changes:
- `runtime/ops/{softmax,log_softmax}/` — max-stable softmax + log_softmax WGSL kernels + handlers
- `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources

Reuses the shared Vulkan partitioner (`aten._softmax`/`_log_softmax` already registered); WebGPU kernels only.

Co-authored-with: Claude Code.
ghstack-source-id: 405026047
@exported-using-ghexport

Differential Revision: [D111755118](https://our.internmc.facebook.com/intern/diff/D111755118/)

Author: JCNTH

backends/webgpu/CMakeLists.txt

@@ -53,6 +53,8 @@ set(WEBGPU_SRCS
     runtime/ops/index/Index.cpp
     runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp
     runtime/ops/mm/Mm.cpp
+    runtime/ops/log_softmax/LogSoftmax.cpp
+    runtime/ops/softmax/Softmax.cpp
     runtime/ops/bmm/Bmm.cpp
     runtime/ops/linear/Linear.cpp
 )

backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp

@@ -0,0 +1,231 @@
+/*
+ * 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/softmax/log_softmax_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// Uniform layout matching the WGSL Params struct (16-byte aligned).
+struct LogSoftmaxParams {
+  uint32_t outer;
+  uint32_t r;
+  uint32_t inner;
+  uint32_t pad_;
+};
+static_assert(
+    sizeof(LogSoftmaxParams) == 16,
+    "LogSoftmaxParams must be 16 bytes");
+
+// Decompose dims into [outer, R, inner] for a reduction along `dim`.
+void decompose(
+    const std::vector<int64_t>& dims,
+    int64_t dim,
+    uint32_t& outer,
+    uint32_t& r,
+    uint32_t& inner) {
+  const int64_t ndim = static_cast<int64_t>(dims.size());
+  if (dim < 0) {
+    dim += ndim;
+  }
+  if (ndim == 0 || dim < 0 || dim >= ndim) {
+    throw std::runtime_error("WebGPU log_softmax: dim out of range");
+  }
+  uint64_t o = 1, in = 1;
+  for (int64_t d = 0; d < dim; ++d) {
+    o *= static_cast<uint64_t>(dims[d]);
+  }
+  for (int64_t d = dim + 1; d < ndim; ++d) {
+    in *= static_cast<uint64_t>(dims[d]);
+  }
+  outer = static_cast<uint32_t>(o);
+  r = static_cast<uint32_t>(dims[dim]);
+  inner = static_cast<uint32_t>(in);
+}
+
+void log_softmax_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int in_id = args.at(0);
+  const int dim_id = args.at(1);
+  const int out_id = args.at(3);
+
+  WGPUDevice device = graph.device();
+
+  const auto& in = graph.get_tensor(in_id);
+  const auto& out = graph.get_tensor(out_id);
+
+  if (in.dims.empty()) {
+    throw std::runtime_error("WebGPU log_softmax: scalar input unsupported");
+  }
+  if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int) {
+    throw std::runtime_error("WebGPU log_softmax: dim must be an int");
+  }
+  const int64_t dim = graph.get_int(dim_id);
+
+  uint32_t outer = 0, r = 0, inner = 0;
+  decompose(in.dims, dim, outer, r, inner);
+  if (outer == 0 || r == 0 || inner == 0) {
+    throw std::runtime_error("WebGPU log_softmax: zero-sized reduction");
+  }
+
+  uint64_t numel = 1;
+  for (int64_t d : in.dims) {
+    numel *= static_cast<uint64_t>(d);
+  }
+  if (in.nbytes != numel * sizeof(float) ||
+      out.nbytes != numel * sizeof(float)) {
+    throw std::runtime_error(
+        "WebGPU log_softmax: fp32-only (byte-size mismatch)");
+  }
+
+  const uint64_t lines = static_cast<uint64_t>(outer) * inner;
+  if (lines > UINT32_MAX) {
+    throw std::runtime_error(
+        "WebGPU log_softmax: line count exceeds dispatch limit");
+  }
+
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kLogSoftmaxWorkgroupSizeX);
+  const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, static_cast<uint32_t>(lines), wg_size, "log_softmax");
+
+  LogSoftmaxParams params = {};
+  params.outer = outer;
+  params.r = r;
+  params.inner = inner;
+
+  WGPUBufferDescriptor uniform_desc = {};
+  uniform_desc.size = sizeof(LogSoftmaxParams);
+  uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
+  uniform_desc.mappedAtCreation = true;
+  WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc);
+  void* mapped =
+      wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(LogSoftmaxParams));
+  std::memcpy(mapped, &params, sizeof(LogSoftmaxParams));
+  wgpuBufferUnmap(uniform_buffer);
+  graph.add_uniform_buffer_bytes(sizeof(LogSoftmaxParams));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kLogSoftmaxWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[3] = {};
+  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_Storage;
+  entries[2].binding = 2;
+  entries[2].visibility = WGPUShaderStage_Compute;
+  entries[2].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 3;
+  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[3] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = in.buffer;
+  bg_entries[0].size = in.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = out.buffer;
+  bg_entries[1].size = out.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = uniform_buffer;
+  bg_entries[2].size = sizeof(LogSoftmaxParams);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 3;
+  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, "log_softmax"});
+
+  // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch.
+  WGPUBuffer params_buf = uniform_buffer;
+  const uint64_t build_numel = static_cast<uint64_t>(outer) * r * inner;
+  graph.add_tensor_resize_hook(
+      in_id,
+      [in_id, out_id, dim, build_numel, wg_size, dispatch_idx, params_buf](
+          WebGPUGraph& g) {
+        const auto& d = g.cur_dims(in_id);
+        uint32_t o = 0, rr = 0, n = 0;
+        decompose(std::vector<int64_t>(d.begin(), d.end()), dim, o, rr, n);
+        if (o == 0u || rr == 0u || n == 0u) {
+          throw std::runtime_error(
+              "WebGPU log_softmax: live zero-sized reduction");
+        }
+        if (static_cast<uint64_t>(o) * rr * n > build_numel) {
+          throw std::runtime_error(
+              "WebGPU log_softmax: live numel exceeds build max");
+        }
+        const uint64_t live_lines = static_cast<uint64_t>(o) * n;
+        LogSoftmaxParams p = {};
+        p.outer = o;
+        p.r = rr;
+        p.inner = n;
+        wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
+        g.dispatch_at(dispatch_idx).workgroup_count_x =
+            utils::compute_1d_workgroup_count(
+                g.device(),
+                static_cast<uint32_t>(live_lines),
+                wg_size,
+                "log_softmax(resize)");
+        g.set_cur_dims(out_id, std::vector<int64_t>(d.begin(), d.end()));
+      });
+
+  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(uniform_buffer);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten._log_softmax.default, log_softmax_impl);
+}
+
+} // namespace executorch::backends::webgpu