Kernels

Download notebookView on GitHubRun when the site was built.

The compiled gather and scatter kernels behind Field: their versions, how one is chosen, and calling them directly.

Every kernel is a folder in mpifield/kernels/ with four versions taking the same arguments: Pyccel (compiled on first use, with OpenMP), NumPy (used without Pyccel), CUDA (for CuPy arrays) and Metal (Apple-silicon GPUs, in float32). The kernel objects come from cunumpy’s Kernel.from_folder, which picks the version from the arrays it is given.

import time

import cunumpy as xp
import numpy as np

from mpifield import LINEAR, BlockGrid, Field
from mpifield.kernels.catalog import KERNEL_NAMES, KERNELS, metal_kernel

for name in KERNEL_NAMES:
    kernel = KERNELS[name]
    print(f"{name:20s} {kernel.implementations}")

gather_linear_1d     ('pyccel', 'numpy', 'python', 'cuda')
gather_linear_2d     ('pyccel', 'numpy', 'python', 'cuda')
gather_linear_3d     ('pyccel', 'numpy', 'python', 'cuda')
scatter_linear_1d    ('pyccel', 'numpy', 'python', 'cuda')
scatter_linear_2d    ('pyccel', 'numpy', 'python', 'cuda')
scatter_linear_3d    ('pyccel', 'numpy', 'python', 'cuda')
gather_quadratic_1d  ('pyccel', 'numpy', 'python', 'cuda')
gather_quadratic_2d  ('pyccel', 'numpy', 'python', 'cuda')
gather_quadratic_3d  ('pyccel', 'numpy', 'python', 'cuda')
scatter_quadratic_1d ('pyccel', 'numpy', 'python', 'cuda')
scatter_quadratic_2d ('pyccel', 'numpy', 'python', 'cuda')
scatter_quadratic_3d ('pyccel', 'numpy', 'python', 'cuda')

Calling a kernel directly

The kernels take component-major markers, the grid with its components last (the linear 3D ones flattened to (nx * ny * nz, ncomp)) and the grid geometry as scalars: lower corner, node spacing, node counts and periodic flags. Field.gather and scatter work these out from the field’s stored block; a code that holds its own arrays, e.g. one patch of a refined mesh, can call the kernels itself, or a weighting’s gather_block and scatter_block with a BlockGrid for the array.

gather = KERNELS["gather_linear_2d"]
print(gather.host_parameters())

grid = np.random.default_rng(1).normal(size=(33, 17, 2))  # (nx, ny, ncomp)
points = np.random.default_rng(2).random((2, 1000)) * np.array([[1.0], [0.5]])
result = np.zeros((2, 1000))
gather(points, grid, result, 0.0, 0.0, 1 / 32, 0.5 / 16, 33, 17, 0, 0)
print(result[:, :3])

['positions', 'field', 'result', 'lower_x', 'lower_y', 'cell_size_x', 'cell_size_y', 'num_x', 'num_y', 'periodic_x', 'periodic_y']
[[ 1.18849616  0.38936531  0.57400037]
 [-0.64661871 -0.37098306 -1.44202414]]

block = BlockGrid.from_geometry((0.0, 0.0), (1 / 32, 0.5 / 16), (33, 17))
same = np.zeros((2, 1000))
LINEAR.gather_block(points, grid, same, block)
print(np.array_equal(same, result))

True

Choosing the host version

cunumpy.kernels.use_host_kernel_implementation switches between the compiled and the NumPy versions, e.g. to compare them; they agree to round-off. The compiled version is an order of magnitude faster, and parallel over the cores with OpenMP (OMP_NUM_THREADS).

field = Field(
    ndim=2,
    lower_bound=(0.0, 0.0),
    upper_bound=(1.0, 0.5),
    num_gridpoints=(33, 17),
    num_components=2,
)
field.set_field_data(grid)
markers = np.random.default_rng(3).random((2, 500_000)) * np.array([[1.0], [0.5]])
for implementation in ("pyccel", "numpy"):
    with xp.kernels.use_host_kernel_implementation(implementation):
        field.gather(markers)  # compile or warm up
        start = time.perf_counter()
        values = field.gather(markers)
        print(
            f"{implementation:7s} {1e9 * (time.perf_counter() - start) / markers.shape[1]:6.1f} ns per marker"
        )

pyccel     4.3 ns per marker
numpy    113.9 ns per marker

GPUs

With CuPy arrays (the gpu extra) the same calls run the CUDA versions, on fields whose storage is on the GPU (cunumpy.set_backend("cupy")). On an Apple-silicon Mac, metal_kernel(name) returns the Metal version (the metal extra, MLX), which takes NumPy arrays and computes in float32.

try:
    metal_gather = metal_kernel("gather_linear_2d")
except RuntimeError as error:
    print(error)
else:
    on_gpu = np.zeros((2, 1000))
    metal_gather(points, grid, on_gpu, 0.0, 0.0, 1 / 32, 0.5 / 16, 33, 17, 0, 0)
    print("largest difference to the CPU:", np.abs(on_gpu - result).max())

Metal kernels need MLX and a Metal GPU (mpifield[metal])