EnzymeAD / EnzymeAD/Enzyme-JAX
Lower Linear Algebra factorizations to StableHLO
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
Instead of using `stablehlo.custom_call` to call LAPACK or TPU kernels, it would be cool to code the numerical methods on top of other StableHLO ops. This has some nice advantages:
- Shardy could know how to automatically shard the linalg factorizations
- Posibly AD could work on linalg factorizations without new rules (although new rules might be better)
- Interesting to see if there's any performance improvement using the opt passes implemented in this repo
- Portable implementations to platforms whose backend do not yet support some of these methods (e.g. cuSOLVER doesn't support all methods available in the LAPACK reference)
# To do
- [ ] LU
- [ ] QR
- [ ] Eigendecomposition
- [ ] SVD
- [ ] QR iteration (LAPACK's `gesvd`)
- [ ] Divide and conquer (LAPACK's `gesdd`)
- [ ] Jacobi method (LAPACK's `gesjv` and `gesvj`)
- [ ] QR-based dynamically weighted Halley iteration (QDWH, used by Jax)
- https://epubs.siam.org/doi/abs/10.1137/120876605
- https://epubs.siam.org/doi/abs/10.1137/090774999
- [ ] Lanczos algorithm for top-k eigenvalues
Contributor guide
No contributing guide indexed for this repository
Research direction
The issue names no files, tests, or concrete entry point. Start by narrowing the checklist to one factorization and locating the existing stablehlo.custom_call path for LAPACK or TPU kernels. Done would mean a portable implementation using StableHLO operations, with its sharding, autodiff, portability, and performance implications evaluated.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100