Skip to content

perf: optimize topk with heap selection, k=1 linear scan, and contiguous fast path - #36

Merged
CodeWithKyrian merged 1 commit into
mainfrom
perf/topk-improvements
Jun 17, 2026
Merged

CodeWithKyrian merged 1 commit into
mainfrom
perf/topk-improvements

Conversation

@CodeWithKyrian

Copy link
Copy Markdown
Contributor

This PR replaces the O(n log n) full-sort topk algorithm with O(n log k) binary-heap selection for moderate k, an O(n) linear scan for k=1 (greedy sampling), and a zero-copy contiguous fast path that bypasses extract_array allocation. Flank refactoring extracts a shared is_c_contiguous helper and a define_extract_view macro to eliminate code duplication.

Motivation and Context

topk() was a bottleneck in ML inference pipelines — taking up to 15ms per call (37% of a 40ms model step) while other operations ran in 1-5ms. The algorithm sorted every lane in full (O(n log n)) even when only k=1 or k=50 elements were needed. Additionally, every call allocated and copied the entire array via extract_array regardless of whether the data was already contiguous.

For a flat 152k-element array with k=50 (a real downstream sampling scenario), the full sort did ~2.6M comparisons. The heap-based approach does ~152k comparisons — a 17x algorithmic improvement, translating to roughly 18x wall-clock speedup (11.3ms → 0.63ms).

What's Changed

  • Binary-heap topk selection (O(n log k)) replaces full sort when k × 4 < n
  • k=1 linear scan via lane iterator — O(n) with no per-element bounds checks, eliminating sort entirely for greedy decoding
  • Contiguous fast path: creates ArrayViewD directly from raw pointer instead of calling extract_array, eliminating one full-array allocation and copy per call
  • is_c_contiguous helper extracted to shared location, replacing 3 duplicate implementations
  • define_extract_view macro added — zero-copy immutable view extraction, symmetric with existing define_extract_view_mut
  • topk.rs simplified from 796 to 373 lines via topk_axis_arm/topk_flat_arm macros

Breaking Changes

None

@CodeWithKyrian
CodeWithKyrian merged commit 8d1bf3e into main Jun 17, 2026
14 checks passed
@CodeWithKyrian
CodeWithKyrian deleted the perf/topk-improvements branch June 17, 2026 09:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant