Skip to content

Commit b64c095

Browse files
bmehta001Copilot
andcommitted
Validate runtime split sizes are non-negative
When split sizes come from a runtime tensor input (not an attribute), they were not validated for non-negative values. Crafted negative split sizes that sum to the correct total could pass the sum check and create tensors with negative dimensions. The attribute path already had this validation. Added non-negative check in both CPU PrepareForCompute and CUDA PrepareForComputeLocal. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 5c0a19b commit b64c095

2 files changed

Lines changed: 6 additions & 0 deletions

File tree

onnxruntime/core/providers/cpu/tensor/split.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,9 @@ class SplitBase {
6565
if (split_size_sum == -1) {
6666
split_size_sum = std::accumulate(split_sizes.cbegin(), split_sizes.cend(), 0LL);
6767
}
68+
ORT_RETURN_IF_NOT(std::all_of(split_sizes.cbegin(), split_sizes.cend(),
69+
[](int64_t value) { return value >= 0; }),
70+
"Invalid value in 'split' input. All values must be >= 0");
6871
if (split_sizes.size() != static_cast<size_t>(num_outputs) || split_size_sum != split_dim_size)
6972
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
7073
"Cannot split using values in 'split' attribute. Axis=", axis_,

onnxruntime/core/providers/cuda/tensor/split.cc

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,9 @@ Status SplitKernel::PrepareForComputeLocal(const TensorShape& input_shape,
8989
if (split_size_sum == -1) {
9090
split_size_sum = std::accumulate(split_sizes.cbegin(), split_sizes.cend(), 0LL);
9191
}
92+
ORT_RETURN_IF_NOT(std::all_of(split_sizes.cbegin(), split_sizes.cend(),
93+
[](int64_t value) { return value >= 0; }),
94+
"Invalid value in 'split' input. All values must be >= 0");
9295
if (split_sizes.size() != static_cast<size_t>(num_outputs) || split_size_sum != split_dim_size) {
9396
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
9497
"Cannot split using values in 'split' attribute. Axis=", axis_,

0 commit comments

Comments
 (0)