microsoft / microsoft/onnxruntime

CUDA ScatterElements reductions dispatch on element size, giving wrong results for several registered types

Open
#32,061 1 comment 1 reaction 2 assignees Claimed by @tianleiwu View on GitHub
ep:CUDA
Dominant language
C++
Stars
21.9k
Forks
4.2k
Avg merge
4d 11h
Merged PRs (30d)
184

Description

### Describe the issue

The CUDA `ScatterElements` kernel selects its compute type by **element size** rather than by the tensor's actual element type:

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cuda/tensor/scatter_elements.cc#L139-L145

```cpp
// Use element size instead of concrete types so we can specialize less template functions to reduce binary size.
int dtype = GetElementType(input_tensor->DataType()->Size());
...
utils::MLTypeCallDispatcher t_disp(dtype);
```

where `GetElementType` maps sizes to types:

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cuda/tensor/gather_elements.cc#L105-L118

```cpp
case sizeof(MLFloat16): return ONNX_NAMESPACE::TensorProto_DataType_FLOAT16;
case sizeof(float): return ONNX_NAMESPACE::TensorProto_DataType_FLOAT;
case sizeof(double): return ONNX_NAMESPACE::TensorProto_DataType_DOUBLE;
```

This is sound for `reduction='none'`, which only moves bytes. It is **not** sound once a reduction is requested, because the reduction functors perform typed arithmetic on the dispatched type:

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu#L100-L104

```cpp
template
struct FuncAdd {
__device__ __inline__ void operator()(T* start_addr, size_t index, T value) const {
atomic_add(start_addr + index, value);
}
};
```

So a `bfloat16` tensor is added as if its bits were `float16`, an `int32` tensor as if its bits were `float`, and so on. The kernel registers `DataTypeImpl::AllFixedSizeTensorTypes()` (which resolves to `element_type_lists::AllFixedSizeIRv4`, including `BFloat16`, `int16_t`, `uint16_t`, `int32_t`, `uint32_t`, `int64_t`, `uint64_t`), and `ComputeInternal` validates only shapes, ranks and that data/updates types match — there is no check that the type is one of the four the dispatcher can actually represent.

Affected combinations, for `reduction` in `add`/`mul`/`min`/`max`:

| tensor type | size | dispatched as | result |
|---|---|---|---|
| `BFloat16` | 2 | `MLFloat16` | wrong — fp16 arithmetic on bf16 bits |
| `int16_t`, `uint16_t` | 2 | `MLFloat16` | wrong — fp16 arithmetic on integer bits |
| `int32_t`, `uint32_t` | 4 | `float` | wrong — float arithmetic on integer bits |
| `int64_t`, `uint64_t` | 8 | `double` | wrong — double arithmetic on integer bits |
| `uint8_t` | 1 | `int8_t` | `add`/`mul` fine (two's complement), `min`/`max` wrong (signed comparison) |

`reduction='none'` is unaffected for all of these, and `float`/`double`/`MLFloat16`/`int8_t` are correct throughout.

The failure is silent — no error, just wrong numbers.

### Urgency

Not blocking me. Filing because it is a silent-wrong-answer path rather than a crash, so it is unlikely to be noticed by users who hit it.

I noticed this while adding `float16`/`bfloat16` support to the **CPU** `ScatterElements` reductions in #32025. That PR makes CPU compute bfloat16 reductions correctly, so once it lands CPU and CUDA will disagree for bfloat16 — CPU right, CUDA wrong. That is a consequence of this issue rather than of that PR, but it seemed worth mentioning so the two are not confused.

### Platform

Other

### Execution Provider

CUDA

### Possible fix

Dispatching on the concrete type would be the straightforward fix, at the cost of the binary-size saving the current comment is deliberately buying. A cheaper alternative that keeps that saving: keep the element-size dispatch for `reduction='none'`, and for the reduction paths either dispatch on the real type or reject the types that cannot be represented, so the failure is loud rather than silent.

I have not written a patch and have no CUDA hardware to test on, so I have deliberately not guessed at which trade-off you would prefer. Everything above is from reading the sources on `main`; I have not executed it.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.