Allocating computations of a machine learning network in a machine learning accelerator

A non-transitory computer readable storage medium storing instructions for implementing a machine learning network on a machine learning accelerator (MLA), the MLA comprising one or more meshes of interconnected Tiles implemented on a semiconductor die, the instructions when executed by one or more processors causing the one or more processors to perform steps including:

• receiving a description of the machine learning network, the machine learning network comprising a plurality of interconnected layers and the description specifying computations performed to implement each of the layers;
• and
• generating a computer program that implements the machine learning network on the MLA, wherein the computer program comprises Tile instructions for execution by the Tiles to implement the machine learning network, and wherein generating the computer program comprises:
• for each of one or more of the layers in the machine learning network:
• determining, for that layer, a partial computation metric based on the computations performed to implement that layer;
• allocating a group of Tiles to execute the Tile instructions implementing the computations of that layer, wherein a number of Tiles in the group is based on the partial computation metric for that layer;
• and
• outputting the computer program.

A machine learning accelerator device comprising:

• one or more meshes of interconnected Tiles implemented on a semiconductor die;
• a controller comprising a processor and a non-transitory computer-readable storage medium for storing a computer program executable by the processor, wherein the computer program when executed causes the processor to perform steps including:
• obtaining an allocation of Tile instructions of the computer program to be executed by different ones of the Tiles, the allocation causing different blocks of Tiles to perform computations associated with different layers of a machine learning network, wherein the different blocks are sized based on respective corresponding proportions of respective partial computation metrics for the different layers relative to an overall computation metric for the machine-learned model;
• and
• distributing the instructions to the Tiles based on the allocation.