tree 18559cc5af91810bae55822bc3c4faa72eba8717
parent 9f8b9cf19d5887b889a14f3768eb3fb5800bf804
author Ben Vanik <benvanik@google.com> 1638220478 -0800
committer GitHub <noreply@github.com> 1638220478 -0800
gpgsig -----BEGIN PGP SIGNATURE-----
 
 wsBcBAABCAAQBQJhpUK+CRBK7hj4Ov3rIwAATpMIADgKZs0rd9kZ/MU4Qwrjzevh
 QqfYLlWrsxOFCjSo8U7/FMnpdUJJciCrL2svo2Npzy6S57k93+u+B/NXrDo3RWgb
 mJKDHHMcM4042XDs9H/EmYb6b94IHm3yzBqZ9RqIGjYEpIxkYYFwYmjE+gSh1aMH
 33JzhT2zqSvIGWxxfGtSwo5e6cTT1y6nv+8i7DqcJqcFVDDM9S4uOTGjlqmxuCea
 TKas0C5ax0zwL6b27Lh2pO1DTLXxLDWXwE1cjMhnC8QATbyP1ig5kFAZHUOr2jsV
 OYWHr6Ex8+Cln4KAjuxjHnbFyJsWrOlW5EoYIG57OHlJ2qNa32SvNmeD+1nrs2Y=
 =+GJZ
 -----END PGP SIGNATURE-----
 

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