Dynamic partitioning

1. A matrix computation unit comprising a systolic array of cells arranged along a first and second dimension, wherein the systolic array of cells comprises a first plurality of cells, each cell of the first plurality of cells comprising:

a weight register configured to store a weight input;

a plurality of activation registers, each activation register of the plurality of activation registers configured to store a corresponding activation input;

multiplexer circuitry communicatively coupled to the plurality of activation registers and configured to select, from the plurality of activation registers, one of the activation inputs as a selected activation input; and

multiplication circuitry communicatively coupled to the weight register and to the multiplexer, wherein the multiplication circuitry is configured to output a product of the weight input and the selected activation input,

wherein the matrix computation unit further comprises a first value loader communicatively coupled to a first cell of the first plurality of cells, wherein the first value loader is configured to provide a plurality of activation inputs to the plurality of activation registers within the first cell, respectively, and

wherein the first cell is configured to pass the plurality of activation inputs to a next adjacent cell of the first plurality of cells.

2. The matrix computation unit of claim 1 , wherein each cell of the first plurality of cells comprises a vertical partition register communicatively coupled to the multiplexer circuitry and configured to store a multiplexer control signal.

3. The matrix computation unit of claim 2 , comprising a first mux select register communicatively coupled to a first cell of the first plurality of cells, the first mux select register being configured to provide a first multiplexer control signal to a first vertical partition register of the first cell.

4. The matrix computation unit of claim 3 , wherein the first cell is configured to pass the first multiplexer control signal to a next adjacent cell of the first plurality of cells.

5. The matrix computation unit of claim 1 , wherein the first value loader is configured to simultaneously provide the plurality of activation inputs to the plurality of activation registers within the first cell, respectively.

6. A matrix computation unit comprising a systolic array of cells arranged along a first and second dimension, wherein the systolic array of cells comprises a first plurality of cells, each cell of the first plurality of cells comprising:

a weight register configured to store a weight input;

multiplication circuitry communicatively coupled to the weight register, wherein the multiplication circuitry is configured to output a product of the weight input and an activation input;

a sum-in register configured to store a previously summed value;

summation circuity communicatively coupled to the multiplication circuitry and to the sum-in register, wherein the summation circuitry is configured to output a sum of the product and the previously summed value;

demultiplexer circuitry communicatively coupled to the summation circuitry, wherein the demultiplexer circuitry is configured to send the sum from the summation circuitry to another sum-in register in a second adjacent cell of the systolic array of cells along the second dimension or send the sum to an escape line depending on a value of a demultiplexer control signal received by the demultiplexer.

7. The matrix computation unit of claim 6 , wherein each cell of the first plurality of cells comprises a horizontal partition register communicatively coupled to the demultiplexer circuitry, wherein the horizontal partition register is configured to store a demultiplexer control signal and to send the stored demultiplexer control signal to the demultiplexer.

8. The matrix computation unit of claim 7 , comprising a first demux select register communicatively coupled to a first cell of the first plurality of cells, the first demux select register being configured to provide the first demultiplexer control signal to a first horizontal partition register of the first cell.

9. The matrix computation unit of claim 6 , comprising a plurality of sum registers, wherein each escape line of the first plurality of cells is communicatively coupled to a different corresponding sum register of the plurality of sum registers.

10. The matrix computation unit of claim 6 , wherein each escape line of the first plurality of cells is communicatively coupled to a vector computation unit.

11. The matrix computation unit of claim 1 , wherein each cell of the first plurality of cells comprises:

a sum-in register; and

summation circuitry configured to

receive, from the sum-in register, an accumulated value,

receive, from the multiplication circuitry, the product of the weight input and the

selected activation input, and

sum the accumulated value and the product.

12. A method for performing computations using a matrix computation unit comprising a systolic array of cells arranged along a first and a second dimension, the method comprising:

sending a first plurality of activation inputs to a first plurality of activation registers, respectively, within a first cell of the systolic array of cells;

sending a first weight input to the first cell;

sending a first multiplexer control signal to first multiplexer circuitry within the first cell, the first multiplexer circuitry being communicatively coupled to the first plurality of activation registers;

selecting, at the first multiplexer circuitry, a first activation input from the first plurality of activation registers, based on the first multiplexer control signal;

obtaining a first product of the first weight input and the first activation input;

sending the first plurality of activation inputs from the first plurality of activation registers to a second plurality of activation registers, respectively, within a second cell of the systolic array of cells.

13. The method of claim 12 ,

sending a second weight input to the second cell;

sending a second multiplexer control signal to second multiplexer circuitry within the second cell, the second multiplexer circuitry being communicatively coupled to the second plurality of activation registers;

selecting, at the second multiplexer circuitry, a second activation input from the second plurality of activation registers, based on the second multiplexer control signal;

obtaining a second product of the first weight input and the first activation input.

14. A method for performing computations using a matrix computation unit comprising a systolic array of cells arranged along a first and a second dimension, the method comprising:

sending a first activation input to a first cell of the systolic array of cells;

sending a first weight input to the first cell;

obtaining a first product of the first weight input and the first activation input;

obtaining a first accumulated sum based on the first product;

sending the first accumulated sum to demultiplexer circuitry within the first cell;

receiving a demultiplexer control signal at the demultiplexer circuitry; and

sending the first accumulated sum from the demultiplexer circuitry within the first cell to another sum-in register in a second adjacent cell along the second dimension or sending the accumulated sum to an escape line, depending on the value of the demultiplexer control signal.

15. The method of claim 14 , comprising:

receiving, at the first cell, from a third adjacent cell along the second dimension, a previously summed value; and

obtaining the first accumulated sum by combining the previously summed value with the first product at summation circuitry within the first cell.

16. The method of claim 15 , wherein sending the accumulated sum to the escape line comprises sending the accumulated sum to a sum register communicatively coupled to the escape line.

17. The method of claim 15 , wherein sending the accumulated sum to the escape line comprises sending the accumulated sum to a vector computation unit communicatively coupled to the escape line.

18. A method for dynamically partitioning a matrix computation unit comprising a systolic array of cells arranged along a first dimension and a second dimension, the method comprising:

sending a plurality of activation inputs to each cell of a first plurality of cells along the first dimension;

selecting, at multiplexer circuitry within a first cell of the first plurality of cells, a first activation input from the plurality of activation inputs, and selecting at multiplexer circuitry within a second cell of the first plurality of cells a second activation input from the plurality of activation inputs;

obtaining, from each cell of the first plurality of cells along the first dimension, an accumulated sum based on the activation input selected in the cell;

sending the accumulated sum obtained in the first cell to demultiplexer circuitry within the first cell; and

sending the accumulated sum from the demultiplexer circuitry within the first cell to an escape line or to another sum-in register in another adjacent cell along the second dimension, depending on a value of a demultiplexer control signal.

19. The method of claim 18 , further comprising:

sending, to the multiplexer circuitry within the first cell, a first multiplexer control signal, wherein selecting the first activation input is based on a value of the first multiplexer control signal; and

sending, to the multiplexer circuitry within the second cell, a second different multiplexer control signal, wherein selecting the second activation input is based on a value of the second different multiplexer control signal.

20. The method of claim 18 , comprising sending the demultiplexer control signal to the demultiplexer circuitry, wherein the demultiplexer circuitry sends the accumulated sum to the escape line when the demultiplexer control signal has a first value, and wherein the demultiplexer circuitry sends the accumulated sum to the another sum-in register when the demultiplexer control signal has a second different value.