commit 38c6418477361a8edd831a96bfcfeb9b29e112d1
author Ben Vanik <benvanik@google.com> Mon Nov 29 13:14:38 2021 -0800
committer GitHub <noreply@github.com> Mon Nov 29 13:14:38 2021 -0800
tree 18559cc5af91810bae55822bc3c4faa72eba8717
parent 9f8b9cf19d5887b889a14f3768eb3fb5800bf804

Exposes control over the partitioning favor mode on the command line. (#7755)

The `#stream.partitioning_config` attribute can be added to the
module/functions/etc to control this explicitly and otherwise the
`-iree-stream-partitioning-favor=` flag will be used as a default.

- `debug` disables all partitioning for concurrency (and in the future
   when partitioning command buffers for concurrency).
- `min-peak-memory` attempts to preserve the debug-level memory
   consumption by only optimistically scheduling work concurrently.
- `max-concurrency` reorders ops to try to extract the maximum amount
   of concurrency allowed by the program - it's not a fantastic algorithm
   and in the future we can reduce the memory consumption.

Models primarily designed for single-threaded execution should favor
`min-peak-memory` while those for multithreaded/GPU should favor
`max-concurrency`. We can extend these modes in the future and make them
smarter: for example, min-peak-memory should perform reordering to
minimize live ranges.

Examples from the JAX mnist_train model:
`-iree-stream-partitioning-favor=debug`:
  12MB transients, 0 concurrent dispatches
`-iree-stream-partitioning-favor=min-peak-memory`:
  12MB transients, ~80 concurrent dispatches
`-iree-stream-partitioning-favor=max-concurrency`:
  22MB transients, ~350 concurrent dispatches