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..f4fb9952706 --- /dev/null +++ b/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 +#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/softmax/Softmax.cpp b/backends/webgpu/runtime/ops/softmax/Softmax.cpp new file mode 100644 index 00000000000..72b48cd8bdd --- /dev/null +++ b/backends/webgpu/runtime/ops/softmax/Softmax.cpp @@ -0,0 +1,220 @@ +/* + * 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/log_softmax_wgsl.h b/backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h new file mode 100644 index 00000000000..d3a848bf198 --- /dev/null +++ b/backends/webgpu/runtime/ops/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 softmax.wgsl - DO NOT EDIT. +// wgsl-sha256: 3e5550d3bf913a4e9e871b1474934278f81b4d635504aceb8ab4e096ef18889a +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.wgsl b/backends/webgpu/runtime/ops/softmax/softmax.wgsl new file mode 100644 index 00000000000..9ff728a6a01 --- /dev/null +++ b/backends/webgpu/runtime/ops/softmax/softmax.wgsl @@ -0,0 +1,44 @@ +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); + } + $if LOG: + let log_sum = log(ssum); + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let idx = base + r * params.inner_; + $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 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