From 265715abfded2ffdf8a332dcaef373c64bfec4a5 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 16 Jul 2026 14:49:00 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 1 + .../runtime/ops/leaky_relu/LeakyRelu.cpp | 100 ++++++++++++++++++ .../runtime/ops/leaky_relu/leaky_relu.wgsl | 24 +++++ .../runtime/ops/leaky_relu/leaky_relu_wgsl.h | 48 +++++++++ 4 files changed, 173 insertions(+) create mode 100644 backends/webgpu/runtime/ops/leaky_relu/LeakyRelu.cpp create mode 100644 backends/webgpu/runtime/ops/leaky_relu/leaky_relu.wgsl create mode 100644 backends/webgpu/runtime/ops/leaky_relu/leaky_relu_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 89f4ef62742..e0d62e51179 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -64,6 +64,7 @@ set(WEBGPU_SRCS runtime/ops/constant_pad_nd/ConstantPadNd.cpp runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp runtime/ops/max_pool2d/MaxPool2d.cpp + runtime/ops/leaky_relu/LeakyRelu.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/leaky_relu/LeakyRelu.cpp b/backends/webgpu/runtime/ops/leaky_relu/LeakyRelu.cpp new file mode 100644 index 00000000000..18af535eb14 --- /dev/null +++ b/backends/webgpu/runtime/ops/leaky_relu/LeakyRelu.cpp @@ -0,0 +1,100 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +struct LeakyReluParams { + uint32_t num_elements; + float neg_slope; + uint32_t _pad[2]; +}; + +// aten.leaky_relu.default: fp32 elementwise; mirrors Vulkan leaky_relu. +void leaky_relu_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 2) { + throw std::runtime_error("WebGPU leaky_relu: expected >=2 args"); + } + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + const float neg_slope = + (args.size() >= 3) ? utils::scalar_or(graph, args.at(1), 0.01f) : 0.01f; + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("leaky_relu: null buffer binding"); + } + if (in_tensor.nbytes != out_tensor.nbytes || + out_tensor.nbytes % sizeof(float) != 0) { + throw std::runtime_error("leaky_relu: fp32-only / size mismatch"); + } + const uint32_t num_elements = + static_cast(out_tensor.nbytes / sizeof(float)); + + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, num_elements, kLeakyReluWorkgroupSizeX, "leaky_relu"); + + auto constants = utils::make_grid_constants(grid); + + LeakyReluParams params = {}; + params.num_elements = num_elements; + params.neg_slope = neg_slope; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(LeakyReluParams)); + graph.add_uniform_buffer_bytes(sizeof(LeakyReluParams)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kLeakyReluWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(LeakyReluParams)}, + }, + constants.data(), + constants.size()); + + WebGPUDispatch dispatch{}; + dispatch.pipeline = bundle.pipeline; + dispatch.bind_group = bundle.bind_group; + dispatch.workgroup_count_x = grid.count_x; + dispatch.workgroup_count_y = grid.count_y; + graph.add_dispatch(dispatch); + + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.leaky_relu.default, leaky_relu_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/leaky_relu/leaky_relu.wgsl b/backends/webgpu/runtime/ops/leaky_relu/leaky_relu.wgsl new file mode 100644 index 00000000000..317333a7e5b --- /dev/null +++ b/backends/webgpu/runtime/ops/leaky_relu/leaky_relu.wgsl @@ -0,0 +1,24 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + neg_slope: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// leaky_relu(x) = max(x,0) + neg_slope*min(x,0); fp32 elementwise, 2D-spill. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.num_elements) { + return; + } + let x = input[i]; + output[i] = max(x, 0.0) + params.neg_slope * min(x, 0.0); +} diff --git a/backends/webgpu/runtime/ops/leaky_relu/leaky_relu_wgsl.h b/backends/webgpu/runtime/ops/leaky_relu/leaky_relu_wgsl.h new file mode 100644 index 00000000000..48e5b925fb1 --- /dev/null +++ b/backends/webgpu/runtime/ops/leaky_relu/leaky_relu_wgsl.h @@ -0,0 +1,48 @@ +/* + * 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 leaky_relu.wgsl - DO NOT EDIT. +// wgsl-sha256: 98add607747a0f47521645f6cb046c500d59709f6641846a0099c5acd3b6a66b +inline constexpr const char* kLeakyReluWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + neg_slope: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// leaky_relu(x) = max(x,0) + neg_slope*min(x,0); fp32 elementwise, 2D-spill. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.num_elements) { + return; + } + let x = input[i]; + output[i] = max(x, 0.0) + params.neg_slope * min(x, 0.0); +} +)"; + +inline constexpr uint32_t kLeakyReluWorkgroupSizeX = 256; +inline constexpr uint32_t kLeakyReluWorkgroupSizeY = 1; +inline constexpr uint32_t kLeakyReluWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu