Dynamism RFC #1668
Dynamism RFC #1668
Conversation
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 - 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.
| I64ElementsAttr:$static_edge_padding_low, | ||
| I64ElementsAttr:$static_edge_padding_high, | ||
| I64ElementsAttr:$static_interior_padding, | ||
| Variadic<0DTensorOf<HLO_DimensionValue>>:$dynamic_edge_padding_low, |
There was a problem hiding this comment.
I am concerned about the usage of variadic arguments. It can make writing Declarative Rewrite Rule (DRR) more difficult (if not impossible). I am not sure if this will have some impact on the PDL dialect.
If we want to go through this approach, we should think about how we can assist declarative rewrite patterns.
Another way I can think of to unify static/dynamic shape, is to replace both of them with 1DTensorOf<HLO_DimensionValue> and provide a utility function that converts variadic arguments into a Tensor:
// Static-shape
%static_edge_padding_low = "stablehlo.constant"() {value = dense<[...]> : tensor<4xi64>} : () -> tensor<4xi64>
%static_edge_padding_high = "stablehlo.constant"() {value = dense<[...]> : tensor<4xi64>} : () -> tensor<4xi64>
%static_interior_padding = "stablehlo.constant"() {value = dense<[...]> : tensor<4xi64>} : () -> tensor<4xi64>
%pad = "stablehlo.pad"(%operand, %pad_value, %static_edge_padding_low, %static_eduge_padding_high, %static_interior_padding)
// Dynamic-shape
%edge_padding_low = "stablehlo.make_tensor"(some, value, computed, dynamically) : (i64, i64, i64, i64) -> tensor<4xi64>
%edge_padding_high = "stablehlo.make_tensor"(some, value, computed, dynamically) : (i64, i64, i64, i64) -> tensor<4xi64>
%interior_padding = "stablehlo.make_tensor"(some, value, computed, dynamically) : (i64, i64, i64, i64) -> tensor<4xi64>
%pad = "stablehlo.pad"(%operand, %pad_value, %edge_padding_low, %eduge_padding_high, %interior_padding)
The benefit of this approach is that the static-shape version looks similar to the prior arts (e.g. TFLite). Though, I am not sure if this approach has been discussed or not.
| 1) Affirm that only shape computations that use StableHLO operations are | ||
| supported in StableHLO portable artifacts (this can be changed in future | ||
| RFCs, but is out of scope for this one). | ||
| 2) Use `Variadic<0DTensorOf<HLO_DimensionValue>>` instead of |
There was a problem hiding this comment.
Could you elaborate on using HLO_DimensionValue over e.g. HLO_Int? Is the rationale for using HLO_DimensionValue that other dialects convert to Index instead of integer element type?
Following the proposal in P5 ("variadic number of 0-dimensional tensors of integer type") verbatim, I'd think the proposal would use Variadic<0DTensorOf<HLO_Int>> instead. I'm asking since the StableHLO spec currently does not define index element type even though the StableHLO dialect uses it.
There was a problem hiding this comment.
Could you elaborate on using HLO_DimensionValue over e.g. HLO_Int? Is the rationale for using HLO_DimensionValue that other dialects convert to Index instead of integer element type?
That is true. For example, Shape dialect replies heavily on index.
There was a problem hiding this comment.
"I'd think the proposal would use Variadic<0DTensorOf<HLO_Int>> instead". You're right. Using HLO_DimensionValue was an oversight on my part, this should be HLO_Int.
"Is the rationale for using HLO_DimensionValue that other dialects convert to Index instead of integer element type?". Yes, interoperability with other dialects was a factor in the design process, but in the end I decided to keep it out of scope of this RFC, and O2 goes into details about that. tl;dr is that switching to index would be a significant amount of work, which I think is best left for a follow-up RFC (this one is already huge).
| * [(P3) Ratify the existing convention for relaxed constraints already implemented in the StableHLO dialect](#p3). | ||
| * [(P4) Enable dynamism for all size-related program elements but keep all axis-related program elements static](#p4). | ||
| * [(P5) Unify static and dynamic versions of StableHLO operations](#p5). | ||
| * [(P6) Represent shape computations as StableHLO operations on 0-dimensional tensors and drop support for unranked dynamism](#p6). |
There was a problem hiding this comment.
After further discussion, P6 has been removed from this RFC, it may still be proposed in a future proposal, but some confidence should be built in the usability of this approach prior to making it the de-facto means for dynamic ops. I.e. this structure of op may be difficult to write declarative rewriter patterns for.
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 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.