[论文解读] Accelerated Charged Particle Tracking with Graph Neural Networks on FPGAs

Aneesh Heintz, Vesal Razavimaleki+19 more|arXiv (Cornell University)|Nov 30, 2020

Parallel Computing and Optimization Techniques参考文献 30被引用 31

一句话总结

本文提出两种 FPGA 实现（OpenCL 协同处理和基于 hls4ml 的）用于带电粒子跟踪的图神经网络，取得相对于 CPU 的显著加速，并有在 LHC 的 FPGA 基 L1 触发中潜在应用。

ABSTRACT

We develop and study FPGA implementations of algorithms for charged particle tracking based on graph neural networks. The two complementary FPGA designs are based on OpenCL, a framework for writing programs that execute across heterogeneous platforms, and hls4ml, a high-level-synthesis-based compiler for neural network to firmware conversion. We evaluate and compare the resource usage, latency, and tracking performance of our implementations based on a benchmark dataset. We find a considerable speedup over CPU-based execution is possible, potentially enabling such algorithms to be used effectively in future computing workflows and the FPGA-based Level-1 trigger at the CERN Large Hadron Collider.

研究动机与目标

• 通过异构硬件满足严格的延迟和数据吞吐需求，推动高能物理中的加速跟踪。
• 在 FPGA 上改编并实现用于分段分类的图神经网络。
• 在基准数据 TrackML 上评估资源使用、延迟和物理性能。
• 展示在 LHC 在线触发工作流程中集成 FPGA 基跟踪的潜力。

提出的方法

• 两种针对 FPGA 的 GNN 实现，针对事件-互作网络 IN 模型用于在图嵌入的探测器击中上进行分段分类。
• OpenCL 实现采用 CPU-FPGA 协同处理，利用 FPGA 加速的矩阵乘法和统一尺寸的输入图填充。
• hls4ml 实现将神经网络翻译为 FPGA 固件，带流水线、流输入和可配置的重用因子以控制延迟和并行度。
• 边块和节点块由带 ReLU 激活的小型多层感知机组成，并有一个用于边分类的 Sigmoid 输出。
• 输入包括节点特征（r、phi、z）和边特征（Delta r、Delta phi、Delta z、Delta R）用于一个模型；另一种变体仅使用基本的边特征。
• 性能指标包括资源使用、延迟，以及 ROC-AUC（AUC），作为比特精度和模型规模的函数。

实验结果

研究问题

• RQ1 是否可以在 FPGA 硬件上使用 OpenCL 和 hls4ml 高效实现基于 GNN 的用于跟踪重建的分段分类？
• RQ2 在比较 OpenCL 协同处理与 hls4ml FPGA 实现时，资源、延迟与物理性能之间存在哪些权衡？
• RQ3 模型精度和重用因子如何影响 FPGA 实现的延迟和 ROC 性能？
• RQ4 在类似 TrackML 的数据集上，这些 FPGA 实现相对于 CPU 推断潜在的加速有多大？
• RQ5 在具有亚微秒要求的 LHC Level-1 触发系统中，这些 FPGA 方法是否可行？

主要发现

• OpenCL FPGA 实现实现全事件图的延迟在 10 ms 到 1 s 范围内，包括数据传输和输入/输出。
• hls4ml 实现目标是超低延迟，对于较小、分区的图， FPGA 执行延迟约为 650 ns 至 1 μs。
• 同一模型的 CPU 推断显著较慢，例如 graph_nets TensorFlow 实现的 pT>2 GeV 图约 27 ms，PyTorch 的 pT>1 GeV 约 86 ms，说明 FPGA 实现具显著加速。
• OpenCL 资源使用随数据精度降低（8、16、32 位）而下降，延迟随最小 pT 和事件大小变化，显示在协同处理下对数据规模的灵活处理。
• hls4ml 模型在固定点表示约 12 位总比特下 reproduces 完整 FP32 性能，延迟在 650 ns–1 μs 范围；更高的重用因子增加延迟但降低资源使用。
• 相较于仅 CPU 的工作流，FPGA 方法提供显著加速，正在进行优化资源使用和进一步降低延迟，以提升 OpenCL 基工作流。