This is an issue for tracking our progress towards porting all "custom calls" within JAX core to the new XLA FFI.

Motivation

This project is worthwhile for a few reasons, including:

1. It gets us closer to being ready to deprecate the legacy custom call API and the related JAX internals.
2. When ported to the FFI, custom calls support export, shape polymorphism, and ahead-of-time lowering without access to the relevant device.
3. The FFI supports useful features like device-specific scratch space allocators and (soon) access to the XLA thread pool for CPU parallelization.
4. The lowering logic can be significantly simplified, and in most cases the lowering rule is now the same for GPU and CPU.

Progress

CPU

The main CPU kernels are all in jaxlib/cpu and @pparuzel has been leading the work to port these kernels:

• LU
  • Port LU Decomposition to XLA's FFI #21580
  • Activate LU Decomposition to XLA's FFI #22707
• QR
  • Port QR Factorization to XLA's FFI #21592
  • Determine LAPACK workspace during QR Factorization Kernel runtime #23029
• Householder
  • Port Householder Product to XLA's FFI #21597
• Cholesky
  • Port Cholesky Factorization to XLA's FFI #21574
  • Activate Cholesky Factorization Kernel to XLA's FFI #21727
• SVD
  • Port Singular Value Decomposition to XLA's FFI #21589
  • Determine LAPACK workspaces during SVD kernel runtime #22912
  • Activate Singular Value Decomposition to XLA's FFI #22833
• Eigenvalue decomposition
  • Port Eigenvalue Decompositions to XLA's FFI #21604
• Hessenberg Decomposition
• Tridiagonal Reduction

GPU

The GPU kernels all live in jaxlib/gpu:

BLAS:

Merged with solver below so that decisions about batching can be made in the kernel (details).

linalg:

• Cholesky update
• LU pivots to permutations
  • Ported LuPivotsToPermutation to the typed XLA FFI #20931
  • Update GPU implementation of lu_pivots_to_permutation to infer the permutation size directly from the input dimensions, instead of using an input parameter. #22911
  • Add support for shape polymorphism with lu_pivots_to_permutation. #22931

PRNG:

• threefry2x32
  • Ported threefry2x32 for GPU to the typed XLA FFI #20997

RNN:

• RNNForward
• RNNBackward

Solver:

• LU
  • Port Getrf GPU kernel from custom call to FFI. #22795
  • Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel. #23053
  • Enable FFI implementation of GPU Getrf FFI handler. #22796
  • Add support for using dynamic shapes with GPU LU decomposition. #22893
• QR
• Sparse QR
• Householder
• Eigenvalue decomposition
• SVD
• Tridiagonal reduction

Sparse:

...

Triton:

...