Skip to content

Commit 2041bbe

Browse files
JCNTHJCNTH
committed
[ExecuTorch][WebGPU] Generate logical and arithmetic binary shader variants
Pull Request resolved: #21451 **Generate byte-identical logical and arithmetic binary variants from shared WGSL families** Logical AND/OR and four arithmetic binary kernels duplicated shader skeletons and broadcast logic. This consolidates logical AND/OR behind one packed-Boolean family and minimum/pow/floor_divide/mul into the existing binary family, with a permanent mixed-rank broadcast contract. Key changes: - Generate logical AND/OR from one operator-token manifest. - Generate minimum, pow, floor_divide, and mul beside the existing div/sub variants. - Lock same-shape and mixed-rank expressions, exact payloads/workgroups, PTE delegation, and broadcast boundary cases. No runtime C++ dispatch, bindings, pipeline construction, workgroups, or expanded shader payloads change. Four standalone WGSL inputs are removed, and future compatible variants require manifest entries instead of copied kernels. This follows the Vulkan binary-family pattern. Co-authored-with: Claude Code. ghstack-source-id: 411961479 @exported-using-ghexport Differential Revision: [D113979789](https://our.internmc.facebook.com/intern/diff/D113979789/)
1 parent 87d8bb5 commit 2041bbe

31 files changed

Lines changed: 402 additions & 404 deletions

‎backends/webgpu/runtime/WebGPUShaderRegistry.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@
1717
#include <executorch/backends/webgpu/runtime/ops/avg_pool2d/avg_pool2d_wgsl.h>
1818
#include <executorch/backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h>
1919
#include <executorch/backends/webgpu/runtime/ops/binary_op/binary_div_wgsl.h>
20+
#include <executorch/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h>
21+
#include <executorch/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h>
22+
#include <executorch/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h>
23+
#include <executorch/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h>
2024
#include <executorch/backends/webgpu/runtime/ops/binary_op/binary_sub_wgsl.h>
2125
#include <executorch/backends/webgpu/runtime/ops/bitwise_not/bitwise_not_wgsl.h>
2226
#include <executorch/backends/webgpu/runtime/ops/bmm/bmm_tiled_wgsl.h>
@@ -50,7 +54,6 @@
5054
#include <executorch/backends/webgpu/runtime/ops/extrema/amin_wgsl.h>
5155
#include <executorch/backends/webgpu/runtime/ops/fill/fill_wgsl.h>
5256
#include <executorch/backends/webgpu/runtime/ops/flip/flip_wgsl.h>
53-
#include <executorch/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h>
5457
#include <executorch/backends/webgpu/runtime/ops/fused_ce/fused_ce_wgsl.h>
5558
#include <executorch/backends/webgpu/runtime/ops/gather/gather_wgsl.h>
5659
#include <executorch/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h>
@@ -65,20 +68,17 @@
6568
#include <executorch/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw_wgsl.h>
6669
#include <executorch/backends/webgpu/runtime/ops/linear_q8ta_q8csw/linear_q8ta_q8csw_wgsl.h>
6770
#include <executorch/backends/webgpu/runtime/ops/linear_qcs4w/qcs4w_linear_wgsl.h>
68-
#include <executorch/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h>
69-
#include <executorch/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h>
71+
#include <executorch/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h>
72+
#include <executorch/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h>
7073
#include <executorch/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h>
71-
#include <executorch/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h>
7274
#include <executorch/backends/webgpu/runtime/ops/mm/mm_tiled_wgsl.h>
7375
#include <executorch/backends/webgpu/runtime/ops/mm/mm_vec4_wgsl.h>
7476
#include <executorch/backends/webgpu/runtime/ops/mm/mm_wgsl.h>
75-
#include <executorch/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h>
7677
#include <executorch/backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h>
7778
#include <executorch/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce_wgsl.h>
7879
#include <executorch/backends/webgpu/runtime/ops/native_group_norm/group_norm_wgsl.h>
7980
#include <executorch/backends/webgpu/runtime/ops/permute/permute_wgsl.h>
8081
#include <executorch/backends/webgpu/runtime/ops/pixel_shuffle/pixel_shuffle_wgsl.h>
81-
#include <executorch/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h>
8282
#include <executorch/backends/webgpu/runtime/ops/q8ta_add/q8ta_add_wgsl.h>
8383
#include <executorch/backends/webgpu/runtime/ops/q8ta_conv2d/q8ta_conv2d_wgsl.h>
8484
#include <executorch/backends/webgpu/runtime/ops/q8ta_conv2d_dw/q8ta_conv2d_dw_wgsl.h>

backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h renamed to backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
namespace executorch::backends::webgpu {
1414

15-
// @generated from binary_floor_divide.wgsl - DO NOT EDIT.
15+
// @generated from binary_op.wgsl - DO NOT EDIT.
1616
// wgsl-sha256: baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3
1717
inline constexpr const char* kBinaryFloorDivideWGSL = R"(
1818
@group(0) @binding(0) var<storage, read> input1: array<f32>;

backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h renamed to backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
namespace executorch::backends::webgpu {
1414

15-
// @generated from binary_minimum.wgsl - DO NOT EDIT.
15+
// @generated from binary_op.wgsl - DO NOT EDIT.
1616
// wgsl-sha256: 929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972
1717
inline constexpr const char* kBinaryMinimumWGSL = R"(
1818
@group(0) @binding(0) var<storage, read> input1: array<f32>;

backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h renamed to backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
namespace executorch::backends::webgpu {
1414

15-
// @generated from binary_mul.wgsl - DO NOT EDIT.
15+
// @generated from binary_op.wgsl - DO NOT EDIT.
1616
// wgsl-sha256: d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812
1717
inline constexpr const char* kBinaryMulWGSL = R"(
1818
@group(0) @binding(0) var<storage, read> input1: array<f32>;

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

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,14 @@ override wg_size: u32 = 64u;
1616
$if USE_ALPHA:
1717
override alpha: f32 = 1.0;
1818

19-
fn op(a: f32, b: f32) -> f32 {
20-
return ${OP_EXPR};
21-
}
19+
$if INLINE:
20+
@compute @workgroup_size(wg_size, 1, 1)
21+
$else:
22+
fn op(a: f32, b: f32) -> f32 {
23+
return ${OP_EXPR};
24+
}
2225

23-
@compute @workgroup_size(wg_size, 1, 1)
26+
@compute @workgroup_size(wg_size, 1, 1)
2427
fn main(
2528
@builtin(global_invocation_id) gid: vec3<u32>,
2629
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
@@ -30,6 +33,8 @@ fn main(
3033
return;
3134
}
3235

36+
$if INLINE:
37+
// Fast path: every input dim matches the output dim -> elementwise.
3338
var same = true;
3439
for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
3540
if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] ||
@@ -38,10 +43,15 @@ fn main(
3843
}
3944
}
4045
if (same) {
41-
output[idx] = op(input1[idx], input2[idx]);
46+
$if INLINE:
47+
output[idx] = ${SAME_EXPR};
48+
$else:
49+
output[idx] = op(input1[idx], input2[idx]);
4250
return;
4351
}
4452

53+
$if INLINE:
54+
// Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize.
4555
var rem = idx;
4656
var l1: u32 = 0u;
4757
var l2: u32 = 0u;
@@ -51,5 +61,8 @@ fn main(
5161
l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u];
5262
l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u];
5363
}
54-
output[idx] = op(input1[l1], input2[l2]);
64+
$if INLINE:
65+
output[idx] = ${BROADCAST_EXPR};
66+
$else:
67+
output[idx] = op(input1[l1], input2[l2]);
5568
}

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,33 @@ binary_op:
22
parameter_names_with_default_values:
33
OP_EXPR: a + alpha * b
44
USE_ALPHA: 1
5+
INLINE: 0
6+
SAME_EXPR: input1[idx] + input2[idx]
7+
BROADCAST_EXPR: input1[l1] + input2[l2]
58
shader_variants:
69
- NAME: binary_div
710
OP_EXPR: a / b
811
USE_ALPHA: 0
912
- NAME: binary_sub
1013
OP_EXPR: a - alpha * b
1114
USE_ALPHA: 1
15+
- NAME: binary_minimum
16+
USE_ALPHA: 0
17+
INLINE: 1
18+
SAME_EXPR: min(input1[idx], input2[idx])
19+
BROADCAST_EXPR: min(input1[l1], input2[l2])
20+
- NAME: binary_pow
21+
USE_ALPHA: 0
22+
INLINE: 1
23+
SAME_EXPR: pow(input1[idx], input2[idx])
24+
BROADCAST_EXPR: pow(input1[l1], input2[l2])
25+
- NAME: binary_floor_divide
26+
USE_ALPHA: 0
27+
INLINE: 1
28+
SAME_EXPR: floor(input1[idx] / input2[idx])
29+
BROADCAST_EXPR: floor(input1[l1] / input2[l2])
30+
- NAME: binary_mul
31+
USE_ALPHA: 0
32+
INLINE: 1
33+
SAME_EXPR: input1[idx] * input2[idx]
34+
BROADCAST_EXPR: input1[l1] * input2[l2]

backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h renamed to backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
namespace executorch::backends::webgpu {
1414

15-
// @generated from binary_pow.wgsl - DO NOT EDIT.
15+
// @generated from binary_op.wgsl - DO NOT EDIT.
1616
// wgsl-sha256: a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c
1717
inline constexpr const char* kBinaryPowWGSL = R"(
1818
@group(0) @binding(0) var<storage, read> input1: array<f32>;

‎backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
1010
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
1111
#include <executorch/backends/webgpu/runtime/ops/binary/BinaryOp.h>
12-
#include <executorch/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h>
12+
#include <executorch/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h>
1313

1414
#include <stdexcept>
1515
#include <vector>

‎backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl‎

Lines changed: 0 additions & 51 deletions
This file was deleted.

‎backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
1010
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
1111
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
12-
#include <executorch/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h>
12+
#include <executorch/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h>
1313

1414
#include <webgpu/webgpu.h>
1515

0 commit comments

Comments
 (0)