EnzymeAD / EnzymeAD/Enzyme-JAX

Lower Linear Algebra factorizations to StableHLO

Open
#1,083 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.