CuPy: NumPy and SciPy for GPU
github.com
github.com
However, the AMD-GPU compatibility for CuPy is quite an attractive feature.
Also: You can just mix match all those functions and tensors thanks to the __cuda_array_interface__.
Code written for CuPy looks similar to numpy but very different from Jax.
As a sidenote, it is funny how this gets released in 2024, and not in say 2014...
Also CuPy was first released in 2015, this post is just a reminder for people that such things exist.
In 2024, with AI you can do these kind of projects very fast.
https://data-apis.org/array-api/
So it's possible to write array API code that consumes arrays from any of those libraries and delegate computation to them without having to explicitly import any of them in your source code.
The only limitation for now is that PyTorch (and to some lower extent cupy as well) array API compliance is still incomplete and in practice one needs to go through this compatibility layer (hopefully temporarily):
The last time I think this happen at market-scale was early 3d accelerator APIs? Glide/opengl/directx. Which has been a minute! (To a lesser extent CPU vectorization extensions)
Curious how much of Nvidia's successful strategy was driven by people who were there during that period.
Powerful first mover flywheel: build high performing hardware that allows you to define an API -> people write useful software that targets your API, because you have the highest performance -> GOTO 10 (because now more software is standardized on your API, so you can build even more performant hardware to optimize its operations)
https://scikit-learn.org/stable/modules/array_api.html
Disclosure: I'm a CuPy maintainer.
You could checkout some of EuroCC's courses. That should get you up to speed. https://www.eurocc-access.eu/services/training/
Is it changing though? Not only do PCIe interfaces keep doubling in performance, but CPU-GPU memory coherence is a thing.
I guess it depends on your target: 8x H100s across a PCIe bridge is going to have quite different costs vs an APU (which have gotten to be quite powerful, not even mentioning MI300a)
Last I checked (a couple months ago) it wasn't quite there, but I totally agree in principle. I've not gotten it to work on my Radeons yet.
I know you can enable ROCm for other hardware as well, but it's not supported and quite hit or miss. I've had limited success with running stuff against ROCm on unsupported cards, mainly having issues with memory management IIRC.
I think that every library required to build cupy is available in the universe repositories, though I've never tried building it myself.
[1]: https://salsa.debian.org/rocm-team/community/team-project/-/...
I've recently had to implement a few kernels to lower the memory footprint and runtime of some pytorch function : it's been really nice because numba kernels have type hints support (as opposed to raw cupy kernels).
And I recently learned that CuPy has a JIT compiler now if you prefer Python syntax over C++. https://docs.cupy.dev/en/stable/user_guide/kernel.html#jit-k...
In Python? Perhaps. Generally? No, it isn't. Try: https://github.com/eyalroz/cuda-api-wrappers/
Full power of the CUDA APIs including all runtime compilation options etc.
(Yes, I wrote that...)
import cupy as cp
vector_add = cp.RawKernel("""
extern "C" __global__
void vector_add(const float *A, const float *B, float *C, int num_elements) {
int i = blockDim.x * blockIdx.x + threadIdx.x;
if (i < num_elements) { C[i] = A[i] + B[i]; }
}
""", "vector_add")
num_elements = 50_000
block_size = 256
# round up to next multiple of block_size
grid_size = (num_elements + block_size - 1) // block_size
a = cp.random.rand(num_elements, dtype=cp.float32)
b = cp.random.rand(num_elements, dtype=cp.float32)
c = cp.zeros(num_elements, dtype=cp.float32)
args = (a, b, c, num_elements)
print(f"[Vector addition of {num_elements} elements]")
print(f"CUDA kernel launch with {grid_size} blocks of {block_size} threads each")
vector_add((grid_size,), (block_size,), args)
incorrect = cp.abs(a + b - c) > 1e-5
if cp.any(incorrect):
print("Result verification failed at element", cp.argmax(incorrect))
print("Test PASSED")
print("SUCCESS")
Although I have to concede that the automatic grid size computation in cuda-api-wrappers is nice.
A few marketing tips for your README:
* Put a code example directly at the top. You want to present the selling points of your library to the reader as fast as possible. For reference, look at the CuPy README https://github.com/cupy/cupy?tab=readme-ov-file#cupy--numpy-... which immediately shows reader what it is good for. Your README starts with lots of text, but nobody reads text anymore these days. A link to examples is almost at the end, and then the examples are deeply nested.
* The first links in the README should link to your own library, for example to documentation or examples. You do not want to lead the reader away from your GitHub page.
* Add syntax highlighting with "cpp" after triple backticks:
```cpp
<code here>
```This difference between the libraries makes your program more terse; however, you lose control over where your buffers are, from where they're accessible, when they get copied around and how etc. You can't even tell - from looking at the program source - whether the buffers will be "managed memory" accessed and copied page-by-page, or rather a copy will be made from system memory to device-global memory.
So, in my book, it is not as easy to access and control CUDA with cuPy. But - it is easier for a user who "needs numpy for GPUs", and does not care about the nitty-gritty, to write their program and get things done. Your program demostrates both of these points.
I should mention that I wrote my library with the hope that others will use it to build higher-level-abstraction libraries and apps. One could use it to create a cuCpp library that would be very numpy-like but for C++, a parallel of NumCpp [2].
Thanks for the tips regarding the README, I'll fix it up.
----
[1] : cuda-api-wrappers does offer a couple of utility classes like a poor man's span for pre-C++17, and a span+unique_ptr combo - which is beyond wrapping CUDA's APIs, but still doesn't quite "do" thing.
Actually that's a bit of a lie, because of the economy of primary device context reference counts (it's quite annoying if you need to do it well and not leak resources) and the context stack. So, let's say it does as little as possible behind the scenes... :-(
Would love to get any feedback from the community.
Has a much larger community, a big push from Google Research, and unlike PFN's Chainer (of which CuPy is the computational base), is not semi-abandoned.
Kind of sad to see CuPy/Chainer eco-system die: not only did they pioneer the PyTorch programming model, but also stuck to Numpy API like JAX does (though the AD is layered on top in Chainer IIRC).
"This is a research project, not an official Google product. Expect bugs and sharp edges. Please help by trying it out, reporting bugs, and letting us know what you think!"
disclaimer in its readme. This is quite scary, especially coming from Google which is known to abandon projects out of the blue
Open projects on github often (at least superficially) require specific versions of Cuda Toolkit (and all the specialty nvidia packages e.g. cudann), Tensorflow, etc, and changing the default versions of these for each little project, or step in a processing chain, is ridiculous.
pyenv et al have really made local, project specific versions of python packages much easier to manage. But I haven't seen a similar type solution for cuda toolkit and associated packages, and the solutions I've encountered seem terribly hacky..but I'm sure though that this is a common issue, so what do people do?
If anyone reading this message has encountered a roadblock while installing CuPy, please reach out. I'd be glad to help you.
(and you can't just add another index for pip to look for if you want to use python build so it has to be explicitly linked to the right wheel, which absolutely sucks especially since you cannot get the CUDA version from pypi)
NB I’m using Miniforge.
pip install cupy-cuda12xJax also provides back-propagation wherever possible, so you can optimize.
Cool to see that it's still kicking!
- JAX Windows support is lacking
- CuPy is much closer to CUDA than JAX, so you can get better performance
- CuPy is generally more mature than JAX (fewer bugs)
- CuPy is more flexible thanks to cp.RawKernel
- (For those familiar with NumPy) CuPy is closer to NumPy than jax.numpy
But CuPy does not support automatic gradient computation, so if you do deep learning, use JAX instead. Or PyTorch, if you do not trust Google to maintain a project for a prolonged period of time https://killedbygoogle.com/
Just prepare the input on NumPy or CuPy, and then you can just feed it to NumPy APIs. NumPy functions will handle itself if the input is NumPy ndarray, or dispatch the execution to CuPy if the input is CuPy ndarray.
CuPy is an (almost) drop-in replacement for NumPy, so the following works surprisingly often:
if use_cpu:
import numpy as np
else:
import cupy as npThis is the problem with these kind of methods. It works, until it doesn't in an unknown way.
DL is major use case; is CuPy planning on adding auto gradient comp?
It’s very easy to replace some slow NumPy/SciPy calls with appropriate CuPy calls, with sometimes literally a 1000x performance boost from like 10min work. It’s also easy to write “hybrid code” where you can switch between NumPy and CuPy depending on what’s available.
This was basically the code needed:
import scipy.linalg as la
if cuda:
import cupy as cp
import cupy.linalg as cla
ε = cp.asnumpy(cla.eigvalsh(cp.asarray(H)))
else:
ε = la.eigvalsh(H)
The eigenvalues of this matrix is the answer to “what are the energies of each stable electron state in this system”. If you know how many electrons you have (they tend to fill the lowest energy states they can at zero temperature), and you know what temperature you have (which gives you the probability of each “excited” state being occupied), then you can say a lot about the system. For instance, you can say what physical state lowers the “free energy” of the electrons at a given temperature (which can be used to predict phase transitions and spin configurations), or what is the “density of states” (which can be used to predict electronic resistance). You can also obtain the system’s entropy from the eigenvalues alone.
There are however many cases where you might need eigenvectors too, since they usually provide all the spatial information about “where in your system is this stuff happening”. When I need the eigenvectors, CuPy is still hundreds of times faster on my hardware, but the gap is just not as extreme as it was for pure eigenvalue calculation in my benchmarks.
Jax, pytorch, vanilla TF, triton. They just don’t cut it
EDIT: Or rather, all the solvers under jax's `scipy.sparse.linalg` all support multiple right hand sides.
“array or tree of arrays” sounds very general, probably even better than an old fashioned 2D array.
Ahh, that's just Jax's concept of pytrees. It was something that they invented to make it easier (this is how I view it, not complete) to pass complex objects to function but still be able to easily consider them as a concatenated vector for AD etc.. E.g. a common pattern is to pass parameters `p` to a function and then internally break them into their physical interpretations, e.g. `mass = p[0]`, `velocity = p[1]`. Pytrees let you just use something like a dictionary `p = {'mass' = 1.0, 'velocity = 1.0'}`, which is a stylistically more natural structure to pass around, and then jax is structured to understand later when AD'ing or otherwise that you're doing so with respect to the 'leaves' of the tree, or the values of the mass and velocity.
Hopefully someone corrects me if I'm not right about this. I'm hardly 100% on Jax's vision on PyTrees.
As an aside, just a list of right hand sides `[b1, b2, ..., bm]` is valid.
https://docs.cupy.dev/en/stable/reference/generated/cupyx.sc...
or https://docs.cupy.dev/en/stable/reference/generated/cupyx.sc... if your linear system is sparse.
But whether that works well depends on the problem you are trying to solve.
But I was hoping to try out something like ILU+bicgstab on the GPU and the python-verse seems like it has the lowest barrier-to-entry for just playing around.
https://github.com/pyamg/pyamg
https://github.com/ddemidov/amgcl
But I only have to deal with positive definite systems, so YMMV.
I am not sure whether those libraries can deal with multiple right-hand sides, but most complexity is in the preconditioners anyway.