(perf) core: fuse a squeezed unit matmul axis and a unicast Mul+Add pair - #2728
(perf) core: fuse a squeezed unit matmul axis and a unicast Mul+Add pair#2728czoli1976 wants to merge 3 commits into
Conversation
|
Not bad at all, @kali |
|
@kali , sorry, folding in an improvement |
|
🔴 Bench vs main — 1 speed regression(s) · Reference: 2026-08-26 morning nightly run (0d old) · full report → run Speed — evaltime · prefill · decode
Improvements
|
| Δ | metric | device | main → PR |
|---|---|---|---|
| arm_ml_kws_cnn_m load · pass |
cortex-a9 |
82 ms → 117 ms | |
| arm_ml_kws_cnn_m load+optimize · pass |
cortex-a9 |
128 ms → 177 ms | |
| en_tdnn_15M RSS @ ready · 2600ms |
cortex-a55 |
111 MB → 117 MB |
|
Added a third commit,
It is smaller than it looks: the clone is a flat memcpy while the copy that The perf table in the description is re-measured with all three commits in. |
|
Reference: 2026-08-26 morning nightly run (0d old) · full report → run Speed — evaltime · prefill · decode no inference-speed regressions Improvements
|
| Δ | metric | device | main → PR |
|---|---|---|---|
| arm_ml_kws_cnn_m load · pass |
cortex-a9 |
82 ms → 94 ms | |
| hey_snips_v4_model17 load · 2sec |
cortex-a7 |
4.21 s → 4.79 s | |
| hey_snips_v4_model17 load+optimize · 2sec |
cortex-a7 |
4.98 s → 5.66 s | |
| mobilenet_v2_1 RSS @ ready · pass_mt |
cortex-a53 |
52.5 MB → 58.8 MB | |
| mobilenet_v2_1 RSS @ ready · pass |
cortex-a55 |
52.4 MB → 57.8 MB | |
| mobilenet_v2_1 RSS @ ready · pass_mt |
cortex-a55 |
52.5 MB → 57.6 MB | |
| arm_ml_kws_cnn_m load · pass |
cortex-a53 |
48 ms → 52 ms | |
| mobilenet_v2_1 RSS @ load · pass |
cortex-a55 |
67.5 MB → 71.8 MB | |
| mobilenet_v2_1 RSS @ load · pass_mt |
cortex-a53 |
67.7 MB → 71.8 MB | |
| mobilenet_v2_1 RSS @ load · pass_mt |
cortex-a55 |
67.5 MB → 71.3 MB | |
| arm_ml_kws_cnn_m load+optimize · pass |
cortex-a53 |
74 ms → 78 ms |
|
|
e181d95 to
3e44da0
Compare
`absorb_squeeze` refused any reshape that dropped the m or n axis, so a matmul whose n axis is the unit axis of a `[1, m]` output kept a separate `IntoShape`, which in turn blocked the successor bias add from fusing into the store. A pure squeeze only ever removes unit axes, and the store already reads a zero stride and an extent of 1 for a `None` axis, so drop the guard and let the squeezed axis become `None`, tracking it through the fused specs' own store specs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A `Mul` and an `Add` over the same accumulator each walked it end to end, reading and writing the whole tensor twice for arithmetic that fits in one traversal. Fuse the pair into `OptMulAddUnicast` when both unicast operands share a period, so one walk covers them together. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`move_axis` consumes its input and materialises the permuted result into a freshly allocated tensor, so cloning beforehand copied the data twice. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
|
3e44da0 to
247da09
Compare
|
Reference: 2026-08-28 morning nightly run (0d old) · full report → run Speed — evaltime · prefill · decode no inference-speed regressions Improvements
|
| Δ | metric | device | main → PR |
|---|---|---|---|
| arm_ml_kws_cnn_m load · pass |
beaglev-ahead |
75 ms → 104 ms | |
| arm_ml_kws_cnn_m load+optimize · pass |
beaglev-ahead |
111 ms → 142 ms | |
| arm_ml_kws_cnn_m load · pass |
cortex-a9 |
81 ms → 98 ms | |
| arm_ml_kws_cnn_m load+optimize · pass |
cortex-a9 |
127 ms → 145 ms | |
| mobilenet_v2_1 RSS @ ready · pass_mt |
cortex-a53 |
52.7 MB → 59.3 MB | |
| mobilenet_v2_1 RSS @ ready · pass_mt |
cortex-a55 |
52.7 MB → 58.4 MB | |
| mobilenet_v2_1 RSS @ ready · pass |
cortex-a55 |
53.1 MB → 58.4 MB | |
| arm_ml_kws_cnn_m load · pass |
cortex-a53 |
47 ms → 51 ms | |
| mobilenet_v2_1 RSS @ load · pass_mt |
cortex-a53 |
67.5 MB → 72.2 MB | |
| mobilenet_v2_1 RSS @ load · pass |
cortex-a55 |
67.6 MB → 71.3 MB | |
| mobilenet_v2_1 RSS @ load · pass_mt |
cortex-a55 |
67.9 MB → 71.3 MB |
2 participants
Three changes that cut copies and node count in the optimised graph. All are
pure plumbing — outputs are bit-identical on every model tested.
absorb_squeezerefused to drop a unit m or n axisOptMatMul::absorb_squeezebailed on any reshape that removed the m or n axis.But a pure squeeze only ever removes unit axes, and the store already handles a
Noneaxis:compute_stridesreads a zero stride for it and the extent fallsback to 1, which is exactly what a squeezed unit axis means. So the guard was
turning down reshapes it could have absorbed.
It costs two nodes, not one. A
[1, m]GEMV output reshaped to[m]kept itsIntoShape, and that node then sat between the matmul and its bias add,blocking the
OptBinUnicast(Add)→AddUnicast-in-the-store fusion that wouldotherwise have fired.
The guard is dropped, both axes now go through a
squeeze_c_axishelper, andrm_c_axistracks the change into the fused specs' own store specs — includingAddUnicast's, which previously was not updated at all (unreachable while theguard stood).
A unicast
Mulfeeding a unicastAddwalked the accumulator twicex * gthen+ breads and writes the whole tensor twice for arithmetic thatfits in one traversal.
OptBinUnicast::fusenow folds the pair intoOptMulAddUnicastwhen both unicast operands share a period, so a single walkcovers them. It is codegen-only, so nothing new has to serialise.
This is the affine tail of a normalisation and the gate update of a GRU cell —
it fires on
LayerNormalization.y_internalandgrucell/Add_2among others.MoveAxiscopied the tensor twiceAxisOp::change_tensor'sMovearm didtensor.clone().move_axis(..). Butmove_axisconsumes its input and goes throughpermute_axes→Tensor::from_datum,which allocates a fresh natural-strided tensor and copies into it. The clone was a
second, redundant full copy on every
MoveAxisnode.Unlike the two rules above this one is not shape-specific — it applies to any
transpose in any model, which is why
dtlnand the DeepFilterNet3 submodels moveat all here.
Measured
p50 ms/frame, median of 5 interleaved runs, single-threaded, aarch64, against
this PR's merge base.
The last two rows are noise, not regressions:
gtcrn_simpleis the same graph asgtcrn(+0.5%) anddf_decruns in 98 µs, where a 0.7 µs swing is a roundingartifact of the median.
The wins concentrate on DPDFNet, whose
Gemm-based GRU cells emit exactly the
[1, m]-GEMV-then-reshape-then-bias chainthe first rule unblocks. On
dpdfnet8_48khz_hrthat is 1901 nodes down to 1103,with
OptAddUnicast148 → 50 andIntoShape184 → 154. GTCRN and DTLN do notform the chain and land inside noise, as expected.
Outputs compared whole-stream, every frame: max absolute difference exactly
0on all nine models above, including the two that gain nothing.Full workspace suite green,
cargo fmtandcargo clippyclean. No newunsafe.🍍