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
remove
remove(mut self, idx: Int)
Remove element by global index, replacing with dead value.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!