Skip to content

Commit 27a7c51

Browse files
JCNTHJCNTH
committed
[ExecuTorch][WebGPU] Add elementwise binary ops (div, sub) to the WebGPU backend
Pull Request resolved: #20921 Add broadcast-aware `aten.div.Tensor` + `aten.sub.Tensor` elementwise handlers for the training tail. Key changes: - `runtime/ops/{div,sub}/` — broadcast elementwise WGSL kernels + handlers - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Reuses the shared Vulkan partitioner (`aten.div.Tensor`/`sub.Tensor` already registered); WebGPU kernels only. Co-authored-with: Claude Code. ghstack-source-id: 405026052 @exported-using-ghexport Differential Revision: [D111755128](https://our.internmc.facebook.com/intern/diff/D111755128/)
1 parent 0b343a9 commit 27a7c51

7 files changed

Lines changed: 680 additions & 0 deletions

File tree

‎backends/webgpu/CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,8 @@ set(WEBGPU_SRCS
5656
runtime/ops/log_softmax/LogSoftmax.cpp
5757
runtime/ops/softmax/Softmax.cpp
5858
runtime/ops/bmm/Bmm.cpp
59+
runtime/ops/div/BinaryOp.cpp
60+
runtime/ops/sub/BinaryOp.cpp
5961
runtime/ops/linear/Linear.cpp
6062
)
6163

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <cstdint>
12+
13+
namespace executorch::backends::webgpu {
14+
15+
// @generated from binary_op.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 233876a29c95d830f88fb22f2e3638b3b91e078b8f957b3d2c199352bdf40f93
17+
inline constexpr const char* kBinaryDivWGSL = R"(
18+
@group(0) @binding(0) var<storage, read> input1: array<f32>;
19+
@group(0) @binding(1) var<storage, read> input2: array<f32>;
20+
@group(0) @binding(2) var<storage, read_write> output: array<f32>;
21+
22+
struct TensorMeta {
23+
ndim: u32,
24+
numel: u32,
25+
sizes: vec4<u32>,
26+
strides: vec4<u32>,
27+
}
28+
@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
29+
@group(0) @binding(4) var<uniform> in1_meta: TensorMeta;
30+
@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
31+
32+
override wg_size: u32 = 64u;
33+
34+
fn op(a: f32, b: f32) -> f32 {
35+
return a / b;
36+
}
37+
38+
@compute @workgroup_size(wg_size, 1, 1)
39+
fn main(
40+
@builtin(global_invocation_id) gid: vec3<u32>,
41+
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
42+
// 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
43+
let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
44+
if (idx >= out_meta.numel) {
45+
return;
46+
}
47+
48+
var same = true;
49+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
50+
if (in1_meta.sizes[d] != out_meta.sizes[d] ||
51+
in2_meta.sizes[d] != out_meta.sizes[d]) {
52+
same = false;
53+
}
54+
}
55+
if (same) {
56+
output[idx] = op(input1[idx], input2[idx]);
57+
return;
58+
}
59+
60+
var rem = idx;
61+
var l1: u32 = 0u;
62+
var l2: u32 = 0u;
63+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
64+
let coord = rem / out_meta.strides[d];
65+
rem = rem % out_meta.strides[d];
66+
l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
67+
l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
68+
}
69+
output[idx] = op(input1[l1], input2[l2]);
70+
}
71+
)";
72+
73+
inline constexpr uint32_t kBinaryDivWorkgroupSizeX = 64;
74+
inline constexpr uint32_t kBinaryDivWorkgroupSizeY = 1;
75+
inline constexpr uint32_t kBinaryDivWorkgroupSizeZ = 1;
76+
77+
} // namespace executorch::backends::webgpu

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
@group(0) @binding(0) var<storage, read> input1: array<f32>;
2+
@group(0) @binding(1) var<storage, read> input2: array<f32>;
3+
@group(0) @binding(2) var<storage, read_write> output: array<f32>;
4+
5+
struct TensorMeta {
6+
ndim: u32,
7+
numel: u32,
8+
sizes: vec4<u32>,
9+
strides: vec4<u32>,
10+
}
11+
@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
12+
@group(0) @binding(4) var<uniform> in1_meta: TensorMeta;
13+
@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
14+
15+
override wg_size: u32 = 64u;
16+
$if USE_ALPHA:
17+
override alpha: f32 = 1.0;
18+
19+
fn op(a: f32, b: f32) -> f32 {
20+
return ${OP_EXPR};
21+
}
22+
23+
@compute @workgroup_size(wg_size, 1, 1)
24+
fn main(
25+
@builtin(global_invocation_id) gid: vec3<u32>,
26+
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
27+
// 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
28+
let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
29+
if (idx >= out_meta.numel) {
30+
return;
31+
}
32+
33+
var same = true;
34+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
35+
if (in1_meta.sizes[d] != out_meta.sizes[d] ||
36+
in2_meta.sizes[d] != out_meta.sizes[d]) {
37+
same = false;
38+
}
39+
}
40+
if (same) {
41+
output[idx] = op(input1[idx], input2[idx]);
42+
return;
43+
}
44+
45+
var rem = idx;
46+
var l1: u32 = 0u;
47+
var l2: u32 = 0u;
48+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
49+
let coord = rem / out_meta.strides[d];
50+
rem = rem % out_meta.strides[d];
51+
l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
52+
l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
53+
}
54+
output[idx] = op(input1[l1], input2[l2]);
55+
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
binary_op:
2+
parameter_names_with_default_values:
3+
OP_EXPR: a + alpha * b
4+
USE_ALPHA: 1
5+
shader_variants:
6+
- NAME: binary_div
7+
OP_EXPR: a / b
8+
USE_ALPHA: 0
9+
- NAME: binary_sub
10+
OP_EXPR: a - alpha * b
11+
USE_ALPHA: 1

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

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <cstdint>
12+
13+
namespace executorch::backends::webgpu {
14+
15+
// @generated from binary_op.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 496b343cef6838c8316686916a1c8fd971afc7a9abfc9abca4eb28db60611371
17+
inline constexpr const char* kBinarySubWGSL = R"(
18+
@group(0) @binding(0) var<storage, read> input1: array<f32>;
19+
@group(0) @binding(1) var<storage, read> input2: array<f32>;
20+
@group(0) @binding(2) var<storage, read_write> output: array<f32>;
21+
22+
struct TensorMeta {
23+
ndim: u32,
24+
numel: u32,
25+
sizes: vec4<u32>,
26+
strides: vec4<u32>,
27+
}
28+
@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
29+
@group(0) @binding(4) var<uniform> in1_meta: TensorMeta;
30+
@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
31+
32+
override wg_size: u32 = 64u;
33+
override alpha: f32 = 1.0;
34+
35+
fn op(a: f32, b: f32) -> f32 {
36+
return a - alpha * b;
37+
}
38+
39+
@compute @workgroup_size(wg_size, 1, 1)
40+
fn main(
41+
@builtin(global_invocation_id) gid: vec3<u32>,
42+
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
43+
// 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
44+
let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
45+
if (idx >= out_meta.numel) {
46+
return;
47+
}
48+
49+
var same = true;
50+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
51+
if (in1_meta.sizes[d] != out_meta.sizes[d] ||
52+
in2_meta.sizes[d] != out_meta.sizes[d]) {
53+
same = false;
54+
}
55+
}
56+
if (same) {
57+
output[idx] = op(input1[idx], input2[idx]);
58+
return;
59+
}
60+
61+
var rem = idx;
62+
var l1: u32 = 0u;
63+
var l2: u32 = 0u;
64+
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
65+
let coord = rem / out_meta.strides[d];
66+
rem = rem % out_meta.strides[d];
67+
l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
68+
l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
69+
}
70+
output[idx] = op(input1[l1], input2[l2]);
71+
}
72+
)";
73+
74+
inline constexpr uint32_t kBinarySubWorkgroupSizeX = 64;
75+
inline constexpr uint32_t kBinarySubWorkgroupSizeY = 1;
76+
inline constexpr uint32_t kBinarySubWorkgroupSizeZ = 1;
77+
78+
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)