Argmax

Experimental C++20 API, package version 0.1.0. Link SimaKernels::a65_argmax. The caller contract below describes supported behavior; implementation notes describe this release, not a stable dispatch or binary ABI promise.

Header: simaai/kernels/ops/argmax.h (A65/argmax/argmax.h). Typed request: ops::argmax::Request.

Contract

Property	Supported
Input	One logical tensor, rank 1–6
Dense or positive-stride input	INT8, UINT8, BF16, FP32
MLA input	INT8, UINT8, BF16; MlaTiled or MlaCBlocked
Output	One contiguous INT32 tensor containing axis indices
Parameters	axis (default -1), keepdims (default false)
Workspace	None
Execution	Synchronous, single caller thread, reusable immutable plan

Provide both input and output shapes. The output drops the reduction axis or sets it to one for keepdims=true. Reducing rank one without keepdims produces Shape{}: one scalar index, not an empty tensor. Negative axes are normalized against the logical rank; unused slots in six-dimension storage are not axes.

Equal values return the first logical index. Positive and negative zero compare equal; if NaNs occur, the first NaN wins, regardless of sign or payload. BF16 and FP32 comparison uses integer ordering of their IEEE encodings, without widening the tensor or depending on FP rounding, flush mode or default-NaN mode.

The reduction extent must be positive and its largest index must fit INT32. Zero extents on other axes give a legal, zero-output no-op. Positive strides must be scalar-aligned and pass the conservative non-overlapping-span check; negative, zero/broadcast and overlapping views are rejected. Output striding, FP16 input, flattened/all-axis reductions and INT64 output are not supported.

INT8/UINT8 compare raw stored values. There are no quantization parameters: common positive affine quantization preserves ordering; different scales or zero points across compared values do not necessarily do so. Handle that distinction before selecting this kernel.

Parameters

axis defaults to -1; keepdims defaults to false. Normalize a negative axis against the logical rank, not the six-slot capacity. There is no per-tensor or per-channel quantization parameter.

Numerical behavior

Open diagram at full size.

For {2,3,17} with axis=-1, the output is {2,3}; with keepdims=true, it is {2,3,1}. A rank-one reduction without keepdims yields Shape{}, which stores one scalar index—not zero elements.

Examples along one reduction axis:

Values	Winning logical index	Reason
[4,7,7]	1	First equal maximum
[-0,+0]	0	Signed zeros compare equal
[2,NaN,NaN,9]	1	First NaN

Changing padding bytes must not change any of these results. Quantized values with one shared positive affine scale preserve ordering; values with different per-channel scales or zero points need a different comparison contract.

The result is an INT32 logical index, not a value or physical-storage offset. Compare only logical elements. The first logical index wins ordinary ties, including +0 versus -0; the first NaN wins if any NaN is present. No arithmetic rounding or dequantization step is inserted. This behavior is defined independently of the physical layout and of the SIMD reduction tree.

C++ call

Place the includes and namespace aliases at file scope; the remaining statements belong in a simaai::kernels::Status-returning caller function. input_views and output_views are caller-supplied arrays of ConstHostBuffer and HostBuffer in port order, with accessible capacities and alignment from the checked plan. Keep their storage alive until synchronous execution returns.

#include <simaai/kernels/ops/argmax.h>
#include <array>

namespace k = simaai::kernels;
namespace op = k::ops::argmax;
const std::array input_specs{k::TensorSpec{{2, 3}, k::DType::BFloat16}};
const std::array output_specs{k::TensorSpec{{2}, k::DType::Int32}};
auto checked = op::check_inference(
    {input_specs, output_specs, op::Params{.axis = -1, .keepdims = false}});
if (const auto *error = std::get_if<k::Error>(&checked)) return *error;
const auto &plan = std::get<op::Plan>(checked);

return op::execute(plan, input_views, output_views);

See the correctness tests (A65/argmax/tests/correctness.cpp) for concrete storage, independent expected results and failure cases.

See calling kernels for the shared lifecycle.

Memory and lifecycle

For other layouts, allocate from plan.input_requirements() and plan.output_requirements(), not logical element count alone. Bindings contain CPU-accessible addresses and available byte spans; they neither own nor map the memory. Every call checks arity, accessible span arithmetic, minimum size, alignment and overlap before writing. The caller must supply genuinely accessible memory and retain its lifetime for the call; the kernel cannot validate mappings from a pointer alone. Runtime owns mapping, cache coherency and device fences. Distinct outputs may execute concurrently using one plan; concurrent calls must not race on the same output or modify inputs.

MLA storage

• MlaTiled uses the explicit D/H/W/C axis mapping and tile geometry. Outer axes remain independent; per-pixel BF16 low/high byte planes are read directly.
• MlaCBlocked uses trailing H/W/C, with earlier axes as consecutive volumes. Rank two means H/C and rank one means C. Each channel tile is independently padded to 16; within it, storage is C16-block/H/W/lane, with BF16 low/high byte stripes. Channel tiles such as 17 are legal and not treated as 32 logical values.
• These are different physical formats. Dtype still describes the scalar value; padding never participates in the reduction. Neither format is re-tessellated or materialized into a dense intermediate by Argmax.

Implementation notes

Preparation chooses a dtype/layout/axis implementation, not a generic per-value interpreter. Dense channel reductions use a SIMD maximum followed by a first winner search. Short axes 2–32 have compile-time specializations; tiny rows vectorize independent outputs. FP32 rows containing exactly two four-lane vectors use zero- and NaN-canonical integer keys, one unsigned maximum reduction and a first-element fast return. Other floating-point rows using the maximum-then-search route retain their maximum/minimum scans and scalar NaN/zero classification. Short specializations using this route, and dense rows longer than 32 elements for all four dtypes, pack each comparison mask into a scalar word to find its first matching lane. The search visits vectors in logical order, including a bounded, possibly overlapping final vector. It adds the lane offset to the full logical index, preserving the first maximum, first NaN and signed-zero ties. Long rows retain their existing extrema scan without a separate first-element fast return.

Dense non-channel reductions reuse cache lines across adjacent outputs. Their full 16-output groups classify complete BF16 or byte vectors before widening comparison keys; index accumulators remain 32-bit. Dense FP32 and BF16 independent-output groups use nonsaturating integer absolute value in their zero-canonical key mapping. This maps both zero signs equally without a separate zero comparison; the NaN classification remains explicit.

On AArch64, preparation also selects a neighboring-output SIMD route for noncanonical positive-stride input when the reduction axis is not the last logical axis, the last-axis byte stride equals the scalar width, and that axis contains at least 16 elements. The route handles 16 adjacent outputs per reduction-axis walk, followed by four-output and scalar tails. Other dimensions may have any validated positive strides, including a different physical order. It advances their logical coordinates once per complete last-axis row. Canonical contiguous views retain the dense route; singleton reductions return zero.

Pixel-tiled spatial reductions also use neighboring-output SIMD on AArch64 when C maps to the last logical axis, both its logical extent and channel tile width are at least 16, and the reduction axis maps to D, H or W. Groups of 16, four and one channel stay within each physical channel tile. Preparation supplies tile skips; execution decodes the first and last reduction-tile locations before reducing the group. This handles compact-channel alignment gaps, padded channels and permuted spatial mappings without reading padding.

Pixel-tiled channel reductions visit spatial tiles rather than decoding each output pixel. Prepared tile skips supply the full and final channel-tile bases; pixel addresses advance within each tile, and output strides omit the reduced logical axis. This supports any logical channel-axis position, spatial-axis mapping and outer dimensions. Channel runs retain ascending logical order and the existing first-winner reduction. Compact alignment gaps and channel padding are excluded from comparisons.

Other strided and pixel-tiled cases retain the generic path. It advances nonreduced coordinates with a mixed-radix counter in logical row-major order, rather than dividing a flat output index for each dimension. It resets the reduced-axis coordinate between outputs; singleton dimensions participate in carries and empty outputs never enter the traversal.

Channel-blocked inputs use physical channel, spatial or outer-volume walks, avoiding per-value coordinate decoding. On AArch64, spatial reductions process all 16 valid lanes of a full C16 stripe in one shared tile/axis walk, using four independent four-lane accumulators. Partial stripes use four-lane and scalar reductions without comparing padding. Both the strided neighboring-output route and the shared C16 walk retain 32-bit indices for every input dtype, including BF16 and byte inputs; they do not limit reduction indices to 16 bits. Only a strictly greater ordered value replaces a winner, preserving the first logical index across tiles and tails. Outer-volume reductions initialize each winner from logical axis index zero, then compare later elements using the same strict update rule. This preserves first-index ties, NaN precedence and signed-zero equality without comparing the first value against an artificial initial key.

BF16 and FP32 retain integer-only NaN and signed-zero comparisons. BF16 blocked paths combine byte planes in registers rather than detessellating a tensor. The BF16 blocked channel path can reuse up to eight run descriptors on the stack; its first-winner search packs comparison masks without changing the preceding extrema scan or first-element fast return.

For INT8/UINT8 channel-blocked reductions on AArch64, preparation selects a separate descriptor-based implementation only when the channel tiles contain at most eight physical C16 stripes in total and every active width tile contains at least two pixels. Execution builds the bounded descriptors once per spatial tile row and reuses them across its pixels. Descriptors retain ascending logical channel order, carry each stripe's valid extent and preserve the existing byte comparison and first-winner index updates. Padding does not participate. Plans with more than eight stripes, width-one tiles or any width-one final tail use direct channel/stripe traversal for the entire execution. Route selection occurs during preparation, not in the execution loop.

These fixed, bounded stack descriptors are not caller workspace. There is no heap allocation or tensor-sized intermediate during execution.

See build and test for default correctness and latency tests.

Tests and qualification

The default build includes a65_argmax_contract and a65_argmax_benchmark.

The independent correctness suite is in A65/argmax/tests/correctness.cpp (A65/argmax/tests/correctness.cpp). See build and test for commands and the distinction between host development checks, A65 measurements and baseline qualification.