jz 0.5.1 → 0.7.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (86) hide show
  1. package/README.md +315 -318
  2. package/bench/README.md +369 -0
  3. package/bench/bench.svg +102 -0
  4. package/cli.js +104 -30
  5. package/dist/interop.js +1 -0
  6. package/dist/jz.js +6024 -0
  7. package/index.d.ts +126 -0
  8. package/index.js +362 -84
  9. package/interop.js +159 -186
  10. package/jz.svg +5 -0
  11. package/jzify/arguments.js +97 -0
  12. package/jzify/bundler.js +382 -0
  13. package/jzify/classes.js +328 -0
  14. package/jzify/hoist-vars.js +177 -0
  15. package/jzify/index.js +51 -0
  16. package/jzify/names.js +37 -0
  17. package/jzify/switch.js +106 -0
  18. package/jzify/transform.js +349 -0
  19. package/layout.js +184 -0
  20. package/module/array.js +377 -176
  21. package/module/collection.js +639 -145
  22. package/module/console.js +55 -43
  23. package/module/core.js +264 -153
  24. package/module/date.js +15 -3
  25. package/module/function.js +73 -6
  26. package/module/index.js +2 -1
  27. package/module/json.js +226 -61
  28. package/module/math.js +551 -187
  29. package/module/number.js +327 -60
  30. package/module/object.js +474 -184
  31. package/module/regex.js +255 -25
  32. package/module/schema.js +24 -6
  33. package/module/simd.js +117 -0
  34. package/module/string.js +621 -226
  35. package/module/symbol.js +1 -1
  36. package/module/timer.js +9 -14
  37. package/module/typedarray.js +459 -73
  38. package/package.json +58 -14
  39. package/src/abi/index.js +39 -23
  40. package/src/abi/string.js +38 -41
  41. package/src/ast.js +460 -0
  42. package/src/autoload.js +29 -24
  43. package/src/bridge.js +111 -0
  44. package/src/compile/analyze-scans.js +824 -0
  45. package/src/compile/analyze.js +1600 -0
  46. package/src/compile/emit-assign.js +411 -0
  47. package/src/compile/emit.js +3512 -0
  48. package/src/compile/flow-types.js +103 -0
  49. package/src/{compile.js → compile/index.js} +778 -128
  50. package/src/{infer.js → compile/infer.js} +62 -98
  51. package/src/compile/loop-divmod.js +155 -0
  52. package/src/{narrow.js → compile/narrow.js} +409 -106
  53. package/src/compile/peel-stencil.js +261 -0
  54. package/src/compile/plan/advise.js +370 -0
  55. package/src/compile/plan/common.js +150 -0
  56. package/src/compile/plan/index.js +122 -0
  57. package/src/compile/plan/inline.js +682 -0
  58. package/src/compile/plan/literals.js +1199 -0
  59. package/src/compile/plan/loops.js +472 -0
  60. package/src/compile/plan/scope.js +649 -0
  61. package/src/compile/program-facts.js +423 -0
  62. package/src/ctx.js +220 -64
  63. package/src/ir.js +589 -172
  64. package/src/kind-traits.js +132 -0
  65. package/src/kind.js +524 -0
  66. package/src/op-policy.js +57 -0
  67. package/src/{optimize.js → optimize/index.js} +1432 -473
  68. package/src/optimize/vectorize.js +4660 -0
  69. package/src/param-reps.js +65 -0
  70. package/src/parse.js +44 -0
  71. package/src/{prepare.js → prepare/index.js} +649 -205
  72. package/src/prepare/lift-iife.js +149 -0
  73. package/src/reps.js +116 -0
  74. package/src/resolve.js +12 -3
  75. package/src/static.js +208 -0
  76. package/src/type.js +651 -0
  77. package/src/{assemble.js → wat/assemble.js} +275 -55
  78. package/src/wat/optimize.js +3938 -0
  79. package/transform.js +21 -0
  80. package/wasi.js +47 -5
  81. package/src/analyze.js +0 -3818
  82. package/src/emit.js +0 -3040
  83. package/src/jzify.js +0 -1580
  84. package/src/plan.js +0 -2132
  85. package/src/vectorize.js +0 -1088
  86. /package/src/{codegen.js → wat/codegen.js} +0 -0
package/src/vectorize.js DELETED
@@ -1,1088 +0,0 @@
1
- import { findBodyStart } from './ir.js'
2
-
3
- /**
4
- * Lane-local SIMD-128 vectorizer.
5
- *
6
- * Recognizes inner loops of shape:
7
- * for (let i = 0; i < N; i++) arr[i] = f(arr[i], …)
8
- * where every body op is "lane-pure" — its k-th lane output depends only
9
- * on k-th lane inputs. Lifts the body to SIMD-128, prefixed before the
10
- * original (now tail) loop. Original loop runs the remainder.
11
- *
12
- * Design:
13
- * • Lane-purity is a structural property, not a benchmark match. The op
14
- * whitelist is the single source of truth (one entry per (lane-type, op)).
15
- * • Lift is mechanical. The recognizer either matches the structure — in
16
- * which case lifting is unambiguous — or skips. No bench-specific
17
- * heuristics.
18
- * • Tail loop is the original WAT, untouched. If anything regresses the
19
- * SIMD recognizer just doesn't match, never miscompiles.
20
- *
21
- * Match conditions:
22
- * 1. (block $brk (loop $L (br_if $brk !cond) BODY (i = i+1) (br $L)))
23
- * 2. cond is `(i32.lt_s i BOUND)` or `i32.lt_u`; BOUND is loop-invariant.
24
- * 3. All loads/stores in BODY use address `(add base (shl i K))` where
25
- * base is loop-invariant and K matches the elem stride. Optional
26
- * enclosing `local.tee` is allowed (and reused).
27
- * 4. All loads share the same opcode → defines lane type.
28
- * 5. All other ops in BODY are in the lane-pure whitelist for that type.
29
- * 6. Each non-induction local in BODY is either purely loop-invariant
30
- * (only read) or purely lane-local (first action is a write). Never
31
- * both — that's a loop-carried scalar (reduction / stencil) → bail.
32
- *
33
- * Lift produces, before the original block:
34
- * (local.set $__simd_bound{N} (i32.and BOUND (i32.const ~(LANES-1))))
35
- * (block $__simd_brk{N}
36
- * (loop $__simd_loop{N}
37
- * (br_if $__simd_brk{N} (i32.eqz (i32.lt_s i $__simd_bound{N})))
38
- * <body lifted op-by-op; lane-local locals routed to v128 shadows>
39
- * (local.set $i (i32.add i (i32.const LANES)))
40
- * (br $__simd_loop{N})))
41
- *
42
- * The original block runs immediately after with i pre-advanced; its own
43
- * `i < BOUND` guard handles the tail.
44
- */
45
-
46
-
47
-
48
- const isArr = Array.isArray
49
-
50
- const exprEq = (a, b) => JSON.stringify(a) === JSON.stringify(b)
51
- const localGetName = n => isArr(n) && n[0] === 'local.get' && typeof n[1] === 'string' ? n[1] : null
52
- const f64Zero = n => isArr(n) && n[0] === 'f64.const' && Number(n[1]) === 0
53
-
54
- const matchF64MulLocals = n => {
55
- if (!isArr(n) || n[0] !== 'f64.mul') return null
56
- const a = localGetName(n[1])
57
- const b = localGetName(n[2])
58
- return a && b ? [a, b] : null
59
- }
60
-
61
- const matchAccumStep = (n, acc) => {
62
- if (!isArr(n) || n[0] !== 'local.set' || n[1] !== acc) return null
63
- const e = n[2]
64
- if (!isArr(e) || e[0] !== 'f64.add') return null
65
- if (localGetName(e[1]) === acc) return matchF64MulLocals(e[2])
66
- if (localGetName(e[2]) === acc) return matchF64MulLocals(e[1])
67
- return null
68
- }
69
-
70
- const matchDotStore = (n, acc) => {
71
- if (!isArr(n) || n[0] !== 'local.set' || typeof n[1] !== 'string') return null
72
- const e = n[2]
73
- if (localGetName(e) === acc) return { out: n[1], addend: null }
74
- if (!isArr(e) || e[0] !== 'f64.add') return null
75
- if (localGetName(e[1]) === acc) return { out: n[1], addend: e[2] }
76
- if (localGetName(e[2]) === acc) return { out: n[1], addend: e[1] }
77
- return null
78
- }
79
-
80
- const matchF64DotSeq = (stmts, i) => {
81
- const reset = stmts[i]
82
- if (!isArr(reset) || reset[0] !== 'local.set' || typeof reset[1] !== 'string' || !f64Zero(reset[2])) return null
83
- const acc = reset[1]
84
- const left = [], right = []
85
- for (let k = 0; k < 4; k++) {
86
- const pair = matchAccumStep(stmts[i + 1 + k], acc)
87
- if (!pair) return null
88
- left.push(pair[0])
89
- right.push(pair[1])
90
- }
91
- const store = matchDotStore(stmts[i + 5], acc)
92
- return store ? { end: i + 6, acc, left, right, ...store } : null
93
- }
94
-
95
- const f64x2Pair = (lo, hi) => ['f64x2.replace_lane', 1, ['f64x2.splat', ['local.get', lo]], ['local.get', hi]]
96
-
97
- const dotPairExpr = (a, pairs) => {
98
- let expr = ['f64x2.mul', ['f64x2.splat', ['local.get', a[0]]], pairs[0]]
99
- for (let i = 1; i < 4; i++) {
100
- expr = ['f64x2.add', expr, ['f64x2.mul', ['f64x2.splat', ['local.get', a[i]]], pairs[i]]]
101
- }
102
- return expr
103
- }
104
-
105
- const vectorizeStraightLineF64DotPairsIn = (node, fnLocals, freshIdRef, newLocalDecls) => {
106
- if (!isArr(node)) return
107
- for (let i = 0; i < node.length; i++) {
108
- const child = node[i]
109
- if (isArr(child)) vectorizeStraightLineF64DotPairsIn(child, fnLocals, freshIdRef, newLocalDecls)
110
- }
111
- const addendTemps = new Map()
112
- const pairTemps = new Map()
113
- for (let i = 0; i < node.length;) {
114
- const a = matchF64DotSeq(node, i)
115
- if (!a) { i++; continue }
116
- const b = matchF64DotSeq(node, a.end)
117
- if (!b || a.acc !== b.acc || !exprEq(a.left, b.left) || !exprEq(a.addend, b.addend) ||
118
- fnLocals.get(a.out) !== 'f64' || fnLocals.get(b.out) !== 'f64') {
119
- i++
120
- continue
121
- }
122
- const v = `$__dot2_${freshIdRef.next++}`
123
- newLocalDecls.push(['local', v, 'v128'])
124
- fnLocals.set(v, 'v128')
125
- let prefix = []
126
- let addend = a.addend
127
- if (addend) {
128
- const key = JSON.stringify(addend)
129
- let tmp = addendTemps.get(key)
130
- if (!tmp) {
131
- tmp = `$__dotadd_${freshIdRef.next++}`
132
- addendTemps.set(key, tmp)
133
- newLocalDecls.push(['local', tmp, 'f64'])
134
- fnLocals.set(tmp, 'f64')
135
- prefix = [['local.set', tmp, addend]]
136
- }
137
- addend = ['local.get', tmp]
138
- }
139
- const pairs = []
140
- for (let k = 0; k < 4; k++) {
141
- const key = `${a.right[k]}\0${b.right[k]}`
142
- let tmp = pairTemps.get(key)
143
- if (!tmp) {
144
- tmp = `$__dotpair_${freshIdRef.next++}`
145
- pairTemps.set(key, tmp)
146
- newLocalDecls.push(['local', tmp, 'v128'])
147
- fnLocals.set(tmp, 'v128')
148
- prefix.push(['local.set', tmp, f64x2Pair(a.right[k], b.right[k])])
149
- }
150
- pairs.push(['local.get', tmp])
151
- }
152
- const dot = dotPairExpr(a.left, pairs)
153
- const expr = addend ? ['f64x2.add', dot, ['f64x2.splat', addend]] : dot
154
- node.splice(i, b.end - i,
155
- ...prefix,
156
- ['local.set', v, expr],
157
- ['local.set', a.out, ['f64x2.extract_lane', 0, ['local.get', v]]],
158
- ['local.set', b.out, ['f64x2.extract_lane', 1, ['local.get', v]]],
159
- )
160
- i += prefix.length + 3
161
- }
162
- }
163
-
164
- // ---- Lane type tables ------------------------------------------------------
165
-
166
- const LANE_INFO = {
167
- i8: { lanes: 16, strideLog2: 0, stride: 1, splat: 'i8x16.splat', constOp: 'i32.const' },
168
- i16: { lanes: 8, strideLog2: 1, stride: 2, splat: 'i16x8.splat', constOp: 'i32.const' },
169
- i32: { lanes: 4, strideLog2: 2, stride: 4, splat: 'i32x4.splat', constOp: 'i32.const' },
170
- i64: { lanes: 2, strideLog2: 3, stride: 8, splat: 'i64x2.splat', constOp: 'i64.const' },
171
- f32: { lanes: 4, strideLog2: 2, stride: 4, splat: 'f32x4.splat', constOp: 'f32.const' },
172
- f64: { lanes: 2, strideLog2: 3, stride: 8, splat: 'f64x2.splat', constOp: 'f64.const' },
173
- }
174
-
175
- // Narrow loads/stores (i32.load8_u etc.) define i8 / i16 lane types — values
176
- // computed in i32 then truncated by store{8,16}, which matches i{8,16}xN wrap
177
- // semantics exactly.
178
- const LOAD_OPS = {
179
- 'i32.load8_u': 'i8', 'i32.load8_s': 'i8',
180
- 'i32.load16_u': 'i16','i32.load16_s': 'i16',
181
- 'i32.load': 'i32', 'i64.load': 'i64', 'f32.load': 'f32', 'f64.load': 'f64',
182
- }
183
- const STORE_OPS = {
184
- 'i32.store8': 'i8', 'i32.store16': 'i16',
185
- 'i32.store': 'i32', 'i64.store': 'i64', 'f32.store': 'f32', 'f64.store': 'f64',
186
- }
187
-
188
- // scalar op → SIMD op. shamtScalar:true means second operand stays scalar i32.
189
- //
190
- // For i8/i16 lanes the SCALAR ops are i32.* — wasm has no native i8/i16 ops,
191
- // values flow as i32 and the trailing store{8,16} truncates. i{8,16}x{N}.add
192
- // wraps within each lane the same way, so the observable result matches.
193
- // Note: wasm SIMD has no i8x16.mul, so multiplication on byte arrays bails.
194
- const LANE_PURE = {
195
- // Right shifts intentionally omitted for narrow lanes: scalar emits
196
- // i32.shr_{s,u} on a load8/load16 i32 (zero- or sign-extended), while
197
- // i{8,16}x{N}.shr_{s,u} treats lanes as their narrow type. The two diverge
198
- // when load and shift signedness mismatch (e.g. load8_u + shr_s on byte
199
- // 0xFF: scalar=0x7F, SIMD=0xFF). Safe set excludes shr_*.
200
- i8: new Map([
201
- ['i32.add', { simd: 'i8x16.add' }],
202
- ['i32.sub', { simd: 'i8x16.sub' }],
203
- ['i32.and', { simd: 'v128.and' }],
204
- ['i32.or', { simd: 'v128.or' }],
205
- ['i32.xor', { simd: 'v128.xor' }],
206
- ['i32.shl', { simd: 'i8x16.shl', shamtScalar: true }],
207
- ]),
208
- i16: new Map([
209
- ['i32.add', { simd: 'i16x8.add' }],
210
- ['i32.sub', { simd: 'i16x8.sub' }],
211
- ['i32.mul', { simd: 'i16x8.mul' }],
212
- ['i32.and', { simd: 'v128.and' }],
213
- ['i32.or', { simd: 'v128.or' }],
214
- ['i32.xor', { simd: 'v128.xor' }],
215
- ['i32.shl', { simd: 'i16x8.shl', shamtScalar: true }],
216
- ]),
217
- i32: new Map([
218
- ['i32.add', { simd: 'i32x4.add' }],
219
- ['i32.sub', { simd: 'i32x4.sub' }],
220
- ['i32.mul', { simd: 'i32x4.mul' }],
221
- ['i32.and', { simd: 'v128.and' }],
222
- ['i32.or', { simd: 'v128.or' }],
223
- ['i32.xor', { simd: 'v128.xor' }],
224
- ['i32.shl', { simd: 'i32x4.shl', shamtScalar: true }],
225
- ['i32.shr_s', { simd: 'i32x4.shr_s', shamtScalar: true }],
226
- ['i32.shr_u', { simd: 'i32x4.shr_u', shamtScalar: true }],
227
- ]),
228
- i64: new Map([
229
- ['i64.add', { simd: 'i64x2.add' }],
230
- ['i64.sub', { simd: 'i64x2.sub' }],
231
- ['i64.mul', { simd: 'i64x2.mul' }],
232
- ['i64.and', { simd: 'v128.and' }],
233
- ['i64.or', { simd: 'v128.or' }],
234
- ['i64.xor', { simd: 'v128.xor' }],
235
- ['i64.shl', { simd: 'i64x2.shl', shamtScalar: true }],
236
- ['i64.shr_s', { simd: 'i64x2.shr_s', shamtScalar: true }],
237
- ['i64.shr_u', { simd: 'i64x2.shr_u', shamtScalar: true }],
238
- ]),
239
- f32: new Map([
240
- ['f32.add', { simd: 'f32x4.add' }],
241
- ['f32.sub', { simd: 'f32x4.sub' }],
242
- ['f32.mul', { simd: 'f32x4.mul' }],
243
- ['f32.div', { simd: 'f32x4.div' }],
244
- ['f32.min', { simd: 'f32x4.min' }],
245
- ['f32.max', { simd: 'f32x4.max' }],
246
- ['f32.neg', { simd: 'f32x4.neg' }],
247
- ['f32.abs', { simd: 'f32x4.abs' }],
248
- ['f32.sqrt', { simd: 'f32x4.sqrt' }],
249
- ]),
250
- f64: new Map([
251
- ['f64.add', { simd: 'f64x2.add' }],
252
- ['f64.sub', { simd: 'f64x2.sub' }],
253
- ['f64.mul', { simd: 'f64x2.mul' }],
254
- ['f64.div', { simd: 'f64x2.div' }],
255
- ['f64.min', { simd: 'f64x2.min' }],
256
- ['f64.max', { simd: 'f64x2.max' }],
257
- ['f64.neg', { simd: 'f64x2.neg' }],
258
- ['f64.abs', { simd: 'f64x2.abs' }],
259
- ['f64.sqrt', { simd: 'f64x2.sqrt' }],
260
- ]),
261
- }
262
-
263
- // Horizontal reductions: associative+commutative ops applied to one
264
- // loop-carried accumulator. Each entry maps the SCALAR op (which is also
265
- // the op used to combine the SIMD result back into the accumulator at the
266
- // end) to its SIMD lane op, lane extractor, and identity element.
267
- //
268
- // Floats (add) are not strictly associative — vectorized order produces
269
- // ulp-level differences from scalar order. Acceptable for typical use
270
- // (reductions over typed arrays of well-conditioned data); strict-equal
271
- // callers must keep the pass off.
272
- //
273
- // Narrow lanes (i8/i16) intentionally absent: `s += a[i]` with a u8/u16
274
- // load expands the value to i32 before the add, so the accumulator's lane
275
- // type is always wider than the load's element type. That widening would
276
- // require pairwise/extending-add ops (i16x8.extadd_pairwise_*) — separate
277
- // recognizer.
278
- const REDUCE_OPS = {
279
- i32: {
280
- 'i32.add': { simd: 'i32x4.add', extract: 'i32x4.extract_lane', laneType: 'i32', constNode: ['i32.const', 0] },
281
- 'i32.xor': { simd: 'v128.xor', extract: 'i32x4.extract_lane', laneType: 'i32', constNode: ['i32.const', 0] },
282
- 'i32.and': { simd: 'v128.and', extract: 'i32x4.extract_lane', laneType: 'i32', constNode: ['i32.const', -1] },
283
- 'i32.or': { simd: 'v128.or', extract: 'i32x4.extract_lane', laneType: 'i32', constNode: ['i32.const', 0] },
284
- },
285
- i64: {
286
- 'i64.add': { simd: 'i64x2.add', extract: 'i64x2.extract_lane', laneType: 'i64', constNode: ['i64.const', 0] },
287
- 'i64.xor': { simd: 'v128.xor', extract: 'i64x2.extract_lane', laneType: 'i64', constNode: ['i64.const', 0] },
288
- 'i64.and': { simd: 'v128.and', extract: 'i64x2.extract_lane', laneType: 'i64', constNode: ['i64.const', -1] },
289
- 'i64.or': { simd: 'v128.or', extract: 'i64x2.extract_lane', laneType: 'i64', constNode: ['i64.const', 0] },
290
- },
291
- f32: {
292
- 'f32.add': { simd: 'f32x4.add', extract: 'f32x4.extract_lane', laneType: 'f32', constNode: ['f32.const', 0] },
293
- },
294
- f64: {
295
- 'f64.add': { simd: 'f64x2.add', extract: 'f64x2.extract_lane', laneType: 'f64', constNode: ['f64.const', 0] },
296
- },
297
- }
298
-
299
- // op-name → REDUCE entry across all lane types (the op-name itself encodes
300
- // the lane type prefix, e.g. `i32.add` ⇒ i32 lanes).
301
- const REDUCE_OP_LOOKUP = (() => {
302
- const m = new Map()
303
- for (const lt of Object.keys(REDUCE_OPS))
304
- for (const op of Object.keys(REDUCE_OPS[lt]))
305
- m.set(op, REDUCE_OPS[lt][op])
306
- return m
307
- })()
308
-
309
- // ---- Recognizer ------------------------------------------------------------
310
-
311
- function isLocalGet(node, name) {
312
- return isArr(node) && node[0] === 'local.get' && (name == null || node[1] === name)
313
- }
314
- function isI32Const(node) {
315
- return isArr(node) && node[0] === 'i32.const'
316
- }
317
- function constNum(node) {
318
- if (!isI32Const(node)) return null
319
- const v = node[1]
320
- return typeof v === 'number' ? v : (typeof v === 'string' ? parseInt(v, 10) : null)
321
- }
322
-
323
- /**
324
- * Match increment shape `(local.set $X (i32.add (local.get $X) (i32.const 1)))`.
325
- * Returns $X or null.
326
- */
327
- function matchInc1(stmt) {
328
- if (!isArr(stmt) || stmt[0] !== 'local.set' || stmt.length !== 3) return null
329
- const x = stmt[1]
330
- const v = stmt[2]
331
- if (!isArr(v) || v[0] !== 'i32.add' || v.length !== 3) return null
332
- if (!isLocalGet(v[1], x)) return null
333
- if (constNum(v[2]) !== 1) return null
334
- return x
335
- }
336
-
337
- /**
338
- * Match `(br_if $LABEL (i32.eqz (i32.lt_{s,u} (local.get $I) BOUND)))`.
339
- * Returns { ind, bound } or null.
340
- */
341
- function matchExitBrIf(stmt, label) {
342
- if (!isArr(stmt) || stmt[0] !== 'br_if' || stmt[1] !== label) return null
343
- const cond = stmt[2]
344
- if (!isArr(cond) || cond[0] !== 'i32.eqz') return null
345
- const cmp = cond[1]
346
- if (!isArr(cmp) || (cmp[0] !== 'i32.lt_s' && cmp[0] !== 'i32.lt_u')) return null
347
- if (!isLocalGet(cmp[1])) return null
348
- return { ind: cmp[1][1], bound: cmp[2] }
349
- }
350
-
351
- /**
352
- * Walk node, collect set of local names that are written via local.set/local.tee
353
- * anywhere within. Used to detect loop-invariant locals.
354
- */
355
- function collectWrites(node, out) {
356
- if (!isArr(node)) return
357
- const op = node[0]
358
- if ((op === 'local.set' || op === 'local.tee') && typeof node[1] === 'string') {
359
- out.add(node[1])
360
- }
361
- for (let i = 0; i < node.length; i++) collectWrites(node[i], out)
362
- }
363
-
364
- /**
365
- * Return the FIRST kind of access for `name` in straight-line walk order.
366
- * 'write' — local.set/local.tee seen first
367
- * 'read' — local.get seen first
368
- * null — not referenced
369
- */
370
- function firstAccess(node, name) {
371
- if (!isArr(node)) return null
372
- const op = node[0]
373
- // Walk children first — operands evaluate before the op. For local.set/tee
374
- // the VALUE child (idx 2) runs before the write, so a `local.get name` in
375
- // the value of `local.set name` is a read-before-write.
376
- if ((op === 'local.set' || op === 'local.tee') && node[1] === name) {
377
- if (node.length >= 3) {
378
- const r = firstAccess(node[2], name)
379
- if (r) return r
380
- }
381
- return 'write'
382
- }
383
- if (op === 'local.get' && node[1] === name) return 'read'
384
- for (let i = 1; i < node.length; i++) {
385
- const r = firstAccess(node[i], name)
386
- if (r) return r
387
- }
388
- return null
389
- }
390
-
391
- /**
392
- * Match the index-offset operand of a lane address — the part that scales the
393
- * induction variable inside `(i32.add base OFFSET)`. OFFSET is one of:
394
- * (i32.shl (local.get IND) (i32.const K)) → strideLog2 K
395
- * (local.get IND) → strideLog2 0
396
- * (local.tee $T <either of the above>) → records $T
397
- * (local.get $T) where $T is a recorded offset-tee → that tee's strideLog2
398
- *
399
- * The tee'd form arises from CSE: a map loop `b[i] = f(a[i])` over two distinct
400
- * base pointers shares one `i << K` offset (`(local.tee $T (i32.shl i K))` in
401
- * the first address, `(local.get $T)` in the second). `offsetTees` (Map
402
- * name→strideLog2) carries that across calls.
403
- *
404
- * Returns { strideLog2, teeName?: string } or null.
405
- */
406
- function matchLaneOffset(off, ind, offsetTees) {
407
- if (isArr(off) && off[0] === 'local.get' && typeof off[1] === 'string' &&
408
- offsetTees && offsetTees.has(off[1])) {
409
- return { strideLog2: offsetTees.get(off[1]), teeName: null }
410
- }
411
- let teeName = null
412
- let n = off
413
- if (isArr(n) && n[0] === 'local.tee' && n.length === 3) { teeName = n[1]; n = n[2] }
414
- // (i32.shl (local.get ind) (i32.const K))
415
- if (isArr(n) && n[0] === 'i32.shl' && n.length === 3 && isLocalGet(n[1], ind)) {
416
- const k = constNum(n[2])
417
- if (k != null && k >= 0 && k <= 3) return { strideLog2: k, teeName }
418
- }
419
- // (local.get ind) — stride 1
420
- if (isLocalGet(n, ind)) return { strideLog2: 0, teeName }
421
- return null
422
- }
423
-
424
- /**
425
- * Match an address expression `(i32.add base OFFSET)`, with optional outer
426
- * `(local.tee $A ...)`. OFFSET is matched by matchLaneOffset (which also
427
- * accepts a CSE'd `(local.tee $T (i32.shl ind K))` / `(local.get $T)` pair).
428
- * Also accepts `(local.get $A)` when $A is a previously-recorded address tee.
429
- *
430
- * Returns { strideLog2, base, teeName?, offsetTeeName?, viaLocal? } or null.
431
- * `strideLog2` = K for i32.shl form, 0 for plain add form.
432
- * `base` is the loop-invariant base subtree.
433
- */
434
- function matchLaneAddr(addr, ind, addrLocals, offsetTees) {
435
- let teeName = null
436
- let n = addr
437
- // (local.get $A) where $A holds a previously-tee'd FULL lane-address.
438
- if (isArr(n) && n[0] === 'local.get' && typeof n[1] === 'string' && addrLocals && addrLocals.has(n[1])) {
439
- const e = addrLocals.get(n[1])
440
- return { strideLog2: e.strideLog2, base: e.base, teeName: null, viaLocal: n[1] }
441
- }
442
- if (isArr(n) && n[0] === 'local.tee' && n.length === 3) {
443
- teeName = n[1]
444
- n = n[2]
445
- }
446
- if (!isArr(n) || n[0] !== 'i32.add' || n.length !== 3) return null
447
- const a = n[1], b = n[2]
448
- const off = matchLaneOffset(b, ind, offsetTees)
449
- if (!off) return null
450
- return { strideLog2: off.strideLog2, base: a, teeName, offsetTeeName: off.teeName }
451
- }
452
-
453
- /**
454
- * A scalar i32 local that is ONLY ever assigned a lane offset — `(i32.shl ind K)`
455
- * (or bare `ind` for stride 0) — is a CSE'd offset shared across base pointers.
456
- * Returns the consistent strideLog2, or null if any write to it diverges.
457
- * This soundness check backs every `(local.get $T)` resolved via `offsetTees`.
458
- */
459
- function _offsetLocalStride(body, name, ind) {
460
- let stride = null, found = false, ok = true
461
- function walk(n) {
462
- if (!isArr(n)) return
463
- if ((n[0] === 'local.tee' || n[0] === 'local.set') && n[1] === name && n.length === 3) {
464
- found = true
465
- const v = n[2]
466
- let k = null
467
- if (isArr(v) && v[0] === 'i32.shl' && v.length === 3 && isLocalGet(v[1], ind)) {
468
- k = constNum(v[2])
469
- if (k == null || k < 0 || k > 3) ok = false
470
- } else if (isLocalGet(v, ind)) {
471
- k = 0
472
- } else ok = false
473
- if (k != null) {
474
- if (stride == null) stride = k
475
- else if (stride !== k) ok = false
476
- }
477
- }
478
- for (let i = 1; i < n.length; i++) walk(n[i])
479
- }
480
- for (const s of body) walk(s)
481
- return found && ok ? stride : null
482
- }
483
-
484
- // ---- Recognize a (block (loop)) pair --------------------------------------
485
-
486
- /**
487
- * Try to vectorize the inner loop. Returns the replacement node array
488
- * (synthetic outer block) or null on no match.
489
- */
490
- function tryVectorize(blockNode, fnLocals, freshIdRef) {
491
- if (!isArr(blockNode) || blockNode[0] !== 'block') return null
492
- // Find label and inner loop.
493
- let blockLabel = null
494
- let loopIdx = -1, loopNode = null
495
- for (let i = 1; i < blockNode.length; i++) {
496
- const c = blockNode[i]
497
- if (typeof c === 'string' && c.startsWith('$') && blockLabel == null && i === 1) {
498
- blockLabel = c; continue
499
- }
500
- if (isArr(c) && c[0] === 'loop') {
501
- if (loopNode) return null // multiple loops
502
- loopIdx = i; loopNode = c
503
- } else if (isArr(c)) {
504
- return null // foreign content alongside the loop
505
- }
506
- }
507
- if (!loopNode || !blockLabel) return null
508
-
509
- // Loop layout: ['loop', '$label', ...stmts]
510
- const loopLabel = typeof loopNode[1] === 'string' && loopNode[1].startsWith('$') ? loopNode[1] : null
511
- if (!loopLabel) return null
512
-
513
- // Find induction increment + back-branch at the END.
514
- let endIdx = loopNode.length - 1
515
- if (!(isArr(loopNode[endIdx]) && loopNode[endIdx][0] === 'br' && loopNode[endIdx][1] === loopLabel)) return null
516
- const incIdx = endIdx - 1
517
- const incVar = matchInc1(loopNode[incIdx])
518
- if (!incVar) return null
519
-
520
- // First stmt must be the exit br_if.
521
- const exitInfo = matchExitBrIf(loopNode[2], blockLabel)
522
- if (!exitInfo) return null
523
- if (exitInfo.ind !== incVar) return null
524
-
525
- // Body = stmts between exit and increment.
526
- const body = []
527
- for (let i = 3; i < incIdx; i++) body.push(loopNode[i])
528
-
529
- // Bound must be loop-invariant. For now, accept (local.get $L) where $L
530
- // is declared but not written inside the body, OR (i32.const N).
531
- let bound = exitInfo.bound
532
- let boundLocal = null
533
- if (isArr(bound) && bound[0] === 'local.get' && typeof bound[1] === 'string') {
534
- boundLocal = bound[1]
535
- } else if (isI32Const(bound)) {
536
- // ok
537
- } else {
538
- return null
539
- }
540
-
541
- // Detect lane type from the FIRST load in body.
542
- let laneType = null
543
- let stride = -1
544
- const loadStoreSites = [] // {parent, idx, kind:'load'|'store'}
545
- // Address tees: name → {strideLog2, base}. A `(local.tee NAME (lane-addr))`
546
- // both validates the load's address AND records NAME so the matching store's
547
- // `(local.get NAME)` is accepted as the same lane address.
548
- const addrLocals = new Map()
549
- // Offset tees: name → strideLog2. A CSE'd `i << K` shared across base
550
- // pointers (map loops over distinct arrays). Soundness re-checked post-scan.
551
- const offsetTees = new Map()
552
-
553
- function scanForLoadsStores(node, parent, pi) {
554
- if (!isArr(node)) return true
555
- const op = node[0]
556
- if (LOAD_OPS[op]) {
557
- if (laneType == null) {
558
- laneType = LOAD_OPS[op]
559
- stride = LANE_INFO[laneType].stride
560
- } else if (LOAD_OPS[op] !== laneType) {
561
- return false
562
- }
563
- const m = matchLaneAddr(node[1], incVar, addrLocals, offsetTees)
564
- if (!m) return false
565
- if ((1 << m.strideLog2) !== stride) return false
566
- if (m.teeName) addrLocals.set(m.teeName, { strideLog2: m.strideLog2, base: m.base })
567
- if (m.offsetTeeName) offsetTees.set(m.offsetTeeName, m.strideLog2)
568
- loadStoreSites.push({ parent, idx: pi, kind: 'load' })
569
- return true
570
- }
571
- if (STORE_OPS[op]) {
572
- const sty = STORE_OPS[op]
573
- if (laneType != null && sty !== laneType) return false
574
- if (laneType == null) { laneType = sty; stride = LANE_INFO[laneType].stride }
575
- const m = matchLaneAddr(node[1], incVar, addrLocals, offsetTees)
576
- if (!m) return false
577
- if ((1 << m.strideLog2) !== stride) return false
578
- if (m.teeName) addrLocals.set(m.teeName, { strideLog2: m.strideLog2, base: m.base })
579
- if (m.offsetTeeName) offsetTees.set(m.offsetTeeName, m.strideLog2)
580
- loadStoreSites.push({ parent, idx: pi, kind: 'store' })
581
- // Recurse into VALUE child (idx 2) — it's data, not address.
582
- if (!scanForLoadsStores(node[2], node, 2)) return false
583
- return true
584
- }
585
- // local.set/tee of an address local outside a load/store context (e.g.
586
- // `(local.set $a (i32.add base (i32.shl i 2)))` as a standalone stmt) —
587
- // record so a later `(local.get $a)` resolves.
588
- if ((op === 'local.set' || op === 'local.tee') && typeof node[1] === 'string' && node.length === 3) {
589
- const valM = matchLaneAddr(['local.tee', node[1], node[2]], incVar, addrLocals, offsetTees)
590
- if (valM && valM.teeName) {
591
- addrLocals.set(valM.teeName, { strideLog2: valM.strideLog2, base: valM.base })
592
- }
593
- // Standalone offset compute: `(local.set $t (i32.shl i K))`.
594
- const offM = matchLaneOffset(node[2], incVar, offsetTees)
595
- if (offM) offsetTees.set(node[1], offM.strideLog2)
596
- }
597
- // Recurse into all children
598
- for (let i = 1; i < node.length; i++) {
599
- if (!scanForLoadsStores(node[i], node, i)) return false
600
- }
601
- return true
602
- }
603
- for (const stmt of body) {
604
- if (!scanForLoadsStores(stmt, null, -1)) return null
605
- }
606
- if (!laneType) return null // no memory ops — vectorizing buys nothing
607
- if (loadStoreSites.length === 0) return null
608
-
609
- // Soundness gate for offset-tee resolution: every `(local.get $T)` we
610
- // accepted as `i << K` is only valid if EVERY write of $T is that offset.
611
- for (const [name, k] of offsetTees) {
612
- if (_offsetLocalStride(body, name, incVar) !== k) return null
613
- }
614
-
615
- // Classify all locals referenced in body.
616
- // - induction var (incVar): exempt
617
- // - bound local (if any): must be invariant
618
- // - each other local: first access must not be a read-then-written pattern
619
- const writes = new Set()
620
- for (const s of body) collectWrites(s, writes)
621
- if (boundLocal && writes.has(boundLocal)) return null // bound varies in body → bail
622
-
623
- const localKind = new Map() // name → 'lane' | 'invariant' | 'addr'
624
- // Walk to collect ALL referenced names
625
- const referenced = new Set()
626
- const collectRefs = (n) => {
627
- if (!isArr(n)) return
628
- const op = n[0]
629
- if ((op === 'local.get' || op === 'local.set' || op === 'local.tee') && typeof n[1] === 'string')
630
- referenced.add(n[1])
631
- for (let i = 1; i < n.length; i++) collectRefs(n[i])
632
- }
633
- for (const s of body) collectRefs(s)
634
-
635
- for (const name of referenced) {
636
- if (name === incVar) continue
637
- if (writes.has(name)) {
638
- // Must be lane-local: first access is a write.
639
- let firstKind = null
640
- for (const s of body) {
641
- const k = firstAccess(s, name)
642
- if (k) { firstKind = k; break }
643
- }
644
- if (firstKind === 'read') return null // loop-carried (reduction or stencil)
645
- // Discriminate lane-data vs address-tee. Address tees hold i32 addresses,
646
- // not vector data. We classify by checking the local's declared type.
647
- const decl = fnLocals.get(name)
648
- if (decl === 'i32' && (offsetTees.has(name) || _isAddressLocal(body, name, incVar))) {
649
- localKind.set(name, 'addr')
650
- } else {
651
- localKind.set(name, 'lane')
652
- }
653
- } else {
654
- localKind.set(name, 'invariant')
655
- }
656
- }
657
-
658
- // Build lifted body. If anything fails to lift, bail.
659
- const newLanedLocals = new Map() // origName → { laneName, simdType }
660
- const ctx = { laneType, incVar, localKind, newLanedLocals, fail: false, failReason: null }
661
- const lifted = []
662
- for (const s of body) {
663
- const r = liftStmt(s, ctx)
664
- if (ctx.fail) return null
665
- if (r != null) {
666
- if (Array.isArray(r) && r[0] === '__seq__') lifted.push(...r.slice(1))
667
- else lifted.push(r)
668
- }
669
- }
670
- if (lifted.length === 0) return null
671
-
672
- // Generate fresh names
673
- const id = freshIdRef.next++
674
- const simdBoundName = `$__simd_bound${id}`
675
- const simdBrkLabel = `$__simd_brk${id}`
676
- const simdLoopLabel = `$__simd_loop${id}`
677
-
678
- const info = LANE_INFO[laneType]
679
- const lanes = info.lanes
680
- const mask = -lanes // bit pattern ~(lanes-1) in i32 two's complement
681
-
682
- // Build SIMD prefix block.
683
- const boundExpr = boundLocal
684
- ? ['local.get', boundLocal]
685
- : bound // i32.const N
686
- const simdBlock = ['block', simdBrkLabel,
687
- ['loop', simdLoopLabel,
688
- ['br_if', simdBrkLabel,
689
- ['i32.eqz', ['i32.lt_s', ['local.get', incVar], ['local.get', simdBoundName]]]],
690
- ...lifted,
691
- ['local.set', incVar,
692
- ['i32.add', ['local.get', incVar], ['i32.const', lanes]]],
693
- ['br', simdLoopLabel]
694
- ]
695
- ]
696
-
697
- // Bound setup: simdBoundName = bound & ~(lanes-1)
698
- const boundSetup = ['local.set', simdBoundName,
699
- ['i32.and', boundExpr, ['i32.const', mask]]]
700
-
701
- // Synthetic outer wrapper — has no result, no label, just sequences.
702
- // The original block is preserved unchanged as the tail.
703
- const wrapper = ['block', boundSetup, simdBlock, blockNode]
704
-
705
- // Locals to add to function header.
706
- const newLocalDecls = [
707
- ['local', simdBoundName, 'i32'],
708
- ...[...newLanedLocals.values()].map(({ laneName }) => ['local', laneName, 'v128'])
709
- ]
710
-
711
- return { wrapper, newLocalDecls }
712
- }
713
-
714
- // ---- Reduction recognizer -------------------------------------------------
715
- //
716
- // Matches inner loops of shape:
717
- // for (let i = 0; i < N; i++) S = OP(S, EXPR(arr[i], ...))
718
- // where OP is associative+commutative (REDUCE_OPS table) and EXPR is lane-
719
- // pure (operates on the loaded element with at most loop-invariant data).
720
- // S is a SCALAR loop-carried accumulator — exempt from the lane-local
721
- // "first access must be a write" check.
722
- //
723
- // Lift:
724
- // acc = splat(IDENTITY)
725
- // for (i = 0; i < bound & ~(L-1); i += L) acc = OP_v(acc, lifted EXPR)
726
- // S = OP(S, horizontal_reduce(acc))
727
- // <original scalar tail handles the remainder>
728
- //
729
- // Float adds are not strictly associative — vectorized reduction differs
730
- // from scalar reduction by ulps. Acceptable when bit-exact equality is not
731
- // required (which it isn't, by spec, in JS engines either).
732
- function tryReduceVectorize(blockNode, fnLocals, freshIdRef) {
733
- if (!isArr(blockNode) || blockNode[0] !== 'block') return null
734
-
735
- // Match outer (block (loop)) structure. Same loop-shape as tryVectorize.
736
- let blockLabel = null
737
- let loopNode = null
738
- for (let i = 1; i < blockNode.length; i++) {
739
- const c = blockNode[i]
740
- if (typeof c === 'string' && c.startsWith('$') && blockLabel == null && i === 1) { blockLabel = c; continue }
741
- if (isArr(c) && c[0] === 'loop') {
742
- if (loopNode) return null
743
- loopNode = c
744
- } else if (isArr(c)) return null
745
- }
746
- if (!loopNode || !blockLabel) return null
747
- const loopLabel = typeof loopNode[1] === 'string' && loopNode[1].startsWith('$') ? loopNode[1] : null
748
- if (!loopLabel) return null
749
- const endIdx = loopNode.length - 1
750
- if (!(isArr(loopNode[endIdx]) && loopNode[endIdx][0] === 'br' && loopNode[endIdx][1] === loopLabel)) return null
751
- const incIdx = endIdx - 1
752
- const incVar = matchInc1(loopNode[incIdx])
753
- if (!incVar) return null
754
- const exitInfo = matchExitBrIf(loopNode[2], blockLabel)
755
- if (!exitInfo) return null
756
- if (exitInfo.ind !== incVar) return null
757
-
758
- // Body must be a single statement: the accumulator update.
759
- if (incIdx - 3 !== 1) return null
760
- const stmt = loopNode[3]
761
- if (!isArr(stmt) || stmt[0] !== 'local.set' || stmt.length !== 3) return null
762
- const accName = stmt[1]
763
- if (typeof accName !== 'string') return null
764
- const rhs = stmt[2]
765
- if (!isArr(rhs) || rhs.length !== 3) return null
766
- const opName = rhs[0]
767
- const reduceEntry = REDUCE_OP_LOOKUP.get(opName)
768
- if (!reduceEntry) return null
769
- if (!isLocalGet(rhs[1], accName)) return null
770
- const exprNode = rhs[2]
771
-
772
- // Accumulator's declared local type must match the lane element type.
773
- const accType = fnLocals.get(accName)
774
- if (accType !== reduceEntry.laneType) return null
775
-
776
- // Bound classification (same as tryVectorize).
777
- let bound = exitInfo.bound
778
- let boundLocal = null
779
- if (isArr(bound) && bound[0] === 'local.get' && typeof bound[1] === 'string') boundLocal = bound[1]
780
- else if (!isI32Const(bound)) return null
781
-
782
- // Scan EXPR for lane-aligned loads. Stores forbidden. Re-references of
783
- // accName forbidden (the accumulator only appears in the outer wrapper).
784
- const laneType = reduceEntry.laneType
785
- const stride = LANE_INFO[laneType].stride
786
- const addrLocals = new Map()
787
- const offsetTees = new Map()
788
- let loadCount = 0
789
- function scanExpr(node) {
790
- if (!isArr(node)) return true
791
- const op = node[0]
792
- if (LOAD_OPS[op]) {
793
- if (LOAD_OPS[op] !== laneType) return false
794
- const m = matchLaneAddr(node[1], incVar, addrLocals, offsetTees)
795
- if (!m) return false
796
- if ((1 << m.strideLog2) !== stride) return false
797
- if (m.teeName) addrLocals.set(m.teeName, { strideLog2: m.strideLog2, base: m.base })
798
- if (m.offsetTeeName) offsetTees.set(m.offsetTeeName, m.strideLog2)
799
- loadCount++
800
- return true
801
- }
802
- if (STORE_OPS[op]) return false
803
- if (op === 'local.set' || op === 'local.tee') return false // no intermediates
804
- if (op === 'local.get' && node[1] === accName) return false
805
- for (let i = 1; i < node.length; i++) if (!scanExpr(node[i])) return false
806
- return true
807
- }
808
- if (!scanExpr(exprNode)) return null
809
- if (loadCount === 0) return null
810
- // Soundness gate for offset-tee resolution (see tryVectorize).
811
- for (const [name, k] of offsetTees) {
812
- if (_offsetLocalStride([exprNode], name, incVar) !== k) return null
813
- }
814
-
815
- // Classify locals referenced in EXPR. Anything not the induction var or an
816
- // address-tee is invariant (we forbade local.set/tee in scanExpr).
817
- const referenced = new Set()
818
- const collectRefs = (n) => {
819
- if (!isArr(n)) return
820
- if (n[0] === 'local.get' && typeof n[1] === 'string') referenced.add(n[1])
821
- for (let i = 1; i < n.length; i++) collectRefs(n[i])
822
- }
823
- collectRefs(exprNode)
824
- const localKind = new Map()
825
- for (const name of referenced) {
826
- if (name === incVar) continue
827
- if (addrLocals.has(name) || offsetTees.has(name)) { localKind.set(name, 'addr'); continue }
828
- localKind.set(name, 'invariant')
829
- }
830
- for (const name of addrLocals.keys()) localKind.set(name, 'addr')
831
- for (const name of offsetTees.keys()) localKind.set(name, 'addr')
832
-
833
- const ctx = { laneType, incVar, localKind, newLanedLocals: new Map(), fail: false, failReason: null }
834
- const liftedExpr = liftExprV(exprNode, ctx)
835
- if (ctx.fail) return null
836
- if (ctx.newLanedLocals.size > 0) return null
837
-
838
- // Synthesize SIMD prefix block + horizontal reduce + (preserved scalar tail).
839
- const id = freshIdRef.next++
840
- const simdBoundName = `$__simd_bound${id}`
841
- const simdAccName = `$__simd_acc${id}`
842
- const simdBrkLabel = `$__simd_brk${id}`
843
- const simdLoopLabel = `$__simd_loop${id}`
844
- const info = LANE_INFO[laneType]
845
- const lanes = info.lanes
846
- const mask = -lanes
847
- const boundExpr = boundLocal ? ['local.get', boundLocal] : bound
848
-
849
- const initAcc = ['local.set', simdAccName, [info.splat, reduceEntry.constNode]]
850
- const simdBlock = ['block', simdBrkLabel,
851
- ['loop', simdLoopLabel,
852
- ['br_if', simdBrkLabel,
853
- ['i32.eqz', ['i32.lt_s', ['local.get', incVar], ['local.get', simdBoundName]]]],
854
- ['local.set', simdAccName,
855
- [reduceEntry.simd, ['local.get', simdAccName], liftedExpr]],
856
- ['local.set', incVar, ['i32.add', ['local.get', incVar], ['i32.const', lanes]]],
857
- ['br', simdLoopLabel]
858
- ]
859
- ]
860
-
861
- // Horizontal fold: scalar.op(extract 0, extract 1, …, extract L-1).
862
- let horiz = [reduceEntry.extract, 0, ['local.get', simdAccName]]
863
- for (let k = 1; k < lanes; k++) {
864
- horiz = [opName, horiz, [reduceEntry.extract, k, ['local.get', simdAccName]]]
865
- }
866
- const mergeBack = ['local.set', accName, [opName, ['local.get', accName], horiz]]
867
- const boundSetup = ['local.set', simdBoundName, ['i32.and', boundExpr, ['i32.const', mask]]]
868
-
869
- const wrapper = ['block', boundSetup, initAcc, simdBlock, mergeBack, blockNode]
870
- const newLocalDecls = [
871
- ['local', simdBoundName, 'i32'],
872
- ['local', simdAccName, 'v128'],
873
- ]
874
- return { wrapper, newLocalDecls }
875
- }
876
-
877
- // Scalar locals that are ALWAYS computed as `(i32.add base (i32.shl ind K))`
878
- // or aliased to such an address are "address tees", not lane data. They stay
879
- // scalar i32 in the lifted body.
880
- function _isAddressLocal(body, name, ind) {
881
- let onlyAsAddrTee = true
882
- let foundTee = false
883
- function walk(n) {
884
- if (!isArr(n)) return
885
- if (n[0] === 'local.tee' && n[1] === name) {
886
- foundTee = true
887
- // Check the value is a lane-address shape
888
- const m = matchLaneAddr(['local.tee', name, n[2]], ind)
889
- if (!m) onlyAsAddrTee = false
890
- return
891
- }
892
- if (n[0] === 'local.set' && n[1] === name) {
893
- // A set-not-tee: check value shape
894
- const m = matchLaneAddr(['local.tee', name, n[2]], ind)
895
- if (!m) onlyAsAddrTee = false
896
- foundTee = true
897
- return
898
- }
899
- for (let i = 1; i < n.length; i++) walk(n[i])
900
- }
901
- for (const s of body) walk(s)
902
- return foundTee && onlyAsAddrTee
903
- }
904
-
905
- // ---- Lifter ----------------------------------------------------------------
906
-
907
- function getOrAllocLanedLocal(name, ctx) {
908
- let r = ctx.newLanedLocals.get(name)
909
- if (!r) {
910
- r = { laneName: `${name}__v`, origName: name }
911
- ctx.newLanedLocals.set(name, r)
912
- }
913
- return r
914
- }
915
-
916
- /** Lift a statement. Returns lifted stmt, or null to skip, or ['__seq__', ...] for multiple. */
917
- function liftStmt(stmt, ctx) {
918
- if (!isArr(stmt)) {
919
- // Bare strings like "drop" — produced by stack-form WAT. We unwrap value-blocks
920
- // separately so an isolated "drop" should not appear here, but tolerate it.
921
- if (stmt === 'drop') return null
922
- ctx.fail = true; return null
923
- }
924
- const op = stmt[0]
925
-
926
- if (op === 'local.set' && typeof stmt[1] === 'string' && stmt.length === 3) {
927
- const name = stmt[1]
928
- const kind = ctx.localKind.get(name)
929
- if (kind === 'addr') {
930
- // Address-only local: lift the value as-is (it's i32 arithmetic on ind).
931
- return ['local.set', name, stmt[2]]
932
- }
933
- if (kind === 'lane') {
934
- const { laneName } = getOrAllocLanedLocal(name, ctx)
935
- const v = liftExprV(stmt[2], ctx)
936
- if (ctx.fail) return null
937
- return ['local.set', laneName, v]
938
- }
939
- ctx.fail = true; return null
940
- }
941
-
942
- if (STORE_OPS[op]) {
943
- const simdStore = 'v128.store'
944
- const addr = stmt[1] // we leave addresses as-is (scalar i32 expressions)
945
- const val = liftExprV(stmt[2], ctx)
946
- if (ctx.fail) return null
947
- // Handle memarg if present (last positional after addr/val): unlikely in
948
- // pre-watr IR for this shape; bail if more than 3 children.
949
- if (stmt.length !== 3) { ctx.fail = true; return null }
950
- return [simdStore, addr, val]
951
- }
952
-
953
- // (block (result T) STMTS... TAIL_EXPR) followed by sibling "drop" — we get
954
- // the block alone here; the "drop" is a separate sibling and is returned as
955
- // null by the next call. Strip the wrapper, lift the inner stmts; the
956
- // dropped-tail expr is discarded.
957
- if (op === 'block') {
958
- // Block may be: ['block', LABEL?, RESULT?, ...stmts]
959
- let i = 1
960
- if (typeof stmt[i] === 'string' && stmt[i].startsWith('$')) i++
961
- const hasResult = isArr(stmt[i]) && stmt[i][0] === 'result'
962
- if (hasResult) i++
963
- const inner = stmt.slice(i)
964
- const stmts = hasResult ? inner.slice(0, inner.length - 1) : inner
965
- const out = ['__seq__']
966
- for (const s of stmts) {
967
- const lifted = liftStmt(s, ctx)
968
- if (ctx.fail) return null
969
- if (lifted == null) continue
970
- if (Array.isArray(lifted) && lifted[0] === '__seq__') out.push(...lifted.slice(1))
971
- else out.push(lifted)
972
- }
973
- return out
974
- }
975
-
976
- // Standalone expression-as-statement (e.g. a load that gets dropped) — bail.
977
- ctx.fail = true; return null
978
- }
979
-
980
- /** Lift a value expression into v128 context. */
981
- function liftExprV(expr, ctx) {
982
- if (!isArr(expr)) { ctx.fail = true; return null }
983
- const op = expr[0]
984
- const info = LANE_INFO[ctx.laneType]
985
-
986
- // Loads → v128.load (preserving address, including any local.tee).
987
- if (LOAD_OPS[op]) {
988
- if (LOAD_OPS[op] !== ctx.laneType) { ctx.fail = true; return null }
989
- return ['v128.load', expr[1]]
990
- }
991
-
992
- // Constants → splat.
993
- if (op === info.constOp) {
994
- return [info.splat, expr]
995
- }
996
-
997
- // local.get
998
- if (op === 'local.get' && typeof expr[1] === 'string') {
999
- const name = expr[1]
1000
- const kind = ctx.localKind.get(name)
1001
- if (kind === 'lane') {
1002
- const { laneName } = getOrAllocLanedLocal(name, ctx)
1003
- return ['local.get', laneName]
1004
- }
1005
- if (kind === 'invariant') {
1006
- return [info.splat, ['local.get', name]]
1007
- }
1008
- if (kind === 'addr' || name === ctx.incVar) {
1009
- ctx.fail = true; return null // can't be in a value position
1010
- }
1011
- ctx.fail = true; return null
1012
- }
1013
-
1014
- // Lane-pure op?
1015
- const table = LANE_PURE[ctx.laneType]
1016
- const entry = table?.get(op)
1017
- if (entry) {
1018
- const a = liftExprV(expr[1], ctx)
1019
- if (ctx.fail) return null
1020
- if (entry.shamtScalar) {
1021
- // Second operand stays scalar i32 — must be const or invariant local.
1022
- const b = expr[2]
1023
- if (!isI32Const(b) && !(isArr(b) && b[0] === 'local.get' && ctx.localKind.get(b[1]) === 'invariant')) {
1024
- ctx.fail = true; return null
1025
- }
1026
- return [entry.simd, a, b]
1027
- }
1028
- if (expr.length === 2) { // unary (neg, abs, sqrt)
1029
- return [entry.simd, a]
1030
- }
1031
- const b = liftExprV(expr[2], ctx)
1032
- if (ctx.fail) return null
1033
- return [entry.simd, a, b]
1034
- }
1035
-
1036
- ctx.fail = true; return null
1037
- }
1038
-
1039
- // ---- Pass entry ------------------------------------------------------------
1040
-
1041
- /**
1042
- * Walk a function looking for vectorizable (block (loop)) pairs, in-place.
1043
- * Adds new locals to the function header.
1044
- */
1045
- export function vectorizeLaneLocal(fn) {
1046
- if (!isArr(fn) || fn[0] !== 'func') return
1047
- const bodyStart = findBodyStart(fn)
1048
- if (bodyStart < 0) return
1049
-
1050
- // Build local-name → wasm-type map.
1051
- const fnLocals = new Map()
1052
- for (let i = 2; i < bodyStart; i++) {
1053
- const d = fn[i]
1054
- if (isArr(d) && d[0] === 'local' && typeof d[1] === 'string' && typeof d[2] === 'string') {
1055
- fnLocals.set(d[1], d[2])
1056
- } else if (isArr(d) && d[0] === 'param' && typeof d[1] === 'string' && typeof d[2] === 'string') {
1057
- fnLocals.set(d[1], d[2])
1058
- }
1059
- }
1060
-
1061
- const freshIdRef = { next: 0 }
1062
- const newLocalDeclsAll = []
1063
-
1064
- vectorizeStraightLineF64DotPairsIn(fn, fnLocals, freshIdRef, newLocalDeclsAll)
1065
-
1066
- // Walk body recursively. Process inner-most matches first (post-order)
1067
- // so we don't try to vectorize an outer loop whose inner is the lane-local one.
1068
- function walk(parent, idx) {
1069
- const node = parent[idx]
1070
- if (!isArr(node)) return
1071
- for (let i = 0; i < node.length; i++) {
1072
- if (isArr(node[i])) walk(node, i)
1073
- }
1074
- if (node[0] === 'block') {
1075
- const r = tryVectorize(node, fnLocals, freshIdRef)
1076
- ?? tryReduceVectorize(node, fnLocals, freshIdRef)
1077
- if (r) {
1078
- parent[idx] = r.wrapper
1079
- newLocalDeclsAll.push(...r.newLocalDecls)
1080
- }
1081
- }
1082
- }
1083
- for (let i = bodyStart; i < fn.length; i++) walk(fn, i)
1084
-
1085
- if (newLocalDeclsAll.length) {
1086
- fn.splice(bodyStart, 0, ...newLocalDeclsAll)
1087
- }
1088
- }