Repository navigation
Datadeps: Build DAG specs faster - #776
antoniocaarvalh wants to merge 2 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.
|
The prefetch test I added assumed there would be other workers around, so it failed in the jobs that run with |
This branch has not been deployed
1 participant
I was looking at where Datadeps spends its time before a region's tasks start running, and with the schedule cache on, building the DAGSpec ended up being most of the planning time, most of it going to
dag_add_task!asking the owner of each remote chunk for its aliasing one round trip after another, and the rest to dynamic dispatch, sincevalue(arg)isn't inferrable when the spec is untyped.This does two things about it. Before planning,
prefetch_chunk_aliasing!goes through the region's tasks, collects every remote chunk they use and resolves their aliasing with oneremotecall_fetchper owning worker, all of them at the same time, seeding the region's aliasing memo with the answers, which is pretty much whatbuild_aliasing_parallelalready does in hierarchical mode, and only theDRefgoes over the wire and only the aliasing comes back, because sending the wholeChunkmade the exchange about 3x slower. The other part puts the per-argument work indag_add_task!behind function barriers and maps over the typed argument tuples withf::F, so Julia specializes on each argument's type instead of dispatching dynamically, and a cache hit now looks up each task by its id instead of scanning the region's tasks for every one of them.On a 4096x4096 matmul with 512x512 tiles and 4 workers, building the DAGSpec with the plan already cached went from 75 ms to 4.7 ms per region, and a Cholesky from 23 ms to 2.3 ms, with about a third of the allocations.
Two details worth knowing: each
DRefis sent only once per worker even when the region uses it under more than one dep_mod, because MemPool'sDRefserializer ignores whatserialize_cycle_headerreturns, so the sameDReftwice in a single message corrupts the stream and the worker either errors or hangs (I'm going to open an issue on MemPool about that one), and the prefetch is skipped under MPI, wherealiasingis a collective that has to run in the same order on every rank.I added a testset for the prefetch ("Flat prefetch of chunk aliasing") and ran the datadeps tests, all the array tests and a few others on this branch locally, all passing.