Skip to content

Commit 820c86f

Browse files
JCNTHJCNTH
committed
Update
[ghstack-poisoned]
2 parents 73005fd + 1d1cf7a commit 820c86f

12 files changed

Lines changed: 55 additions & 107 deletions

File tree

backends/webgpu/runtime/ops/div/binary_div_wgsl.h renamed to backends/webgpu/runtime/ops/binary_op/binary_div_wgsl.h

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212

1313
namespace executorch::backends::webgpu {
1414

15-
// @generated from binary_div.wgsl - DO NOT EDIT.
16-
// wgsl-sha256: fd0732a2aac4ba396c506331a91541e892751c7ae4cc9e4b35f45983c7b9addd
15+
// @generated from binary_op.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 233876a29c95d830f88fb22f2e3638b3b91e078b8f957b3d2c199352bdf40f93
1717
inline constexpr const char* kBinaryDivWGSL = R"(
1818
@group(0) @binding(0) var<storage, read> input1: array<f32>;
1919
@group(0) @binding(1) var<storage, read> input2: array<f32>;
@@ -31,6 +31,10 @@ struct TensorMeta {
3131
3232
override wg_size: u32 = 64u;
3333
34+
fn op(a: f32, b: f32) -> f32 {
35+
return a / b;
36+
}
37+
3438
@compute @workgroup_size(wg_size, 1, 1)
3539
fn main(
3640
@builtin(global_invocation_id) gid: vec3<u32>,
@@ -49,7 +53,7 @@ fn main(
4953
}
5054
}
5155
if (same) {
52-
output[idx] = input1[idx] / input2[idx];
56+
output[idx] = op(input1[idx], input2[idx]);
5357
return;
5458
}
5559
@@ -62,7 +66,7 @@ fn main(
6266
l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
6367
l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
6468
}
65-
output[idx] = input1[l1] / input2[l2];
69+
output[idx] = op(input1[l1], input2[l2]);
6670
}
6771
)";
6872

backends/webgpu/runtime/ops/div/binary_div.wgsl renamed to backends/webgpu/runtime/ops/binary_op/binary_op.wgsl

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,12 @@ struct TensorMeta {
1313
@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
1414

1515
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+
}
1622

1723
@compute @workgroup_size(wg_size, 1, 1)
1824
fn main(
@@ -32,7 +38,7 @@ fn main(
3238
}
3339
}
3440
if (same) {
35-
output[idx] = input1[idx] / input2[idx];
41+
output[idx] = op(input1[idx], input2[idx]);
3642
return;
3743
}
3844

@@ -45,5 +51,5 @@ fn main(
4551
l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
4652
l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
4753
}
48-
output[idx] = input1[l1] / input2[l2];
54+
output[idx] = op(input1[l1], input2[l2]);
4955
}

‎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/sub/binary_sub_wgsl.h renamed to backends/webgpu/runtime/ops/binary_op/binary_sub_wgsl.h

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212

1313
namespace executorch::backends::webgpu {
1414

15-
// @generated from binary_sub.wgsl - DO NOT EDIT.
16-
// wgsl-sha256: 3f1b1253c826770b5f9f69d1072aa63b88450096893687ecbf23156204890f6a
15+
// @generated from binary_op.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 496b343cef6838c8316686916a1c8fd971afc7a9abfc9abca4eb28db60611371
1717
inline constexpr const char* kBinarySubWGSL = R"(
1818
@group(0) @binding(0) var<storage, read> input1: array<f32>;
1919
@group(0) @binding(1) var<storage, read> input2: array<f32>;
@@ -32,6 +32,10 @@ struct TensorMeta {
3232
override wg_size: u32 = 64u;
3333
override alpha: f32 = 1.0;
3434
35+
fn op(a: f32, b: f32) -> f32 {
36+
return a - alpha * b;
37+
}
38+
3539
@compute @workgroup_size(wg_size, 1, 1)
3640
fn main(
3741
@builtin(global_invocation_id) gid: vec3<u32>,
@@ -50,7 +54,7 @@ fn main(
5054
}
5155
}
5256
if (same) {
53-
output[idx] = input1[idx] - alpha * input2[idx];
57+
output[idx] = op(input1[idx], input2[idx]);
5458
return;
5559
}
5660
@@ -63,7 +67,7 @@ fn main(
6367
l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
6468
l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
6569
}
66-
output[idx] = input1[l1] - alpha * input2[l2];
70+
output[idx] = op(input1[l1], input2[l2]);
6771
}
6872
)";
6973

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
1111
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
1212
#include <executorch/backends/webgpu/runtime/ops/TensorMeta.h>
13-
#include <executorch/backends/webgpu/runtime/ops/div/binary_div_wgsl.h>
13+
#include <executorch/backends/webgpu/runtime/ops/binary_op/binary_div_wgsl.h>
1414

1515
#include <webgpu/webgpu.h>
1616

‎backends/webgpu/runtime/ops/log_softmax/LogSoftmax.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/log_softmax/log_softmax_wgsl.h>
12+
#include <executorch/backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h>
1313

1414
#include <webgpu/webgpu.h>
1515

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

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

backends/webgpu/runtime/ops/log_softmax/log_softmax_wgsl.h renamed to backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212

1313
namespace executorch::backends::webgpu {
1414

15-
// @generated from log_softmax.wgsl - DO NOT EDIT.
16-
// wgsl-sha256: 0f6d1e9eea8d0ec8246c2c4f0b0be5c03deb4a73790896b8a0b093fc5be16114
15+
// @generated from softmax.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 3e5550d3bf913a4e9e871b1474934278f81b4d635504aceb8ab4e096ef18889a
1717
inline constexpr const char* kLogSoftmaxWGSL = R"(
1818
struct Params {
1919
outer_: u32,
@@ -38,7 +38,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
3838
let oo = t / params.inner_;
3939
let ii = t % params.inner_;
4040
let base = oo * params.r_ * params.inner_ + ii;
41-
// Online flash-style softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX).
41+
// Online (flash-style) softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX).
4242
var mx: f32 = -3.4e38;
4343
var ssum: f32 = 0.0;
4444
for (var r: u32 = 0u; r < params.r_; r = r + 1u) {

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,13 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
3232
}
3333
ssum = ssum + exp(v - mx);
3434
}
35+
$if LOG:
36+
let log_sum = log(ssum);
3537
for (var r: u32 = 0u; r < params.r_; r = r + 1u) {
3638
let idx = base + r * params.inner_;
37-
out[idx] = exp(inp[idx] - mx) / ssum;
39+
$if LOG:
40+
out[idx] = (inp[idx] - mx) - log_sum;
41+
$else:
42+
out[idx] = exp(inp[idx] - mx) / ssum;
3843
}
3944
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
softmax:
2+
parameter_names_with_default_values:
3+
LOG: 0
4+
shader_variants:
5+
- NAME: softmax
6+
LOG: 0
7+
- NAME: log_softmax
8+
LOG: 1

0 commit comments

Comments
 (0)