by Evgeni Burovski, Ralf Gommers and Mario Lezcano

Quansight engineers have implemented support for tracing through NumPy code via torch.compile in PyTorch 2.1. This feature leverages PyTorch’s compiler to generate efficient fused vectorized code without having to modify your original NumPy code. Even more, it also allows for executing NumPy code on CUDA just by running it through torch.compile under torch.device("cuda")!

In this post, we go over how to use this feature and give a few tips and tricks to make the most out of it.

Compiling NumPy code into Parallel C++

We take as our running example one step in a K-Means algorithm. This piece of code is borrowed from this NumPy book

import numpy as np

def kmeans(X, means):
    return np.argmin(np.linalg.norm(X - means[:, None], axis=2), axis=0)

We create a synthetic dataset with 20M random 2-D points. We can see that, given that the means are chosen appropriately, the function returns the correct cluster for all of them

npts = 10_000_000
X = np.repeat([[5, 5], [10, 10]], [npts, npts], axis=0)
X = X + np.random.randn(*X.shape)  # 2 distinct "blobs"
means = np.array([[5, 5], [10, 10]])
np_pred = kmeans(X, means)

Benchmarking this function gives us a baseline of 1.26s on an AMD 3970X CPU.

Compiling this function is now as easy as wrapping it with torch.compile and executing it with the example inputs

import torch

compiled_fn = torch.compile(kmeans)
compiled_pred = compiled_fn(X, means)
assert np.allclose(np_pred, compiled_pred)

The compiled function yields a 9x speed-up when running it on 1 core. Even better, as opposed to NumPy, our generated code does take advantage of all the cores in a processor. As such, when we run it on 32 cores, we get a 57x speed-up. Note that PyTorch always uses all the available cores unless explicitly restricted, so this is the default behavior you get when using torch.compile.

We may inspect the generated C++ code by running the script with the environment variable TORCH_LOGS=output_code. When doing so, we can see that torch.compile was able to compile the broadcasting and the two reductions into just one for-loop, and parallelize it using OpenMP

extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr0) {
    #pragma omp parallel num_threads(32)
    #pragma omp for
    for(long i0=0L; i0<20000000L; i0+=1L) {
        auto tmp0 = in_ptr0[2L*i0];
        auto tmp1 = in_ptr1[0L];
        auto tmp5 = in_ptr0[1L + (2L*i0)];
        auto tmp6 = in_ptr1[1L];
        // Rest of the kernel omitted for brevity

Compiling NumPy code into CUDA

Compiling our code so that it runs on CUDA is as simple as setting the default device to be CUDA

with torch.device("cuda"):
    cuda_pred = compiled_fn(X, means)
assert np.allclose(np_pred, cuda_pred)

By inspecting the generated code via TORCH_LOGS=output_code, we see that, rather than generating CUDA code directly, torch.compile generates rather readable triton code

def triton_(in_ptr0, in_ptr1, out_ptr0, XBLOCK : tl.constexpr):
    xnumel = 20000000
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (2*x0), xmask)
    tmp1 = tl.load(in_ptr1 + (0))
    // Rest of the kernel omitted for brevity

Running this small snippet on an RTX 2060 gives an 8x speed-up over the original NumPy code. This is something, but it is not particularly impressive, given the speed-ups we have seen on CPU. Let’s have a look into how to squeeze the most out of our GPU via a couple minor changes.

float64 vs float32. Many GPUs, in particular consumer-grade ones, are rather sluggish when running operations on float64. For this reason, changing the data generation to float32, the original NumPy code just gets a bit faster, about a 9%, but our CUDA code gets 40% faster, yielding a 11x speed-up over the plain NumPy code.

torch.compile, by default, respects the NumPy semantics, and as such, it uses np.float64 as its default dtype for all its creation ops. As discussed, this can hinder performance, so it is possible to change this default by setting

from torch._dynamo import config
config.numpy_default_float = "float32"

CPU <> CUDA copies. An 11x speed-up is good, but it is not even close to the CPU numbers. This is caused by a small transformation that torch.compile does behind the scenes. The code above takes NumPy arrays and returns NumPy arrays. All of these arrays are on CPU, but the computations are performed on the GPU. This means that every time the function is called, torch.compile has to copy all these arrays from CPU to the GPU, and then copy the result back to CPU to preserve the original semantics. There is no native solution to this issue in NumPy, as NumPy does not have the notion of a device. That being said, we can work around it by creating a wrapper to this function so that it accepts PyTorch tensors and returns PyTorch tensors.

@torch.compile
def tensor_fn(X, means):
    X, means = X.numpy(), means.numpy()
    ret = kmeans(X, means)
    return torch.from_numpy(ret)

def cuda_fn(X, means):
    with torch.device("cuda"):
        return tensor_fn(X, means)

This function now takes tensors in CUDA memory and returns tensors in CUDA memory, but the function itself is written in NumPy! torch.compile uses the numpy() and the from_numpy() calls as hints, and optimizes them away, and internally it simply works with PyTorch tensors without moving the memory at all. When we keep the tensors in CUDA and perform the computations in float32, we see a 200x speed-up over the initial NumPy implementation on float32 arrays.

Mixing NumPy and PyTorch. In this example, we had to write a small adaptor to convert tensors to ndarrays and then back to tensors. In programs that mix PyTorch and NumPy converting a tensor into an ndarray is often implemented as x.detach().cpu().numpy(), or simply x.numpy(force=True). Since when running under torch.compile we can run NumPy code in CUDA, we can implement this conversion pattern as call to x.numpy(), as we did above. Doing so and running the resulting code under device("cuda") will generate efficient CUDA code from original NumPy calls without copying the data from CUDA to CPU at all. Note that the resulting code does not run without torch.compile. For it to run in eager mode one would need to rollback to x.numpy(force=True).

Further Speed-up tricks

General advice. The CUDA code we have shown is already quite efficient, but it is true that the running example is rather short. When dealing with larger programs, we may need to tweak parts of it to make it more efficient. A good place to start is the multiple tutorials and FAQs for torch.compile. This showcases a number of ways to inspect the tracing process, and how to identify problematic code that may cause slowdowns.

Advice when compiling NumPy code. NumPy, even if rather similar to PyTorch, is often used very differently. It is rather common to perform computations in NumPy and then do an if/else depending on values within the array, or perform operations in-place, perhaps via boolean masks. These constructions, while supported by torch.compile, hamper its performance. Changes like writing the code in a branchless way to avoid graph breaks, or avoiding in-place ops can go a long way.

To write fast NumPy code, it is best to avoid loops, but sometimes they are unavoidable. When tracing through a loop, torch.compile will try to fully unroll it. This is sometimes desirable, but sometimes it may not even be possible, like when we have a dynamic stopping condition, like in a while loop. In these cases, it may be best to just compile the body of the loop, perhaps a few iterations at a time (loop unrolling).

Debugging NumPy code. Debugging is rather tricky when a compiler is involved. To figure out whether an error you are hitting is a torch.compile error, or an error from the program, you can execute your NumPy program without torch.compile by replacing the NumPy import by import torch._numpy as np. This is should just be used for debugging purposes and is in no way a replacement for the PyTorch API, as it is much slower and, as a private API, may change without notice. See also this FAQ for other tricks.

Differences between NumPy and torch.compile NumPy

NumPy scalars. NumPy returns NumPy scalars in almost any case where PyTorch would return a 0-D tensor (e.g. from np.sum). Under torch.compile, NumPy scalars are treated as 0-D arrays. This is just fine in most cases. The only case when their behavior diverges is when NumPy scalars are implicitly used as Python scalars. For example,

>>> np.asarray(2) * [1, 2, 3]  # 0-D array is an array-like
array([2, 4, 6])
>>> u = np.int32(2)
>>> u * [1, 2, 3]              # scalar decays into a Python int
[1, 2, 3, 1, 2, 3]
>>> torch.compile(lambda: u * [1, 2, 3])()
array([2, 4, 6])               # acts as a 0-D array, not as a scalar ?!?!

If we compile the first two lines, we see that torch.compile treats u as a 0-D array. To recover the eager semantics, we just need to make the casting explicit

>>> torch.compile(lambda: int(u) * [1, 2, 3])()
[1, 2, 3, 1, 2, 3]

Type promotion and versioning. NumPy’s type promotion rules may be, at times, a bit surprising

>>> np.zeros(1, dtype=np.int8) + 127
array([127], dtype=int8)
>>> np.zeros(1, dtype=np.int8) + 128
array([128], dtype=int16)

NumPy 2.0 is changing these rules to follow others that are closer to those PyTorch. The relevant technical document is NEP 50. torch.compile went ahead and implemented NEP 50 rather than the about-to-be-deprecated rules.

In general, NumPy within torch.compile follows NumPy 2.0 pre-release.

Beyond NumPy: SciPy and scikit-learn

In parallel to this effort of making torch.compile understand NumPy code, other Quansight engineers have designed and proposed a way to support PyTorch tensors within scikit-learn and SciPy. This was received enthusiastically by other maintainers from these libraries, as it was shown that using PyTorch as a backend would often yield considerable speed-ups. Both projects have now merged initial support for PyTorch tensors across a number of APIs and submodules.

This sets the stepping stone to move towards a future where PyTorch tensors can be used within other libraries in the Python data ecosystem. Even more, this will enable running these other libraries on GPUs and even compiling code mixing these libraries and PyTorch, similar to what we have been discussed in this post.

If you want to learn more about this effort, how to use it, or how to help moving it forward, see this other blogpost.

Conclusion

PyTorch has committed since its inception to be a framework compatible with the rest of the Python ecosystem. Enabling compiling NumPy programs, and establishing the tools necessary to do the same for other prominent libraries are two more steps in this direction. Quansight and Meta continue working hand on hand, improving the compatibility between PyTorch and the rest of the ecosystem.

From Quansight, we would like to thank Mengwei, Voz, and Ed for their invaluable help in integrating our work with torch.compile. We would also like to thank Meta for funding this project as well as previous work on improving NumPy compatibility within PyTorch, and the project that led to supporting PyTorch within scikit-learn and SciPy. These are giant leaps towards consolidating PyTorch as the framework of choice within the open source Python data ecosystem.