diff --git a/onnxruntime/core/providers/xnnpack/math/gemm.cc b/onnxruntime/core/providers/xnnpack/math/gemm.cc index f8992b13c8f5d..4b31b78eccfdf 100644 --- a/onnxruntime/core/providers/xnnpack/math/gemm.cc +++ b/onnxruntime/core/providers/xnnpack/math/gemm.cc @@ -121,12 +121,11 @@ Gemm::Gemm(const OpKernelInfo& info) : GemmBase(info), XnnpackKernel(info, /*ena C_matrix_exists_ = C_arg && C_arg->Exists(); - // A - MxK + // A - MxK. M is not cached here: it can be a dynamic dimension that is only known at + // Compute() time, so it is read from the input tensor shape instead. if (trans_A_ == CblasNoTrans) { - M_ = shapeA->dim(0).dim_value() > 1 ? shapeA->dim(0).dim_value() : 1; K_ = shapeA->dim(1).dim_value(); } else { - M_ = shapeA->dim(1).dim_value(); K_ = shapeA->dim(0).dim_value() > 1 ? shapeA->dim(0).dim_value() : 1; } // B - KxN @@ -208,10 +207,11 @@ Status Gemm::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr, Status Gemm::Compute(OpKernelContext* context) const { pthreadpool_t threadpool = GetThreadPool(); const auto* A = context->Input(0); - auto Y = context->Output(0, {M_, N_}); + const int64_t M = trans_A_ == CblasNoTrans ? A->Shape()[0] : A->Shape()[1]; + auto Y = context->Output(0, {M, N_}); // if input is empty tensor, return as nothing need to be calculated and we've set the shape for the output - if (M_ == 0 || N_ == 0) { + if (M == 0 || N_ == 0) { return Status::OK(); } @@ -221,7 +221,7 @@ Status Gemm::Compute(OpKernelContext* context) const { } xnn_status status = reshape_func(op0_.get(), // Number of rows to multiply - trans_A_ == CblasNoTrans ? M_ : K_, + trans_A_ == CblasNoTrans ? M : K_, threadpool); if (status != xnn_status_success) { diff --git a/onnxruntime/core/providers/xnnpack/math/gemm.h b/onnxruntime/core/providers/xnnpack/math/gemm.h index 954aab0698b9c..7dfe075b0508c 100644 --- a/onnxruntime/core/providers/xnnpack/math/gemm.h +++ b/onnxruntime/core/providers/xnnpack/math/gemm.h @@ -29,7 +29,6 @@ class Gemm : protected GemmBase, public XnnpackKernel { private: const Tensor* B_{nullptr}; - int64_t M_ = -1; int64_t K_ = -1; int64_t N_ = -1; diff --git a/onnxruntime/test/providers/xnnpack/xnnpack_basic_test.cc b/onnxruntime/test/providers/xnnpack/xnnpack_basic_test.cc index 62a0f28ad3b10..dc9727c75f7b7 100644 --- a/onnxruntime/test/providers/xnnpack/xnnpack_basic_test.cc +++ b/onnxruntime/test/providers/xnnpack/xnnpack_basic_test.cc @@ -692,6 +692,39 @@ TEST(XnnpackEP, TestGemm_EmptyC_NoSegfault) { }); } +// Regression test for https://github.com/microsoft/onnxruntime/issues/31153. +// A Gemm whose M dimension is symbolic used to have M cached as 1 by the kernel +// constructor (dim_value() returns 0 for a symbolic dim), so Compute() produced a +// [1, N] output and only multiplied the first row. M must come from the input tensor +// at Compute() time. +TEST(XnnpackEP, TestGemm_DynamicM) { + // a_shape is the runtime shape fed for the symbolic M; it must have M > 1 to catch the bug. + const std::vector a_shape = {10, 3}; + const std::vector b_shape = {3, 4}; + auto modelBuilder = [&](ModelTestBuilder& builder) { + auto* input_a = builder.MakeSymbolicInput({std::string("dynamic_m"), b_shape[0]}); + auto* input_b = builder.MakeInitializer(b_shape, -1.f, 1.f); + auto* input_c = builder.MakeInitializer({b_shape[1]}, -1.f, 1.f); + auto* output_arg = builder.MakeOutput(); + auto& gemm_node = builder.AddNode("Gemm", {input_a, input_b, input_c}, {output_arg}); + gemm_node.AddAttribute("alpha", 1.0f); + gemm_node.AddAttribute("beta", 1.0f); + gemm_node.AddAttribute("transA", static_cast(0)); + gemm_node.AddAttribute("transB", static_cast(0)); + + // MakeSymbolicInput doesn't register a feed, so add one ourselves. + OrtValue input_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], a_shape, + builder.rand_gen_.Uniform(a_shape, -1.f, 1.f), &input_value); + builder.feeds_.insert(std::make_pair(input_a->Name(), input_value)); + }; + RunModelTest(modelBuilder, "xnnpack_test_graph_gemm_dynamic_m", + { + ExpectedEPNodeAssignment::All, + 1e-4f /* fp32_abs_err */, + }); +} + #endif } // namespace test