Ordering computations of a machine learning network in a machine learning accelerator for efficient memory usage

A method for generating a computer program to implement a machine learning network on a machine learning accelerator (MLA), the MLA comprising one or more meshes of interconnected processing elements implemented on a semiconductor die, the method comprising:

• receiving a description of the machine learning network, the machine learning network comprising a plurality of interconnected layers;
• determining for each of a plurality of outputs of a layer of the machine learning network, respective corresponding partial networks that each compute one of the outputs and a set of intermediate outputs from prior layers on which the output is dependent;
• determining an order to implement the respective partial networks based on an overlap between the partial networks;
• allocating computations of the machine learning network to the processing elements;
• generating a set of instructions that implement the machine learning network on the processing elements;
• scheduling execution of the instructions to implement the respective partial networks according to the determined order;
• and
• outputting the computer program.

A system comprising a machine learning accelerator (MLA) implemented on a semiconductor die, the MLA comprising:

• an on-chip memory system configured to store data used in computations for implementing a machine learning network;
• and
• a mesh of interconnected processing elements coupled to the on-chip memory system, the processing elements including integrated L1 memory, the processing elements configured to execute instructions of a computer program to implement the machine learning network by sequentially implementing a set of partial networks corresponding to respective outputs of a layer of the machine learning network, each of the partial networks including computations for producing the output of the layer and a set of intermediate outputs from one or more prior layer on which the output is dependent, the instructions when executed causing the processing elements to:
• for a first partial network of the machine learning network, obtaining a first set of intermediate outputs from a memory and perform a first computation on the first set of intermediate outputs of the first partial network to generate the first output;
• and
• following completion of the first computation, computing, for a second partial network of the machine learning network, at least one of the second set of intermediate outputs of the second partial network and removing at least one of the first set of intermediate outputs from the memory.