Skip to content

Datadeps: Allocate uninitialized slots with one round trip - #777

Open
antoniocaarvalh wants to merge 1 commit into
JuliaParallel:aot-schedulers-rebasedfrom
antoniocaarvalh:datadeps-direct-slots
Open

antoniocaarvalh wants to merge 1 commit into
JuliaParallel:aot-schedulers-rebasedfrom
antoniocaarvalh:datadeps-direct-slots

Conversation

@antoniocaarvalh

Copy link
Copy Markdown

Following the DAGSpec PR (#776), I measured where the rest of a flat-mode region's launch loop goes, and about 90% of it was creating slots: each one 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, so it's about seven round trips in series per slot, while the task submissions themselves were only 3-4% of the loop. On a 4096x4096 matmul with 512x512 tiles and 4 workers under RoundRobin, that was 2.8 s of a 3.4 s loop for 461 slots.

When the slot is only going to be allocated, which is the case whenever slot_may_be_uninit lets the copy-to phase fill it, the source worker has nothing to contribute, so generate_slot! now does the cache lookups itself and allocates with a single round trip to the destination (alloc_slot_with_aliasing!), which also brings the new slot's aliasing back and memoizes it so planning doesn't have to ask for it again. direct_uninit_slot decides when that's possible from the chunk alone: only under Distributed, since with MPI the cache is replicated and every rank takes part, and only for dense arrays that aren't wrappers and whose size the chunk's domain already gives, with everything else going through move_rewrap as before.

With that, the RoundRobin launch loop on the same matmul got about 30-40% faster, and Greedy about 10%, since it creates far fewer slots. What's left is mostly the one allocation round trip per slot, which waits behind the region's copies on the destination, and the follow-up I'm working on reserves that memory up front so the loop doesn't wait on it at all.

The new path only changes how the slot gets allocated, not where or with what: for a GPU destination it ends up in the same alloc_uninit the old path used, but I don't have a GPU to test that here. I ran the datadeps tests and all the array tests on this branch locally, plus a few others, all passing.

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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant