Metadata-Version: 2.4
Name: jaxmg
Version: 0.0.2
Summary: JAXMg provides a C++ interface between JAX and cuSolverMg, NVIDIA's multi-GPU linear solver. 
Author: Roeland Wiersema
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/therooler/jaxmg
Project-URL: Issues, https://github.com/therooler/jaxmg/issues
Project-URL: Docs, https://therooler.github.io/jaxmg/
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.6.1; sys_platform == "darwin"
Requires-Dist: jax[cuda]>=0.6.1; sys_platform == "linux"
Provides-Extra: docs
Requires-Dist: jax>=0.6.2; extra == "docs"
Requires-Dist: mkdocs==1.6.1; extra == "docs"
Requires-Dist: mkdocs-material==9.6.23; extra == "docs"
Requires-Dist: mkdocstrings==0.30.1; extra == "docs"
Requires-Dist: mkdocs-exclude==1.0.2; extra == "docs"
Requires-Dist: mkdocs-ipynb==0.1.1; extra == "docs"
Requires-Dist: hippogriffe==0.2.2; extra == "docs"
Provides-Extra: dev
Requires-Dist: matplotlib>=3.10.3; extra == "dev"
Requires-Dist: jax[cuda12]>=0.6.2; extra == "dev"
Requires-Dist: pytest==8.2.1; extra == "dev"
Requires-Dist: mkdocs==1.6.1; extra == "dev"
Requires-Dist: mkdocs-material==9.6.23; extra == "dev"
Requires-Dist: mkdocstrings==0.30.1; extra == "dev"
Requires-Dist: mkdocs-exclude==1.0.2; extra == "dev"
Requires-Dist: mkdocs-ipynb==0.1.1; extra == "dev"
Requires-Dist: hippogriffe==0.2.2; extra == "dev"
Dynamic: license-file

<p align="center">
    <img src="docs/_static/logo.png" alt="Jaxmg" width="300">
</p>

# JAXMg: A distributed linear solver in JAX with cuSolverMg

[![Docs](https://img.shields.io/badge/docs-site-blue?style=flat-square)](https://therooler.github.io/jaxmg/)
[![Releases](https://img.shields.io/github/v/release/therooler/jaxmg?style=flat-square)](https://github.com/therooler/jaxmg/releases)
[![Continuous integration](https://github.com/therooler/jaxmg/actions/workflows/ci-build.yaml/badge.svg)](https://github.com/therooler/jaxmg/actions/workflows/ci-build.yaml)

# JAXMg
JAXMg provides a C++ interface between [JAX](https://github.com/google/jax) and [cuSolverMg](https://docs.nvidia.com/cuda/cusolver/index.html#using-the-cuSolverMg-api), NVIDIA’s multi-GPU linear solver.  We provide a jittable API for the following routines.

- [cusolverMgPotrs](https://docs.nvidia.com/cuda/cusolver/index.html#cusolvermgpotrs-deprecated): Solves the system of linear equations: $Ax=b$ where $A$ is an $N\times N$ symmetric (Hermitian) positive-definite matrix via a Cholesky decomposition 
- [cusolverMgPotrs](https://docs.nvidia.com/cuda/cusolver/index.html#cusolvermgpotri-deprecated): Computes the inverse of an $N\times N$ symmetric (Hermitian) positive-definite matrix via a Cholesky decomposition.
- [cusolverMgPotrs](https://docs.nvidia.com/cuda/cusolver/index.html#cusolvermgsyevd-deprecated): Computes eigenvalues and eigenvectors of an $N\times N$ symmetric (Hermitian) matrix.

For more details, see the [API](api/potrs.md).

The provided binary is compiled with:

| Component | Version |
|---|---:|
| **GCC** | 11.5.0 |
| **CUDA** | 12.8.0 |
| **cuDNN** | 9.2.0.82-12 |

!!! Compatibility
    We require JAX>=0.6.0, since it ships with CUDA 12.x binaries, which this package relies on. No local version of CUDA is required.


## Installation

Clone the repository and install with:

```bash
pip install ".[cuda]"
```

This will install a GPU compatible version of JAX. 

To verify the installation (requires at least one GPU) run

```bash
pytest 
```
There are two types of tests:

1. SPMD tests: Single Process Multiple GPU tests.
3. MPMD: Multiple Processes Multiple GPU tests.

### cuSolverMp
As of CUDA 13, there is a new distributed linear algebra library called [cuSolverMp](https://docs.nvidia.com/cuda/cusolvermp/) with similar capabilities as cuSolverMg, that does support multi-node computations as well as >16 devices. Given the similarities in syntax, it should be straightforward to eventually switch to this API. This will require sharding data into a cyclic 2D form and handling the solver orchestration with MPI.

## Citations
(Citation details will be available soon.)

## Acknowledgements
I acknowledge support from the Flatiron Institute. The Flatiron Institute is a
division of the Simons Foundation.
