A Sigmoid After MatMul: Why the Compiler Still Refuses to Fuse It

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

“If convolution can fuse, surely the Sigmoid after this matrix multiplication can too?”

Sometimes it can. Sometimes the compiler quite deliberately says no. A useful optimization article should show more than the attractive graph after a successful merge. It should explain nodes intentionally left in place, because they often mark real semantic boundaries.

Direct Adjacency and Semantic Barriers

The simplest teaching graph is:

1

X:[B,M,K] -- MatMul(W:[K,N]) --> U:[B,M,N] -- Sigmoid --> Y

If only Sigmoid uses U, the target supports the corresponding postprocessing, and the intermediate type and LUT parameters are complete, a fused compute node can produce Y as its final output. The historical implementation added an optional LUT input, a function identifier, and an intermediate-type description, completed address and parameter loading, and checked results with graph-level and code-generation tests.

Put this graph on the same page, however, and the question changes:

1

MatMul -> Requantize -> Sigmoid

Requantize is not a sticky note saying “just passing through; ignore me.” It can change scale, zero point, integer width, rounding, and saturation. Suppose an intermediate real value is 0.6 and a requantization boundary maps it to an integer representative with step size 1. The activation may receive 1; applying Sigmoid directly to 0.6 gives a different result. Exact decimal evaluations are unnecessary to see that the function input has changed.

If fused hardware postprocessing can reproduce this conversion exactly, broader fusion could be designed. That requires an additional proof and parameter representation. The verified historical pattern matched only the direct producer, so it retained the original graph when Requantize intervened. Tests explicitly locked down that behavior. Rejection is not a lack of courage; it marks the end of the current proof.

Clip is another gate. Clipping to [-1,1] before Sigmoid is plainly different from applying Sigmoid directly to arbitrary original values. A large positive input is first limited to 1, so the activation output no longer approaches 1. Recognizing a familiar activation name does not entitle an optimization to reach through intervening nodes and grab it.

Users, Shapes, and Backend Constraints

Multiple users form a separate boundary. If U feeds Sigmoid on one edge and is returned or passed to another operator on another, a fused producer can replace only the first relationship without changing the other value. Historical rules required a single user, and tests retained a separate unary branch. Such negative cases express design intent more directly than a large random model: this implementation deliberately avoids duplicated computation and dual-output fusion.

Already-fused nodes raise another small but important issue. Consider Y=Sigmoid(Sigmoid(MatMul(...))). Repeating the optimization pass must not overwrite an existing fusion field with the second activation, turning two function applications into one. The historical rule stopped when the producer already had a LUT. This also supports idempotence: running the same optimization again should not continue damaging semantics.

Restricting the activation kind matters as well. Similar piecewise-linear graph nodes may represent GELU, Tanh, or other functions, while the inspected MatMul postprocessing path enabled only one target function. Convolution postprocessing supporting a function does not establish that a matrix-multiplication descriptor has the same protocol. Capability must be demonstrated for each path, not inherited by analogy.

Shape checks reveal an engineering compromise. Historical tests included intermediate and final types with different leading unit dimensions but equal element counts, and the rule allowed certain representational differences on that basis. Generalizing still requires checking linear element correspondence and memory layout. A reshape or permute concealed behind “equal element count” could make the fused computation write matrix results at incorrect coordinates. Equal totals are a checkpoint, not a universal pass.

Code generation validated incomplete states again: a LUT required its metadata; no LUT permitted no lingering fusion attributes; intermediate values needed supported storage types; and table addresses needed prior allocation. It also rejected grouped code-generation cases not implemented in the inspected version. This restriction belongs in the design account because graph fusion and later tiling or grouping do not naturally commute.

Possible engineering approaches include querying target capability earlier, retaining a legal alternative during grouping, or running fusion once the final scheduling structure is known. Every ordering affects other optimization opportunities.

Parameter Completeness and Validation

At parameter level, fusion needs more than a Boolean. It requires coherent table addresses, header parameters, intermediate encoding, final output encoding, function selection, and postprocessing enablement. The historical kernel loaded LUT parameters before matrix computation. Tests inspected descriptors, address records, parameter records, and final instructions, checking that the independent activation instruction disappeared. Ordinary MatMul retained comparison tests to avoid additional parameter loads when fusion was absent.

A model-level script also built a small MatMul+Sigmoid graph and checked structure before and after lowering, fusion attributes, and code-generation artifacts. This is closer to the complete flow than hand-writing final backend IR, because it can reveal information lost during import or quantization. It remains mainly a structural and artifact check. The script’s existence does not mean it was executed in this task or proves numerical accuracy for every input.

Three groups of numerical experiments are suggested. First, sample densely in Sigmoid’s central region, where it changes more rapidly, to see how quantization errors propagate through its slope. Second, cover the saturating regions on both sides and check range handling. Third, place inputs close to rounding boundaries and compare fused and original two-node paths. Matching intermediate encodings is necessary to compare the same function.

Performance Costs and Equivalence Boundaries

For performance, start with a simple ledger. Does the unfused path write U into a larger memory? If so, reading U back may be significant; if U remains on chip, the benefit has a different structure. Is the LUT loaded repeatedly? Can matrix computation overlap with postprocessing? Does fusion alter tile sizes and trigger more weight or input rereads? One fewer node answers none of these automatically.

With very small M and N, launch and parameter-loading costs may dominate. With large K, arithmetic may dwarf postprocessing costs. For output sizes sensitive to bandwidth, avoiding intermediate movement may matter more. These are distinct testable hypotheses. Record generated instructions, actual traffic, peak memory, and end-to-end latency together, rather than reporting only one kernel’s time.

Further performance study should first state observable hypotheses about work removed by fusion. Suppose U contains E logical elements of b bytes each. An independent activation reads U once; if matrix computation must also write U to a memory level first, fusion might avoid that write as well. This estimates an idealized upper bound on data volume, not directly on time: accesses may overlap computation or be covered by caches. Stating potential bytes is firmer than guessing speed from node count.

The LUT’s lifetime changes cost too. If repeated matrix calls use the same table, can table contents and parameters be reused? Loading them again for every call may mean that fusion removes a task while retaining most preparation cost. Conversely, if several fused nodes share mutable parameter registers, scheduling must prevent them from overwriting one another. The historical commit supplied the load chain for a single invocation; reuse across tasks needs separate design.

Numerical validation can deliberately approach boundaries. First use a high-precision reference to find inputs placing the matrix result near a quantization half-step, then run the original two-node and fused paths and compare outputs. Random inputs may rarely hit those positions, yet an extra rounding step is particularly visible there. Sigmoid saturation at the tails can flatten some upstream errors, while the central region better reveals small input changes. Both regions matter.

Two meanings of “better precision” should also be distinguished. Keeping a higher-precision intermediate after fusion may bring the result closer to the real-number formula. But if the original graph explicitly requires intermediate quantization, that change may violate strict equivalence. An optimizer should decide in advance whether it aims for bitwise agreement, agreement within an error bound, or permission for particular reassociations and quantization-boundary changes. The historical evidence mainly establishes structural retention of intermediate-type information; it does not contain a complete numerical-equivalence theorem.

For unsupported grouped backend paths, error placement affects usability. If final scheduling discovers the restriction, the diagnostic should explain that “this fused computation cannot yet be generated in this scheduling structure,” rather than vaguely blaming matrix multiplication. Users or upstream passes can then recognize options: disable this fusion, retain the original graph, or change grouping. A specific rejection reason makes a correct fallback possible.

Consider another pass-ordering counterexample. Removing a redundant reshape first may make MatMul and Sigmoid directly adjacent, opening fusion. Inserting a necessary requantize first may intentionally close that opportunity. Pass ordering is therefore not simply a matter of fusing as early as possible. It should follow semantic boundaries and target capability, with final-artifact checks guarding against accidental expansion of matching rules after ordering changes.

When the compiler refuses fusion, walk the graph first: are the nodes directly adjacent, what conversions intervene, how many users exist, is the producer already fused, and does the final scheduling path support it? With explicit answers, refusal itself becomes part of optimization design. A reliable fusion pass knows both when it may combine operations and when it should keep its hands off.

Series contents · Previous · Next

• ← Previous Post
• Next Post →