[ExecuTorch][WebGPU] Add elementwise binary ops (div, sub) to the WebGPU backend

backends/webgpu/CMakeLists.txt

@@ -56,6 +56,8 @@ set(WEBGPU_SRCS
     runtime/ops/log_softmax/LogSoftmax.cpp
     runtime/ops/softmax/Softmax.cpp
     runtime/ops/bmm/Bmm.cpp
+    runtime/ops/div/BinaryOp.cpp
+    runtime/ops/sub/BinaryOp.cpp
     runtime/ops/linear/Linear.cpp
 )
 

backends/webgpu/runtime/ops/binary_op/binary_div_wgsl.h

@@ -0,0 +1,77 @@
+/*
+ * 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 binary_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: 233876a29c95d830f88fb22f2e3638b3b91e078b8f957b3d2c199352bdf40f93
+inline constexpr const char* kBinaryDivWGSL = R"(
+@group(0) @binding(0) var<storage, read> input1: array<f32>;
+@group(0) @binding(1) var<storage, read> input2: array<f32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(4) var<uniform> in1_meta: TensorMeta;
+@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
+
+override wg_size: u32 = 64u;
+
+fn op(a: f32, b: f32) -> f32 {
+  return a / b;
+}
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= out_meta.numel) {
+        return;
+    }
+
+    var same = true;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        if (in1_meta.sizes[d] != out_meta.sizes[d] ||
+            in2_meta.sizes[d] != out_meta.sizes[d]) {
+            same = false;
+        }
+    }
+    if (same) {
+        output[idx] = op(input1[idx], input2[idx]);
+        return;
+    }
+
+    var rem = idx;
+    var l1: u32 = 0u;
+    var l2: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
+        l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
+    }
+    output[idx] = op(input1[l1], input2[l2]);
+}
+)";
+
+inline constexpr uint32_t kBinaryDivWorkgroupSizeX = 64;
+inline constexpr uint32_t kBinaryDivWorkgroupSizeY = 1;
+inline constexpr uint32_t kBinaryDivWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/ops/binary_op/binary_op.wgsl

@@ -0,0 +1,55 @@
+@group(0) @binding(0) var<storage, read> input1: array<f32>;
+@group(0) @binding(1) var<storage, read> input2: array<f32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(4) var<uniform> in1_meta: TensorMeta;
+@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
+
+override wg_size: u32 = 64u;
+$if USE_ALPHA:
+  override alpha: f32 = 1.0;
+
+fn op(a: f32, b: f32) -> f32 {
+  return ${OP_EXPR};
+}
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= out_meta.numel) {
+        return;
+    }
+
+    var same = true;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        if (in1_meta.sizes[d] != out_meta.sizes[d] ||
+            in2_meta.sizes[d] != out_meta.sizes[d]) {
+            same = false;
+        }
+    }
+    if (same) {
+        output[idx] = op(input1[idx], input2[idx]);
+        return;
+    }
+
+    var rem = idx;
+    var l1: u32 = 0u;
+    var l2: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
+        l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
+    }
+    output[idx] = op(input1[l1], input2[l2]);
+}

backends/webgpu/runtime/ops/binary_op/binary_op.yaml

@@ -0,0 +1,11 @@
+binary_op:
+  parameter_names_with_default_values:
+    OP_EXPR: a + alpha * b
+    USE_ALPHA: 1
+  shader_variants:
+    - NAME: binary_div
+      OP_EXPR: a / b
+      USE_ALPHA: 0
+    - NAME: binary_sub
+      OP_EXPR: a - alpha * b
+      USE_ALPHA: 1

backends/webgpu/runtime/ops/binary_op/binary_sub_wgsl.h

@@ -0,0 +1,78 @@
+/*
+ * 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 binary_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: 496b343cef6838c8316686916a1c8fd971afc7a9abfc9abca4eb28db60611371
+inline constexpr const char* kBinarySubWGSL = R"(
+@group(0) @binding(0) var<storage, read> input1: array<f32>;
+@group(0) @binding(1) var<storage, read> input2: array<f32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(4) var<uniform> in1_meta: TensorMeta;
+@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
+
+override wg_size: u32 = 64u;
+override alpha: f32 = 1.0;
+
+fn op(a: f32, b: f32) -> f32 {
+  return a - alpha * b;
+}
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= out_meta.numel) {
+        return;
+    }
+
+    var same = true;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        if (in1_meta.sizes[d] != out_meta.sizes[d] ||
+            in2_meta.sizes[d] != out_meta.sizes[d]) {
+            same = false;
+        }
+    }
+    if (same) {
+        output[idx] = op(input1[idx], input2[idx]);
+        return;
+    }
+
+    var rem = idx;
+    var l1: u32 = 0u;
+    var l2: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
+        l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
+    }
+    output[idx] = op(input1[l1], input2[l2]);
+}
+)";
+
+inline constexpr uint32_t kBinarySubWorkgroupSizeX = 64;
+inline constexpr uint32_t kBinarySubWorkgroupSizeY = 1;
+inline constexpr uint32_t kBinarySubWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu