codon/docs/advanced/gpu.md

Codon supports GPU programming through a native GPU backend. Currently, only Nvidia devices are supported. Here is a simple example:

import gpu

@gpu.kernel
def hello(a, b, c):
    i = gpu.thread.x
    c[i] = a[i] + b[i]

a = [i for i in range(16)]
b = [2*i for i in range(16)]
c = [0 for _ in range(16)]

hello(a, b, c, grid=1, block=16)
print(c)

which outputs:

[0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45]

The same code can be written using Codon's @par syntax:

a = [i for i in range(16)]
b = [2*i for i in range(16)]
c = [0 for _ in range(16)]

@par(gpu=True)
for i in range(16):
    c[i] = a[i] + b[i]

print(c)

Below is a more comprehensive example for computing the Mandelbrot set, and plotting it using NumPy/Matplotlib:

from python import numpy as np
from python import matplotlib.pyplot as plt
import gpu

MAX    = 1000  # maximum Mandelbrot iterations
N      = 4096  # width and height of image
pixels = [0 for _ in range(N * N)]

def scale(x, a, b):
    return a + (x/N)*(b - a)

@gpu.kernel
def mandelbrot(pixels):
    idx = (gpu.block.x * gpu.block.dim.x) + gpu.thread.x
    i, j = divmod(idx, N)
    c = complex(scale(j, -2.00, 0.47), scale(i, -1.12, 1.12))
    z = 0j
    iteration = 0

    while abs(z) <= 2 and iteration < MAX:
        z = z**2 + c
        iteration += 1

    pixels[idx] = int(255 * iteration/MAX)

mandelbrot(pixels, grid=(N*N)//1024, block=1024)
plt.imshow(np.array(pixels).reshape(N, N))
plt.show()

The GPU version of the Mandelbrot code is about 450 times faster than an equivalent CPU version.

GPU kernels are marked with the @gpu.kernel annotation, and compiled specially in Codon's backend. Kernel functions can use the vast majority of features supported in Codon, with a couple notable exceptions:

• Exception handling is not supported inside the kernel, meaning kernel code should not throw or catch exceptions. raise statements inside the kernel are marked as unreachable and optimized out.

• Functionality related to I/O is not supported (e.g. you can't open a file in the kernel).

• A few other modules and functions are not allowed, such as the re module (which uses an external regex library) or the os module.

{% hint style="warning" %} The GPU module is under active development. APIs and semantics might change between Codon releases. {% endhint %}

Invoking the kernel

The kernel can be invoked via a simple call with added grid and block parameters. These parameters define the grid and block dimensions, respectively. Recall that GPU execution involves a grid of (X x Y x Z) blocks where each block contains (x x y x z) executing threads. Device-specific restrictions on grid and block sizes apply.

The grid and block parameters can be one of:

• Single integer x, giving dimensions (x, 1, 1)
• Tuple of two integers (x, y), giving dimensions (x, y, 1)
• Tuple of three integers (x, y, z), giving dimensions (x, y, z)
• Instance of gpu.Dim3 as in Dim3(x, y, z), specifying the three dimensions

GPU intrinsics

Codon's GPU module provides many of the same intrinsics that CUDA does:

Codon	Description	CUDA equivalent
gpu.thread.x	x-coordinate of current thread in block	threadId.x
gpu.block.x	x-coordinate of current block in grid	blockIdx.x
gpu.block.dim.x	x-dimension of block	blockDim.x
gpu.grid.dim.x	x-dimension of grid	gridDim.x

The same applies for the y and z coordinates. The *.dim objects are instances of gpu.Dim3.

Math functions

All the functions in the math module are supported in kernel functions, and are automatically replaced with GPU-optimized versions:

import math
import gpu

@gpu.kernel
def hello(x):
    i = gpu.thread.x
    x[i] = math.sqrt(x[i])  # uses __nv_sqrt from libdevice

x = [float(i) for i in range(10)]
hello(x, grid=1, block=10)
print(x)

gives:

[0, 1, 1.41421, 1.73205, 2, 2.23607, 2.44949, 2.64575, 2.82843, 3]

Libdevice

Codon uses libdevice for GPU-optimized math functions. The default libdevice path is /usr/local/cuda/nvvm/libdevice/libdevice.10.bc. An alternative path can be specified via the -libdevice compiler flag.

Working with raw pointers

By default, objects are converted entirely to their GPU counterparts, which have the same data layout as the original objects (although the Codon compiler might perform optimizations by swapping a CPU implementation of a data type with a GPU-optimized implementation that exposes the same API). This preserves all of Codon/Python's standard semantics within the kernel.

It is possible to use a kernel with raw pointers via gpu.raw, which corresponds to how the kernel would be written in C++/CUDA:

import gpu

@gpu.kernel
def hello(a, b, c):
    i = gpu.thread.x
    c[i] = a[i] + b[i]

a = [i for i in range(16)]
b = [2*i for i in range(16)]
c = [0 for _ in range(16)]

# call the kernel with three int-pointer arguments:
hello(gpu.raw(a), gpu.raw(b), gpu.raw(c), grid=1, block=16)
print(c)  # output same as first snippet's

gpu.raw can avoid an extra pointer indirection, but outputs a Codon Ptr object, meaning the corresponding kernel parameters will not have the full list API, instead having the more limited Ptr API (which primarily just supports indexing/assignment).

Object conversions

A hidden API is used to copy objects to and from the GPU device. This API consists of two new magic methods:

• __to_gpu__(self): Allocates the necessary GPU memory and copies the object self to the device.

• __from_gpu__(self, gpu_object): Copies the GPU memory of gpu_object (which is a value returned by __to_gpu__) back to the CPU object self.

For primitive types like int and float, __to_gpu__ simply returns self and __from_gpu__ does nothing. These methods are defined for all the built-in types and are automatically generated for user-defined classes, so most objects can be transferred back and forth from the GPU seamlessly. A user-defined class that makes use of raw pointers or other low-level constructs will have to define these methods for GPU use. Please refer to the gpu module for implementation examples.

@par(gpu=True)

Codon's @par syntax can be used to seamlessly parallelize existing loops on the GPU, without needing to explicitly write them as kernels. For loop nests, the collapse argument can be used to cover the entire iteration space on the GPU. For example, here is the Mandelbrot code above written using @par:

MAX    = 1000  # maximum Mandelbrot iterations
N      = 4096  # width and height of image
pixels = [0 for _ in range(N * N)]

def scale(x, a, b):
    return a + (x/N)*(b - a)

@par(gpu=True, collapse=2)
for i in range(N):
    for j in range(N):
        c = complex(scale(j, -2.00, 0.47), scale(i, -1.12, 1.12))
        z = 0j
        iteration = 0

        while abs(z) <= 2 and iteration < MAX:
            z = z**2 + c
            iteration += 1

        pixels[i*N + j] = int(255 * iteration/MAX)

Note that the gpu=True option disallows shared variables (i.e. assigning out-of-loop variables in the loop body) as well as reductions. The other GPU-specific restrictions described here apply as well.

Troubleshooting

CUDA errors resulting in kernel abortion are printed, and typically arise from invalid code in the kernel, either via using exceptions or using unsupported modules/objects.