For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
pick_simd_width
def pick_simd_width[M: ReduceOp, target: StringSpan[ImmStaticOrigin] = StringSpan("cpu"), PER_THREAD_BUDGET: Int = Int(64), *Ts: DType = *?]() -> Int
Picks the SIMD width a rowwise body should use throughout.
Mirrors elementwise's "pick the smallest dtype's natural SIMD
width" rule: the widest register that fits the smallest element,
with any wider-element dtype expressed via paired SIMD ops at no
perf cost. Computed as max(simd_width_of[T] for T in Ts).
On GPU only, caps the result down when M's storage at W=1
(sizeof[M]) times the max SIMD width would exceed
PER_THREAD_BUDGET bytes per thread. This protects heavy-state
monoids (ArgMax/ArgMin with int64 indices, OnlineLogSumExp /
future Welford with multi-field state) from the GPU BLOCK_SIZE
multiplier blowing past the register budget. Single-field monoids
(Sum/Max/Min/Product) almost always stay under the cap.
Why M at W=1: bodies declare states like M[dtype, 1] for the
budget check; the helper assumes state size scales roughly linearly
with W (true for SIMD field accumulators). Bodies wanting a
sub-1 width or a custom rule can compute W themselves.
Parameters:
- M (
ReduceOp): The monoid type atW=1(sizeofproxy for the GPU storage-budget check). - target (
StringSpan[ImmStaticOrigin]):"cpu"or"gpu"(anything non-CPU = GPU). - PER_THREAD_BUDGET (
Int): GPU per-thread storage cap in bytes. Tunable; default 64B is an empirically reasonable register-pressure ceiling on B200-class GPUs. - *Ts (
DType): Variadic pack of dtypes involved in the body (input, accumulator, output, gamma weights, ...); the helper takes the max natural SIMD width across them.
Returns:
Int: The SIMD width for the monoid M[dtype, W], simd_width=W in
ContextParams, and tile_fn[ws] calls. Always at least 1.