From 57d4b03523d2908ea78944916af57720bf758b99 Mon Sep 17 00:00:00 2001 From: Gopalakrishnan Nallasamy Date: Wed, 8 Apr 2026 16:05:25 -0700 Subject: [PATCH 1/5] Harden OneHot operator input validation and output size computation - Add overflow check in PrepareOutputShape using SafeInt for output size and prefix_dim_size multiplication to prevent unbounded allocation when depth or indices shape would overflow int64 - Guard against division by zero when prefix_dim_size is zero - Add CUDA int32 range validation before fast_divmod to avoid silent truncation in gsl::narrow_cast for suffix_dim_size and depth_val * suffix_dim_size - Check for nullptr from Output() in both CPU and CUDA Compute paths - Add unit tests: depth overflow (two variants), negative depth, depth=1 edge case, scalar-indices rejection (ONNX spec requires rank>=1), and opset 9 coverage --- .../core/providers/cpu/tensor/onehot.cc | 28 ++++++- .../core/providers/cuda/tensor/onehot.cc | 13 +++ .../providers/cpu/tensor/onehot_op_test.cc | 83 +++++++++++++++++++ 3 files changed, 121 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/onehot.cc b/onnxruntime/core/providers/cpu/tensor/onehot.cc index 91d6ada1fb864..a7310d7f29253 100644 --- a/onnxruntime/core/providers/cpu/tensor/onehot.cc +++ b/onnxruntime/core/providers/cpu/tensor/onehot.cc @@ -16,9 +16,12 @@ limitations under the License. #include "core/providers/cpu/tensor/onehot.h" #include "core/common/eigen_common_wrapper.h" +#include "core/common/safeint.h" #include "core/platform/env.h" #include "core/providers/common.h" +#include + #ifndef EIGEN_USE_THREADS #define EIGEN_USE_THREADS #endif @@ -100,11 +103,29 @@ Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val, const output_shape.insert(output_shape.begin() + true_axis, depth_val); - prefix_dim_size = 1; + // Validate that the total output tensor element count does not overflow int64. + { + int64_t total_elements = 1; + for (auto dim : output_shape) { + if (dim > 0 && total_elements > std::numeric_limits::max() / dim) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "OneHot: output tensor size would overflow for the given indices shape " + "and depth value (", + depth_val, ")."); + } + total_elements *= dim; + } + } + + // Use SafeInt for prefix_dim_size computation to guard against overflow. + SafeInt safe_prefix = 1; for (int64_t i = 0; i < true_axis; ++i) { - prefix_dim_size *= indices_dims[onnxruntime::narrow(i)]; + safe_prefix *= indices_dims[onnxruntime::narrow(i)]; } - suffix_dim_size = indices_shape.Size() / prefix_dim_size; + prefix_dim_size = safe_prefix; + + // Guard against division by zero when indices have a zero-sized dimension before the axis. + suffix_dim_size = (prefix_dim_size > 0) ? (indices_shape.Size() / prefix_dim_size) : 0; return Status::OK(); } @@ -166,6 +187,7 @@ Status OneHotOp::Compute(OpKernelContext* p_op_ke // allocate output const auto* values_data = values->Data(); Tensor* output = p_op_kernel_context->Output(0, TensorShape(output_shape)); + ORT_RETURN_IF_NOT(output, "OneHot: failed to allocate output tensor. Output shape may be too large."); // edge case where we have a dim with a value of 0 if (output->Shape().Size() == 0) diff --git a/onnxruntime/core/providers/cuda/tensor/onehot.cc b/onnxruntime/core/providers/cuda/tensor/onehot.cc index e5748e7f22417..7af58fdbb18a7 100644 --- a/onnxruntime/core/providers/cuda/tensor/onehot.cc +++ b/onnxruntime/core/providers/cuda/tensor/onehot.cc @@ -3,6 +3,8 @@ #include "core/providers/cuda/tensor/onehot.h" +#include + using namespace onnxruntime::common; namespace onnxruntime { @@ -55,11 +57,22 @@ Status OneHotOp::ComputeInternal(OpKernelContext* // allocate output const auto* values_data = reinterpret_cast(values->Data()); Tensor* output = ctx->Output(0, TensorShape(output_shape)); + ORT_RETURN_IF_NOT(output, "OneHot: failed to allocate output tensor. Output shape may be too large."); // edge case where we have a dim with a value of 0 if (output->Shape().Size() == 0) return Status::OK(); + // Validate that dimensions used by CUDA kernels fit in int32 range. + // fast_divmod requires int32 operands. + constexpr int64_t kInt32Max = std::numeric_limits::max(); + ORT_RETURN_IF_NOT(suffix_dim_size <= kInt32Max, + "OneHot: suffix dimension size (", suffix_dim_size, + ") exceeds int32 range supported by the CUDA kernel."); + ORT_RETURN_IF_NOT(depth_val <= kInt32Max / std::max(suffix_dim_size, int64_t{1}), + "OneHot: depth (", depth_val, ") * suffix dimension size (", suffix_dim_size, + ") exceeds int32 range supported by the CUDA kernel."); + const fast_divmod fdm_suffix(gsl::narrow_cast(suffix_dim_size)); const auto* indices_data = indices->Data(); auto* output_data = reinterpret_cast(output->MutableData()); diff --git a/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc index 55c247e4c2fea..a9b7566e82e97 100644 --- a/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" #include "test/common/trt_op_test_utils.h" @@ -499,6 +501,87 @@ TEST(OneHotOpTest, DimWithZero) { test.Run(); } +// Test that extremely large depth values that would cause output tensor size overflow are rejected. +TEST(OneHotOpTest, DepthTooLarge_OutputSizeOverflow) { + OpTester test("OneHot", 11); + // indices shape [2, 3] with depth = INT64_MAX causes output shape [2, 3, INT64_MAX] + // which would overflow when computing total element count. + test.AddInput("indices", {2, 3}, {1, 2, 3, 4, 5, 6}); + test.AddInput("depth", {1}, {std::numeric_limits::max()}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {2, 3, 1}, {0, 0, 0, 0, 0, 0}); + // Exclude TensorRT and DML EPs: they fail internally on INT64_MAX depth before our kernel's + // validation runs, producing a different error message. + test.Run(OpTester::ExpectResult::kExpectFailure, "output tensor size would overflow", + {kTensorrtExecutionProvider, kDmlExecutionProvider}); +} + +// Test that a very large depth value that overflows with multi-dimensional indices is rejected. +TEST(OneHotOpTest, DepthTooLarge_OutputSizeOverflow_LargeIndices) { + OpTester test("OneHot", 11); + // indices shape [1000] with depth = INT64_MAX / 500 causes overflow in element count. + const int64_t large_depth = std::numeric_limits::max() / 500; + std::vector indices(1000, 0); + std::vector dummy_output(1000, 0); + test.AddInput("indices", {1000}, indices); + test.AddInput("depth", {1}, {large_depth}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {1000, 1}, dummy_output); + // Exclude TensorRT and DML EPs: they fail internally on overflow-inducing depth before our + // kernel's validation runs. + test.Run(OpTester::ExpectResult::kExpectFailure, "output tensor size would overflow", + {kTensorrtExecutionProvider, kDmlExecutionProvider}); +} + +// Test that a negative depth value is rejected. +TEST(OneHotOpTest, NegativeDepth) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1, 2, 3, 4, 5, 6}); + test.AddInput("depth", {1}, {-5}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {2, 3, 1}, {0, 0, 0, 0, 0, 0}); + // Exclude TensorRT and DML EPs: they reject negative depth with their own error messages rather + // than ours. + test.Run(OpTester::ExpectResult::kExpectFailure, "Depth is negative", + {kTensorrtExecutionProvider, kDmlExecutionProvider}); +} + +// Test minimum valid depth value of 1. +TEST(OneHotOpTest, DepthOne) { + OpTester test("OneHot", 11); + test.AddInput("indices", {3}, {0, 0, 0}); + test.AddInput("depth", {1}, {1}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {3, 1}, {1, 1, 1}); + test.Run(); +} + +// Test scalar (rank-0) indices are rejected per ONNX spec (indices must have rank >= 1). +TEST(OneHotOpTest, ScalarIndicesRejected) { + OpTester test("OneHot", 11); + test.AddInput("indices", {}, {2}); + test.AddInput("depth", {1}, {5}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {5}, {0, 0, 1, 0, 0}); + test.Run(OpTester::ExpectResult::kExpectFailure, "Indices tensor must have rank >= 1"); +} + +// Test with opset 9. +TEST(OneHotOpTest, DefaultAxis_Opset9) { + OpTester test("OneHot", 9); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {2, 3, 10}, + {0, 1, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, + 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, + 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 1, 0, 0, 0}); + test.Run(); +} + #ifdef USE_CUDA TEST(OneHotOpTest, DefaultAxis_int64_MLFloat16_int64 /*indices, output, depth*/) { From d11a8b415d026c92d0634f4a996a8e681ff08bc6 Mon Sep 17 00:00:00 2001 From: Gopalakrishnan Nallasamy Date: Fri, 1 May 2026 17:39:16 -0700 Subject: [PATCH 2/5] Address Copilot review: rank check, plugin shim parity, missing includes - Reject rank-0 indices in PrepareOutputShape (CPU and CUDA plugin shim) per ONNX spec. - Mirror overflow / SafeInt prefix / div-by-zero guards in the CUDA plugin shim PrepareOutputShape. - Add include in CUDA onehot.cc for std::max. - Add core/common/safeint.h include in cuda_kernel_adapter.h for SafeInt. - Loosen ScalarIndicesRejected expected substring to match both ONNX shape-inference and kernel-level errors. --- .../core/providers/cpu/tensor/onehot.cc | 7 ++++ .../cuda/plugin/cuda_kernel_adapter.h | 33 +++++++++++++++++-- .../core/providers/cuda/tensor/onehot.cc | 1 + .../providers/cpu/tensor/onehot_op_test.cc | 4 ++- 4 files changed, 41 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/onehot.cc b/onnxruntime/core/providers/cpu/tensor/onehot.cc index a7310d7f29253..a6541ea3cde13 100644 --- a/onnxruntime/core/providers/cpu/tensor/onehot.cc +++ b/onnxruntime/core/providers/cpu/tensor/onehot.cc @@ -94,6 +94,13 @@ Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val, const const auto& indices_shape = indices->Shape(); const auto indices_dims = indices_shape.GetDims(); const auto indices_num_dims = indices_shape.NumDimensions(); + + // ONNX spec requires indices to have rank >= 1. + if (indices_num_dims == 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "OneHot: indices tensor must have rank >= 1."); + } + output_shape = indices_shape.AsShapeVector(); // output rank is always 1 more than the input rank as a new dimension is added to the input shape diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index b1da0aa816a03..ad1fb99b19298 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -17,6 +17,7 @@ #include "core/common/status.h" #include "core/common/narrow.h" +#include "core/common/safeint.h" #include "core/common/float16.h" #include "core/common/float8.h" #include "core/framework/float4.h" @@ -769,15 +770,41 @@ inline Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val, const auto& indices_shape = indices->Shape(); const auto indices_dims = indices_shape.GetDims(); const auto indices_num_dims = indices_shape.NumDimensions(); + + // ONNX spec requires indices to have rank >= 1. + if (indices_num_dims == 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "OneHot: indices tensor must have rank >= 1."); + } + output_shape = indices_shape.AsShapeVector(); const auto output_rank = static_cast(indices_num_dims) + 1; auto true_axis = HandleNegativeAxis(axis, output_rank); output_shape.insert(output_shape.begin() + true_axis, depth_val); - prefix_dim_size = 1; + + // Validate that the total output tensor element count does not overflow int64. + { + int64_t total_elements = 1; + for (auto dim : output_shape) { + if (dim > 0 && total_elements > std::numeric_limits::max() / dim) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "OneHot: output tensor size would overflow for the given indices shape " + "and depth value (", + depth_val, ")."); + } + total_elements *= dim; + } + } + + // Use SafeInt for prefix_dim_size to guard against overflow. + SafeInt safe_prefix = 1; for (int64_t i = 0; i < true_axis; ++i) { - prefix_dim_size *= indices_dims[narrow(i)]; + safe_prefix *= indices_dims[narrow(i)]; } - suffix_dim_size = indices_shape.Size() / prefix_dim_size; + prefix_dim_size = safe_prefix; + + // Guard against division by zero when indices have a zero-sized dimension before the axis. + suffix_dim_size = (prefix_dim_size > 0) ? (indices_shape.Size() / prefix_dim_size) : 0; return Status::OK(); } diff --git a/onnxruntime/core/providers/cuda/tensor/onehot.cc b/onnxruntime/core/providers/cuda/tensor/onehot.cc index 7af58fdbb18a7..eaf78ed1ea4e8 100644 --- a/onnxruntime/core/providers/cuda/tensor/onehot.cc +++ b/onnxruntime/core/providers/cuda/tensor/onehot.cc @@ -3,6 +3,7 @@ #include "core/providers/cuda/tensor/onehot.h" +#include #include using namespace onnxruntime::common; diff --git a/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc index a9b7566e82e97..ed449b96eacca 100644 --- a/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc @@ -563,7 +563,9 @@ TEST(OneHotOpTest, ScalarIndicesRejected) { test.AddInput("depth", {1}, {5}); test.AddInput("values", {2}, {0, 1}); test.AddOutput("output", {5}, {0, 0, 1, 0, 0}); - test.Run(OpTester::ExpectResult::kExpectFailure, "Indices tensor must have rank >= 1"); + // Match either the ONNX shape-inference error ("Indices tensor must have rank >= 1") or the + // explicit kernel-level rejection ("OneHot: indices tensor must have rank >= 1."). + test.Run(OpTester::ExpectResult::kExpectFailure, "ndices tensor must have rank >= 1"); } // Test with opset 9. From 14d5f4b4f08e8256ac80d3a5492cf437d914a9d6 Mon Sep 17 00:00:00 2001 From: Gopalakrishnan Nallasamy Date: Thu, 7 May 2026 08:52:07 -0700 Subject: [PATCH 3/5] Fix cpplint include order in cpu/onehot.cc --- onnxruntime/core/providers/cpu/tensor/onehot.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/onehot.cc b/onnxruntime/core/providers/cpu/tensor/onehot.cc index a6541ea3cde13..a31c91268ee1b 100644 --- a/onnxruntime/core/providers/cpu/tensor/onehot.cc +++ b/onnxruntime/core/providers/cpu/tensor/onehot.cc @@ -15,13 +15,14 @@ limitations under the License. /* Modifications Copyright (c) Microsoft. */ #include "core/providers/cpu/tensor/onehot.h" + +#include + #include "core/common/eigen_common_wrapper.h" #include "core/common/safeint.h" #include "core/platform/env.h" #include "core/providers/common.h" -#include - #ifndef EIGEN_USE_THREADS #define EIGEN_USE_THREADS #endif From b865764d78fab66b28d7939948b8c30cfd45216c Mon Sep 17 00:00:00 2001 From: Gopalakrishnan Nallasamy Date: Thu, 7 May 2026 11:43:04 -0700 Subject: [PATCH 4/5] Address review: defer depth*suffix int32 check; clarify SafeInt intent --- onnxruntime/core/providers/cpu/tensor/onehot.cc | 2 ++ .../core/providers/cuda/plugin/cuda_kernel_adapter.h | 2 ++ onnxruntime/core/providers/cuda/tensor/onehot.cc | 11 ++++++----- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/onehot.cc b/onnxruntime/core/providers/cpu/tensor/onehot.cc index a31c91268ee1b..c6b42c2f560a9 100644 --- a/onnxruntime/core/providers/cpu/tensor/onehot.cc +++ b/onnxruntime/core/providers/cpu/tensor/onehot.cc @@ -126,6 +126,8 @@ Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val, const } // Use SafeInt for prefix_dim_size computation to guard against overflow. + // SafeInt is defensive here -- the total-element overflow check above already covers this case, + // so a SafeIntException should never fire in practice. SafeInt safe_prefix = 1; for (int64_t i = 0; i < true_axis; ++i) { safe_prefix *= indices_dims[onnxruntime::narrow(i)]; diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index ad1fb99b19298..3be5effbfcf02 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -797,6 +797,8 @@ inline Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val, } // Use SafeInt for prefix_dim_size to guard against overflow. + // SafeInt is defensive here -- the total-element overflow check above already covers this case, + // so a SafeIntException should never fire in practice. SafeInt safe_prefix = 1; for (int64_t i = 0; i < true_axis; ++i) { safe_prefix *= indices_dims[narrow(i)]; diff --git a/onnxruntime/core/providers/cuda/tensor/onehot.cc b/onnxruntime/core/providers/cuda/tensor/onehot.cc index eaf78ed1ea4e8..ef46a20c22df7 100644 --- a/onnxruntime/core/providers/cuda/tensor/onehot.cc +++ b/onnxruntime/core/providers/cuda/tensor/onehot.cc @@ -64,15 +64,12 @@ Status OneHotOp::ComputeInternal(OpKernelContext* if (output->Shape().Size() == 0) return Status::OK(); - // Validate that dimensions used by CUDA kernels fit in int32 range. - // fast_divmod requires int32 operands. + // Validate that suffix_dim_size fits in int32 range. fast_divmod requires int32 operands + // and fdm_suffix is constructed on every code path below. constexpr int64_t kInt32Max = std::numeric_limits::max(); ORT_RETURN_IF_NOT(suffix_dim_size <= kInt32Max, "OneHot: suffix dimension size (", suffix_dim_size, ") exceeds int32 range supported by the CUDA kernel."); - ORT_RETURN_IF_NOT(depth_val <= kInt32Max / std::max(suffix_dim_size, int64_t{1}), - "OneHot: depth (", depth_val, ") * suffix dimension size (", suffix_dim_size, - ") exceeds int32 range supported by the CUDA kernel."); const fast_divmod fdm_suffix(gsl::narrow_cast(suffix_dim_size)); const auto* indices_data = indices->Data(); @@ -90,6 +87,10 @@ Status OneHotOp::ComputeInternal(OpKernelContext* return Status::OK(); } + // depth * suffix is only needed for fdm_depth_suffix on the non-zero-off-value path. + ORT_RETURN_IF_NOT(depth_val <= kInt32Max / std::max(suffix_dim_size, int64_t{1}), + "OneHot: depth (", depth_val, ") * suffix dimension size (", suffix_dim_size, + ") exceeds int32 range supported by the CUDA kernel."); const fast_divmod fdm_depth_suffix(gsl::narrow_cast(depth_val * suffix_dim_size)); OneHotImpl(Stream(ctx), indices_data, fdm_depth_suffix, fdm_suffix, depth_val, From bb615e2fcb7555f4fb2b1647edf9862557714ab4 Mon Sep 17 00:00:00 2001 From: Gopalakrishnan Nallasamy Date: Thu, 7 May 2026 15:38:56 -0700 Subject: [PATCH 5/5] Add explicit include in cuda_kernel_adapter.h --- onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index 3be5effbfcf02..28e94e3efccee 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -15,6 +15,8 @@ #pragma once +#include + #include "core/common/status.h" #include "core/common/narrow.h" #include "core/common/safeint.h"