在撰写本文时,StableHLO 已准备好取代 MHLO/HLO 作为编译器接口。它可以由 TensorFlow、JAX 和 PyTorch 生成,可以由 XLA 和多个第三方 PJRT 插件使用,并且具有 MHLO/HLO 提供的所有公共功能以及其他功能。
本文档介绍了 StableHLO 项目的后续步骤,对问题跟踪器中反映的正在进行的工作进行分类,并将这些工作整理成计划交付成果。
当前里程碑
我们目前的里程碑遵循以下两个主要趋势:
- 最大限度地发挥 StableHLO 对整个 OpenXLA 社区的好处。
- 为所有 OpenXLA 成员统一开发者体验。
- MHLO 废弃:在 2024 年第 4 季度,我们开始探索 MHLO 的内部废弃,将包括规范化和文件夹模式在内的实用传递迁移到 StableHLO。在内部证明迁移过程非常简单后,我们计划分享一份 RFC,其中包含向 StableHLO 进行外部迁移的时间表。这可能发生在 2025 年第 1 季度,我们计划为各团队提供充足的时间和支持,以便他们在 2025 年上半年迁移到 StableHLO。
- 将与硬件无关的优化迁移到 StableHLO:根据上述趋势,我们希望 StableHLO 成为整合与硬件无关的图表简化功能的最佳位置,以便所有 PJRT 插件(包括从 StableHLO 转换为非 XLA 编译器 IR 的插件)都能获得最大优势。实现这一目标的部分工作涉及将 Google AI Edge、JAX-Enzyme 项目以及其他项目中使用的模式整合到 StableHLO 代码库中。部分整合工作已经开始,但大部分工作流将于 2025 年第 1 季度开始并完成。
- OpenXLA 组件化:我们已开始在 openxla/xla 中为 HLO 创建专用组件,这些组件类似于 StableHLO 代码库设置 (ref),并且也开始将所有 OpenXLA 后端移至 PJRT 插件后面。此外,我们还在投入资源来解决我们在这些 PJRT 插件中发现的重大用户体验问题,包括在 StableHLO 插件中实现精确的 StableHLO 版本通信,以便新插件可以立即使用新功能 (ref)。
- 使复合体在端到端工作:在 2024 年第 3 季度,我们向 HLO 添加了复合体,为抽象提供了完整的编译器堆栈支持。在 2024 年第 4 季度,我们向 XLA 内嵌器介绍了复合体,并在 HLO/StableHLO 中添加了一些传递,以便内嵌未知复合体及其分解。我们目前正在研究添加专用于从框架生成复合体的 JAX API(PyTorch API 已存在),以及添加有关如何正确使用复合体的 Colab 文档,预计将于 2024 年第 4 季度完成。
过往里程碑
2024 年上半年,我们发布了 StableHLO v1.0,其中实现了一些高优先级改进,包括完成了公共操作集规范、参考解释器、扩展了向前和向后兼容性以支持设备端部署,以及通过复合操作实现了可扩展性、规范化的动态性支持,发布了包含已评估的标准结果的完整测试数据套件等。