XNNPACK Optimization
Whileone has been making slow incremental changes to XNNPACK that makes the first convolution of common vision models 1.20–1.39x faster on AWS Graviton3, Graviton4 and Graviton5. Every other layer runs at the same speed.
XNNPACK is the CPU engine under LiteRT (TensorFlow Lite), ExecuTorch and MediaPipe. A small win there reaches a very large number of devices.
The problem: a kernel built for wide inputs
On Arm CPUs with the I8MM extension, XNNPACK runs int8 convolutions through matrix-multiply kernels that consume input channels in groups of eight (kr = 8). That suits the deep middle of a network, where layers have dozens or hundreds of channels.
The first layer is different. An RGB image has three channels, so each kernel tap is padded from 3 to 8. Five of every eight multiplies are spent on zeros. The same happens in super-resolution models such as SRCNN, whose first layer reads a single channel. The first layer also runs at full image resolution, so that waste is paid on every pixel of every frame.
The fix: pick the kernel that fits the layer
Most of these CPUs also support the older dot-product instructions, which consume channels in groups of four (kr = 4). Padding 3 channels to 4 wastes one multiply in four instead of five in eight.
The change gives XNNPACK's int8 GEMM configuration an optional small-K companion: a kr = 4 dot-product configuration with its own matching weight packing. When a convolution is created, XNNPACK checks how many input channels it has per group. If they fit in four, it uses the companion; otherwise nothing changes. The decision is made once, at operator creation, so there is no cost at inference time.
The change is about 100 lines across four files, including seven new tests for convolutions with one to four input channels.
Results: 1.2–1.4x where it matters, nothing lost elsewhere
We measured upstream master against the change on three generations of AWS Graviton, built with clang 21. Each run used one pinned thread, five repetitions and two rounds in opposite order.
The other 117 int8 convolution layers in the benchmark showed a median change of 0.0% to −0.3%, which is within noise.
The benchmarking also turned up a compiler issue. GCC builds XNNPACK's I8MM kernel about 2x slower than clang on Graviton4, because one of its optimization passes splits the kernel's register tuples. We traced it to long-standing GCC bug 8996.
@ Whileone
We read the workload, not just the kernel. The gain came from noticing that the first layer's shape doesn't fit the fastest kernel, then choosing a better one per layer.
We measure like a vendor would. That means three CPU generations, pinned threads, alternating rounds, and tables that show the layers that didn't move as well as the ones that did.
We go down the whole stack. One investigation took us from model layers through micro-kernels to a fix in the compiler.
We work upstream. Our changes land in the projects our customers already use, reviewed by their maintainers, so nobody has to carry a private fork.
If your models run on Arm servers or devices and you want them faster, in XNNPACK, Arm Compute Library, llama.cpp or your own kernels, talk to Whileone.
Comments