In October 2023, PyTorch 2.1 added support for compiling NumPy functions into C++ or CUDA code (via Triton), even adding parallelization with OpenMP when possible.

In the PyTorch blogpost, they demonstrated how they achieved up to 200x faster execution times for a simple kmeans one-step algorithm.

Having seen these results, and inspired by this post, I needed to try it and see whether I could achieve speedups of that magnitude with examples from my own code.

Trying it out in genetic algorithms

I’ve recently developed some interest in genetic algorithms as a form of heuristic search. They’re fun to implement, and the results they yield are impressive. However, I’ve noticed that, if implemented inefficiently, execution times may skyrocket, especially with complex fitness functions, large populations, and sophisticated chromosome structures. Naturally, this led me to wonder how much speedup could be obtained by compiling NumPy code with this new method.

Example #1

To start with something simple, I wanted to examine the speedups on my machine using a straightforward yet computationally intensive fitness function: the sum of an array.

Assuming we have a population of 900,000 chromosomes, each consisting of 1,000 genes, and the fitness function for each chromosome is the sum of all its genes, let’s evaluate the execution times using different implementations of this fitness function.

First, here’s the fitness function in pure Python:

population = [random.choices(range(1, 25), k=1000) for _ in range(N_POPULATION)]

def fitness_fn_python(individual):
	return sum(individual)

fitnesses = [fitness_fn_python(individual) for individual in population]

Executing this with a population of 900,000 chromosomes takes 6.91s, not including the time to generate the population. Now, let’s see how much we can accelerate this with numpy and torch.

This fitness function in NumPy is as follows:

population = np.random.randint(1, 25, size=(N_POPULATION, 1000), dtype=np.int8)

def fitness_fn_numpy(individual):
	return np.sum(individual)

fitnesses = np.apply_along_axis(fitness_fn_numpy, 1, population)

It executes in 5.19s which is 33% faster than the pure Python. This is still slow, because it’s not implemented as good as it can be.

A more efficient implementation using NumPy would be:

def fitness_fn_numpy_better(population):
	return np.sum(population, axis=1)

fitnesses = fitness_fn_numpy_better(population)

This version executes in 0.51s, which is 13.5 times faster than the pure Python implementation.

Now let’s compile this function with torch.compile, to see how much of a speedup we can get.

import torch
fitness_fn_compiled = torch.compile(fitness_fn_numpy_better)

fitnesses = fitness_fn_compiled(population)

The initial execution takes 0.59s due to the compilation overhead. However, on subsequent runs, it executes in 0.41s, making it 17 times faster than the pure Python version.

For even faster execution, we can write this function in PyTorch, compile it, and run it on cuda.

First, let’s rewrite the fitness function in PyTorch:

population = torch.randint(1, 25, size=(N_POPULATION, 1000), dtype=torch.int8)

def fitness_fn_torch(population):
	return torch.sum(population, dim=1)

fitnesses = fitness_fn_torch(population)

Uncompiled, this fitness function runs in 2.52s.

When compiled:

fitness_fn_compiled = torch.compile(fitness_fn_torch)
fitnesses = fitness_fn_compiled(population)

the first execution, including compilation time, takes 4.10s. But on subsequent executions, the time is reduced to 0.067s, which is 103 times faster than the pure Python implementation and 9 times faster than the uncompiled NumPy version.

The final performance boost comes from executing this PyTorch-compiled function on cuda:

population = population.to('cuda')
with torch.device('cuda'):
    fitnesses = fitness_fn_compiled(population)

This achieves an execution time of 0.055s, which is 125 times faster than pure Python and 10 times faster than NumPy, even on a modest laptop GPU.

In this example, maximum performance was achieved by writing and compiling the fitness function in PyTorch, yielding faster execution times than with PyTorch-compiled NumPy code.

Example #2

This is a more real-world example of a numerical optimization problem with constraints. These problems aim to find optimal values that minimize (or maximize) a given mathematical function subject to certain constraints.

Consider the following function:

$f(x, y, z) = x^2 + 2y^2 + 3z^2 + 4xy - 5yz$

with the constraints:

$$\begin{align*} & x \geq 1 \\ & 3 < y \leq 7 \\ & z < 0 \\ & x + y + z \geq 3 \\ & x^2 + z^2 \leq 10 \\ \end{align*}$$

In pure python, the fitness function is the following:

import random
import math

population = [[random.uniform(-25, 25) for _ in range(3)] for _ in range(N_POPULATION)]

def function_to_optimize(individual):
    x, y, z = individual
    return x**2 + 2*y**2 + 3*z**2 + 4*x*y - 5*y*z

def check_constraints(individual):
    x, y, z = individual
    if x < 1:
        return False
    if not (3 < y <= 7):
        return False
    if z >= 0:
        return False
    if (x + y + z) < 3:
        return False
    if (x**2 + z**2) > 10:
        return False
    return True

def fitness_fn_python(individual):
    if not check_constraints(individual):
        return math.inf
    return -function_to_optimize(individual)

fitnesses = [fitness_fn_python(individual) for individual in population]

Executing this for a population of 900,000 chromosomes takes 270ms.

Translating this function to vectorized NumPy:

import numpy as np

population = np.random.uniform(-25, 25, size=(N_POPULATION, 3))

def function_to_optimize(solutions):
    x = solutions[:, 0]
    y = solutions[:, 1]
    z = solutions[:, 2]
    return x**2 + 2*y**2 + 3*z**2 + 4*x*y - 5*y*z

# Vectorized function to check constraints
def check_constraints(solutions):
    x = solutions[:, 0]
    y = solutions[:, 1]
    z = solutions[:, 2]
    constraints = (x >= 1) & (y > 3) & (y <= 7) & (z < 0) & ((x + y + z) >= 3) & ((x**2 + z**2) <= 10)
    return constraints

# Vectorized fitness function
def fitness_function_numpy(solutions):
    # Check constraints for all solutions
    valid_constraints = check_constraints(solutions)
    # Apply constraints
    fitness = np.where(valid_constraints, -function_to_optimize(solutions), np.inf)
    return fitness

fitnesses = fitness_function_numpy(population)

The vectorized NumPy version takes 48ms, which is 5.6 times faster than the pure Python version.

If we compile this NumPy function with torch.compile:

import torch 
population = torch.rand(N_POPULATION, 3) * 50 - 25

check_constraints_compiled = torch.compile(check_constraints)
function_to_optimize_compiled = torch.compile(function_to_optimize)

def fitness_fn_pytorch(population):
    # Check constraints for all solutions
    valid_constraints = check_constraints_compiled(population)
    # Apply constraints
    fitness = np.where(valid_constraints, -function_to_optimize_compiled(population), np.inf)
    return fitness

fitnesses = fitness_fn_pytorch(population)

This compiled function takes 6.2ms to run, excluding compile overhead. This is 44 times faster than pure Python and 7 times faster than standard NumPy.

When executing on CUDA, the function takes 3.7ms to run, which is 73 times faster than pure Python and 13 times faster than NumPy.

Example #3

Similar to the second example, we’ll try to optimize this function:

$f(x, y, z) = \sin(x) + x^2 - 2\cos(y) + y^2 + z^2$

subject to these constraints:

$$\begin{align*} & x \geq 0 \\ & 0 \leq y \leq \pi \\ & z \leq 5 \\ & x + z \geq 1 \\ & \cos(y) + z \geq 0.5 \\ \end{align*}$$

I will not provide the code for this experiment as it’s very similar to the second example, but the execution times are the following:

• 240ms on pure Python.
• 82ms in NumPy, 3x faster than pure Python.
• 36.8ms with torch-compiled NumPy, 6.5x faster than pure Python and 2x faster than NumPy.
• 22.6ms with torch-compiled NumPy on CUDA, 10x faster than pure Python and 3.6x faster than NumPy.

Conclusion

This update to PyTorch makes it possible to get substantial speedups from numerical Python code without having to rewrite everything in another language.

In these examples, the largest gains came from code that was already vectorized. Compiling the NumPy version helped, and rewriting the function in PyTorch gave even better results. The exact speedup will depend on the workload, but the results show that this can be worth trying before moving performance-critical code to C++, Rust, or a custom CUDA implementation.

For machine learning, optimization, and data analysis workloads, being able to test this with only a few changes can save a significant amount of time. The code stays close to the original Python implementation, while the compiled version can take advantage of CPU parallelism or CUDA when the workload allows it.

Compilation overhead still matters, so it is most useful when the same function is executed enough times. But it gives Python programmers another option for speeding up numerical code without adding much complexity to the workflow.

Notes

Limitations in parallelization.

While this new compile function succeeds in detecting parallelizable pieces of code, at the time of writing it had been implemented for less than a month. Sometimes it fails to detect code segments that can be distributed across different cores, especially if they contain conditional logic not expressed in NumPy terms (see how the constraint check was done with np.where clauses).

It’s possible that this gets better in the future, but it’s not a bad idea to make the code easier for TorchInductor in order to get the best speedups possible.

Also, it’s a good idea to check the C++/triton code generated with TORCH_LOGS=output_code env variable.

Compile overhead matters.

When using this method, it’s useful to consider whether the compilation overhead is worth it for a given use case. These example functions have short compile times; however, with more complex tasks, the time it takes to compile could overshadow the time saved during execution.

In the PyTorch docs, we have this example:

model = init_model()
opt = torch.optim.Adam(model.parameters())

def train(mod, data):
    opt.zero_grad(True)
    pred = mod(data[0])
    loss = torch.nn.CrossEntropyLoss()(pred, data[1])
    loss.backward()
    opt.step()

eager_times = []
for i in range(N_ITERS):
    inp = generate_data(16)
    _, eager_time = timed(lambda: train(model, inp))
    eager_times.append(eager_time)
    print(f"eager train time {i}: {eager_time}")
print("~" * 10)

model = init_model()
opt = torch.optim.Adam(model.parameters())
train_opt = torch.compile(train, mode="reduce-overhead")

compile_times = []
for i in range(N_ITERS):
    inp = generate_data(16)
    _, compile_time = timed(lambda: train_opt(model, inp))
    compile_times.append(compile_time)
    print(f"compile train time {i}: {compile_time}")
print("~" * 10)

eager_med = np.median(eager_times)
compile_med = np.median(compile_times)
speedup = eager_med / compile_med
assert(speedup > 1)
print(f"(train) eager median: {eager_med}, compile median: {compile_med}, speedup: {speedup}x")
print("~" * 10)

In which the output is the following:

eager train time 0: 0.3666875
eager train time 1: 0.06788114929199218
eager train time 2: 0.06595875549316406
eager train time 3: 0.06623709106445312
eager train time 4: 0.06622108459472656
eager train time 5: 0.06640013122558594
eager train time 6: 0.06633692932128907
eager train time 7: 0.06608159637451172
eager train time 8: 0.06633171081542968
eager train time 9: 0.06623030090332031
~~~~~~~~~~
skipping cudagraphs due to input mutation
compile train time 0: 285.17234375
compile train time 1: 3.08613623046875
compile train time 2: 0.04830195236206054
compile train time 3: 0.03755414581298828
compile train time 4: 0.037427520751953124
compile train time 5: 0.03738982391357422
compile train time 6: 0.03739823913574219
compile train time 7: 0.03733385467529297
compile train time 8: 0.03738035202026367
compile train time 9: 0.03721017456054688
~~~~~~~~~~

Training times are 1.77 times faster, but the first iteration takes 790 times as long. So, compiling would only be useful if the number of training steps is high enough.

References

Making Python 100x faster with less than 100 lines of Rust - Ohad Ravid

torch.compile Tutorial - William Wen

Compiling NumPy code into C++ or CUDA via torch.compile - PyTorch blog