Just-in-time re-partitioning of feature maps for efficient balancing of compute core workloads

Abstract

Certain aspects of the present disclosure provide techniques for partitioning feature maps to improve machine learning model processing. In one aspect, a method, includes partitioning a feature map row-wise into a plurality of feature sub-maps such that: each respective feature sub-map of the plurality of feature sub-maps is defined with respect to a split row determined based on a dense data element count for each row of the feature map; and each feature sub-map of the plurality of feature sub-maps has a same column dimensionality as the feature map; and assigning each of the plurality of feature sub-maps to one of a plurality of tensor compute units and one of a plurality of tensor feature map memory units for processing in parallel.