GNN框架优化GNNAdvisor: An Adaptive and Efficient Runtime System for GNN Acceleration on GPUs

OSDI’21

Abstract

GNNAdvisor

• 从GNN模型和输入图中寻找与性能相关的特征，作为优化点
• 实现为GNN计算定制的2D工作负载管理，提升GPU利用率和性能
• 利用GPU内存结构，根据GPU内存结构和GNN工作负载的特点协调GNN的执行
• 为了实现自动优化，继承了一个轻量级的分析模型进行参数搜索

1. Introduction

目前，支持GNN训练和推理的工作可以分为两类：

• 图处理系统，融合神经操作
• 神经网络框架，扩展支持基于向量的图操作

但都是初步的、不完善的，有三个方面的缺陷：

• 没有利用GNN的输入信息：GNN模型有不同的层序列、聚合函数的类型、特征向量的维度，输入的图数据也有不同特征。现在的GNN框架（pyg, neugraph, dgl），采用的是一刀切的优化模式，无法最大化优化特定GNN应用
• 没有针对GNN的优化：更新阶段的NN操作是密集的、访存规律的，但聚合阶段的图操作是稀疏的、访存不规律的。当前的GNN框架的图操作来自图处理系统，针对的是标量属性，而GNN中每个节点都有一个embedding，因此应在embedding上积极利用并行。
• 对于不同输入，缺乏适应性的运行时支持：之前的GNN框架依赖python编程接口，使用编译器进行静态优化，但是GNN的一些特征在运行时才会显现出来（比如向量大小、节点度数），因此需要运行时优化。

提出GNNAdvisor

GNNAdvisor使用Pytorch作为前端，在底层是用CUDA写的，并用pytorch wrapper集成到pytorch里面。它可以被看作是具有kernel优化、运行时支持的operator

使用GNNAdvisor实现的两层GCN如下：

import GNNAdvisor as GNNA
import torch

# create a GCN class
class GCN(torch.nn.Module):
    def __init__(self, inDim, hiDim, outDim, nLayers):
        self.layers = torch.nn.ModuleList()
        self.layers.append(GNNA.GCNConv(inDim, 	hiDim))
        for i in range(nLayers - 2):
            layer = GNNA.GCNConv(hiDim, hiDim)
            self.layers.append(layer)
        self.layers.append(GNNA.GCNConv(hiDim, outDim))
        self.softmax = torch.nn.Softmax()
    
    def forward(self, X, graph, param):
        for i in range(len(self.layers)):
            X = self.layers[i](X, graph, param)
            X = self.ReLU(X)
        X = self.softmax(X)
        return X

# define a 2-layer GCN model
model = GCN(inDim=100, hiDim=16, outDim=10, nLayers=2)

# loading graph and extracting input properties
graphObj, inputInfo = GNNA.LoaderExtractor(graphFile, model)

# set runtime parameters automatically
X, 	graph, param = GNNA.Decider(graphObj, inputInfo)

# run model
predict_y = model(X, graph, param)