Hardware acceleration framework for graph neural network quantization

1 . 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.

2 . The method of claim 1 , wherein the segmenting the plurality of feature vectors into the plurality of sub-vectors and grouping the plurality of sub-vectors into a plurality of groups of sub-vectors comprises:

segmenting the plurality of feature vectors into the plurality of sub-vectors of a predetermined sub-vector size; and

grouping the sub-vectors into the plurality of groups of sub-vectors based on an offset of each sub-vector within a corresponding feature vector.

3 . The method of claim 1 , prior to the segmenting the plurality of feature vectors into a plurality of sub-vectors, further comprising:

quantifying a sparsity level of the plurality of feature vectors;

determining a sub-vector size based on the sparsity level, wherein the sparsity level is inversely proportional to the sub-vector size; and

segmenting the plurality of feature vectors into the plurality of sub-vectors with the determined sub-vector size.

4 . The method of claim 1 , wherein each of the plurality of sub-vectors comprises one sub-vector from each of the plurality of feature vectors at a same offset.

5 . The method of claim 1 , wherein the vector clustering comprises K-means or Support Vector Clustering (SVC).

6 . The method of claim 1 , wherein the performing vector clustering on each of the plurality of groups of sub-vectors to generate the plurality of centroids comprises:

within each of the plurality of groups of sub-vectors, determining multiple centroids by executing a vector clustering algorithm, wherein each of the multiple centroids comprises a vector of a same size as each of the plurality of sub-vectors.

7 . The method of claim 1 , wherein the storing the plurality of centroids in the codebook comprises:

indexing the plurality of centroids with a plurality of indices; and

storing mappings between the plurality of indices and the plurality of centroids.

8 . The method of claim 1 , wherein the quantizing the sub-vectors within the each feature vector based on the codebook comprises:

for each of the sub-vectors that belongs to one of the plurality of groups of sub-vectors, determining one or more of the plurality of centroids that correspond to the one group of sub-vectors;

determining a vector distance between the sub-vector and each of the one or more centroids;

identifying one of the one or more centroids having a shortest vector distance to the sub-vector; and

encoding the sub-vector using an index of the one identified centroid, wherein storing the index uses less number of bytes than storing the sub-vector.

9 . The method of claim 1 , wherein the GNN computation comprises:

receiving a request for a feature vector representing a graph node;

obtaining the index map corresponding to the requested feature vector, wherein the index map comprises one or more indices corresponding to one or more of the plurality of centroids;

obtaining the one or more centroids based on the one or more indices and the codebook; and

aggregating the one or more centroids as the requested feature vector of the graph node.

10 . The method of claim 1 , wherein the performing vector clustering on each of the plurality of groups of sub-vectors comprises:

filtering out sub-vectors comprising all zeros from the plurality of sub-vectors;

excluding the sub-vectors comprising all zeros from the vector clustering;

assigning a special centroid to represent all of the sub-vectors comprising all zeros; and

performing the vector clustering on the sub-vectors comprising non-zero data.

11 . The method of claim 1 , wherein the method further comprises:

assigning a higher priority for storing the codebook and a lower priority for storing the assignment table, and

wherein the storing the codebook and the assignment table comprises:

in response to an on-chip memory being sufficient to store both the codebook and the assignment table, storing both codebook and the assignment table in the on-chip memory; and

in response to the on-chip memory being insufficient to store both the codebook and the assignment table, storing the codebook in the on-chip memory and storing the assignment table in an off-chip memory.

12 . The method of claim 1 , wherein the batching comprises:

batching the plurality of nodes into the plurality of batches by applying a hash function to node indices of the plurality of nodes.

13 . The method of claim 1 , wherein the batching comprises:

batching the plurality of nodes into the plurality of batches based on locality of each node within the graph.

14 . A hardware accelerator for accelerating graph neural network (GNN) computation based on product quantization, comprising:

an on-chip memory;

a graph node feature vector segmenting circuitry configured to:

receive 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;

segment the plurality of feature vectors into a plurality of sub-vectors;

group the plurality of sub-vectors into a plurality of groups of sub-vectors;

a feature vector clustering circuitry configured to:

perform vector clustering on each of the plurality of groups of sub-vectors to generate a plurality of centroids;

store the plurality of centroids as a codebook;

an encoding circuitry configured to:

encode 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

store the plurality of index maps as an assignment table and at least a portion of the codebook on an on-chip memory of the hardware accelerator to represent the plurality of nodes for the GNN computation using the hardware accelerator, 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, and

the encoding circuitry is further configured to:

batch the plurality of nodes into a plurality of batches and assigning the plurality of batches into a plurality of processing units for parallel encoding, and generate the codebook and the assignment table for each of the plurality of batches.

15 . The hardware accelerator of claim 14 , wherein the quantizing the sub-vectors within the each feature vector based on the codebook comprises:

for each of the sub-vectors that belongs to one of the plurality of groups of sub-vectors, determining one or more of the plurality of centroids that correspond to the one group of sub-vectors;

determining a vector distance between the sub-vector and each of the one or more centroids;

identifying one of the one or more centroids having a shortest vector distance to the sub-vector; and

encoding the sub-vector using an index of the one identified centroid, wherein storing the index uses less number of bytes than storing the sub-vector.

16 . The hardware accelerator of claim 14 , wherein the GNN computation comprises:

receiving a request for a feature vector representing a graph node;

obtaining the index map corresponding to the requested feature vector, wherein the index map comprises one or more indices corresponding to one or more of the plurality of centroids;

obtaining the one or more centroids based on the one or more indices and the codebook; and

aggregating the one or more centroids as the requested feature vector of the graph node.

17 . The hardware accelerator of claim 14 , wherein to perform vector clustering on each of the plurality of groups of sub-vectors, the feature vector clustering circuitry is further configured to:

filter out sub-vectors comprising all zeros from the plurality of sub-vectors;

exclude the sub-vectors comprising all zeros from the vector clustering;

assign a special centroid to represent all of the sub-vectors comprising all zeros; and

perform the vector clustering on the sub-vectors comprising non-zero data.

18 . The hardware accelerator of claim 14 , wherein the codebook has a higher priority for storage in the on-chip memory of the accelerator than the assignment table does, and the encoding circuitry is further configured to:

in response to an on-chip memory being sufficient to store both the codebook and the assignment table, store both codebook and the assignment table in the on-chip memory; and

in response to the on-chip memory being insufficient to store both the codebook and the assignment table, store the codebook in the on-chip memory and store the assignment table in an off-chip memory.