From bd59c1c0d30db2e2a2761375003be4c845b41425 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:48:08 -0700 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 2 + .../runtime/ops/log_softmax/LogSoftmax.cpp | 226 ++++++++++++++++++ .../runtime/ops/log_softmax/log_softmax.wgsl | 40 ++++ .../ops/log_softmax/log_softmax_wgsl.h | 64 +++++ .../webgpu/runtime/ops/softmax/Softmax.cpp | 219 +++++++++++++++++ .../webgpu/runtime/ops/softmax/softmax.wgsl | 39 +++ .../webgpu/runtime/ops/softmax/softmax_wgsl.h | 63 +++++ 7 files changed, 653 insertions(+) create mode 100644 backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp create mode 100644 backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl create mode 100644 backends/webgpu/runtime/ops/log_softmax/log_softmax_wgsl.h create mode 100644 backends/webgpu/runtime/ops/softmax/Softmax.cpp create mode 100644 backends/webgpu/runtime/ops/softmax/softmax.wgsl create mode 100644 backends/webgpu/runtime/ops/softmax/softmax_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 93e81192e80..ef5a08bdfb7 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/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 ) diff --git a/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp b/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp new file mode 100644 index 00000000000..ac19fd67fdd --- /dev/null +++ b/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp @@ -0,0 +1,226 @@ +/* + * 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 +#include + +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& dims, + int64_t dim, + uint32_t& outer, + uint32_t& r, + uint32_t& inner) { + const int64_t ndim = static_cast(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(dims[d]); + } + for (int64_t d = dim + 1; d < ndim; ++d) { + in *= static_cast(dims[d]); + } + outer = static_cast(o); + r = static_cast(dims[dim]); + inner = static_cast(in); +} + +void log_softmax_impl(WebGPUGraph& graph, const std::vector& 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(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(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(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, ¶ms, 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(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(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(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(o) * rr * n > build_numel) { + throw std::runtime_error( + "WebGPU log_softmax: live numel exceeds build max"); + } + const uint64_t live_lines = static_cast(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(live_lines), wg_size, + "log_softmax(resize)"); + g.set_cur_dims(out_id, std::vector(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 diff --git a/backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl b/backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl new file mode 100644 index 00000000000..b1a7831337b --- /dev/null +++ b/backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl @@ -0,0 +1,40 @@ +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let t = gid.x; + let lines = params.outer_ * params.inner_; + if (t >= lines) { + return; + } + let oo = t / params.inner_; + let ii = t % params.inner_; + let base = oo * params.r_ * params.inner_ + ii; + // Online flash-style softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX). + var mx: f32 = -3.4e38; + var ssum: f32 = 0.0; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let v = inp[base + r * params.inner_]; + if (v > mx) { + ssum = ssum * exp(mx - v); + mx = v; + } + ssum = ssum + exp(v - mx); + } + let log_sum = log(ssum); + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let idx = base + r * params.inner_; + out[idx] = (inp[idx] - mx) - log_sum; + } +} diff --git a/backends/webgpu/runtime/ops/log_softmax/log_softmax_wgsl.h b/backends/webgpu/runtime/ops/log_softmax/log_softmax_wgsl.h new file mode 100644 index 00000000000..45ef76601a1 --- /dev/null +++ b/backends/webgpu/runtime/ops/log_softmax/log_softmax_wgsl.h @@ -0,0 +1,64 @@ +/* + * 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 log_softmax.wgsl - DO NOT EDIT. +// wgsl-sha256: 0f6d1e9eea8d0ec8246c2c4f0b0be5c03deb4a73790896b8a0b093fc5be16114 +inline constexpr const char* kLogSoftmaxWGSL = R"( +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let t = gid.x; + let lines = params.outer_ * params.inner_; + if (t >= lines) { + return; + } + let oo = t / params.inner_; + let ii = t % params.inner_; + let base = oo * params.r_ * params.inner_ + ii; + // Online flash-style softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX). + var mx: f32 = -3.4e38; + var ssum: f32 = 0.0; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let v = inp[base + r * params.inner_]; + if (v > mx) { + ssum = ssum * exp(mx - v); + mx = v; + } + ssum = ssum + exp(v - mx); + } + let log_sum = log(ssum); + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let idx = base + r * params.inner_; + out[idx] = (inp[idx] - mx) - log_sum; + } +} +)"; + +inline constexpr uint32_t kLogSoftmaxWorkgroupSizeX = 256; +inline constexpr uint32_t kLogSoftmaxWorkgroupSizeY = 1; +inline constexpr uint32_t kLogSoftmaxWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/softmax/Softmax.cpp b/backends/webgpu/runtime/ops/softmax/Softmax.cpp new file mode 100644 index 00000000000..b906fa37175 --- /dev/null +++ b/backends/webgpu/runtime/ops/softmax/Softmax.cpp @@ -0,0 +1,219 @@ +/* + * 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 { + +// Uniform layout matching the WGSL Params struct (16-byte aligned). +struct SoftmaxParams { + uint32_t outer; + uint32_t r; + uint32_t inner; + uint32_t pad_; +}; +static_assert(sizeof(SoftmaxParams) == 16, "SoftmaxParams must be 16 bytes"); + +// Decompose dims into [outer, R, inner] for a reduction along `dim`. +void decompose( + const std::vector& dims, + int64_t dim, + uint32_t& outer, + uint32_t& r, + uint32_t& inner) { + const int64_t ndim = static_cast(dims.size()); + if (dim < 0) { + dim += ndim; + } + if (ndim == 0 || dim < 0 || dim >= ndim) { + throw std::runtime_error("WebGPU softmax: dim out of range"); + } + uint64_t o = 1, in = 1; + for (int64_t d = 0; d < dim; ++d) { + o *= static_cast(dims[d]); + } + for (int64_t d = dim + 1; d < ndim; ++d) { + in *= static_cast(dims[d]); + } + outer = static_cast(o); + r = static_cast(dims[dim]); + inner = static_cast(in); +} + +void softmax_impl(WebGPUGraph& graph, const std::vector& 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 softmax: scalar input unsupported"); + } + if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int) { + throw std::runtime_error("WebGPU 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 softmax: zero-sized reduction"); + } + + uint64_t numel = 1; + for (int64_t d : in.dims) { + numel *= static_cast(d); + } + if (in.nbytes != numel * sizeof(float) || + out.nbytes != numel * sizeof(float)) { + throw std::runtime_error("WebGPU softmax: fp32-only (byte-size mismatch)"); + } + + const uint64_t lines = static_cast(outer) * inner; + if (lines > UINT32_MAX) { + throw std::runtime_error( + "WebGPU softmax: line count exceeds dispatch limit"); + } + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kSoftmaxWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, static_cast(lines), wg_size, "softmax"); + + SoftmaxParams params = {}; + params.outer = outer; + params.r = r; + params.inner = inner; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(SoftmaxParams)); + graph.add_uniform_buffer_bytes(sizeof(SoftmaxParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kSoftmaxWGSL, 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(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(SoftmaxParams); + + 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, "softmax"}); + + // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch. + WGPUBuffer params_buf = uniform_buffer; + const uint64_t build_numel = + static_cast(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(d, dim, o, rr, n); + if (o == 0u || rr == 0u || n == 0u) { + throw std::runtime_error("WebGPU softmax: live zero-sized reduction"); + } + // Buffers were sized for build-time numel; reject a larger live size. + if (static_cast(o) * rr * n > build_numel) { + throw std::runtime_error( + "WebGPU softmax: live numel exceeds build max"); + } + const uint64_t live_lines = static_cast(o) * n; + SoftmaxParams 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(live_lines), wg_size, + "softmax(resize)"); + g.set_cur_dims(out_id, 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(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten._softmax.default, softmax_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/softmax/softmax.wgsl b/backends/webgpu/runtime/ops/softmax/softmax.wgsl new file mode 100644 index 00000000000..ce8596ccb8d --- /dev/null +++ b/backends/webgpu/runtime/ops/softmax/softmax.wgsl @@ -0,0 +1,39 @@ +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let t = gid.x; + let lines = params.outer_ * params.inner_; + if (t >= lines) { + return; + } + let oo = t / params.inner_; + let ii = t % params.inner_; + let base = oo * params.r_ * params.inner_ + ii; + // Online (flash-style) softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX). + var mx: f32 = -3.4e38; + var ssum: f32 = 0.0; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let v = inp[base + r * params.inner_]; + if (v > mx) { + ssum = ssum * exp(mx - v); + mx = v; + } + ssum = ssum + exp(v - mx); + } + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let idx = base + r * params.inner_; + out[idx] = exp(inp[idx] - mx) / ssum; + } +} diff --git a/backends/webgpu/runtime/ops/softmax/softmax_wgsl.h b/backends/webgpu/runtime/ops/softmax/softmax_wgsl.h new file mode 100644 index 00000000000..d287f882efd --- /dev/null +++ b/backends/webgpu/runtime/ops/softmax/softmax_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 softmax.wgsl - DO NOT EDIT. +// wgsl-sha256: cdd4171c9f599fb9f828b2b7d8e25c338cbdf3b61d9fe0a4ee4e54c92a8a4f8a +inline constexpr const char* kSoftmaxWGSL = R"( +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let t = gid.x; + let lines = params.outer_ * params.inner_; + if (t >= lines) { + return; + } + let oo = t / params.inner_; + let ii = t % params.inner_; + let base = oo * params.r_ * params.inner_ + ii; + // Online (flash-style) softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX). + var mx: f32 = -3.4e38; + var ssum: f32 = 0.0; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let v = inp[base + r * params.inner_]; + if (v > mx) { + ssum = ssum * exp(mx - v); + mx = v; + } + ssum = ssum + exp(v - mx); + } + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let idx = base + r * params.inner_; + out[idx] = exp(inp[idx] - mx) / ssum; + } +} +)"; + +inline constexpr uint32_t kSoftmaxWorkgroupSizeX = 256; +inline constexpr uint32_t kSoftmaxWorkgroupSizeY = 1; +inline constexpr uint32_t kSoftmaxWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 9e9f7da7c2a3f9416b49671c88af5e76b0f41253 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 16:29:33 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- .../runtime/ops/log_softmax/LogSoftmax.cpp | 2 +- .../runtime/ops/log_softmax/log_softmax.wgsl | 40 ------------------- .../log_softmax_wgsl.h | 6 +-- .../webgpu/runtime/ops/softmax/softmax.wgsl | 7 +++- .../webgpu/runtime/ops/softmax/softmax.yaml | 8 ++++ 5 files changed, 18 insertions(+), 45 deletions(-) delete mode 100644 backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl rename backends/webgpu/runtime/ops/{log_softmax => softmax}/log_softmax_wgsl.h (88%) create mode 100644 backends/webgpu/runtime/ops/softmax/softmax.yaml diff --git a/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp b/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp index e29a2c64ac2..f4fb9952706 100644 --- a/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp +++ b/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl b/backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl deleted file mode 100644 index b1a7831337b..00000000000 --- a/backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl +++ /dev/null @@ -1,40 +0,0 @@ -struct Params { - outer_: u32, - r_: u32, - inner_: u32, - pad_: u32, -}; - -@group(0) @binding(0) var inp: array; -@group(0) @binding(1) var out: array; -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256; - -@compute @workgroup_size(wg_size) -fn main(@builtin(global_invocation_id) gid: vec3) { - let t = gid.x; - let lines = params.outer_ * params.inner_; - if (t >= lines) { - return; - } - let oo = t / params.inner_; - let ii = t % params.inner_; - let base = oo * params.r_ * params.inner_ + ii; - // Online flash-style softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX). - var mx: f32 = -3.4e38; - var ssum: f32 = 0.0; - for (var r: u32 = 0u; r < params.r_; r = r + 1u) { - let v = inp[base + r * params.inner_]; - if (v > mx) { - ssum = ssum * exp(mx - v); - mx = v; - } - ssum = ssum + exp(v - mx); - } - let log_sum = log(ssum); - for (var r: u32 = 0u; r < params.r_; r = r + 1u) { - let idx = base + r * params.inner_; - out[idx] = (inp[idx] - mx) - log_sum; - } -} diff --git a/backends/webgpu/runtime/ops/log_softmax/log_softmax_wgsl.h b/backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h similarity index 88% rename from backends/webgpu/runtime/ops/log_softmax/log_softmax_wgsl.h rename to backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h index 45ef76601a1..d3a848bf198 100644 --- a/backends/webgpu/runtime/ops/log_softmax/log_softmax_wgsl.h +++ b/backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h @@ -12,8 +12,8 @@ namespace executorch::backends::webgpu { -// @generated from log_softmax.wgsl - DO NOT EDIT. -// wgsl-sha256: 0f6d1e9eea8d0ec8246c2c4f0b0be5c03deb4a73790896b8a0b093fc5be16114 +// @generated from softmax.wgsl - DO NOT EDIT. +// wgsl-sha256: 3e5550d3bf913a4e9e871b1474934278f81b4d635504aceb8ab4e096ef18889a inline constexpr const char* kLogSoftmaxWGSL = R"( struct Params { outer_: u32, @@ -38,7 +38,7 @@ fn main(@builtin(global_invocation_id) gid: vec3) { let oo = t / params.inner_; let ii = t % params.inner_; let base = oo * params.r_ * params.inner_ + ii; - // Online flash-style softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX). + // Online (flash-style) softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX). var mx: f32 = -3.4e38; var ssum: f32 = 0.0; for (var r: u32 = 0u; r < params.r_; r = r + 1u) { diff --git a/backends/webgpu/runtime/ops/softmax/softmax.wgsl b/backends/webgpu/runtime/ops/softmax/softmax.wgsl index ce8596ccb8d..9ff728a6a01 100644 --- a/backends/webgpu/runtime/ops/softmax/softmax.wgsl +++ b/backends/webgpu/runtime/ops/softmax/softmax.wgsl @@ -32,8 +32,13 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } ssum = ssum + exp(v - mx); } + $if LOG: + let log_sum = log(ssum); for (var r: u32 = 0u; r < params.r_; r = r + 1u) { let idx = base + r * params.inner_; - out[idx] = exp(inp[idx] - mx) / ssum; + $if LOG: + out[idx] = (inp[idx] - mx) - log_sum; + $else: + out[idx] = exp(inp[idx] - mx) / ssum; } } diff --git a/backends/webgpu/runtime/ops/softmax/softmax.yaml b/backends/webgpu/runtime/ops/softmax/softmax.yaml new file mode 100644 index 00000000000..dfe26ee4cc0 --- /dev/null +++ b/backends/webgpu/runtime/ops/softmax/softmax.yaml @@ -0,0 +1,8 @@ +softmax: + parameter_names_with_default_values: + LOG: 0 + shader_variants: + - NAME: softmax + LOG: 0 + - NAME: log_softmax + LOG: 1