Repository navigation
Datadeps: Carve slots out of per-region arenas - #782
Draft
antoniocaarvalh wants to merge 4 commits into
Draft
antoniocaarvalh wants to merge 4 commits into
antoniocaarvalh wants to merge 4 commits into
Conversation
With the schedule cache on, building a region's DAGSpec was most of its planning time, and most of that went to dag_add_task! asking the owner of each remote chunk for its aliasing one round trip at a time, plus dynamic dispatch on arguments whose type isn't inferrable for an untyped spec. Resolve the aliasing of every remote chunk up front, with one remotecall_fetch per owning worker issued concurrently, and seed the region's aliasing memo with the answers, the same way build_aliasing_parallel already does in hierarchical mode, sending only each chunk's DRef and bringing back only its aliasing. Handle each argument in dag_add_task! behind a function barrier, mapping over typed argument tuples so every call is specialized on its argument's type, and look up a cached plan's tasks by id instead of scanning the region for each one. A 4096x4096 matmul with 512x512 tiles on 4 workers now builds its DAGSpec in 4.7 ms instead of 75 ms with the plan already cached, and a Cholesky in 2.3 ms instead of 23 ms, with about a third of the allocations. Each DRef goes once per worker even when the region uses it under several dep_mods, since MemPool's DRef serializer ignores what serialize_cycle_header returns and the same DRef twice in one message corrupts the stream. The prefetch is skipped under MPI, where aliasing is a collective that has to run in the same order on every rank.
The prefetch test spawned its finished task on the last worker and always expected that task's result in the memo, but in the CI jobs that run with no workers besides the first (-p 0) the task runs locally, so its result is a local chunk that prefetch_chunk_aliasing! is right to leave alone. Only expect it when it ran on another worker.
Creating a slot went through move_rewrap on the source worker, which calls back to the planning process for every lookup in the aliased object cache and then on to the destination to allocate, about seven round trips in series per slot, and in flat mode that was most of the launch loop: a 4096x4096 matmul with 512x512 tiles on 4 workers spent 2.8 s of a 3.4 s loop creating 461 slots under RoundRobin. When the slot is only going to be allocated, because slot_may_be_uninit lets the copy-to phase fill it, the source worker has nothing to add, so generate_slot! now does the cache lookups itself and allocates with a single round trip to the destination that also brings back the new slot's aliasing. This only applies under Distributed and to dense, non-wrapper arrays whose type and size the chunk already tells, and everything else keeps going through move_rewrap. The RoundRobin launch loop on that matmul gets about 30-40% faster.
Even with one round trip per slot, a flat launch loop still waited on an allocation for every slot, queued behind the region's copies on the destination. Instead, before the launch loop, reserve one arena per destination worker, sized exactly from the plan for AOT schedulers and grown on demand for the rest, and hand slots out of it from the planning process as base + offset, so each slot's ArenaRef and aliasing are known without asking anyone. The arenas are freed whole when the region ends, right away when it succeeds, and in the background once its tasks are done when one of them fails, since the scheduler fails a task's dependents outright and a per-slot free task would never run. Only dense CPU arrays of isbits elements use arenas; everything else allocates as before. A 4096x4096 matmul with 512x512 tiles on 4 workers now launches its RoundRobin region in 0.17 s instead of 3 s, and the whole region takes 1.97 s instead of 3.3 s.
jpsamaroo
added a commit
that referenced
this pull request
Oct 23, 2026
Port the allocator from PR #782 (cd28463), replacing its synchronous region scope with explicit DataDepsContext ownership. Async kernels, write-back, pristine prefill and partition-local pools now share a bounded epoch; full drains reclaim it after all inflight work and fences retire. Targeted drains and deferred dirty write-back retain their storage. Exclude arena handles from both slot-cache admission paths and preserve cold cacheable DRef slots. Foreign drains wait for the planner to publish all tasks and retiring cache leases before reclaiming buffers. Logging uses ordinary slots, and expired identity aliasing reads descriptors only. Add poison-on-free coverage for async carry-in, failed siblings, cache reuse, AOT prefill, all drain paths, planning races, budget contention and tiled GEMM. CPU Distributed 83 arena assertions and 24 unchanged allocation bounds pass; local Datadeps, Distributed Datadeps and MPI CPU also pass. Eight helper inference checks return concrete/union types without Core.Box. Paired combined PR773/776/777/782 workload versus eff8dd5, driver GC minima after ten warmups + GC over seven runs (16 64x64 chunks, two 2-thread workers, eight regions): flat sync 766087 / 39004424 B -> 421950 / 22266968 B; flat async 205161 / 10445592 B -> 166791 / 8693008 B. Median times 170.619 -> 91.820 ms and 43.938 -> 36.770 ms. Hierarchical bytes stay within 1%; sync median 93.237 -> 97.852 ms, async 33.994 -> 32.798 ms. Counts are process-local and exclude native arena/worker bytes. Co-authored-by: antoniocaarvalh <antoniocaarvalho2@gmail.com>
This branch has not been deployed
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
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.
This is the arena version of the AOT allocator idea, and it builds on #776 and #777, so only the last commit is new here (I'll rebase once those are in). Before the launch loop, the planning process reserves one block of memory per destination worker, sized exactly from the plan for AOT schedulers and growing on demand from 64 KiB up to 64 MiB for the rest, and then hands out the slots itself as base + offset, so every slot gets an
ArenaRefhandle and a real address and aliasing at plan time without a single round trip in the launch loop. At the end of the region each arena is freed whole instead of spawning a free task per slot, right away when the region succeeds, and from a background task once the region's tasks are done when one of them fails, since the scheduler fails a task's dependents outright and a free task would never get to run, which leaveswait_alland the error behavior exactly as they were. It works in flat and hierarchical mode, and only dense CPU arrays of isbits elements use it, everything else allocates like before.On a 4096x4096 matmul with 512x512 tiles and 4 workers, the flat RoundRobin launch loop went from about 3 s to 0.17 s and the whole region from 3.3 s to 1.97 s, flat Greedy from 1.5 s to 1.31 s and hierarchical RoundRobin from 1.67 s to 1.19 s, Cholesky got 14-34% faster in all four modes, and LU about 15% faster under flat RoundRobin and about the same in the others.
There are still a few things I'm not sure about, which is why this is a draft: slots now live until the end of the region instead of being freed after their last use, so the peak memory of a big region goes up and maybe this should be opt-in, GPU slots aren't covered yet, and since the addresses are known at plan time the arena could be registered once for RDMA, which is what came up in the NVSHMEM discussion. I ran the datadeps tests (with and without extra workers) and the linalg tests on this branch, and the full test suite on an earlier version of these changes.