RFC: fuse bias/activation epilogue through the einsum unmerge reshape (workload-dependent)

[czoli1976]

This is an RFC / discussion draft — not proposed for merge as-is. It implements a real fusion, but the measurements show it's a net win only for memory-bound models and a regression for compute-bound ones, so the question is whether you'd want it behind a cost heuristic (or at all).

Observation

Pointwise conv and Einsum lower to:

OptMatMul → IntoShape(unmerge_out) → elementwise   (per-channel bias add, ReLU, …)

The matmul epilogue can already absorb these (FusedSpec::{BinScalar, BinPerRow, BinPerCol}), and the direct matmul → bias case does fuse in OptMatMul::fuse. But when the unmerge reshape sits between them, fusion is blocked: the reshape-reorder branch in OptMatMul::fuse only handled TypedBinOp/Cast with a uniform operand, not the codegen'd OptBinByScalar/OptBinUnicast per-channel forms.

What this patch does

Extends that reorder to:

• match the OptBinByScalar / OptBinUnicast forms (rebuilding a TypedBinOp to rewire), and
• handle a per-channel constant operand: its single non-unary axis is tracked back through the reshape's AxesMapping (track_axis) to the matmul output axis, and the constant is reshaped onto that axis. A later fuse pass then turns it into a BinPerRow/BinPerCol epilogue.

It only fires when the tracked axis lands on c_m_axis/c_n_axis (where fuse_binary can express it) and bails otherwise.

Correctness

Bit-identical output (max_abs_diff = 0) on every model I tested — DeepFilterNet3 (enc/erb_dec/df_dec), MobileNetV2 (98→44 epilogue ops fused), SqueezeNet (49→18). cargo test -p tract-core passes (252).

Performance — the catch

Paired benchmarks (tract … -O bench, baseline vs this, interleaved to cancel thermal drift) on aarch64 macOS:

Model	Regime	Result
DeepFilterNet3 enc	memory-bound streaming	~10% faster (bias-only); ~5% with ReLU fusion too
DeepFilterNet3 erb_dec / df_dec	memory-bound	~4–5% faster
MobileNetV2	compute-bound CNN	~8–10% slower

The MobileNet regression holds for bias-only (ReLU excluded), so it's the bias fusion itself. Intuition: fusing the epilogue into a compute-bound matmul's hot loop costs more than the cheap standalone vectorized pass it replaces, whereas on memory-bound matmuls (where the bias is a large fraction of runtime) fusing avoids a full output read-modify-write and wins.

Question

Is fusing bias/activation through the unmerge reshape something you'd want — gated by a cost heuristic (e.g. only when the matmul is memory-bound / low arithmetic intensity)? Or is it deliberately left undone for this reason? Happy to rework into a gated version (or drop it) based on your preference — I didn't want to guess a threshold from two data points.

---

🤖 Generated with Claude Code