Techniques for performing matrix computations using hierarchical representations of sparse matrices
One embodiment sets forth a technique for performing matrix operations. The technique includes traversing a tree structure to access one or more non-empty regions within a matrix. The tree structure includes a first plurality of nodes and a second plurality of nodes corresponding to non-empty regions in the matrix. The first plurality of nodes includes a first node representing a first region and one or more second nodes that are children of the first node and represent second region(s) with an equal size formed within the first region. The second plurality of nodes include a third node representing a third region and one or more fourth nodes that are children of the third node and represent fourth region(s) with substantially equal numbers of non-zero matrix values formed within the third region. The technique also includes performing matrix operation(s) based on the non-empty region(s) to generate a matrix operation result.
1 . A computer-implemented method for performing one or more matrix multiplication operations based on a first matrix and a second matrix, the method comprising:
transmitting, to each processing element included in a plurality of processing elements, a different portion of one or more nonzero elements of the first matrix that are determined via a tree structure representing the first matrix;
performing, via the plurality of processing elements, one or more computations based on the different portions of the one or more nonzero elements and data associated with the second matrix to generate one or more intermediate results;
storing the one or more intermediate results in a buffer memory; and
combining the one or more intermediate results from the buffer memory into a final result.
2 . The computer-implemented method of claim 1 , further comprising traversing the tree structure to determine the one or more nonzero elements, wherein the tree structure comprises:
a first plurality of nodes that correspond to a first plurality of regions in the first matrix that have an equal size; and
a second plurality of nodes that correspond to a second plurality of regions in the first matrix that have a substantially equal number of non-zero matrix values.
3 . The computer-implemented method of claim 2 , wherein the second plurality of regions in the first matrix have at least two different sizes.
4 . The computer-implemented method of claim 2 , wherein traversing the tree structure comprises determining an overlap between at least one of (i) an upper left corner of a first region represented by a first node included in the tree structure with a lower right corner of a region of interest, or (ii) an upper left corner of the region of interest with a lower right corner of the first region, and wherein at least one nonzero element included in the one or more nonzero elements resides in the first region.
5 . The computer-implemented method of claim 2 , wherein the first plurality of regions comprise four quadrants of a first region included in the first matrix, and the second plurality of regions are included in one of the first plurality of regions.
6 . The computer-implemented method of claim 1 , wherein the data associated with the second matrix comprises one or more nonzero elements of the second matrix that are determined via another tree structure representing the second matrix.
7 . The computer-implemented method of claim 1 , wherein each processing element included in the plurality of processing elements computes a partial result of the one or more matrix multiplication operations.
8 . The computer-implemented method of claim 1 , wherein performing the one or more computations comprises computing one or more partial matrix multiplication results.
9 . The computer-implemented method of claim 1 , wherein combining the one or more intermediate results comprises performing at least one of a sum operation or an aggregation operation.
10 . One or more non-transitory computer-readable media storing instructions that, when executed by at least one processor, cause the at least one processor to perform the steps of:
transmitting, to each processing element included in a plurality of processing elements, a different portion of one or more nonzero elements of a first matrix that are determined via a tree structure representing the first matrix;
performing, via the plurality of processing elements, one or more computations based on the different portions of the one or more nonzero elements and data associated with a second matrix to generate one or more intermediate results;
storing the one or more intermediate results in a buffer memory; and
combining the one or more intermediate results from the buffer memory into a final result.
11 . The one or more non-transitory computer-readable media of claim 10 , wherein the instructions, when executed by the at least one processor, further cause the at least one processor to perform the step of traversing the tree structure to determine the one or more nonzero elements, wherein the tree structure comprises:
a first plurality of nodes that correspond to a first plurality of regions in the first matrix that have an equal size; and
a second plurality of nodes that correspond to a second plurality of regions in the first matrix that have a substantially equal number of non-zero matrix values.
12 . The one or more non-transitory computer-readable media of claim 11 , wherein the second plurality of regions in the first matrix have at least two different sizes.
13 . The one or more non-transitory computer-readable media of claim 11 , wherein the first plurality of nodes is associated with a first stage that comprises metadata identifying a type of the first plurality of regions and boundaries of the first stage.
14 . The one or more non-transitory computer-readable media of claim 11 , wherein the tree structure further comprises a third plurality of nodes that are children of the second plurality of nodes and store one or more nonzero elements within the second plurality of regions.
15 . The one or more non-transitory computer-readable media of claim 10 , wherein the data associated with the second matrix comprises one or more nonzero elements of the second matrix that are determined via another tree structure representing the second matrix.
16 . The one or more non-transitory computer-readable media of claim 10 , wherein performing the one or more computations comprises computing one or more partial matrix multiplication results.
17 . The one or more non-transitory computer-readable media of claim 10 , wherein combining the one or more intermediate results comprises performing at least one of a sum operation or an aggregation operation.
18 . A system, comprising:
one or more memories storing instructions; and
a plurality of processors that are coupled to the one or more memories and, when executing the instructions, are configured to:
receive different portions of one or more nonzero elements of a first matrix that are determined via a tree structure representing the first matrix,
perform one or more computations based on the different portions of the one or more nonzero elements and data associated with a second matrix to generate one or more intermediate results,
store the one or more intermediate results in a buffer memory, and
combine the one or more intermediate results from the buffer memory into a final result.