As most of you probably know, the OpenXLA StableHLO operation set is intended as a compiler IR for ML computations and a portability layer between ML frameworks and ML compilers. The “stable” in StableHLO refers to the intended stability of changes to the definition, making this operation set a suitable target IR for a wide variety of ML frameworks including TensorFlow, JAX and PyTorch, as well as a suitable source IR for ML compilers including XLA and IREE.
Because sparse tensor algebra is expected to become rather ubiquitous in many ML problems, the Google Sparse Compiler team has filed an RFC that proposes adding sparsity support to StableHLO, in particular by introducing a new sparse tensor type that combines transparently with a restricted subset of StableHLO operations. The proposal is based on our, by now, hopefully, well-known principle that sparsity should be a property, and not as a tedious implementation detail the model writer has to deal with explicitly.
Analogous to any retargetable design, we believe having a standard IR with sparse tensor types will be central in providing a healthy sparse ecosystem, where various ML frameworks map to the same IR with sparsity support, and various back-ends understand this IR and map it to efficient code for different targets.
Sparse JAX --\ /--> sparse-aware backend Target1
PyTorch/sparse tensors --> StableHLO --
Other ML framework -- / +Sparsity \--> sparse-aware backend TargetN
The initial RFC shown above has been recognized by the OpenXLA governance body as something worthwhile to pursuit further. In the next few months, the Google Sparse Compiler is planning to file more concrete proposals on how to extend StableHLO with sparsity support. This will go hand-in-hand with the surface syntax changes we are making to the MLIR sparse tensor types (see e.g. this progress posting).
If you are interested in this idea, please provide feedback and hopefully positive support to the proposals. We are looking forward making the world a sparse-friendlier place!