Feat/array map - #911
Open
jtranq wants to merge 6 commits into
Open
Feat/array map#911jtranq wants to merge 6 commits into
jtranq wants to merge 6 commits into
Conversation
…ray.map.seq
Array.map is written over the tree the language presents. A match on ANode calls blk_half, which allocates a half and copies it, and ANode{xs, ys} calls blk_node, which allocates the join and copies both halves back. Mapping n elements therefore copies O(n log n) words where the map itself only does O(n) work.
An indexed walk avoids that. It gets the size, allocates the result once with Array.new, then reads each cell with Array.get and writes each result with Array.set, all of which lower to direct block operations. That walk is sequential, and it needs Data elements, because get hands an element out while keeping the array and Array.new fills the destination with a copy of the first result.
Array.map keeps the tree walk. Its two recursive calls fork, so subtrees map in parallel, and it stays Type-generic, so Array<U32> to Array<Array<U32>> and back still check. The indexed version is available as Array.map.seq for cheap callbacks over Data elements, where it is about 8 times faster on 2^24 U32 elements, 30 ms against 520 ms on this machine, with peak memory dropping from about 720 MB to 135 MB.
The test covers both paths, including the nested array cases and the boxed element cases.
The minted Array.map instance is now compiled specially. The tree body splits and rebuilds a block at every level, so a map over n elements copies O(n log n) words. When the callback inlined into the ALeaf arm is flat, the compiler walks the block by index instead: it allocates the destination once, moves each source cell out exactly once, emits the leaf body, and writes the result straight into its slot with a raw write. Nothing reads or drops a destination cell, so there is no seed value and no Data element requirement. Differing input and output layouts, boxed elements, nested arrays and padded multi-word records all go through it. This is the flat-callback checkpoint. A callback that needs a continuation or its own fork still takes the tree body, and the JavaScript lane is unchanged, so the optimized traversal is not yet complete and parallel mapping of flat callbacks is not preserved. On a 2^24 element U32 array through an unchanged Array.map call, the compiler's C lane goes from about 0.10 s and 715 MB peak resident memory to 0.03 s and 135 MB, with no blk_half or blk_node in the emitted map.
…leases left elements A packed numeric destination is a BUF block, so its physical class is buf_wcls of the logical class. Allocating it in the logical class and freeing it in the physical one sent every block to a different size list, so a repeated map never reused a block and resident memory grew with each turn. A 20,000 turn map of a 4096 element U32 array went from about 617 MB resident to 1.6 MB, and the loop is a regression test now. The indexed pass also skipped the cleanup a normal body does after an expression. A callback compiled as borrowing its input leaves the map owning that element, so the element was never released and the final shallow free of the array lost it. The loop now runs the same binding analysis a body end does and sinks exactly the bindings the leaf left. The range worker with ordinary callback continuations and preserved parallel execution is still to come. This commit is the two fixes to the flat checkpoint.
The minted Array.map instance is rewritten in the C lane into three defs the emitter compiles like any other. The entry takes the array as a raw word, allocates the destination once in its physical class and walks the whole index range. The walk forks on halves down to a leaf, so its fork is the one the emitter always emits: tasks while a lane grows, frames once it winds back, on the cores and on the device alike. The leaf moves each element out, applies the callback and writes the result raw into its slot, and the source is freed shallow once every element has moved. A range is Nat words nobody owns, so no half is ever an array value. Eleven bodiless intrinsics under Array.map.* do the raw block work, and SYNTH keeps their names.
A callback keeps every shape a body has. A flat one compiles into the leaf's spin loop; one that calls a def needing a continuation becomes a cut with the emitter's own frame; a parallel let inside it goes through anf as in any def. The previous indexed pass replaced the whole map with one loop, so it ran every callback serially, fell back to the copying tree for non-flat callbacks, and crashed the compiler on a parallel let in the callback ("an unbound binder") and on a call to a def with no arguments ("a live call into the law"). Both are regression tests now (array_map_cont).
The leaf size is chosen per instance from the callback: 2^12 elements when it is straight-line C (intrinsics, constructors and defs that neither loop nor fork), 2^3 otherwise, so a small cheap map never forks and an expensive one spreads across the cores. On this machine with 10 threads, a map of 2^24 U32 goes from 0.14 s and 713 MB on main to 0.03 s and 130 MB; 20,000 maps of 4096 elements from 2.75 s to 0.01 s; a heavy callback over 4096 elements stays at 0.12 s on both. The definition, its signature, the interpreter and the JS lane are unchanged. bend2/comp.ts now measures 67011 ttok against the repo gate's 64000 cap.
Each range worker returns the index after its range. A join combined the two halves' indices with Nat.add, which summed two absolute positions into a number nothing read: the close ignores it, so the map was right, but the arithmetic meant nothing and cost a step at every join. The join is now Array.map.join, which forces the low half and returns the high half's end, so the walk's result is the index after the whole range at every level.
The leaf of a lowered Array.map under an expensive callback was 2^3 elements from the first level, so an array of eight or fewer such elements never forked: four callbacks of fifty million steps each ran on one thread in 0.28 s where the tree map ran them on four in 0.14 s. The leaf is now bounded from both sides: at most 2^3 elements, and no wider than leaves a split of 2^6 leaves to fork over, so a small array splits down to single elements and a large one still batches. The cheap callback keeps its rule, a leaf of up to 2^12 and no forced split, since a fork there costs more than the elements it would divide. The four-element map is back to 0.14 s on ten threads and on four, and the 4096-element heavy map, the 2^24 light map and the 20,000 small maps measure as before.
jtranq
marked this pull request as ready for review
September 20, 2026 19:02
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR removes
Array.map'sO(n log n)split/rebuild copying and replaces it with direct destination construction.On a
4096 × 4096transformer weight-initialization workload, the same Bend program drops from 2.01 s to 0.143 s (~14× faster), moving from ~16× slower than C to within 14% of C, while peak RSS falls from 194 MB to 130 MB.Optimize the existing
Array.mapin the native compiler by replacing recursive array splitting and rebuilding with compiler-generated workers that read the source and write directly into a single destination allocation. This addresses a mismatch between the abstraction and the representation. Arrays expose a tree-shaped interface but their native storage is contiguous.The lowering is very similar to Futhark's sequentialisation and array short-circuiting passes, which compile map into indexed traversal over preallocated result storage and push result construction into its final destination to avoid redundant copies. You can read a deeper description of that by Troels Henriksen at the above hyperlink.
Existing callers benefit with no code changes. The public signature remains
Type-generic, including support for affine elements, nested arrays, boxed values, and different input/output layouts. Parallel mapping and callbacks requiring continuations or containing their own forks remain supported. The key observation is that the new algorithm is fundamentally doing less data movement than the old one. So the 3–6% regressions are likely to be overhead introduced by the new worker/scheduling machinery.One basic result I was able to get was with a materializing FizzBuzz implementation. The
Array.mapformulation became about 1.85× faster and substantially reduced its memory use. Currently theArray.mapFizzBuzz formulation is not the fastest way to materialize the result, aList-based implementation already avoided the pathological array splitting and rebuilding, but now theArray.mapimplementation has become roughly 10% faster than theListbased one, making it the fastest materializing Bend formulation I tested.Performance
This removes the pathological overhead of cheap elementwise maps producing an ~8.7× improvement on a large trivial map, but it remains basically performance neutral on compute-heavy callbacks, you can see the other cases are within ~5% of the existing implementation though sometimes mildly slower.
U32elementsThe large cheap map is approximately 8.7× faster. The 4,096-element expensive map retains approximately 3.6× scaling from one to four workers, although its four-worker median is about 6% slower than the tree baseline in this run.
For the same large numeric mapping workload:
That is approximately 33% less peak resident memory. This implementation still uses separate input and output buffers rather than reusing the input in place.
Practical example: transformer weight initialization
As a more representative workload, I tested initializing transformer-style weight matrices using
Array.map.The benchmark initializes
4096 × 4096F32weight matrices using a Kaiming-uniform bound of1 / sqrt(fan_in), with each weight generated deterministically from its global index using a counter-based hash. This gives the map callback a realistic amount of work in hashing, integer-to-float conversion, and a few floating-point operations, while remaining a pure pointwise transformation.One worker
Each row below generates the same total number of weights (
2^27), but groups them into differently sized matrices:2^20weights2^24weights2^26weightsFor the
4096 × 4096case, the whole Bend workload goes from roughly 16× slower than C to within ~14% of it on one worker.Subtracting the index-generation and result-folding baseline makes the effect on
Array.mapitself clearer:2^20weights2^24weights2^26weightsPreviously, most of the cost of this operation was not generating the weights at all: it was recursively splitting the input block and rebuilding the output block. With this PR, each generated weight is instead written directly into its final destination slot.
Parallel execution
Using Bend's default 10 workers:
2^20weights2^24weights2^26weightsThe memory difference becomes especially visible at larger tensor sizes. Producing a 256 MB weight matrix previously reached roughly 3.4 GB of resident memory with 10 workers; the new lowering reduces that to roughly 515 MB.
On this workload, the new one-worker implementation lands within roughly 12–20% of the C reference, while the parallel Bend implementation is roughly competitive with the single-threaded C implementation.