Computer Science Speakers Calendar

View Full Calendar

CS Compiler Seminar: Kum Wu, "PIGEON: Optimizing CUDA Code Generator for End-to-End Training and Inference of Relational Graph Neural Networks"

Event Type
Computer Science
2124 Siebel Center
wifi event
Mar 27, 2023   4:30 pm  

We look forward to seeing you in person on Monday, March 27, at 4:30pm. Join in person at 2124 Siebel Center for Computer Science, 201 N. Goodwin Ave or via zoom,

Title: PIGEON: Optimizing CUDA Code Generator for End-to-End Training and Inference of Relational Graph Neural Networks

Author(s): Kun Wu, Mert Hidayeto─člu, Xiang Song, Sitao Huang, Da Zheng, Israt Nisa, Wen-mei Hwu

Abstract: Relational graph neural networks (RGNNs) are graph neural networks (GNNs) with dedicated structures for modeling the different types of nodes and/or edges in heterogeneous graphs. While RGNNs have been increasingly adopted in many real-world applications due to their versatility and accuracy, they pose performance and system design challenges due to their inherent computation patterns, gap between the programming interface and kernel APIs, and heavy programming efforts in optimizing kernels caused by their coupling with data layout and heterogeneity. To systematically address these challenges, we propose Pigeon, a novel two-level intermediate representation (IR) and its code generator framework, that (a) represents the key properties of the RGNN models to bridge the gap between the programming interface and kernel APIs, (b) decouples model semantics, data layout, and operators-specific optimization from each other to reduce programming efforts, (c) expresses and leverages optimization opportunities in inter-operator transforms, data layout, and operator-specific schedules. By building on one general matrix multiply (GEMM) template and a node/edge traversal template, Pigeon achieves up to 7.8x speed-up in inference and 5.6x speed-up in training compared with the state-of-the-art public systems in select models, i.e., RGCN, RGAT, HGT, when running heterogeneous graphs provided by Deep Graph Library (DGL) and Open Graph Benchmark (OGB). Pigeon also triggers fewer out-of-memory (OOM) errors. In addition, we propose linear operator fusion and compact materialization to further accelerate the system by up to 2.2x. 

Bio: Kun Wu is currently a fourth-year Ph.D. student, advised by Prof. Wen-mei Hwu. His research interest lies in compilers and libraries for graphics processing units and parallel computer architecture. He has contributed to several impactful projects, involving PyTorch-Direct and Pylog.

link for robots only