Use FusedMatMul When Transpose is Between First Dim and Contiguous Batch Dims (#9734)

* fusedmatmul support transpose batches

* fix win build

* fix contrib op md

* more comments

docs/ContribOperators.md

@@ -1272,6 +1272,10 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dd>Whether A should be transposed on the last two dimensions before doing multiplication</dd>
 <dt><tt>transB</tt> : int</dt>
 <dd>Whether B should be transposed on the last two dimensions before doing multiplication</dd>
+<dt><tt>transBatchA</tt> : int</dt>
+<dd>Whether A should be transposed on the 1st dimension and batch dimensions (dim-1 to dim-rank-2) before doing multiplication</dd>
+<dt><tt>transBatchB</tt> : int</dt>
+<dd>Whether B should be transposed on the 1st dimension and batch dimensions (dim-1 to dim-rank-2) before doing multiplication</dd>
 </dl>
 
 #### Inputs

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -282,6 +282,10 @@ void FusedMatMulShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
   bool transa = transAAttr ? static_cast<int>(transAAttr->i()) != 0 : false;
   auto transBAttr = ctx.getAttribute("transB");
   bool transb = transBAttr ? static_cast<int>(transBAttr->i()) != 0 : false;
+  auto trans_batch_a_attr = ctx.getAttribute("transBatchA");
+  bool trans_batch_a = trans_batch_a_attr ? static_cast<int>(trans_batch_a_attr->i()) != 0 : false;
+  auto trans_batch_b_attr = ctx.getAttribute("transBatchB");
+  bool trans_batch_b = trans_batch_b_attr ? static_cast<int>(trans_batch_b_attr->i()) != 0 : false;
   int input1Idx = 0;
   int input2Idx = 1;
   if (!hasInputShape(ctx, input1Idx) || !hasInputShape(ctx, input2Idx)) {
@@ -309,11 +313,13 @@ void FusedMatMulShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
     // for vector input, transa does not make impact on the dim.
     shape0 = shape0_raw;
   } else {
-    for (int i = 0; i < rank0 - 2; ++i) {
+    int start = trans_batch_a ? 1 : 0;
+    int end = trans_batch_a ? rank0 - 1 : rank0 - 2;
+    for (int i = start; i < end; ++i) {
       *shape0.add_dim() = shape0_raw.dim(i);
     }
-    *shape0.add_dim() = shape0_raw.dim(transa ? rank0 - 1 : rank0 - 2);
-    *shape0.add_dim() = shape0_raw.dim(transa ? rank0 - 2 : rank0 - 1);
+    *shape0.add_dim() = shape0_raw.dim(transa ? rank0 - 1 : (trans_batch_a ? 0 : rank0 - 2));
+    *shape0.add_dim() = shape0_raw.dim(transa ? (trans_batch_a ? 0 : rank0 - 2) : rank0 - 1);
   }
 
   auto rank1 = shape1_raw.dim_size();
@@ -321,11 +327,13 @@ void FusedMatMulShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
     // for vector input, transb does not make impact on the dim.
     shape1 = shape1_raw;
   } else {
-    for (int i = 0; i < rank1 - 2; ++i) {
+    int start = trans_batch_b ? 1 : 0;
+    int end = trans_batch_b ? rank1 - 1 : rank1 - 2;
+    for (int i = start; i < end; ++i) {
       *shape1.add_dim() = shape1_raw.dim(i);
     }
-    *shape1.add_dim() = shape1_raw.dim(transb ? rank1 - 1 : rank1 - 2);
-    *shape1.add_dim() = shape1_raw.dim(transb ? rank1 - 2 : rank1 - 1);
+    *shape1.add_dim() = shape1_raw.dim(transb ? rank1 - 1 : (trans_batch_b ? 0 : rank1 - 2));
+    *shape1.add_dim() = shape1_raw.dim(transb ? (trans_batch_b ? 0 : rank1 - 2) : rank1 - 1);
   }
 
   ONNX_NAMESPACE::TensorShapeProto shapeL, shapeR;
@@ -2138,30 +2146,24 @@ Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-
       .SinceVersion(1)
       .Input(0, "A", "N-dimensional matrix A", "T")
       .Input(1, "B", "N-dimensional matrix B", "T")
-      .Attr(
-          "alpha",
-          "Scalar multiplier for the product of the input tensors.",
-          AttributeProto::FLOAT,
-          1.0f)
-      .Attr(
-          "transA",
-          "Whether A should be transposed on the last two dimensions before doing multiplication",
-          AttributeProto::INT,
-          static_cast<int64_t>(0))
-      .Attr(
-          "transB",
-          "Whether B should be transposed on the last two dimensions before doing multiplication",
-          AttributeProto::INT,
-          static_cast<int64_t>(0))
+      .Attr("alpha", "Scalar multiplier for the product of the input tensors.", AttributeProto::FLOAT, 1.0f)
+      .Attr("transA", "Whether A should be transposed on the last two dimensions before doing multiplication",
+            AttributeProto::INT, static_cast<int64_t>(0))
+      .Attr("transB", "Whether B should be transposed on the last two dimensions before doing multiplication",
+            AttributeProto::INT, static_cast<int64_t>(0))
+      .Attr("transBatchA",
+            "Whether A should be transposed on the 1st dimension and batch dimensions (dim-1 to dim-rank-2) before "
+            "doing multiplication",
+            AttributeProto::INT, static_cast<int64_t>(0))
+      .Attr("transBatchB",
+            "Whether B should be transposed on the 1st dimension and batch dimensions (dim-1 to dim-rank-2) before "
+            "doing multiplication",
+            AttributeProto::INT, static_cast<int64_t>(0))
       .Output(0, "Y", "Matrix multiply results", "T")
-      .TypeConstraint(
-          "T",
-          {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
-          "Constrain input and output types to float tensors.")
+      .TypeConstraint("T", {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
+                      "Constrain input and output types to float tensors.")
       .SetDoc(FusedMatMul_doc)
-      .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
-        FusedMatMulShapeInference(ctx);
-      });
+      .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { FusedMatMulShapeInference(ctx); });
 
   ONNX_CONTRIB_OPERATOR_SCHEMA(SparseToDenseMatMul)
       .SetDomain(kMSDomain)

onnxruntime/core/optimizer/matmul_transpose_fusion.cc

@@ -32,14 +32,25 @@ static bool GetTransposePerms(const Node& transpose_node, std::vector<int64_t>&
   return true;
 }
 
-static Node* GetTransposeNodeFromOutput(Graph& graph, NodeArg& node_arg) {
-  Node* trans_node = graph.GetMutableProducerNode(node_arg.Name());
-  if (trans_node == nullptr || trans_node->OpType() != "Transpose") {
-    return nullptr;
-  }
+// is_trans is whether to transpose the 2 dims used to MatMul.
+// is_trans_batch is whether to transpose 1st dim and batch dims.
+// Batch here has the same meaning in CUDA's GemmStridedBatched (other than the training batch concept),
+// i.e., all dims except the 2 dims used to MatMul, including from dim[0] to dim[rank-3].
+// For example (take lhs input as example):
+// is_trans=False, is_trans_batch=False:
+//   the input tensor is in shape [b0,...,bn,M,K], no Transpose (no fuse for this case)
+// is_trans=False, is_trans_batch=True:
+//   the input tensor is in shape [M,b0,...,bn,K], Transpose perm=[1,...,rank-2,0,rank-1]
+// is_trans=True , is_trans_batch=False:
+//   the input tensor is in shape [b0,...,bn,K,M], Transpose perm=[0,...,rank-3,rank-1,rank-2]
+// is_trans=True , is_trans_batch=True:
+//   the input tensor is in shape [K,b0,...,bn,M], Transpose perm=[1,...,rank-2,rank-1,0]
+static Node* GetTransposeNodeFromOutput(Graph& graph, NodeArg& node_arg, bool& is_trans, bool& is_trans_batch) {
+  is_trans = is_trans_batch = false;
 
-  // if the node has Graph output, skip it too
-  if (graph.NodeProducesGraphOutput(*trans_node)) {
+  // Skip if not a Transpose node or it has graph output.
+  Node* trans_node = graph.GetMutableProducerNode(node_arg.Name());
+  if (trans_node == nullptr || trans_node->OpType() != "Transpose" || graph.NodeProducesGraphOutput(*trans_node)) {
     return nullptr;
   }
 
@@ -48,28 +59,34 @@ static Node* GetTransposeNodeFromOutput(Graph& graph, NodeArg& node_arg) {
     return nullptr;
   }
 
-  int64_t rank = perms.size();
+  size_t rank = perms.size();
   if (rank < 2) {
     return nullptr;
   }
 
-  bool is_trans_on_last_two_dims = true;
-  for (int64_t i = 0; i < rank - 2; i++) {
-    if (perms[static_cast<size_t>(i)] != i) {
-      is_trans_on_last_two_dims = false;
-      break;
-    }
-  }
-
-  if (is_trans_on_last_two_dims) {
-    // rank is atleast 2 (checked above) and so it is safe to cast (rank - 2) and (rank - 1) to size_t
-    is_trans_on_last_two_dims = perms[static_cast<size_t>(rank - 2)] == rank - 1 && perms[static_cast<size_t>(rank - 1)] == rank - 2;
-  }
-
-  if (!is_trans_on_last_two_dims) {
+  // We can fuse the Transpose node only when the last axis of original tensor is within the last two dims after transpose.
+  int64_t last_axis = static_cast<int64_t>(rank) - 1;
+  size_t last_axis_index = perms[rank - 1] == last_axis ? rank - 1 : perms[rank - 2] == last_axis ? rank - 2 : rank;
+  if (last_axis_index == rank) {
     return nullptr;
   }
 
+  // Transpose node can be fused to MatMul when the batch dims keep same relative orders before and after transpose.
+  // But if they are not contiguous, after the fusion, we can only use GemmBatched instead of GemmStridedBatched,
+  // which may have perf issue. To keep it simple, we will fuse only when batch dimensions are contiguous.
+  if (rank >= 3) {
+    if (perms[0] != 0 && perms[0] != 1) {
+      return nullptr;
+    }
+    for (size_t i = 0; i < rank - 3; ++i) {
+      if (perms[i] + 1 != perms[i + 1]) {
+        return nullptr;
+      }
+    }
+  }
+
+  is_trans = last_axis_index == rank - 2;
+  is_trans_batch = rank >= 3 && perms[0] == 1;
   return trans_node;
 }
 
@@ -118,9 +135,10 @@ static size_t UpdateConsumerCount(Graph& graph, NodeArg* target, std::unordered_
 */
 static Node* ReorderCastAndTranspose(Graph& graph, Node* cast,
                                      std::unordered_map<NodeArg*, size_t>& consumer_count,
-                                     std::deque<onnxruntime::NodeIndex>& removed_nodes) {
+                                     std::deque<onnxruntime::NodeIndex>& removed_nodes,
+                                     bool& is_trans, bool& is_trans_batch) {
   ORT_ENFORCE(cast != nullptr);
-  auto transpose = GetTransposeNodeFromOutput(graph, *cast->MutableInputDefs()[0]);
+  auto transpose = GetTransposeNodeFromOutput(graph, *cast->MutableInputDefs()[0], is_trans, is_trans_batch);
   if (transpose == nullptr) {
     return nullptr;
   }
@@ -290,26 +308,58 @@ Status MatmulTransposeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_
     if (!IsAllowedFusedMatMulDataType(static_cast<ONNX_NAMESPACE::TensorProto_DataType>(left_type))) {
       continue;
     }
-    auto left = GetTransposeNodeFromOutput(graph, *left_input);
+
+    bool is_trans_left = false;
+    bool is_trans_batch_left = false;
+    Node* left = nullptr;
+    // If it's already a FusedMatMul with transBatchA is true, don't fuse it.
+    if (node.OpType() != "FusedMatMul" || node.GetAttributes().at("transBatchA").i() == 0) {
+      left = GetTransposeNodeFromOutput(graph, *left_input, is_trans_left, is_trans_batch_left);
+      if (!left) {
+        Node* left_node = graph.GetMutableProducerNode(left_input->Name());
+        if (left_node && left_node->OpType() == "Cast") {
+          left = ReorderCastAndTranspose(graph, left_node, consumer_count, removed_nodes, is_trans_left,
+                                         is_trans_batch_left);
+        }
+      }
+    }
 
     NodeArg* right_input = node.MutableInputDefs()[1];
     auto right_type = right_input->TypeAsProto()->tensor_type().elem_type();
     if (!IsAllowedFusedMatMulDataType(static_cast<ONNX_NAMESPACE::TensorProto_DataType>(right_type))) {
       continue;
     }
-    auto right = GetTransposeNodeFromOutput(graph, *right_input);
 
-    if (!left) {
-      Node* left_node = graph.GetMutableProducerNode(left_input->Name());
-      if (left_node && left_node->OpType() == "Cast") {
-        left = ReorderCastAndTranspose(graph, left_node, consumer_count, removed_nodes);
+    bool is_trans_right = false;
+    bool is_trans_batch_right = false;
+    Node* right = nullptr;
+    // If it's already a FusedMatMul with transBatchB is true, don't fuse it.
+    if (node.OpType() != "FusedMatMul" || node.GetAttributes().at("transBatchB").i() == 0) {
+      right = GetTransposeNodeFromOutput(graph, *right_input, is_trans_right, is_trans_batch_right);
+      if (!right) {
+        Node* right_node = graph.GetMutableProducerNode(right_input->Name());
+        if (right_node && right_node->OpType() == "Cast") {
+          right = ReorderCastAndTranspose(graph, right_node, consumer_count, removed_nodes, is_trans_right,
+                                          is_trans_batch_right);
+        }
       }
     }
 
-    if (!right) {
-      Node* right_node = graph.GetMutableProducerNode(right_input->Name());
-      if (right_node && right_node->OpType() == "Cast") {
-        right = ReorderCastAndTranspose(graph, right_node, consumer_count, removed_nodes);
+    // When the rank of two inputs are not equal, we need to pad 1 to one of them for MutMul.
+    // For example, if padding 1 is the "M" dim for left input, set is_trans_batch to true is not correct logically.
+    // To keep it simple, if any side of is_trans_batch is true, we require both side has same rank.
+    if (is_trans_batch_left || is_trans_batch_right) {
+      auto shape_left = left_input->Shape();
+      auto shape_right = right_input->Shape();
+      if (!shape_left || !shape_right || shape_left->dim_size() != shape_right->dim_size()) {
+        if (is_trans_batch_left) {
+          is_trans_left = is_trans_batch_left = false;
+          left = nullptr;
+        }
+        if (is_trans_batch_right) {
+          is_trans_right = is_trans_batch_right = false;
+          right= nullptr;
+        }
       }
     }
 
@@ -339,16 +389,18 @@ Status MatmulTransposeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_
                                       "fused MatMul and Transpose ",
                                       input_defs,
                                       output_defs, {}, kMSDomain);
-    bool transpose_left = (left != nullptr);
-    bool transpose_right = (right != nullptr);
     float alpha = 1.0f;
     if (node.OpType() == "FusedMatMul") {
-      transpose_left ^= static_cast<bool>(node.GetAttributes().at("transA").i());
-      transpose_right ^= static_cast<bool>(node.GetAttributes().at("transB").i());
+      is_trans_left ^= static_cast<bool>(node.GetAttributes().at("transA").i());
+      is_trans_right ^= static_cast<bool>(node.GetAttributes().at("transB").i());
+      is_trans_batch_left ^= static_cast<bool>(node.GetAttributes().at("transBatchA").i());
+      is_trans_batch_right ^= static_cast<bool>(node.GetAttributes().at("transBatchB").i());
       alpha = node.GetAttributes().at("alpha").f();
     }
-    matmul_node.AddAttribute("transA", static_cast<int64_t>(transpose_left));
-    matmul_node.AddAttribute("transB", static_cast<int64_t>(transpose_right));
+    matmul_node.AddAttribute("transA", static_cast<int64_t>(is_trans_left));
+    matmul_node.AddAttribute("transB", static_cast<int64_t>(is_trans_right));
+    matmul_node.AddAttribute("transBatchA", static_cast<int64_t>(is_trans_batch_left));
+    matmul_node.AddAttribute("transBatchB", static_cast<int64_t>(is_trans_batch_right));
     matmul_node.AddAttribute("alpha", alpha);
     // Assign provider to this new node. Provider should be same as the provider for old node.
     matmul_node.SetExecutionProviderType(node.GetExecutionProviderType());

onnxruntime/core/providers/cpu/math/matmul.cc

@@ -169,7 +169,7 @@ Status MatMul<float>::Compute(OpKernelContext* ctx) const {
   const bool trans_b = trans_b_attr_ && b_shape.NumDimensions() != 1;
 
   MatMulComputeHelper helper;
-  ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, trans_a, trans_b));
+  ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, trans_a, trans_b, trans_batch_a_, trans_batch_b_));
   Tensor* y = ctx->Output(0, helper.OutputShape());
 
   // Bail out early if the output is going to be empty
@@ -184,8 +184,8 @@ Status MatMul<float>::Compute(OpKernelContext* ctx) const {
   const size_t M = static_cast<size_t>(helper.M());
   const size_t N = static_cast<size_t>(helper.N());
   const size_t K = static_cast<size_t>(helper.K());
-  const size_t lda = static_cast<int>(trans_a ? M : K);
-  const size_t ldb = static_cast<int>(trans_b ? K : N);
+  const size_t lda = helper.Lda(trans_a);
+  const size_t ldb = helper.Ldb(trans_b);
 
   std::vector<MLAS_SGEMM_DATA_PARAMS> data(max_len);
   for (size_t i = 0; i < max_len; i++) {

onnxruntime/core/providers/cpu/math/matmul.h

@@ -22,14 +22,18 @@ class MatMul<float> final : public OpKernel {
     info.GetAttrOrDefault<int64_t>("transA", &trans_a_attr_, 0);
     info.GetAttrOrDefault<int64_t>("transB", &trans_b_attr_, 0);
     info.GetAttrOrDefault<float>("alpha", &alpha_attr_, 1.0);
+    int64_t trans_batch_a_attr, trans_batch_b_attr;
+    info.GetAttrOrDefault<int64_t>("transBatchA", &trans_batch_a_attr, 0);
+    info.GetAttrOrDefault<int64_t>("transBatchB", &trans_batch_b_attr, 0);
+    trans_batch_a_ = trans_batch_a_attr != 0;
+    trans_batch_b_ = trans_batch_b_attr != 0;
   }
 
   Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc,
                  /*out*/ bool& is_packed,
                  /*out*/ PrePackedWeights* prepacked_weights) override;
 
-  Status UseSharedPrePackedBuffers(std::vector<BufferUniquePtr>& prepacked_buffers,
-                                   int input_idx,
+  Status UseSharedPrePackedBuffers(std::vector<BufferUniquePtr>& prepacked_buffers, int input_idx,
                                    /*out*/ bool& used_shared_buffers) override;
 
   Status Compute(OpKernelContext* context) const override;
@@ -42,6 +46,8 @@ class MatMul<float> final : public OpKernel {
   float alpha_attr_;
   int64_t trans_a_attr_;
   int64_t trans_b_attr_;
+  bool trans_batch_a_;
+  bool trans_batch_b_;
 };
 
 }  // namespace onnxruntime

onnxruntime/core/providers/cpu/math/matmul_helper.h

@@ -22,8 +22,12 @@ inline void TensorShapeCopyDims(const TensorShape& shape, int64_t* dims, size_t
 
 class MatMulComputeHelper {
  public:
-  Status Compute(const TensorShape& left_shape, const TensorShape& right_shape,
-                 bool transa = false, bool transb = false) {
+  // fill_offsets is to control if to fill offsets here.
+  // For CUDA/ROCM kernel when we can use GemmStridedBatched, we don't need to fill the offsets.
+  Status Compute(const TensorShape& orig_left_shape, const TensorShape& orig_right_shape,
+                 bool transa = false, bool transb = false,
+                 bool trans_batch_a = false, bool trans_batch_b = false,
+                 bool fill_offsets = true) {
     // Following numpy.matmul for shape inference:
     // https://docs.scipy.org/doc/numpy/reference/generated/numpy.matmul.html
     // The behavior depends on the arguments in the following way.
@@ -32,8 +36,8 @@ class MatMulComputeHelper {
     // * If the first argument is 1 - D, it is promoted to a matrix by prepending a 1 to its dimensions.After matrix multiplication the prepended 1 is removed.
     // * If the second argument is 1 - D, it is promoted to a matrix by appending a 1 to its dimensions.After matrix multiplication the appended 1 is removed.
 
-    size_t left_num_dims = left_shape.NumDimensions();
-    size_t right_num_dims = right_shape.NumDimensions();
+    size_t left_num_dims = orig_left_shape.NumDimensions();
+    size_t right_num_dims = orig_right_shape.NumDimensions();
     ORT_RETURN_IF_NOT(left_num_dims >= 1 && right_num_dims >= 1, "left_num_dims and right_num_dims must be >= 1");
 
     // Special cases below for right_shape being 2D and left_shape > 2D by flattening left_shape to 2D
@@ -42,22 +46,59 @@ class MatMulComputeHelper {
     // A: [M1, M2, ... K], B: [N, K]^T
     // A: [M1, M2, ... K], B: [1, ..., 1, K, N]
     // A: [M1, M2, ... K], B: [1, ..., 1, N, K]^T
-    if (!transa && left_num_dims >= 2 && right_num_dims >= 2 && left_num_dims >= right_num_dims &&
-        right_shape.SizeToDimension(right_num_dims - 1) == right_shape[right_num_dims - 2]) {
-      M_ = static_cast<ptrdiff_t>(left_shape.SizeToDimension(left_num_dims - 1));
-      K_ = static_cast<ptrdiff_t>(left_shape[left_num_dims - 1]);
-      N_ = static_cast<ptrdiff_t>(transb ? right_shape[right_num_dims - 2] : right_shape[right_num_dims - 1]);
-      output_shape_ = left_shape;
+    if (!transa && !trans_batch_a && !trans_batch_b && left_num_dims >= 2 && right_num_dims >= 2 && left_num_dims >= right_num_dims &&
+        orig_right_shape.SizeToDimension(right_num_dims - 1) == orig_right_shape[right_num_dims - 2]) {
+      M_ = static_cast<ptrdiff_t>(orig_left_shape.SizeToDimension(left_num_dims - 1));
+      K_ = static_cast<ptrdiff_t>(orig_left_shape[left_num_dims - 1]);
+      N_ = static_cast<ptrdiff_t>(transb ? orig_right_shape[right_num_dims - 2] : orig_right_shape[right_num_dims - 1]);
+      output_shape_ = orig_left_shape;
       output_shape_[left_num_dims - 1] = N_;
       output_offsets_ = {0};
       left_offsets_ = {0};
       right_offsets_ = {0};
-      ORT_RETURN_IF_NOT(K_ == right_shape[right_num_dims - 2] ||
-                            (transb && K_ == right_shape[right_num_dims - 1]),
+      ORT_RETURN_IF_NOT((!transb && K_ == orig_right_shape[right_num_dims - 2]) ||
+                            (transb && K_ == orig_right_shape[right_num_dims - 1]),
                         "MatMul dimension mismatch");
       return Status::OK();
     }
 
+    std::vector<int64_t> dims_left(left_num_dims);
+    std::vector<int64_t> dims_right(right_num_dims);
+    orig_left_shape.CopyDims(&dims_left[0], left_num_dims);
+    orig_right_shape.CopyDims(&dims_right[0], right_num_dims);
+    left_stride_factor_ = right_stride_factor_ = 1;
+    left_ld_factor_ = right_ld_factor_ = 1;
+
+    if (trans_batch_a || trans_batch_b) {
+      ORT_ENFORCE(left_num_dims > 2 && left_num_dims == right_num_dims,
+                  "Two inputs should have same rank and rank >= 3 if transBatchA or transBatchB is true");
+      // trans_batch_a means the input tensor is either [M,b0,...,bn,K] or [K,b0,...,bn,M].
+      // Switch the dim[0] and batch dims here.
+      if (trans_batch_a) {
+        int64_t leading_dim = dims_left[0];
+        for (size_t i = 0; i < left_num_dims - 2; ++i) {
+          dims_left[i] = dims_left[i + 1];
+          left_ld_factor_ *= static_cast<int>(dims_left[i]);
+        }
+        dims_left[left_num_dims - 2] = leading_dim;
+        left_stride_factor_ = static_cast<size_t>(leading_dim);
+      }
+      // trans_batch_b means the input tensor is either [K,b0,...,bn,N] or [N,b0,...,bn,K].
+      // Switch the dim[0] and batch dims here.
+      if (trans_batch_b) {
+        int64_t leading_dim = dims_right[0];
+        for (size_t i = 0; i < right_num_dims - 2; ++i) {
+          dims_right[i] = dims_right[i + 1];
+          right_ld_factor_ *= static_cast<int>(dims_right[i]);
+        }
+        dims_right[right_num_dims - 2] = leading_dim;
+        right_stride_factor_ = static_cast<size_t>(leading_dim);
+      }
+    }
+
+    TensorShape left_shape(dims_left);
+    TensorShape right_shape(dims_right);
+
     bool has_1D_input = (left_num_dims == 1 || right_num_dims == 1);
 
     size_t num_input_dims = std::max(left_num_dims, right_num_dims);
@@ -146,7 +187,7 @@ class MatMulComputeHelper {
     output_shape_ = TensorShape(output_dims);
 
     // compute broadcast offsets
-    ComputeBroadcastOffsets();
+    ComputeBroadcastOffsets(fill_offsets);
 
     return Status::OK();
   }
@@ -178,21 +219,8 @@ class MatMulComputeHelper {
     return Status::OK();
   }
 
- private:
-  void ComputeBroadcastOffsets() {
-    num_broadcasted_dims_ = left_padded_dims_.size() - 2;
-
-    if (num_broadcasted_dims_ == 0) {
-      left_offsets_ = {0};
-      right_offsets_ = {0};
-      output_offsets_ = {0};
-      return;
-    }
-
-    left_mat_size_ = M_ * K_;
-    right_mat_size_ = K_ * N_;
-    output_mat_size_ = M_ * N_;
-
+  // Move this piece of code to public so that we don't need to call this if we can use GemmStridedBatched.
+  void FillOffsets() {
     // stride in mats and dims for broadcasting
     left_padded_strides_.resize(num_broadcasted_dims_);
     right_padded_strides_.resize(num_broadcasted_dims_);
@@ -214,6 +242,26 @@ class MatMulComputeHelper {
     RecursiveFill(0, 0, 0, 0);
   }
 
+ private:
+  void ComputeBroadcastOffsets(bool fill_offsets) {
+    num_broadcasted_dims_ = left_padded_dims_.size() - 2;
+
+    if (num_broadcasted_dims_ == 0) {
+      left_offsets_ = {0};
+      right_offsets_ = {0};
+      output_offsets_ = {0};
+      return;
+    }
+
+    left_mat_size_ = M_ * K_ / left_stride_factor_;
+    right_mat_size_ = K_ * N_ / right_stride_factor_;
+    output_mat_size_ = M_ * N_;
+
+    if (fill_offsets) {
+      FillOffsets();
+    }
+  }
+
   void RecursiveFill(size_t idx_dim, size_t idx_left, size_t idx_right, size_t idx_out) {
     if (idx_dim == num_broadcasted_dims_) {
       left_offsets_[idx_out] = idx_left * left_mat_size_;
@@ -260,6 +308,11 @@ class MatMulComputeHelper {
   std::vector<size_t> right_zp_offsets_;
   std::vector<size_t> right_scale_offsets_;
 
+  size_t left_stride_factor_ = 1;
+  size_t right_stride_factor_ = 1;
+  int left_ld_factor_ = 1;
+  int right_ld_factor_ = 1;
+
  public:
   // output shape
   const TensorShape& OutputShape() const {
@@ -281,6 +334,18 @@ class MatMulComputeHelper {
     return K_;
   }
 
+  int Lda(bool is_trans) const {
+    return (is_trans ? static_cast<int>(M_) : static_cast<int>(K_)) * left_ld_factor_;
+  }
+
+  int Ldb(bool is_trans) const {
+    return (is_trans ? static_cast<int>(K_) : static_cast<int>(N_)) * right_ld_factor_; 
+  }
+
+  int Ldc() const {
+    return static_cast<int>(N_);
+  }
+
   // Batched Gemm offsets in left matrices
   const std::vector<size_t>& LeftOffsets() const {
     return left_offsets_;

onnxruntime/core/providers/cuda/math/matmul.cc

@@ -48,7 +48,7 @@ REGISTER_KERNEL_TYPED(BFloat16)
 // StridedBatchedGemm can be used for the following GEMM computation
 // C[pnm] = A[pnk]*B[km] or C[pnm] = A[pnk]*B[pkm]
 static bool CanUseStridedBatchedGemm(const TensorShape& left_shape, const TensorShape& right_shape,
-                                     bool transa, bool transb,
+                                     bool transa, bool transb, bool trans_batch_a, bool trans_batch_b,
                                      int64_t& stride_A, int64_t& stride_B, int64_t& stride_C, int64_t& batch_count) {
   size_t left_num_dims = left_shape.NumDimensions();
   size_t right_num_dims = right_shape.NumDimensions();
@@ -57,25 +57,33 @@ static bool CanUseStridedBatchedGemm(const TensorShape& left_shape, const Tensor
     return false;
   }
 
+  size_t left_leading_axis = trans_batch_a ? 0 : left_num_dims - 2;
+  size_t right_leading_axis = trans_batch_b ? 0 : right_num_dims - 2;
   int64_t left_p = left_shape.SizeToDimension(left_num_dims - 2);
-  int64_t left_k = transa ? left_shape[left_num_dims - 2] : left_shape[left_num_dims - 1];
+  if (trans_batch_a) {
+    left_p = left_p * left_shape[left_num_dims - 2] / left_shape[0];
+  }
+  int64_t left_k = transa ? left_shape[left_leading_axis] : left_shape[left_num_dims - 1];
 
   if (right_num_dims >= 3) {
     int64_t right_p = right_shape.SizeToDimension(right_num_dims - 2);
+    if (trans_batch_b) {
+      right_p = right_p * right_shape[right_num_dims - 2] / right_shape[0];
+    }
     if (left_p != right_p) {
       return false;
     }
   }
 
-  int64_t right_k = transb ? right_shape[right_num_dims - 1] : right_shape[right_num_dims - 2];
+  int64_t right_k = transb ? right_shape[right_num_dims - 1] : right_shape[right_leading_axis];
   if (left_k != right_k) {
     return false;
   }
 
-  int64_t n = transa ? left_shape[left_num_dims - 1] : left_shape[left_num_dims - 2];
-  int64_t m = transb ? right_shape[right_num_dims - 2] : right_shape[right_num_dims - 1];
-  stride_A = n * left_k;
-  stride_B = right_num_dims == 2 ? 0 : right_k * m;
+  int64_t n = transa ? left_shape[left_num_dims - 1] : left_shape[left_leading_axis];
+  int64_t m = transb ? right_shape[right_leading_axis] : right_shape[right_num_dims - 1];
+  stride_A = n * left_k / (trans_batch_a ? left_shape[0] : 1);
+  stride_B = right_num_dims == 2 ? 0 : right_k * m / (trans_batch_b ? right_shape[0] : 1);
   stride_C = n * m;
   batch_count = left_p;
   return true;
@@ -100,7 +108,7 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
   }
 
   MatMulComputeHelper helper;
-  ORT_RETURN_IF_ERROR(helper.Compute(left_X->Shape(), right_X->Shape(), transa, transb));
+  ORT_RETURN_IF_ERROR(helper.Compute(left_X->Shape(), right_X->Shape(), transa, transb, trans_batch_a_, trans_batch_b_, false));
 
   Tensor* Y = ctx->Output(0, helper.OutputShape());
 
@@ -113,9 +121,9 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
 
   cublasOperation_t transA = transa ? CUBLAS_OP_T : CUBLAS_OP_N;
   cublasOperation_t transB = transb ? CUBLAS_OP_T : CUBLAS_OP_N;
-  const int lda = transa ? static_cast<int>(helper.M()) : static_cast<int>(helper.K());
-  const int ldb = transb ? static_cast<int>(helper.K()) : static_cast<int>(helper.N());
-  const int ldc = static_cast<int>(helper.N());
+  const int lda = helper.Lda(transa);
+  const int ldb = helper.Ldb(transb);
+  const int ldc = helper.Ldc();
   int64_t stride_A, stride_B, stride_C, batch_count;
   auto& device_prop = GetDeviceProp();
   if (helper.OutputOffsets().size() == 1) {
@@ -137,7 +145,7 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
         device_prop));
     return Status::OK();
   } else if (CanUseStridedBatchedGemm(left_X->Shape(), right_X->Shape(),
-                                      transa, transb, stride_A, stride_B, stride_C, batch_count)) {
+                                      transa, transb, trans_batch_a_, trans_batch_b_, stride_A, stride_B, stride_C, batch_count)) {
     CUBLAS_RETURN_IF_ERROR(cublasGemmStridedBatchedHelper(Base::CublasHandle(),
                                                           transB,
                                                           transA,
@@ -161,6 +169,8 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
     return Status::OK();
   }
 
+  // Fill offsets when needed.
+  helper.FillOffsets();
   CudaAsyncBuffer<const CudaT*> left_arrays(this, helper.LeftOffsets().size());
   CudaAsyncBuffer<const CudaT*> right_arrays(this, helper.RightOffsets().size());
   CudaAsyncBuffer<CudaT*> output_arrays(this, helper.OutputOffsets().size());

onnxruntime/core/providers/cuda/math/matmul.h

@@ -16,8 +16,9 @@ class MatMul final : public CudaKernel {
       : CudaKernel(info),
         alpha_{info.GetAttrOrDefault<float>("alpha", 1.0f)},
         trans_A_{info.GetAttrOrDefault<int64_t>("transA", 0) != 0},
-        trans_B_{info.GetAttrOrDefault<int64_t>("transB", 0) != 0} {
-  }
+        trans_B_{info.GetAttrOrDefault<int64_t>("transB", 0) != 0},
+        trans_batch_a_{info.GetAttrOrDefault<int64_t>("transBatchA", 0) != 0},
+        trans_batch_b_{info.GetAttrOrDefault<int64_t>("transBatchB", 0) != 0} {}
 
   Status ComputeInternal(OpKernelContext* context) const override;
 
@@ -25,6 +26,8 @@ class MatMul final : public CudaKernel {
   const float alpha_;
   const bool trans_A_;
   const bool trans_B_;
+  const bool trans_batch_a_;
+  const bool trans_batch_b_;
 };
 }  // namespace cuda
 }  // namespace onnxruntime

onnxruntime/core/providers/rocm/math/matmul.cc

@@ -45,7 +45,7 @@ REGISTER_KERNEL_TYPED(MLFloat16)
 // StridedBatchedGemm can be used for the following GEMM computation
 // C[pnm] = A[pnk]*B[km] or C[pnm] = A[pnk]*B[pkm]
 static bool CanUseStridedBatchedGemm(const TensorShape& left_shape, const TensorShape& right_shape,
-                                     bool transa, bool transb,
+                                     bool transa, bool transb, bool trans_batch_a, bool trans_batch_b,
                                      int64_t& stride_A, int64_t& stride_B, int64_t& stride_C, int64_t& batch_count) {
   size_t left_num_dims = left_shape.NumDimensions();
   size_t right_num_dims = right_shape.NumDimensions();
@@ -54,25 +54,33 @@ static bool CanUseStridedBatchedGemm(const TensorShape& left_shape, const Tensor
     return false;
   }
 
+  size_t left_leading_axis = trans_batch_a ? 0 : left_num_dims - 2;
+  size_t right_leading_axis = trans_batch_b ? 0 : right_num_dims - 2;
   int64_t left_p = left_shape.SizeToDimension(left_num_dims - 2);
-  int64_t left_k = transa ? left_shape[left_num_dims - 2] : left_shape[left_num_dims - 1];
+  if (trans_batch_a) {
+    left_p = left_p * left_shape[left_num_dims - 2] / left_shape[0];
+  }
+  int64_t left_k = transa ? left_shape[left_leading_axis] : left_shape[left_num_dims - 1];
 
   if (right_num_dims >= 3) {
     int64_t right_p = right_shape.SizeToDimension(right_num_dims - 2);
+    if (trans_batch_b) {
+      right_p = right_p * right_shape[right_num_dims - 2] / right_shape[0];
+    }
     if (left_p != right_p) {
       return false;
     }
   }
 
-  int64_t right_k = transb ? right_shape[right_num_dims - 1] : right_shape[right_num_dims - 2];
+  int64_t right_k = transb ? right_shape[right_num_dims - 1] : right_shape[right_leading_axis];
   if (left_k != right_k) {
     return false;
   }
 
-  int64_t n = transa ? left_shape[left_num_dims - 1] : left_shape[left_num_dims - 2];
-  int64_t m = transb ? right_shape[right_num_dims - 2] : right_shape[right_num_dims - 1];
-  stride_A = n * left_k;
-  stride_B = right_num_dims == 2 ? 0 : right_k * m;
+  int64_t n = transa ? left_shape[left_num_dims - 1] : left_shape[left_leading_axis];
+  int64_t m = transb ? right_shape[right_leading_axis] : right_shape[right_num_dims - 1];
+  stride_A = n * left_k / (trans_batch_a ? left_shape[0] : 1);
+  stride_B = right_num_dims == 2 ? 0 : right_k * m / (trans_batch_b ? right_shape[0] : 1);
   stride_C = n * m;
   batch_count = left_p;
   return true;
@@ -97,7 +105,7 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
   }
 
   MatMulComputeHelper helper;
-  ORT_RETURN_IF_ERROR(helper.Compute(left_X->Shape(), right_X->Shape(), transa, transb));
+  ORT_RETURN_IF_ERROR(helper.Compute(left_X->Shape(), right_X->Shape(), transa, transb, trans_batch_a_, trans_batch_b_, false));
 
   Tensor* Y = ctx->Output(0, helper.OutputShape());
 
@@ -110,9 +118,9 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
 
   rocblas_operation transA = transa ? rocblas_operation_transpose : rocblas_operation_none;
   rocblas_operation transB = transb ? rocblas_operation_transpose : rocblas_operation_none;
-  const int lda = transa ? static_cast<int>(helper.M()) : static_cast<int>(helper.K());
-  const int ldb = transb ? static_cast<int>(helper.K()) : static_cast<int>(helper.N());
-  const int ldc = static_cast<int>(helper.N());
+  const int lda = helper.Lda(transa);
+  const int ldb = helper.Ldb(transb);
+  const int ldc = helper.Ldc();
   int64_t stride_A, stride_B, stride_C, batch_count;
 
   if (helper.OutputOffsets().size() == 1) {
@@ -133,7 +141,7 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
         ldc));
     return Status::OK();
   } else if (CanUseStridedBatchedGemm(left_X->Shape(), right_X->Shape(),
-                                      transa, transb, stride_A, stride_B, stride_C, batch_count)) {
+                                      transa, transb, trans_batch_a_, trans_batch_b_, stride_A, stride_B, stride_C, batch_count)) {
     ROCBLAS_RETURN_IF_ERROR(rocblasGemmStridedBatchedHelper(Base::RocblasHandle(),
                                                           transB,
                                                           transA,
@@ -155,6 +163,8 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
     return Status::OK();
   }
 
+  // Fill offsets when needed.
+  helper.FillOffsets();
   RocmAsyncBuffer<const HipT*> left_arrays(this, helper.LeftOffsets().size());
   RocmAsyncBuffer<const HipT*> right_arrays(this, helper.RightOffsets().size());
   RocmAsyncBuffer<HipT*> output_arrays(this, helper.OutputOffsets().size());

onnxruntime/test/contrib_ops/fused_matmul_op_test.cc

@@ -91,71 +91,121 @@ std::vector<MatMulTestData<T>> GenerateSimpleTestCases() {
        {0, 4},
        {}});
 
+  test_cases.push_back(
+      {"test 4D",
+       {2, 2, 2, 2},
+       {2, 2, 2, 2},
+       {2, 2, 2, 2},
+       {2, 3, 6, 11, 46, 55, 66, 79, 154, 171, 190, 211, 326, 351, 378, 407}});
+
+  test_cases.push_back(
+      {"test 4D and broadcast",
+       {1, 2, 3, 2},
+       {3, 2, 2, 1},
+       {3, 2, 3, 1},
+       {1, 3, 5, 33, 43, 53, 5, 23, 41, 85, 111, 137, 9, 43, 77, 137, 179, 221}});
+
   return test_cases;
 }
 
-/* Transpose the last two dimentions */
+// [batch,N,M]->[batch,M,N]
 template <typename T>
-static void Transpose(const std::vector<T>& src, std::vector<T>& dst, const int64_t batch, const int64_t N, const int64_t M) {
-  for (int64_t b = 0; b < batch; b++) {
-    for (int64_t n = 0; n < N * M; n++) {
-      int64_t i = n / N;
-      int64_t j = n % N;
-      dst[b * N * M + n] = src[b * N * M + M * j + i];
+static void Transpose(const std::vector<T>& src, std::vector<T>& dst, size_t batch, size_t N, size_t M) {
+  for (size_t b = 0; b < batch; b++) {
+    for (size_t n = 0; n < N * M; n++) {
+      dst[b * N * M + n] = src[b * N * M + M * (n % N) + n / N];
+    }
+  }
+}
+
+// [batch,N,M]->[N,batch,M]
+template <typename T>
+static void TransposeBatch(const std::vector<T>& src, std::vector<T>& dst, size_t batch, size_t N, size_t M) {
+  for (size_t i = 0; i < batch * N; ++i) {
+    size_t src_pos = ((i % batch) * N + i / batch) * M;
+    size_t dst_pos = i * M;
+    for (size_t j = 0; j < M; j++) {
+      dst[dst_pos + j] = src[src_pos + j];
     }
   }
 }
 
 template <typename T>
-void ProcessInputs(const std::vector<int64_t>& input_dims, const std::vector<T>& common_input_vals, bool trans_flag,
-                   std::vector<int64_t>& modified_input_dims, std::vector<T>& input_vals) {
-  auto rank = input_dims.size();
-  ORT_ENFORCE(rank >= 1);
-  int64_t size0 = TensorShape::FromExistingBuffer(input_dims).SizeHelper(0, rank);
-  std::vector<T> input_vals_raw(common_input_vals.cbegin(), common_input_vals.cbegin() + size0);
-  input_vals.resize(size0);
-
-  // transpose on 1-d does not take any effect.
-  if (rank == 1) {
-    trans_flag = false;
+static void TransposeInput(const std::vector<T>& src, std::vector<T>& dst, const std::vector<int64_t>& dims,
+                           std::vector<int64_t>& new_dims, bool is_trans, bool is_trans_batch) {
+  if (dims.size() < 2 || (!is_trans && !is_trans_batch)) return;
+  ORT_ENFORCE(!is_trans_batch || dims.size() >= 3);
+  size_t batch = 1;
+  size_t size = dims.size();
+  new_dims.resize(size);
+  for (size_t i = 0; i < size - 2; ++i) {
+    batch *= static_cast<size_t>(dims[i]);
+    new_dims[i + (is_trans_batch ? 1 : 0)] = dims[i];
   }
-
-  if (trans_flag) {
-    modified_input_dims[rank - 1] = input_dims[rank - 2];
-    modified_input_dims[rank - 2] = input_dims[rank - 1];
-    auto batch_size = TensorShape::FromExistingBuffer(input_dims).SizeHelper(0, rank - 2);
-    Transpose<T>(input_vals_raw, input_vals, batch_size, input_dims[rank - 2], input_dims[rank - 1]);
+  size_t N = static_cast<size_t>(dims[size - 2]);
+  size_t M = static_cast<size_t>(dims[size - 1]);
+  if (is_trans && !is_trans_batch) {
+    new_dims[size - 1] = dims[size - 2];
+    new_dims[size - 2] = dims[size - 1];
+    Transpose<T>(src, dst, batch, N, M);
+  } else if (!is_trans && is_trans_batch) {
+    new_dims[0] = dims[size - 2];
+    TransposeBatch<T>(src, dst, batch, N, M);
   } else {
-    input_vals = input_vals_raw;
+    new_dims[size - 1] = dims[size - 2];
+    new_dims[0] = dims[size - 1];
+    std::vector<T> intermediate(src.size());
+    Transpose<T>(src, intermediate, batch, N, M);       // [batch,N,M]->[batch,M,N]
+    TransposeBatch<T>(intermediate, dst, batch, M, N);  // [batch,M,N]->[M,batch,N]
   }
 }
 
 template <typename T>
-void RunFusedMatMulTest(const char* op_name, int32_t opset_version = 7, bool transa = false, bool transb = false, float alpha = 1.0f, bool is_b_constant = false) {
-  std::vector<T> common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11};
+void ProcessInputs(const std::vector<int64_t>& input_dims, const std::vector<T>& common_input_vals, bool is_trans,
+                   bool is_trans_batch, std::vector<int64_t>& modified_input_dims, std::vector<T>& input_vals) {
+  auto rank = input_dims.size();
+  ORT_ENFORCE(rank >= 1);
+  int64_t size0 = TensorShape::FromExistingBuffer(input_dims).SizeHelper(0, rank);
+  std::vector<T> input_vals_raw(common_input_vals.cbegin(), common_input_vals.cbegin() + size0);
+  input_vals = input_vals_raw;
+  TransposeInput<T>(input_vals_raw, input_vals, input_dims, modified_input_dims, is_trans, is_trans_batch);
+}
+
+template <typename T>
+void RunFusedMatMulTest(const char* op_name, int32_t opset_version = 7, bool transa = false, bool transb = false,
+                        bool is_trans_batch_a = false, bool is_trans_batch_b = false, float alpha = 1.0f,
+                        bool is_b_constant = false) {
+  std::vector<T> common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
   for (auto t : GenerateSimpleTestCases<T>()) {
+    // is_trans_batch requires dim_size >= 3 and same dim size in both sides.
+    if (is_trans_batch_a || is_trans_batch_b) {
+      if (t.input0_dims.size() < 3 || t.input0_dims.size() != t.input1_dims.size()) {
+        continue;
+      }
+    }
+
     OpTester test(op_name, opset_version, onnxruntime::kMSDomain);
 
     std::vector<int64_t> input0_dims(t.input0_dims);
     std::vector<T> input0_vals;
-    ProcessInputs(t.input0_dims, common_input_vals, transa, input0_dims, input0_vals);
+    ProcessInputs(t.input0_dims, common_input_vals, transa, is_trans_batch_a, input0_dims, input0_vals);
 
     std::vector<int64_t> input1_dims(t.input1_dims);
     std::vector<T> input1_vals;
-    ProcessInputs(t.input1_dims, common_input_vals, transb, input1_dims, input1_vals);
+    ProcessInputs(t.input1_dims, common_input_vals, transb, is_trans_batch_b, input1_dims, input1_vals);
 
     test.AddInput<T>("A", input0_dims, input0_vals);
-
     test.AddInput<T>("B", input1_dims, input1_vals, is_b_constant);
 
     test.AddAttribute("transA", (int64_t)transa);
     test.AddAttribute("transB", (int64_t)transb);
+    test.AddAttribute("transBatchA", (int64_t)is_trans_batch_a);
+    test.AddAttribute("transBatchB", (int64_t)is_trans_batch_b);
     test.AddAttribute("alpha", alpha);
 
     if (alpha != 1.0f) {
-      std::transform(
-          t.expected_vals.begin(), t.expected_vals.end(), t.expected_vals.begin(),
-          [alpha](const T& val) -> T { return alpha * val; });
+      std::transform(t.expected_vals.begin(), t.expected_vals.end(), t.expected_vals.begin(),
+                     [alpha](const T& val) -> T { return alpha * val; });
     }
 
     test.AddOutput<T>("Y", t.expected_dims, t.expected_vals);
@@ -183,25 +233,40 @@ TEST(FusedMatMulOpTest, FloatTypeTransposeA) {
 TEST(FusedMatMulOpTest, FloatTypeTransposeB) {
   RunFusedMatMulTest<float>("FusedMatMul", 1, false, true);
   // b is constant. This tests weight packing logic
-  RunFusedMatMulTest<float>("FusedMatMul", 1, false, true, 1.0f, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, false, true, false, false, 1.0f, true);
 }
 
 TEST(FusedMatMulOpTest, FloatTypeTransposeAB) {
   RunFusedMatMulTest<float>("FusedMatMul", 1, true, true);
 
   // b is constant. This tests weight packing logic
-  RunFusedMatMulTest<float>("FusedMatMul", 1, true, true, 1.0f, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, true, false, false, 1.0f, true);
 }
 
 TEST(FusedMatMulOpTest, FloatTypeScale) {
-  RunFusedMatMulTest<float>("FusedMatMul", 1, false, false, 0.5f);
-  RunFusedMatMulTest<float>("FusedMatMul", 1, true, false, 2.0f);
-  RunFusedMatMulTest<float>("FusedMatMul", 1, true, true, 4.0f);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, false, false, false, false, 0.5f);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, false, false, false, 2.0f);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, true, false, false, 4.0f);
 
   // now run tests with b constant.
-  RunFusedMatMulTest<float>("FusedMatMul", 1, false, false, 0.5f, true);
-  RunFusedMatMulTest<float>("FusedMatMul", 1, true, false, 2.0f, true);
-  RunFusedMatMulTest<float>("FusedMatMul", 1, true, true, 4.0f, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, false, false, false, false, 0.5f, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, false, false, false, 2.0f, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, true, false, false, 4.0f, true);
+}
+
+TEST(FusedMatMulOpTest, FloatTypeTransposeBatch) {
+  RunFusedMatMulTest<float>("FusedMatMul", 1, false, false, true, false);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, false, false, false, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, false, false, true, true, 0.5f);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, false, true, false);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, false, false, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, false, true, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, false, true, true, false);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, false, true, false, true, 1.0f, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, false, true, true, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, true, true, false, 2.0f, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, true, false, true);
+  RunFusedMatMulTest<float>("FusedMatMul", 1, true, true, true, true);
 }
 
 }  // namespace transpose_matmul

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -1178,6 +1178,76 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusionWithPreservedTranspose) {
   ASSERT_FALSE(graph.GraphResolveNeeded());
 }
 
+TEST_F(GraphTransformationTests, TransposeMatmulTransBatchFusion) {
+  const std::vector<PathString> model_uris = {
+      MODEL_FOLDER "fusion/transpose_matmul_trans_batch_fusion1.onnx",
+      MODEL_FOLDER "fusion/transpose_matmul_trans_batch_fusion2.onnx",
+      MODEL_FOLDER "fusion/transpose_matmul_trans_batch_fusion3.onnx",
+  };
+  const std::vector<std::pair<int64_t, int64_t>> trans_batch_attrs = {
+      {1, 0},
+      {1, 1},
+      {1, 1},
+  };
+  size_t index = 0;
+  for (const auto& model_uri : model_uris) {
+    std::shared_ptr<Model> p_model;
+    ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+    Graph& graph = p_model->MainGraph();
+
+    onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+    ASSERT_STATUS_OK(graph_transformation_mgr.Register(
+        std::make_unique<MatmulTransposeFusion>(), TransformerLevel::Level1));
+    ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+    std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+    ASSERT_EQ(op_to_count["Transpose"], 0);
+    ASSERT_EQ(op_to_count["MatMul"], 0);
+    ASSERT_EQ(op_to_count["com.microsoft.FusedMatMul"], 1);
+    for (auto& node : graph.Nodes()) {
+      if (node.OpType() == "FusedMatMul") {
+        auto attrs = node.GetAttributes();
+        int64_t trans_batch_a = 0;
+        if (attrs.find("transBatchA") != attrs.end()) {
+          trans_batch_a = attrs.at("transBatchA").i();
+        }
+        int64_t trans_batch_b = 0;
+        if (attrs.find("transBatchB") != attrs.end()) {
+          trans_batch_b = attrs.at("transBatchB").i();
+        }
+        ASSERT_EQ(trans_batch_a, trans_batch_attrs[index].first);
+        ASSERT_EQ(trans_batch_b, trans_batch_attrs[index].second);
+        break;
+      }
+    }
+    ++index;
+  }
+}
+
+TEST_F(GraphTransformationTests, TransposeMatmulTransBatchNoFusion) {
+  const std::vector<PathString> model_uris = {
+      MODEL_FOLDER "fusion/transpose_matmul_trans_batch_fusion_invalid_case1.onnx",
+      MODEL_FOLDER "fusion/transpose_matmul_trans_batch_fusion_invalid_case2.onnx",
+      MODEL_FOLDER "fusion/transpose_matmul_trans_batch_fusion_invalid_case3.onnx",
+  };
+  for (const auto& model_uri : model_uris) {
+    std::shared_ptr<Model> p_model;
+    ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+    Graph& graph = p_model->MainGraph();
+    std::map<std::string, int> orig_op_to_count = CountOpsInGraph(graph);
+
+    onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+    ASSERT_STATUS_OK(graph_transformation_mgr.Register(
+        std::make_unique<MatmulTransposeFusion>(), TransformerLevel::Level1));
+    ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+    std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+    ASSERT_EQ(op_to_count["Transpose"], orig_op_to_count["Transpose"]);
+    ASSERT_EQ(op_to_count["MatMul"], orig_op_to_count["MatMul"]);
+    ASSERT_EQ(op_to_count["com.microsoft.FusedMatMul"], orig_op_to_count["com.microsoft.FusedMatMul"]);
+  }
+}
+
 TEST_F(GraphTransformationTests, Gemm_LeakyRelu_Fusion) {
   auto model_uri = MODEL_FOLDER "gemm_activation_fusion/gemm_activation_fusion.onnx";
 

onnxruntime/test/testdata/transform/fusion/transpose_matmul_gen.py

@@ -252,3 +252,190 @@ def gen_transpose_fusion_invalid_datatype(model_path, datatype):
 
 gen_transpose_fusion_invalid_datatype("transpose_matmul_4d_fusion_invalid_datatype_int32.onnx", TensorProto.INT32)
 gen_transpose_fusion_invalid_datatype("transpose_matmul_4d_fusion_invalid_datatype_int64.onnx", TensorProto.INT64)
+
+
+def gen_transpose_matmul_trans_batch_fusion(model_path):
+    nodes = [
+        helper.make_node(
+            "Transpose",
+            ["input_0"],
+            ["transposed_input_0"],
+            perm = [1, 2, 0]),
+        helper.make_node(
+            "Transpose",
+            ["input_1"],
+            ["transposed_input_1"],
+            perm = [0, 2, 1]),
+        helper.make_node(
+            "MatMul",
+            ["transposed_input_0", "transposed_input_1"],
+            ["output"])
+    ]
+
+    inputs = [
+        helper.make_tensor_value_info(
+            "input_0", TensorProto.FLOAT, ['K', 3, 'M']),
+        helper.make_tensor_value_info(
+            "input_1", TensorProto.FLOAT, [3, 'N', 'K']),
+    ]
+
+    outputs = [
+        helper.make_tensor_value_info(
+            "output", TensorProto.FLOAT, [3, 'M', 'N'])
+    ]
+
+    save(model_path + "1.onnx", nodes, inputs, outputs, [])
+
+    nodes = [
+        helper.make_node(
+            "Transpose",
+            ["input_0"],
+            ["transposed_input_0"],
+            perm = [1, 2, 0, 3]),
+        helper.make_node(
+            "Transpose",
+            ["input_0"],
+            ["transposed_input_1"],
+            perm = [1, 2, 3, 0]),
+        helper.make_node(
+            "MatMul",
+            ["transposed_input_0", "transposed_input_1"],
+            ["output"])
+    ]
+
+    inputs = [
+        helper.make_tensor_value_info(
+            "input_0", TensorProto.FLOAT, ['M', 2, 3, 'K']),
+    ]
+
+    outputs = [
+        helper.make_tensor_value_info(
+            "output", TensorProto.FLOAT, [2, 3, 'M', 'M'])
+    ]
+
+    save(model_path + "2.onnx", nodes, inputs, outputs, [])
+
+    nodes = [
+        helper.make_node(
+            "Transpose",
+            ["input_0"],
+            ["transposed_input_0"],
+            perm = [1, 2, 3, 0]),
+        helper.make_node(
+            "FusedMatMul",
+            ["transposed_input_0", "input_1"],
+            ["output"],
+            "FusedMatMul",
+            "",
+            msdomain.domain,
+            alpha=3.0, transA=1, transBatchB=1)
+    ]
+
+    inputs = [
+        helper.make_tensor_value_info(
+            "input_0", TensorProto.FLOAT, ['M', 2, 3, 'K']),
+        helper.make_tensor_value_info(
+            "input_1", TensorProto.FLOAT, ['K', 2, 3, 'N']),
+    ]
+
+    outputs = [
+        helper.make_tensor_value_info(
+            "output", TensorProto.FLOAT, [2, 3, 'M', 'M'])
+    ]
+
+    save(model_path + "3.onnx", nodes, inputs, outputs, [])
+
+
+gen_transpose_matmul_trans_batch_fusion(
+    "transpose_matmul_trans_batch_fusion")
+
+
+def gen_transpose_matmul_trans_batch_fusion_invalid_cases(model_path):
+    nodes = [
+        helper.make_node(
+            "Transpose",
+            ["input_0"],
+            ["transposed_input_0"],
+            perm = [1, 2, 0]),
+        helper.make_node(
+            "MatMul",
+            ["transposed_input_0", "input_1"],
+            ["output"])
+    ]
+
+    inputs = [
+        helper.make_tensor_value_info(
+            "input_0", TensorProto.FLOAT, ['K', 3, 'M']),
+        helper.make_tensor_value_info(
+            "input_1", TensorProto.FLOAT, [2, 3, 'K', 'N']),
+    ]
+
+    outputs = [
+        helper.make_tensor_value_info(
+            "output", TensorProto.FLOAT, [2, 3, 'M', 'N'])
+    ]
+
+    save(model_path + "1.onnx", nodes, inputs, outputs, [])
+
+    nodes = [
+        helper.make_node(
+            "Transpose",
+            ["input_0"],
+            ["transposed_input_0"],
+            perm = [0, 2, 1, 3]),
+        helper.make_node(
+            "Transpose",
+            ["input_0"],
+            ["transposed_input_1"],
+            perm = [0, 2, 3, 1]),
+        helper.make_node(
+            "MatMul",
+            ["transposed_input_0", "transposed_input_1"],
+            ["output"])
+    ]
+
+    inputs = [
+        helper.make_tensor_value_info(
+            "input_0", TensorProto.FLOAT, [2, 'M', 3, 'K']),
+    ]
+
+    outputs = [
+        helper.make_tensor_value_info(
+            "output", TensorProto.FLOAT, [2, 3, 'M', 'M'])
+    ]
+
+    save(model_path + "2.onnx", nodes, inputs, outputs, [])
+
+    nodes = [
+        helper.make_node(
+            "Transpose",
+            ["input_0"],
+            ["transposed_input_0"],
+            perm = [1, 2, 3, 0]),
+        helper.make_node(
+            "FusedMatMul",
+            ["transposed_input_0", "input_1"],
+            ["output"],
+            "FusedMatMul",
+            "",
+            msdomain.domain,
+            alpha=3.0, transBatchA=1)
+    ]
+
+    inputs = [
+        helper.make_tensor_value_info(
+            "input_0", TensorProto.FLOAT, ['K', 'M', 2, 3]),
+        helper.make_tensor_value_info(
+            "input_1", TensorProto.FLOAT, [2, 3, 'K', 'N']),
+    ]
+
+    outputs = [
+        helper.make_tensor_value_info(
+            "output", TensorProto.FLOAT, [2, 3, 'M', 'M'])
+    ]
+
+    save(model_path + "3.onnx", nodes, inputs, outputs, [])
+
+
+gen_transpose_matmul_trans_batch_fusion_invalid_cases(
+    "transpose_matmul_trans_batch_fusion_invalid_case")