From f20400c7b50570ab951b77c261bd3e6dc886421d Mon Sep 17 00:00:00 2001 From: vraspar Date: Tue, 7 Apr 2026 15:39:07 -0700 Subject: [PATCH 1/3] Fix OOB reads in SoftmaxCrossEntropyLoss via label bounds validation Add bounds checking for label tensor values in SoftmaxCrossEntropyLoss and SoftmaxCrossEntropyLossGrad to prevent out-of-bounds memory reads when processing untrusted ONNX models. The operators used label_data values as array indices into log_prob_data and weight_data buffers without validating they fall within [0, C) where C is the number of classes. A malicious model could embed arbitrary label values causing heap OOB reads. Changes: - Add label value validation in both forward and backward Compute methods - Add weight tensor size validation (weight_shape[0] == C) - Move weight_data[label] access after ignore_index check in grad path to prevent OOB when ignore_index is outside [0, C) - Add regression tests for both forward and backward paths Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/loss/cross_entropy_test.cc | 99 +++++++++++++++++++ .../cpu/loss/softmax_cross_entropy_loss.cc | 37 +++++-- 2 files changed, 130 insertions(+), 6 deletions(-) create mode 100644 orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc diff --git a/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc b/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc new file mode 100644 index 0000000000000..f8220dd1977c5 --- /dev/null +++ b/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc @@ -0,0 +1,99 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime { +namespace test { + +// Regression tests for OOB reads when label values are outside [0, C). + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLarge) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 5, 2}; // 5 is out of range [0, 5) + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_NegativeLabel) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-100)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, -1, 2}; // -1 is out of range (and != ignore_index) + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLargeWithWeights) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 100, 2}; // 100 is out of range + std::vector weight_data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f}; + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddInput("weight", {5}, weight_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLarge) { + OpTester test("SoftmaxCrossEntropyLossGrad", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector dY_data = {1.0f}; + std::vector log_prob_data(3 * 5, -1.6094f); + std::vector index_data = {0, 5, 2}; // 5 is out of range [0, 5) + + test.AddInput("dY", {}, dY_data); + test.AddInput("log_prob", {3, 5}, log_prob_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("dX", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLargeWithWeights) { + OpTester test("SoftmaxCrossEntropyLossGrad", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector dY_data = {1.0f}; + std::vector log_prob_data(3 * 5, -1.6094f); + std::vector index_data = {0, 5, 2}; // 5 is out of range [0, 5) + std::vector weight_data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f}; + + test.AddInput("dY", {}, dY_data); + test.AddInput("log_prob", {3, 5}, log_prob_data); + test.AddInput("index", {3}, index_data); + test.AddInput("weight", {5}, weight_data); + test.AddOutput("dX", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc index c74bf06a77d6e..7cec8bd0b6946 100644 --- a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc +++ b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc @@ -56,12 +56,17 @@ void GetNDCFromLogitAndLabelShape(const TensorShape& logit_shape, const TensorSh void VerifyLogitWeightAndLabelShape(const TensorShape& logit_shape, const TensorShape& label_shape, const TensorShape* weight_shape) { - ORT_ENFORCE(nullptr == weight_shape || 1 == weight_shape->NumDimensions(), "Weights tensor is not 1-D."); - const size_t label_dims = label_shape.NumDimensions(); + ORT_ENFORCE(label_dims >= 1, "label must be at least 1-D."); + ORT_ENFORCE(logit_shape.NumDimensions() >= 2, "logit must be at least 2-D."); ORT_ENFORCE(logit_shape.NumDimensions() == label_dims + 1, "logit_shape must be (1 + label_shape)"); + ORT_ENFORCE(nullptr == weight_shape || 1 == weight_shape->NumDimensions(), "Weights tensor is not 1-D."); + ORT_ENFORCE(nullptr == weight_shape || (*weight_shape)[0] == logit_shape[1], + "Weight tensor size (", (weight_shape ? (*weight_shape)[0] : 0), + ") must equal the number of classes (", logit_shape[1], ")"); + ORT_ENFORCE(label_shape[0] == logit_shape[0], "The shape of logit and label does not match"); if (label_dims >= 2) { @@ -147,6 +152,16 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const } const T2* label_data = label.template Data(); + + // Validate label values are within [0, C) to prevent out-of-bounds reads. + for (int64_t i = 0; i < N_D; i++) { + if (ignore_index != label_data[i]) { + ORT_RETURN_IF(label_data[i] < 0 || label_data[i] >= C, + "SoftmaxCrossEntropyLoss: label value ", label_data[i], + " at index ", i, " is out of range [0, ", C, ")"); + } + } + T1* loss_data = loss->template MutableData(); std::vector shifted_logit(narrow(n_d_c)); ORT_ENFORCE(n_d_c <= static_cast(std::numeric_limits::max())); @@ -267,6 +282,16 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co const T1* dY_data = dY.template Data(); const T1* log_prob_data = log_prob.template Data(); const T2* label_data = label.template Data(); + + // Validate label values are within [0, C) to prevent out-of-bounds reads. + for (int64_t i = 0; i < N_D; i++) { + if (ignore_index != label_data[i]) { + ORT_RETURN_IF(label_data[i] < 0 || label_data[i] >= C, + "SoftmaxCrossEntropyLossGrad: label value ", label_data[i], + " at index ", i, " is out of range [0, ", C, ")"); + } + } + Tensor* d_logit = context->Output(0, probability_shape); T1* d_logit_data = d_logit->template MutableData(); std::memset(d_logit_data, 0, narrow(sizeof(T1) * N_D)); @@ -299,11 +324,11 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co int64_t row = index / C; int64_t col = index % C; T2 label_sample = label_data[row]; - T1 weight_smaple = weight_data[label_sample] * dY_data[row]; if (ignore_index == label_sample) { d_logit_data[index] = 0; } else { - d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_smaple; + T1 weight_sample = weight_data[label_sample] * dY_data[row]; + d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_sample; } } }); @@ -330,11 +355,11 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co int64_t row = index / C; int64_t col = index % C; T2 label_sample = label_data[row]; - T1 weight_smaple = weight_data[label_sample] * dY_scaled; if (ignore_index == label_sample) { d_logit_data[index] = 0; } else { - d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_smaple; + T1 weight_sample = weight_data[label_sample] * dY_scaled; + d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_sample; } } }); From 5699dd17128a12c4c7e1c8dcbd4746ed75e0707a Mon Sep 17 00:00:00 2001 From: vraspar Date: Tue, 12 May 2026 21:26:33 +0000 Subject: [PATCH 2/3] Address review: inline label bounds check, SafeInt overflow guards, return Status on bad input - Fold label range check into the 3 forward per-sample loops and drop the separate validation pass. - Backward keeps a single upfront serial check (parallel-for lambdas cannot return Status) with a comment explaining why. - Wrap N_D * C in SafeInt and use gsl::narrow on per-loop counts; catch overflow/narrowing and return INVALID_ARGUMENT. - Convert IsScalar and Eigen::Index size ORT_ENFORCE to ORT_RETURN_IF_NOT / ORT_RETURN_IF in both forward and backward. - Fix wrong-sized memset in backward: size from N_D to N_D * C via probability_shape.Size(). - Add tests for: weighted SUM forward path, int32 labels (fwd+grad), 4-D logit, and Internal/InternalGrad variants with runtime ignore_index. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/loss/cross_entropy_test.cc | 115 ++++++++++++++++++ .../cpu/loss/softmax_cross_entropy_loss.cc | 92 +++++++++----- 2 files changed, 177 insertions(+), 30 deletions(-) diff --git a/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc b/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc index f8220dd1977c5..8f72bdccdd589 100644 --- a/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc +++ b/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc @@ -1,6 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include + #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" @@ -59,6 +62,61 @@ TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLargeWithWeights) { test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); } +// Covers the weighted, non-MEAN forward loop (the second per-sample loop in Compute). +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLargeWithWeightsSumReduction) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("sum")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 7, 2}; // 7 is out of range [0, 5) + std::vector weight_data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f}; + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddInput("weight", {5}, weight_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +// int32 label type — kernel is registered for both int32_t and int64_t. +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLargeInt32) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 5, 2}; // 5 is out of range [0, 5) + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +// Higher-dimensional inputs: logit [N, C, D1, D2], label [N, D1, D2]. +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLargeHighDim) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + // [N=2, C=4, D1=2, D2=3] -> label shape [2, 2, 3] -> 12 label entries. + std::vector X_data(2 * 4 * 2 * 3, 1.0f); + std::vector index_data = {0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 99, 3}; // 99 is out of range + std::vector log_prob_init(2 * 4 * 2 * 3, 0.0f); + + test.AddInput("X", {2, 4, 2, 3}, X_data); + test.AddInput("index", {2, 2, 3}, index_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {2, 4, 2, 3}, log_prob_init); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLarge) { OpTester test("SoftmaxCrossEntropyLossGrad", 1, onnxruntime::kMSDomain); test.AddAttribute("reduction", std::string("mean")); @@ -95,5 +153,62 @@ TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLargeWithWeights) { test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); } +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLargeInt32) { + OpTester test("SoftmaxCrossEntropyLossGrad", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector dY_data = {1.0f}; + std::vector log_prob_data(3 * 5, -1.6094f); + std::vector index_data = {0, 5, 2}; + + test.AddInput("dY", {}, dY_data); + test.AddInput("log_prob", {3, 5}, log_prob_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("dX", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +// SoftmaxCrossEntropyLossInternal shares the same Compute as SoftmaxCrossEntropyLoss but +// is registered separately under kMSDomain v1 with an optional runtime ignore_index input. +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternal_LabelTooLarge) { + OpTester test("SoftmaxCrossEntropyLossInternal", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 5, 2}; + int64_t ignore_index_val = -1; + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + // weight is optional; pass empty input to skip and still provide ignore_index input below. + test.AddOptionalInputEdge(); + test.AddInput("ignore_index", {}, &ignore_index_val, 1); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternalGrad_LabelTooLarge) { + OpTester test("SoftmaxCrossEntropyLossInternalGrad", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + + std::vector dY_data = {1.0f}; + std::vector log_prob_data(3 * 5, -1.6094f); + std::vector index_data = {0, 5, 2}; + int64_t ignore_index_val = -1; + + test.AddInput("dY", {}, dY_data); + test.AddInput("log_prob", {3, 5}, log_prob_data); + test.AddInput("index", {3}, index_data); + test.AddOptionalInputEdge(); // weight + test.AddInput("ignore_index", {}, &ignore_index_val, 1); + test.AddOutput("dX", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + } // namespace test } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc index 7cec8bd0b6946..00278166ee24b 100644 --- a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc +++ b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc @@ -3,6 +3,7 @@ #include "core/util/math.h" #include "core/util/math_cpuonly.h" +#include "core/common/safeint.h" #include "core/providers/common.h" #include #include "core/providers/cpu/math/matmul_helper.h" @@ -108,7 +109,7 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const const Tensor* p_ignore_index = context->Input(3); int64_t ignore_index = ignore_index_; if (p_ignore_index) { - ORT_ENFORCE(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar."); + ORT_RETURN_IF_NOT(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar."); ignore_index = *(p_ignore_index->template Data()); } @@ -135,9 +136,23 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const logit_data = (*transpose_output.GetMutable()).template Data(); } - const int n_d = gsl::narrow_cast(N_D); - const int c = gsl::narrow_cast(C); - const uint64_t n_d_c = N_D * C; + // Convert N_D, C to int and compute N_D * C with overflow / truncation checks. + // gsl::narrow throws on truncation; SafeInt throws on overflow. Translate both into INVALID_ARGUMENT. + int n_d = 0; + int c = 0; + uint64_t n_d_c = 0; + try { + n_d = gsl::narrow(N_D); + c = gsl::narrow(C); + n_d_c = static_cast(SafeInt(N_D) * static_cast(C)); + } catch (const std::exception& e) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLoss: invalid logit dimensions N_D=", N_D, + ", C=", C, " (", e.what(), ")"); + } + ORT_RETURN_IF(n_d_c > static_cast(std::numeric_limits::max()), + "SoftmaxCrossEntropyLoss: N_D * C (", n_d_c, ") exceeds Eigen::Index max."); + Tensor* loss = context->Output(0, reduction_ == ReductionType::NONE ? TensorShape(label.Shape()) : TensorShape({})); T1* log_prob_data; std::vector log_prob_data_buffer(0); @@ -153,18 +168,8 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const const T2* label_data = label.template Data(); - // Validate label values are within [0, C) to prevent out-of-bounds reads. - for (int64_t i = 0; i < N_D; i++) { - if (ignore_index != label_data[i]) { - ORT_RETURN_IF(label_data[i] < 0 || label_data[i] >= C, - "SoftmaxCrossEntropyLoss: label value ", label_data[i], - " at index ", i, " is out of range [0, ", C, ")"); - } - } - T1* loss_data = loss->template MutableData(); std::vector shifted_logit(narrow(n_d_c)); - ORT_ENFORCE(n_d_c <= static_cast(std::numeric_limits::max())); ComputeShareSoftmaxCrossEntropyCPU(n_d, c, static_cast(n_d_c), logit_data, shifted_logit.data(), log_prob_data); std::vector loss_sample_buffer(0); @@ -190,19 +195,27 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const // Compute weighed loss for each sample while summing weights for unignored target/label values. if (reduction_ == ReductionType::MEAN) { for (ptrdiff_t i = 0; i < n_d; i++) { - if (ignore_index == label_data[i]) { + const T2 label_sample = label_data[i]; + if (ignore_index == label_sample) { loss_sample[i] = 0; } else { - loss_sample[i] = -log_prob_data[i * c + label_data[i]] * weight_data[label_data[i]]; - sum_weight += weight_data[label_data[i]]; + ORT_RETURN_IF(label_sample < 0 || label_sample >= C, + "SoftmaxCrossEntropyLoss: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + loss_sample[i] = -log_prob_data[i * c + label_sample] * weight_data[label_sample]; + sum_weight += weight_data[label_sample]; } } } else { for (ptrdiff_t i = 0; i < n_d; i++) { - if (ignore_index == label_data[i]) { + const T2 label_sample = label_data[i]; + if (ignore_index == label_sample) { loss_sample[i] = 0; } else { - loss_sample[i] = -log_prob_data[i * c + label_data[i]] * weight_data[label_data[i]]; + ORT_RETURN_IF(label_sample < 0 || label_sample >= C, + "SoftmaxCrossEntropyLoss: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + loss_sample[i] = -log_prob_data[i * c + label_sample] * weight_data[label_sample]; } } } @@ -220,10 +233,14 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const // Compute loss for each sample while counting unignored target/label values. int unignored_samples = 0; for (ptrdiff_t i = 0; i < n_d; i++) { - if (ignore_index == label_data[i]) { + const T2 label_sample = label_data[i]; + if (ignore_index == label_sample) { loss_sample[i] = 0; } else { - loss_sample[i] = -log_prob_data[i * c + label_data[i]]; + ORT_RETURN_IF(label_sample < 0 || label_sample >= C, + "SoftmaxCrossEntropyLoss: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + loss_sample[i] = -log_prob_data[i * c + label_sample]; unignored_samples += 1; } } @@ -267,7 +284,7 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co const Tensor* p_ignore_index = context->Input(4); int64_t ignore_index = ignore_index_; if (p_ignore_index) { - ORT_ENFORCE(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar."); + ORT_RETURN_IF_NOT(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar."); ignore_index = *(p_ignore_index->template Data()); } @@ -279,22 +296,37 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co int64_t N_D = 0; int64_t C = 0; GetNDCFromLogitAndLabelShape(probability_shape, label_shape, N_D, C); + + // Compute N_D * C once with overflow protection; reused by every parallel-for below. + ptrdiff_t n_d_c = 0; + try { + n_d_c = narrow(static_cast(SafeInt(N_D) * C)); + } catch (const std::exception& e) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLossGrad: invalid logit dimensions N_D=", N_D, + ", C=", C, " (", e.what(), ")"); + } + const T1* dY_data = dY.template Data(); const T1* log_prob_data = log_prob.template Data(); const T2* label_data = label.template Data(); // Validate label values are within [0, C) to prevent out-of-bounds reads. + // Done as a single up-front O(N_D) pass: ORT_RETURN_IF cannot escape from inside the + // ThreadPool::TryParallelFor lambdas below, and the cost is negligible compared to the + // O(N_D * C) parallel gradient computation that follows. for (int64_t i = 0; i < N_D; i++) { - if (ignore_index != label_data[i]) { - ORT_RETURN_IF(label_data[i] < 0 || label_data[i] >= C, - "SoftmaxCrossEntropyLossGrad: label value ", label_data[i], + const T2 label_sample = label_data[i]; + if (ignore_index != label_sample) { + ORT_RETURN_IF(label_sample < 0 || label_sample >= C, + "SoftmaxCrossEntropyLossGrad: label value ", label_sample, " at index ", i, " is out of range [0, ", C, ")"); } } Tensor* d_logit = context->Output(0, probability_shape); T1* d_logit_data = d_logit->template MutableData(); - std::memset(d_logit_data, 0, narrow(sizeof(T1) * N_D)); + std::memset(d_logit_data, 0, narrow(sizeof(T1) * probability_shape.Size())); OrtValue transpose_output; TensorShapeVector new_shape; std::vector permutations; @@ -317,7 +349,7 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co if (reduction_ == ReductionType::NONE) { concurrency::ThreadPool::TryParallelFor( - tp, narrow(N_D * C), cost, + tp, n_d_c, cost, [&label_data, &weight_data, &d_logit_data, &log_prob_data, ignore_index, C, &dY_data]( std::ptrdiff_t begin, std::ptrdiff_t end) { for (std::ptrdiff_t index = begin; index != end; ++index) { @@ -348,7 +380,7 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co } concurrency::ThreadPool::TryParallelFor( - tp, narrow(N_D * C), cost, + tp, n_d_c, cost, [&label_data, &weight_data, dY_scaled, &d_logit_data, &log_prob_data, ignore_index, C]( std::ptrdiff_t begin, std::ptrdiff_t end) { for (std::ptrdiff_t index = begin; index != end; ++index) { @@ -367,7 +399,7 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co } else { if (reduction_ == ReductionType::NONE) { concurrency::ThreadPool::TryParallelFor( - tp, narrow(N_D * C), cost, + tp, n_d_c, cost, [&label_data, &d_logit_data, &log_prob_data, ignore_index, C, &dY_data]( std::ptrdiff_t begin, std::ptrdiff_t end) { for (std::ptrdiff_t index = begin; index != end; ++index) { @@ -396,7 +428,7 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co } concurrency::ThreadPool::TryParallelFor( - tp, narrow(N_D * C), cost, + tp, n_d_c, cost, [&label_data, &d_logit_data, &log_prob_data, ignore_index, C, &dY_scaled]( std::ptrdiff_t begin, std::ptrdiff_t end) { for (std::ptrdiff_t index = begin; index != end; ++index) { From cbba7a355b01d947ea90a2af2d63000c9e7c4848 Mon Sep 17 00:00:00 2001 From: vraspar Date: Tue, 12 May 2026 23:54:02 +0000 Subject: [PATCH 3/3] Fix no-exceptions build: use deform_conv overflow pattern and return INVALID_ARGUMENT for label bounds --- .../cpu/loss/softmax_cross_entropy_loss.cc | 111 +++++++++++------- 1 file changed, 69 insertions(+), 42 deletions(-) diff --git a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc index 00278166ee24b..b575344219d47 100644 --- a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc +++ b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc @@ -3,7 +3,6 @@ #include "core/util/math.h" #include "core/util/math_cpuonly.h" -#include "core/common/safeint.h" #include "core/providers/common.h" #include #include "core/providers/cpu/math/matmul_helper.h" @@ -11,6 +10,8 @@ #include "core/providers/cpu/controlflow/scan_utils.h" #include "orttraining/training_ops/cpu/loss/cross_entropy.h" #include "orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.h" +#include + #include namespace onnxruntime { @@ -117,6 +118,12 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const const TensorShape label_shape{label.Shape()}; VerifyLogitWeightAndLabelShape(logit_shape, label_shape, p_weight ? &p_weight->Shape() : nullptr); + // Pre-empt the divide-by-zero in GetNDCFromLogitAndLabelShape (which does + // C = logit_shape.Size() / N_D). VerifyLogitWeightAndLabelShape only checks + // rank, not Size > 0 — e.g. label shape [3, 0, 5] is rank 2 but empty. + ORT_RETURN_IF_NOT(label_shape.Size() > 0, + "SoftmaxCrossEntropyLoss: label tensor must not be empty."); + // N_D = N * D1 * D2...D*K int64_t N_D = 0; int64_t C = 0; @@ -136,22 +143,26 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const logit_data = (*transpose_output.GetMutable()).template Data(); } - // Convert N_D, C to int and compute N_D * C with overflow / truncation checks. - // gsl::narrow throws on truncation; SafeInt throws on overflow. Translate both into INVALID_ARGUMENT. - int n_d = 0; - int c = 0; - uint64_t n_d_c = 0; - try { - n_d = gsl::narrow(N_D); - c = gsl::narrow(C); - n_d_c = static_cast(SafeInt(N_D) * static_cast(C)); - } catch (const std::exception& e) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "SoftmaxCrossEntropyLoss: invalid logit dimensions N_D=", N_D, - ", C=", C, " (", e.what(), ")"); - } - ORT_RETURN_IF(n_d_c > static_cast(std::numeric_limits::max()), - "SoftmaxCrossEntropyLoss: N_D * C (", n_d_c, ") exceeds Eigen::Index max."); + // Validate N_D, C and compute N_D * C with overflow checks. + // Uses the no-exceptions-safe pattern from deform_conv_attributes.h: explicit + // pre-multiply guards via ORT_RETURN_IF_NOT, then a plain multiply, then narrow_cast. + constexpr int64_t kInt64Max = std::numeric_limits::max(); + constexpr int64_t kIntMax = static_cast(std::numeric_limits::max()); + constexpr int64_t kEigenIndexMax = static_cast(std::numeric_limits::max()); + + ORT_RETURN_IF_NOT(N_D > 0 && C > 0, + "SoftmaxCrossEntropyLoss: N_D and C must be positive (got N_D=", N_D, ", C=", C, ")."); + ORT_RETURN_IF_NOT(N_D <= kIntMax && C <= kIntMax, + "SoftmaxCrossEntropyLoss: N_D=", N_D, ", C=", C, " exceed int max."); + ORT_RETURN_IF_NOT(N_D <= kInt64Max / C, + "SoftmaxCrossEntropyLoss: N_D * C overflows int64 (N_D=", N_D, ", C=", C, ")."); + const int64_t n_d_c_i64 = N_D * C; + ORT_RETURN_IF_NOT(n_d_c_i64 <= kEigenIndexMax, + "SoftmaxCrossEntropyLoss: N_D * C (", n_d_c_i64, ") exceeds Eigen::Index max."); + + const int n_d = gsl::narrow_cast(N_D); + const int c = gsl::narrow_cast(C); + const uint64_t n_d_c = static_cast(n_d_c_i64); Tensor* loss = context->Output(0, reduction_ == ReductionType::NONE ? TensorShape(label.Shape()) : TensorShape({})); T1* log_prob_data; @@ -199,9 +210,11 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const if (ignore_index == label_sample) { loss_sample[i] = 0; } else { - ORT_RETURN_IF(label_sample < 0 || label_sample >= C, - "SoftmaxCrossEntropyLoss: label value ", label_sample, - " at index ", i, " is out of range [0, ", C, ")"); + if (label_sample < 0 || label_sample >= C) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLoss: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + } loss_sample[i] = -log_prob_data[i * c + label_sample] * weight_data[label_sample]; sum_weight += weight_data[label_sample]; } @@ -212,9 +225,11 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const if (ignore_index == label_sample) { loss_sample[i] = 0; } else { - ORT_RETURN_IF(label_sample < 0 || label_sample >= C, - "SoftmaxCrossEntropyLoss: label value ", label_sample, - " at index ", i, " is out of range [0, ", C, ")"); + if (label_sample < 0 || label_sample >= C) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLoss: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + } loss_sample[i] = -log_prob_data[i * c + label_sample] * weight_data[label_sample]; } } @@ -237,9 +252,11 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const if (ignore_index == label_sample) { loss_sample[i] = 0; } else { - ORT_RETURN_IF(label_sample < 0 || label_sample >= C, - "SoftmaxCrossEntropyLoss: label value ", label_sample, - " at index ", i, " is out of range [0, ", C, ")"); + if (label_sample < 0 || label_sample >= C) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLoss: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + } loss_sample[i] = -log_prob_data[i * c + label_sample]; unignored_samples += 1; } @@ -292,35 +309,45 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co const TensorShape label_shape{label.Shape()}; VerifyLogitWeightAndLabelShape(probability_shape, label_shape, p_weight ? &p_weight->Shape() : nullptr); + // Pre-empt the divide-by-zero in GetNDCFromLogitAndLabelShape (see forward Compute). + ORT_RETURN_IF_NOT(label_shape.Size() > 0, + "SoftmaxCrossEntropyLossGrad: label tensor must not be empty."); + // N_D = N * D1 * D2...D*K int64_t N_D = 0; int64_t C = 0; GetNDCFromLogitAndLabelShape(probability_shape, label_shape, N_D, C); - // Compute N_D * C once with overflow protection; reused by every parallel-for below. - ptrdiff_t n_d_c = 0; - try { - n_d_c = narrow(static_cast(SafeInt(N_D) * C)); - } catch (const std::exception& e) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "SoftmaxCrossEntropyLossGrad: invalid logit dimensions N_D=", N_D, - ", C=", C, " (", e.what(), ")"); - } + // Compute N_D * C once with overflow checks; reused by every parallel-for below. + // Backward keeps N_D and C as int64_t in lambdas, so only the product needs to fit. + // TryParallelFor takes std::ptrdiff_t, which is 32-bit on 32-bit minimal builds — + // use ptrdiff_t max as the ceiling, not int64_t max. + constexpr int64_t kInt64Max = std::numeric_limits::max(); + constexpr int64_t kPtrdiffMax = static_cast(std::numeric_limits::max()); + + ORT_RETURN_IF_NOT(N_D > 0 && C > 0, + "SoftmaxCrossEntropyLossGrad: N_D and C must be positive (got N_D=", N_D, ", C=", C, ")."); + ORT_RETURN_IF_NOT(N_D <= kInt64Max / C, + "SoftmaxCrossEntropyLossGrad: N_D * C overflows int64 (N_D=", N_D, ", C=", C, ")."); + const int64_t n_d_c_i64 = N_D * C; + ORT_RETURN_IF_NOT(n_d_c_i64 <= kPtrdiffMax, + "SoftmaxCrossEntropyLossGrad: N_D * C (", n_d_c_i64, ") exceeds ptrdiff_t max."); + const std::ptrdiff_t n_d_c = static_cast(n_d_c_i64); const T1* dY_data = dY.template Data(); const T1* log_prob_data = log_prob.template Data(); const T2* label_data = label.template Data(); // Validate label values are within [0, C) to prevent out-of-bounds reads. - // Done as a single up-front O(N_D) pass: ORT_RETURN_IF cannot escape from inside the - // ThreadPool::TryParallelFor lambdas below, and the cost is negligible compared to the - // O(N_D * C) parallel gradient computation that follows. + // Done as a single up-front O(N_D) pass: returning Status from inside a + // ThreadPool::TryParallelFor lambda would only return from the lambda, not from + // Compute. Cost is negligible vs the O(N_D * C) parallel gradient that follows. for (int64_t i = 0; i < N_D; i++) { const T2 label_sample = label_data[i]; - if (ignore_index != label_sample) { - ORT_RETURN_IF(label_sample < 0 || label_sample >= C, - "SoftmaxCrossEntropyLossGrad: label value ", label_sample, - " at index ", i, " is out of range [0, ", C, ")"); + if (ignore_index != label_sample && (label_sample < 0 || label_sample >= C)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLossGrad: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); } }