NUMA-aware SAXPY

This example demonstrates how to define and run a SAXPY kernel (single-precision Y[i] = a * X[i] + Y[i]) such that it runs efficiently on a system with multiple memory domains (NUMA) with the multithreaded CPU backend. (You likely will need to fine-tune the value of N on your system of interest if you care about the particular measurement.)

# EXCLUDE FROM TESTING
# Run with `POCL_AFFINITY=1` to pin the threads that run the kernels, e.g.
#   POCL_AFFINITY=1 julia -t 128 examples/numa_aware.jl
using BenchmarkTools
using KernelAbstractions

@kernel function saxpy_kernel(a, @Const(X), Y)
    I = @index(Global)
    @inbounds Y[I] = a * X[I] + Y[I]
end

@kernel function fill_kernel(A, x)
    I = @index(Global)
    @inbounds A[I] = x
end

"""
  measure_membw(; kwargs...) -> membw, flops

Estimate the memory bandwidth (GB/s) by performing a time measurement of a
SAXPY kernel. Returns the memory bandwidth (GB/s) and the compute (GFLOP/s).
"""
function measure_membw(
        backend = CPU(); verbose = true, N = 1024 * 500_000, dtype = Float32,
        init = :parallel,
    )
    bytes = 3 * sizeof(dtype) * N # num bytes transferred in SAXPY
    flops = 2 * N # num flops in SAXY
    workgroup_size = 1024

    a = dtype(3.1415)
    X = KernelAbstractions.allocate(backend, dtype, N)
    Y = KernelAbstractions.allocate(backend, dtype, N)
    if init == :serial
        # The calling thread touches all the memory first
        fill!(X, dtype(1))
        fill!(Y, dtype(2))
    else
        # The threads that run the kernels touch the memory first
        fill_kernel(backend, workgroup_size)(X, dtype(1), ndrange = size(X))
        fill_kernel(backend, workgroup_size)(Y, dtype(2), ndrange = size(Y))
        KernelAbstractions.synchronize(backend)
    end

    t = @belapsed begin
        kernel = saxpy_kernel($backend, $workgroup_size, $(size(Y)))
        kernel($a, $X, $Y, ndrange = $(size(Y)))
        KernelAbstractions.synchronize($backend)
    end evals = 2 samples = 10

    mem_rate = bytes * 1.0e-9 / t # GB/s
    flop_rate = flops * 1.0e-9 / t # GFLOP/s

    if verbose
        println("\tMemory Bandwidth (GB/s): ", round(mem_rate; digits = 2))
        println("\tCompute (GFLOP/s): ", round(flop_rate; digits = 2))
    end
    return mem_rate, flop_rate
end

measure_membw(CPU());

# On a system with multiple NUMA domains, this places all the memory in the domain of the calling thread
# measure_membw(CPU(); init=:serial);

Important remarks:

The CPU backend runs kernels with PoCL on PoCL's own threads, not on Julia's. That determines how to apply the usual advice for NUMA systems:

  1. Pin the threads that run the kernels. Set PoCL's POCL_AFFINITY=1 environment variable before the backend is first used, which pins each of PoCL's threads to a core. Tools that pin Julia's threads, like ThreadPinning.jl, don't affect the threads that run the kernels.
  2. Choose the number of threads with julia -t N or the JULIA_KA_CPU_THREADS environment variable; see CPU.
  3. Initialize your data in parallel(!), with a kernel. Under the "NUMA first-touch policy" a page of memory is placed in the memory domain of the thread that first writes to it. KernelAbstractions.zeros(backend, dtype, N) and fill! write from the calling thread, which places all the memory in that thread's domain. The example instead allocates with KernelAbstractions.allocate and writes the initial values with a kernel (init = :parallel), so that the threads that run the computational kernel touch the memory first.

Demonstration:

So far, the example has only been measured with KernelAbstractions 0.10 on a system with a single memory domain, where pinning and the initialization don't change where memory is placed. With 16 Julia threads on an AMD Ryzen 9 5950X (16 physical cores, 1 NUMA domain), one gets the following numbers (comments for demonstration purposes):

Memory Bandwidth (GB/s): 41.19 # POCL_AFFINITY=1, init = :parallel
Compute (GFLOP/s): 6.87

Memory Bandwidth (GB/s): 41.22 # POCL_AFFINITY=1, init = :serial
Compute (GFLOP/s): 6.87

Memory Bandwidth (GB/s): 41.24 # POCL_AFFINITY unset, init = :parallel
Compute (GFLOP/s): 6.87

Memory Bandwidth (GB/s): 41.16 # POCL_AFFINITY unset, init = :serial
Compute (GFLOP/s): 6.86

As expected for a single memory domain, the four configurations don't differ: the kernel is limited by the memory bandwidth of the system. How much pinning and parallel initialization gain on a system with multiple memory domains has not been measured with the PoCL-based CPU backend yet.