[XNNPACK EP] Read Gemm M from the input tensor at Compute time

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<Tensor>(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) {

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;
 

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<int64_t> a_shape = {10, 3};
+  const std::vector<int64_t> b_shape = {3, 4};
+  auto modelBuilder = [&](ModelTestBuilder& builder) {
+    auto* input_a = builder.MakeSymbolicInput<float>({std::string("dynamic_m"), b_shape[0]});
+    auto* input_b = builder.MakeInitializer<float>(b_shape, -1.f, 1.f);
+    auto* input_c = builder.MakeInitializer<float>({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<int64_t>(0));
+    gemm_node.AddAttribute("transB", static_cast<int64_t>(0));
+
+    // MakeSymbolicInput doesn't register a feed, so add one ourselves.
+    OrtValue input_value;
+    CreateMLValue<float>(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], a_shape,
+                         builder.rand_gen_.Uniform<float>(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