In recent years, Graph Neural Networks (GNNs) have ignited a surge of innovation, significantly enhancing the processing of geometric data structures such as graphs, point clouds, and meshes. As the domain continues to evolve, a series of frameworks and libraries are being developed to push GNN efficiency to new heights. While graph-centric libraries have achieved success in the past, the advent of efficient tensor compilers has highlighted the urgent need for tensor-centric libraries. Yet, efficient tensor-centric frameworks for GNNs remain scarce due to unique challenges and limitations encountered when implementing segment reduction in GNN contexts. We introduce GeoT, a cutting-edge tensor-centric library designed specifically for GNNs via efficient segment reduction. GeoT debuts innovative parallel algorithms that not only introduce new design principles but also expand the available design space. Importantly, GeoT is engineered for straightforward fusion within a computation graph, ensuring compatibility with contemporary tensor-centric machine learning frameworks and compilers. Setting a new performance benchmark, GeoT marks a considerable advancement by showcasing an average operator speedup of 1.80x and an end-to-end speedup of 1.68x.
翻译:近年来,图神经网络(GNN)引发了创新浪潮,显著提升了对图、点云和网格等几何数据结构的处理能力。随着该领域的持续发展,一系列框架和库被开发出来以推动GNN效率达到新高度。尽管以图为中心的库在过去取得了成功,但高效张量编译器的出现凸显了对张量中心库的迫切需求。然而,由于在GNN场景中实现分段归约时面临独特的挑战和限制,高效的张量中心GNN框架仍然稀缺。我们提出了GeoT,一个通过高效分段归约专门为GNN设计的尖端张量中心库。GeoT首次引入创新并行算法,不仅确立了新设计原则,还扩展了可用设计空间。重要的是,GeoT被设计为可轻松融合到计算图中,确保与当代张量中心机器学习框架和编译器的兼容性。通过展现平均算子加速1.80倍和端到端加速1.68倍,GeoT树立了新的性能标杆,标志着重大进展。