Examples

We provide a complete walkthrough of the JAXMg workflow, from launching a distributed JAX program to running a cuSOLVERMp routine.

Ready-to-run scripts are provided in the repository's examples directory.

1. Configure Distributed Execution

Start by initializing distributed JAX, assigning one GPU to each process, and constructing the two-dimensional device mesh used by the matrix.

Configure distributed execution

2. Choose a tile size \(T_A\)

Choose a cuSOLVERMp tile size that divides both dimensions of each local matrix shard. Choosing an appropriate tile size can substantially improve runtime and reduce memory use.

Choose a tile size

3. Run a solver

• Cholesky solve solves a symmetric or Hermitian positive-definite linear system with potrs.
• LU solve solves a general nonsingular linear system with lu_solve.
• Symmetric or Hermitian eigensolve computes eigenvalues and optional eigenvectors with syevd.