Broadcasting: Sometimes the Right Address Is the One That Does Not Move

Series contents · Operators and Layouts · 阅读中文版

“The framework does this in one line. Why does the backend need so many cases?”

The mathematical definition of broadcasting is compact: align dimensions from the right; corresponding lengths must be equal, or one must be 1. A machine instruction also needs to know where to read, how much to read at once, how far each address advances next time, and which input uses vector reads. Mathematics says “reuse this value.” The machine asks “how many times, and then move how many bytes?”

Translate Broadcasting into Address Traces

Consider a main tensor A:[2,3,5]. Adding B:[1,1,5] uses the same B for every length-5 vector:

1
2
3

for outer in range(2):
    for row in range(3):
        Y[outer,row,:] = A[outer,row,:] + B[0,0,:]

Assuming compact storage for now, with e bytes per element, the main input and output advance by 5e each iteration, B has stride 0, and the loop runs 6 times. Broadcasting has become “one fixed vector address and two moving addresses.” Compactness is an explicit assumption; the next article addresses alignment.

With B:[2,3,1], each row reuses its own scalar, with B advancing by e bytes each time. With B:[1,3,5], every outer block reuses the same entire matrix, giving B an outer stride of 0. With B:[2,1,1], one scalar covers a whole matrix block: the main input advances by a block, and B advances to its next scalar. All are broadcasting mathematically, but their loop plans differ.

The historical implementation classified these cases into a native path, loops for right-side broadcasting, and corresponding paths after swapping operands when allowed. It then constructed block shapes, loop counts, input and output strides, and vector or scalar source modes. This is easier to review than scattering shape checks throughout an emitter: classification says what the case is, planning says how addresses move, and emission says how the plan becomes instructions.

The condition for swapping operands must be explicit. Add and Mul can use commutativity to put the full tensor on one side and reuse a loop. Sub cannot. A tiny example suffices: with full tensor [4,9] and broadcast scalar 2, A-2=[2,7], whereas 2-A=[-2,-7]. Shapes are identical and both are easy to emit, but the result changes sign. The inspected subtraction classifier did not copy addition’s operand-swapping path. That constraint matters.

Middle Axes and Lower-Rank Inputs

Broadcasting along a middle axis often exceeds what a single loop can express. With A:[2,3,5] and B:[1,3,1], each of the three middle-axis values expands over the last dimension and is reused over the first. A backend that handles last-axis vector broadcasting well can first swap the last two axes:

1
2
3
4

A' = transpose(A, [0,2,1])   # [2,5,3]
B' = transpose(B, [0,2,1])   # [1,1,3]
Y' = elementwise(A', B')
Y  = transpose(Y', [0,2,1])

An elementwise operator can be composed with the same coordinate permutation applied to both inputs and the output. This particular swap is its own inverse; arbitrary permutations are not, so a general implementation must calculate the true inverse. The historical middle-axis rule chose this transformation and restricted subtraction to cases where the full input occupied the permitted position.

What does it cost? If every Permute physically moves data, a cheap Add acquires several additional memory accesses. Prearranged constants or merged neighboring permutations may reduce that cost. Correctness and performance need independent acceptance criteria.

Next, a two-dimensional model walks into the test room. The expression [3,5] * [5] is mathematically valid, but a classifier accepting only three-dimensional shapes may return “unsupported” or even choose an incorrect ordinary path. A historical multiplication change normalized shapes by prepending 1s: [3,5] -> [1,3,5] and [5] -> [1,1,5]. Dimensions align from the right, so the extra dimensions belong at the front, not the back.

The scope of this fix deserves precision. The inspected change targeted multiplication classification and its broadcast emission entry point. It does not establish that low-rank handling was unified for every elementwise operator. A common misreading of history is “a shared concept was fixed in one place, so every module was fixed.” If classification code remains duplicated, the regression checklist must still be organized by operator.

Single-element multiplication is even more interesting. Shape [1,1,1] qualifies as both a scalar shape and a vector of length 1. There is no high-level contradiction, but a machine instruction may forbid both inputs from using scalar source mode. The historical fix retained one vector source for ordinary single-element multiplication and rejected two-scalar combinations at instruction construction. A mathematical scalar and a machine’s scalar addressing mode are different concepts.

Why check at two levels? Kernel selection chooses a legal mode so the normal path succeeds; instruction construction prevents other callers from creating illegal encodings. Related tests checked both textual instructions and binary source-type bits, and added negative cases expected to throw. Correct text alone does not prove correct encoding; an encoder that rejects bad inputs alone does not prove that ordinary users can compile successfully.

Address Bounds and Regression Design

Address fields also have representational limits. Even with correct loop logic, a stride can exceed its target field. The worst response is truncation, causing a large model to read an entirely different address. Better options are an early unsupported diagnostic or explicit decomposition into smaller loop plans. For packed low-bit elements, byte width cannot casually be represented by an integer either; the historical path conservatively rejected those cases. Support limits should be visible, not hidden inside integer conversions.

A useful broadcast test is a coordinate-mapping test. Make the main input vary with all three coordinates and the smaller input vary only with retained coordinates, so expected values are easy to calculate. For example, A[a,b,c]=100a+10b+c and B[b]=b+1. A result varying along the wrong axis produces a clear pattern. Fully random tests complement coverage, but these labels make axis mistakes easier to locate.

Regression combinations should also address mathematically valid cases in which both operands need expansion, whether the output shape actually equals the broadcast result, a loop count of 1, low-rank padding, swapping Add and Mul operands, subtraction’s left-right asymmetry, single-element source modes, and diagnostics for unsupported ranks. The historical code mainly implemented plans with one full-shaped operand and a limited selection of smaller shapes. It should not be advertised as a complete arbitrary-broadcasting framework.

Suggested performance measurements separate arithmetic work, loop-control instructions, parameter loading, and actual data movement. Broadcasting a very short vector many times may spend most of its time on loop control; a large broadcast handled by a native vector mode is a different case. Without measurements, we can enumerate costs but cannot conclude that broadcasting is “almost free.”

General References and Quantization Constraints

One mathematically valid case is easy for an implementation to miss: neither input has the complete output shape. Teaching inputs A:[2,1,5] and B:[1,3,1] produce [2,3,5]. A is reused along the middle axis; B is reused along the first and last axes. A classifier looking only for “one side exactly matches the output” finds neither. Per-axis compatibility does not justify silently selecting the old native path.

Possible responses include a more general multidimensional address plan, explicitly expanding an input first, or rejecting the case at the current stage. Each has a cost. Explicit expansion adds memory and movement; multidimensional loops add control and descriptor complexity; rejection limits model coverage. An article should separate “mathematically broadcastable” from “expressible by the current backend plan.” The inspected historical code provided bounded classification, and this account does not repackage it as a universal system.

For a general address reference, assign each input an effective stride for every output axis: zero if the input length on that axis is one, otherwise its actual physical stride. The input address for output coordinate i is then the base address plus the sum of all i_axis*effective_stride_axis terms. This is an excellent independent reference, although implementing it directly as an element-by-element loop may be slow. An optimized block plan should prove that it generates the same addresses with the same repetition counts.

Do not trust the loop count merely because it came from multiplication. Dynamic or zero-length dimensions, product overflow, and signedness conversions can produce unexpected counter values. Many loops that “execute once, then decrement and check” assume a count of at least one. With zero, the machine program may execute once anyway or even keep looping. If upstream rules forbid such dimensions, document that invariant; otherwise, reject them during planning or generate a dedicated empty-tensor path.

Swapping quantized operands also means more than swapping two addresses. Addition operands can have different scales and zero points. When the full tensor moves to a designated side, parameter construction must use the operands in their new semantic positions. Historical Add and Mul code constructed parameters after deciding the left and right values; that ordering matters. Generate parameters first and swap only addresses afterward, and mathematical commutativity will not rearrange the parameter table for you.

Finally, splitting a long vector into blocks depends on tail handling. When the logical length is not divisible by block length, does the implementation shorten the final block, use a mask, or permit reads from defined padding? The smaller broadcast input and output must follow consistent tail rules. The inspected changes mainly concern whole-block shapes and address loops; they do not establish safety for every arbitrary tail. Test “exactly a whole block” alongside “one element more,” before a new vector length brings back a familiar visitor: the first section works, and the last one disappears.

Once broadcasting is translated into address trajectories, debugging becomes surprisingly plain: where does iteration 0 read, where does iteration 1 read, which side should stay still, and when should the row change? The machine does not understand the statement “this axis is broadcast.” It faithfully executes every address increment. The compiler’s job is to make those increments match the mathematics.

Series contents · Previous · Next

• ← Previous Post
• Next Post →