Tiles, halos and partial sums

Partition tensors to fit storage

Choose output height Th,width Tw,input-channel count Tc and output-channel count Tm for one tile. An input halo extends beyond the output tile. For stride S,dilation D and kernel R, required input height is (Th−1)S+(R−1)D+1; apply the horizontal counterpart for width. Boundary padding may be generated instead of fetched.

Splitting input channels requires preserving each output partial sum across all channel tiles. Rounding each tile early changes the operation. Add bias once in the first channel tile, then continue the same sum. Changing output-channel tiles also changes weights and biases.

Worked storage budget

For a 3×3, unit-stride, unit-dilation operation, let Th=Tw=2,Tc=2,Tm=2. The halo is 4×4, requiring 32 int8 input bytes. Weights use 2×2×3×3=36 bytes. Eight 32-bit partial sums use 32 bytes, and two 32-bit biases use eight. Total:108 bytes. Scales, addresses, validity and control storage are excluded and need a separate allowance.

Double-buffering input and weights requires 2×(32+36)+32+8=176 bytes, beyond a 128-byte budget. Reducing Tc to one gives input sixteen and weights eighteen bytes. Double-buffered storage becomes 2×34+32+8=108 bytes. Preserve partial sums through the two input-channel tiles. Fitting buffers does not prove that transfers can overlap computation; check the schedule separately.

Stride two enlarges the same 2×2 output tile's input span to 5×5. Adjacent tiles may share halo values, but independent reloads still duplicate traffic. Distinguish unique tensor bytes, a lower bound, from actual tile transfers.

If M=5 and Tm=2, the last tile has one valid output channel. Mask the extra weight and output access. Apply equivalent masks to final spatial and input-channel tiles. NVIDIA CUTLASS documentation describes a real implicit-GEMM approach that generates tiles without a full expanded matrix. We calculate storage and partial sums here rather than executing that CUDA implementation.