Google announced a new open-source machine learning compiler at its recent Next ‘22 conference. The compiler was built in collaboration with the leading names in machine learning and affiliated hardware, with the goal of lessening incompatibility across frameworks and hardware.
The OpenXLA Project will be headed by Google and includes AMD, Arm, Apple, AWS, Intel, Meta, Nvidia, and other AI/ML developers. It works with leading machine learning frameworks such as TensorFlow, PyTouch, and JAX and ensures that they run at optimal rates on hardware such as GPU, CPU, and ML accelerators.
“ML development is often stymied by incompatibilities between frameworks and hardware, forcing developers to compromise on technologies when building ML solutions,” said vice president and GM of Infrastructure at Google Cloud, Sachin Gupta. “OpenXLA will address this challenge by letting ML developers build their models on leading frameworks and execute them with high performance across hardware backends. This flexibility will let developers make the right choice for their project, rather than being locked into decisions by closed systems.”
The OpenXLA project is currently working on decoupling XLA from TensorFlow, which will allow further collaboration and third-party development. One of the first tasks is to build StableHLO, a portable ML compute operation that would serve as a portability layer for frameworks and compliers.
Google will work with its partners to build out OpenXLA, with some of the key objectives, including promoting XLA industry collaboration, solicit input on OpenXLA technical direction, create new XLA repository or organization independent of the hardware and framework, and set up governance outside of TensorFlow.
Google has been one of the largest contributors to the Cloud Native Computing Foundation (CNCF), most notably donating Kubernetes to The Linux Foundation, which then graduated it to CNCF to maintain the project. While it is not confirmed that CNCF will have some role in the future of OpenXLA, Google does mention it is a lead contributor to the project in the second paragraph of its OpenXLA announcement.