Optimizations

[IanButterworth]

Note

This PR was written by Claude (Claude Code), from a performance review requested and supervised by @IanButterworth.

Follow-up to #131 (now merged; this branch has been rebased onto master). This PR reduces post-processing cost when many detections are kept — the stage that scales with detection count — and picks up the two forward-pass wins that survived measurement. Every commit was verified against the stored reference results for v2_COCO, v2_tiny_COCO, v3_tiny_COCO and v3_COCO (bit-identical for the post-processing commits; < 1e-6 for the batchnorm fold), plus the NMS unit tests and the soft-NMS fuzz check.

All numbers below were re-measured after the rebase, on current master vs this branch, same machine back-to-back (M2 Pro, CPU, single image, Julia 1.12, current deps incl. TimerOutputs v1).

End-to-end benchmark

Typical settings (detect_thresh=0.5, overlap_thresh=0.5, default allocator, min of 15):

Model	master	this PR	Δ
v3_416_COCO	336.7 ms	273.7 ms	−19%
v3_tiny_416_COCO	41.5 ms	32.9 ms	−21%
v4_416_COCO	589.3 ms	505.1 ms	−14%

Post-processing benchmark

Worst case: v3_416_COCO with all 10,647 candidates kept (detect_thresh=0.0, overlap_thresh=1.0, disallow_bumper=true), TimerOutputs stage breakdown (min of 5):

Stage	master	this PR
processing outputs	10.9 ms / 18.5 MiB	11.0 ms / 5.4 MiB
filter detections	730 µs / 9.1 MiB	190 µs / 4.5 MiB
nms	13.7 ms / 53.9 MiB	11.4 ms / 14.2 MiB
post-processing total	25.4 ms / 81.5 MiB	22.7 ms / 24.1 MiB

The 3.4× allocation reduction is the headline for dense scenes: in continuous-inference loops, post-processing GC pressure compounds.

Post-processing commits

1. Fuse post-processing broadcasts in place — the ~17 per-head transform statements were view[...] = RHS, each materializing the RHS array before copying; .= fuses them into single passes.
2. Restructure NMS: single global sort, allocation-free IoU indexing — the big one, two parts:
   • perform_detection_nms scanned all columns per batch (findall), then per distinct class (Set + findall), then ran a per-group sortperm + matrix copy through nested index-vector views. Now one stable sortperm by (batch, class, −score) plus one gather makes every (batch, class) group a contiguous, already-sorted column range.
   • nms! allocated an index slice per suppression round to build its IoU view — O(N²) bytes per group, the bulk of the old 54 MiB. bboxiou! gained a column-index variant so rounds read through a view of the persistent index vector with zero per-round allocation.
   • ⚠️ One visible change (also in the CHANGELOG): output columns are now ordered by batch, then ascending class id, then descending score; classes were previously in first-appearance order. Within-class ordering is unchanged (stable sort).
3. Flatten output blocks to attribute-major with a single allocation — extend_for_attributes (zeros + cat) followed by permutedims copied each head's data three times; now a single permutedims! writes straight into the first a rows of the final (a+4)-row buffer. GPU arrays keep the dense cat+permutedims path via a new fast_scalar_indexing trait (CUDAExt sets it false for CuArray).
4. Gather kept detections without materializing the concatenation — keepdetections needed cat(outweights...) plus a boolean-mask gather (two full copies); a new vector method counts and copies kept columns directly from the per-head matrices, with function barriers keeping the loops concretely typed. GPU falls back to the cat path via the same trait.

Forward pass

Measurement first (v3_416, CPU): the pass is gemm-bound, not allocation-bound — the default bumper path already recycles the ~890 MiB of forward allocations down to ~260 KiB with no time change, Julia threads do nothing (conv threading is BLAS-side: 1 BLAS thread = 842 ms, 6 = 317 ms), and AppleAccelerate drops v3_416 from ~317 ms to ~207 ms (−35%) with zero code changes. Two commits follow from this:

5. Fold batchnorm into conv weights at load time (darknet's fuse_conv_batchnorm) — BN at inference is a per-channel affine map, so it folds exactly into the kernel and bias, deleting the BatchNorm pass over every conv output. This is the bulk of the end-to-end wins in the first table. Max abs deviation vs stored refs < 1e-6, so no reference updates. Side benefit: the positional Flux.BatchNorm internal constructor — a compat hazard across Flux versions flagged in the original review — is no longer used.
6. README: CPU performance tips — AppleAccelerate on Apple silicon (~35%), MKL on Intel, BLAS threads as the parallelism knob, and batching for throughput.

Considered and deliberately not done

• The remaining ~11 ms in "processing outputs" is real compute (~900k sigmoid/exp evaluations); faster approximations would change numerics.
• The O(n²) IoU work when nothing gets suppressed is inherent to hard NMS at overlap_thresh=1.0; at realistic thresholds the loop shrinks as boxes are culled.
• In-place conv! into the preallocated W buffers: mostly moot since the pass isn't allocation-bound.
• A Metal.jl backend for Apple silicon is the big structural option (GPU-class speedup) but needs its own extension work — the post-processing kernels are CUDA-specific today.

Caveats

• The GPU fallback paths (trait-gated) are review-verified but untested on CUDA hardware, same as Fix yolov2 support and other bugs found in deep review #131. The batchnorm fold applies before weights move to the GPU, so the CUDA path gets it too (and one less layer per conv block).

🤖 Generated with Claude Code