Hardware acceleration framework for graph neural network quantization

A computer-implemented method for accelerating graph neural network (GNN) computation using product quantization, comprising:

• receiving, by a graph neural network (GNN) accelerator on a computing device, a graph comprising a plurality of nodes as an input for the GNN computation, wherein the plurality of nodes are respectively represented by a plurality of feature vectors;
• segmenting, by the GNN accelerator, the plurality of feature vectors into a plurality of sub-vectors;
• grouping, by the GNN accelerator, the plurality of sub-vectors into a plurality of groups of sub-vectors;
• performing, by the GNN accelerator, vector clustering on each of the plurality of groups of sub-vectors to generate a plurality of centroids,
• storing, by the GNN accelerator, the plurality of centroids as a codebook;
• encoding, by the GNN accelerator, each of the plurality of feature vectors into an index map to obtain a plurality of index maps by quantizing sub-vectors within the each feature vector based on the codebook, wherein the index map occupies a smaller memory space than the each feature vector does;
• and
• storing, by the GNN accelerator, the plurality of index maps as an assignment table and storing at least a portion of the codebook on an on-chip memory of the GNN accelerator to represent the plurality of nodes for the GNN computation using the GNN accelerator on the computing device,
• wherein storing the assignment table and the codebook saves memory space and reduces memory access frequency than directly storing the plurality of feature vectors, thereby accelerating the GNN computation,
• wherein the method further comprises batching the plurality of nodes into a plurality of batches and assigning the plurality of batches into a plurality of processing units, and generating the codebook and the assignment table for each of the plurality of batches.