Support float16 and bfloat16 in CPU ScatterElements reductions

[Alexander Novikov (novikov-alexander)]

Description

ScatterElements with reduction='add' or 'mul' throws ORT_NOT_IMPLEMENTED on the CPU EP for MLFloat16 and BFloat16, and 'min'/'max' throws for BFloat16. This implements them.

The stubs exist for a mechanical reason: the generic functors use compound assignment ((*a) += (*b)), and neither half type defines operator+= or operator*=. So add and mul are now evaluated in float and rounded back to half on each update:

template <>
struct Func_Add<MLFloat16> {
  void operator()(MLFloat16* a, const MLFloat16* b) const {
    *a = MLFloat16(a->ToFloat() + b->ToFloat());
  }
};

Func_Min<BFloat16> and Func_Max<BFloat16> are removed rather than implemented. Nothing forced those two stubs — BFloat16 already has the comparison operators the generic functors use, so it now falls through to the generic path exactly as MLFloat16 has been doing all along (MinReduction_MLFloat16 and MaxReduction_MLFloat16 have been passing on the CPU EP unmodified).

Accumulation precision: rounding to half after every update, not accumulating in float and rounding once. This matches the other backends that implement these reductions — CUDA (cuda/atomic/common.cuh, native atomicAdd(half) on sm_70+ and a CAS loop that widens only for the single pairwise add below that) and WebGPU (webgpu/tensor/scatter_elements.cc, oldF16 + value in f16). It is also the natural fit for a functor that receives one element pair at a time.

No kernel registration or type-constraint changes, so docs/OperatorKernels.md does not need regenerating.

Motivation and Context

docs/OperatorKernels.md already lists tensor(float16) and tensor(bfloat16) for CPU ScatterElements, because the table is generated from the registrations and both types are in element_type_lists::All. The op therefore advertised support it did not have: a float16 model using reduction='add' fails at inference time with NOT_IMPLEMENTED rather than falling back to anything.

This is the shape most affected in practice — scatter-based pooling and embedding-style graphs that run in half precision end up needing a float32 cast around the scatter purely to work around the CPU kernel.

Verification

Built and ran locally on macOS/arm64:

• onnxruntime_provider_test --gtest_filter="ScatterElements.*" — 18/18 pass (10 pre-existing, 8 new/enabled).
• All 8 new tests fail against the unmodified kernel and pass with it — verified by reverting only scatter.cc, rebuilding, and re-running: exactly those 8 fail while the 10 pre-existing tests still pass. None of them pass vacuously.
• Full onnxruntime_provider_test sweep: 5644 tests, 5456 passed, 0 failures — unchanged from baseline. This includes GatherElements, which shares ScatterData via GatherElementsGradImpl.

Tests

AddReduction_MLFloat16 already existed but was compiled out behind #if defined(CUDA_VERSION) with the comment "Operation on float16 (MLFloat16) is not implemented on CPU". That guard is removed and the test now runs.

Added: fp16 mul and axis-1 add; bf16 add, mul, min and max; and AddReduction_MLFloat16_RoundsAfterEachUpdate, which pins the accumulation semantics — it accumulates eight 0.25 updates onto 1024.0, where one ULP is 1.0 in binary16, so per-update rounding leaves 1024 while a float accumulator rounded once at the end would produce 1026. The gap is far outside the fp16 comparison tolerance, so the test genuinely discriminates between the two designs rather than merely passing.

Note for reviewers: the bfloat16 mul expectations use smaller factors than the float test, because bfloat16 keeps only 8 bits of significand and -343 would round to -344.

One thing worth flagging separately

While checking cross-backend behaviour I noticed the CUDA kernel dispatches ScatterElements on element size (gather_elements.cc, GetElementType), which maps any 2-byte type to FLOAT16. CUDA therefore appears to compute bfloat16 ScatterElements as float16. That is out of scope here and I have not changed it, but with this PR the CPU and CUDA results for bfloat16 would differ. Happy to open a separate issue if that is useful.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Alexander Novikov (novikov-alexander)]

Pushed ff3edae, scoping five of the new tests to the CPU EP before CI gets to them.

Full CI has not run on this PR yet — it currently shows a single check — but it did run on #32026, and it failed there for exactly this reason: a test that asserted a property of the CPU kernel rather than of the operator was running on every provider. The same two patterns are present here, so I have fixed them rather than waiting.

• AddReduction_MLFloat16_RoundsAfterEachUpdate pins the accumulation precision. ONNX does not specify the intermediate precision, so a provider that accumulated in float and rounded once at the end would be equally correct and would produce 1026 rather than 1024. That is a property of this kernel, so the test now runs on CPU only.
• The four bfloat16 tests would fail on the CUDA plugin EP for a reason unrelated to this PR: that kernel picks its compute type by element size, so bfloat16 is treated as float16 and add/mul are computed on misread bits. I filed that separately as CUDA ScatterElements reductions dispatch on element size, giving wrong results for several registered types #32061. I scoped min and max too, conservatively — the bit patterns of both formats are monotonic for same-sign values so those two may well survive the misinterpretation, but I have no CUDA hardware to confirm it and would rather not find out from CI. Worth widening once CUDA ScatterElements reductions dispatch on element size, giving wrong results for several registered types #32061 is fixed.

The float16 add/mul tests are unchanged and still run on every provider.

I have applied the same scoping to the bfloat16 tests in #32065, which adds the equivalent support to ScatterND.