mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
[DML EP] Add MatMul + SoftMax fusion (#15240)
This commit is contained in:
parent
7c927bb95c
commit
3be5bfe363
14 changed files with 272 additions and 3 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)|
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)>());
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
83
onnxruntime/core/optimizer/matmul_activation_fusion.cc
Normal file
83
onnxruntime/core/optimizer/matmul_activation_fusion.cc
Normal 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
|
||||
18
onnxruntime/core/optimizer/matmul_activation_fusion.h
Normal file
18
onnxruntime/core/optimizer/matmul_activation_fusion.h
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)},
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Reference in a new issue