IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo struct

ArgMax

struct ArgMax[dtype: DType, W: Int = simd_width_of[dtype]()]

Argmax reduction monoid: tracks the (value, axis-index) pair with the largest value. Ties break to the lower index.

Per-tile work is two selects into SIMD-wide acc_values/acc_indices (no per-tile horizontal reduce); the horizontal reduce + lane-find happens once at reduce. Partial tiles get identity-padded via Self.pad — -inf (or MIN) for values, Int64.MAX for indices — so padded lanes never win the tie-break.

W defaults to 1 on GPU because ArgMax storage scales as (value_size + 8) * W per thread (int64 indices), and GPU's BLOCK_SIZE multiplier turns large W into severe register pressure. CPU defaults to simd_width_of[dtype] for full bulk-path SIMD parallelism.

Parameters​

  • ​dtype (DType): The value dtype.
  • ​W (Int): SIMD width of the per-lane accumulator. Defaults to 1 on GPU (low register pressure) and simd_width_of[dtype] on CPU (full SIMD parallelism).

Fields​

  • ​best (Scalar[dtype]): Final scalar best value. Body reads this after join_parallel.
  • ​best_idx (Int): Final scalar best index. Body reads this after join_parallel.
  • ​acc_values (SIMD[dtype, W]): Per-lane running max during SIMD-wide tile accumulation.
  • ​acc_indices (SIMD[DType.int64, W]): Per-lane axis indices corresponding to acc_values.

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, ReduceOp, RegisterPassable, TrivialRegisterPassable

comptime members​

Single​

comptime Single = ArgMax[dtype, Int(1)]

width​

comptime width = W

Methods​

__init__​

def __init__() -> Self

Identity: scalars at -inf/MIN and Int.MAX; SIMD accs at the same identities. First real (value, idx) always wins under the lower-idx tie-break.

__getitem__​

def __getitem__(self, j: Int) -> Self.Single

Returns lane j as a width-1 monoid.

Returns:

Self.Single

__setitem__​

def __setitem__(mut self, j: Int, s: ArgMax[dtype, Int(1)])

Writes width-1 monoid s into lane j.

accumulate​

def accumulate[val_dtype: DType, w: Int](mut self, val: SIMD[val_dtype, w], idx: SIMD[DType.int64, w] = 0)

Folds a SIMD tile via lane-wise compare-and-select. val and idx are padded to width Self.W via Self.pad — padded value lanes -inf/MIN, padded index lanes Int64.MAX — so they always lose to any real candidate.

le (not lt) preserves the lower-idx tie-break: a tied value with a higher idx keeps the current (lower-idx) acc. idx is the scaffolder-built per-lane axis-position vector — the monoid never adds an iota.

Parameters:

  • ​val_dtype (DType): The tile's source dtype.
  • ​w (Int): SIMD width of the tile.

Args:

join​

def join(mut self, other: Self)

Sequential combine. Element-wise SIMD merge of the acc; scalar merge of the (best, best_idx) tail.

Tie-symmetric: when values tie, the smaller index wins regardless of which side it came from — required for correctness under cross-thread tree reduction, where the pair-ordering isn't index-ordered.

Args:

  • ​other (Self): The state to combine into self.

reduce​

def reduce(self) -> Self.Single

SIMD-tree collapse (faster than the default lane-fold for a value+index select): max value in lane 0, lowest index among the max lanes. join_parallel folds lane 0 into (best, best_idx).

Returns:

Self.Single

join_parallel​

def join_parallel[R: Reducer](mut self, reducer: R)

Folds the within-thread acc at lane 0 into the scalar (best, best_idx), resets the SIMD acc so the cross-thread join can't double-count it, combines (best, best_idx) across participants via reducer.generic, and leaves the winning index in acc_indices[0] so the body's emit reads acc_indices.slice[tile_width] uniformly (cooperative: lane 0; tiled: per-column, where join_parallel is skipped).

At W == 1 lane 0 is the running argmax directly; at W > 1 reduce reduced it there first.

Parameters:

  • ​R (Reducer): The parallel reducer.

Args:

  • ​reducer (R): The reducer instance.

Models

Cookbook

Puzzles

Blog

Community

Contact