sync 6fdf6301e2bb

Browse files

Files changed (6) hide show

README.md +2 -2
build/webgpu/manifest.json +8 -20
build/webgpu/metadata.json +7 -7
build/webgpu/test.json +2 -2
build/webgpu/unary-scalar.wgsl.jinja +98 -115
build/webgpu/unary-vec4.wgsl.jinja +90 -118

README.md CHANGED Viewed

@@ -39,14 +39,14 @@ See the [ONNX `Tan` spec](https://onnx.ai/onnx/operators/onnx__Tan.html) for the
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
 - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
 - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/manifest.json CHANGED Viewed

@@ -11,15 +11,11 @@
       "id": "same_layout_vec4",
       "priority": 20,
       "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
-      "derive": {
-        "scalar": "dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
-        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\""
-      },
       "passes": [
         {
           "id": "main",
-          "name": "Tan.vec4",
           "shader": "unary-vec4.wgsl.jinja",
           "derive": {
             "op": "\"tan\"",
@@ -37,14 +33,14 @@
     {
       "id": "elementwise",
       "when": ["numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
-      "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "Tan",
           "shader": "unary-scalar.wgsl.jinja",
           "derive": { "op": "\"tan\"", "itemsPerInvocation": 4 },
-          "bindings": ["x_2", "y_2", "params_2_unary"],
           "dispatch": {
             "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -57,17 +53,9 @@
   "bindings": {
     "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
     "y": { "buffer": "storage", "elementType": "$vectorScalar" },
-    "params_unary": {
-      "buffer": "uniform",
-      "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }],
-      "name": "params"
-    },
-    "x_2": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
-    "y_2": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
-    "params_2_unary": {
-      "buffer": "uniform",
-      "name": "params",
-      "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
-    }
   }
 }

       "id": "same_layout_vec4",
       "priority": 20,
       "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
+      "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
+          "name": "Tan.Vec4",
           "shader": "unary-vec4.wgsl.jinja",
           "derive": {
             "op": "\"tan\"",
     {
       "id": "elementwise",
       "when": ["numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
+      "derive": { "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Tan",
           "shader": "unary-scalar.wgsl.jinja",
           "derive": { "op": "\"tan\"", "itemsPerInvocation": 4 },
+          "bindings": ["x_scalar", "y_scalar", "params_2_unary"],
           "dispatch": {
             "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
   "bindings": {
     "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
     "y": { "buffer": "storage", "elementType": "$vectorScalar" },
+    "params_unary": { "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }], "name": "params" },
+    "x_scalar": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
+    "y_scalar": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
+    "params_2_unary": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
   }
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "ai.onnx.Tan",
-  "id": "_ai_onnx_tan_webgpu_ea47a00",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -8,15 +8,15 @@
     "algorithm": "sha256",
     "files": {
       "bench.json": "UdERmcCeq6Elm16WZSEs0wCOLlz5Gg64dtn0Wviv3DU=",
-      "manifest.json": "P0urm4KuqNzj8+8CUgeB6Y75va7RIxhENjBqrgu+1FM=",
-      "test.json": "nr23CSBkeNglZRyckIYhrEX1pE4F9/leWOqRXvPJAPo=",
-      "unary-scalar.wgsl.jinja": "EfW2qA5VRzSLMb5zpcr5BhhOJA7T9hBytO00uoqt5kE=",
-      "unary-vec4.wgsl.jinja": "dnEOe8UV0joGHXg9Tl9Ejlycadsg9UxkcGwqtfTFI1c="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": { "same_layout_vec4": ["unary-vec4.wgsl.jinja"], "elementwise": ["unary-scalar.wgsl.jinja"] }
   }
 }

 {
   "name": "ai.onnx.Tan",
+  "id": "_ai_onnx_tan_webgpu_60ac5b0",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "algorithm": "sha256",
     "files": {
       "bench.json": "UdERmcCeq6Elm16WZSEs0wCOLlz5Gg64dtn0Wviv3DU=",
+      "manifest.json": "ExZsBKPXQx2r4Ij2LuPuAlwp+XH245/1agGxiZBfrPs=",
+      "test.json": "ZWt1K92ZoMAO2Y73HdN+RBDk+F1zgL8uSgKfLSrsjeg=",
+      "unary-scalar.wgsl.jinja": "FLarw3ZNNJ8cFjwPvaZgM9cxPxCCwWEga17KNxRl7Zg=",
+      "unary-vec4.wgsl.jinja": "CUYHc97trSgiXBBrwy4I+xtoLNSsNxj5qWP+fZi7D4M="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": { "same_layout_vec4": ["unary-vec4.wgsl.jinja"], "elementwise": ["unary-scalar.wgsl.jinja"] }
   }
 }

build/webgpu/test.json CHANGED Viewed

@@ -213,7 +213,7 @@
       },
       "outputs": { "y": { "dtype": "float16", "shape": [7], "tolerance": 0.05, "relTolerance": 0.01 } },
       "provenance": {
-        "notes": "Float16 arguments at the top of the range force argument reduction in float32; the tolerance is scaled to the result magnitude with headroom over the reduction residue."
       }
     },
     {
@@ -227,7 +227,7 @@
       },
       "outputs": { "y": { "dtype": "float16", "shape": [8], "tolerance": 0.05, "relTolerance": 0.01 } },
       "provenance": {
-        "notes": "Large float16 arguments on the vec4 lane path force argument reduction in float32; the tolerance is scaled to the result magnitude with headroom over the reduction residue."
       }
     }
   ]

       },
       "outputs": { "y": { "dtype": "float16", "shape": [7], "tolerance": 0.05, "relTolerance": 0.01 } },
       "provenance": {
+        "notes": "Float16 arguments near the format's maximum (65,504) require argument reduction in float32 before evaluating tan; tolerance is scaled to the result magnitude with headroom for the reduction residue."
       }
     },
     {
       },
       "outputs": { "y": { "dtype": "float16", "shape": [8], "tolerance": 0.05, "relTolerance": 0.01 } },
       "provenance": {
+        "notes": "Float16 arguments up to 40,000, mixed with values near 1, require argument reduction in float32 before evaluating tan across two vec4 words; tolerance is scaled to the result magnitude with headroom for the reduction residue."
       }
     }
   ]

build/webgpu/unary-scalar.wgsl.jinja CHANGED Viewed

@@ -12,29 +12,77 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   for (var i = begin; i < end; i = i + 1u) {
 {%- endmacro %}
 {% macro flat_tail_close() %}
-  }
-{% endmacro %}
-// Scalar unary elementwise implementation. Specialization emits only the
-// selected operation and any numerical helper it requires.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-{% macro emit_is_nan_f32() %}
-fn is_nan_f32(value: f32) -> bool {
-  let bits = bitcast<u32>(value);
-  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
-}{% endmacro %}
-{% macro emit_is_inf_f32() %}
 fn is_inf_f32(value: f32) -> bool {
   return (bitcast<u32>(value) & 0x7fffffffu) == 0x7f800000u;
-}{% endmacro %}
-{% macro emit_nan_from_inf_f32() %}
 fn nan_from_inf_f32(value: f32) -> f32 {
   return bitcast<f32>(bitcast<u32>(value) | 0x00400000u);
-}{% endmacro %}
-{% macro emit_reduce_pio2() %}
 fn reduce_pio2(ax: f32) -> Pio2 {
   let ix = bitcast<u32>(ax);
   let e = (ix >> 23u) & 0xffu;
@@ -98,97 +146,32 @@ fn reduce_pio2(ax: f32) -> Pio2 {
   out.octant = n & 3u;
   out.r = phi * PIO2_F;
   return out;
-}{% endmacro %}
-{% macro emit_reduce_pio2_fast() %}
-fn reduce_pio2_fast(ax: f32) -> Pio2 {
-  // Cody-Waite reduction for ordinary magnitudes. Splitting pi/2 keeps the
-  // residual accurate even when n is thousands; larger inputs use Payne-Hanek.
-  let nFloat = floor(fma(ax, INV_PIO2_F, 0.5));
-  var residual = fma(-nFloat, PIO2_HI_F, ax);
-  residual = fma(-nFloat, PIO2_LO_F, residual);
-  var out: Pio2;
-  out.octant = u32(nFloat) & 3u;
-  out.r = residual;
-  return out;
-}{% endmacro %}
-{% macro emit_trig_reduction_support() %}
-{% set needPreciseTrigCentered = false %}
-{% set needPreciseTrigTwoPi = false %}
-// Backend-stable f32 sine/cosine core.
-//
-// Shader transcendental accuracy is implementation-defined. Each path retains
-// its available phase representation, reduces it to [-pi, pi], then evaluates
-// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
-// and degree 12 for cosine.
-{% if needPreciseTrigCentered %}
-const PRECISE_TRIG_PI: f32 = 3.141592653589793;
-{% endif %}
-{% if needPreciseTrigTwoPi %}
-const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
-{% endif %}
-{% if needPreciseTrigCentered %}
-const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
-{% endif %}
-// Input must be in [-pi/2, pi/2]. Returns (cos(x), sin(x)).
-fn precise_sincos_half_pi(x: f32) -> vec2<f32> {
-  let x2 = x * x;
-  var sinPolynomial = 1.6059043836821613e-10;
-  sinPolynomial = fma(sinPolynomial, x2, -2.505210838544172e-8);
-  sinPolynomial = fma(sinPolynomial, x2, 2.7557319223985893e-6);
-  sinPolynomial = fma(sinPolynomial, x2, -1.984126984126984e-4);
-  sinPolynomial = fma(sinPolynomial, x2, 8.333333333333333e-3);
-  sinPolynomial = fma(sinPolynomial, x2, -1.6666666666666666e-1);
-  let sine = x * fma(sinPolynomial, x2, 1.0);
-  var cosPolynomial = 2.08767569878681e-9;
-  cosPolynomial = fma(cosPolynomial, x2, -2.755731922398589e-7);
-  cosPolynomial = fma(cosPolynomial, x2, 2.48015873015873e-5);
-  cosPolynomial = fma(cosPolynomial, x2, -1.388888888888889e-3);
-  cosPolynomial = fma(cosPolynomial, x2, 4.1666666666666664e-2);
-  cosPolynomial = fma(cosPolynomial, x2, -5.0e-1);
-  let cosine = fma(cosPolynomial, x2, 1.0);
-  return vec2<f32>(cosine, sine);
 }
-{% if needPreciseTrigCentered %}
-// Input must be in [-pi, pi]. Returns (cos(x), sin(x)).
-fn precise_sincos_centered(x: f32) -> vec2<f32> {
-  var folded = x;
-  var cosineSign = 1.0;
-  if (folded > PRECISE_TRIG_HALF_PI) {
-    folded = PRECISE_TRIG_PI - folded;
-    cosineSign = -1.0;
-  } else if (folded < -PRECISE_TRIG_HALF_PI) {
-    folded = -PRECISE_TRIG_PI - folded;
-    cosineSign = -1.0;
   }
-  let value = precise_sincos_half_pi(folded);
-  return vec2<f32>(cosineSign * value.x, value.y);
 }
-{% endif %}
-// 24-bit words of 2/pi used by the exact Payne-Hanek large-input reducer.
-const TWO_OVER_PI: array<u32, 14> = array<u32, 14>(
-  0xa2f983u, 0x6e4e44u, 0x1529fcu, 0x2757d1u, 0xf534ddu, 0xc0db62u,
-  0x95993cu, 0x439041u, 0xfe5163u, 0xabdebbu, 0xc561b7u, 0x246e3au,
-  0x424dd2u, 0xe00649u
-);
-const PIO2_F: f32 = 1.5707963267948966;
-const INV_PIO2_F: f32 = 0.6366197723675814;
-const PIO2_HI_F: f32 = 1.570796251296997;
-const PIO2_LO_F: f32 = 7.549789415861596e-8;
-const REDUCE_THRESHOLD: f32 = 1.0e4;
-struct Pio2 { octant: u32, r: f32 };
-{{ emit_reduce_pio2_fast() }}
-{{ emit_reduce_pio2() }}
-{%- endmacro %}
 {% macro emit_tan_accurate() %}
 fn tan_accurate(x: f32) -> f32 {
   if (is_nan_f32(x)) {
@@ -214,17 +197,17 @@ fn tan_accurate(x: f32) -> f32 {
   let tn = sn / cs;
   // tan is odd: transfer x's sign bit so tan(-0.0) stays -0.0 (x < 0.0 misses it).
   return bitcast<f32>(bitcast<u32>(tn) ^ (bitcast<u32>(x) & 0x80000000u));
-}{% endmacro %}
 {{ emit_trig_reduction_support() }}
-{{ emit_is_nan_f32() }}
-{{ emit_is_inf_f32() }}
-{{ emit_nan_from_inf_f32() }}
-{{ emit_tan_accurate() }}
 {{ flat_tail_open() }}
   y[i] = {{ scalar }}(tan_accurate(f32(x[i])));
-{{ flat_tail_close() -}}
 }

   for (var i = begin; i < end; i = i + 1u) {
 {%- endmacro %}
 {% macro flat_tail_close() %}
+  }{% endmacro %}
+// Scalar unary elementwise kernel: applies the selected operation, with the
+// numerical helpers it uses, to each element.
 {{ env.wgsl.resourceDeclarations }}
+{% macro emit_inf_classifiers() %}
 fn is_inf_f32(value: f32) -> bool {
   return (bitcast<u32>(value) & 0x7fffffffu) == 0x7f800000u;
+}
 fn nan_from_inf_f32(value: f32) -> f32 {
   return bitcast<f32>(bitcast<u32>(value) | 0x00400000u);
+}
+{% endmacro %}
+{% macro emit_trig_reduction_support() %}
+// Backend-stable f32 sine/cosine core.
+//
+// Shader transcendental accuracy is implementation-defined. Each path retains
+// its available phase representation, reduces it to [-pi, pi], then evaluates
+// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
+// and degree 12 for cosine.
+// Input must be in [-pi/2, pi/2]. Returns (cos(x), sin(x)).
+fn precise_sincos_half_pi(x: f32) -> vec2<f32> {
+  let x2 = x * x;
+  var sinPolynomial = 1.6059043836821613e-10;
+  sinPolynomial = fma(sinPolynomial, x2, -2.505210838544172e-8);
+  sinPolynomial = fma(sinPolynomial, x2, 2.7557319223985893e-6);
+  sinPolynomial = fma(sinPolynomial, x2, -1.984126984126984e-4);
+  sinPolynomial = fma(sinPolynomial, x2, 8.333333333333333e-3);
+  sinPolynomial = fma(sinPolynomial, x2, -1.6666666666666666e-1);
+  let sine = x * fma(sinPolynomial, x2, 1.0);
+  var cosPolynomial = 2.08767569878681e-9;
+  cosPolynomial = fma(cosPolynomial, x2, -2.755731922398589e-7);
+  cosPolynomial = fma(cosPolynomial, x2, 2.48015873015873e-5);
+  cosPolynomial = fma(cosPolynomial, x2, -1.388888888888889e-3);
+  cosPolynomial = fma(cosPolynomial, x2, 4.1666666666666664e-2);
+  cosPolynomial = fma(cosPolynomial, x2, -5.0e-1);
+  let cosine = fma(cosPolynomial, x2, 1.0);
+  return vec2<f32>(cosine, sine);
+}
+// 24-bit words of 2/pi used by the exact Payne-Hanek large-input reducer.
+const TWO_OVER_PI: array<u32, 14> = array<u32, 14>(
+  0xa2f983u, 0x6e4e44u, 0x1529fcu, 0x2757d1u, 0xf534ddu, 0xc0db62u,
+  0x95993cu, 0x439041u, 0xfe5163u, 0xabdebbu, 0xc561b7u, 0x246e3au,
+  0x424dd2u, 0xe00649u
+);
+const PIO2_F: f32 = 1.5707963267948966;
+const INV_PIO2_F: f32 = 0.6366197723675814;
+const PIO2_HI_F: f32 = 1.570796251296997;
+const PIO2_LO_F: f32 = 7.549789415861596e-8;
+const REDUCE_THRESHOLD: f32 = 1.0e4;
+struct Pio2 { octant: u32, r: f32 };
+fn reduce_pio2_fast(ax: f32) -> Pio2 {
+  // Cody-Waite reduction for ordinary magnitudes. Splitting pi/2 keeps the
+  // residual accurate even when n is thousands; larger inputs use Payne-Hanek.
+  let nFloat = floor(fma(ax, INV_PIO2_F, 0.5));
+  var residual = fma(-nFloat, PIO2_HI_F, ax);
+  residual = fma(-nFloat, PIO2_LO_F, residual);
+  var out: Pio2;
+  out.octant = u32(nFloat) & 3u;
+  out.r = residual;
+  return out;
+}
 fn reduce_pio2(ax: f32) -> Pio2 {
   let ix = bitcast<u32>(ax);
   let e = (ix >> 23u) & 0xffu;
   out.octant = n & 3u;
   out.r = phi * PIO2_F;
   return out;
 }
+{% endmacro %}
+{% macro emit_sin_accurate() %}
+fn sin_accurate(x: f32) -> f32 {
+  if (is_nan_f32(x)) {
+    return x;
+  }
+  if (is_inf_f32(x)) {
+    return nan_from_inf_f32(x);
+  }
+  let ax = abs(x);
+  var red: Pio2;
+  if (ax < REDUCE_THRESHOLD) { red = reduce_pio2_fast(ax); }
+  else { red = reduce_pio2(ax); }
+  let reduced = precise_sincos_half_pi(red.r);
+  var s: f32;
+  switch (red.octant) {
+    case 0u: { s = reduced.y; }
+    case 1u: { s = reduced.x; }
+    case 2u: { s = -reduced.y; }
+    default: { s = -reduced.x; }
   }
+  // sin is odd: transfer x's sign bit so sin(-0.0) stays -0.0 (x < 0.0 misses it).
+  return bitcast<f32>(bitcast<u32>(s) ^ (bitcast<u32>(x) & 0x80000000u));
 }
+{% endmacro %}
 {% macro emit_tan_accurate() %}
 fn tan_accurate(x: f32) -> f32 {
   if (is_nan_f32(x)) {
   let tn = sn / cs;
   // tan is odd: transfer x's sign bit so tan(-0.0) stays -0.0 (x < 0.0 misses it).
   return bitcast<f32>(bitcast<u32>(tn) ^ (bitcast<u32>(x) & 0x80000000u));
+}
+{% endmacro %}
 {{ emit_trig_reduction_support() }}
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{{ emit_inf_classifiers() }}
+{{ emit_sin_accurate() if op == "sin" else emit_tan_accurate() }}
 {{ flat_tail_open() }}
   y[i] = {{ scalar }}(tan_accurate(f32(x[i])));
+{{ flat_tail_close() }}
 }

build/webgpu/unary-vec4.wgsl.jinja CHANGED Viewed

@@ -2,20 +2,72 @@
 // component.
 {{ env.wgsl.resourceDeclarations }}
-{% macro emit_is_nan_f32() %}
-fn is_nan_f32(value: f32) -> bool {
-  let bits = bitcast<u32>(value);
-  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
-}{% endmacro %}
-{% macro emit_is_inf_f32() %}
 fn is_inf_f32(value: f32) -> bool {
   return (bitcast<u32>(value) & 0x7fffffffu) == 0x7f800000u;
-}{% endmacro %}
-{% macro emit_nan_from_inf_f32() %}
 fn nan_from_inf_f32(value: f32) -> f32 {
   return bitcast<f32>(bitcast<u32>(value) | 0x00400000u);
-}{% endmacro %}
-{% macro emit_reduce_pio2() %}
 fn reduce_pio2(ax: f32) -> Pio2 {
   let ix = bitcast<u32>(ax);
   let e = (ix >> 23u) & 0xffu;
@@ -79,97 +131,8 @@ fn reduce_pio2(ax: f32) -> Pio2 {
   out.octant = n & 3u;
   out.r = phi * PIO2_F;
   return out;
-}{% endmacro %}
-{% macro emit_reduce_pio2_fast() %}
-fn reduce_pio2_fast(ax: f32) -> Pio2 {
-  // Cody-Waite reduction for ordinary magnitudes. Splitting pi/2 keeps the
-  // residual accurate even when n is thousands; larger inputs use Payne-Hanek.
-  let nFloat = floor(fma(ax, INV_PIO2_F, 0.5));
-  var residual = fma(-nFloat, PIO2_HI_F, ax);
-  residual = fma(-nFloat, PIO2_LO_F, residual);
-  var out: Pio2;
-  out.octant = u32(nFloat) & 3u;
-  out.r = residual;
-  return out;
-}{% endmacro %}
-{% macro emit_trig_reduction_support() %}
-{% set needPreciseTrigCentered = false %}
-{% set needPreciseTrigTwoPi = false %}
-// Backend-stable f32 sine/cosine core.
-//
-// Shader transcendental accuracy is implementation-defined. Each path retains
-// its available phase representation, reduces it to [-pi, pi], then evaluates
-// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
-// and degree 12 for cosine.
-{% if needPreciseTrigCentered %}
-const PRECISE_TRIG_PI: f32 = 3.141592653589793;
-{% endif %}
-{% if needPreciseTrigTwoPi %}
-const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
-{% endif %}
-{% if needPreciseTrigCentered %}
-const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
-{% endif %}
-// Input must be in [-pi/2, pi/2]. Returns (cos(x), sin(x)).
-fn precise_sincos_half_pi(x: f32) -> vec2<f32> {
-  let x2 = x * x;
-  var sinPolynomial = 1.6059043836821613e-10;
-  sinPolynomial = fma(sinPolynomial, x2, -2.505210838544172e-8);
-  sinPolynomial = fma(sinPolynomial, x2, 2.7557319223985893e-6);
-  sinPolynomial = fma(sinPolynomial, x2, -1.984126984126984e-4);
-  sinPolynomial = fma(sinPolynomial, x2, 8.333333333333333e-3);
-  sinPolynomial = fma(sinPolynomial, x2, -1.6666666666666666e-1);
-  let sine = x * fma(sinPolynomial, x2, 1.0);
-  var cosPolynomial = 2.08767569878681e-9;
-  cosPolynomial = fma(cosPolynomial, x2, -2.755731922398589e-7);
-  cosPolynomial = fma(cosPolynomial, x2, 2.48015873015873e-5);
-  cosPolynomial = fma(cosPolynomial, x2, -1.388888888888889e-3);
-  cosPolynomial = fma(cosPolynomial, x2, 4.1666666666666664e-2);
-  cosPolynomial = fma(cosPolynomial, x2, -5.0e-1);
-  let cosine = fma(cosPolynomial, x2, 1.0);
-  return vec2<f32>(cosine, sine);
-}
-{% if needPreciseTrigCentered %}
-// Input must be in [-pi, pi]. Returns (cos(x), sin(x)).
-fn precise_sincos_centered(x: f32) -> vec2<f32> {
-  var folded = x;
-  var cosineSign = 1.0;
-  if (folded > PRECISE_TRIG_HALF_PI) {
-    folded = PRECISE_TRIG_PI - folded;
-    cosineSign = -1.0;
-  } else if (folded < -PRECISE_TRIG_HALF_PI) {
-    folded = -PRECISE_TRIG_PI - folded;
-    cosineSign = -1.0;
-  }
-  let value = precise_sincos_half_pi(folded);
-  return vec2<f32>(cosineSign * value.x, value.y);
 }
-{% endif %}
-// 24-bit words of 2/pi used by the exact Payne-Hanek large-input reducer.
-const TWO_OVER_PI: array<u32, 14> = array<u32, 14>(
-  0xa2f983u, 0x6e4e44u, 0x1529fcu, 0x2757d1u, 0xf534ddu, 0xc0db62u,
-  0x95993cu, 0x439041u, 0xfe5163u, 0xabdebbu, 0xc561b7u, 0x246e3au,
-  0x424dd2u, 0xe00649u
-);
-const PIO2_F: f32 = 1.5707963267948966;
-const INV_PIO2_F: f32 = 0.6366197723675814;
-const PIO2_HI_F: f32 = 1.570796251296997;
-const PIO2_LO_F: f32 = 7.549789415861596e-8;
-const REDUCE_THRESHOLD: f32 = 1.0e4;
-struct Pio2 { octant: u32, r: f32 };
-{{ emit_reduce_pio2_fast() }}
-{{ emit_reduce_pio2() }}
-{%- endmacro %}
 {% macro emit_tan_accurate() %}
 fn tan_accurate(x: f32) -> f32 {
   if (is_nan_f32(x)) {
@@ -195,43 +158,52 @@ fn tan_accurate(x: f32) -> f32 {
   let tn = sn / cs;
   // tan is odd: transfer x's sign bit so tan(-0.0) stays -0.0 (x < 0.0 misses it).
   return bitcast<f32>(bitcast<u32>(tn) ^ (bitcast<u32>(x) & 0x80000000u));
-}{% endmacro %}
-{{ emit_is_nan_f32() }}
-{{ emit_trig_reduction_support() }}
-{{ emit_is_inf_f32() }}
-{{ emit_nan_from_inf_f32() }}
 {{ emit_tan_accurate() }}
-{% set vec4PerThread = vec4PerThread %}
 {% if vec4PerThread > 1 %}
 const ITEMS: u32 = {{ vec4PerThread }}u;
-{% endif %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
-  // per-axis dispatch fold width (the dispatch caps x and spills the rest into y).
 {% if vec4PerThread > 1 %}
   // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
   // access consecutive words on every step, while each lane can keep several
   // independent loads in flight.
-  let tid = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   let span = (params.count + ITEMS - 1u) / ITEMS;
   for (var j = 0u; j < ITEMS; j = j + 1u) {
     let i = tid + j * span;
     if (i >= params.count) {
       break;
     }
 {% else %}
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (i >= params.count) {
-    return;
-  }
 {% endif %}
   let xv = x[i];
   let fv = vec4<f32>(xv);
   y[i] = {{ vectorScalar }}(vec4<f32>(tan_accurate(fv.x), tan_accurate(fv.y), tan_accurate(fv.z), tan_accurate(fv.w)));

 // component.
 {{ env.wgsl.resourceDeclarations }}
+{% macro emit_inf_classifiers() %}
 fn is_inf_f32(value: f32) -> bool {
   return (bitcast<u32>(value) & 0x7fffffffu) == 0x7f800000u;
+}
 fn nan_from_inf_f32(value: f32) -> f32 {
   return bitcast<f32>(bitcast<u32>(value) | 0x00400000u);
+}
+{% endmacro %}
+{% macro emit_trig_reduction_support() %}
+// Backend-stable f32 sine/cosine core.
+//
+// Shader transcendental accuracy is implementation-defined. Each path retains
+// its available phase representation, reduces it to [-pi, pi], then evaluates
+// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
+// and degree 12 for cosine.
+// Input must be in [-pi/2, pi/2]. Returns (cos(x), sin(x)).
+fn precise_sincos_half_pi(x: f32) -> vec2<f32> {
+  let x2 = x * x;
+  var sinPolynomial = 1.6059043836821613e-10;
+  sinPolynomial = fma(sinPolynomial, x2, -2.505210838544172e-8);
+  sinPolynomial = fma(sinPolynomial, x2, 2.7557319223985893e-6);
+  sinPolynomial = fma(sinPolynomial, x2, -1.984126984126984e-4);
+  sinPolynomial = fma(sinPolynomial, x2, 8.333333333333333e-3);
+  sinPolynomial = fma(sinPolynomial, x2, -1.6666666666666666e-1);
+  let sine = x * fma(sinPolynomial, x2, 1.0);
+  var cosPolynomial = 2.08767569878681e-9;
+  cosPolynomial = fma(cosPolynomial, x2, -2.755731922398589e-7);
+  cosPolynomial = fma(cosPolynomial, x2, 2.48015873015873e-5);
+  cosPolynomial = fma(cosPolynomial, x2, -1.388888888888889e-3);
+  cosPolynomial = fma(cosPolynomial, x2, 4.1666666666666664e-2);
+  cosPolynomial = fma(cosPolynomial, x2, -5.0e-1);
+  let cosine = fma(cosPolynomial, x2, 1.0);
+  return vec2<f32>(cosine, sine);
+}
+// 24-bit words of 2/pi used by the exact Payne-Hanek large-input reducer.
+const TWO_OVER_PI: array<u32, 14> = array<u32, 14>(
+  0xa2f983u, 0x6e4e44u, 0x1529fcu, 0x2757d1u, 0xf534ddu, 0xc0db62u,
+  0x95993cu, 0x439041u, 0xfe5163u, 0xabdebbu, 0xc561b7u, 0x246e3au,
+  0x424dd2u, 0xe00649u
+);
+const PIO2_F: f32 = 1.5707963267948966;
+const INV_PIO2_F: f32 = 0.6366197723675814;
+const PIO2_HI_F: f32 = 1.570796251296997;
+const PIO2_LO_F: f32 = 7.549789415861596e-8;
+const REDUCE_THRESHOLD: f32 = 1.0e4;
+struct Pio2 { octant: u32, r: f32 };
+fn reduce_pio2_fast(ax: f32) -> Pio2 {
+  // Cody-Waite reduction for ordinary magnitudes. Splitting pi/2 keeps the
+  // residual accurate even when n is thousands; larger inputs use Payne-Hanek.
+  let nFloat = floor(fma(ax, INV_PIO2_F, 0.5));
+  var residual = fma(-nFloat, PIO2_HI_F, ax);
+  residual = fma(-nFloat, PIO2_LO_F, residual);
+  var out: Pio2;
+  out.octant = u32(nFloat) & 3u;
+  out.r = residual;
+  return out;
+}
 fn reduce_pio2(ax: f32) -> Pio2 {
   let ix = bitcast<u32>(ax);
   let e = (ix >> 23u) & 0xffu;
   out.octant = n & 3u;
   out.r = phi * PIO2_F;
   return out;
 }
+{% endmacro %}
 {% macro emit_tan_accurate() %}
 fn tan_accurate(x: f32) -> f32 {
   if (is_nan_f32(x)) {
   let tn = sn / cs;
   // tan is odd: transfer x's sign bit so tan(-0.0) stays -0.0 (x < 0.0 misses it).
   return bitcast<f32>(bitcast<u32>(tn) ^ (bitcast<u32>(x) & 0x80000000u));
+}
+{% endmacro %}
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{{ emit_trig_reduction_support() }}
+{{ emit_inf_classifiers() }}
 {{ emit_tan_accurate() }}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+{% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endif %}{% endmacro %}
 {% if vec4PerThread > 1 %}
 const ITEMS: u32 = {{ vec4PerThread }}u;
+{% endif %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {% if vec4PerThread > 1 %}
   // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
   // access consecutive words on every step, while each lane can keep several
   // independent loads in flight.
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "tid", "") }}
   let span = (params.count + ITEMS - 1u) / ITEMS;
+  // Lanes at or past span exist only because the dispatch rounds up to whole
+  // workgroups. Lane span + k would start on lane k's second group and rewrite
+  // up to ITEMS - 1 groups another lane already stored.
+  if (tid >= span) {
+    return;
+  }
   for (var j = 0u; j < ITEMS; j = j + 1u) {
     let i = tid + j * span;
     if (i >= params.count) {
       break;
     }
 {% else %}
+{{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
 {% endif %}
   let xv = x[i];
   let fv = vec4<f32>(xv);
   y[i] = {{ vectorScalar }}(vec4<f32>(tan_accurate(fv.x), tan_accurate(fv.y), tan_accurate(fv.z), tan_accurate(fv.w)));