Just Take the Mean—Why Is the Compiler Writing Loops?

Series contents · Numerics and Quantization · 阅读中文版

“The operator is one line: take the mean.” A colleague slides over a sheet of paper. It really does contain one line. Half an hour later, the whiteboard is covered in axes, strides, temporary buffers, and two loop labels. Someone laughs: “Before we calculate the average, shall we distribute the work evenly?”

The explanation is simple: a model describes mathematical sets, while a vector engine accepts addresses and fixed shapes. Translating “along this axis” into “read here, advance this many bytes, stop there” is the compiler’s job.

1. One Mean, Three Very Different Access Patterns

Consider a contiguous teaching tensor with shape [P,R,Q]=[3,5,7], reduced along the middle axis R. The desired result is:

1

out[p,q] = sum(input[p,r,q], r=0..R-1) / R

For fixed p and q, consecutive contributors to the reduction are Q elements apart. Mathematically they belong to one group; physically they are not one adjacent row.

Reducing the last axis Q reads contiguous segments. Reducing R and Q together lets us combine them into a single contiguous segment. Reducing only R must preserve every position along Q. These cases contain the same number of input elements, but their access patterns differ. An engine designed for trailing-axis reductions does not automatically support the third case simply because the high-level operator has the same name.

An early historical implementation tried to handle these differences in code generation: it grouped adjacent reduction axes, emitted runtime loops for middle-axis reductions, and planned multiple passes for discontiguous axes. Later changes narrowed the low-level support boundary and moved nonnative-axis transformations into the graph layer. This establishes a change in design direction; it does not establish that the first approach is slow on every device.

2. What a Loop Really Maintains Is an Address Invariant

Suppose a kernel processes one p plane at a time. Input elements occupy b bytes, and output elements occupy d bytes. At the end of each iteration, the address increments should be:

1
2

next_input  = current_input  + R * Q * b
next_output = current_output + Q * d

The input and output strides differ, and so may their element sizes. Forgetting that an input axis has been collapsed often produces a correct first block followed by misplaced results. Testing only P=1 gives this bug an excellent hiding place.

The loop invariant can be stated plainly: at the start of iteration p, the input points to the first element of that plane, the output points to its Q results, and the loop counter equals the number of planes remaining. Every address update should follow from this statement. The branch condition also needs a precise definition: after the last computation, do we exit before advancing, or advance first? Even an address beyond the end that is never dereferenced may affect a checked descriptor or later reuse.

Historical changes added branches and address stepping; other changes later removed reduction-loop configuration. A new operator may appear to add one mathematical operation while actually expanding the instruction-generation infrastructure. Once that infrastructure exists, it also brings maintenance obligations: unique labels, register usage, and reproducible output.

3. Multipass Means Agree over the Reals, but Not Necessarily Bit for Bit

For a regular rectangular array, taking the mean along one axis and then another equals a single mean over all the corresponding elements in exact real arithmetic. Fixed-point computation, however, often rounds after each pass and may saturate.

Take the teaching integer groups [0,1] and [0,3]. Their exact overall mean is 1. If each group mean is first rounded down, the intermediate results are 0 and 1. Averaging those and rounding down again gives 0. The discrepancy comes from two rounding steps, not from an incorrect mean formula. A real implementation must be analyzed using its own rounding rule; this example only demonstrates that staged rounding is not associative.

The basic relationship for a quantized mean is:

1
2

real_input = s_in * (q_in - z_in)
q_out ≈ round[(s_in / s_out) * sum(q_in - z_in) / R] + z_out

Both the reciprocal 1/R and the scale ratio need a representation. Approximating them separately and then multiplying is not necessarily equivalent to combining them into one fixed-point coefficient; the choice affects error and register constraints. Historical diffs show changes to where the reciprocal and input/output scale parameters were assigned. Having a parameter is insufficient: it must occupy the position the execution unit actually uses.

Choosing intermediate types for a multipass algorithm raises several questions. Do we retain a wider accumulator? Does each pass apply a zero point? Who chooses the intermediate scale? If one pass saturates, can a later reduction in magnitude recover the lost information? The answer to that last question is usually no.

4. Why Move the Complexity into the Graph?

Another approach first reorders the axes, turning a nonnative reduction into a supported trailing-axis reduction, and then restores the output shape. The low-level implementation only needs to validate a finite set of legal patterns and generate one parameter set.

The benefit is not “less code means faster execution.” The shape transformations become visible in inspectable IR, where existing permutation optimization, buffer planning, and scheduling can act on them. If several passes are hidden inside one code-generation function, the scheduler may see only an opaque large node and miss the lifetimes of its internal temporaries.

The cost is equally clear: an explicit permutation may move substantial amounts of data. Whether it can be absorbed into a producer, a consumer, or a layout interpretation determines whether it is worthwhile. Fewer low-level branches do not imply a shorter runtime.

One possible selection model, not yet measured, is:

1
2

T_loop    ≈ T_read + P * (T_reduce + T_address + T_branch)
T_reorder ≈ T_permute_in + T_native_reduce + T_permute_out

If reordering requires a full read and write while the reduction itself is small, reordering may dominate. If each loop iteration does very little work, instruction and parameter-loading costs may also matter. These should become measurable hypotheses, not a vote based on operator counts.

5. Make Regression Tests Uncomfortable for Both Addresses and Numbers

The history contains compilation scripts and IR transformation tests. Their existence alone does not establish complete numerical or hardware-performance validation. Structural tests answer “what was generated?” Reference computation answers “is the result correct?” Performance experiments answer “why is it fast or slow?” Each kind of evidence must stand on its own.

Look again at that sheet with its one-line formula. The formula was honest; it simply omitted memory’s habits. The pleasure of compiler engineering lies in translating those habits accurately enough that the next maintainer can understand them.

6. A Manual Review That Needs No Hardware

Write out one iteration’s address calculations in a table, deliberately choosing different input and output widths. Keep three outer blocks, each with five rows and seven columns. Let each input element occupy one byte and each output element two bytes. The first input starts at the block beginning; the next input block should be thirty-five bytes ahead. Each iteration produces seven output elements, so the next output block should be fourteen bytes ahead. Adding thirty-five on both sides may still produce a correct first iteration, but the next two leave gaps or even overrun the buffer. Treating output elements as one byte instead causes later writes to overwrite part of the preceding result.

Label the address units in the table as well. Some descriptors take byte addresses; some fields take element counts or counts of fixed-size blocks. Two variables both named stride need not have the same unit. Write units alongside the derivation and centralize conversion at the relevant boundaries. This makes double conversion easy to spot—for example, multiplying by element size inside a function when its caller already did so.

When reviewing a multipass plan, record more than the current shape: record which original axes each current axis represents. Reducing a later group before an earlier one makes it easier to preserve the positions of unprocessed axes. A different order requires updating the axis mapping accordingly. Reusing original axis numbers on a tensor whose rank has already decreased can make the algorithm work only by accident on a particular shape.

Include quantization in the same table: each pass’s input scale, output scale, zero point, temporary width, and rounding behavior. If a pass preserves the input scale, ask whether this protects resolution or merely makes copying the type convenient. If it uses the final scale immediately, check whether small values are lost too early. Keeping separate shape and numerical plans, then reconciling them at every intermediate tensor, exposes gaps more readily than repeatedly searching the code for multiplier.

Finally, give every plan the same assertion: each output element receives exactly the contributor set specified by the original reduction, with every input contribution appearing once. Correct addresses can still produce incorrect values if contributor sets overlap. Equal element counts with different groupings do not define the same mean. This set-based view puts loop and permutation strategies on one common review sheet.

Back to series contents · Next

• ← Previous Post
• Next Post →