nv-legate / cunumeric

An Aspiring Drop-In Replacement for NumPy at Scale
https://docs.nvidia.com/cunumeric/24.06/
Apache License 2.0
605 stars 69 forks source link

cuNumeric

cuNumeric is a Legate library that aims to provide a distributed and accelerated drop-in replacement for the NumPy API on top of the Legion runtime. Using cuNumeric you do things like run the final example of the Python CFD course completely unmodified on 2048 A100 GPUs in a DGX SuperPOD and achieve good weak scaling.

drawing

cuNumeric works best for programs that have very large arrays of data that cannot fit in the memory of a single GPU or a single node and need to span multiple nodes and GPUs. While our implementation of the current NumPy API is still incomplete, programs that use unimplemented features will still work (assuming enough memory) by falling back to the canonical NumPy implementation.

If you have questions, please contact us at legate(at)nvidia.com.

Installation

cuNumeric is available from conda on the legate channel. Please make sure you have at least conda version 24.1 installed, then create a new environment containing cuNumeric:

conda create -n myenv -c conda-forge -c legate cunumeric

or install it into an existing environment:

conda install -c conda-forge -c legate cunumeric

Once installed, you can verify the installation by running one of the examples from the cuNumeric repository, for instance:

$ legate examples/black_scholes.py
Running black scholes on 10K options...
Elapsed Time: 129.017 ms

Only linux-64 packages are available at the moment.

The default package contains GPU support, and is compatible with CUDA >= 11.8 (driver >= 520), and Volta or later GPU architectures. There are also CPU-only packages available, which will be automatically selected when installing on a machine without GPUs available. See https://nv-legate.github.io/cunumeric for details about manually forcing different install configurations, or building cuNumeric from source.

Usage and Execution

Using cuNumeric as a replacement for NumPy is easy. Users only need to replace:

import numpy as np

with:

import cunumeric as np

These programs can then be run by the Legate driver script described in the Legate Core documentation.

legate cunumeric_program.py

For execution with multiple nodes (assuming Legate Core is installed with networking support) users can supply the --nodes option. For execution with GPUs, users can use the --gpus flags to specify the number of GPUs to use per node. We encourage all users to familiarize themselves with these resource flags as described in the Legate Core documentation or simply by passing --help to the legate driver script.

You can use test.py to run the test suite. Invoke the script directly or through standard python; the script will invoke the legate driver script internally. Check out test.py --help for further options.

Supported and Planned Features

cuNumeric is currently a work in progress and we are gradually adding support for additional NumPy operators. Unsupported NumPy operations will provide a warning that we are falling back to canonical NumPy. Please report unimplemented features that are necessary for attaining good performance so that we can triage them and prioritize implementation appropriately. The more users that report an unimplemented feature, the more we will prioritize it. Please include a pointer to your code if possible too so we can see how you are using the feature in context.

Supported Types and Dimensions

cuNumeric currently supports the following NumPy types: float16, float32, float64, int16, int32, int64, uint16, uint32, uint64, bool, complex64, and complex128.

cuNumeric supports up to 4D arrays by default, you can adjust this setting by installing legate.core with a larger --max-dim.

Documentation

The cuNumeric documentation can be found here.

Future Directions

There are three primary directions that we plan to investigate with cuNumeric going forward:

We are open to comments, suggestions, and ideas.

Contributing

See the discussion of contributing in CONTRIBUTING.md.

Known Issues