Skip to content

readme: add notes on cuda versions #144

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 6, 2025
Merged

readme: add notes on cuda versions #144

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,9 +78,7 @@ The non-Python dependencies that you'll need are:

- [FFTW](https://www.fftw.org),
- [OpenMP](https://www.openmp.org) (for CPU, optional),
- CUDA (for GPU, >= 11.8), and

Older versions of CUDA may work, but they are untested.
- CUDA (for GPU, >= 11.8)

Below we provide some example workflows for installing the required dependencies:

Expand Down Expand Up @@ -133,6 +131,16 @@ export CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_ARCHITECTURES=60;70;80;90 -DJAX_FINU
```
</details>

#### Notes on CUDA versions
While jax-finufft may build with a wide range of CUDA
versions, the resulting binaries may not be compatible with JAX (resulting in
odd runtime errors, like failed cuDNN or cuBLAS initialization). For the greatest
chance of success, we recommend building with the same version as JAX was built with.
To discover that, one can look at the requirements in [JAX's `build` directory](https://github.com/jax-ml/jax/tree/main/build)
(be sure to select the git tag for your version of JAX). Similarly, we encourage installing
`jax[cuda12-local]` so JAX and jax-finufft use the same CUDA libraries.

Depending on how challenging the installation is, users might want to run jax-finufft in a container. The [`.devcontainer`](./.devcontainer) directory is a good starting point for this.

#### Configuring the build
There are several important CMake variables that control aspects of the jax-finufft and (cu)finufft builds. These include:
Expand Down Expand Up @@ -336,7 +344,7 @@ are currently only listed in source code form in
This package, developed by Dan Foreman-Mackey is licensed under the Apache
License, Version 2.0, with the following copyright:

Copyright 2021-2024 The Simons Foundation, Inc.
Copyright 2021-2025 The Simons Foundation, Inc.

If you use this software, please cite the primary references listed on the
[FINUFFT docs](https://finufft.readthedocs.io/en/latest/refs.html).
2 changes: 1 addition & 1 deletion src/jax_finufft/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Copyright 2021, 2022, 2023 The Simons Foundation, Inc.
# Copyright 2021-2025 The Simons Foundation, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
Expand Down