cuNumeric.jl : Multi-GPU Array Programamming

cuNumeric brings array programming to multiple GPUs. Built on kernels from NVIDIA’s cuPyNumeric and accelerated by Julia macro-metaprogramming, cuNumeric provides performance competitive with bespoke multi-GPU libraries.

cuNumeric enables users to write high-level code that leverages an entire GPU node without having to worry about MPI, CUDA or NCCL. Furthermore, cuNumeric composes well with external libraries such as OrdinaryDiffEq.jl, TensorOperations.jl and Krylov.jl. In general, any program which is designed around generic array operations should work. See our docs page for examples.

There are two key differences from Base Julia:

  • To maximize performance we provide the @accelerate macro which fuses operations and elides unneeded temporary arrays. (docs)
  • Not all array code maps onto distributed hardware well. For example, we bring the @allowscalar macro from CUDA.jl to selectively enable scalar indexing of arrays. We also add @allowpromotion to allow silent promotion to a wider type and @allowfetch so that scalars automatically become a Julia Number instead of staying on device. These macros can modify the behavior of external functions (i.e., in another library). (docs)

Quick Example:

using cuNumeric

integrand(x) = exp(-x^2)

@accelerate function monte_carlo(N, x_max)
    Ω = 2 * x_max
    raw_samples = cuNumeric.rand(N)
    samples = @. Ω * raw_samples - x_max
    return (Ω / N) * sum(integrand.(samples))
end

N = 1_000_000
x_max = 10.0f0
estimate = monte_carlo(N, x_max)

println("Monte-Carlo Estimate: $(fetch(estimate))")
Composability Benchmarks Selected Kernel Benchmarks

The composability benchmark shows the performance of conjugate gradient in Krylov.jl and the 2D heat diffusion equation using OrdinaryDiffEq.jl solvers. On one GPU (left panel) we match CUDA.jl for large problem sizes while surpassing Dagger.jl. As we add more GPUs, cuNumeric achieves great weak scaling.

This kernel benchmarks show the performance of cuNumeric on applications with varying communication patterns. In all cases except multi-grid (MG), cuNumeric scales well and achieves performance competitive with JACC while retaining an array programming style.

Future Plans

Composability with the wider Julia ecosystem is our primary focus. If any of your applications do not work please let us know by opening an issue! You can also send an email to david@krasow.dev.

@krasow is a co-developer of cuNumeric.

4 Likes