Method and system for balanced-weight sparse convolution processing

A system for balanced-weight sparse convolution processing, comprising one or more processors and one or more non-transitory computer-readable memories coupled to the one or more processors and configured with instructions executable by the one or more processors to cause the system to perform operations comprising:

• obtaining an input tensor and a plurality of filters, wherein the input tensor and the plurality of filters have a same number of channels;
• segmenting the input tensor into a plurality of sub-tensors according to a number of available processors;
• pruning each of the plurality of filters so that every predetermined number of channels comprise a fixed number of non-zero weights;
• segmenting each of the plurality of filters into a plurality of the sub-filters, wherein each of the plurality of the sub-filters comprises the same number of non-zero weights;
• respectively assigning the plurality of sub-tensors and the plurality of sub-filters to the available processors, wherein each of the plurality of sub-filters is represented in a compact memory layout storing the non-zero weights of the sub-filter as index-value pairs, each of the index-value pairs comprising a channel-dimension index, a width-dimension index, and a corresponding non-zero weight;
• for each of the available processors, iterating each of the index-value pairs in the compact memory layout of the assigned plurality of sub-filters, and identifying a corresponding input value in the assigned sub-tensor at a location identified by the channel-dimension index and the width-dimension index of the index-value pair to perform a multiply-and-accumulate (MAC) operation to generate a partial sum;
• reassigning the plurality of sub-filters to the available processors by rotating the plurality of sub-filters among the plurality of processors;
• and
• accumulating a plurality of the partial sums generated from each of the plurality of processors to obtain an output tensor.

A non-transitory computer-readable storage medium for balanced-weight sparse convolution processing, configured with instructions executable by one or more processors to cause the one or more processors to perform operations comprising:

• obtaining an input tensor and a plurality of filters, wherein the input tensor and the plurality of filters have a same number of channels;
• segmenting the input tensor into a plurality of sub-tensors according to a number of available processors;
• pruning each of the plurality of filters so that every predetermined number of channels comprise a fixed number of non-zero weights;
• segmenting each of the plurality of filters into a plurality of the sub-filters, wherein each of the plurality of the sub-filters comprises the same number of non-zero weights;
• respectively assigning the plurality of sub-tensors and the plurality of sub-filters to the available processors, wherein each of the plurality of sub-filters is represented in a compact memory layout storing the non-zero weights of the sub-filter as index-value pairs, each of the index-value pairs comprising a channel-dimension index, a width-dimension index, and a corresponding non-zero weight;
• for each of the available processors, iterating each of the index-value pairs in the compact memory layout of the assigned plurality of sub-filters, and identifying a corresponding input value in the assigned sub-tensor at a location identified by the channel-dimension index and the width-dimension index of the index-value pair to perform a multiply-and-accumulate (MAC) operation to generate a partial sum;
• reassigning the plurality of sub-filters to the available processors by rotating the plurality of sub-filters among the plurality of processors;
• and
• accumulating a plurality of the partial sums generated from each of the plurality of processors to obtain an output tensor.