[DML EP] Add MatMul + SoftMax fusion (#15240)

This commit is contained in:
Patrice Vignola 2023-04-11 08:31:04 -07:00 committed by GitHub
parent 7c927bb95c
commit 3be5bfe363
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 272 additions and 3 deletions

View file

@ -33,6 +33,7 @@ Do not modify directly.*
* <a href="#com.microsoft.FusedConv">com.microsoft.FusedConv</a>
* <a href="#com.microsoft.FusedGemm">com.microsoft.FusedGemm</a>
* <a href="#com.microsoft.FusedMatMul">com.microsoft.FusedMatMul</a>
* <a href="#com.microsoft.FusedMatMulActivation">com.microsoft.FusedMatMulActivation</a>
* <a href="#com.microsoft.GatedRelativePositionBias">com.microsoft.GatedRelativePositionBias</a>
* <a href="#com.microsoft.GatherND">com.microsoft.GatherND</a>
* <a href="#com.microsoft.Gelu">com.microsoft.Gelu</a>
@ -1797,6 +1798,63 @@ This version of the operator has been available since version 1 of the 'com.micr
</dl>
### <a name="com.microsoft.FusedMatMulActivation"></a><a name="com.microsoft.fusedmatmulactivation">**com.microsoft.FusedMatMulActivation**</a>
Executes the same operation as FusedMatMul, but also has an activation function fused to its output.
#### Version
This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
#### Attributes
<dl>
<dt><tt>activation</tt> : string (required)</dt>
<dd></dd>
<dt><tt>activation_alpha</tt> : float</dt>
<dd></dd>
<dt><tt>activation_axis</tt> : int</dt>
<dd></dd>
<dt><tt>activation_beta</tt> : float</dt>
<dd></dd>
<dt><tt>activation_gamma</tt> : float</dt>
<dd></dd>
<dt><tt>alpha</tt> : float</dt>
<dd>Scalar multiplier for the product of the input tensors.</dd>
<dt><tt>transA</tt> : int</dt>
<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
<dl>
<dt><tt>A</tt> : T</dt>
<dd>N-dimensional matrix A</dd>
<dt><tt>B</tt> : T</dt>
<dd>N-dimensional matrix B</dd>
</dl>
#### Outputs
<dl>
<dt><tt>Y</tt> : T</dt>
<dd>Matrix multiply results</dd>
</dl>
#### Type Constraints
<dl>
<dt><tt>T</tt> : tensor(float16), tensor(float), tensor(double), tensor(bfloat16)</dt>
<dd>Constrain input and output types to float tensors.</dd>
</dl>
### <a name="com.microsoft.GatedRelativePositionBias"></a><a name="com.microsoft.gatedrelativepositionbias">**com.microsoft.GatedRelativePositionBias**</a>
query_layer = (query_layer + query_bias).reshape(batch_size, seq_len, num_heads, head_size).transpose(1, 2)

View file

@ -1180,6 +1180,7 @@ Do not modify directly.*
|DequantizeLinear|*in* x:**T1**<br> *in* x_scale:**T2**<br> *in* x_zero_point:**T1**<br> *out* y:**T2**|1+|**T1** = tensor(float)<br/> **T2** = tensor(uint8)|
|EmbedLayerNormalization|*in* input_ids:**T1**<br> *in* segment_ids:**T1**<br> *in* word_embedding:**T**<br> *in* position_embedding:**T**<br> *in* segment_embedding:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* mask:**T1**<br> *in* position_ids:**T1**<br> *out* output:**T**<br> *out* mask_index:**T1**<br> *out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
|FusedMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|FusedMatMulActivation|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|Gelu|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|GroupNorm|*in* X:**T**<br> *in* gamma:**M**<br> *in* beta:**M**<br> *out* Y:**T**|1+|**M** = tensor(float), tensor(float16)<br/> **T** = tensor(float), tensor(float16)|
|NhwcConv|*in* X:**T**<br> *in* W:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|

View file

@ -1695,6 +1695,10 @@ constexpr 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";
constexpr const char* FusedMatMulActivation_doc = R"DOC(
Executes the same operation as FusedMatMul, but also has an activation function fused to its output.
)DOC";
ONNX_MS_OPERATOR_SET_SCHEMA(TransposeMatMul, 1,
OpSchema()
.Input(0, "A", "N-dimensional matrix A", "T")
@ -1747,6 +1751,53 @@ ONNX_MS_OPERATOR_SET_SCHEMA(FusedMatMul, 1,
.SetDoc(FusedMatMul_doc)
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { FusedMatMulShapeInference(ctx); }));
ONNX_MS_OPERATOR_SET_SCHEMA(FusedMatMulActivation, 1,
OpSchema()
.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("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))
.Attr(
"activation",
"",
AttributeProto::STRING)
.Attr(
"activation_alpha",
"",
AttributeProto::FLOAT,
OPTIONAL_VALUE)
.Attr(
"activation_beta",
"",
AttributeProto::FLOAT,
OPTIONAL_VALUE)
.Attr(
"activation_gamma",
"",
AttributeProto::FLOAT,
OPTIONAL_VALUE)
.Attr(
"activation_axis",
"",
AttributeProto::INT,
OPTIONAL_VALUE)
.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(FusedMatMulActivation_doc)
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { FusedMatMulShapeInference(ctx); }));
ONNX_MS_OPERATOR_SET_SCHEMA(SparseToDenseMatMul, 1,
OpSchema()
.Input(0, "A", "2-dimensional sparse matrix A. Either COO or CSR format", "T")

View file

@ -67,6 +67,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FastGelu);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FusedConv);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FusedGemm);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FusedMatMul);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FusedMatMulActivation);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatherND);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Gelu);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QuickGelu);
@ -161,6 +162,7 @@ class OpSet_Microsoft_ver1 {
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FusedConv)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FusedGemm)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FusedMatMul)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FusedMatMulActivation)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatherND)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Gelu)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QuickGelu)>());

View file

@ -45,6 +45,7 @@
#include "core/optimizer/identical_children_consolidation.h"
#include "core/optimizer/identity_elimination.h"
#include "core/optimizer/layer_norm_fusion.h"
#include "core/optimizer/matmul_activation_fusion.h"
#include "core/optimizer/matmul_add_fusion.h"
#include "core/optimizer/matmul_integer_to_float.h"
#include "core/optimizer/matmul_scale_fusion.h"
@ -183,6 +184,7 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
#ifndef DISABLE_CONTRIB_OPS
const InlinedHashSet<std::string_view> cpu_ep = {onnxruntime::kCpuExecutionProvider};
#endif
const InlinedHashSet<std::string_view> dml_ep = {onnxruntime::kDmlExecutionProvider};
switch (level) {
case TransformerLevel::Level1: {
// RewriteRule optimizations are the simplest (they generally remove unnecessary nodes and are cheap to run)
@ -200,7 +202,7 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
// Put ConstantSharing before CommonSubexpressionElimination by intention as it can create more opportunities for
// CSE. For example, if A and B nodes both do Add operation with a same value but different initializers, by
// default, CSE will not merge them, because the different initializers are represented by different NodeArg.
if (session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsDisableDoubleQDQRemover, "0") == "0"){
if (session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsDisableDoubleQDQRemover, "0") == "0") {
transformers.emplace_back(std::make_unique<DoubleQDQPairsRemover>());
}
transformers.emplace_back(std::make_unique<ConstantSharing>());
@ -302,6 +304,7 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
transformers.emplace_back(std::make_unique<QuickGeluFusion>(cpu_cuda_dml_rocm_eps));
transformers.emplace_back(std::make_unique<MatMulScaleFusion>(cpu_cuda_dml_rocm_eps));
transformers.emplace_back(std::make_unique<MatMulActivationFusion>(dml_ep));
// GeluApproximation has side effects which may change results. It needs to be manually enabled,
// or alternatively the model can be updated offline using a model conversion script

View file

@ -0,0 +1,83 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/optimizer/initializer.h"
#include "core/optimizer/matmul_activation_fusion.h"
#include "core/graph/graph_utils.h"
using namespace ONNX_NAMESPACE;
using namespace ::onnxruntime::common;
namespace onnxruntime {
namespace {
// Don't check if the op is Deprecated. In ONNX Runtime's world, there is no deprecation.
bool IsSupportedOptypeVersionAndDomain(const Node& node, const std::string& op_type,
std::initializer_list<ONNX_NAMESPACE::OperatorSetVersion> versions,
std::string_view domain) {
return (node.OpType() == op_type && graph_utils::MatchesOpSinceVersion(node, versions) &&
graph_utils::MatchesOpSetDomain(node, domain));
}
// If the op has multiple versions, here we require it must have a single implementation that can work across all the
// versions. Because in the fusion, we discarded the op version information.
bool IsFusableActivation(const Node& node) {
return IsSupportedOptypeVersionAndDomain(node, "Softmax", {1, 11, 13}, kOnnxDomain);
}
} // namespace
Status MatMulActivationFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
const logging::Logger& logger) const {
GraphViewer graph_viewer(graph);
const auto& order = graph_viewer.GetNodesInTopologicalOrder();
for (auto index : order) {
auto* node_ptr = graph.GetNode(index);
if (!node_ptr)
continue; // node was removed
auto& node = *node_ptr;
ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level, logger));
if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "FusedMatMul", {1}, kMSDomain) ||
!graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders()) || node.GetOutputEdgesCount() != 1) {
continue;
}
const Node& next_node = *(node.OutputNodesBegin());
if (!IsFusableActivation(next_node) || next_node.GetExecutionProviderType() != node.GetExecutionProviderType()) {
continue;
}
if (graph.NodeProducesGraphOutput(node)) {
continue;
}
Node& act_node = *graph.GetNode(next_node.Index()); // get mutable reference
Node& fused_node = graph.AddNode(graph.GenerateNodeName(node.Name() + "_FusedActivation"), "FusedMatMulActivation",
node.Description() + " with activation " + act_node.OpType(),
node.MutableInputDefs(), {}, &node.GetAttributes(), kMSDomain);
// Add a new attribute to specify the activation type
fused_node.AddAttribute("activation", act_node.OpType());
// Assign provider to this new node. Provider should be same as the provider for old node.
fused_node.SetExecutionProviderType(node.GetExecutionProviderType());
// Add optional attributes for activations
const NodeAttributes& attrs = act_node.GetAttributes();
for (const auto& attr : attrs) {
AttributeProto fused_node_attr(attr.second);
fused_node_attr.set_name("activation_" + attr.first);
fused_node.AddAttributeProto(std::move(fused_node_attr));
}
// move output definitions and edges from act_node to fused_node. delete node and act_node.
graph_utils::FinalizeNodeFusion(graph, {node, act_node}, fused_node);
modified = true;
}
return Status::OK();
}
} // namespace onnxruntime

View file

@ -0,0 +1,18 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/optimizer/graph_transformer.h"
namespace onnxruntime {
class MatMulActivationFusion : public GraphTransformer {
public:
MatMulActivationFusion(const InlinedHashSet<std::string_view>& compatible_execution_providers = {}) noexcept
: GraphTransformer("MatMulActivationFusion", compatible_execution_providers) {}
Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override;
};
} // namespace onnxruntime

View file

@ -57,8 +57,8 @@ public:
std::vector<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
std::vector<DML_TENSOR_DESC> outputDescs = GetDmlOutputDescs();
std::optional<ActivationOperatorDesc> fusedActivation = FusionHelpers::TryGetFusedActivationDesc(kernelInfo);
DML_OPERATOR_DESC fusedActivationDmlDesc = fusedActivation ? fusedActivation->GetDmlDesc() : DML_OPERATOR_DESC();
std::optional<ActivationOperatorDescWrapper> fusedActivation = FusionHelpers::TryGetGraphFusedActivationDesc(kernelInfo);
DML_OPERATOR_DESC fusedActivationDmlDesc = fusedActivation ? fusedActivation->desc.GetDmlDesc() : DML_OPERATOR_DESC();
const float alpha = kernelInfo.GetOptionalAttribute<float>(AttrName::Alpha, 1.0f);
@ -79,5 +79,6 @@ public:
};
DML_OP_DEFINE_CREATION_FUNCTION(FusedMatMul, DmlOperatorFusedMatMul);
DML_OP_DEFINE_CREATION_FUNCTION(FusedMatMulActivation, DmlOperatorFusedMatMul);
} // namespace Dml

View file

@ -312,6 +312,7 @@ DML_OP_EXTERN_CREATION_FUNCTION(Affine);
DML_OP_EXTERN_CREATION_FUNCTION(Dropout);
DML_OP_EXTERN_CREATION_FUNCTION(MatMul);
DML_OP_EXTERN_CREATION_FUNCTION(FusedMatMul);
DML_OP_EXTERN_CREATION_FUNCTION(FusedMatMulActivation);
DML_OP_EXTERN_CREATION_FUNCTION(Cast);
DML_OP_EXTERN_CREATION_FUNCTION(CastLike15);
DML_OP_EXTERN_CREATION_FUNCTION(MemcpyFromHost);
@ -858,6 +859,7 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation
{REG_INFO_MS( 1, Gelu, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
{REG_INFO_MS( 1, BiasGelu, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
{REG_INFO_MS( 1, FusedMatMul, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
{REG_INFO_MS( 1, FusedMatMulActivation, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
{REG_INFO_MS( 1, QLinearSigmoid, typeNameListDefault, supportedTypeListQLinearSigmoid, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryQLinearSigmoid)},
{REG_INFO_MS( 1, Attention, typeNameListAttention, supportedTypeListAttention, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryAttention)},

View file

@ -348,6 +348,44 @@ namespace Dml
return activation;
}
std::optional<ActivationOperatorDescWrapper> TryGetGraphFusedActivationDesc(const MLOperatorKernelCreationContext& kernelInfo)
{
if (!kernelInfo.HasAttribute(AttrName::GraphFusedActivation, MLOperatorAttributeType::String))
{
return std::nullopt; // No activation
}
auto activationName = kernelInfo.GetAttribute(AttrName::GraphFusedActivation);
ActivationOperatorDescWrapper activation = {};
if (activationName == "Softmax")
{
const uint32_t onnxDimCount = gsl::narrow_cast<uint32_t>(kernelInfo.GetTensorShapeDescription().GetInputTensorShape(0).size());
int onnxAxis = HandleNegativeAxis(kernelInfo.GetOptionalAttribute<int>(AttrName::GraphFusedAxis, -1), onnxDimCount);
auto dmlAdjustedAxis = GetDmlAdjustedAxis(onnxAxis, onnxDimCount, kernelInfo.GetTensorShapeDescription().GetInputTensorDimensionCount(0));
// If the axis is supported by Softmax, use this version instead since it's more likely to be supported by metacommands
if (dmlAdjustedAxis == onnxDimCount - 1)
{
activation.desc.activationType = DML_OPERATOR_ACTIVATION_SOFTMAX;
}
else
{
activation.desc.activationType = DML_OPERATOR_ACTIVATION_SOFTMAX1;
activation.dmlAxes.push_back(dmlAdjustedAxis);
activation.desc.params.softmax1.Axes = activation.dmlAxes.data();
activation.desc.params.softmax1.AxisCount = gsl::narrow_cast<uint32_t>(activation.dmlAxes.size());
}
}
else
{
ML_INVALID_ARGUMENT("Unsupported activation function.");
}
return activation;
}
/*static*/ std::string GetFusedAttributeName(std::string_view name)
{
return std::string("fused_").append(name);

View file

@ -7,6 +7,12 @@ namespace Dml
{
constexpr float DefaultEpsilon = 0.00001f;
struct ActivationOperatorDescWrapper
{
ActivationOperatorDesc desc;
std::vector<uint32_t> dmlAxes;
};
namespace ActivationHelper
{
float GetDefaultAlpha(DML_OPERATOR_TYPE function);
@ -41,6 +47,7 @@ namespace Dml
bool IsFusableActivationOperator(std::string_view opType, std::string_view domain, int sinceVersion);
std::optional<ActivationOperatorDesc> TryGetFusedActivationDesc(const MLOperatorKernelCreationContext& kernelInfo);
std::optional<ActivationOperatorDescWrapper> TryGetGraphFusedActivationDesc(const MLOperatorKernelCreationContext& kernelInfo);
// Produces names for attributes added to fused kernels. This effectively prepends a string to distinguish ONNX
// attributes from those added dynamically via operator fusion. For example, this function would be used to

View file

@ -114,6 +114,9 @@ namespace AttrName
static constexpr const char* Activation = "activation";
static constexpr const char* Groups = "groups";
static constexpr const char* GraphFusedActivation = "activation";
static constexpr const char* GraphFusedAxis = "activation_axis";
} // namespace AttrName
namespace AttrValue

View file

@ -1654,6 +1654,7 @@ using ShapeInferenceHelper_DmlFusedMeanVarianceNormalization = GetOutputShapeAsI
using ShapeInferenceHelper_DmlFusedGemm = GemmHelper;
using ShapeInferenceHelper_DmlFusedMatMul = MatMulHelper;
using ShapeInferenceHelper_FusedMatMul = FusedMatMulHelper;
using ShapeInferenceHelper_FusedMatMulActivation = FusedMatMulHelper;
using ShapeInferenceHelper_DmlFusedAdd = GetBroadcastedOutputShapeHelper;
using ShapeInferenceHelper_DmlFusedSum = GetBroadcastedOutputShapeHelper;

View file

@ -406,6 +406,7 @@ namespace OperatorHelper
static const int sc_sinceVer_Gelu = 1;
static const int sc_sinceVer_BiasGelu = 1;
static const int sc_sinceVer_FusedMatMul = 1;
static const int sc_sinceVer_FusedMatMulActivation = 1;
static const int sc_sinceVer_QLinearSigmoid = 1;
static const int sc_sinceVer_Attention = 1;
static const int sc_sinceVer_SkipLayerNormalization = 1;