This page was generated from unit-5.6-gpu/commonGPU.ipynb.

5.6 Common GPU programming from Python¶

One kernel source, written in a small macro vocabulary, compiled at run-time for whichever backend is loaded: CUDA, Metal, or a host reference backend that is always available and compiles the kernel with the system C++ compiler.

Getting started
Reference: the kernel syntax
Using tinybla

Getting started¶

ngsolve.gpu picks up whichever backend this build has - Metal, CUDA, or the host reference device if neither is there - so GetGPUDevice() never returns None.

[1]:

import numpy as np
from ngsolve.gpu import *

dev = GetGPUDevice()
print(dev, " backend =", backend, " fp64 =", dev.has_float64,
      " max group =", dev.max_threads_per_group)

GPUDevice CPU reference  backend = host  fp64 = True  max group = 1024

/usr/lib/python3.12/dist-packages/ngsolve/gpu.py:62: UserWarning: ngsolve.ngscuda is present but could not be loaded: ngscuda was imported, but CUDA support is not enabled in this build of NGSolve.
  warnings.warn(f"{_module} is present but could not be loaded: {_e}")

A kernel. KERNEL declares the entry point, GLOBAL_IN/GLOBAL are device buffers, VALUE a scalar passed by value, and GLOBAL_ID_X is the work-item index. Prepend GPUKernelPrelude, which defines these for the current backend.

[2]:

src = '''
KERNEL(saxpy, GLOBAL_IN(float,x), GLOBAL(float,y), VALUE(float,a), VALUE(int,n))
{
  int i = GLOBAL_ID_X;
  if (i < n) y[i] = a*x[i] + y[i];
}
'''

lib = dev.CompileSource(GPUKernelPrelude + src)
q   = dev.DefaultQueue()

NewBuffer(n, dtype) allocates device memory for n elements of a numpy dtype, and NewBuffer(array) takes dtype, size and contents from an existing array. H2D uploads a numpy array (converted to the buffer’s dtype), D2H(n) reads n values back, all counted in elements. A Python float argument to a kernel arrives as float, an int as int; a numpy scalar such as np.float64(2.0) arrives with its own type. Launch checks every argument against the KERNEL(...) declaration - a float64 buffer in a float slot, a wrong scalar type or a missing argument raises with the kernel and argument name (see kernel.signature).

The launch geometry is as in CUDA - groups threadgroups of groupsize work items. Launches are asynchronous, Finish() waits.

[3]:

n, tg = 1000, 256
ngroups = (n+tg-1)//tg
hx = np.arange(n, dtype=np.float32)

x = dev.NewBuffer(hx)                          # dtype and contents from the array
y = dev.NewBuffer(n, np.float32)
y.H2D(np.ones(n))

q.Launch(lib.GetKernel("saxpy"), groups=[ngroups], groupsize=[tg], args=[x, y, 2.0, n])
q.Finish()

print("device :", y.D2H(5))
print("numpy  :", 2*hx[:5]+1)

device : [1. 3. 5. 7. 9.]
numpy  : [1. 3. 5. 7. 9.]

The point of the common syntax: the same source runs on the host reference backend, so a GPU result can be checked against a reference that shares no code with it.

[4]:

def run(device):
    lib = device.CompileSource(GPUKernelPrelude + src)
    q   = device.DefaultQueue()
    xb, yb = device.NewBuffer(n, np.float32), device.NewBuffer(n, np.float32)
    xb.H2D(hx)
    yb.H2D(np.ones(n, dtype=np.float32))
    q.Launch(lib.GetKernel("saxpy"), groups=[ngroups], groupsize=[tg], args=[xb, yb, 2.0, n])
    q.Finish()
    return yb.D2H(n)

on_device = run(dev)
on_host   = run(GetCPUDevice())
print("max difference :", np.abs(on_device-on_host).max())

max difference : 0.0

Reference: the kernel syntax¶

GPUKernelPrelude defines one vocabulary that expands to CUDA, Metal or plain C++.

Entry point and arguments. Arguments are positional: the i-th declared argument is the i-th entry of the args list of Launch.

`KERNEL(name, args...)`	declares an entry point
`GLOBAL(T,x)`	device buffer, read and write
`GLOBAL_IN(T,x)`	device buffer, read only
`GLOBAL_ATOMIC(T,x)`	device buffer written with `ATOMIC_ADD`
`VALUE(T,a)`	scalar passed by value

Indices. All exist as _X, _Y and _Z.

`GLOBAL_ID_X`	index of this work item in the whole grid
`LOCAL_ID_X`	index inside its threadgroup
`GROUP_ID_X`	index of the threadgroup
`GROUP_SIZE_X`	work items per threadgroup
`NUM_GROUPS_X`	number of threadgroups

Group-shared memory and synchronisation.

`SHARED(T,s,N)`	array shared by the threadgroup
`SHARED_2D(T,s,N,M)`	two-dimensional variant
`BARRIER()`	wait until the whole group arrives
`ATOMIC_ADD(p,val)`	atomic accumulate into a `GLOBAL_ATOMIC` buffer

Types for helper functions.

`GPU_FUNC`	qualifier for a function called from a kernel
`GLOBAL_PTR(T)`, `LOCAL_PTR(T)`	pointer into global / group-shared memory
`ATOMIC_PTR(T)`	pointer into a `GLOBAL_ATOMIC` buffer

Note that uint and float4 are not portable spellings - use unsigned, and index buffers element by element.

A kernel using group-shared memory and a barrier: each group reduces its slice of the input to one partial sum.

[5]:

reduce_src = '''
KERNEL(blocksum, GLOBAL_IN(float,x), GLOBAL(float,out), VALUE(int,n))
{
  SHARED(float, s, 256);
  int lid = LOCAL_ID_X;
  int gid = GLOBAL_ID_X;
  s[lid] = (gid < n) ? x[gid] : 0.0f;
  BARRIER();
  for (int stride = 128; stride > 0; stride >>= 1)
    { if (lid < stride) s[lid] += s[lid+stride]; BARRIER(); }
  if (lid == 0) out[GROUP_ID_X] = s[0];
}
'''

libr = dev.CompileSource(GPUKernelPrelude + reduce_src)
out  = dev.NewBuffer(ngroups, np.float32)
q.Launch(libr.GetKernel("blocksum"), groups=[ngroups], groupsize=[tg], args=[x, out, n])
q.Finish()

partial = out.D2H(ngroups)
print("partial sums :", partial)
print("device total :", partial.sum(), "  numpy :", hx.sum())

partial sums : [ 32640.  98176. 163712. 204972.]
device total : 499500.0   numpy : 499500.0

And one where many work items accumulate into the same locations, so the target is declared GLOBAL_ATOMIC and written with ATOMIC_ADD.

[6]:

hist_src = '''
KERNEL(histogram, GLOBAL_IN(float,x), GLOBAL_ATOMIC(float,bins),
                  VALUE(int,n), VALUE(int,nbins))
{
  int i = GLOBAL_ID_X;
  if (i < n) ATOMIC_ADD(&bins[i % nbins], x[i]);
}
'''

nbins = 4
libh  = dev.CompileSource(GPUKernelPrelude + hist_src)
bins  = dev.NewBuffer(nbins, np.float32)
bins.H2D(np.zeros(nbins, dtype=np.float32))

q.Launch(libh.GetKernel("histogram"), groups=[ngroups], groupsize=[tg],
         args=[x, bins, n, nbins])
q.Finish()

print("device :", bins.D2H(nbins))
print("numpy  :", np.array([hx[i::nbins].sum() for i in range(nbins)]))

device : [124500. 124750. 125000. 125250.]
numpy  : [124500. 124750. 125000. 125250.]

MemType.Shared is addressable from host and device, MemType.Device lives only on the accelerator. On unified memory shared costs nothing; on a discrete card a device-only buffer avoids keeping a host copy. D2H works for both; a shared buffer can also be viewed as a numpy array without a copy.

[7]:

for mt in [MemType.Shared, MemType.Device]:
    b = dev.NewBuffer(n, np.float32, mt)
    print(f"{str(mt):18s} host_visible = {b.host_visible}")

ib = dev.NewBuffer(n, dtype=np.int32)
ib.H2D(np.arange(n))
print(ib, " first entries :", np.asarray(ib)[:5])   # zero-copy view of a shared buffer

MemType.Shared     host_visible = True
MemType.Device     host_visible = True
GPUBuffer, size = 1000, dtype = int32  first entries : [0 1 2 3 4]

Complex numbers¶

Complex<float> and Complex<double> are part of the prelude, the same struct on every backend with the layout of std::complex, so numpy complex arrays transfer as they are. Arithmetic with + - * /, also mixed with a real of the same precision, conj, real, imag, Norm (\(|z|^2\)) and abs. A buffer accumulated atomically is declared GLOBAL_ATOMIC_COMPLEX(T,y) with the real type and written with ATOMIC_ADD_COMPLEX(y, i, val). On Apple GPUs only the float variant exists.

[8]:

cplx_src = '''
KERNEL(cplx_axpy, GLOBAL_IN(Complex<float>,x), GLOBAL(Complex<float>,y),
                  VALUE(Complex<float>,a), VALUE(int,n))
{
  int i = GLOBAL_ID_X;
  if (i < n) y[i] = a*x[i] + conj(y[i]);
}

KERNEL(cplx_dot, GLOBAL_IN(Complex<float>,x), GLOBAL_IN(Complex<float>,y),
                 GLOBAL_ATOMIC_COMPLEX(float,res), VALUE(int,n))
{
  int i = GLOBAL_ID_X;
  if (i < n) ATOMIC_ADD_COMPLEX(res, 0, conj(x[i])*y[i]);
}
'''
libc = dev.CompileSource(GPUKernelPrelude + cplx_src)
rng = np.random.default_rng(0)
zx = (rng.standard_normal(n) + 1j*rng.standard_normal(n)).astype(np.complex64)
zy = (rng.standard_normal(n) + 1j*rng.standard_normal(n)).astype(np.complex64)
a = np.complex64(0.3-0.7j)
bx, by, bres = dev.NewBuffer(zx), dev.NewBuffer(zy), dev.NewBuffer(np.zeros(1, np.complex64))
q.Launch(libc.GetKernel("cplx_dot"),  groups=[ngroups], groupsize=[tg], args=[bx, by, bres, n])
q.Launch(libc.GetKernel("cplx_axpy"), groups=[ngroups], groupsize=[tg], args=[bx, by, a, n])
q.Finish()
print("dot  :", bres.D2H(1)[0], "  numpy :", np.vdot(zx, zy))
print("axpy : max error", np.abs(by.D2H(n) - (a*zx + np.conj(zy))).max())

dot  : (-33.117584+0.507411j)   numpy : (-33.117584+0.50742435j)
axpy : max error 0.0

Using tinybla¶

tinybla is a small linear algebra library that lives inside the kernel: fixed size vectors and matrices in registers, views on buffers, and tile products across a threadgroup. Prepend TinyBlaPrelude after GPUKernelPrelude and open the namespace.

Per work item. Vec<S,T> and Mat<H,W,T> are held in registers, one per work item.

`Vec<S,T>`, `Mat<H,W,T>`	fixed size, `v(i)` and `m(i,j)`
`+`, `-`, `*`	with a scalar, a vector or a matrix
`Trans`, `Det`, `Cof`, `Inv`	`Det`/`Cof`/`Inv` for sizes 1, 2, 3
`ToVec`, `ToMat<H,W>`	reinterpret between the two
`v.Range<FIRST,NEXT>()`, `v.SetRange<FIRST,NEXT>(w)`	sub-vectors

Views on memory. BareMatrix is a pointer plus a leading dimension - no storage of its own - over a buffer or a SHARED_2D array.

`MakeBareMatrix<RowMajor>(ptr, ld)`	view with explicit leading dimension
`MakeBareMatrix<RowMajor>(shared2d)`	view of a `SHARED_2D` array
`m.SubMatrix(r,c)`	view starting at that entry
`m.Transpose()`	swaps `RowMajor` and `ColMajor`, no data moved
`m.ShiftRows(n)`, `m.ShiftCols(n)`	move the origin

A kernel with one small matrix per work item - invert a 3x3 and apply it, the shape of a geometry transformation in an FE kernel. Checked against the host backend and numpy.

[9]:

tb_src = '''
using namespace tinybla;

KERNEL(transform, GLOBAL_IN(float,jac), GLOBAL_IN(float,v), GLOBAL(float,out), VALUE(int,n))
{
  int i = GLOBAL_ID_X;
  if (i >= n) return;

  Mat<3,3,float> F;
  for (int r = 0; r < 3; r++)
    for (int c = 0; c < 3; c++)
      F(r,c) = jac[9*i + 3*r + c];

  Vec<3,float> u;
  for (int k = 0; k < 3; k++) u(k) = v[3*i+k];

  Vec<3,float> res = Inv(F) * u;
  for (int k = 0; k < 3; k++) out[3*i+k] = res(k);
  out[3*n+i] = Det(F);
}
'''

m = 64
rng = np.random.default_rng(0)
J = rng.random((m,3,3)).astype(np.float32) + np.eye(3, dtype=np.float32)
V = rng.random((m,3)).astype(np.float32)

def run_tb(device):
    lib = device.CompileSource(GPUKernelPrelude + TinyBlaPrelude + tb_src)
    qq  = device.DefaultQueue()
    bj, bv, bo = device.NewBuffer(J), device.NewBuffer(V), device.NewBuffer(4*m, np.float32)
    qq.Launch(lib.GetKernel("transform"), groups=[1], groupsize=[64], args=[bj, bv, bo, m])
    qq.Finish()
    r = bo.D2H(4*m)
    return r[:3*m].reshape(m,3), r[3*m:]

sol, det = run_tb(dev)
sol_host, det_host = run_tb(GetCPUDevice())

print("Inv(F)*u   vs host :", np.abs(sol-sol_host).max(),
      "  vs numpy :", np.abs(sol-np.linalg.solve(J, V[...,None])[...,0]).max())
print("Det(F)     vs host :", np.abs(det-det_host).max(),
      "  vs numpy :", np.abs(det-np.linalg.det(J)).max())

Inv(F)*u   vs host : 0.0   vs numpy : 1.7881393e-07
Det(F)     vs host : 0.0   vs numpy : 4.7683716e-07

Tiles across a threadgroup¶

WarpMatrix<H,W,T> is one tile held collectively by a simdgroup (32 lanes), which is how the matrix-free operators reach the tensor hardware. It is built from a BareMatrix and the work-item index, accumulates products with AddMM<K>, and is written back with Store.

`WarpMatrix<H,W,T> t = 0.0f;`	a zero tile, `H%8==0`, `W%4==0`
`WarpMatrix<H,W,T> t(baremat, tid);`	load a tile
`t.AddMM<K>(a, b, tid)`	`t += a*b`, with `a` H x K and `b` K x W
`t.Store(baremat, tid)`	write the tile back

This one needs a real GPU: the shuffle primitives it is built on are implemented for Metal only, so the host reference backend cannot compile it.

[10]:

warp_src = '''
using namespace tinybla;

KERNEL(tilemul, GLOBAL_IN(float,a), GLOBAL_IN(float,b), GLOBAL(float,c))
{
  unsigned tid = LOCAL_ID_X;
  auto ma = MakeBareMatrix<RowMajor>(a, 16u);   // 8 x 16
  auto mb = MakeBareMatrix<RowMajor>(b,  8u);   // 16 x 8
  auto mc = MakeBareMatrix<RowMajor>(c,  8u);   // 8 x 8

  WarpMatrix<8,8,float> sum = 0.0f;
  sum.AddMM<16>(ma, mb, tid);
  sum.Store(mc, tid);
}
'''

K = 16
A = rng.random((8,K)).astype(np.float32)
B = rng.random((K,8)).astype(np.float32)

try:
    libw = dev.CompileSource(GPUKernelPrelude + TinyBlaPrelude + warp_src)
    ba, bb, bc = dev.NewBuffer(8*K, np.float32), dev.NewBuffer(K*8, np.float32), dev.NewBuffer(64, np.float32)
    ba.H2D(A.reshape(-1))
    bb.H2D(B.reshape(-1))
    q.Launch(libw.GetKernel("tilemul"), groups=[1], groupsize=[32], args=[ba, bb, bc])
    q.Finish()
    print("max error vs numpy :", np.abs(bc.D2H(64).reshape(8,8) - A@B).max())
except Exception as e:
    print("WarpMatrix not available on this backend:", str(e).splitlines()[0])

max error vs numpy : 4.7683716e-07

The kernel sources used here are also available as a standalone script in ngstd/examples/gpukernel.py.