search

apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.46k stars 3.41k forks source link

[Disco][QoL] Implement broadcast/scatter methods for Session #17035

Closed Lunderberg closed 2 months ago

Lunderberg commented 2 months ago

Prior to this commit, use of the disco.Session API to broadcast or scatter an array required several steps from the caller.

  1. Allocate memory on worker0
  2. Transfer data from the controller to worker0
  3. Allocate memory on each worker
  4. Broadcast/scatter data from worker0 to all workers

While exposing these steps is necessary for performance, especially when used repeatedly, it can be tedious/error-prone to use for initialization that is only performed once.

This commit adds utility methods Session.broadcast and Session.scatter, which are implemented in terms of the existing lower-level methods Session.broadcast_from_worker0 and Session.scatter_from_worker0. These methods perform the transfer from the controller to worker0, and from worker0 to all other workers.

Lunderberg commented 2 months ago

This PR is identical to the now-closed PR https://github.com/apache/tvm/pull/17011, as a CI restart prevented the CI from running on the original PR.