Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Cos

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
ff9bb7a
·
verified ·
1 Parent(s): 0f33bb4

sync 6fdf6301e2bb

Browse files
Files changed (6) hide show
  1. README.md +3 -3
  2. build/webgpu/bench.json +1 -1
  3. build/webgpu/manifest.json +10 -25
  4. build/webgpu/metadata.json +7 -7
  5. build/webgpu/unary-scalar.wgsl.jinja +65 -102
  6. build/webgpu/unary-vec4.wgsl.jinja +79 -103
README.md CHANGED
@@ -39,14 +39,14 @@ See the [ONNX `Cos` spec](https://onnx.ai/onnx/operators/onnx__Cos.html) for the
39
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
40
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
41
  - [`test.json`](build/webgpu/test.json) — correctness cases
42
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
43
  - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
44
  - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
45
 
46
  ## Use with `@huggingface/kernels`
47
 
48
  ```sh
49
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
50
  ```
51
 
52
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
@@ -60,5 +60,5 @@ Replace each `*Data` placeholder with a typed array containing the corresponding
60
  import { getKernel } from "@huggingface/kernels";
61
 
62
  const kernel = await getKernel("webgpu-kernels/ai.onnx.Cos", { version: 1 });
63
- const { y } = await kernel({ x: { data: xData, shape: [] } });
64
  ```
 
39
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
40
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
41
  - [`test.json`](build/webgpu/test.json) — correctness cases
42
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
43
  - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
44
  - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
45
 
46
  ## Use with `@huggingface/kernels`
47
 
48
  ```sh
49
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
50
  ```
51
 
52
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
60
  import { getKernel } from "@huggingface/kernels";
61
 
62
  const kernel = await getKernel("webgpu-kernels/ai.onnx.Cos", { version: 1 });
63
+ const { y } = await kernel({ x: { data: xData, shape: [3] } });
64
  ```
build/webgpu/bench.json CHANGED
@@ -6,7 +6,7 @@
6
  "outputs": { "y": { "dtype": "float32", "shape": [1048576] } }
7
  },
8
  {
9
- "name": "cos-f32-1m-vec4-healthy",
10
  "preset": "smoke",
11
  "vars": { "dtype": "float32", "count": 1048576 },
12
  "inputs": { "x": { "shape": [1048576], "dtype": "float32", "dist": "normal", "seed": 8801, "scale": 3 } },
 
6
  "outputs": { "y": { "dtype": "float32", "shape": [1048576] } }
7
  },
8
  {
9
+ "name": "cos-f32-1m-vec4-control",
10
  "preset": "smoke",
11
  "vars": { "dtype": "float32", "count": 1048576 },
12
  "inputs": { "x": { "shape": [1048576], "dtype": "float32", "dist": "normal", "seed": 8801, "scale": 3 } },
build/webgpu/manifest.json CHANGED
@@ -11,20 +11,13 @@
11
  "id": "same_layout_vec4",
12
  "priority": 20,
13
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
14
- "derive": {
15
- "scalar": "dtypes.T",
16
- "usesF16": "dtypes.T == \"f16\"",
17
- "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\""
18
- },
19
  "passes": [
20
  {
21
  "id": "main",
22
- "name": "Cos.vec4",
23
  "shader": "unary-vec4.wgsl.jinja",
24
- "derive": {
25
- "op": "\"cos\"",
26
- "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1"
27
- },
28
  "bindings": ["x", "y", "params_unary"],
29
  "dispatch": {
30
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -37,14 +30,14 @@
37
  {
38
  "id": "elementwise",
39
  "when": ["numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
40
- "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
41
  "passes": [
42
  {
43
  "id": "main",
44
  "name": "Cos",
45
  "shader": "unary-scalar.wgsl.jinja",
46
- "derive": { "op": "\"cos\"", "itemsPerInvocation": 4 },
47
- "bindings": ["x_2", "y_2", "params_2_unary"],
48
  "dispatch": {
49
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
50
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -57,17 +50,9 @@
57
  "bindings": {
58
  "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
59
  "y": { "buffer": "storage", "elementType": "$vectorScalar" },
60
- "params_unary": {
61
- "buffer": "uniform",
62
- "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }],
63
- "name": "params"
64
- },
65
- "x_2": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
66
- "y_2": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
67
- "params_2_unary": {
68
- "buffer": "uniform",
69
- "name": "params",
70
- "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
71
- }
72
  }
73
  }
 
11
  "id": "same_layout_vec4",
12
  "priority": 20,
13
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
14
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
 
 
 
 
15
  "passes": [
16
  {
17
  "id": "main",
18
+ "name": "Cos.Vec4",
19
  "shader": "unary-vec4.wgsl.jinja",
20
+ "derive": { "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1" },
 
 
 
21
  "bindings": ["x", "y", "params_unary"],
22
  "dispatch": {
23
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
 
30
  {
31
  "id": "elementwise",
32
  "when": ["numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
33
+ "derive": { "scalar": "dtypes.T" },
34
  "passes": [
35
  {
36
  "id": "main",
37
  "name": "Cos",
38
  "shader": "unary-scalar.wgsl.jinja",
39
+ "derive": { "itemsPerInvocation": 4 },
40
+ "bindings": ["x_scalar", "y_scalar", "params_2_unary"],
41
  "dispatch": {
42
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
43
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
50
  "bindings": {
51
  "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
52
  "y": { "buffer": "storage", "elementType": "$vectorScalar" },
53
+ "params_unary": { "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }], "name": "params" },
54
+ "x_scalar": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
55
+ "y_scalar": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
56
+ "params_2_unary": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
 
 
 
 
 
 
 
 
57
  }
58
  }
build/webgpu/metadata.json CHANGED
@@ -1,22 +1,22 @@
1
  {
2
  "name": "ai.onnx.Cos",
3
- "id": "_ai_onnx_cos_webgpu_b8cb940",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "hZB5P81WgcOr8dQM++HnJ3SeV7se1ce5xBKtyIc7wAU=",
11
- "manifest.json": "zfXDJZ9MmehTi2+hmO6LmOc6d0G/3MoiI/BBrdGNbec=",
12
  "test.json": "nCOlkbGujYFc4xvpMBa8GIhuE9gsEvUwrN7hbkN056A=",
13
- "unary-scalar.wgsl.jinja": "u8Oh5cH6SLO8tlNHu5CQY6yrH7ZuJnUDAV2p++iIGTc=",
14
- "unary-vec4.wgsl.jinja": "h+TvunU/ZWWcEFtLA63R/pj1TbBBTYdsbcKWfpAWVbY="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
18
  "webgpu": {
19
- "manifestSpec": "2.0",
20
  "variants": { "same_layout_vec4": ["unary-vec4.wgsl.jinja"], "elementwise": ["unary-scalar.wgsl.jinja"] }
21
  }
22
  }
 
1
  {
2
  "name": "ai.onnx.Cos",
3
+ "id": "_ai_onnx_cos_webgpu_43056d9",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "hQhjoj8wcbNpLeBeF0PH4cjXV17E0VP8GqW/qAStCUM=",
11
+ "manifest.json": "0wr6OjAyBgjiVOTZCIrdZwSMK5IxNkj1Yfp5cPvDClg=",
12
  "test.json": "nCOlkbGujYFc4xvpMBa8GIhuE9gsEvUwrN7hbkN056A=",
13
+ "unary-scalar.wgsl.jinja": "KTa10jy8LXkELYpluF6G57rK6wGHITy56lRj1JUko94=",
14
+ "unary-vec4.wgsl.jinja": "Jsd6Oa/zVPucZErigW8VbhRKF5Jrsqfb70Sl9f3Stu0="
15
  }
16
  },
17
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
18
  "webgpu": {
19
+ "manifestSpec": "2.1",
20
  "variants": { "same_layout_vec4": ["unary-vec4.wgsl.jinja"], "elementwise": ["unary-scalar.wgsl.jinja"] }
21
  }
22
  }
build/webgpu/unary-scalar.wgsl.jinja CHANGED
@@ -12,16 +12,68 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
12
  for (var i = begin; i < end; i = i + 1u) {
13
  {%- endmacro %}
14
  {% macro flat_tail_close() %}
15
- }
16
- {% endmacro %}
17
-
18
- // Scalar unary elementwise implementation. Specialization emits only the
19
- // selected operation and any numerical helper it requires.
20
- {% if usesF16 %}
21
- enable f16;
22
- {% endif %}
23
  {{ env.wgsl.resourceDeclarations }}
24
- {% macro emit_reduce_pio2() %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25
  fn reduce_pio2(ax: f32) -> Pio2 {
26
  let ix = bitcast<u32>(ax);
27
  let e = (ix >> 23u) & 0xffu;
@@ -85,97 +137,8 @@ fn reduce_pio2(ax: f32) -> Pio2 {
85
  out.octant = n & 3u;
86
  out.r = phi * PIO2_F;
87
  return out;
88
- }{% endmacro %}
89
- {% macro emit_reduce_pio2_fast() %}
90
- fn reduce_pio2_fast(ax: f32) -> Pio2 {
91
- // Cody-Waite reduction for ordinary magnitudes. Splitting pi/2 keeps the
92
- // residual accurate even when n is thousands; larger inputs use Payne-Hanek.
93
- let nFloat = floor(fma(ax, INV_PIO2_F, 0.5));
94
- var residual = fma(-nFloat, PIO2_HI_F, ax);
95
- residual = fma(-nFloat, PIO2_LO_F, residual);
96
- var out: Pio2;
97
- out.octant = u32(nFloat) & 3u;
98
- out.r = residual;
99
- return out;
100
- }{% endmacro %}
101
- {% macro emit_trig_reduction_support() %}
102
- {% set needPreciseTrigCentered = false %}
103
- {% set needPreciseTrigTwoPi = false %}
104
- // Backend-stable f32 sine/cosine core.
105
- //
106
- // Shader transcendental accuracy is implementation-defined. Each path retains
107
- // its available phase representation, reduces it to [-pi, pi], then evaluates
108
- // explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
109
- // and degree 12 for cosine.
110
- {% if needPreciseTrigCentered %}
111
- const PRECISE_TRIG_PI: f32 = 3.141592653589793;
112
- {% endif %}
113
- {% if needPreciseTrigTwoPi %}
114
- const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
115
- {% endif %}
116
- {% if needPreciseTrigCentered %}
117
- const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
118
- {% endif %}
119
-
120
- // Input must be in [-pi/2, pi/2]. Returns (cos(x), sin(x)).
121
- fn precise_sincos_half_pi(x: f32) -> vec2<f32> {
122
- let x2 = x * x;
123
-
124
- var sinPolynomial = 1.6059043836821613e-10;
125
- sinPolynomial = fma(sinPolynomial, x2, -2.505210838544172e-8);
126
- sinPolynomial = fma(sinPolynomial, x2, 2.7557319223985893e-6);
127
- sinPolynomial = fma(sinPolynomial, x2, -1.984126984126984e-4);
128
- sinPolynomial = fma(sinPolynomial, x2, 8.333333333333333e-3);
129
- sinPolynomial = fma(sinPolynomial, x2, -1.6666666666666666e-1);
130
- let sine = x * fma(sinPolynomial, x2, 1.0);
131
-
132
- var cosPolynomial = 2.08767569878681e-9;
133
- cosPolynomial = fma(cosPolynomial, x2, -2.755731922398589e-7);
134
- cosPolynomial = fma(cosPolynomial, x2, 2.48015873015873e-5);
135
- cosPolynomial = fma(cosPolynomial, x2, -1.388888888888889e-3);
136
- cosPolynomial = fma(cosPolynomial, x2, 4.1666666666666664e-2);
137
- cosPolynomial = fma(cosPolynomial, x2, -5.0e-1);
138
- let cosine = fma(cosPolynomial, x2, 1.0);
139
-
140
- return vec2<f32>(cosine, sine);
141
- }
142
- {% if needPreciseTrigCentered %}
143
-
144
- // Input must be in [-pi, pi]. Returns (cos(x), sin(x)).
145
- fn precise_sincos_centered(x: f32) -> vec2<f32> {
146
- var folded = x;
147
- var cosineSign = 1.0;
148
- if (folded > PRECISE_TRIG_HALF_PI) {
149
- folded = PRECISE_TRIG_PI - folded;
150
- cosineSign = -1.0;
151
- } else if (folded < -PRECISE_TRIG_HALF_PI) {
152
- folded = -PRECISE_TRIG_PI - folded;
153
- cosineSign = -1.0;
154
- }
155
- let value = precise_sincos_half_pi(folded);
156
- return vec2<f32>(cosineSign * value.x, value.y);
157
  }
158
- {% endif %}
159
-
160
-
161
- // 24-bit words of 2/pi used by the exact Payne-Hanek large-input reducer.
162
- const TWO_OVER_PI: array<u32, 14> = array<u32, 14>(
163
- 0xa2f983u, 0x6e4e44u, 0x1529fcu, 0x2757d1u, 0xf534ddu, 0xc0db62u,
164
- 0x95993cu, 0x439041u, 0xfe5163u, 0xabdebbu, 0xc561b7u, 0x246e3au,
165
- 0x424dd2u, 0xe00649u
166
- );
167
- const PIO2_F: f32 = 1.5707963267948966;
168
- const INV_PIO2_F: f32 = 0.6366197723675814;
169
- const PIO2_HI_F: f32 = 1.570796251296997;
170
- const PIO2_LO_F: f32 = 7.549789415861596e-8;
171
- const REDUCE_THRESHOLD: f32 = 1.0e4;
172
-
173
- struct Pio2 { octant: u32, r: f32 };
174
-
175
- {{ emit_reduce_pio2_fast() }}
176
-
177
- {{ emit_reduce_pio2() }}
178
- {%- endmacro %}
179
  {% macro emit_cos_accurate() %}
180
  fn cos_accurate(x: f32) -> f32 {
181
  let ax = abs(x); // cos is even
@@ -195,11 +158,11 @@ fn cos_accurate(x: f32) -> f32 {
195
  case 2u: { return -reduced.x; }
196
  default: { return reduced.y; }
197
  }
198
- }{% endmacro %}
 
199
  {{ emit_trig_reduction_support() }}
200
-
201
  {{ emit_cos_accurate() }}
202
  {{ flat_tail_open() }}
203
  y[i] = {{ scalar }}(cos_accurate(f32(x[i])));
204
- {{ flat_tail_close() -}}
205
  }
 
12
  for (var i = begin; i < end; i = i + 1u) {
13
  {%- endmacro %}
14
  {% macro flat_tail_close() %}
15
+ }{% endmacro %}
16
+ // Scalar unary elementwise kernel: applies the selected operation, with the
17
+ // numerical helpers it uses, to each element.
 
 
 
 
 
18
  {{ env.wgsl.resourceDeclarations }}
19
+
20
+ {% macro emit_trig_reduction_support() %}
21
+ // Backend-stable f32 sine/cosine core.
22
+ //
23
+ // Shader transcendental accuracy is implementation-defined. Each path retains
24
+ // its available phase representation, reduces it to [-pi, pi], then evaluates
25
+ // explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
26
+ // and degree 12 for cosine.
27
+
28
+ // Input must be in [-pi/2, pi/2]. Returns (cos(x), sin(x)).
29
+ fn precise_sincos_half_pi(x: f32) -> vec2<f32> {
30
+ let x2 = x * x;
31
+
32
+ var sinPolynomial = 1.6059043836821613e-10;
33
+ sinPolynomial = fma(sinPolynomial, x2, -2.505210838544172e-8);
34
+ sinPolynomial = fma(sinPolynomial, x2, 2.7557319223985893e-6);
35
+ sinPolynomial = fma(sinPolynomial, x2, -1.984126984126984e-4);
36
+ sinPolynomial = fma(sinPolynomial, x2, 8.333333333333333e-3);
37
+ sinPolynomial = fma(sinPolynomial, x2, -1.6666666666666666e-1);
38
+ let sine = x * fma(sinPolynomial, x2, 1.0);
39
+
40
+ var cosPolynomial = 2.08767569878681e-9;
41
+ cosPolynomial = fma(cosPolynomial, x2, -2.755731922398589e-7);
42
+ cosPolynomial = fma(cosPolynomial, x2, 2.48015873015873e-5);
43
+ cosPolynomial = fma(cosPolynomial, x2, -1.388888888888889e-3);
44
+ cosPolynomial = fma(cosPolynomial, x2, 4.1666666666666664e-2);
45
+ cosPolynomial = fma(cosPolynomial, x2, -5.0e-1);
46
+ let cosine = fma(cosPolynomial, x2, 1.0);
47
+
48
+ return vec2<f32>(cosine, sine);
49
+ }
50
+
51
+ // 24-bit words of 2/pi used by the exact Payne-Hanek large-input reducer.
52
+ const TWO_OVER_PI: array<u32, 14> = array<u32, 14>(
53
+ 0xa2f983u, 0x6e4e44u, 0x1529fcu, 0x2757d1u, 0xf534ddu, 0xc0db62u,
54
+ 0x95993cu, 0x439041u, 0xfe5163u, 0xabdebbu, 0xc561b7u, 0x246e3au,
55
+ 0x424dd2u, 0xe00649u
56
+ );
57
+ const PIO2_F: f32 = 1.5707963267948966;
58
+ const INV_PIO2_F: f32 = 0.6366197723675814;
59
+ const PIO2_HI_F: f32 = 1.570796251296997;
60
+ const PIO2_LO_F: f32 = 7.549789415861596e-8;
61
+ const REDUCE_THRESHOLD: f32 = 1.0e4;
62
+
63
+ struct Pio2 { octant: u32, r: f32 };
64
+
65
+ fn reduce_pio2_fast(ax: f32) -> Pio2 {
66
+ // Cody-Waite reduction for ordinary magnitudes. Splitting pi/2 keeps the
67
+ // residual accurate even when n is thousands; larger inputs use Payne-Hanek.
68
+ let nFloat = floor(fma(ax, INV_PIO2_F, 0.5));
69
+ var residual = fma(-nFloat, PIO2_HI_F, ax);
70
+ residual = fma(-nFloat, PIO2_LO_F, residual);
71
+ var out: Pio2;
72
+ out.octant = u32(nFloat) & 3u;
73
+ out.r = residual;
74
+ return out;
75
+ }
76
+
77
  fn reduce_pio2(ax: f32) -> Pio2 {
78
  let ix = bitcast<u32>(ax);
79
  let e = (ix >> 23u) & 0xffu;
 
137
  out.octant = n & 3u;
138
  out.r = phi * PIO2_F;
139
  return out;
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
140
  }
141
+ {% endmacro %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
142
  {% macro emit_cos_accurate() %}
143
  fn cos_accurate(x: f32) -> f32 {
144
  let ax = abs(x); // cos is even
 
158
  case 2u: { return -reduced.x; }
159
  default: { return reduced.y; }
160
  }
161
+ }
162
+ {% endmacro %}
163
  {{ emit_trig_reduction_support() }}
 
164
  {{ emit_cos_accurate() }}
165
  {{ flat_tail_open() }}
166
  y[i] = {{ scalar }}(cos_accurate(f32(x[i])));
167
+ {{ flat_tail_close() }}
168
  }
build/webgpu/unary-vec4.wgsl.jinja CHANGED
@@ -2,7 +2,63 @@
2
  // component.
3
  {{ env.wgsl.resourceDeclarations }}
4
 
5
- {% macro emit_reduce_pio2() %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
  fn reduce_pio2(ax: f32) -> Pio2 {
7
  let ix = bitcast<u32>(ax);
8
  let e = (ix >> 23u) & 0xffu;
@@ -66,97 +122,8 @@ fn reduce_pio2(ax: f32) -> Pio2 {
66
  out.octant = n & 3u;
67
  out.r = phi * PIO2_F;
68
  return out;
69
- }{% endmacro %}
70
- {% macro emit_reduce_pio2_fast() %}
71
- fn reduce_pio2_fast(ax: f32) -> Pio2 {
72
- // Cody-Waite reduction for ordinary magnitudes. Splitting pi/2 keeps the
73
- // residual accurate even when n is thousands; larger inputs use Payne-Hanek.
74
- let nFloat = floor(fma(ax, INV_PIO2_F, 0.5));
75
- var residual = fma(-nFloat, PIO2_HI_F, ax);
76
- residual = fma(-nFloat, PIO2_LO_F, residual);
77
- var out: Pio2;
78
- out.octant = u32(nFloat) & 3u;
79
- out.r = residual;
80
- return out;
81
- }{% endmacro %}
82
- {% macro emit_trig_reduction_support() %}
83
- {% set needPreciseTrigCentered = false %}
84
- {% set needPreciseTrigTwoPi = false %}
85
- // Backend-stable f32 sine/cosine core.
86
- //
87
- // Shader transcendental accuracy is implementation-defined. Each path retains
88
- // its available phase representation, reduces it to [-pi, pi], then evaluates
89
- // explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
90
- // and degree 12 for cosine.
91
- {% if needPreciseTrigCentered %}
92
- const PRECISE_TRIG_PI: f32 = 3.141592653589793;
93
- {% endif %}
94
- {% if needPreciseTrigTwoPi %}
95
- const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
96
- {% endif %}
97
- {% if needPreciseTrigCentered %}
98
- const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
99
- {% endif %}
100
-
101
- // Input must be in [-pi/2, pi/2]. Returns (cos(x), sin(x)).
102
- fn precise_sincos_half_pi(x: f32) -> vec2<f32> {
103
- let x2 = x * x;
104
-
105
- var sinPolynomial = 1.6059043836821613e-10;
106
- sinPolynomial = fma(sinPolynomial, x2, -2.505210838544172e-8);
107
- sinPolynomial = fma(sinPolynomial, x2, 2.7557319223985893e-6);
108
- sinPolynomial = fma(sinPolynomial, x2, -1.984126984126984e-4);
109
- sinPolynomial = fma(sinPolynomial, x2, 8.333333333333333e-3);
110
- sinPolynomial = fma(sinPolynomial, x2, -1.6666666666666666e-1);
111
- let sine = x * fma(sinPolynomial, x2, 1.0);
112
-
113
- var cosPolynomial = 2.08767569878681e-9;
114
- cosPolynomial = fma(cosPolynomial, x2, -2.755731922398589e-7);
115
- cosPolynomial = fma(cosPolynomial, x2, 2.48015873015873e-5);
116
- cosPolynomial = fma(cosPolynomial, x2, -1.388888888888889e-3);
117
- cosPolynomial = fma(cosPolynomial, x2, 4.1666666666666664e-2);
118
- cosPolynomial = fma(cosPolynomial, x2, -5.0e-1);
119
- let cosine = fma(cosPolynomial, x2, 1.0);
120
-
121
- return vec2<f32>(cosine, sine);
122
- }
123
- {% if needPreciseTrigCentered %}
124
-
125
- // Input must be in [-pi, pi]. Returns (cos(x), sin(x)).
126
- fn precise_sincos_centered(x: f32) -> vec2<f32> {
127
- var folded = x;
128
- var cosineSign = 1.0;
129
- if (folded > PRECISE_TRIG_HALF_PI) {
130
- folded = PRECISE_TRIG_PI - folded;
131
- cosineSign = -1.0;
132
- } else if (folded < -PRECISE_TRIG_HALF_PI) {
133
- folded = -PRECISE_TRIG_PI - folded;
134
- cosineSign = -1.0;
135
- }
136
- let value = precise_sincos_half_pi(folded);
137
- return vec2<f32>(cosineSign * value.x, value.y);
138
  }
139
- {% endif %}
140
-
141
-
142
- // 24-bit words of 2/pi used by the exact Payne-Hanek large-input reducer.
143
- const TWO_OVER_PI: array<u32, 14> = array<u32, 14>(
144
- 0xa2f983u, 0x6e4e44u, 0x1529fcu, 0x2757d1u, 0xf534ddu, 0xc0db62u,
145
- 0x95993cu, 0x439041u, 0xfe5163u, 0xabdebbu, 0xc561b7u, 0x246e3au,
146
- 0x424dd2u, 0xe00649u
147
- );
148
- const PIO2_F: f32 = 1.5707963267948966;
149
- const INV_PIO2_F: f32 = 0.6366197723675814;
150
- const PIO2_HI_F: f32 = 1.570796251296997;
151
- const PIO2_LO_F: f32 = 7.549789415861596e-8;
152
- const REDUCE_THRESHOLD: f32 = 1.0e4;
153
-
154
- struct Pio2 { octant: u32, r: f32 };
155
-
156
- {{ emit_reduce_pio2_fast() }}
157
-
158
- {{ emit_reduce_pio2() }}
159
- {%- endmacro %}
160
  {% macro emit_cos_accurate() %}
161
  fn cos_accurate(x: f32) -> f32 {
162
  let ax = abs(x); // cos is even
@@ -176,37 +143,46 @@ fn cos_accurate(x: f32) -> f32 {
176
  case 2u: { return -reduced.x; }
177
  default: { return reduced.y; }
178
  }
179
- }{% endmacro %}
 
180
  {{ emit_trig_reduction_support() }}
181
  {{ emit_cos_accurate() }}
182
-
183
- {% set vec4PerThread = vec4PerThread %}
 
 
 
 
 
 
 
 
184
  {% if vec4PerThread > 1 %}
185
  const ITEMS: u32 = {{ vec4PerThread }}u;
186
- {% endif %}
187
 
 
188
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
189
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
190
- // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
191
- // per-axis dispatch fold width (the dispatch caps x and spills the rest into y).
192
  {% if vec4PerThread > 1 %}
193
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
194
  // access consecutive words on every step, while each lane can keep several
195
  // independent loads in flight.
196
- let tid = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
197
  let span = (params.count + ITEMS - 1u) / ITEMS;
 
 
 
 
 
 
198
  for (var j = 0u; j < ITEMS; j = j + 1u) {
199
  let i = tid + j * span;
200
  if (i >= params.count) {
201
  break;
202
  }
203
  {% else %}
204
- let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
205
- if (i >= params.count) {
206
- return;
207
- }
208
  {% endif %}
209
-
210
  let xv = x[i];
211
  let fv = vec4<f32>(xv);
212
  y[i] = {{ vectorScalar }}(vec4<f32>(cos_accurate(fv.x), cos_accurate(fv.y), cos_accurate(fv.z), cos_accurate(fv.w)));
 
2
  // component.
3
  {{ env.wgsl.resourceDeclarations }}
4
 
5
+ {% macro emit_trig_reduction_support() %}
6
+ // Backend-stable f32 sine/cosine core.
7
+ //
8
+ // Shader transcendental accuracy is implementation-defined. Each path retains
9
+ // its available phase representation, reduces it to [-pi, pi], then evaluates
10
+ // explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
11
+ // and degree 12 for cosine.
12
+
13
+ // Input must be in [-pi/2, pi/2]. Returns (cos(x), sin(x)).
14
+ fn precise_sincos_half_pi(x: f32) -> vec2<f32> {
15
+ let x2 = x * x;
16
+
17
+ var sinPolynomial = 1.6059043836821613e-10;
18
+ sinPolynomial = fma(sinPolynomial, x2, -2.505210838544172e-8);
19
+ sinPolynomial = fma(sinPolynomial, x2, 2.7557319223985893e-6);
20
+ sinPolynomial = fma(sinPolynomial, x2, -1.984126984126984e-4);
21
+ sinPolynomial = fma(sinPolynomial, x2, 8.333333333333333e-3);
22
+ sinPolynomial = fma(sinPolynomial, x2, -1.6666666666666666e-1);
23
+ let sine = x * fma(sinPolynomial, x2, 1.0);
24
+
25
+ var cosPolynomial = 2.08767569878681e-9;
26
+ cosPolynomial = fma(cosPolynomial, x2, -2.755731922398589e-7);
27
+ cosPolynomial = fma(cosPolynomial, x2, 2.48015873015873e-5);
28
+ cosPolynomial = fma(cosPolynomial, x2, -1.388888888888889e-3);
29
+ cosPolynomial = fma(cosPolynomial, x2, 4.1666666666666664e-2);
30
+ cosPolynomial = fma(cosPolynomial, x2, -5.0e-1);
31
+ let cosine = fma(cosPolynomial, x2, 1.0);
32
+
33
+ return vec2<f32>(cosine, sine);
34
+ }
35
+
36
+ // 24-bit words of 2/pi used by the exact Payne-Hanek large-input reducer.
37
+ const TWO_OVER_PI: array<u32, 14> = array<u32, 14>(
38
+ 0xa2f983u, 0x6e4e44u, 0x1529fcu, 0x2757d1u, 0xf534ddu, 0xc0db62u,
39
+ 0x95993cu, 0x439041u, 0xfe5163u, 0xabdebbu, 0xc561b7u, 0x246e3au,
40
+ 0x424dd2u, 0xe00649u
41
+ );
42
+ const PIO2_F: f32 = 1.5707963267948966;
43
+ const INV_PIO2_F: f32 = 0.6366197723675814;
44
+ const PIO2_HI_F: f32 = 1.570796251296997;
45
+ const PIO2_LO_F: f32 = 7.549789415861596e-8;
46
+ const REDUCE_THRESHOLD: f32 = 1.0e4;
47
+
48
+ struct Pio2 { octant: u32, r: f32 };
49
+
50
+ fn reduce_pio2_fast(ax: f32) -> Pio2 {
51
+ // Cody-Waite reduction for ordinary magnitudes. Splitting pi/2 keeps the
52
+ // residual accurate even when n is thousands; larger inputs use Payne-Hanek.
53
+ let nFloat = floor(fma(ax, INV_PIO2_F, 0.5));
54
+ var residual = fma(-nFloat, PIO2_HI_F, ax);
55
+ residual = fma(-nFloat, PIO2_LO_F, residual);
56
+ var out: Pio2;
57
+ out.octant = u32(nFloat) & 3u;
58
+ out.r = residual;
59
+ return out;
60
+ }
61
+
62
  fn reduce_pio2(ax: f32) -> Pio2 {
63
  let ix = bitcast<u32>(ax);
64
  let e = (ix >> 23u) & 0xffu;
 
122
  out.octant = n & 3u;
123
  out.r = phi * PIO2_F;
124
  return out;
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
125
  }
126
+ {% endmacro %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
127
  {% macro emit_cos_accurate() %}
128
  fn cos_accurate(x: f32) -> f32 {
129
  let ax = abs(x); // cos is even
 
143
  case 2u: { return -reduced.x; }
144
  default: { return reduced.y; }
145
  }
146
+ }
147
+ {% endmacro %}
148
  {{ emit_trig_reduction_support() }}
149
  {{ emit_cos_accurate() }}
150
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
151
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
152
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
153
+ // per-axis workgroup fold width.
154
+ {% if bound == "" %}
155
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% else %}
156
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
157
+ if ({{ name }} >= {{ bound }}) {
158
+ return;
159
+ }{% endif %}{% endmacro %}
160
  {% if vec4PerThread > 1 %}
161
  const ITEMS: u32 = {{ vec4PerThread }}u;
 
162
 
163
+ {% endif %}
164
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
165
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
 
166
  {% if vec4PerThread > 1 %}
167
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
168
  // access consecutive words on every step, while each lane can keep several
169
  // independent loads in flight.
170
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "tid", "") }}
171
  let span = (params.count + ITEMS - 1u) / ITEMS;
172
+ // Lanes at or past span exist only because the dispatch rounds up to whole
173
+ // workgroups. Lane span + k would start on lane k's second group and rewrite
174
+ // up to ITEMS - 1 groups another lane already stored.
175
+ if (tid >= span) {
176
+ return;
177
+ }
178
  for (var j = 0u; j < ITEMS; j = j + 1u) {
179
  let i = tid + j * span;
180
  if (i >= params.count) {
181
  break;
182
  }
183
  {% else %}
184
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
 
 
 
185
  {% endif %}
 
186
  let xv = x[i];
187
  let fv = vec4<f32>(xv);
188
  y[i] = {{ vectorScalar }}(vec4<f32>(cos_accurate(fv.x), cos_accurate(fv.y), cos_accurate(fv.z), cos_accurate(fv.w)));