google / google/heir

feat: StableHLO as a front end for HEIR

Open
#738 3 comments 7 reactions 1 assignee Claimed by @asraa View on GitHub
Dominant language
MLIR
Stars
906
Forks
171
Avg merge
4d 12h
Merged PRs (30d)
32

Description

I've recently learned about StableHLO, and been pretty convinced that it should be a frontend to HEIR:

1. StableHLO is an open-source first project aiming to be a standard inside and
outside of Google. It aims to be a portability layer between ML frameworks
and ML compilers and is currently used by TensorFlow, JAX, PyTorch and XLA,
IREE, and more.

2. StableHLO has lowerings to standard MLIR without the use of bufferization,
preserving the original tensors. This would create a pathway for high-level
ML programs to RLWE schemes that utilize types and passes based on tensor
types.

3. This would also enable a frontend for quantized PyTorch models (for example,
Zama takes QAT PyTorch and ingests them into concrete-ml through the ONNX
format).

4. Support for quantizing TensorFlow models using StableHLO quantization

5. Support for qKeras models through a qKeras compilation to HLO. HLO has
parity guarantees with StableHLO. qKeras offers full integer quantization,
so would enable us to quantize models more efficiently than the existing use
of TensorFlowLite quantization.

### Will this replace TOSA?

Probably not, and it would be nice to keep both representations.

### StableHLO to standard MLIR

I have some internal code that lowers stableHLO to standard MLIR (using, func, affine loops, tensor, arith - notably not memref). Some of it uses passes with tensorflow's XLA compiler right now, so I'll attach a PR with the added dep, and perhaps create a standalone tool depending on feedback.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.