Amd gpu support added to awkward array for faster physics data analysis
Bridging the Vendor Gap: Enabling AMD GPU Support for Awkward Array via ROCm/HIP for the HL-LHC Era
Distributed, Parallel, and Cluster ComputingProgramming Languages
Summary
Data analysis for physics experiments at the HL-LHC needs to be much faster and use different types of graphics cards, not just Nvidia’s. The authors created a bridge that lets a popular data tool called Awkward Array work well on AMD GPUs, which are used in some supercomputers. They found simple copying of Nvidia-focused code ran much slower on AMD and fixed this by optimizing the code specifically for AMD’s hardware quirks. Their approach speeds up certain computations by up to 12 times compared to many CPU cores, without changing how users program with Awkward Array.
What this means in practice
- •For high energy physics data teams: Enable faster processing of HL-LHC physics datasets using AMD GPUs with optimized Awkward Array kernels that match CUDA performance.
- •For gpu software engineers: Provide a set of optimization patterns to adapt CUDA kernels for AMD hardware, improving performance on irregular, nested data structures.
Authors
Ianna Osborne, Maxym Naumchyk, Tai Sakuma, Andres Rios-Tascon, Peter Elmer
Abstract
The High-Luminosity LHC (HL-LHC) will demand order-of-magnitude gains in analysis throughput, and increasingly those gains must come from GPUs that are not made by a single vendor. Leadership-class systems such as El Capitan, Frontier and LUMI are built on AMD accelerators, yet the Scikit-HEP analysis stack---and Awkward Array in particular---has grown up CUDA-first. We report on $rawkward$, a Rust-backed kernel engine that adds a ROCm/HIP backend for Awkward Array's nested, jagged, variable-length data structures. Our central finding is that a naive source-level port of CUDA kernels to HIP loses $5$--$10\times$ in performance on irregular kernels, because AMD's $64$-lane wavefronts, higher register pressure and more expensive divergence behave fundamentally differently from NVIDIA's $32$-thread warps. We show that a small, reusable set of optimization patterns---loop flattening, $128$-bit vectorized loads, splitting fused kernels, and profile-guided launch configuration---recovers CUDA-class performance without changing the public API. A Rust macro-and-match dispatch layer keeps a single, backend-agnostic call site while emitting vendor-specific kernel strategies, and the type system enforces buffer-size and lifetime correctness at compile time. On a two-socket AMD Instinct MI210 node we measure GPU speedups from $1.03\times$ (bandwidth-bound $sum$) up to $12.5\times$ ($count$) over $128$ EPYC~7763 CPU cores, and the Rust CPU kernels match or beat on aggregate the incumbent C++ kernels (geometric-mean runtime ratio $0.37\times$ across twelve kernels). We argue that these patterns constitute a practical recipe for performance-portable, vendor-agnostic HEP analysis kernels.