#28014 Harden OneHot operator input validation and output size computation

Harden OneHot operator input validation and output size computation

Summary

This PR tightens input validation and output-shape computation for the OneHot operator on the CPU and CUDA execution providers so that invalid or extreme inputs fail cleanly with an INVALID_ARGUMENT status instead of triggering overflow, silent truncation, or division by zero.

Changes

1. Overflow check in PrepareOutputShape using SafeInt

   • The output tensor size computation now uses a checked multiplication over the output dimensions, and the prefix_dim_size multiplication loop is rewritten with SafeInt<int64_t>. This prevents unbounded allocation attempts when a large depth value (or large indices shape) would overflow int64_t.

2. Guard against division by zero when prefix_dim_size is zero

   • suffix_dim_size is now computed as (prefix_dim_size > 0) ? (indices_shape.Size() / prefix_dim_size) : 0, avoiding an integer division by zero when an indices dimension before the axis is zero.

3. CUDA int32 range validation before fast_divmod

   • Before narrowing to int for fast_divmod (which requires int32 operands), the kernel now validates that suffix_dim_size and depth_val * suffix_dim_size fit in int32_t. Previously gsl::narrow_cast<int> could silently truncate, producing wrong divmod operands.

4. nullptr check on Output() in both CPU and CUDA Compute paths

   • If allocation of the output tensor fails (e.g., because the requested shape is too large for the allocator), both kernels now return a descriptive error instead of dereferencing a null pointer.

5. Unit tests in onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc:

   • DepthTooLarge_OutputSizeOverflow — depth = INT64_MAX, indices = [2,3].
   • DepthTooLarge_OutputSizeOverflow_LargeIndices — depth = INT64_MAX / 500, indices = [1000].
   • NegativeDepth — negative depth is rejected.
   • DepthOne — minimum valid depth = 1 edge case.
   • ScalarIndicesRejected — rank-0 indices are rejected per the ONNX spec (indices rank ≥ 1).
   • DefaultAxis_Opset9 — opset 9 coverage for the default-axis path.
   • The three negative-path tests exclude kTensorrtExecutionProvider and kDmlExecutionProvider because those EPs fail with their own (different) error messages before our kernel validation runs.

Files touched

onnxruntime/core/providers/cpu/tensor/onehot.cc        | 28 +++++++-
onnxruntime/core/providers/cuda/tensor/onehot.cc       | 13 ++++
onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc | 83 ++++++++++++++++++++++

Testing

Built onnxruntime_provider_test on Windows (Release) and ran OneHotOpTest.*:

[==========] 34 tests from 1 test suite ran.
[  PASSED  ] 34 tests.

Overflow tests produce the expected error, for example:

OneHot: output tensor size would overflow for the given indices shape and depth value (9223372036854775807).