The optimal memory layout usually depends upon which dimensions are the reduction dimensions (for 2d convnet convolution, the spatial dimension on which there i

The optimal memory layout usually depends upon which dimensions are the reduction dimensions (for 2d convnet convolution, the spatial dimension on which there is also reduction is likely less important than the batch or channel dimension). Thus, the optimal memory layout for forward and backward passes usually differs a lot, but transposing between the different layouts on the fly has high cost. Other alternatives beyond just permuting the dimensions include strip mining/tiling and raising/sinking dimensions, techniques which come from loop nest analysis (and correspond exactly to what one would do with the loops of the code): e.g., translating NCHW -> N(C/4)(4)HW -> N(C/4)HW(4) for vectorization purposes, where we turn a 4 dimensional array into a 5 dimensional array, with the innermost dimension being a set of 4 contiguous channels which is sunk into the loop nest, and is amenable to vectorization. Since many of these kernels are hand-tuned or generated by library vendors, there is likely not much of a choice available, but there are likely many other more optimal memory layouts out there that would require machine learning-driven compilation or mathematical optimization techniques like polyhedral compilation to explore and discover.