PyTorch CPU fold path eyed for parallel col2im speedups
A proposed change would process output channels concurrently and cut redundant writes on common Fold shapes.
A PyTorch contributor has proposed speeding up the CPU implementation of col2im, the routine behind torch.nn.functional.fold, by running output channels in parallel and writing each covered location once when tiles do not overlap.
Fold rebuilds spatial tensors from columnized patches and is used in vision and other workloads that stay on CPU. The existing path walks NCHW channels serially. The candidate uses PyTorch’s channel-level parallel loop for larger outputs, pulls width checks out of the hot path, and keeps the same addition order and signed-zero behavior so results stay bit-identical where that matters.
On an Intel Core Ultra 7 host with four threads, the author reported roughly 10× lower median times for several FP32 Fold shapes (for example about 1.5 ms to 0.11 ms on 16-channel 368×368 with 2×2 kernels). Checks across seven dtypes, one and four threads, gradients, and sanitizers were described as clean; a full tree build and upstream regression suite were still pending.
The change is scoped to the native CPU col2im helpers. Reviewers still need to decide whether that narrow surface is enough for merge.