[ExecuTorch][WebGPU] Add relu op (shared unary handler; sigmoid adopts make_compute_pipeline)

backends/webgpu/runtime/ops/batch_norm/BatchNorm.cpp

@@ -0,0 +1,172 @@
+/*
+ * 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/batch_norm/batch_norm_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <stdexcept>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct BNParams {
+  uint32_t num_elements;
+  uint32_t C;
+  uint32_t HW;
+  uint32_t has_weight;
+  uint32_t has_bias;
+  float eps;
+  uint32_t _p0;
+  uint32_t _p1;
+};
+static_assert(sizeof(BNParams) == 32, "BNParams must be 32 bytes");
+
+// _native_batch_norm_legit_no_training: per-channel affine from running stats.
+void batch_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  if (args.size() < 8) {
+    throw std::runtime_error("WebGPU batch_norm: expected >=8 args");
+  }
+  const int in_id = args.at(0);
+  const int weight_id = args.at(1);
+  const int bias_id = args.at(2);
+  const int mean_id = args.at(3);
+  const int var_id = args.at(4);
+  const double eps = graph.get_double(args.at(6));
+  const std::vector<int>& outs = graph.get_value_list(args.at(args.size() - 1));
+  if (outs.empty()) {
+    throw std::runtime_error("WebGPU batch_norm: empty out ValueList");
+  }
+  const int out_id = outs.at(0);
+
+  WGPUDevice device = graph.device();
+  const auto& in = graph.get_tensor(in_id);
+  const auto& out = graph.get_tensor(out_id);
+  const auto& mean = graph.get_tensor(mean_id);
+  const auto& var = graph.get_tensor(var_id);
+
+  if (in.dims.size() != 4 || out.dims.size() != 4) {
+    throw std::runtime_error("WebGPU batch_norm: expected 4D in/out");
+  }
+  const uint32_t N = static_cast<uint32_t>(in.dims[0]);
+  const uint32_t C = static_cast<uint32_t>(in.dims[1]);
+  const uint32_t HW =
+      static_cast<uint32_t>(in.dims[2]) * static_cast<uint32_t>(in.dims[3]);
+
+  uint64_t out_numel = uint64_t(N) * C * HW;
+  if (out.nbytes != out_numel * sizeof(float)) {
+    throw std::runtime_error("WebGPU batch_norm: fp32-only (byte mismatch)");
+  }
+  if (out_numel > 0xFFFFFFFFull) {
+    throw std::runtime_error("WebGPU batch_norm: output too large");
+  }
+
+  const bool has_weight =
+      graph.get_value_type(weight_id) == WebGPUGraph::ValueType::Tensor;
+  const bool has_bias =
+      graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor;
+
+  // Per-channel params are indexed [0, C); guard their byte sizes so an
+  // undersized buffer can't drive an out-of-bounds shader read. mean/var are
+  // required; weight/bias only when present.
+  const uint64_t c_bytes = uint64_t(C) * sizeof(float);
+  if (mean.nbytes != c_bytes || var.nbytes != c_bytes) {
+    throw std::runtime_error(
+        "WebGPU batch_norm: mean/var must be fp32 of length C");
+  }
+  if (has_weight && graph.get_tensor(weight_id).nbytes != c_bytes) {
+    throw std::runtime_error(
+        "WebGPU batch_norm: weight must be fp32 of length C");
+  }
+  if (has_bias && graph.get_tensor(bias_id).nbytes != c_bytes) {
+    throw std::runtime_error(
+        "WebGPU batch_norm: bias must be fp32 of length C");
+  }
+
+  utils::DispatchGrid grid = utils::compute_dispatch_grid(
+      device,
+      static_cast<uint32_t>(out_numel),
+      kBatchNormWorkgroupSizeX,
+      "batch_norm");
+
+  BNParams params = {};
+  params.num_elements = static_cast<uint32_t>(out_numel);
+  params.C = C;
+  params.HW = HW;
+  params.has_weight = has_weight ? 1u : 0u;
+  params.has_bias = has_bias ? 1u : 0u;
+  params.eps = static_cast<float>(eps);
+
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, sizeof(BNParams));
+  graph.add_uniform_buffer_bytes(sizeof(BNParams));
+
+  utils::OptionalBinding wbind = utils::make_optional_binding(
+      device,
+      has_weight,
+      has_weight ? graph.get_tensor(weight_id).buffer : nullptr,
+      has_weight ? graph.get_tensor(weight_id).nbytes : 0);
+  utils::OptionalBinding bbind = utils::make_optional_binding(
+      device,
+      has_bias,
+      has_bias ? graph.get_tensor(bias_id).buffer : nullptr,
+      has_bias ? graph.get_tensor(bias_id).nbytes : 0);
+
+  auto constants = utils::make_grid_constants(grid);
+
+  utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+      device,
+      kBatchNormWGSL,
+      {
+          {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes},
+          {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes},
+          {2,
+           WGPUBufferBindingType_ReadOnlyStorage,
+           wbind.buffer,
+           wbind.nbytes},
+          {3,
+           WGPUBufferBindingType_ReadOnlyStorage,
+           bbind.buffer,
+           bbind.nbytes},
+          {4, WGPUBufferBindingType_ReadOnlyStorage, mean.buffer, mean.nbytes},
+          {5, WGPUBufferBindingType_ReadOnlyStorage, var.buffer, var.nbytes},
+          {6, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(BNParams)},
+      },
+      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);
+  if (wbind.owned_dummy != nullptr) {
+    wgpuBufferRelease(wbind.owned_dummy);
+  }
+  if (bbind.owned_dummy != nullptr) {
+    wgpuBufferRelease(bbind.owned_dummy);
+  }
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(
+      aten._native_batch_norm_legit_no_training.default, batch_norm_impl);
+}
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/ops/batch_norm/batch_norm.wgsl

@@ -0,0 +1,41 @@
+@group(0) @binding(0) var<storage, read_write> out: array<f32>;
+@group(0) @binding(1) var<storage, read> inp: array<f32>;
+@group(0) @binding(2) var<storage, read> weight: array<f32>;
+@group(0) @binding(3) var<storage, read> bias: array<f32>;
+@group(0) @binding(4) var<storage, read> mean: array<f32>;
+@group(0) @binding(5) var<storage, read> var_: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  C: u32,
+  HW: u32,
+  has_weight: u32,
+  has_bias: u32,
+  eps: f32,
+  _p0: u32,
+  _p1: u32,
+}
+@group(0) @binding(6) var<uniform> params: Params;
+
+override wg_size: u32 = 256;
+override stride_x: u32 = 4294967295u;
+
+// Inference batch_norm NCHW fp32: per-channel affine from running stats.
+@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 c = (i / params.HW) % params.C;
+  let inv = inverseSqrt(var_[c] + params.eps);
+  var g: f32 = 1.0;
+  if (params.has_weight == 1u) {
+    g = weight[c];
+  }
+  var b: f32 = 0.0;
+  if (params.has_bias == 1u) {
+    b = bias[c];
+  }
+  out[i] = (inp[i] - mean[c]) * inv * g + b;
+}

backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h

@@ -0,0 +1,65 @@
+/*
+ * 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 batch_norm.wgsl - DO NOT EDIT.
+// wgsl-sha256: ac7208cec0fe1454698b9b16985172c1ef34c994880b568238eaeaf1d62c0342
+inline constexpr const char* kBatchNormWGSL = R"(
+@group(0) @binding(0) var<storage, read_write> out: array<f32>;
+@group(0) @binding(1) var<storage, read> inp: array<f32>;
+@group(0) @binding(2) var<storage, read> weight: array<f32>;
+@group(0) @binding(3) var<storage, read> bias: array<f32>;
+@group(0) @binding(4) var<storage, read> mean: array<f32>;
+@group(0) @binding(5) var<storage, read> var_: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  C: u32,
+  HW: u32,
+  has_weight: u32,
+  has_bias: u32,
+  eps: f32,
+  _p0: u32,
+  _p1: u32,
+}
+@group(0) @binding(6) var<uniform> params: Params;
+
+override wg_size: u32 = 256;
+override stride_x: u32 = 4294967295u;
+
+// Inference batch_norm NCHW fp32: per-channel affine from running stats.
+@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 c = (i / params.HW) % params.C;
+  let inv = inverseSqrt(var_[c] + params.eps);
+  var g: f32 = 1.0;
+  if (params.has_weight == 1u) {
+    g = weight[c];
+  }
+  var b: f32 = 0.0;
+  if (params.has_bias == 1u) {
+    b = bias[c];
+  }
+  out[i] = (inp[i] - mean[c]) * inv * g + b;
+}
+)";
+
+inline constexpr uint32_t kBatchNormWorkgroupSizeX = 256;
+inline constexpr uint32_t kBatchNormWorkgroupSizeY = 1;
+inline constexpr uint32_t kBatchNormWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp

@@ -84,61 +84,22 @@ void dispatch_bool_op(
       utils::make_uniform(device, &params, sizeof(BoolOpParams));
   graph.add_uniform_buffer_bytes(sizeof(BoolOpParams));
 
-  WGPUShaderSourceWGSL wgsl_desc = {};
-  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
-  wgsl_desc.code = {wgsl, WGPU_STRLEN};
-  WGPUShaderModuleDescriptor shader_desc = {};
-  shader_desc.nextInChain = &wgsl_desc.chain;
-  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
-
-  WGPUBindGroupLayoutEntry entries[3] = {};
-  entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
-  entries[1].buffer.type = WGPUBufferBindingType_Storage;
-  entries[2].buffer.type = WGPUBufferBindingType_Uniform;
-  for (uint32_t i = 0; i < 3; i++) {
-    entries[i].binding = i;
-    entries[i].visibility = WGPUShaderStage_Compute;
-  }
-
-  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);
-
-  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 = self_tensor.buffer;
-  bg_entries[0].size = self_bind_size;
-  bg_entries[1].binding = 1;
-  bg_entries[1].buffer = out_tensor.buffer;
-  bg_entries[1].size = out_bind_size;
-  bg_entries[2].binding = 2;
-  bg_entries[2].buffer = params_buf;
-  bg_entries[2].size = sizeof(BoolOpParams);
-
-  WGPUBindGroupDescriptor bg_desc = {};
-  bg_desc.layout = bgl;
-  bg_desc.entryCount = 3;
-  bg_desc.entries = bg_entries;
-  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+  utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+      device,
+      wgsl,
+      {
+          {0,
+           WGPUBufferBindingType_ReadOnlyStorage,
+           self_tensor.buffer,
+           self_bind_size},
+          {1, WGPUBufferBindingType_Storage, out_tensor.buffer, out_bind_size},
+          {2, WGPUBufferBindingType_Uniform, params_buf, sizeof(BoolOpParams)},
+      },
+      &wg_size_constant,
+      1);
 
   const size_t dispatch_idx =
-      graph.add_dispatch({pipeline, bind_group, workgroup_count});
+      graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count});
 
   WGPUBuffer p_buf = params_buf;
   auto resize =
@@ -158,9 +119,6 @@ void dispatch_bool_op(
       };
   graph.add_tensor_resize_hook(self_id, resize);
 
-  wgpuShaderModuleRelease(shader);
-  wgpuBindGroupLayoutRelease(bgl);
-  wgpuPipelineLayoutRelease(pipeline_layout);
   graph.own_uniform_buffer(params_buf);
 }