Add JAX Support to NPBench and Implement JAX Benchmarks #31
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Define the JaxFramework class for implementing block_until_ready() calls after kernel computation for correctness of profiling, convert np.ndarray to jax Array in the copy_func()
Define the JaxFramework class for implementing block_until_ready() calls after kernel computation for correctness of profiling, convert np.ndarray to jax Array in the copy_func()
Previously >10x slower, now around same runtime as numpy.
Previously was up to 70x slower, now it's faster for smaller sizes and comparable to numpy for bigger ones.
Previously was up to 90x slower, now it's up to 40x faster than numpy.
Previously was up to 3x slower, now comparable to numpy.
Add JAX as a supported framework
alexnick83
reviewed
Feb 3, 2025
| @@ -0,0 +1,106 @@ | |||
| # Copyright 2021 ETH Zurich and the NPBench authors. All rights reserved. | |||
| import pathlib | |||
| import jax.numpy as jnp | |||
There was a problem hiding this comment.
You have to put this in a try/except block, otherwise JAX becomes a mandatory dependency.
tbennun
requested changes
Feb 3, 2025
npbench/benchmarks/spmv/spmv_jax.py
Outdated
Comment on lines
7
to
12
| def spmv(A_row, A_col, A_val, x): | ||
| dim = A_row.size - 1 # needed because for the "paper" test size, scipy auto-infers the dims wrong | ||
| matrix_in_csr_format = scipy.sparse.csr_matrix((A_val, A_col, A_row), shape=(dim, dim)) | ||
| matrix_in_bcoo_format = jax_sparse.BCOO.from_scipy_sparse(matrix_in_csr_format) | ||
|
|
||
| return matrix_in_bcoo_format @ x |
There was a problem hiding this comment.
Could you please clarify this implementation?
- I am not sure the use of scipy or library calls would be allowed here without the
_lib.pysuffix (like you added for other applications) - This implementation changes the sparsity representation, which is BCOO instead of CSR
There was a problem hiding this comment.
We observed a very significant speedup (40x compared to numpy) with the current implementation that uses JAX's internal BCOO format and hence used that as the final implementation.
But you are right, having it in the lib implementation is better for fair comparison. Done.
The BCOO implementation has been added to spmv_jax_lib.py and spmv_jax.py has been reverted to a previous version that mimicked the numpy implementation in JAX.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Overview:
This PR introduces JAX as a supported framework in NPBench, providing JAX implementations for all existing benchmarks. This enhancement allows NPBench to compare the performance of JAX against other frameworks, for all implemented benchmarks in NPBench.
Changes Introduced:
framework_info/jax.jsonto define framework-specific configuration settings for JAX.npbench/benchmarks/<bench_parent>/<bench_name>_jax.py.npbench/benchmarks/<bench_parent>/<bench_name>_jax_lib.pynpbench/infrastructure/__init__.pyto include JAX framework initialization.npbench/infrastructure/jax_framework.pyto handle:block_until_ready()andjnp.array()).jax_framework.pyrelated to incorrect naming of the current framework during benchmarking.Motivation:
The addition of JAX support enhances NPBench by:
Testing and Validation:
Contributors: