Other Abstraction Layers

OpenMP, SYCL, RAJA, JAX, PyTorch, Taichi, Julia, LLVM, autodiff

Lars Pastewka

Learning objectives

After this lecture you can

  • place Kokkos in the wider landscape of GPU programming models
  • recognize the main alternatives: OpenMP target, RAJA, JAX, PyTorch, Taichi, Julia
  • explain why LLVM is the common infrastructure behind most of them
  • pick an appropriate abstraction layer for a given project

One source, many backends

Kokkos does this – what else does?

The landscape

Approach Examples Strength
Directives OpenMP target Incremental adoption
C++ libraries Kokkos, RAJA, SYCL Performance control
Python frameworks JAX, PyTorch, Taichi Productivity, autodiff
Language-level Julia High-level + fast

OpenMP target offloading

#pragma omp target teams distribute parallel for \
    map(to: a[0:N], b[0:N]) map(from: c[0:N])
for (int i = 0; i < N; i++) {
    c[i] = a[i] + b[i];
}
  • Add pragmas to existing code \(\rightarrow\) GPU execution
  • map clauses control data movement
  • Supported by GCC 12+, Clang/LLVM, nvc++, amdclang++
  • Standardized but performance portability is harder than Kokkos

RAJA (LLNL)

RAJA::forall<RAJA::cuda_exec<256>>(
    RAJA::RangeSegment(0, N),
    [=] RAJA_DEVICE(int i) { c[i] = a[i] + b[i]; });
  • Same idea as Kokkos: execution policies as template parameters
  • Focuses on loop execution only – uses separate libraries for memory (CHAI, Umpire)
  • Widely used in LLNL production codes

JAX (Google)

import jax, jax.numpy as jnp

@jax.jit
def f(a, b):
    return a + b

a = jnp.ones(1_000_000)
c = f(a, a)  # Compiled via XLA, runs on GPU
  • NumPy-compatible API + JIT compilation via XLA
  • jax.grad(f) \(\rightarrow\) automatic differentiation
  • jax.vmap \(\rightarrow\) auto-vectorization, jax.pmap \(\rightarrow\) multi-GPU
  • Functional programming model (no side effects)

PyTorch (Meta)

import torch
x = torch.tensor([2.0], requires_grad=True, device='cuda')
y = x**2 + 3*x
y.backward()
print(x.grad)  # tensor([7.])
  • Dominant in deep learning, increasingly used in scientific computing
  • Eager execution by default + torch.compile for JIT optimization
  • Autograd engine for reverse-mode automatic differentiation
  • CUDA + ROCm backends

Taichi

@ti.kernel
def advance(f: ti.template()):
    for i, j in f:  # parallelized over the grid
        ...
  • Python DSL for spatial computing: particles, grids, meshes
  • @ti.kernel compiles via LLVM \(\rightarrow\) CUDA / Vulkan / Metal
  • Built-in sparse data structures and automatic differentiation

Julia

  • A language designed for numerical computing, compiled via LLVM
  • KernelAbstractions.jl = Kokkos-like portability (CUDA, ROCm, CPU)
  • Enzyme.jl = compiler-level AD through GPU kernels
  • Performance comparable to C/C++ with Python-like syntax

The role of LLVM

  • Common infrastructure behind Clang, Julia, Taichi, Enzyme
  • Same back ends generate GPU code for many front-end languages
  • AMD ROCm is entirely LLVM-based; NVIDIA moving toward LLVM too

Automatic differentiation

Forward mode: propagate derivatives alongside computation (\(n\) inputs \(\rightarrow\) \(n\) passes)

Reverse mode: record computation graph, traverse backward (\(m\) outputs \(\rightarrow\) \(m\) passes)

  • ML training: reverse mode = backpropagation (millions of params, one scalar loss)
  • Differentiable physics: differentiate through entire simulations
  • Enzyme: AD at the LLVM IR level \(\rightarrow\) works for any LLVM language, including GPU kernels

Comparison summary

Language AD GPU backends Compilation
Kokkos C++ No CUDA, HIP, SYCL AOT
OpenMP C/C++/Fortran No CUDA, HIP AOT
RAJA C++ No CUDA, HIP AOT
JAX Python Yes CUDA, ROCm, TPU JIT (XLA)
PyTorch Python Yes CUDA, ROCm Eager+JIT
Taichi Python Yes CUDA, Vulkan JIT (LLVM)
Julia Julia Yes CUDA, ROCm JIT (LLVM)

Which should you pick?

  • C++ codebase + portability \(\rightarrow\) Kokkos / RAJA
  • Python + autodiff \(\rightarrow\) JAX / PyTorch
  • Pragmas on legacy code \(\rightarrow\) OpenMP target
  • Vendor-locked maximum performance \(\rightarrow\) CUDA / HIP

Summary

  • Many roads to the GPU: directives, C++ template libraries, Python JIT frameworks, whole languages
  • Nearly all of them funnel through LLVM IR to reach multiple hardware back ends
  • Python frameworks add automatic differentiation – C++ layers generally do not
  • The right choice depends on your codebase, your need for autodiff, and how portable you must be

Reading