Skip to content

Commit

Permalink
Dynamism RFC (#1881)
Browse files Browse the repository at this point in the history
The current design of dynamism in MHLO and StableHLO has been
practically useful. There are success stories of it connecting JAX,
PyTorch and TensorFlow to a number of compilers, in a mix of research
and production environments. This RFC aims to leverage existing support
for dynamism in the StableHLO dialect, discuss improvements to the
existing design and then formalize the improved design in the StableHLO
opset.

The main challenge with writing this RFC was that it affects the entire
opset. The current design involves a lot of corner cases, so it took
about a year of practical evaluation by the author - primarily within
JAX native serialization, but also correlating with other ongoing and
prior projects - to distill the design into just a few general design
principles. I think I'm happy with the outcome, but please take a look
at the "Summary" section for what that entails.

This RFC addresses a considerable chunk of community feedback on the
existing design, but some feedback is deliberately left out of scope for
this RFC to enable incremental progress while areas which require
additional alignment are developed in parallel. See sections "Community
feedback" and "Out of scope" for details.

The only open question at the moment is interoperability with HLO with
respect to bounded dynamism. The proposed representation for bounded
dynamic types in StableHLO is in parity with the representation for
bounded dynamic types in HLO. However, I'm not sure whether this parity
covers 100% of bounded dynamism functionality. For example: 1) there
appears to be a mismatch in how broadcasts are represented, 2) there is
misalignment in representations of dynamic windows (HLO has a high-level
representation: VALID and SAME, whereas StableHLO expects the producers
to dynamically compute window sizes). Nonetheless, I think that this
shouldn't block the initial review of the RFC, since there's a lot of
stuff to discuss - and in the meanwhile, I'll be working on confirming
interoperability with HLO.

Finally, I'd like to acknowledge Smit Hinsu's work on the [Bounded
Dynamism RFC](#194) from Q4
2022, which was superseded by this work. The representation for bounded
dynamic types in the StableHLO dialect was designed and implemented by
Smit, and Smit's proposal to allow bounded dynamic types everywhere is
compatible with the more general proposal from this RFC to enable
dynamism for all size-related program elements. Furthermore, Smit
contributed the formal spec for get_dimension_size as well as the
informal spec for set_dimension_size.
  • Loading branch information
GleasonK authored Feb 5, 2024
1 parent d50b6fc commit 229b9a3
Showing 1 changed file with 662 additions and 0 deletions.
Loading

0 comments on commit 229b9a3

Please sign in to comment.