diff --git a/onnxruntime/core/providers/cpu/math/matmul_helper.h b/onnxruntime/core/providers/cpu/math/matmul_helper.h index 9da7509eea2c6..41739dd6abbef 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_helper.h +++ b/onnxruntime/core/providers/cpu/math/matmul_helper.h @@ -168,6 +168,10 @@ class MatMulComputeHelper { if (num_output_dims == 0) { // for left and right being both vector, output is scalar thus no shape ORT_RETURN_IF_NOT(M_ == 1 && N_ == 1, "M_ == 1 && N_ == 1 was false"); + // Both inputs are 1D vectors: left(K) dot right(K) => scalar. + // Validate K dimensions match to prevent out-of-bounds reads. + ORT_RETURN_IF_NOT(K_ == right_shape[0], + "MatMul dimension mismatch"); } else { if (left_num_dims == 1) { ORT_RETURN_IF_NOT(num_dims_with_pad - 1 == num_output_dims, "num_dims_with_pad - 1 != num_output_dims"); diff --git a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc index 3562d8b009002..6f03f875673d0 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc @@ -504,5 +504,43 @@ TEST(MatmulIntegerOpTest, SharedPrepackedWeights) { } #endif +// Regression test: 1D×1D inputs with mismatched K dimensions must fail +// instead of causing an out-of-bounds read in the MLAS backend. +TEST(MatmulIntegerOpTest, MatMulInteger_1D_DimensionMismatch) { + OpTester test("MatMulInteger", 10); + // Omit shape info so ONNX shape inference won't reject 1D inputs at graph build time. + test.AddShapeToTensorData(false); + // A is [5] (K=5), B is [1] (K=1) — vector×vector K dimensions don't match. + test.AddInput("T1", {5}, {11, 7, 3, 10, 6}); + test.AddInput("T2", {1}, {1}); + test.AddInput("a_zero_point", {}, {0}); + test.AddInput("b_zero_point", {}, {0}); + test.AddOutput("T3", {}, {0}); + + // Restrict to CPU — other EPs may not support 1D quantized inputs. + std::vector> cpu_only; + cpu_only.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "MatMul dimension mismatch", + {}, nullptr, &cpu_only); +} + +// Valid 1D×1D dot product: both vectors have the same K, producing scalar output. +TEST(MatmulIntegerOpTest, MatMulInteger_1D_Valid) { + OpTester test("MatMulInteger", 10); + // Omit shape info so ONNX shape inference won't reject 1D inputs at graph build time. + test.AddShapeToTensorData(false); + // A=[2,3], B=[4,5] => scalar dot = 2*4 + 3*5 = 23 + test.AddInput("T1", {2}, {2, 3}); + test.AddInput("T2", {2}, {4, 5}); + test.AddInput("a_zero_point", {}, {0}); + test.AddInput("b_zero_point", {}, {0}); + test.AddOutput("T3", {}, {23}); + + // Restrict to CPU — other EPs may not support 1D quantized inputs. + std::vector> cpu_only; + cpu_only.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &cpu_only); +} + } // namespace test } // namespace onnxruntime