Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.LogSoftmax

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
f7ca565
·
verified ·
1 Parent(s): 8bfadd9

sync c928d21e6cc1

Browse files
Files changed (5) hide show
  1. README.md +1 -1
  2. build/webgpu/manifest.json +39 -22
  3. build/webgpu/metadata.json +5 -5
  4. build/webgpu/softmax-online.wgsl.jinja +6 -2
  5. build/webgpu/{softmax-strided-online-lane.wgsl.jinja → softmax-strided-online-coop.wgsl.jinja} +69 -25
README.md CHANGED
@@ -55,7 +55,7 @@ Default values (overridable per request):
55
  - [`softmax-online.wgsl.jinja`](build/webgpu/softmax-online.wgsl.jinja)
56
  - [`softmax-row-stage-strided-vec4.wgsl.jinja`](build/webgpu/softmax-row-stage-strided-vec4.wgsl.jinja)
57
  - [`softmax-row-stage.wgsl.jinja`](build/webgpu/softmax-row-stage.wgsl.jinja)
58
- - [`softmax-strided-online-lane.wgsl.jinja`](build/webgpu/softmax-strided-online-lane.wgsl.jinja)
59
  - [`softmax-strided-packed4-tail.wgsl.jinja`](build/webgpu/softmax-strided-packed4-tail.wgsl.jinja)
60
 
61
  ## Use with `@huggingface/kernels`
 
55
  - [`softmax-online.wgsl.jinja`](build/webgpu/softmax-online.wgsl.jinja)
56
  - [`softmax-row-stage-strided-vec4.wgsl.jinja`](build/webgpu/softmax-row-stage-strided-vec4.wgsl.jinja)
57
  - [`softmax-row-stage.wgsl.jinja`](build/webgpu/softmax-row-stage.wgsl.jinja)
58
+ - [`softmax-strided-online-coop.wgsl.jinja`](build/webgpu/softmax-strided-online-coop.wgsl.jinja)
59
  - [`softmax-strided-packed4-tail.wgsl.jinja`](build/webgpu/softmax-strided-packed4-tail.wgsl.jinja)
60
 
61
  ## Use with `@huggingface/kernels`
build/webgpu/manifest.json CHANGED
@@ -24,6 +24,8 @@
24
  },
25
  "tunables": {
26
  "WORKGROUP_SIZE": 256,
 
 
27
  "BLOCK_COLS": 1024,
28
  "LONGROW_SPLIT_MIN_COLS": 65536,
29
  "STRIDED_CAPACITY_MIN_ROWS": 12288
@@ -44,7 +46,9 @@
44
  "axisDispatchFits": "ceilDiv(axisRows, device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
45
  "axisThreadDispatchFits": "ceilDiv(ceilDiv(axisRows, tunables.WORKGROUP_SIZE), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
46
  "lastAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 1 and (attrs.axis == -1 or attrs.axis == ranks.x - 1) and f16Ok(dtypes.T)",
47
- "stridedAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 2 and attrs.axis + ranks.x >= 0 and attrs.axis < ranks.x and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)"
 
 
48
  },
49
  "bindingSets": {
50
  "longrowBlockStats": [
@@ -271,7 +275,8 @@
271
  "combineSubgroups": "device.features.has(\"subgroups\") and dim(shapes.x, ranks.x - 1) >= 1024",
272
  "scalar": "dtypes.T",
273
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
274
- "usesF16": "dtypes.T == \"f16\""
 
275
  },
276
  "passes": [
277
  {
@@ -290,7 +295,8 @@
290
  "constants": {
291
  "combineSubgroups": "device.features.has(\"subgroups\")",
292
  "scalar": "dtypes.T",
293
- "usesF16": "dtypes.T == \"f16\""
 
294
  },
295
  "passes": [
296
  {
@@ -339,7 +345,8 @@
339
  "combineSubgroups": false,
340
  "scalar": "dtypes.T",
341
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
342
- "usesF16": "dtypes.T == \"f16\""
 
343
  },
344
  "passes": [
345
  {
@@ -356,7 +363,12 @@
356
  "priority": 14,
357
  "tunables": { "WORKGROUP_SIZE": 32 },
358
  "when": ["dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) <= 32", "ceilDiv(rows(shapes.x, ranks.x - 1), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension", "lastAxisOk"],
359
- "constants": { "combineSubgroups": false, "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
 
 
 
 
 
360
  "passes": [
361
  {
362
  "id": "main",
@@ -478,6 +490,28 @@
478
  }
479
  ]
480
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
481
  {
482
  "id": "strided_online_packed4_tail",
483
  "priority": 18,
@@ -549,23 +583,6 @@
549
  "dispatch": { "threads": "numel(shapes.x)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
550
  }
551
  ]
552
- },
553
- {
554
- "id": "strided_online_capacity_fallback",
555
- "priority": -1,
556
- "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisThreadDispatchFits"],
557
- "passes": [
558
- {
559
- "id": "main",
560
- "name": "LogSoftmax.StridedOnlineCapacityFallback",
561
- "source": {
562
- "shader": "softmax-strided-online-lane.wgsl.jinja",
563
- "inputs": { "xShape": "shapes.x", "axis": "axis", "op": "\"logsoftmax\"" }
564
- },
565
- "bindings": "stridedOnline",
566
- "dispatch": { "threads": "axisRows", "workgroupSize": "tunables.WORKGROUP_SIZE" }
567
- }
568
- ]
569
  }
570
  ]
571
  }
 
24
  },
25
  "tunables": {
26
  "WORKGROUP_SIZE": 256,
27
+ "ROW_MIN_WORKGROUP": 64,
28
+ "ROW_COLS_PER_LANE": 8,
29
  "BLOCK_COLS": 1024,
30
  "LONGROW_SPLIT_MIN_COLS": 65536,
31
  "STRIDED_CAPACITY_MIN_ROWS": 12288
 
46
  "axisDispatchFits": "ceilDiv(axisRows, device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
47
  "axisThreadDispatchFits": "ceilDiv(ceilDiv(axisRows, tunables.WORKGROUP_SIZE), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
48
  "lastAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 1 and (attrs.axis == -1 or attrs.axis == ranks.x - 1) and f16Ok(dtypes.T)",
49
+ "stridedAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 2 and attrs.axis + ranks.x >= 0 and attrs.axis < ranks.x and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)",
50
+ "stridedCoopTileCols": "min(128, max(8, pow2ceil(ceilDiv(axisOuter * axisInner, 256))))",
51
+ "rowWorkgroupSize": "min(tunables.WORKGROUP_SIZE, max(tunables.ROW_MIN_WORKGROUP, pow2ceil(ceilDiv(lastCols, tunables.ROW_COLS_PER_LANE))))"
52
  },
53
  "bindingSets": {
54
  "longrowBlockStats": [
 
275
  "combineSubgroups": "device.features.has(\"subgroups\") and dim(shapes.x, ranks.x - 1) >= 1024",
276
  "scalar": "dtypes.T",
277
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
278
+ "usesF16": "dtypes.T == \"f16\"",
279
+ "workgroupSize": "rowWorkgroupSize"
280
  },
281
  "passes": [
282
  {
 
295
  "constants": {
296
  "combineSubgroups": "device.features.has(\"subgroups\")",
297
  "scalar": "dtypes.T",
298
+ "usesF16": "dtypes.T == \"f16\"",
299
+ "workgroupSize": "rowWorkgroupSize"
300
  },
301
  "passes": [
302
  {
 
345
  "combineSubgroups": false,
346
  "scalar": "dtypes.T",
347
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
348
+ "usesF16": "dtypes.T == \"f16\"",
349
+ "workgroupSize": "tunables.WORKGROUP_SIZE"
350
  },
351
  "passes": [
352
  {
 
363
  "priority": 14,
364
  "tunables": { "WORKGROUP_SIZE": 32 },
365
  "when": ["dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) <= 32", "ceilDiv(rows(shapes.x, ranks.x - 1), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension", "lastAxisOk"],
366
+ "constants": {
367
+ "combineSubgroups": false,
368
+ "scalar": "dtypes.T",
369
+ "usesF16": "dtypes.T == \"f16\"",
370
+ "workgroupSize": "tunables.WORKGROUP_SIZE"
371
+ },
372
  "passes": [
373
  {
374
  "id": "main",
 
490
  }
491
  ]
492
  },
493
+ {
494
+ "id": "strided_online_coop",
495
+ "priority": 19,
496
+ "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisThreadDispatchFits"],
497
+ "passes": [
498
+ {
499
+ "id": "main",
500
+ "name": "LogSoftmax.StridedOnlineCoop",
501
+ "source": {
502
+ "shader": "softmax-strided-online-coop.wgsl.jinja",
503
+ "inputs": {
504
+ "xShape": "shapes.x",
505
+ "axis": "axis",
506
+ "op": "\"logsoftmax\"",
507
+ "tileCols": "stridedCoopTileCols"
508
+ }
509
+ },
510
+ "bindings": "stridedOnline",
511
+ "dispatch": { "workgroups": "ceilDiv(axisOuter * axisInner, stridedCoopTileCols)" }
512
+ }
513
+ ]
514
+ },
515
  {
516
  "id": "strided_online_packed4_tail",
517
  "priority": 18,
 
583
  "dispatch": { "threads": "numel(shapes.x)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
584
  }
585
  ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
586
  }
587
  ]
588
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.LogSoftmax",
3
- "id": "_ai_onnx_logsoftmax_webgpu_0804fe4",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,19 +8,19 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "rqFaX5XmJwQz0QQ1ET5FRMCuYEMotdcOJR+z6Y57a/0=",
11
- "manifest.json": "QDbHBdv1drmn3+Mqxdqx24xatM7c38mQxNZrDzDPmM8=",
12
  "softmax-longrow-normalize.wgsl.jinja": "AamVdoPDaWPQ3sk5loxtcuNu7oftC5J7hsPwLRDA7M0=",
13
  "softmax-longrow-stats.wgsl.jinja": "K0PI+9cmRgjJsOOcBRyugN0Dm7yxBy0xvyPamEpCgxo=",
14
  "softmax-normalize.wgsl.jinja": "eekHyttLyBk2X/mG54dUS2kIc5KycR3grH6000pqh1A=",
15
  "softmax-online-packed-rows.wgsl.jinja": "94Xmz8Mtj/ji7XpIQUpsDdO6uShKJtWknzEI0Er+sT8=",
16
- "softmax-online.wgsl.jinja": "qL/179d83HTFuyCUkfHqE0MBbtFMOEl8lzAH7aWmUXA=",
17
  "softmax-row-stage-strided-vec4.wgsl.jinja": "l1byT6xoLl6Q7kSPd/uws6X5NsS2PwlSv2TBMCioVZA=",
18
  "softmax-row-stage.wgsl.jinja": "dxlM6y02DwXeLeavHDiq70h0NCRZgz+bniCltyEgRX4=",
19
- "softmax-strided-online-lane.wgsl.jinja": "BTLdB6LamPLNJsNmsllObl/9Y/Cx1EbuwdqeYpcppXg=",
20
  "softmax-strided-packed4-tail.wgsl.jinja": "z7+TYeybaE/yXds/ODe4wMyosZeIrCtTfUmZsA2eIXo=",
21
  "test.json": "uvZ0xakkYLGTptboBDjo8khdaQQ2h+QWiu4Z/wLWmfQ="
22
  }
23
  },
24
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
25
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.LogSoftmax" }
26
  }
 
1
  {
2
  "name": "ai.onnx.LogSoftmax",
3
+ "id": "_ai_onnx_logsoftmax_webgpu_00d56e1",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "rqFaX5XmJwQz0QQ1ET5FRMCuYEMotdcOJR+z6Y57a/0=",
11
+ "manifest.json": "m0YO7FppNrFhQRavPhUB1ruckiRUpkx9OaKbZWrZpdc=",
12
  "softmax-longrow-normalize.wgsl.jinja": "AamVdoPDaWPQ3sk5loxtcuNu7oftC5J7hsPwLRDA7M0=",
13
  "softmax-longrow-stats.wgsl.jinja": "K0PI+9cmRgjJsOOcBRyugN0Dm7yxBy0xvyPamEpCgxo=",
14
  "softmax-normalize.wgsl.jinja": "eekHyttLyBk2X/mG54dUS2kIc5KycR3grH6000pqh1A=",
15
  "softmax-online-packed-rows.wgsl.jinja": "94Xmz8Mtj/ji7XpIQUpsDdO6uShKJtWknzEI0Er+sT8=",
16
+ "softmax-online.wgsl.jinja": "s3GyunMaKcOIzUKhI2TUkb650oKqApFW3EASngfE/Ro=",
17
  "softmax-row-stage-strided-vec4.wgsl.jinja": "l1byT6xoLl6Q7kSPd/uws6X5NsS2PwlSv2TBMCioVZA=",
18
  "softmax-row-stage.wgsl.jinja": "dxlM6y02DwXeLeavHDiq70h0NCRZgz+bniCltyEgRX4=",
19
+ "softmax-strided-online-coop.wgsl.jinja": "bZkpOM+peRv2E0z0iT4V6bJTEsCHBpiF1ScPwENDFlU=",
20
  "softmax-strided-packed4-tail.wgsl.jinja": "z7+TYeybaE/yXds/ODe4wMyosZeIrCtTfUmZsA2eIXo=",
21
  "test.json": "uvZ0xakkYLGTptboBDjo8khdaQQ2h+QWiu4Z/wLWmfQ="
22
  }
23
  },
24
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
25
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.LogSoftmax" }
26
  }
build/webgpu/softmax-online.wgsl.jinja CHANGED
@@ -12,7 +12,11 @@ enable subgroups;
12
  // the shared online-softmax combiner) and a fused second sweep writes
13
  // the result. The epilogue is either softmax or log-softmax, and vec4 row access
14
  // is used when rows are vectorized.
15
- const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 
 
 
 
16
  // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
17
  // `m - m` finite so an empty lane / all--inf row contributes the exact
18
  // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
@@ -174,7 +178,7 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
174
  {% endif %}
175
 
176
 
177
- @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
178
  fn main(@builtin(workgroup_id) wid: vec3<u32>,
179
  @builtin(num_workgroups) nwg: vec3<u32>,
180
  @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups %},
 
12
  // the shared online-softmax combiner) and a fused second sweep writes
13
  // the result. The epilogue is either softmax or log-softmax, and vec4 row access
14
  // is used when rows are vectorized.
15
+ //
16
+ // `workgroupSize` is a variant constant, not the op-wide tunable: the merge tree
17
+ // costs a barrier per level whatever the row length, so a short row wants fewer,
18
+ // busier lanes and a long row wants the widest workgroup the device allows.
19
+ const WG: u32 = {{ workgroupSize }}u;
20
  // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
21
  // `m - m` finite so an empty lane / all--inf row contributes the exact
22
  // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
 
178
  {% endif %}
179
 
180
 
181
+ @compute @workgroup_size({{ workgroupSize }})
182
  fn main(@builtin(workgroup_id) wid: vec3<u32>,
183
  @builtin(num_workgroups) nwg: vec3<u32>,
184
  @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups %},
build/webgpu/{softmax-strided-online-lane.wgsl.jinja → softmax-strided-online-coop.wgsl.jinja} RENAMED
@@ -1,8 +1,13 @@
1
- // Scratch-free online softmax/log-softmax for a non-last reduction axis. One
2
- // invocation owns one (outer, inner) row, so neighboring lanes make coalesced
3
- // reads at every reduction step. This is both the large-inner fast path and the
4
- // capacity-safe fallback when row-stat scratch cannot be allocated.
 
 
 
5
  {% set operation = source.op if source.op is defined else "softmax" %}
 
 
6
  {% if T == "f16" %}
7
  enable f16;
8
  {% endif %}
@@ -84,33 +89,72 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
84
  return select(value - maxValue, 0.0, equalFiniteMax);
85
  }
86
 
87
- const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
88
- @compute @workgroup_size(WGS)
89
- fn main(@builtin(global_invocation_id) gid: vec3<u32>,
90
- @builtin(num_workgroups) nwg: vec3<u32>) {
91
- let row = gid.x + gid.y * nwg.x * WGS;
92
- if (row >= STRIDED_ROWS) { return; }
93
 
94
- var row_max = -{{ "FLT_MAX" if operation == "logsoftmax" else "F32_MAX" }};
95
- var row_sum = 0.0;
96
- for (var k = 0u; k < REDUCE_SIZE; k += 1u) {
97
- let value = f32(x[element_offset(row, k)]);
98
- if (value > row_max) {
99
- row_sum = row_sum * exp(row_max - value) + 1.0;
100
- row_max = value;
101
- } else {
102
- row_sum += exp(value - row_max);
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
103
  }
 
 
104
  }
 
105
 
 
 
106
  let sum_bits = bitcast<u32>(row_sum);
107
  let sum_is_nan = (sum_bits & 0x7f800000u) == 0x7f800000u && (sum_bits & 0x007fffffu) != 0u;
108
- for (var k = 0u; k < REDUCE_SIZE; k += 1u) {
109
- let offset = element_offset(row, k);
110
- var out = shifted_value(f32(x[offset]), row_max) - log(row_sum);
111
- if (row_max > FLT_MAX || sum_is_nan) {
112
- out = bitcast<f32>(bitcast<u32>(row_sum) | 0x7fc00000u);
 
 
 
113
  }
114
- y[offset] = {{ T }}(out);
115
  }
116
  }
 
1
+ // Cooperative online softmax/log-softmax for a non-last reduction axis. A
2
+ // workgroup owns TILE_COLS consecutive (outer, inner) rows and splits the
3
+ // reduction walk across AXIS_LANES invocations per row, so short-row shapes
4
+ // still fill the machine and every load stays coalesced across the column
5
+ // lanes regardless of whether the inner extent divides four. Per-lane online
6
+ // partials merge in workgroup memory with the same update form as the serial
7
+ // walk, so NaN and +Inf rows reproduce the single-lane kernel's outputs.
8
  {% set operation = source.op if source.op is defined else "softmax" %}
9
+ {% set tileCols = source.tileCols %}
10
+ {% set axisLanes = (256 / tileCols) | int %}
11
  {% if T == "f16" %}
12
  enable f16;
13
  {% endif %}
 
89
  return select(value - maxValue, 0.0, equalFiniteMax);
90
  }
91
 
92
+ const TILE_COLS: u32 = {{ tileCols }}u;
93
+ const AXIS_LANES: u32 = {{ axisLanes }}u;
94
+ {% set negInit = "-FLT_MAX" if operation == "logsoftmax" else "-F32_MAX" %}
 
 
 
95
 
96
+ var<workgroup> part_max: array<f32, TILE_COLS * AXIS_LANES>;
97
+ var<workgroup> part_sum: array<f32, TILE_COLS * AXIS_LANES>;
98
+
99
+ @compute @workgroup_size(TILE_COLS, AXIS_LANES, 1)
100
+ fn main(@builtin(workgroup_id) wg: vec3<u32>,
101
+ @builtin(num_workgroups) nwg: vec3<u32>,
102
+ @builtin(local_invocation_id) lid: vec3<u32>) {
103
+ let tile = wg.x + wg.y * nwg.x;
104
+ let row = tile * TILE_COLS + lid.x;
105
+ let row_live = row < STRIDED_ROWS;
106
+
107
+ var lane_max = {{ negInit }};
108
+ var lane_sum = 0.0;
109
+ if (row_live) {
110
+ for (var k = lid.y; k < REDUCE_SIZE; k += AXIS_LANES) {
111
+ let value = f32(x[element_offset(row, k)]);
112
+ if (value > lane_max) {
113
+ lane_sum = lane_sum * exp(lane_max - value) + 1.0;
114
+ lane_max = value;
115
+ } else {
116
+ lane_sum += exp(value - lane_max);
117
+ }
118
+ }
119
+ }
120
+ part_max[lid.y * TILE_COLS + lid.x] = lane_max;
121
+ part_sum[lid.y * TILE_COLS + lid.x] = lane_sum;
122
+ workgroupBarrier();
123
+
124
+ // One lane per column merges the axis partials serially, in the same
125
+ // (max, rescale, add) form as the walk above: an empty partial
126
+ // ({{ negInit }}, 0) is the identity, a NaN partial poisons the sum the
127
+ // way a NaN element does, and +Inf carries into the merged max.
128
+ if (lid.y == 0u) {
129
+ var merged_max = part_max[lid.x];
130
+ var merged_sum = part_sum[lid.x];
131
+ for (var t = 1u; t < AXIS_LANES; t += 1u) {
132
+ let m2 = part_max[t * TILE_COLS + lid.x];
133
+ let s2 = part_sum[t * TILE_COLS + lid.x];
134
+ if (m2 > merged_max) {
135
+ merged_sum = merged_sum * exp(merged_max - m2) + s2;
136
+ merged_max = m2;
137
+ } else {
138
+ merged_sum += s2 * exp(m2 - merged_max);
139
+ }
140
  }
141
+ part_max[lid.x] = merged_max;
142
+ part_sum[lid.x] = merged_sum;
143
  }
144
+ workgroupBarrier();
145
 
146
+ let row_max = part_max[lid.x];
147
+ let row_sum = part_sum[lid.x];
148
  let sum_bits = bitcast<u32>(row_sum);
149
  let sum_is_nan = (sum_bits & 0x7f800000u) == 0x7f800000u && (sum_bits & 0x007fffffu) != 0u;
150
+ if (row_live) {
151
+ for (var k = lid.y; k < REDUCE_SIZE; k += AXIS_LANES) {
152
+ let offset = element_offset(row, k);
153
+ var out = shifted_value(f32(x[offset]), row_max) - log(row_sum);
154
+ if (row_max > FLT_MAX || sum_is_nan) {
155
+ out = bitcast<f32>(bitcast<u32>(row_sum) | 0x7fc00000u);
156
+ }
157
+ y[offset] = {{ T }}(out);
158
  }
 
159
  }
160
  }