[NNAPI EP] Add some support for MatMul with batch inputs (#12261)

MatMul allows multiplying batches of matrices. This change enables limited support of batch inputs in the NNAPI EP.

Some limitations:
- Broadcasting is not supported. A and B must have the same leading dimensions.
- Only float inputs are supported. QDQ MatMul or QLinearMatMul with batch inputs is not supported yet.

Note that NNAPI's ANEURALNETWORKS_BATCH_MATMUL is pretty much what we need, but it is only available from NNAPI feature level 6. This change composes a bunch of NNAPI operations to achieve a similar result but this is not ideal.

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc

@@ -1,21 +1,23 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "core/providers/nnapi/nnapi_builtin/builders/helper.h"
+
+#include <functional>
 #include <iostream>
+#include <numeric>
 #include <string>
 #include <vector>
 
-#include "helper.h"
-
-#include "core/common/safeint.h"
 #include "core/common/logging/logging.h"
+#include "core/common/safeint.h"
 #include "core/framework/tensorprotoutils.h"
-#include "core/graph/graph.h"
 #include "core/graph/graph_viewer.h"
+#include "core/graph/graph.h"
 #include "core/providers/common.h"
+#include "core/providers/nnapi/nnapi_builtin/builders/op_support_checker.h"
 #include "core/providers/shared/node_unit/node_unit.h"
 #include "core/providers/shared/utils/utils.h"
-#include "op_support_checker.h"
 
 namespace onnxruntime {
 namespace nnapi {
@@ -377,6 +379,13 @@ std::string Shape2String(const std::vector<uint32_t>& shape) {
   return os.str();
 }
 
+uint32_t ShapeSize(const Shape& shape, size_t begin_idx, size_t end_idx) {
+  ORT_ENFORCE(begin_idx <= end_idx && begin_idx <= shape.size(),
+              "Invalid indices: begin [", begin_idx, "], end [", end_idx, "], shape size [", shape.size(), "]");
+  return std::accumulate(shape.begin() + begin_idx, shape.begin() + end_idx,
+                         SafeInt<uint32_t>{1}, std::multiplies<SafeInt<uint32_t>>{});
+}
+
 bool CheckIsInitializer(const InitializedTensorSet& initializers, const NodeUnit& node_unit,
                         const std::string& input_name, const char* input_description) {
   if (!Contains(initializers, input_name)) {

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h

@@ -4,6 +4,7 @@
 #pragma once
 
 #include <string>
+#include <vector>
 #include "core/graph/basic_types.h"
 #include "core/providers/nnapi/nnapi_builtin/nnapi_lib/NeuralNetworksTypes.h"
 
@@ -158,6 +159,11 @@ bool IsValidSupportedNodeGroup(const std::vector<const Node*>& supported_node_gr
 // Get string representation of a Shape
 std::string Shape2String(const std::vector<uint32_t>& shape);
 
+uint32_t ShapeSize(const Shape& shape, size_t begin_idx, size_t end_idx);
+inline uint32_t ShapeSize(const Shape& shape) {
+  return ShapeSize(shape, 0, shape.size());
+}
+
 // Check the given input is an initializer tensor
 // input_name is the name of the initializer
 // input_description is the string describing the input in the output message (if any)

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc

@@ -13,6 +13,7 @@
 #include "core/providers/shared/utils/utils.h"
 #include "core/providers/shared/node_unit/node_unit.h"
 #include "core/providers/cpu/tensor/slice_helper.h"
+#include "core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h"
 #include "helper.h"
 #include "model_builder.h"
 #include "op_support_checker.h"
@@ -29,37 +30,9 @@ struct OpBuilderRegistrations {
   std::unordered_map<std::string, const IOpBuilder*> op_builder_map;
 };
 
-#define ADD_SCALAR_OPERAND(model_builder, input_indices, scalar_value)             \
-  do {                                                                             \
-    uint32_t _index = 0;                                                           \
-    ORT_RETURN_IF_ERROR(model_builder.AddOperandFromScalar(scalar_value, _index)); \
-    input_indices.push_back(_index);                                               \
-  } while (0)
-
-Status AddTransposeOperator(ModelBuilder& model_builder,
-                            const std::string& input,
-                            const std::string& perm_name,
-                            std::vector<int32_t> perm,
-                            const std::string& output) {
-  auto& shaper(model_builder.GetShaper());
-  const auto& operand_indices(model_builder.GetOperandIndices());
-  const auto& operand_types(model_builder.GetOperandTypes());
-
-  std::vector<uint32_t> input_indices;
-  input_indices.push_back(operand_indices.at(input));  // input
-
-  Shape perm_dimen = {SafeInt<uint32_t>(perm.size())};
-  OperandType perm_operand_type(Type::TENSOR_INT32, perm_dimen);
-  ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(perm_name, perm.data(), perm_operand_type));
-  uint32_t perm_idx = operand_indices.at(perm_name);
-
-  input_indices.push_back(perm_idx);  // permutation
-  ORT_RETURN_IF_ERROR(shaper.Transpose(input, perm, output));
-  OperandType output_operand_type = operand_types.at(input);
-  output_operand_type.SetDimensions(shaper[output]);
-  return model_builder.AddOperation(ANEURALNETWORKS_TRANSPOSE, input_indices, {output},
-                                    {output_operand_type});
-}
+#define ADD_SCALAR_OPERAND(model_builder, input_indices, scalar_value)          \
+  ORT_RETURN_IF_ERROR(onnxruntime::nnapi::op_builder_helpers::AddScalarOperand( \
+      model_builder, input_indices, scalar_value))
 
 static Status AddBinaryOperator(int32_t op_type,
                                 ModelBuilder& model_builder,
@@ -776,7 +749,7 @@ Status TransposeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, co
 
   std::string perm_name = model_builder.GetUniqueName(node_unit.Name() + input + "perm");
 
-  ORT_RETURN_IF_ERROR(AddTransposeOperator(model_builder, input, perm_name, perm, output));
+  ORT_RETURN_IF_ERROR(op_builder_helpers::AddNnapiTranspose(model_builder, input, perm_name, perm, output));
 
   return Status::OK();
 }
@@ -905,16 +878,9 @@ bool ReshapeOpBuilder::IsQuantizedOp(const NodeUnit& node_unit) const {
     model_builder.RegisterOperand(output, operand_indices.at(input), output_operand_type);
   } else {
     // We still need to perform a reshape here
-    // Add input
-    std::vector<uint32_t> input_indices;
-    input_indices.push_back(operand_indices.at(input));
-    // Add new shape
-    Shape shape_dimen = {static_cast<uint32_t>(shape.size())};
     std::string shape_name = model_builder.GetUniqueName(node_unit.Name() + input + "newshape");
-    OperandType shape_operand_type(Type::TENSOR_INT32, shape_dimen);
-    ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(shape_name, shape.data(), shape_operand_type));
-    input_indices.push_back(operand_indices.at(shape_name));
-    ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_RESHAPE, input_indices, {output}, {output_operand_type}));
+    ORT_RETURN_IF_ERROR(op_builder_helpers::AddNnapiReshape(model_builder, input, shape_name, shape, output,
+                                                            &shaper[output]));
   }
 
   return Status::OK();
@@ -1722,6 +1688,11 @@ bool GemmOpBuilder::IsQuantizedOp(const NodeUnit& node_unit) const {
 }
 
 void GemmOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const NodeUnit& node_unit) const {
+  if (op_builder_helpers::IsSupportedBatchMatMul(node_unit, model_builder.GetNNAPIFeatureLevel())) {
+    // no initializers to skip for batch matmul
+    return;
+  }
+
   const auto& inputs = node_unit.Inputs();
   if (IsQuantizedOp(node_unit)) {
     if (node_unit.OpType() == "QLinearMatMul" || node_unit.OpType() == "MatMul") {                 // QLinearMatMul/QDQMatMul
@@ -1758,6 +1729,10 @@ void GemmOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Nod
 }
 
 Status GemmOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const NodeUnit& node_unit) const {
+  if (op_builder_helpers::IsSupportedBatchMatMul(node_unit, model_builder.GetNNAPIFeatureLevel())) {
+    return op_builder_helpers::BuildBatchMatMul(model_builder, node_unit);
+  }
+
   auto& shaper(model_builder.GetShaper());
   const auto& operand_indices(model_builder.GetOperandIndices());
   const auto& operand_types(model_builder.GetOperandTypes());

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.cc

@@ -0,0 +1,356 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h"
+
+#include <algorithm>
+
+#include "gsl/gsl"
+
+#include "core/common/safeint.h"
+#include "core/graph/node_arg.h"
+#include "core/providers/common.h"
+#include "core/providers/nnapi/nnapi_builtin/builders/helper.h"
+#include "core/providers/nnapi/nnapi_builtin/nnapi_lib/NeuralNetworksWrapper.h"
+
+namespace onnxruntime::nnapi::op_builder_helpers {
+
+using android::nn::wrapper::OperandType, android::nn::wrapper::Type;
+
+Status AddNnapiTranspose(ModelBuilder& model_builder,
+                         const std::string& data_input,
+                         const std::string& perm_input, const std::vector<int32_t>& perm,
+                         const std::string& output) {
+  auto& shaper(model_builder.GetShaper());
+  const auto& operand_indices(model_builder.GetOperandIndices());
+  const auto& operand_types(model_builder.GetOperandTypes());
+
+  std::vector<uint32_t> input_indices;
+  input_indices.push_back(operand_indices.at(data_input));  // input
+
+  Shape perm_dimen = {SafeInt<uint32_t>(perm.size())};
+  OperandType perm_operand_type(Type::TENSOR_INT32, perm_dimen);
+  ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(perm_input, perm.data(), perm_operand_type));
+  uint32_t perm_idx = operand_indices.at(perm_input);
+
+  input_indices.push_back(perm_idx);  // permutation
+  ORT_RETURN_IF_ERROR(shaper.Transpose(data_input, perm, output));
+  OperandType output_operand_type = operand_types.at(data_input);
+  output_operand_type.SetDimensions(shaper[output]);
+  return model_builder.AddOperation(ANEURALNETWORKS_TRANSPOSE, input_indices, {output},
+                                    {output_operand_type});
+}
+
+Status AddNnapiReshape(ModelBuilder& model_builder,
+                       const std::string& data_input,
+                       const std::string& shape_input, const std::vector<int32_t>& shape_value,
+                       const std::string& output, const Shape* output_shape) {
+  if (output_shape == nullptr) {
+    auto& shaper = model_builder.GetShaper();
+    ORT_RETURN_IF_ERROR(shaper.Reshape(data_input, shape_value, output));
+    output_shape = &shaper[output];
+  }
+
+  const auto& operand_indices = model_builder.GetOperandIndices();
+  const auto& operand_types = model_builder.GetOperandTypes();
+
+  // Add input
+  std::vector<uint32_t> input_indices;
+  input_indices.push_back(operand_indices.at(data_input));
+
+  // Add new shape
+  const Shape shape_dimen{static_cast<uint32_t>(shape_value.size())};
+  const OperandType shape_operand_type{Type::TENSOR_INT32, shape_dimen};
+  ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(shape_input, shape_value.data(),
+                                                                      shape_operand_type));
+  input_indices.push_back(operand_indices.at(shape_input));
+
+  // For reshape, the output type should be the same as the input type except the shape is different
+  OperandType output_operand_type{operand_types.at(data_input)};
+  output_operand_type.SetDimensions(*output_shape);
+  ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_RESHAPE,
+                                                 input_indices, {output}, {output_operand_type}));
+
+  return Status::OK();
+}
+
+Status AddNnapiSplit(ModelBuilder& model_builder,
+                     const std::string& input,
+                     int32_t axis,
+                     const std::vector<std::string>& outputs) {
+  const auto& operand_indices = model_builder.GetOperandIndices();
+  const auto& operand_types = model_builder.GetOperandTypes();
+  auto& shaper = model_builder.GetShaper();
+
+  const auto input_rank = shaper[input].size();
+  axis = static_cast<int32_t>(HandleNegativeAxis(axis, input_rank));
+
+  ORT_RETURN_IF_ERROR(shaper.Split(input, axis, outputs));
+
+  std::vector<uint32_t> input_indices;
+  input_indices.push_back(operand_indices.at(input));
+  ORT_RETURN_IF_ERROR(AddScalarOperand(model_builder, input_indices, axis));
+  const auto count = gsl::narrow<int32_t>(outputs.size());
+  ORT_RETURN_IF_ERROR(AddScalarOperand(model_builder, input_indices, count));
+
+  const OperandType& input_operand_type = operand_types.at(input);
+  std::vector<OperandType> output_operand_types;
+  output_operand_types.reserve(count);
+  std::transform(outputs.begin(), outputs.end(), std::back_inserter(output_operand_types),
+                 [&](const std::string& output) {
+                   OperandType output_operand_type = input_operand_type;
+                   output_operand_type.SetDimensions(shaper[output]);
+                   return output_operand_type;
+                 });
+
+  ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_SPLIT,
+                                                 input_indices, outputs, output_operand_types));
+
+  return Status::OK();
+}
+
+bool IsSupportedBatchMatMul(const NodeUnit& node_unit, int32_t nnapi_feature_level) {
+  // Currently, batch MatMul is composed of various operations including ANEURALNETWORKS_SPLIT which requires
+  // ANEURALNETWORKS_FEATURE_LEVEL_3.
+  const auto min_nnapi_feature_level = ANEURALNETWORKS_FEATURE_LEVEL_3;
+  if (nnapi_feature_level < min_nnapi_feature_level) {
+    LOGS_DEFAULT(VERBOSE) << "Minimum NNAPI feature level required: " << min_nnapi_feature_level
+                          << ", actual: " << nnapi_feature_level;
+    return false;
+  }
+
+  // Only support non-QDQ MatMul for now.
+  // TODO could be expanded to support QLinearMatMul and QDQ MatMul
+  if (node_unit.UnitType() != NodeUnit::Type::SingleNode ||
+      node_unit.OpType() != "MatMul") {
+    LOGS_DEFAULT(VERBOSE) << "Unsupported op type: "
+                          << (node_unit.UnitType() == NodeUnit::Type::QDQGroup ? "QDQ " : "") << node_unit.OpType();
+    return false;
+  }
+
+  const auto& inputs = node_unit.Inputs();
+
+  // Verify shapes
+  // A and B should have at least three dimensions* and have the same leading dimensions except for the last two.
+  // [*] Having two dimensions is valid for a MatMul but for simplicity we don't support it in the current batch
+  // MatMul implementation. That case is handled by the regular Gemm/MatMul op building logic.
+  Shape a_shape;
+  if (!GetShape(inputs[0].node_arg, a_shape)) {
+    return false;
+  }
+
+  Shape b_shape;
+  if (!GetShape(inputs[1].node_arg, b_shape)) {
+    return false;
+  }
+
+  if (a_shape.size() < 3 ||
+      a_shape.size() != b_shape.size() ||
+      !std::equal(a_shape.begin(), a_shape.end() - 2,
+                  b_shape.begin(), b_shape.end() - 2)) {
+    LOGS_DEFAULT(VERBOSE)
+        << "A and B must have at least three dimensions and have the same leading dimensions except for the last two. "
+        << "A shape: " << Shape2String(a_shape) << ", B shape: " << Shape2String(b_shape);
+    return false;
+  }
+
+  // Verify type
+  int32_t a_type;
+  if (!GetType(inputs[0].node_arg, a_type)) {
+    return false;
+  }
+
+  // Only support float for now.
+  // TODO could be expanded to support other types
+  if (a_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
+    LOGS_DEFAULT(VERBOSE) << "Unsupported element data type: " << a_type;
+    return false;
+  }
+
+  return true;
+}
+
+Status BuildBatchMatMul(ModelBuilder& model_builder, const NodeUnit& node_unit) {
+  // we will implement batch MatMul by composing NNAPI operations
+  // this could be replaced with ANEURALNETWORKS_BATCH_MATMUL when that is more widely supported
+
+  // assuming A and B have at least three dimensions and the same leading dimensions other than the last two
+
+  const auto& inputs = node_unit.Inputs();
+
+  Shape a_shape;
+  ORT_RETURN_IF_NOT(GetShape(inputs[0].node_arg, a_shape), "Failed to get A's shape.");
+
+  Shape b_shape;
+  ORT_RETURN_IF_NOT(GetShape(inputs[1].node_arg, b_shape), "Failed to get B's shape.");
+
+  const std::string& a = inputs[0].node_arg.Name();
+  const std::string& b = inputs[1].node_arg.Name();
+
+  const std::string& output = node_unit.Outputs()[0].node_arg.Name();
+
+  std::vector<std::string> gemm_a_inputs{a};
+  std::vector<std::string> gemm_b_inputs{b};
+
+  const auto m = a_shape[a_shape.size() - 2],
+             k = a_shape[a_shape.size() - 1],
+             n = b_shape[b_shape.size() - 1];
+
+  const bool reshape_leading_dimensions = a_shape.size() > 3;
+  const auto batch_size = ShapeSize(a_shape, 0, a_shape.size() - 2);
+
+  auto add_reshape = [&model_builder](const std::string& input, const Shape& new_shape,
+                                      const std::string& output) -> Status {
+    const std::string new_shape_name = model_builder.GetUniqueName(input + "/new_shape");
+    std::vector<int32_t> new_shape_i32{};
+    new_shape_i32.reserve(new_shape.size());
+    std::transform(new_shape.begin(), new_shape.end(), std::back_inserter(new_shape_i32),
+                   [](uint32_t d) { return gsl::narrow<int32_t>(d); });
+    ORT_RETURN_IF_ERROR(AddNnapiReshape(model_builder, input, new_shape_name, new_shape_i32, output, nullptr));
+    return Status::OK();
+  };
+
+  auto add_reshape_generate_output = [&model_builder, &add_reshape](const std::string& input, const Shape& new_shape,
+                                                                    std::string& output) -> Status {
+    std::string reshaped = model_builder.GetUniqueName(input + "/reshaped");
+    ORT_RETURN_IF_ERROR(add_reshape(input, new_shape, reshaped));
+    output = std::move(reshaped);
+    return Status::OK();
+  };
+
+  // collapse leading dimensions to a single one
+  if (reshape_leading_dimensions) {
+    const Shape a_new_shape_value = {batch_size, m, k},
+                b_new_shape_value = {batch_size, k, n};
+    std::string a_reshaped, b_reshaped;
+
+    ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_a_inputs.front(), a_new_shape_value, a_reshaped));
+    gemm_a_inputs.front() = a_reshaped;
+
+    ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_b_inputs.front(), b_new_shape_value, b_reshaped));
+    gemm_b_inputs.front() = b_reshaped;
+  }
+
+  // transpose B
+  {
+    const std::string b_new_perm = model_builder.GetUniqueName(b + "/new_perm"),
+                      b_transposed = model_builder.GetUniqueName(b + "/transposed");
+    ORT_RETURN_IF_ERROR(AddNnapiTranspose(model_builder, gemm_b_inputs.front(), b_new_perm, {0, 2, 1}, b_transposed));
+    gemm_b_inputs.front() = b_transposed;
+  }
+
+  // split batch
+  {
+    auto add_split = [&model_builder, batch_size](const std::string& input,
+                                                  std::vector<std::string>& outputs_result) -> Status {
+      std::vector<std::string> outputs;
+      outputs.reserve(batch_size);
+      for (size_t i = 0; i < batch_size; ++i) {
+        outputs.push_back(model_builder.GetUniqueName(MakeString(input, "/split_", i)));
+      }
+      ORT_RETURN_IF_ERROR(AddNnapiSplit(model_builder, input, 0, outputs));
+      outputs_result = std::move(outputs);
+      return Status::OK();
+    };
+
+    std::vector<std::string> a_split_outputs;
+    ORT_RETURN_IF_ERROR(add_split(gemm_a_inputs.front(), a_split_outputs));
+    gemm_a_inputs = std::move(a_split_outputs);
+
+    std::vector<std::string> b_split_outputs;
+    ORT_RETURN_IF_ERROR(add_split(gemm_b_inputs.front(), b_split_outputs));
+    gemm_b_inputs = std::move(b_split_outputs);
+  }
+
+  // GEMM per matrix pair
+  std::vector<std::string> gemm_outputs;
+  gemm_outputs.reserve(batch_size);
+  {
+    const std::string bias = model_builder.GetUniqueName(node_unit.Name() + "/zero_bias");
+    {
+      if (model_builder.GetOperandTypes().at(b).type != Type::TENSOR_FLOAT32) {
+        ORT_NOT_IMPLEMENTED("Only float input is supported now.");
+      }
+      const Shape bias_shape{n};
+      const std::vector<float> buffer(n, 0.0f);
+      const OperandType bias_operand_type(Type::TENSOR_FLOAT32, bias_shape);
+      ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(bias, buffer.data(), bias_operand_type));
+    }
+
+    auto add_fc = [&model_builder, &bias](const std::string& a, const std::string& b_transposed,
+                                          const std::string& output) -> Status {
+      const auto& operand_indices = model_builder.GetOperandIndices();
+      const auto& operand_types = model_builder.GetOperandTypes();
+      auto& shaper = model_builder.GetShaper();
+      std::vector<uint32_t> input_indices;
+      input_indices.push_back(operand_indices.at(a));             // A
+      input_indices.push_back(operand_indices.at(b_transposed));  // B'
+      input_indices.push_back(operand_indices.at(bias));          // C
+      int32_t fuse_code = ANEURALNETWORKS_FUSED_NONE;
+      ORT_RETURN_IF_ERROR(AddScalarOperand(model_builder, input_indices, fuse_code));
+
+      ORT_RETURN_IF_ERROR(shaper.FC(a, b_transposed, output));
+      const OperandType output_operand_type(operand_types.at(a).type, shaper[output]);
+      ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_FULLY_CONNECTED, input_indices,
+                                                     {output}, {output_operand_type}));
+      return Status::OK();
+    };
+
+    for (uint32_t i = 0; i < batch_size; ++i) {
+      const auto &gemm_a_input = gemm_a_inputs[i],
+                 &gemm_b_input = gemm_b_inputs[i];
+
+      // make inputs 2D ([1, x, y] -> [x, y])
+      std::string a_2d, b_transposed_2d;
+      ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_a_input, Shape{m, k}, a_2d));
+      ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_b_input, Shape{n, k}, b_transposed_2d));
+
+      const std::string gemm_output = model_builder.GetUniqueName(MakeString(node_unit.Name(), "/gemm_", i));
+      ORT_RETURN_IF_ERROR(add_fc(a_2d, b_transposed_2d, gemm_output));
+
+      // reshape output for concatenation ([x, y] -> [1, x, y])
+      std::string gemm_output_3d;
+      ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_output, Shape{1, m, n}, gemm_output_3d));
+
+      gemm_outputs.push_back(gemm_output_3d);
+    }
+  }
+
+  // concat batch
+  const std::string joined_gemm_output =
+      reshape_leading_dimensions ? model_builder.GetUniqueName(node_unit.Name() + "/joined_gemm_output") : output;
+  {
+    auto add_concat = [&model_builder](const std::vector<std::string>& inputs,
+                                       const std::string& output) -> Status {
+      const auto& operand_indices = model_builder.GetOperandIndices();
+      const auto& operand_types = model_builder.GetOperandTypes();
+      auto& shaper = model_builder.GetShaper();
+      std::vector<uint32_t> input_indices;
+      input_indices.reserve(inputs.size() + 1);
+      std::transform(inputs.begin(), inputs.end(), std::back_inserter(input_indices),
+                     [&operand_indices](const std::string& input) { return operand_indices.at(input); });
+      const int32_t axis = 0;
+      ORT_RETURN_IF_ERROR(AddScalarOperand(model_builder, input_indices, axis));
+      ORT_RETURN_IF_ERROR(shaper.Concat(inputs, axis, output));
+      OperandType output_operand_type = operand_types.at(inputs[0]);
+      output_operand_type.SetDimensions(shaper[output]);
+      ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_CONCATENATION,
+                                                     input_indices, {output}, {output_operand_type}));
+      return Status::OK();
+    };
+
+    ORT_RETURN_IF_ERROR(add_concat(gemm_outputs, joined_gemm_output));
+  }
+
+  // reshape to original dimensions
+  if (reshape_leading_dimensions) {
+    Shape new_shape = a_shape;
+    new_shape[new_shape.size() - 2] = m;
+    new_shape[new_shape.size() - 1] = n;
+    ORT_RETURN_IF_ERROR(add_reshape(joined_gemm_output, new_shape, output));
+  }
+
+  return Status::OK();
+}
+
+}  // namespace onnxruntime::nnapi::op_builder_helpers

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h

@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+#include "core/common/common.h"
+#include "core/providers/nnapi/nnapi_builtin/builders/helper.h"
+#include "core/providers/nnapi/nnapi_builtin/builders/model_builder.h"
+#include "core/providers/shared/node_unit/node_unit.h"
+
+namespace onnxruntime::nnapi::op_builder_helpers {
+
+// adds a scalar operand to the NNAPI model and appends its index to `input_indices`
+template <typename T>
+Status AddScalarOperand(ModelBuilder& model_builder, std::vector<uint32_t>& input_indices, T scalar_value) {
+  uint32_t index = 0;
+  ORT_RETURN_IF_ERROR(model_builder.AddOperandFromScalar(std::move(scalar_value), index));
+  input_indices.push_back(index);
+  return Status::OK();
+}
+
+// adds ANEURALNETWORKS_TRANSPOSE operation
+Status AddNnapiTranspose(ModelBuilder& model_builder,
+                         const std::string& data_input,
+                         const std::string& perm_input,
+                         const std::vector<int32_t>& perm,
+                         const std::string& output);
+
+// adds ANEURALNETWORKS_RESHAPE operation
+Status AddNnapiReshape(ModelBuilder& model_builder,
+                       const std::string& data_input,
+                       const std::string& shape_input, const std::vector<int32_t>& shape_value,
+                       const std::string& output, const Shape* output_shape);
+
+// adds ANEURALNETWORKS_SPLIT operation
+Status AddNnapiSplit(ModelBuilder& model_builder,
+                     const std::string& input,
+                     int32_t axis,
+                     const std::vector<std::string>& outputs);
+
+// checks whether batch MatMul in the given NodeUnit is supported by NNAPI EP
+bool IsSupportedBatchMatMul(const NodeUnit& node_unit, int32_t nnapi_feature_level);
+
+// builds a batch MatMul in the NNAPI model from the given NodeUnit
+// note: the pre-conditions of this function are checked in IsSupportedBatchMatMul()
+Status BuildBatchMatMul(ModelBuilder& model_builder, const NodeUnit& node_unit);
+
+}  // namespace onnxruntime::nnapi::op_builder_helpers

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc

@@ -8,9 +8,10 @@
 #include "core/framework/tensorprotoutils.h"
 #include "core/graph/graph.h"
 #include "core/providers/common.h"
+#include "core/providers/nnapi/nnapi_builtin/builders/helper.h"
+#include "core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h"
 #include "core/providers/shared/node_unit/node_unit.h"
 #include "core/providers/shared/utils/utils.h"
-#include "helper.h"
 
 namespace onnxruntime {
 namespace nnapi {
@@ -1431,6 +1432,16 @@ int GemmOpSupportChecker::GetMinSupportedOpSet(const NodeUnit& node_unit) const
 
 bool GemmOpSupportChecker::IsOpSupportedImpl(const InitializedTensorSet& initializers, const NodeUnit& node_unit,
                                              const OpSupportCheckParams& params) const {
+  // check batch matmul first, then fall back to checking single gemm/matmul
+  {
+    const bool is_supported_batch_matmul =
+        op_builder_helpers::IsSupportedBatchMatMul(node_unit, params.android_feature_level);
+    LOGS_DEFAULT(VERBOSE) << "Supported batch matmul: [" << is_supported_batch_matmul << "]";
+    if (is_supported_batch_matmul) {
+      return true;
+    }
+  }
+
   const auto& op_type = node_unit.OpType();
   const auto& inputs = node_unit.Inputs();
   const bool is_qlinear_matmul = op_type == "QLinearMatMul";

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc

@@ -123,6 +123,11 @@ Status Shaper::Concat(const std::vector<std::string>& input_names,
   SHAPER_FUNC(Concat, input_names, axis, output_name);
 }
 
+Status Shaper::Split(const std::string& input_name, int32_t axis,
+                     const std::vector<std::string>& output_names) {
+  SHAPER_FUNC(Split, input_name, axis, output_names);
+}
+
 Status Shaper::Squeeze(const std::string& input_name,
                        const std::vector<int32_t>& axes,
                        const std::string& output_name) {
@@ -362,13 +367,6 @@ Status Shaper::ConcatImpl(const std::vector<std::string>& input_names,
   std::vector<Shape> dimens;
   for (const auto& input_name : input_names) {
     const Shape& dimen = shape_map_.at(input_name);
-    if (!dimens.empty()) {
-      for (size_t i = 0; i < dimens[0].size(); i++) {
-        if ((int32_t)i == axis)
-          continue;
-      }
-    }
-
     dimens.push_back(dimen);
   }
 
@@ -388,6 +386,24 @@ Status Shaper::ConcatImpl(const std::vector<std::string>& input_names,
   return Status::OK();
 }
 
+Status Shaper::SplitImpl(const std::string& input_name, int32_t axis,
+                         const std::vector<std::string>& output_names) {
+  const auto& input_shape = shape_map_.at(input_name);
+  const auto count = output_names.size();
+
+  ORT_RETURN_IF_NOT(input_shape[axis] % count == 0,
+                    "count [", count, "] does not evenly divide dimension ", axis, " [", input_shape[axis], "]");
+
+  Shape output_shape = input_shape;
+  output_shape[axis] = input_shape[axis] / count;
+
+  for (const auto& output_name : output_names) {
+    shape_map_[output_name] = output_shape;
+  }
+
+  return Status::OK();
+}
+
 Status Shaper::SqueezeImpl(const std::string& input_name,
                            const std::vector<int32_t>& axes,
                            const std::string& output_name) {

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h

@@ -57,6 +57,8 @@ class Shaper {
 
   common::Status Concat(const std::vector<std::string>& input_names, const int32_t axis, const std::string& output_name);
 
+  common::Status Split(const std::string& input_name, int32_t axis, const std::vector<std::string>& output_names);
+
   common::Status Squeeze(const std::string& input_name, const std::vector<int32_t>& axes, const std::string& output_name);
 
   common::Status DepthToSpace(const std::string& input_name,
@@ -122,6 +124,7 @@ class Shaper {
   common::Status IdentityImpl(const std::string& input_name, const std::string& output_name);
   common::Status FCImpl(const std::string& input1_name, const std::string& input2_name, const std::string& output_name);
   common::Status ConcatImpl(const std::vector<std::string>& input_names, const int32_t axis, const std::string& output_name);
+  common::Status SplitImpl(const std::string& input_name, int32_t axis, const std::vector<std::string>& output_names);
   common::Status SqueezeImpl(const std::string& input_names, const std::vector<int32_t>& axes, const std::string& output_name);
   common::Status DepthToSpaceImpl(const std::string& input_names,
                                   const int32_t blocksize,

onnxruntime/test/providers/cpu/math/matmul_test.cc

@@ -4,6 +4,7 @@
 #include "gtest/gtest.h"
 #include "test/providers/provider_test_utils.h"
 #include "test/common/cuda_op_test_utils.h"
+#include "test/common/tensor_op_test_utils.h"
 #include "default_providers.h"
 
 namespace onnxruntime {
@@ -99,21 +100,49 @@ std::vector<MatMulTestData<T>> GenerateTestCases() {
        {3, 0},
        {}});
 
+  test_cases.push_back(
+      {"test 3D batch",
+       {3, 1, 3},
+       {3, 3, 2},
+       {3, 1, 2},
+       {
+           // clang-format off
+            10,  13,
+           100, 112,
+           298, 319,
+           // clang-format on
+       }});
+
+  test_cases.push_back(
+      {"test 4D batch",
+       {2, 2, 1, 3},
+       {2, 2, 3, 2},
+       {2, 2, 1, 2},
+       {
+           // clang-format off
+            10,  13,
+           100, 112,
+           298, 319,
+           604, 634,
+           // clang-format on
+       }});
+
   return test_cases;
 }
 
 template <typename T>
 void RunMatMulTest(int32_t opset_version, bool is_a_constant, bool is_b_constant) {
-  std::vector<T> common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11};
   for (auto t : GenerateTestCases<T>()) {
+    SCOPED_TRACE("test case: " + t.name);
+
     OpTester test("MatMul", opset_version);
 
     int64_t size0 = TensorShape::FromExistingBuffer(t.input0_dims).SizeHelper(0, t.input0_dims.size());
-    std::vector<T> input0_vals(common_input_vals.cbegin(), common_input_vals.cbegin() + size0);
+    std::vector<T> input0_vals = ValueRange<T>(size0);
     test.AddInput<T>("A", t.input0_dims, input0_vals, is_a_constant);
 
     int64_t size1 = TensorShape::FromExistingBuffer(t.input1_dims).SizeHelper(0, t.input1_dims.size());
-    std::vector<T> input1_vals(common_input_vals.cbegin(), common_input_vals.cbegin() + size1);
+    std::vector<T> input1_vals = ValueRange<T>(size1);
     test.AddInput<T>("B", t.input1_dims, input1_vals, is_b_constant);
 
     test.AddOutput<T>("Y", t.expected_dims, t.expected_vals);
@@ -121,7 +150,7 @@ void RunMatMulTest(int32_t opset_version, bool is_a_constant, bool is_b_constant
     // OpenVINO EP: Disabled temporarily matmul broadcasting not fully supported
     // Disable TensorRT because of unsupported data type
     std::unordered_set<std::string> excluded_providers{kTensorrtExecutionProvider, kOpenVINOExecutionProvider};
-    if (is_b_constant) {
+    if (t.name == "test 2D empty input") {
       // NNAPI: currently fails for the "test 2D empty input" case
       excluded_providers.insert(kNnapiExecutionProvider);
     }
@@ -136,7 +165,6 @@ void RunMatMulTest(int32_t opset_version) {
 
 TEST(MathOpTest, MatMulFloatType) {
   RunMatMulTest<float>(7, false, false);
-  RunMatMulTest<float>(7, false, true);
 }
 
 TEST(MathOpTest, MatMulDoubleType) {
@@ -190,7 +218,7 @@ TEST(MathOpTest, MatMul_Float16) {
   test.AddInput<MLFloat16>("A", {2, 4}, f_A);
   test.AddInput<MLFloat16>("B", {4, 3}, f_B);
   test.AddOutput<MLFloat16>("Y", {2, 3}, f_Y);
-  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});  //TensorRT: fp16 is not supported
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});  // TensorRT: fp16 is not supported
 }
 #endif
 
@@ -213,7 +241,7 @@ TEST(MathOpTest, MatMul_BFloat16) {
   execution_providers.push_back(DefaultCudaExecutionProvider());
 #elif USE_ROCM
   execution_providers.push_back(DefaultRocmExecutionProvider());
-#endif 
+#endif
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
 }
 #endif