Inference Systems · Graph Compilation

Tile Selection: Cost Models Against Measurement

A fused kernel still needs a tile shape and a loop order, and the compiler must choose them either by estimating cost or by measuring candidates on the device.

Fusion decides which operators run in one kernel. It does not decide how that kernel walks its iteration space. A fused group implementing a matrix multiply with a bias and an activation has at least four free parameters: the tile extents along the output rows and columns, the size of the reduction step, and the order in which the resulting loops nest. Each combination produces a correct kernel. Their throughput differs by a large factor, and which combination is best depends on the extents, the dtype, and the specific device.

This is the selection problem. The compiler must pick one configuration per kernel, and it has two ways to do it: estimate the cost analytically, or generate candidates and time them.

The structure of the search space

For a tiled matrix multiply computing $C = AB$ with $C$ of shape $M \times N$ and reduction depth $K$, the standard decomposition fixes a block tile $[B_M, B_N]$ assigned to one thread block, a per-thread tile $[T_M, T_N]$ within it, and a reduction step $K_S$. The number of blocks is $\lceil M/B_M \rceil \times \lceil N/B_N \rceil$, and each block streams through $K$ in steps of $K_S$, loading a $[B_M, K_S]$ panel of $A$ and a $[K_S, B_N]$ panel of $B$ into shared memory.

The constraints that bound the space are hardware limits, and they are checkable before any code is generated:

  • Shared memory per block must hold both panels: $4(B_M K_S + K_S B_N)$ bytes for float32. On a device with 48 KB of shared memory per block, that caps $B_M K_S + K_S B_N$ at 12,288 elements.
  • Registers per thread must hold the accumulator tile plus the fragments being multiplied: roughly $T_M T_N$ accumulators, which at 255 registers per thread caps $T_M T_N$ in the low hundreds.
  • The block count must be large enough to fill the device. A tile so large that only one block fits per multiprocessor leaves the rest idle, and a tile so small that the block count exceeds the resident capacity adds tail effects.
def candidates(M, N, K, smem_bytes=48 * 1024, max_threads=1024):
    """Enumerate feasible (BM, BN, KS, TM, TN) tiles for a float32 matmul."""
    out = []
    for BM in (32, 64, 128, 256):
        for BN in (32, 64, 128, 256):
            for KS in (8, 16, 32, 64):
                panels = 4 * (BM * KS + KS * BN)      # bytes of A and B panels
                if panels > smem_bytes:
                    continue
                for TM in (2, 4, 8):
                    for TN in (2, 4, 8):
                        threads = (BM // TM) * (BN // TN)
                        if not (32 <= threads <= max_threads):
                            continue
                        if threads % 32:                  # must fill whole warps
                            continue
                        if TM * TN > 128:                 # accumulator pressure
                            continue
                        out.append((BM, BN, KS, TM, TN))
    return out

The enumeration above is a few hundred to a few thousand configurations depending on the bounds. That is small enough to measure and far too large to measure often: each measurement requires generating, compiling, and running a kernel, and a graph with thirty distinct fused groups would need thirty searches.

Analytical models and where they break

A cost model predicts the time of a configuration from its tile parameters without running it. The simplest useful model treats the kernel as bounded by whichever of two resources saturates first: bytes moved from global memory, or arithmetic operations issued.

For the matrix multiply, each thread block loads a $[B_M, K_S]$ panel of $A$ and a $[K_S, B_N]$ panel of $B$ on every step of the reduction, so over the full depth $K$ it moves $4K(B_M + B_N)$ bytes while performing $2 B_M B_N K$ operations. The arithmetic intensity of the block is therefore $B_M B_N / (2(B_M + B_N))$ operations per byte, which does not depend on $K$ at all: a larger block tile raises intensity only through the ratio of its area to its perimeter. Total traffic across the grid is that per-block figure multiplied by $\lceil M/B_M \rceil \lceil N/B_N \rceil$ if no two blocks ever share a cached panel, and the whole kernel performs $2MNK$ operations. For large square matrices the resulting intensity is high enough that the kernel is compute-bound, and the model predicts the time as operations divided by peak throughput. For a thin matrix, such as a single-row matrix-vector product, the intensity collapses, the kernel is bandwidth-bound, and the model predicts bytes divided by bandwidth.

That model is accurate to within a factor of two on large shapes and wrong in a specific, instructive way on small ones. It does not account for the latency of filling the pipeline at the start of each block, which dominates when each block does only a few reduction steps. It does not account for shared-memory bank conflicts, which depend on the exact stride of the panels and on whether the inner dimension is padded. And it does not account for the occupancy the tile implies, which is a function of registers and shared memory together and which decides whether a memory stall is hidden by another resident block.

Compilers that use analytical models therefore add correction terms for occupancy and for the tail, and they accept a residual error on shapes far from the ones the model was calibrated on. The practical consequence is that the model chooses well among configurations that differ by a factor of two, and cannot distinguish configurations that differ by ten percent.

Measurement, and what it must control

Autotuning abandons prediction and measures. The compiler generates a kernel per candidate, runs each on representative inputs, records the time, and keeps the fastest. The result is exact for the measured shape on the measured device, which is precisely its limitation.

Three things must be controlled for the measurement to mean anything. The first is the clock: a device that has downclocked for thermal reasons reports a longer time for whichever candidate ran last, so candidates must be interleaved or the measurement repeated. The second is the input data: a kernel’s timing can depend on the values it processes when it contains a branch, as a sparse or masked path does, so the tuning inputs must resemble the real ones in distribution and not merely in shape. The third is the launch overhead: a kernel that runs for a few microseconds is dominated by the cost of launching it, so a measured difference between candidates may reflect enqueue timing rather than execution. Timing many repetitions and dividing is the standard remedy.

The cost of measurement is proportional to the number of candidates times the time per measurement. A search over a thousand candidates at ten milliseconds each takes ten seconds per kernel, which is acceptable at model build time and unacceptable at request time. This is why autotuning results are cached: keyed by the operator’s shape and dtype, the device model, and the library version, and shipped with the artifact or stored in a tuning database that persists across deployments.

Hybrid strategies

The two approaches compose. The common arrangement is a coarse analytical filter followed by measurement over the survivors. The filter uses the resource bounds above to discard configurations that cannot fit, then uses the cost model to rank what remains and keeps the top handful. Measurement then distinguishes among those, which is where the model is weakest and where measurement is cheapest because the candidate set is small.

A second arrangement inverts the roles: measure a small set of configurations on a representative shape, fit a model to those measurements, and use the fitted model to predict the best configuration for shapes that were never measured. This is the approach taken by learned cost models, which train a regressor on measured kernel times across many shapes and then evaluate it in microseconds at compile time. The model is only as good as its coverage, and a shape outside the training distribution gets a prediction with no guarantee.

The third arrangement is to avoid choosing at all. A kernel can be generated with a small runtime dispatch: compile two or three variants and select among them at launch based on the actual extents, which are known then even if they were not known at compile time. The dispatch costs a branch and keeps the artifact small, and it converts a compile-time decision into a run-time one that has better information.

Selection is per device, and devices change

The tuned configuration is a property of the hardware, not of the model. A grid sized so that one block lands on each multiprocessor of a small device under-fills a device with twice as many multiprocessors, and a tile so large that only one block fits on a multiprocessor leaves the rest of that unit’s resources unused for the duration. A reduction step tuned to a particular shared-memory latency is not the best step on an architecture with a different ratio between shared memory and register bandwidth. Compilers therefore key their tuning databases on the device, and a model artifact that embeds tuned parameters is tied to the device it was tuned on.

The consequence for a serving fleet with mixed hardware is that the same graph compiles to different artifacts per device class, and the artifacts cannot be swapped. A fleet that adds a new accelerator generation re-runs the tuning for its models, which is a build-time cost paid once per model per device rather than per request.

Selection is also where the compiler’s assumptions about shapes become concrete. A kernel tuned for a batch of thirty-two is not the kernel that should run for a batch of one, and the difference is not a small correction: at batch one the matrix multiply is a matrix-vector product whose best configuration has a completely different tile shape. A compiler that tunes for one shape and serves another has optimised the wrong problem, and no amount of fusion or memory planning compensates for it.

References

  1. Ansor: Generating High-Performance Tensor Programs for Deep Learning — Zheng et al., OSDI 2020
  2. Learning to Optimize Tensor Programs — Chen et al., NeurIPS 2018
  3. TVM: An Automated End-to-End Optimizing Compiler for Deep Learning — Chen et al., OSDI 2018
  4. Roofline: An Insightful Visual Performance Model for Multicore Architectures — Williams, Waterman, Patterson; Communications of the ACM, 2009