Cloud TPU has two padding rules on batch_size and feature_size of convolution operations, to minimize memory overhead and maximize computational efficiency (from here).
- The total batch size should be a multiple of 64 (8 per TPU core), and feature dimensions should be a multiple of 128,
or
- The total batch size should be a multiple of 1024 (128 per TPU core), and feature dimensions should be a multiple of 8.
If batch size and feature don't conform to the rules, padding occurs. According to the profiling results, the second one (batch_size/core -> 128, feature/core -> 8) is used.
I want to ask the rationale for these rules. As far as I know, the MXU unit is 128x128 systolic array since TPUv2. Why not pad both pre core batch size and feature to 128?