Accelerating the Inference of the Exa.TrkX Pipeline

TL;DR

This paper presents a GPU-accelerated C++ implementation of the Exa.TrkX particle tracking pipeline using optimized libraries and ONNX Runtime, significantly improving inference speed and integration with existing software.

Contribution

It introduces a C++ implementation with GPU acceleration and optimized graph algorithms, enhancing speed and compatibility for particle tracking in high energy physics.

Findings

Reduced event latency compared to Python implementation

Efficient GPU-based graph building with CUDA

Memory usage optimized for large datasets

Abstract

Recently, graph neural networks (GNNs) have been successfully used for a variety of particle reconstruction problems in high energy physics, including particle tracking. The Exa.TrkX pipeline based on GNNs demonstrated promising performance in reconstructing particle tracks in dense environments. It includes five discrete steps: data encoding, graph building, edge filtering, GNN, and track labeling. All steps were written in Python and run on both GPUs and CPUs. In this work, we accelerate the Python implementation of the pipeline through customized and commercial GPU-enabled software libraries, and develop a C++ implementation for inferencing the pipeline. The implementation features an improved, CUDA-enabled fixed-radius nearest neighbor search for graph building and a weakly connected component graph algorithm for track labeling. GNNs and other trained deep learning models are converted to ONNX and inferenced via the ONNX Runtime C++ API. The complete C++ implementation of the pipeline allows integration with existing tracking software. We report the memory usage and average event latency tracking performance of our implementation applied to the TrackML benchmark dataset.