Running on a GPU

Everything in the earlier examples runs on the CPU by default. This example samples the same g g > t t~ g mapping from the first example on a CUDA device instead. The mapping itself does not change. Only the Context it runs against, and the array library used to hand it data, are different: on a GPU, arrays are PyTorch tensors rather than NumPy arrays, since MadSpace interfaces with whichever array library the caller uses.

This requires a MadSpace build with the CUDA backend enabled (--cuda to madspace/install.py; the pre-built wheels for Linux include it). The same pattern works for HIP, using hip_device in place of cuda_device.

Selecting the device

cuda_device raises if the CUDA backend was not built into the installation:

import torch
import madspace as ms

device = ms.cuda_device()
context = ms.Context(device=device, thread_count=32)

Building a runtime bound to the context

mapping.map_forward is a CPU-only convenience, so a specific device needs the underlying Function and FunctionRuntime explicitly:

masses = [0.0, 0.0, 173.0, 173.0, 0.0]  # g g -> t t~ g
mapping = ms.PhaseSpaceMapping(masses, 13000.0, mode="rambo")

forward = ms.FunctionRuntime(mapping.forward_function(), context)

Sampling with PyTorch

Calling the runtime with a CUDA tensor keeps the whole computation on the device. The outputs come back in the same order as map_forward uses, just as a plain tuple rather than a namedtuple:

r = torch.rand((100000, mapping.random_dim()), device="cuda", dtype=torch.float64)
momenta, x1, x2, det = forward(r)
print(momenta.shape, momenta.device)

torch.Size([100000, 5, 4]) cuda:0

momenta, x1, x2 and det are ordinary CUDA tensors. Move them to the host with momenta.cpu() once the computation that needs them, such as a matrix-element evaluation, is also done.