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# pin Julia's threads before the first launch:#   julia -t 128,0 -e 'using ThreadPinning; pinthreads(:numa); include("examples/numa_aware.jl")'using BenchmarkToolsusing 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] = xend"""  measure_membw(; kwargs...) -> membw, flopsEstimate the memory bandwidth (GB/s) by performing a time measurement of aSAXPY 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_rateendmeasure_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 on Julia's threads (see CPU). For NUMA locality:

1. Pin Julia's threads before the first launch with ThreadPinning.jl. pinthreads(:numa) spreads threads across NUMA domains; pinthreads(:cores) follows the machine's topology. Both pin only the default pool. Start Julia with -t N,0 or pass threadpool = :all to pin small launches on the launching thread too. Without pinning, threads can migrate away from the memory they first touched.
2. Choose the number of threads with julia -t N or the JULIA_KA_CPU_THREADS environment variable; see CPU.
3. Initialize data with a parallel kernel. Under the NUMA first-touch policy, a memory page belongs to the domain of the thread that first writes it. fill! places pages in the calling thread's domain. The example uses KernelAbstractions.allocate followed by an initialization kernel (init = :parallel). Launches with as many workgroups start each block of them on the same thread, as Threads.@threads :static does, so later kernels mostly work on the data their thread touched first. KernelAbstractions.zeros(backend, dtype, N) and KernelAbstractions.ones also use a kernel for arrays of bits types of at least 1 MiB; other arrays are filled on the calling thread.

Demonstration:

On a system with two AMD EPYC 7502 (2 × 32 cores, 2 NUMA domains), the example (N = 512_000_000, workgroups of 1024), started with -t T,0 for each thread count T, gets the following memory bandwidth, in GB/s:

A plain Threads.@threads :static loop with its threads pinned by ThreadPinning.jl gets 240 GB/s with 64 threads. Initializing serially puts all the memory in one domain. Without pinning, the threads that touch the memory first aren't reliably the ones that work on it later, and the numbers vary from run to run with where the operating system places the threads.

On a system with a single memory domain, pinning and the initialization don't change where memory is placed: on an AMD Ryzen 9 9950X with 8 or 16 threads, all configurations get 53–55 GB/s.