Accelerating the Inference of the Exa.TrkX Pipeline

Alina Lazar, Xiangyang Ju, Daniel Murnane, Paolo Calafiura, Steven Farrell, Yaoyuan Xu, Maria Spiropulu, Jean Roch Vlimant, Giuseppe Cerati, Lindsey Gray, Thomas Klijnsma, Jim Kowalkowski, Markus Atkinson, Mark Neubauer, Gage DeZoort, Savannah Thais, Shih Chieh Hsu, Adam Aurisano, Jeremy Hewes, Alexandra BallowNirajan Acharya, Chun Yi Wang, Emma Liu, Alberto Lucas

Research output: Contribution to journalConference articlepeer-review

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.

Original languageEnglish (US)
Article number012008
JournalJournal of Physics: Conference Series
Volume2438
Issue number1
DOIs
StatePublished - 2023
Event20th International Workshop on Advanced Computing and Analysis Techniques in Physics Research, ACAT 2021 - Daejeon, Virtual, Korea, Republic of
Duration: Nov 29 2021Dec 3 2021

ASJC Scopus subject areas

  • General Physics and Astronomy

Fingerprint

Dive into the research topics of 'Accelerating the Inference of the Exa.TrkX Pipeline'. Together they form a unique fingerprint.

Cite this