Schedule-aware tensor distribution module

A deep neural network system, comprising:

• one or more neural network accelerators comprising:
• a plurality of processing engines coupled together and configured to perform arithmetic operations in support of an inference performed using the deep neural network system, wherein the plurality of processing engines are implemented using a plurality of processing elements;
• a schedule-aware tensor data distribution circuitry, comprising:
• a load finite state machine configured to load tensor data into the plurality of processing engines in a load phase;
• a drain finite state machine that is configured to:
• extract output data from the plurality of processing engines in an extraction phase;
• adjust the extracted output data based at least in part on a schedule for a next layer after a current layer to output the extracted output data by reshaping the extracted output data for storage in memory to reduce a number of accesses to the memory for the next layer by changing a shape of how the extracted output data is stored in the memory based at least in part on a type of the next layer;
• and
• store the adjusted extracted output data to a plurality of memory banks.

A method, comprising:

• for a plurality of layers of a deep neural network inference:
• updating configuration descriptor registers used to describe a schedule for a current layer of the plurality of layers;
• based at least in part on the configuration descriptor registers, loading tensor data into a plurality of processing engines;
• using the loaded tensor data, performing computations using the plurality of processing engines for the deep neural network inference;
• after the computations are performed, extracting output data;
• and
• reorganizing the extracted output data to memory banks based on a schedule for a next layer of the plurality of layers by reshaping the extracted output data for storage in the memory banks to reduce a number of accesses to the memory banks for the next layer by changing a shape of how the extracted output data is stored in the memory banks based at least in part on a type of the next layer.

A neural network accelerator, comprising:

• a memory;
• a plurality of processing engines coupled together and configured to perform arithmetic operations in support of an inference performed using the neural network accelerator, wherein the plurality of processing engines are implemented using a plurality of processing elements;
• and
• a distribution circuitry configured to:
• load tensor data into the plurality of processing engines in a load phase;
• extract output data from the plurality of processing engines in an extraction phase,
• wherein extracting the output data from the plurality of processing engines may be performed in a row-wise or column-wise organization;
• reorganize the extracted output data based at least in part on a schedule for a next layer after a current layer to output the extracted output data by reshaping the extracted output data for storage in the memory to reduce a number of accesses to the memory for the next layer by changing a shape of how the extracted output data is stored in the memory based at least in part on a type of the next layer, wherein changing the shape comprises reorganizing the extracted output data to a column-wise organization from the row-wise organization or to the row-wise organization from the column-wise organization based on a specification of a next layer;
• and
• store the reorganized extracted output data to the memory.