Optimising Workload DAGs for a Work-Stealing AMT: A Compiler for the DVM

Imec is developing the Distributed Virtual Machine (DVM): an Asynchronous Many-Task (AMT) system built on a distributed Write-Once-Read-Many (WORM) memory model, using random work stealing to maximise parallelism and resource utilisation. Workloads are represented as Directed Acyclic Graphs (DAGs): a neural network is ultimately such a DAG, and any pure function (in the functional-programming sense, as in Haskell) maps directly onto one.

The performance of such a system depends critically on how the DAG is structured. Tasks should be small enough to expose parallelism and let work stealing balance the load, yet because stealing moves tasks to arbitrary nodes, the data they consume or produce must travel over the network, so overly fine tasks inflate communication cost. This PhD develops graph algorithms that transform the initial workload DAG into a semantically equivalent DAG optimised for the target runtime, jointly optimising task granularity, inter-task data partitioning, and synchronisation — without constraining the dynamic scheduler. A key tension is that global optimisation requires a global view, whereas work stealing is inherently decentralised; the approach therefore transforms the DAG statically while keeping scheduling dynamic and work-stealing based.

Several transformation classes will be investigated, underpinned by a cost model of compute, memory and network. Task/kernel fusion, as in XLA and NVIDIA's Legate/Diffuse on the Legion AMT, aggregates sequential tasks to eliminate unnecessary data movement. Activation checkpointing trades memory against recomputation in backpropagation — especially relevant under WORM, where recomputation is safe and conflict-free but the absence of in-place updates raises memory pressure. Collective operations such as all-reduce introduce synchronisation barriers that undermine the AMT; by redrawing the DAG these effects can be significantly relaxed, which again requires a global view. Existing systems provide a starting point but each addresses only part of the problem: Legate optimises the DAG only locally and relies on static load balancing rather than work stealing, XLA has a global view but targets static, non-AMT backends, and related work such as FlexFlow, Alpa, TASO, DaCe and classical partitioners like METIS covers individual aspects only.