|
2 | 2 |
|
3 | 3 | At the time of writing, StableHLO is ready to supersede MHLO/HLO as compiler |
4 | 4 | interface. It can be produced by TensorFlow, JAX and PyTorch, it can be consumed |
5 | | -by XLA and IREE, and it has all the public features provided by MHLO/HLO |
6 | | -as well as additional functionality. |
| 5 | +by XLA and several 3p PJRT plugins, and it has all the public features provided |
| 6 | +by MHLO/HLO as well as additional functionality. |
7 | 7 |
|
8 | 8 | This document describes the next steps for the StableHLO project, categorizing |
9 | 9 | the ongoing work reflected in the issue tracker and arranging this work into |
10 | 10 | planned deliverables. |
11 | 11 |
|
12 | | -## Milestones |
13 | | - |
14 | | -In 2023, we are planning two big milestones: 1) StableHLO v0.9 which will |
15 | | -provide an initial version of the opset and initial compatibility |
16 | | -guarantees, 2) StableHLO v1.0 which will implement high-priority improvements |
17 | | -and start providing full compatibility guarantees. |
18 | | - |
19 | | -**StableHLO v0.9** will mirror MHLO/HLO, augmented with a specification for |
20 | | -statically-shaped ops and initial compatibility guarantees. Per the |
21 | | -[compatibility RFC](https://github.com/openxla/stablehlo/blob/main/rfcs/20220912-compatibility.md), |
22 | | -this release will provide 1 month of forward and backward compatibility. These |
23 | | -modest guarantees will enable gaining experience with dialect evolution and |
24 | | -allow some time for cleanup before full guarantees go into effect. We are |
25 | | -planning to release StableHLO v0.9 in Q1 2023. |
26 | | - |
27 | | -**StableHLO v1.0** will implement high-priority improvements, including |
28 | | -cleaning up the frontend contract (with the goal that StableHLO programs only |
29 | | -include ops from the StableHLO dialect, rather than today's mixture of dialects |
30 | | -and unregistered attributes) and providing a reference implementation. We are |
31 | | -planning to release StableHLO v1.0 in H2 2023. |
32 | | - |
33 | | -## Workstreams |
34 | | - |
35 | | -In order to organize the development towards the aforementioned milestones, |
36 | | -we have categorized the tickets in the issue tracker into multiple workstreams |
37 | | -and tied these workstreams to the milestones. A limited number of tickets |
38 | | -(less than 10%) are not assigned to any specific workstream and aren't part of |
39 | | -any specific milestone. |
40 | | - |
41 | | -(P0) The |
42 | | -[Compatibility Implementation](https://github.com/orgs/openxla/projects/4) |
43 | | -workstream is dedicated to implementing |
44 | | -[the compatibility RFC](https://github.com/openxla/stablehlo/blob/main/rfcs/20220912-compatibility.md) |
45 | | -along with a compatibility test suite. Most of this work is expected to be |
46 | | -completed in StableHLO v0.9, and the rest will be done in StableHLO v1.0. |
47 | | - |
48 | | -(P0) The |
49 | | -[Frontend Contract](https://github.com/orgs/openxla/projects/6) workstream |
50 | | -consists of implementing 100% of the features which are used by StableHLO |
51 | | -frontends but are not yet in the StableHLO specification. The goal of this |
52 | | -workstream is to ensure that StableHLO programs only include ops from the |
53 | | -StableHLO dialect, rather than today's mixture of dialects and unregistered |
54 | | -attributes. We are planning to complete all or almost all work in this |
55 | | -workstream in StableHLO v1.0. |
56 | | - |
57 | | -(P0) The |
58 | | -[Reference Implementation](https://github.com/orgs/openxla/projects/7) |
59 | | -workstream organizes the work on implementing |
60 | | -[an interpreter](https://github.com/openxla/stablehlo/blob/main/docs/reference.md) |
61 | | -for 100% of StableHLO ops as defined in the StableHLO specification. We are |
62 | | -planning to complete all or almost all work in this workstream in |
63 | | -StableHLO v1.0. |
64 | | - |
65 | | -(P0) The [Documentation](https://github.com/orgs/openxla/projects/12) workstream |
66 | | -is dedicated to providing all the information that StableHLO producers or |
67 | | -consumers might need. The StableHLO specification is a major deliverable, as |
68 | | -well as a reference for the StableHLO API and StableHLO serialization format. |
69 | | -Critical pieces of the workstream will be delivered in StableHLO v1.0, with |
70 | | -lower-priority items addressed on a rolling basis. |
71 | | - |
72 | | -(P1) The [Conformance Suite](https://github.com/orgs/openxla/projects/8) |
73 | | -workstream consists of delivering a test suite that compares a reference |
74 | | -implementation with implementations provided by StableHLO backends. Tests for |
75 | | -the reference implementation will provide a conformance suite of sorts, so this |
76 | | -workstream does not have P0 priority. However, further augmenting this suite |
77 | | -with additional interesting test cases will likely be a useful area for future |
78 | | -work. |
79 | | - |
80 | | -(P1) The |
81 | | -[Specification Compliance](https://github.com/orgs/openxla/projects/9) |
82 | | -workstream ensures that 100% of StableHLO ops are implemented in the StableHLO |
83 | | -dialect as defined in the StableHLO specification. The StableHLO dialect is |
84 | | -already reasonably compliant, so this workstream does not have P0 priority, |
85 | | -but a lot of minor items still remain (especially in the corner cases of |
86 | | -verifier implementation) and will be addressed on a rolling basis. |
87 | | - |
88 | | -(P1) The [New Features](https://github.com/orgs/openxla/projects/10) workstream |
89 | | -concludes the StableHLO roadmap and consists of a ragtag collection of new |
90 | | -functionality for the StableHLO opset (not the StableHLO dialect or the |
91 | | -StableHLO bindings - that would be other workstreams). A few of these new |
92 | | -features are planned for delivery in StableHLO v1.0, but |
93 | | -the majority are currently lower-priority items which are not part of any |
94 | | -specific milestone. |
95 | | - |
96 | | -(P1) The [Public API](https://github.com/orgs/openxla/projects/5) workstream is |
97 | | -dedicated to delivering C/C++/Python bindings for the StableHLO dialect. |
98 | | -Existing C++/Python bindings are already fairly reasonable, so this workstream |
99 | | -does not have P0 priority. However, quite a bit of work still remains to be |
100 | | -done, especially around providing stability for these bindings - which is |
101 | | -something that is not currently covered by the compatibility RFC but will |
102 | | -likely be a useful area for future work. |
| 12 | +## Current Milestones |
| 13 | + |
| 14 | +Our current milestones follow two main trends: |
| 15 | + |
| 16 | +1. Maximize the benefit of StableHLO for the entire OpenXLA community. |
| 17 | +2. Unify the developer experience for all OpenXLA members. |
| 18 | + |
| 19 | +- **MHLO Deprecation**: In Q4'24 we have began exploring internal deprecation of |
| 20 | + MHLO, migrating useful passes including canonicalization and folder patterns |
| 21 | + to StableHLO. Once the migration process is proven trivial internally, we plan |
| 22 | + to share an RFC with timelines for external migrations to StableHLO. This will |
| 23 | + likely happen in Q1'25, and we plan to give ample time and support to teams to |
| 24 | + migrate to StableHLO in H1'25. |
| 25 | +- **Migrate hardware independent optimizations to StableHLO**: Following the |
| 26 | + trend above, we want StableHLO to be the best place to consolidate hardware |
| 27 | + independent graph simplifications, so that all PJRT plugins including those |
| 28 | + which convert from StableHLO to a non-XLA compiler IR can see max benefits. |
| 29 | + Part of this goal involves consolidating patterns used in |
| 30 | + [Google AI Edge](https://io.google/2024/explore/18c47ed9-a8f7-4cd5-aec2-80457d839942/), |
| 31 | + the [JAX-Enzyme](https://github.com/EnzymeAD/Enzyme-JAX) project, and other |
| 32 | + projects all in the StableHLO repo. Some of this consolidation has already |
| 33 | + began, but the workstream will largely pick up and complete in Q1'25. |
| 34 | +- **OpenXLA Componentization**: We have began creating dedicated components in |
| 35 | + openxla/xla for HLO which resembles the StableHLO repo setup |
| 36 | + ([ref](https://github.com/openxla/xla/tree/main/xla/hlo)), as well as started |
| 37 | + moving all OpenXLA backends behind PJRT plugins. We are additionally investing |
| 38 | + in fixing prominent UX issues we discover in these PJRT plugins, including |
| 39 | + things like having precise StableHLO version communication in StableHLO |
| 40 | + plugins, so new features can be used immediately by new plugins |
| 41 | + ([ref](https://github.com/openxla/xla/commit/84ad5fab88e3979b9c43ce93089e0ef537d14b88)). |
| 42 | +- **Make composites work e2e**: In Q3'24 we added composites to HLO, enabling |
| 43 | + full compiler stack support for abstractions. In Q4'24 we taught the XLA |
| 44 | + inliner about composites and added passes in HLO/StableHLO for inlining |
| 45 | + unknown composites with their decompositions. We are now investigating adding |
| 46 | + dedicated JAX APIs for generating composites from the framework (PyTorch APIs |
| 47 | + already exist), as well as adding Colab documentation on how to properly use |
| 48 | + composites, to be completed in Q4'24. |
| 49 | + |
| 50 | +## Past Milestones |
| 51 | + |
| 52 | +In H1 2024, we released **StableHLO v1.0** which implemented high-priority |
| 53 | +improvements, including finishing the public opset spec, reference interpreter, |
| 54 | +extending forward and backward compatibility to support on-device deployment |
| 55 | +with compatibility unit testing, extensibility via composite operations, spec'ed |
| 56 | +dynamism support, released a full testdata suite with evaluated golden results, |
| 57 | +and more. |
0 commit comments