Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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