Mojo struct

TopKHeap

struct TopKHeap[T: DType, largest: Bool, M: Int]

Fixed-capacity register heap for per-thread top-M tracking.

Stores up to M (value, index) pairs in registers. During the scan phase, a cached threshold provides O(1) rejection of non-competitive elements. All internal loops are compile-time unrolled to keep data in registers on GPU. Indices are stored as Int32 to reduce register pressure for large block sizes.

Fields

vals (InlineArray[Scalar[T], M]):
idxs (InlineArray[Int32, M]):
threshold (Scalar[T]):

Implemented traits

AnyType, ImplicitlyDestructible

Methods

`init`

__init__(out self)

`insert`

insert(mut self, val: Scalar[T], idx: Int)

Insert an element, evicting the worst if full.

`best`

best(self) -> TopK_2[T, largest]

Return the best element, ties broken by smallest index.

Returns a dead TopK_2 (p=-1) when all entries are exhausted.

Returns:

TopK_2[T, largest]

`remove`

remove(mut self, idx: Int)

Remove element by global index, replacing with dead value.

Fields​

Implemented traits​

Methods​

__init__​

insert​

best​

remove​