Batch matrix multiplication operations in a machine learning accelerator

A One or more computer-readable non-transitory storage media embodying software that is operable when executed to, by a computing system for accelerating machine-learning computations comprising a plurality of tensor processor clusters, each comprising a plurality of tensor processor units, an activation direct-memory access (DMA), and a weight DMA:

• determine that a batch matrix multiplication (bmm) operation between a first activation tensor and a second activation tensor needs to be performed, wherein the first activation tensor and the second activation tensor are distributed in activation buffers of N tensor processor units;
• collect, at the weight DMA, the second activation tensor in two blocks from the activation buffers of the N tensor processor units;
• split, at the weight DMA, each of the two blocks of the second activation tensor into a most significant bits (MSB) tile and a least significant bits (LSB) tile, wherein the MSB tile comprises most significant four bits of each byte in the block, and wherein the LSB tile comprises least significant four bits of each byte in the block;
• load the second activation tensor to weight buffers of the N tensor processor units by filling a first entry of each weight buffer of each of the N tensor processor units with contents of the MSB tiles of the two blocks and filling a second entry of the weight buffer with contents of the LSB tiles of the two blocks;
• and
• generate a bmm result using the first activation tensor distributed in the activation buffers of the N tensor processor units and the second activation tensor in the weight buffers of the N tensor processor units.

A method comprising, by a computing system for accelerating machine-learning computations comprising a plurality of tensor processor clusters, each comprising a plurality of tensor processor units, an activation direct-memory access (DMA), and a weight DMA:

• determining that a batch matrix multiplication (bmm) operation between a first activation tensor and a second activation tensor needs to be performed, wherein the first activation tensor and the second activation tensor are distributed in activation buffers of N tensor processor units;
• collecting, at the weight DMA, the second activation tensor in two blocks from the activation buffers of the N tensor processor units;
• splitting, at the weight DMA, each of the two blocks of the second activation tensor into a most significant bits (MSB) tile and a least significant bits (LSB) tile, wherein the MSB tile comprises most significant four bits of each byte in the block, and wherein the LSB tile comprises least significant four bits of each byte in the block;
• loading the second activation tensor to weight buffers of the N tensor processor units by filling a first entry of each weight buffer of each of the N tensor processor units with contents of the MSB tiles of the two blocks and filling a second entry of the weight buffer with contents of the LSB tiles of the two blocks;
• and
• generating a bmm result using the first activation tensor distributed in the activation buffers of the N tensor processor units and the second activation tensor in the weight buffers of the N tensor processor units.