Kokkos: Execution

Parallel loops, execution policies, hardware mapping

Lars Pastewka

Learning objectives

After this lecture, you can

explain why the memory access pattern must depend on the architecture, and why pragmas cannot fix this
describe what Kokkos abstracts: execution spaces and memory spaces
decompose any parallel loop into pattern, execution policy and body
write Kokkos parallel_for and parallel_reduce kernels with range and MD-range policies

The problem: CPUs and GPUs want different access patterns

Figure 1: CPU: one thread visits consecutive addresses in time (caching). GPU: consecutive threads visit consecutive addresses (coalescing).

total += data[element*vectorSize + i];   // CPU-fast: consecutive i per thread
total += data[i*numElements + element];  // GPU-fast: consecutive threads, same i

A CPU-friendly array layout drops GPU performance by >10× – and vice versa.

For performance, the memory access pattern must depend on the architecture.

What is Kokkos?

Figure 2: One C++ source, compiled through a backend of your choice, runs on any vendor’s hardware.

C++ hardware abstraction layer: single source → NVIDIA, AMD, Intel GPUs + CPUs
Abstracts execution spaces (where code runs) + memory spaces (where data lives)
Battle-tested: developed at 5 US National Labs, used in Trilinos, LAMMPS, 100+ projects

Why not just OpenMP / OpenACC?

#pragma omp target data map(...)
#pragma omp teams num_teams(...) num_threads(...) private(...)
#pragma omp distribute
for (int element = 0; element < numElements; ++element) {
    // ...
    #pragma omp parallel for
    for (int qp = 0; qp < numQPs; ++qp) { /* ... */ }
}

Portable code ≠ performance-portable code.

The pragmas compile everywhere, but each backend behaves differently – and pragmas can’t change the data layout, which is the bigger problem. Kokkos can (next lecture).

Patterns, Policies, Bodies

The Kokkos vocabulary for any parallel loop:

Kokkos::parallel_for( RangePolicy<>(0, N) , [=](const int i) { /* work */ } );

Token	Role	Examples
`parallel_for`	Pattern – structure of the computation	`for`, `reduce`, `scan`, task graph
`RangePolicy<>(0, N)`	Execution policy – how work is scheduled	range, MD-range, team, dynamic
`[=](const int i) {...}`	Body – the per-iteration work	your code

You change where code runs by changing the policy, not the body.

“Parallel for”: same loop, three ways

// Serial
for (int64_t i = 0; i < N; ++i) { /* body */ }

// OpenMP (CPU only)
#pragma omp parallel for
for (int64_t i = 0; i < N; ++i) { /* body */ }

// Kokkos (CPU and GPU)
Kokkos::parallel_for(N, [=] (const int64_t i) { /* body */ });

Same conceptual difficulty as OpenMP – the annotations just go in a different place. Unlike OpenMP, the same source runs on GPU.

Initialization

#include <Kokkos_Core.hpp>

int main(int argc, char* argv[]) {
    Kokkos::initialize(argc, argv);
    // ... your code ...
    Kokkos::finalize();
    return 0;
}

Runtime options: --kokkos-num-threads=4, --kokkos-devices=cuda

`parallel_for` with range

Kokkos::parallel_for("label", N,
    KOKKOS_LAMBDA(const int i) {
        result(i) = a(i) + b(i);
    });

Each iteration \(i\) = one GPU thread
Runtime chooses block size (typically 128 or 256)
KOKKOS_LAMBDA wraps a C++ lambda for portability

Multidimensional loops: `MDRangePolicy`

Kokkos::parallel_for("init_2d",
    Kokkos::MDRangePolicy({0, 0}, {Nx, Ny}),
    KOKKOS_LAMBDA(const int i, const int j) {
        u(i, j) = 0.1 * i * dx + 0.2 * j * dy;
    });

First and second arguments of MDRangePolicy give lower and upper iteration bounds.

Reductions

double total = 0.0;
Kokkos::parallel_reduce("sum", N,
    KOKKOS_LAMBDA(const int i, double& lsum) {
        lsum += data(i);
    }, total);

parallel_reduce handles thread-safe accumulation
Also supports Min, Max, custom reducers

Execution spaces

An execution space is the Kokkos name for “where this kernel runs” – pick one per backend you compiled for:

Kokkos space	Backend
`Kokkos::Serial`	Single CPU thread
`Kokkos::Threads`	CPU threads
`Kokkos::OpenMP`	OpenMP
`Kokkos::Cuda`	NVIDIA GPU
`Kokkos::HIP`	AMD GPU
`Kokkos::SYCL`	Intel GPU

Selected at compile time via CMake options.

Summary

Memory access patterns must match the architecture: CPUs reward sequential access per thread, GPUs reward sequential access across threads
Pragma models are portable but not performance-portable; Kokkos abstracts execution and memory
Every parallel loop = pattern (parallel_for, parallel_reduce, …) + policy (range, MD-range, team) + body (lambda)
KOKKOS_LAMBDA + Kokkos::initialize/finalize are all the boilerplate you need
The execution space (Serial, OpenMP, CUDA, HIP, SYCL) is chosen at compile time