saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add FusedMatMul contrib op (#5213)

Browse source 
* bug fix transformer

* fuse cpu kernel for transposescalematmul and matmul

* fuse transpose_scale_matmul cpu kernel with matmul

* fix test

* Add FusedMatMul Contrib Op

* fix test

* fix typo

* plus more updates per review
This commit is contained in:
Ashwini Khade 2020-09-23 12:17:50 -07:00 committed by GitHub
parent fe7e7bfe60
commit 16220f3848
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
29 changed files with 343 additions and 290 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
View file

@ -21,6 +21,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Range
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, WordConvEmbedding);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GatherND);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, TransposeMatMul);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FusedMatMul);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MurmurHash3);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, MaxpoolWithMask);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Pad);
@ -158,6 +159,7 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GatherND)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MurmurHash3)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, TransposeMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FusedMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, MaxpoolWithMask)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Pad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Unique)>,

26
onnxruntime/contrib_ops/cpu/fused_matmul.cc Normal file
View file

@ -0,0 +1,26 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/cpu/math/matmul.h"
namespace onnxruntime {
namespace contrib {
ONNX_OPERATOR_KERNEL_EX(
TransposeMatMul,
kMSDomain,
1,
kCpuExecutionProvider,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
MatMul<float>);
ONNX_OPERATOR_KERNEL_EX(
FusedMatMul,
kMSDomain,
1,
kCpuExecutionProvider,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
MatMul<float>);
} // namespace contrib
} // namespace onnxruntime

63
onnxruntime/contrib_ops/cpu/transpose_scale_matmul.cc
View file

@ -1,63 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "contrib_ops/cpu/transpose_scale_matmul.h"
#include "core/providers/cpu/math/matmul_helper.h"
#include "core/util/math.h"
namespace onnxruntime {
namespace contrib {
ONNX_OPERATOR_KERNEL_EX(
TransposeMatMul,
kMSDomain,
1,
kCpuExecutionProvider,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
TransposeMatMul);
TransposeMatMul::TransposeMatMul(const OpKernelInfo& info)
: OpKernel{info} {
ORT_THROW_IF_ERROR(info.GetAttr("alpha", &alpha_attr_));
ORT_THROW_IF_ERROR(info.GetAttr("transA", &trans_a_attr_));
ORT_THROW_IF_ERROR(info.GetAttr("transB", &trans_b_attr_));
}
Status TransposeMatMul::Compute(OpKernelContext* context) const {
concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool();
const Tensor* A = context->Input<Tensor>(0);
const Tensor* B = context->Input<Tensor>(1);
// match CUDA kernel implementation, ignore transpose for vectors
const bool trans_a = trans_a_attr_ && A->Shape().NumDimensions() != 1;
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));
Tensor* Y = context->Output(0, helper.OutputShape());
// Bail out early if the output is going to be empty
if (Y->Shape().Size() == 0)
return Status::OK();
const size_t num_offsets = helper.OutputOffsets().size();
for (size_t i = 0; i < num_offsets; ++i) {
math::Gemm<float, concurrency::ThreadPool>(
trans_a ? CblasTrans : CblasNoTrans,
trans_b ? CblasTrans : CblasNoTrans,
helper.M(), helper.N(), helper.K(),
alpha_attr_,
A->Data<float>() + helper.LeftOffsets()[i],
B->Data<float>() + helper.RightOffsets()[i],
0.0f,
Y->MutableData<float>() + helper.OutputOffsets()[i],
thread_pool);
}
return Status::OK();
}
} // namespace contrib
} // namespace onnxruntime

23
onnxruntime/contrib_ops/cpu/transpose_scale_matmul.h
View file

@ -1,23 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
class TransposeMatMul final : public OpKernel {
public:
TransposeMatMul(const OpKernelInfo& info);
Status Compute(OpKernelContext* context) const override;
private:
float alpha_attr_;
int64_t trans_a_attr_, trans_b_attr_;
};
} // namespace contrib
} // namespace onnxruntime

6
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
View file

@ -20,6 +20,9 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, TransposeMatMul);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, TransposeMatMul);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, TransposeMatMul);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FusedMatMul);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, FusedMatMul);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, FusedMatMul);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Rfft);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Rfft);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Rfft);
@ -92,6 +95,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, TransposeMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, TransposeMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, TransposeMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FusedMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, FusedMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, FusedMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Rfft)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Rfft)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Rfft)>,

14
onnxruntime/contrib_ops/cuda/math/transpose_scale_matmul.cc → onnxruntime/contrib_ops/cuda/math/fused_matmul.cc
View file

@ -7,9 +7,9 @@ namespace onnxruntime {
namespace contrib {
namespace cuda {
#define REGISTER_KERNEL_TYPED(T) \
#define REGISTER_KERNEL_TYPED(op_name, T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
TransposeMatMul, \
op_name, \
kMSDomain, \
1, \
T, \
@ -18,9 +18,13 @@ namespace cuda {
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
onnxruntime::cuda::MatMul<T>);
REGISTER_KERNEL_TYPED(float)
REGISTER_KERNEL_TYPED(double)
REGISTER_KERNEL_TYPED(MLFloat16)
REGISTER_KERNEL_TYPED(TransposeMatMul, float)
REGISTER_KERNEL_TYPED(TransposeMatMul, double)
REGISTER_KERNEL_TYPED(TransposeMatMul, MLFloat16)
REGISTER_KERNEL_TYPED(FusedMatMul, float)
REGISTER_KERNEL_TYPED(FusedMatMul, double)
REGISTER_KERNEL_TYPED(FusedMatMul, MLFloat16)
} // namespace cuda
} // namespace contrib

254
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

@ -198,6 +198,115 @@ void ValidateTypeAndShapeForScaleAndZP(ONNX_NAMESPACE::InferenceContext& ctx, in
}
}
void FusedMatMulShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
auto transAAttr = ctx.getAttribute("transA");
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;
int input1Idx = 0;
int input2Idx = 1;
if (!hasInputShape(ctx, input1Idx) || !hasInputShape(ctx, input2Idx)) {
return;
}
const auto shape0_raw = getInputShape(ctx, input1Idx);
const auto shape1_raw = getInputShape(ctx, input2Idx);
if (shape0_raw.dim_size() == 0 || shape1_raw.dim_size() == 0) {
fail_shape_inference("Input tensors of wrong rank (0).");
}
// numpy transpose on a vector does not change anything.
if (shape0_raw.dim_size() == 1) {
transa = false;
}
if (shape1_raw.dim_size() == 1) {
transb = false;
}
ONNX_NAMESPACE::TensorShapeProto shape0, shape1;
auto rank0 = shape0_raw.dim_size();
if (rank0 == 1) {
// for vector input, transa does not make impact on the dim.
shape0 = shape0_raw;
} else {
for (int i = 0; i < rank0 - 2; ++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);
}
auto rank1 = shape1_raw.dim_size();
if (rank1 == 1) {
// for vector input, transb does not make impact on the dim.
shape1 = shape1_raw;
} else {
for (int i = 0; i < rank1 - 2; ++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);
}
ONNX_NAMESPACE::TensorShapeProto shapeL, shapeR;
// First promote each shape to at least rank-2. This logic is
// specific to matmul, not generic broadcasting.
{
if (shape0.dim_size() == 1) {
shapeL.add_dim()->set_dim_value(1);
*shapeL.add_dim() = shape0.dim(0);
} else {
*shapeL.mutable_dim() = shape0.dim();
}
if (shape1.dim_size() == 1) {
*shapeR.add_dim() = shape1.dim(0);
shapeR.add_dim()->set_dim_value(1);
} else {
*shapeR.mutable_dim() = shape1.dim();
}
}
// Check for compatible matrix multiply dimensions
{
auto dimL = shapeL.dim(shapeL.dim_size() - 1);
auto dimR = shapeR.dim(shapeR.dim_size() - 2);
if (dimL.has_dim_value() && dimR.has_dim_value() &&
dimL.dim_value() != dimR.dim_value()) {
fail_shape_inference("Incompatible dimensions for matrix multiplication");
}
}
ONNX_NAMESPACE::TensorShapeProto resultShape;
// Now call out to generic multidimensional broadcasting for
// the broadcastable prefixes.
{
ONNX_NAMESPACE::TensorShapeProto prefixShapeL, prefixShapeR;
for (int i = 0; i < shapeL.dim_size() - 2; ++i) {
*prefixShapeL.add_dim() = shapeL.dim(i);
}
for (int i = 0; i < shapeR.dim_size() - 2; ++i) {
*prefixShapeR.add_dim() = shapeR.dim(i);
}
bidirectionalBroadcastShapeInference(
prefixShapeL, prefixShapeR, resultShape);
}
// Back to matmul-specific. Add the trailing dimensions back in.
{
if (shape0.dim_size() != 1) {
*resultShape.add_dim() = shapeL.dim(shapeL.dim_size() - 2);
}
if (shape1.dim_size() != 1) {
*resultShape.add_dim() = shapeR.dim(shapeR.dim_size() - 1);
}
}
updateOutputShape(ctx, 0, resultShape);
}
std::function<void(OpSchema&)> QLinearMathDocGenerator(const char* name, const char* additionalDocumentation) {
return [=](OpSchema& schema) {
std::string doc = R"DOC(
@ -1775,13 +1884,17 @@ Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-
});
static const char* TransposeMatMul_doc = R"DOC(
Deprecated. Going forward FusedMatMul should be used. This OP will be supported for backward compatibility.
Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.matmul.html
)DOC";
static const char* FusedMatMul_doc = R"DOC(
Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.matmul.html
)DOC";
ONNX_CONTRIB_OPERATOR_SCHEMA(TransposeMatMul)
.SetDomain(kMSDomain)
.SinceVersion(1)
.SetDoc("TransposeMatMul")
.Input(0, "A", "N-dimensional matrix A", "T")
.Input(1, "B", "N-dimensional matrix B", "T")
.Attr(
@ -1806,112 +1919,37 @@ Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-
"Constrain input and output types to float tensors.")
.SetDoc(TransposeMatMul_doc)
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
auto transAAttr = ctx.getAttribute("transA");
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;
int input1Idx = 0;
int input2Idx = 1;
if (!hasInputShape(ctx, input1Idx) || !hasInputShape(ctx, input2Idx)) {
return;
}
FusedMatMulShapeInference(ctx);
});
const auto shape0_raw = getInputShape(ctx, input1Idx);
const auto shape1_raw = getInputShape(ctx, input2Idx);
if (shape0_raw.dim_size() == 0 || shape1_raw.dim_size() == 0) {
fail_shape_inference("Input tensors of wrong rank (0).");
}
// numpy transpose on a vector does not change anything.
if (shape0_raw.dim_size() == 1) {
transa = false;
}
if (shape1_raw.dim_size() == 1) {
transb = false;
}
ONNX_NAMESPACE::TensorShapeProto shape0, shape1;
auto rank0 = shape0_raw.dim_size();
if (rank0 == 1) {
// for vector input, transa does not make impact on the dim.
shape0 = shape0_raw;
} else {
for (int i = 0; i < rank0 - 2; ++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);
}
auto rank1 = shape1_raw.dim_size();
if (rank1 == 1) {
// for vector input, transb does not make impact on the dim.
shape1 = shape1_raw;
} else {
for (int i = 0; i < rank1 - 2; ++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);
}
ONNX_NAMESPACE::TensorShapeProto shapeL, shapeR;
// First promote each shape to at least rank-2. This logic is
// specific to matmul, not generic broadcasting.
{
if (shape0.dim_size() == 1) {
shapeL.add_dim()->set_dim_value(1);
*shapeL.add_dim() = shape0.dim(0);
} else {
*shapeL.mutable_dim() = shape0.dim();
}
if (shape1.dim_size() == 1) {
*shapeR.add_dim() = shape1.dim(0);
shapeR.add_dim()->set_dim_value(1);
} else {
*shapeR.mutable_dim() = shape1.dim();
}
}
// Check for compatible matrix multiply dimensions
{
auto dimL = shapeL.dim(shapeL.dim_size() - 1);
auto dimR = shapeR.dim(shapeR.dim_size() - 2);
if (dimL.has_dim_value() && dimR.has_dim_value() &&
dimL.dim_value() != dimR.dim_value()) {
fail_shape_inference("Incompatible dimensions for matrix multiplication");
}
}
ONNX_NAMESPACE::TensorShapeProto resultShape;
// Now call out to generic multidimensional broadcasting for
// the broadcastable prefixes.
{
ONNX_NAMESPACE::TensorShapeProto prefixShapeL, prefixShapeR;
for (int i = 0; i < shapeL.dim_size() - 2; ++i) {
*prefixShapeL.add_dim() = shapeL.dim(i);
}
for (int i = 0; i < shapeR.dim_size() - 2; ++i) {
*prefixShapeR.add_dim() = shapeR.dim(i);
}
bidirectionalBroadcastShapeInference(
prefixShapeL, prefixShapeR, resultShape);
}
// Back to matmul-specific. Add the trailing dimensions back in.
{
if (shape0.dim_size() != 1) {
*resultShape.add_dim() = shapeL.dim(shapeL.dim_size() - 2);
}
if (shape1.dim_size() != 1) {
*resultShape.add_dim() = shapeR.dim(shapeR.dim_size() - 1);
}
}
updateOutputShape(ctx, 0, resultShape);
ONNX_CONTRIB_OPERATOR_SCHEMA(FusedMatMul)
.SetDomain(kMSDomain)
.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))
.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.")
.SetDoc(FusedMatMul_doc)
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
FusedMatMulShapeInference(ctx);
});
ONNX_CONTRIB_OPERATOR_SCHEMA(ReduceSumInteger)
@ -2755,11 +2793,11 @@ Example 4:
propagateShapeAndTypeFromFirstInput(ctx);
propagateElemTypeFromInputToOutput(ctx, 0, 0);
auto type = ctx.getAttribute("stash_type")->i();
if (ctx.getNumOutputs() > 1){
if (ctx.getNumOutputs() > 1) {
auto output_type = ctx.getOutputType(1);
output_type->mutable_tensor_type()->set_elem_type(static_cast<int32_t>(type));
}
if (ctx.getNumOutputs() > 2){
if (ctx.getNumOutputs() > 2) {
auto output_type = ctx.getOutputType(2);
output_type->mutable_tensor_type()->set_elem_type(static_cast<int32_t>(type));
}

5
onnxruntime/core/optimizer/matmul_scale_fusion.cc
View file

@ -187,6 +187,7 @@ Status ProcessNode(
}
if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "MatMul", {9, 13}) &&
!graph_utils::IsSupportedOptypeVersionAndDomain(node, "FusedMatMul", {1}, kMSDomain) &&
!graph_utils::IsSupportedOptypeVersionAndDomain(node, "TransposeMatMul", {1}, kMSDomain)) {
return Status::OK();
}
@ -205,7 +206,7 @@ Status ProcessNode(
}
NodeAttributes fused_node_attrs =
node.OpType() == "TransposeMatMul" ? node.GetAttributes() : NodeAttributes{};
(node.OpType() == "TransposeMatMul") || (node.OpType() == "FusedMatMul") ? node.GetAttributes() : NodeAttributes{};
{
ONNX_NAMESPACE::AttributeProto& alpha_attr = fused_node_attrs["alpha"];
@ -240,7 +241,7 @@ Status ProcessNode(
Node& matmul_scale_node = graph.AddNode(
graph.GenerateNodeName(node.Name() + "_FusedMatMulAndScale"),
"TransposeMatMul",
"FusedMatMul",
"Fused MatMul and Scale",
fused_node_inputs,
fused_node_outputs,

4
onnxruntime/core/optimizer/matmul_scale_fusion.h
View file

@ -7,11 +7,11 @@ namespace onnxruntime {
/**
* Fuses MatMul with surrounding scales (multiplies or divides) by a constant
* scalar into TransposeMatMul (which supports scaling the product).
* scalar into FusedMatMul (which supports scaling the product).
*
* For example, given matrices A and B and constant scalars t, u, and v:
* Mul(v, MatMul(Mul(t, A), Mul(u, B))
* -> TransposeMatMul(A, B, alpha=t*u*v)
* -> FusedMatMul(A, B, alpha=t*u*v)
*
* Note: Since both leading and following scales may be fused into a single
* scale, the order and number of mathematical operations may change. This may

7
onnxruntime/core/optimizer/matmul_transpose_fusion.cc
View file

@ -97,7 +97,8 @@ Status MatmulTransposeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_
ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level, logger));
if ((!graph_utils::IsSupportedOptypeVersionAndDomain(node, "MatMul", {9, 13}) &&
!graph_utils::IsSupportedOptypeVersionAndDomain(node, "TransposeMatMul", {1, 13}, kMSDomain)) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(node, "FusedMatMul", {1}, kMSDomain) &&
!graph_utils::IsSupportedOptypeVersionAndDomain(node, "TransposeMatMul", {1}, kMSDomain)) ||
!graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders())) {
continue;
}
@ -130,14 +131,14 @@ Status MatmulTransposeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_
const std::vector<NodeArg*> output_defs{node.MutableOutputDefs()[0]};
Node& matmul_node = graph.AddNode(graph.GenerateNodeName("MatMul_With_Transpose"),
"TransposeMatMul",
"FusedMatMul",
"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() == "TransposeMatMul") {
if ((node.OpType() == "TransposeMatMul") || (node.OpType() == "FusedMatMul")) {
transpose_left ^= static_cast<bool>(node.GetAttributes().at("transA").i());
transpose_right ^= static_cast<bool>(node.GetAttributes().at("transB").i());
alpha = node.GetAttributes().at("alpha").f();

95
onnxruntime/core/providers/cpu/math/matmul.cc
View file

@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/framework/op_kernel.h"
#include "core/providers/cpu/math/matmul.h"
#include "core/providers/cpu/math/matmul_helper.h"
#include "core/util/math.h"
#include "core/util/math_cpuonly.h"
@ -9,31 +9,6 @@
namespace onnxruntime {
template <typename T>
class MatMul final : public OpKernel {
public:
MatMul(const OpKernelInfo& info) : OpKernel(info) {}
Status Compute(OpKernelContext* context) const override;
};
#if !defined(USE_MKLML_FOR_BLAS)
template <>
class MatMul<float> final : public OpKernel {
public:
MatMul(const OpKernelInfo& info) : OpKernel(info) {}
Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override;
Status Compute(OpKernelContext* context) const override;
private:
TensorShape b_shape_;
BufferUniquePtr packed_b_;
};
#endif
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
MatMul,
1, 8,
@ -117,7 +92,6 @@ Status MatMul<T>::Compute(OpKernelContext* ctx) const {
}
#if !defined(USE_MKLML_FOR_BLAS)
Status MatMul<float>::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) {
is_packed = false;
@ -130,8 +104,11 @@ Status MatMul<float>::PrePack(const Tensor& tensor, int input_idx, bool& is_pack
return Status::OK();
}
const size_t K = static_cast<size_t>(b_shape_[0]);
const size_t N = static_cast<size_t>(b_shape_[1]);
const bool trans_b = trans_b_attr_ && b_shape_.NumDimensions() != 1;
const size_t K = trans_b ? static_cast<size_t>(b_shape_[1])
: static_cast<size_t>(b_shape_[0]);
const size_t N = trans_b ? static_cast<size_t>(b_shape_[0])
: static_cast<size_t>(b_shape_[1]);
const size_t packed_b_size = MlasGemmPackBSize(N, K);
if (packed_b_size == 0) {
@ -141,20 +118,31 @@ Status MatMul<float>::PrePack(const Tensor& tensor, int input_idx, bool& is_pack
auto alloc = Info().GetAllocator(0, OrtMemTypeDefault);
auto* packed_b_data = alloc->Alloc(packed_b_size);
packed_b_ = BufferUniquePtr(packed_b_data, BufferDeleter(alloc));
MlasGemmPackB(CblasNoTrans, N, K, tensor.Data<float>(), N, packed_b_data);
MlasGemmPackB(trans_b ? CblasTrans : CblasNoTrans,
N,
K,
tensor.Data<float>(),
static_cast<int>(trans_b ? K : N),
packed_b_data);
is_packed = true;
}
return Status::OK();
}
#endif
Status MatMul<float>::Compute(OpKernelContext* ctx) const {
concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
const Tensor* a = ctx->Input<Tensor>(0);
const Tensor* b = packed_b_ ? nullptr : ctx->Input<Tensor>(1);
const auto& b_shape = b ? b->Shape() : b_shape_;
// match CUDA kernel implementation, ignore transpose for vectors
const bool trans_a = trans_a_attr_ && a->Shape().NumDimensions() != 1;
const bool trans_b = trans_b_attr_ && b_shape.NumDimensions() != 1;
MatMulComputeHelper helper;
ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b ? b->Shape() : b_shape_));
ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, trans_a, trans_b));
Tensor* y = ctx->Output(0, helper.OutputShape());
// Bail out early if the output is going to be empty
@ -168,34 +156,39 @@ Status MatMul<float>::Compute(OpKernelContext* ctx) const {
// TODO: replace it with GemmBatch for performance, it's OK for now as GemmBatch unrolls as well
size_t max_len = helper.OutputOffsets().size();
for (size_t i = 0; i < max_len; i++) {
#if !defined(USE_MKLML_FOR_BLAS)
if (packed_b_) {
MlasGemm(CblasNoTrans,
static_cast<size_t>(helper.M()),
static_cast<size_t>(helper.N()),
static_cast<size_t>(helper.K()),
1.0f,
a_data + helper.LeftOffsets()[i],
static_cast<size_t>(helper.K()),
packed_b_.get(),
0.0f,
y_data + helper.OutputOffsets()[i],
static_cast<size_t>(helper.N()),
thread_pool);
} else {
math::MatMul(
static_cast<int>(helper.M()),
static_cast<int>(helper.N()),
static_cast<int>(helper.K()),
MlasGemm(
trans_a ? CblasTrans : CblasNoTrans,
static_cast<size_t>(helper.M()),
static_cast<size_t>(helper.N()),
static_cast<size_t>(helper.K()),
alpha_attr_,
a_data + helper.LeftOffsets()[i],
b_data + helper.RightOffsets()[i],
static_cast<size_t>(trans_a ? helper.M() : helper.K()),
packed_b_.get(),
0.0f,
y_data + helper.OutputOffsets()[i],
static_cast<size_t>(helper.N()),
thread_pool);
continue;
}
#endif
math::Gemm<float, concurrency::ThreadPool>(
trans_a ? CblasTrans : CblasNoTrans,
trans_b ? CblasTrans : CblasNoTrans,
helper.M(),
helper.N(),
helper.K(),
alpha_attr_,
a_data + helper.LeftOffsets()[i],
b_data + helper.RightOffsets()[i],
0.0f,
y_data + helper.OutputOffsets()[i],
thread_pool);
}
return Status::OK();
}
#endif
} // namespace onnxruntime

43
onnxruntime/core/providers/cpu/math/matmul.h Normal file
View file

@ -0,0 +1,43 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/framework/op_kernel.h"
namespace onnxruntime {
template <typename T>
class MatMul final : public OpKernel {
public:
MatMul(const OpKernelInfo& info) : OpKernel(info) {}
Status Compute(OpKernelContext* context) const override;
};
template <>
class MatMul<float> final : public OpKernel {
public:
MatMul(const OpKernelInfo& info) : OpKernel(info) {
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);
}
#if !defined(USE_MKLML_FOR_BLAS)
Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override;
#endif
Status Compute(OpKernelContext* context) const override;
private:
TensorShape b_shape_;
BufferUniquePtr packed_b_;
// For FusedMatMul and TransposeMatMul contrib ops
float alpha_attr_;
int64_t trans_a_attr_;
int64_t trans_b_attr_;
};
} // namespace onnxruntime

55
onnxruntime/test/contrib_ops/transpose_scale_matmul_op_test.cc → onnxruntime/test/contrib_ops/fused_matmul_op_test.cc
View file

@ -131,10 +131,10 @@ void ProcessInputs(const std::vector<int64_t>& input_dims, const std::vector<T>&
}
template <typename T>
void RunTransposeMatMulTest(int32_t opset_version = 7, bool transa = false, bool transb = false, float alpha = 1.0f) {
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};
for (auto t : GenerateSimpleTestCases<T>()) {
OpTester test("TransposeMatMul", opset_version, onnxruntime::kMSDomain);
OpTester test(op_name, opset_version, onnxruntime::kMSDomain);
std::vector<int64_t> input0_dims(t.input0_dims);
std::vector<T> input0_vals;
@ -146,7 +146,7 @@ void RunTransposeMatMulTest(int32_t opset_version = 7, bool transa = false, bool
test.AddInput<T>("A", input0_dims, input0_vals);
test.AddInput<T>("B", input1_dims, input1_vals);
test.AddInput<T>("B", input1_dims, input1_vals, is_b_constant);
test.AddAttribute("transA", (int64_t)transa);
test.AddAttribute("transB", (int64_t)transb);
@ -165,32 +165,57 @@ void RunTransposeMatMulTest(int32_t opset_version = 7, bool transa = false, bool
}
}
TEST(TransposeMatMulOpTest, FloatTypeNoTranspose) {
RunTransposeMatMulTest<float>(1);
TEST(FusedMatMulOpTest, FloatTypeNoTranspose) {
RunFusedMatMulTest<float>("FusedMatMul", 1);
}
#ifdef USE_CUDA // double support only implemented in CUDA kernel
TEST(TransposeMatMulOpTest, DoubleTypeNoTranspose) {
RunTransposeMatMulTest<double>(1);
RunFusedMatMulTest<double>("TransposeMatMul", 1);
}
TEST(FusedMatMulOpTest, DoubleTypeNoTranspose) {
RunFusedMatMulTest<double>("FusedMatMul", 1);
}
#endif
TEST(TransposeMatMulOpTest, FloatTypeTransposeA) {
RunTransposeMatMulTest<float>(1, true, false);
TEST(FusedMatMulOpTest, FloatTypeTransposeA) {
RunFusedMatMulTest<float>("FusedMatMul", 1, true, false);
}
TEST(TransposeMatMulOpTest, FloatTypeTransposeB) {
RunTransposeMatMulTest<float>(1, false, true);
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);
}
TEST(TransposeMatMulOpTest, FloatTypeTransposeAB) {
RunTransposeMatMulTest<float>(1, true, 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);
}
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);
// 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);
}
TEST(TransposeMatMulOpTest, FloatTypeScale) {
RunTransposeMatMulTest<float>(1, false, false, 0.5f);
RunTransposeMatMulTest<float>(1, true, false, 2.0f);
RunTransposeMatMulTest<float>(1, true, true, 4.0f);
RunFusedMatMulTest<float>("TransposeMatMul", 1, false, false, 0.5f);
RunFusedMatMulTest<float>("TransposeMatMul", 1, true, false, 2.0f);
RunFusedMatMulTest<float>("TransposeMatMul", 1, true, true, 4.0f);
// now run tests with b constant.
RunFusedMatMulTest<float>("TransposeMatMul", 1, false, false, 0.5f, true);
RunFusedMatMulTest<float>("TransposeMatMul", 1, true, false, 2.0f, true);
RunFusedMatMulTest<float>("TransposeMatMul", 1, true, true, 4.0f, true);
}
} // namespace transpose_matmul

28
onnxruntime/test/optimizer/graph_transform_test.cc
View file

@ -711,7 +711,7 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusion) {
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Transpose"] == 0);
ASSERT_TRUE(op_to_count["MatMul"] == 0);
ASSERT_TRUE(op_to_count["TransposeMatMul"] == 1);
ASSERT_TRUE(op_to_count["FusedMatMul"] == 1);
}
TEST_F(GraphTransformationTests, TransposeMatmulFusionOnTwoTranspose) {
@ -728,10 +728,10 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusionOnTwoTranspose) {
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Transpose"] == 0);
ASSERT_TRUE(op_to_count["MatMul"] == 0);
ASSERT_TRUE(op_to_count["TransposeMatMul"] == 1);
ASSERT_TRUE(op_to_count["FusedMatMul"] == 1);
auto& node = *graph.Nodes().begin();
ASSERT_TRUE(node.OpType() == "TransposeMatMul");
ASSERT_TRUE(node.OpType() == "FusedMatMul");
ASSERT_TRUE(static_cast<bool>(node.GetAttributes().at("transA").i()));
ASSERT_TRUE(static_cast<bool>(node.GetAttributes().at("transB").i()));
}
@ -750,10 +750,10 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusionOnThreeTranspose) {
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Transpose"] == 0);
ASSERT_TRUE(op_to_count["MatMul"] == 0);
ASSERT_TRUE(op_to_count["TransposeMatMul"] == 1);
ASSERT_TRUE(op_to_count["FusedMatMul"] == 1);
auto& node = *graph.Nodes().begin();
ASSERT_TRUE(node.OpType() == "TransposeMatMul");
ASSERT_TRUE(node.OpType() == "FusedMatMul");
ASSERT_FALSE(static_cast<bool>(node.GetAttributes().at("transA").i()));
ASSERT_TRUE(static_cast<bool>(node.GetAttributes().at("transB").i()));
}
@ -776,7 +776,7 @@ TEST_F(GraphTransformationTests, TransposeMatmulNoFusionOnInvalidPerm) {
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_EQ(op_to_count["Transpose"], 1);
ASSERT_EQ(op_to_count["MatMul"], 1);
ASSERT_EQ(op_to_count["TransposeMatMul"], 0);
ASSERT_EQ(op_to_count["FusedMatMul"], 0);
}
}
@ -791,7 +791,7 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusionFromTransposeMatMul) {
auto transpose_scale_matmul_node =
std::find_if(
graph.Nodes().cbegin(), graph.Nodes().cend(),
[](const Node& node) { return node.Name() == "TransposeMatMul"; });
[](const Node& node) { return node.Name() == "FusedMatMul"; });
ASSERT_NE(transpose_scale_matmul_node, graph.Nodes().cend());
expected_alpha = transpose_scale_matmul_node->GetAttributes().at("alpha").f();
}
@ -804,10 +804,10 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusionFromTransposeMatMul) {
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["TransposeMatMul"], 1);
ASSERT_EQ(op_to_count["FusedMatMul"], 1);
auto& transpose_scale_matmul_node = *graph.Nodes().begin();
ASSERT_EQ(transpose_scale_matmul_node.OpType(), "TransposeMatMul");
ASSERT_EQ(transpose_scale_matmul_node.OpType(), "FusedMatMul");
ASSERT_FALSE(static_cast<bool>(transpose_scale_matmul_node.GetAttributes().at("transA").i()));
ASSERT_FALSE(static_cast<bool>(transpose_scale_matmul_node.GetAttributes().at("transB").i()));
ASSERT_EQ(transpose_scale_matmul_node.GetAttributes().at("alpha").f(), expected_alpha);
@ -827,7 +827,7 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusionWithPreservedTranspose) {
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_EQ(op_to_count["Transpose"], 1);
ASSERT_EQ(op_to_count["MatMul"], 0);
ASSERT_EQ(op_to_count["TransposeMatMul"], 1);
ASSERT_EQ(op_to_count["FusedMatMul"], 1);
ASSERT_FALSE(graph.GraphResolveNeeded());
}
@ -3177,17 +3177,17 @@ TEST_F(GraphTransformationTests, MatMulScaleFusionFusableModels) {
EXPECT_EQ(transformed_op_counts["Mul"], 0);
EXPECT_EQ(transformed_op_counts["Div"], 0);
EXPECT_EQ(transformed_op_counts["MatMul"], 0);
EXPECT_EQ(transformed_op_counts["TransposeMatMul"], 1);
EXPECT_EQ(transformed_op_counts["FusedMatMul"], 1);
// check combined scale, individual scales should all have the same value
const float scale_value = 3.0f;
const int num_scales =
original_op_counts["Mul"] + original_op_counts["Div"] + original_op_counts["TransposeMatMul"];
original_op_counts["Mul"] + original_op_counts["Div"] + original_op_counts["FusedMatMul"];
auto fused_node = std::find_if(
graph.Nodes().cbegin(), graph.Nodes().cend(),
[](const Node& node) { return node.OpType() == "TransposeMatMul"; });
[](const Node& node) { return node.OpType() == "FusedMatMul"; });
ASSERT_NE(fused_node, graph.Nodes().cend());
auto alpha_attr = fused_node->GetAttributes().find("alpha");
@ -3225,7 +3225,7 @@ TEST_F(GraphTransformationTests, MatMulScaleFusionReusedInputScale) {
EXPECT_EQ(transformed_op_counts["Mul"], 0);
EXPECT_EQ(transformed_op_counts["Div"], 0);
EXPECT_EQ(transformed_op_counts["MatMul"], 0);
EXPECT_EQ(transformed_op_counts["TransposeMatMul"], 2);
EXPECT_EQ(transformed_op_counts["FusedMatMul"], 2);
});
}

2
onnxruntime/test/testdata/transform/fusion/matmul_scale_gen.py vendored
View file

@ -31,7 +31,7 @@ def save(model_path, nodes, inputs, outputs, initializers):
def gen(model_path,
use_transpose_matmul,
scale_input_0, scale_input_1, scale_output):
matmul_op = "TransposeMatMul" if use_transpose_matmul else "MatMul"
matmul_op = "FusedMatMul" if use_transpose_matmul else "MatMul"
matmul_domain = "com.microsoft" if use_transpose_matmul else ""
matmul_attrs = {"alpha": scale_value} if use_transpose_matmul else {}

BIN
onnxruntime/test/testdata/transform/fusion/matmul_scale_in0.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/matmul_scale_in0_in1.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/matmul_scale_in0_in1_out.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/matmul_scale_int32.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/matmul_scale_reused_input_scale.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/matmul_scale_transposescalematmul_in0_in1_out.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/matmul_scale_unfusable_div_not_scale.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/matmul_scale_unfusable_scale_not_constant.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/matmul_scale_unfusable_scale_not_scalar.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/transpose_matmul_2d_fusion_from_transpose_scale_matmul.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/transpose_matmul_2d_fusion_with_preserved_transpose.onnx vendored
View file

Binary file not shown.

BIN
onnxruntime/test/testdata/transform/fusion/transpose_matmul_4d_fusion_invalid_default_perm.onnx vendored
View file

Binary file not shown.

4
onnxruntime/test/testdata/transform/fusion/transpose_matmul_gen.py vendored
View file

@ -32,10 +32,10 @@ def gen_from_transpose_scale_matmul(model_path):
["input_0"],
["transposed_input_0"]),
helper.make_node(
"TransposeMatMul",
"FusedMatMul",
["transposed_input_0", "input_1"],
["output"],
"TransposeMatMul",
"FusedMatMul",
"",
msdomain.domain,
alpha=3.0, transA=1)

2
orttraining/orttraining/core/graph/gradient_builder.cc
View file

@ -360,7 +360,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetMatMulGradient) {
} else {
ArgDef pre_reduce_grad_1 = IA("PreReduceGrad1");
result.push_back(
NodeDef(OpDef{"TransposeMatMul", kMSDomain, 1},
NodeDef(OpDef{"FusedMatMul", kMSDomain, 1},
{A, GO(0)},
{pre_reduce_grad_1},
{{"transA", MakeAttribute("transA", int64_t(1))}}));