Improving performance using sparse linear algebra
Graph Neural Networks (GNNs) rely on message passing, traditionally implemented via a gather/scatter mechanism. In GraphNeuralNetworks.jl, this approach materializes dense feature matrices on graph edges, leading to severe memory bottlenecks and reduced performance as graph sizes scale. This project addresses this limitation by redefining message passing as non-materializing Sparse Matrix-Matrix Multiplications (SpMM). By treating the graph as a sparse mathematical operator, this project will unlock significant computational efficiency and memory savings on CUDA GPUs. The core objectives include integrating COO storage, upstreaming patches to CUDA.jl, and refactoring the GNNlib API. To ensure continuous momentum during upstream review cycles, the project features a secondary track focused on upgrading the Automatic Differentiation (AD) pipeline by fully integrating Mooncake.jl and Enzyme.jl, alongside expanding the ecosystem with new datasets and new tutorials / improved documentation.
Project details
Technologies
Not listed in the archive