diff --git a/onnxruntime/core/providers/cpu/ml/linearregressor.cc b/onnxruntime/core/providers/cpu/ml/linearregressor.cc index 3f877ec7e6d0e..0c5b6580274bc 100644 --- a/onnxruntime/core/providers/cpu/ml/linearregressor.cc +++ b/onnxruntime/core/providers/cpu/ml/linearregressor.cc @@ -86,6 +86,18 @@ Status LinearRegressor::Compute(OpKernelContext* ctx) const { ptrdiff_t num_batches = input_shape.NumDimensions() <= 1 ? 1 : narrow(input_shape[0]); ptrdiff_t num_features = input_shape.NumDimensions() <= 1 ? narrow(input_shape.Size()) : narrow(input_shape[1]); + + // Coefficients are treated as a [num_targets, num_features] matrix. + // Validate size to prevent out-of-bounds reads in the GEMM backend. + // Use SafeInt for overflow-safe multiplication. + if (num_targets_ < 0 || + coefficients_.size() != SafeInt(num_targets_) * static_cast(num_features)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "LinearRegressor: coefficients attribute size (", coefficients_.size(), + ") does not match targets (", num_targets_, + ") * input features (", num_features, ")"); + } + Tensor& Y = *ctx->Output(0, {num_batches, num_targets_}); concurrency::ThreadPool* tp = ctx->GetOperatorThreadPool(); diff --git a/onnxruntime/test/providers/cpu/ml/linearregressor_test.cc b/onnxruntime/test/providers/cpu/ml/linearregressor_test.cc index c6eb3f5f5777f..bd70051b65f06 100644 --- a/onnxruntime/test/providers/cpu/ml/linearregressor_test.cc +++ b/onnxruntime/test/providers/cpu/ml/linearregressor_test.cc @@ -85,5 +85,18 @@ INSTANTIATE_TEST_SUITE_P( LinearRegressorParam("SOFTMAX_ZERO", {3.442477e-14f, 1.f, 1.670142e-05f, 1.f, 1.0f, 0.f}, 2) )); + +// Regression test: coefficients size must match targets * num_features. +// A mismatch previously caused an out-of-bounds read in MLAS SGEMM packing. +TEST(LinearRegressorTest, CoefficientsSizeMismatch) { + OpTester test("LinearRegressor", 1, onnxruntime::kMLDomain); + // 1 coefficient but input has 2 features and targets=1 → expects 2 coefficients + test.AddAttribute("coefficients", std::vector{1.0f}); + test.AddAttribute("targets", int64_t{1}); + test.AddInput("X", {1, 2}, {1.f, 2.f}); + test.AddOutput("Y", {1, 1}, {0.f}); + test.Run(OpTester::ExpectResult::kExpectFailure, "coefficients attribute size"); +} + } // namespace test } // namespace onnxruntime