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.
@article{arxiv.2202.06929,
title = {Accelerating the Inference of the Exa.TrkX Pipeline},
author = {Alina Lazar and Xiangyang Ju and Daniel Murnane and Paolo Calafiura and Steven Farrell and Yaoyuan Xu and Maria Spiropulu and Jean-Roch Vlimant and Giuseppe Cerati and Lindsey Gray and Thomas Klijnsma and Jim Kowalkowski and Markus Atkinson and Mark Neubauer and Gage DeZoort and Savannah Thais and Shih-Chieh Hsu and Adam Aurisano and V Hewes and Alexandra Ballow and Nirajan Acharya and Chun-yi Wang and Emma Liu and Alberto Lucas},
journal= {arXiv preprint arXiv:2202.06929},
year = {2023}
}
Comments
Proceedings submission to ACAT2021 Conference, 7 pages