[ExecuTorch][WebGPU] Add leaky_relu op

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})

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 <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/leaky_relu/leaky_relu_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <stdexcept>
+
+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<int>& 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<uint32_t>(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, &params, 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

backends/webgpu/runtime/ops/leaky_relu/leaky_relu.wgsl

@@ -0,0 +1,24 @@
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  neg_slope: f32,
+  _p0: u32,
+  _p1: u32,
+}
+@group(0) @binding(2) var<uniform> 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<u32>) {
+  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);
+}

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 <cstdint>
+
+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<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  neg_slope: f32,
+  _p0: u32,
+  _p1: u32,
+}
+@group(0) @binding(2) var<uniform> 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<u32>) {
+  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