diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index a982999f0a..d8630076c1 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1272,6 +1272,10 @@ This version of the operator has been available since version 1 of the 'com.micr
Whether A should be transposed on the last two dimensions before doing multiplication
transB : int
Whether B should be transposed on the last two dimensions before doing multiplication
+
transBatchA : int
+
Whether A should be transposed on the 1st dimension and batch dimensions (dim-1 to dim-rank-2) before doing multiplication
+
transBatchB : int
+
Whether B should be transposed on the 1st dimension and batch dimensions (dim-1 to dim-rank-2) before doing multiplication
#### Inputs diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index cf4e7efd8d..502244fa4b 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -282,6 +282,10 @@ void FusedMatMulShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { bool transa = transAAttr ? static_cast(transAAttr->i()) != 0 : false; auto transBAttr = ctx.getAttribute("transB"); bool transb = transBAttr ? static_cast(transBAttr->i()) != 0 : false; + auto trans_batch_a_attr = ctx.getAttribute("transBatchA"); + bool trans_batch_a = trans_batch_a_attr ? static_cast(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(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(0)) - .Attr( - "transB", - "Whether B should be transposed on the last two dimensions before doing multiplication", - AttributeProto::INT, - static_cast(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(0)) + .Attr("transB", "Whether B should be transposed on the last two dimensions before doing multiplication", + AttributeProto::INT, static_cast(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(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(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) diff --git a/onnxruntime/core/optimizer/matmul_transpose_fusion.cc b/onnxruntime/core/optimizer/matmul_transpose_fusion.cc index 1f6e546e84..25ce5ba150 100644 --- a/onnxruntime/core/optimizer/matmul_transpose_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_transpose_fusion.cc @@ -32,14 +32,25 @@ static bool GetTransposePerms(const Node& transpose_node, std::vector& 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(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(rank - 2)] == rank - 1 && perms[static_cast(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(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& consumer_count, - std::deque& removed_nodes) { + std::deque& 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(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(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(node.GetAttributes().at("transA").i()); - transpose_right ^= static_cast(node.GetAttributes().at("transB").i()); + is_trans_left ^= static_cast(node.GetAttributes().at("transA").i()); + is_trans_right ^= static_cast(node.GetAttributes().at("transB").i()); + is_trans_batch_left ^= static_cast(node.GetAttributes().at("transBatchA").i()); + is_trans_batch_right ^= static_cast(node.GetAttributes().at("transBatchB").i()); alpha = node.GetAttributes().at("alpha").f(); } - matmul_node.AddAttribute("transA", static_cast(transpose_left)); - matmul_node.AddAttribute("transB", static_cast(transpose_right)); + matmul_node.AddAttribute("transA", static_cast(is_trans_left)); + matmul_node.AddAttribute("transB", static_cast(is_trans_right)); + matmul_node.AddAttribute("transBatchA", static_cast(is_trans_batch_left)); + matmul_node.AddAttribute("transBatchB", static_cast(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()); diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc index ad9e34cc67..ec395cf018 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.cc +++ b/onnxruntime/core/providers/cpu/math/matmul.cc @@ -169,7 +169,7 @@ Status MatMul::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::Compute(OpKernelContext* ctx) const { const size_t M = static_cast(helper.M()); const size_t N = static_cast(helper.N()); const size_t K = static_cast(helper.K()); - const size_t lda = static_cast(trans_a ? M : K); - const size_t ldb = static_cast(trans_b ? K : N); + const size_t lda = helper.Lda(trans_a); + const size_t ldb = helper.Ldb(trans_b); std::vector data(max_len); for (size_t i = 0; i < max_len; i++) { diff --git a/onnxruntime/core/providers/cpu/math/matmul.h b/onnxruntime/core/providers/cpu/math/matmul.h index 4f312855f9..96e461673b 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.h +++ b/onnxruntime/core/providers/cpu/math/matmul.h @@ -22,14 +22,18 @@ class MatMul final : public OpKernel { info.GetAttrOrDefault("transA", &trans_a_attr_, 0); info.GetAttrOrDefault("transB", &trans_b_attr_, 0); info.GetAttrOrDefault("alpha", &alpha_attr_, 1.0); + int64_t trans_batch_a_attr, trans_batch_b_attr; + info.GetAttrOrDefault("transBatchA", &trans_batch_a_attr, 0); + info.GetAttrOrDefault("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& prepacked_buffers, - int input_idx, + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, int input_idx, /*out*/ bool& used_shared_buffers) override; Status Compute(OpKernelContext* context) const override; @@ -42,6 +46,8 @@ class MatMul 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 diff --git a/onnxruntime/core/providers/cpu/math/matmul_helper.h b/onnxruntime/core/providers/cpu/math/matmul_helper.h index 2ba1e2a234..6988225ccd 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_helper.h +++ b/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(left_shape.SizeToDimension(left_num_dims - 1)); - K_ = static_cast(left_shape[left_num_dims - 1]); - N_ = static_cast(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(orig_left_shape.SizeToDimension(left_num_dims - 1)); + K_ = static_cast(orig_left_shape[left_num_dims - 1]); + N_ = static_cast(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 dims_left(left_num_dims); + std::vector 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(dims_left[i]); + } + dims_left[left_num_dims - 2] = leading_dim; + left_stride_factor_ = static_cast(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(dims_right[i]); + } + dims_right[right_num_dims - 2] = leading_dim; + right_stride_factor_ = static_cast(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 right_zp_offsets_; std::vector 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(M_) : static_cast(K_)) * left_ld_factor_; + } + + int Ldb(bool is_trans) const { + return (is_trans ? static_cast(K_) : static_cast(N_)) * right_ld_factor_; + } + + int Ldc() const { + return static_cast(N_); + } + // Batched Gemm offsets in left matrices const std::vector& LeftOffsets() const { return left_offsets_; diff --git a/onnxruntime/core/providers/cuda/math/matmul.cc b/onnxruntime/core/providers/cuda/math/matmul.cc index 1a233aec96..c7632dcdc5 100644 --- a/onnxruntime/core/providers/cuda/math/matmul.cc +++ b/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::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::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(helper.M()) : static_cast(helper.K()); - const int ldb = transb ? static_cast(helper.K()) : static_cast(helper.N()); - const int ldc = static_cast(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::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::ComputeInternal(OpKernelContext* ctx) const { return Status::OK(); } + // Fill offsets when needed. + helper.FillOffsets(); CudaAsyncBuffer left_arrays(this, helper.LeftOffsets().size()); CudaAsyncBuffer right_arrays(this, helper.RightOffsets().size()); CudaAsyncBuffer output_arrays(this, helper.OutputOffsets().size()); diff --git a/onnxruntime/core/providers/cuda/math/matmul.h b/onnxruntime/core/providers/cuda/math/matmul.h index ac3d863e44..f857fa1ee2 100644 --- a/onnxruntime/core/providers/cuda/math/matmul.h +++ b/onnxruntime/core/providers/cuda/math/matmul.h @@ -16,8 +16,9 @@ class MatMul final : public CudaKernel { : CudaKernel(info), alpha_{info.GetAttrOrDefault("alpha", 1.0f)}, trans_A_{info.GetAttrOrDefault("transA", 0) != 0}, - trans_B_{info.GetAttrOrDefault("transB", 0) != 0} { - } + trans_B_{info.GetAttrOrDefault("transB", 0) != 0}, + trans_batch_a_{info.GetAttrOrDefault("transBatchA", 0) != 0}, + trans_batch_b_{info.GetAttrOrDefault("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 diff --git a/onnxruntime/core/providers/rocm/math/matmul.cc b/onnxruntime/core/providers/rocm/math/matmul.cc index 0329917f3d..4e6e8a125d 100644 --- a/onnxruntime/core/providers/rocm/math/matmul.cc +++ b/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::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::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(helper.M()) : static_cast(helper.K()); - const int ldb = transb ? static_cast(helper.K()) : static_cast(helper.N()); - const int ldc = static_cast(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::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::ComputeInternal(OpKernelContext* ctx) const { return Status::OK(); } + // Fill offsets when needed. + helper.FillOffsets(); RocmAsyncBuffer left_arrays(this, helper.LeftOffsets().size()); RocmAsyncBuffer right_arrays(this, helper.RightOffsets().size()); RocmAsyncBuffer output_arrays(this, helper.OutputOffsets().size()); diff --git a/onnxruntime/test/contrib_ops/fused_matmul_op_test.cc b/onnxruntime/test/contrib_ops/fused_matmul_op_test.cc index 7bf3fc676d..c1c7a75a18 100644 --- a/onnxruntime/test/contrib_ops/fused_matmul_op_test.cc +++ b/onnxruntime/test/contrib_ops/fused_matmul_op_test.cc @@ -91,71 +91,121 @@ std::vector> 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 -static void Transpose(const std::vector& src, std::vector& 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& src, std::vector& 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 +static void TransposeBatch(const std::vector& src, std::vector& 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 -void ProcessInputs(const std::vector& input_dims, const std::vector& common_input_vals, bool trans_flag, - std::vector& modified_input_dims, std::vector& input_vals) { - auto rank = input_dims.size(); - ORT_ENFORCE(rank >= 1); - int64_t size0 = TensorShape::FromExistingBuffer(input_dims).SizeHelper(0, rank); - std::vector 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& src, std::vector& dst, const std::vector& dims, + std::vector& 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(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(input_vals_raw, input_vals, batch_size, input_dims[rank - 2], input_dims[rank - 1]); + size_t N = static_cast(dims[size - 2]); + size_t M = static_cast(dims[size - 1]); + if (is_trans && !is_trans_batch) { + new_dims[size - 1] = dims[size - 2]; + new_dims[size - 2] = dims[size - 1]; + Transpose(src, dst, batch, N, M); + } else if (!is_trans && is_trans_batch) { + new_dims[0] = dims[size - 2]; + TransposeBatch(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 intermediate(src.size()); + Transpose(src, intermediate, batch, N, M); // [batch,N,M]->[batch,M,N] + TransposeBatch(intermediate, dst, batch, M, N); // [batch,M,N]->[M,batch,N] } } template -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 common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}; +void ProcessInputs(const std::vector& input_dims, const std::vector& common_input_vals, bool is_trans, + bool is_trans_batch, std::vector& modified_input_dims, std::vector& input_vals) { + auto rank = input_dims.size(); + ORT_ENFORCE(rank >= 1); + int64_t size0 = TensorShape::FromExistingBuffer(input_dims).SizeHelper(0, rank); + std::vector input_vals_raw(common_input_vals.cbegin(), common_input_vals.cbegin() + size0); + input_vals = input_vals_raw; + TransposeInput(input_vals_raw, input_vals, input_dims, modified_input_dims, is_trans, is_trans_batch); +} + +template +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 common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; for (auto t : GenerateSimpleTestCases()) { + // 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 input0_dims(t.input0_dims); std::vector 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 input1_dims(t.input1_dims); std::vector 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("A", input0_dims, input0_vals); - test.AddInput("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("Y", t.expected_dims, t.expected_vals); @@ -183,25 +233,40 @@ TEST(FusedMatMulOpTest, FloatTypeTransposeA) { TEST(FusedMatMulOpTest, FloatTypeTransposeB) { RunFusedMatMulTest("FusedMatMul", 1, false, true); // b is constant. This tests weight packing logic - RunFusedMatMulTest("FusedMatMul", 1, false, true, 1.0f, true); + RunFusedMatMulTest("FusedMatMul", 1, false, true, false, false, 1.0f, true); } TEST(FusedMatMulOpTest, FloatTypeTransposeAB) { RunFusedMatMulTest("FusedMatMul", 1, true, true); // b is constant. This tests weight packing logic - RunFusedMatMulTest("FusedMatMul", 1, true, true, 1.0f, true); + RunFusedMatMulTest("FusedMatMul", 1, true, true, false, false, 1.0f, true); } TEST(FusedMatMulOpTest, FloatTypeScale) { - RunFusedMatMulTest("FusedMatMul", 1, false, false, 0.5f); - RunFusedMatMulTest("FusedMatMul", 1, true, false, 2.0f); - RunFusedMatMulTest("FusedMatMul", 1, true, true, 4.0f); + RunFusedMatMulTest("FusedMatMul", 1, false, false, false, false, 0.5f); + RunFusedMatMulTest("FusedMatMul", 1, true, false, false, false, 2.0f); + RunFusedMatMulTest("FusedMatMul", 1, true, true, false, false, 4.0f); // now run tests with b constant. - RunFusedMatMulTest("FusedMatMul", 1, false, false, 0.5f, true); - RunFusedMatMulTest("FusedMatMul", 1, true, false, 2.0f, true); - RunFusedMatMulTest("FusedMatMul", 1, true, true, 4.0f, true); + RunFusedMatMulTest("FusedMatMul", 1, false, false, false, false, 0.5f, true); + RunFusedMatMulTest("FusedMatMul", 1, true, false, false, false, 2.0f, true); + RunFusedMatMulTest("FusedMatMul", 1, true, true, false, false, 4.0f, true); +} + +TEST(FusedMatMulOpTest, FloatTypeTransposeBatch) { + RunFusedMatMulTest("FusedMatMul", 1, false, false, true, false); + RunFusedMatMulTest("FusedMatMul", 1, false, false, false, true); + RunFusedMatMulTest("FusedMatMul", 1, false, false, true, true, 0.5f); + RunFusedMatMulTest("FusedMatMul", 1, true, false, true, false); + RunFusedMatMulTest("FusedMatMul", 1, true, false, false, true); + RunFusedMatMulTest("FusedMatMul", 1, true, false, true, true); + RunFusedMatMulTest("FusedMatMul", 1, false, true, true, false); + RunFusedMatMulTest("FusedMatMul", 1, false, true, false, true, 1.0f, true); + RunFusedMatMulTest("FusedMatMul", 1, false, true, true, true); + RunFusedMatMulTest("FusedMatMul", 1, true, true, true, false, 2.0f, true); + RunFusedMatMulTest("FusedMatMul", 1, true, true, false, true); + RunFusedMatMulTest("FusedMatMul", 1, true, true, true, true); } } // namespace transpose_matmul diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index abc080a31a..34920c043c 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/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 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> trans_batch_attrs = { + {1, 0}, + {1, 1}, + {1, 1}, + }; + size_t index = 0; + for (const auto& model_uri : model_uris) { + std::shared_ptr 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(), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + std::map 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 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 p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + std::map orig_op_to_count = CountOpsInGraph(graph); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + ASSERT_STATUS_OK(graph_transformation_mgr.Register( + std::make_unique(), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + std::map 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"; diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_gen.py b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_gen.py index 1f2bc05551..2904a618bb 100644 --- a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_gen.py +++ b/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") diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion1.onnx b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion1.onnx new file mode 100644 index 0000000000..ff11d6f263 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion1.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion2.onnx b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion2.onnx new file mode 100644 index 0000000000..1266ca28fd Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion2.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion3.onnx b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion3.onnx new file mode 100644 index 0000000000..741c73cced Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion3.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion_invalid_case1.onnx b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion_invalid_case1.onnx new file mode 100644 index 0000000000..65f31f641a Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion_invalid_case1.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion_invalid_case2.onnx b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion_invalid_case2.onnx new file mode 100644 index 0000000000..a2ea9b758a Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion_invalid_case2.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion_invalid_case3.onnx b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion_invalid_case3.onnx new file mode 100644 index 0000000000..76e09a048d Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_trans_batch_fusion_invalid_case3.onnx differ