Parallel processing for signal generation neural networks

1. A method comprising:

receiving weight matrices of a layer of a plurality of layers of a signal generation neural network, wherein each layer of one or more layers in the neural network has a residual connection to a subsequent layer, and wherein each layer of the plurality of layers has a skip connection, wherein for each layer:

a respective first weight matrix comprises values for a first matrix operation of the layer, and

a residual weight matrix and a skip weight matrix each comprise values for another matrix operation of the layer, wherein each of the other matrix operations depends on a result of the first matrix operation;

interleaving rows of the first weight matrix for the layer by assigning groups of rows of the first weight matrix to respective thread blocks of a plurality of thread blocks, each thread block being a computation unit for execution by an independent processing unit of a plurality of independent processing units of a parallel processing device, wherein the first weight matrix has a number of rows that is a multiple of a number of independent processing units of the parallel processing device, each independent processing unit being a streaming multiprocessor;

assigning groups of rows of the residual weight matrix to a first subset of the plurality of thread blocks and assigning groups of rows of the skip weight matrix to a second subset of the plurality of thread blocks that does not overlap with the first subset of the plurality of thread blocks;

receiving, by the layer, an input vector;

performing, by the plurality of thread blocks, the first matrix operation corresponding to the first weight matrix substantially in parallel;

after performing the first matrix operation corresponding to the first weight matrix:

performing, by the first subset of the plurality of thread blocks, the other matrix operation corresponding to the residual weight matrix substantially in parallel with performing, by the second subset of the plurality of thread blocks, the other matrix operation corresponding to the skip weight matrix; and

outputting a result computed from the first subset of the plurality of thread blocks to the residual connection in the network and outputting a result computed from the second subset of the plurality of thread blocks to the skip connection in the network.

2. The method of claim 1 , further comprising performing, by the plurality of thread blocks, a synchronization after performing the first matrix operation.

3. The method of claim 1 , wherein the first subset of the plurality of thread blocks comprises half of the available thread blocks, and wherein the second subset of the plurality of thread blocks comprises another half of the available thread blocks.

4. The method of claim 1 , wherein interleaving rows of the first weight matrix for the layer comprises assigning, to each thread block, at least one row of the first weight matrix.

5. The method of claim 1 , wherein the residual weight matrix and the skip weight matrix each have a number of columns that is a multiple of a number of threads within a warp of the parallel processing device.

6. A system comprising: one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform operations comprising:

receiving weight matrices of a layer of a plurality of layers of a signal generation neural network, wherein each layer of one or more layers in the neural network has a residual connection to a subsequent layer, and wherein each layer of the plurality of layers has a skip connection, wherein for each layer:

a respective first weight matrix comprises values for a first matrix operation of the layer, and

a residual weight matrix and a skip weight matrix each comprise values for another matrix operation of the layer, wherein each of the other matrix operations depends on a result of the first matrix operation;

interleaving rows of the first weight matrix for the layer by assigning groups of rows of the first weight matrix to respective thread blocks of a plurality of thread blocks, each thread block being a computation unit for execution by an independent processing unit of a plurality of independent processing units of a parallel processing device, wherein the first weight matrix has a number of rows that is a multiple of a number of independent processing units of the parallel processing device, each independent processing unit being a streaming multiprocessor;

assigning groups of rows of the residual weight matrix to a first subset of the plurality of thread blocks and assigning groups of rows of the skip weight matrix to a second subset of the plurality of thread blocks that does not overlap with the first subset of the plurality of thread blocks;

receiving, by the layer, an input vector;

performing, by the plurality of thread blocks, the first matrix operation corresponding to the first weight matrix substantially in parallel;

after performing the first matrix operation corresponding to the first weight matrix:

performing, by the first subset of the plurality of thread blocks, the other matrix operation corresponding to the residual weight matrix substantially in parallel with performing, by the second subset of the plurality of thread blocks, the other matrix operation corresponding to the skip weight matrix; and

outputting a result computed from the first subset of the plurality of thread blocks to the residual connection in the network and outputting a result computed from the second subset of the plurality of thread blocks to the skip connection in the network.

7. The system of claim 6 , wherein the operations further comprise performing, by the plurality of thread blocks, a synchronization after performing the first matrix operation.

8. The system of claim 6 , wherein the first subset of the plurality of thread blocks comprises half of the available thread blocks, and wherein the second subset of the plurality of thread blocks comprises another half of the available thread blocks.

9. The system of claim 6 , wherein interleaving rows of the first weight matrix for the layer comprises assigning, to each thread block, at least one row of the first weight matrix.

10. The system of claim 6 , wherein the residual weight matrix and the skip weight matrix each have a number of columns that is a multiple of a number of threads within a warp of the parallel processing device.

11. A non-transitory computer storage medium encoded with a computer program, the program comprising instructions that are operable, when executed by data processing apparatus, to cause the data processing apparatus to perform operations comprising:

receiving weight matrices of a layer of a plurality of layers of a signal generation neural network, wherein each layer of one or more layers in the neural network has a residual connection to a subsequent layer, and wherein each layer of the plurality of layers has a skip connection, wherein for each layer:

a respective first weight matrix comprises values for a first matrix operation of the layer, and

a residual weight matrix and a skip weight matrix each comprise values for another matrix operation of the layer, wherein each of the other matrix operations depends on a result of the first matrix operation;

interleaving rows of the first weight matrix for the layer by assigning groups of rows of the first weight matrix to respective thread blocks of a plurality of thread blocks, each thread block being a computation unit for execution by an independent processing unit of a plurality of independent processing units of a parallel processing device, wherein the first weight matrix has a number of rows that is a multiple of a number of independent processing units of the parallel processing device, each independent processing unit being a streaming multiprocessor;

assigning groups of rows of the residual weight matrix to a first subset of the plurality of thread blocks and assigning groups of rows of the skip weight matrix to a second subset of the plurality of thread blocks that does not overlap with the first subset of the plurality of thread blocks;

receiving, by the layer, an input vector;

performing, by the plurality of thread blocks, the first matrix operation corresponding to the first weight matrix substantially in parallel;

after performing the first matrix operation corresponding to the first weight matrix:

performing, by the first subset of the plurality of thread blocks, the other matrix operation corresponding to the residual weight matrix substantially in parallel with performing, by the second subset of the plurality of thread blocks, the other matrix operation corresponding to the skip weight matrix; and

outputting a result computed from the first subset of the plurality of thread blocks to the residual connection in the network and outputting a result computed from the second subset of the plurality of thread blocks to the skip connection in the network.

12. The non-transitory computer storage medium of claim 11 , wherein the operations further comprise performing, by the plurality of thread blocks, a synchronization after performing the first matrix operation.

13. The non-transitory computer storage medium of claim 11 , wherein the first subset of the plurality of thread blocks comprises half of the available thread blocks, and wherein the second subset of the plurality of thread blocks comprises another half of the available thread blocks.

14. The non-transitory computer storage medium of claim 11 , wherein interleaving rows of the first weight matrix for the layer comprises assigning, to each thread block, at least one row of the first weight matrix.