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 fusion support for Dnnl execution provider (#9897)

Browse source 
* Op fusion support added

In addition the following op fusions are detected
 - ConvRelu
 - MatMulAdd

This change includes
  - Change abstraction of Subgraph + node + tensor to support delete insert
    modify
  - add nodearg class to establish connection from tensor to node
  - add graphtransformer class to support fusion
  - add topological sort to ensure propoer node ordering after fusion
  - add convrelu + matmuladd primitive to support execution of fused nodes
  - Fix FusionResolution with missing tensors

    when fusing, if the target node contains fewer tensors then original
    patterns (Gelu and FastGelu ignores many initializers), potentially delete them
    also from inputs and initializers

    Also check tensor has no producer and consumer before deleting

Signed-off-by: Wang <zhaoyang.wang@intel.com>

* Gelu and FastGelu Fusion for DNNL EP

The basics of the Gelu/FastGelu code is modeled after:
   - core/optimizer/fast_gelu_fusion.cc and
   - core/optimizer/gelu_fusion.cc

OneDNN does not have support for 'Erf' unless it is part of 'Gelu'.

This results in detecting 'Gelu' fusion twice. Once when detecting
if the 'Erf' Operator is supported and again in the subgraph transformer
code. The capability code is finding the Gelu using onnxruntime:GraphViewer
and onnxruntime::Node.  While the transformer code is using DnnlSubgraph
and DnnlNode.  This results in two parts of code looking for the same
pattern but unfortanatly having little code reuse.

This also adds support for Biased versions of Gelu and FastGelu if they already
exist in a model.

Signed-off-by: George Nash <george.nash@intel.com>

* Code Clean Up

Signed-off-by: Wang <zhaoyang.wang@intel.com>

Co-authored-by: Wang <zhaoyang.wang@intel.com>
This commit is contained in:
George Nash 2021-12-02 18:11:19 -08:00 committed by GitHub
parent 06e63218be
commit 1c38ceda49
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 1597 additions and 98 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
onnxruntime/core/providers/dnnl/dnnl_execution_provider.cc
View file

@ -10,6 +10,7 @@
#include "dnnl_execution_provider.h"
#include "dnnl_fwd.h"
#include "dnnl_node_capability.h"
#include "core/providers/dnnl/subgraph/dnnl_subgraph_transformer.h"
#include <iomanip>
#include <fstream>
@ -50,6 +51,11 @@ DNNLExecutionProvider::DNNLExecutionProvider(const DNNLExecutionProviderInfo& in
if (!debug_log_env.empty()) {
debug_log_ = (std::stoi(debug_log_env) == 0 ? false : true);
}
const std::string fusion_env = onnxruntime::GetEnvironmentVar("ORT_DNNL_ENABLE_FUSION");
if (!fusion_env.empty()) {
enable_fusion_ = (std::stoi(fusion_env) == 0 ? false : true);
}
} // namespace onnxruntime
DNNLExecutionProvider::~DNNLExecutionProvider() {
@ -292,6 +298,11 @@ Status DNNLExecutionProvider::Compile(const std::vector<Node*>& fused_nodes,
auto dnnl_subgraph = std::make_unique<ort_dnnl::DnnlSubgraph>(ort_dnnl::DnnlSubgraph(*graph_body_viewer.get()));
subgraphs_.emplace(fused_node->Name(), std::move(dnnl_subgraph));
//apply transformation to subgraph
if (enable_fusion_) {
ort_dnnl::DnnlGraphTransformer().Apply(*subgraphs_[fused_node->Name()].get());
}
//subgraph primitive
auto dnnl_subgraph_primitive = std::make_unique<ort_dnnl::DnnlSubgraphPrimitive>(*subgraphs_[fused_node->Name()].get());
{

2
onnxruntime/core/providers/dnnl/dnnl_execution_provider.h
View file

@ -46,6 +46,8 @@ class DNNLExecutionProvider : public IExecutionProvider {
// dump subgraphs to onnx format for debugging purpose
bool dump_subgraphs_ = false;
bool debug_log_ = false;
//enable fusion by default
bool enable_fusion_ = true;
};
} // namespace onnxruntime

240
onnxruntime/core/providers/dnnl/dnnl_node_capability.cc
View file

@ -4,6 +4,11 @@
#include "dnnl_node_capability.h"
#include "dnnl.hpp"
#ifndef _USE_MATH_DEFINES
#define _USE_MATH_DEFINES
#endif
#include <cmath>
namespace onnxruntime {
// DnnlDefaultNodeCapability class
//-------------------------------------
@ -443,6 +448,18 @@ bool DnnlElementwiseCapability::IsDimensionSupported(const Node* node) const {
return true;
}
bool IsScalar(const NodeArg* node_arg) {
if (node_arg->Shape()->dim_size() == 0) {
return true;
}
for (int j = 0; j < node_arg->Shape()->dim_size(); ++j) {
if (node_arg->Shape()->dim(j).dim_value() != 1) {
return false;
}
}
return true;
}
// DnnlPowNodeCapability class
//-------------------------------------
bool DnnlPowNodeCapability::Supported(const Node* node, const GraphViewer& graph_viewer) const {
@ -468,13 +485,8 @@ bool DnnlPowNodeCapability::IsDimensionSupported(const Node* node, const GraphVi
if (!graph_viewer.IsConstantInitializer(node_inputs[1]->Name(), true)) {
return false;
}
if (node_inputs[1]->Shape()->dim_size() == 0) {
return true;
}
for (int j = 0; j < node_inputs[1]->Shape()->dim_size(); ++j) {
if (node_inputs[1]->Shape()->dim(j).dim_value() != 1) {
return false;
}
if (!IsScalar(node_inputs[1])) {
return false;
}
}
@ -551,4 +563,218 @@ bool DnnlSqueezeNodeCapability::IsDimensionSupported(const Node* node, const Gra
return true;
}
bool DnnlErfNodeCapability::Supported(const Node* node, const GraphViewer& graph_viewer) const {
ORT_UNUSED_PARAMETER(graph_viewer);
if (!IsTypeSupported(node)) return false;
if (!IsErfPartOfGelu(node, graph_viewer)) return false;
return true;
}
bool DnnlErfNodeCapability::IsInitilizedWithExpectedValue(const GraphViewer& graph_viewer, const NodeArg* node_arg, float expected_value) const {
//TypeAsProto()->tensor_type().elem_type()
if ((ORT_DataType)node_arg->TypeAsProto()->tensor_type().elem_type() == type_float32) {
const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr;
graph_viewer.GetInitializedTensor(node_arg->Name(), tensor_proto);
const float* val = reinterpret_cast<const float*>(tensor_proto->raw_data().data());
// Check for NaN and Inf
if (std::isnan(val[0]) || std::isinf(val[0])) {
if (std::isinf(val[0]) && std::isinf(expected_value) && (std::signbit(val[0]) == std::signbit(expected_value))) {
return true;
}
return false;
}
const float atol = 1e-8f;
const float rtol = 1e-5f;
float diff = std::abs(val[0] - expected_value);
if (diff > (atol + rtol * std::abs(expected_value))) {
return false;
}
return true;
}
return false;
}
const Node* DnnlErfNodeCapability::FirstParentByType(const Node& node, const std::string& parent_type) const {
for (auto it = node.InputNodesBegin(); it != node.InputNodesEnd(); ++it) {
if ((*it).OpType().compare(parent_type) == 0) {
return &(*it);
}
}
return nullptr;
}
bool DnnlErfNodeCapability::IsNodeFusable(const Node* node, const GraphViewer& graph_viewer) const {
if (nullptr == node) {
return false;
}
if (1 != node->GetOutputEdgesCount() &&
1 != node->OutputDefs().size()) {
return false;
}
// Check if the NodeArg outputs to the graph
if (std::find(graph_viewer.GetOutputs().begin(), graph_viewer.GetOutputs().end(), node->OutputDefs()[0]) != graph_viewer.GetOutputs().end()) {
return false;
}
return true;
}
/*
OneDNN only suports Erf if it is part of Gelu. Gelu is only possible thanks to fusion.
This code only runs when and Erf node is detected. So we check to see if the Erf is is part
a Gelu starting from the Erf. This means checking the inputs befor and after Erf. The
following pattern and some code were directly referenced from code\optimizer\gelu_fusion.cc
Subgraphs like the following fuse into Gelu.
Subgraph pattern 1:
+-------Mul(0.5)---------------------+
| |
| v
[root] --> Div -----> Erf --> Add --> Mul ==>
(B=1.4142...) (1)
Subgraph pattern 2:
+------------------------------------+
| |
| v
[root] --> Div -----> Erf --> Add --> Mul -->Mul ==>
(B=1.4142...) (1) (0.5)
After Fusion:
[root]--> Gelu ==>
*/
bool DnnlErfNodeCapability::IsErfPartOfGelu(const Node* node, const GraphViewer& graph_viewer) const {
// if DNNL fusion is not enabled then Erf is not supported.
const std::string fusion_env = onnxruntime::GetEnvironmentVar("ORT_DNNL_ENABLE_FUSION");
if (!fusion_env.empty() && std::stoi(fusion_env) == 0) {
return false;
}
if (node->InputDefs().size() != 1) {
return false;
}
if (!IsNodeFusable(node, graph_viewer)) {
return false;
}
const Node* div_node = FirstParentByType(*node, "Div");
if (!IsNodeFusable(div_node, graph_viewer)) {
return false;
}
if (!_binary.Supported(div_node, graph_viewer)) {
return false;
}
auto divisor = div_node->InputDefs()[1];
if (divisor->Shape() != nullptr) {
if (!graph_viewer.IsConstantInitializer(divisor->Name(), true)) {
return false;
}
if (!IsScalar(divisor)) {
return false;
}
// Some Bert models uses this approximation of SQRT2 in the Gelu function
float approximated_sqrt_two = 1.4142099618911743f;
if (!IsInitilizedWithExpectedValue(graph_viewer, divisor, approximated_sqrt_two) &&
!IsInitilizedWithExpectedValue(graph_viewer, divisor, static_cast<float>(M_SQRT2))) {
return false;
}
}
std::vector<const Node*> add_nodes;
for (auto i = node->OutputNodesBegin(); i != node->OutputNodesEnd(); ++i) {
add_nodes.push_back(&(*i));
}
if (add_nodes.size() != 1 && add_nodes[0]->OpType() != "Add") {
return false;
}
if (!_binary.Supported(add_nodes[0], graph_viewer)) {
return false;
}
if (!IsNodeFusable(add_nodes[0], graph_viewer)) {
return false;
}
// check that the Add Op node_arg is 1.0f
bool is_add_input0 = node->OutputDefs()[0]->Name() == add_nodes[0]->InputDefs()[0]->Name();
auto add_val = add_nodes[0]->InputDefs()[is_add_input0 ? 1 : 0];
if (add_val->Shape() != nullptr) {
if (!graph_viewer.IsConstantInitializer(add_val->Name(), true)) {
return false;
}
if (!IsScalar(add_val)) {
return false;
}
if (!IsInitilizedWithExpectedValue(graph_viewer, add_val, 1.0f)) {
return false;
}
}
std::vector<const Node*> mul1_nodes;
for (auto i = add_nodes[0]->OutputNodesBegin(); i != add_nodes[0]->OutputNodesEnd(); ++i) {
mul1_nodes.push_back(&(*i));
}
if (mul1_nodes.size() != 1 && mul1_nodes[0]->OpType() != "Mul") {
return false;
}
if (!_binary.Supported(mul1_nodes[0], graph_viewer)) {
return false;
}
// Subgraph pattern 1
// Don't return if subgraph pattern 1 is not found instead we dropdown and look for pattern 2
{
auto mul2_node = FirstParentByType(*mul1_nodes[0], "Mul");
if (mul2_node != nullptr) {
bool is_mul2_input0 = div_node->InputDefs()[0]->Name() == mul2_node->InputDefs()[0]->Name();
bool is_mul2_input1 = div_node->InputDefs()[0]->Name() == mul2_node->InputDefs()[1]->Name();
if (is_mul2_input0 ^ is_mul2_input1) {
auto mul2_val = mul2_node->InputDefs()[is_mul2_input0 ? 1 : 0];
if (mul2_val->Shape() != nullptr) {
if (IsScalar(mul2_val) &&
IsInitilizedWithExpectedValue(graph_viewer, mul2_val, 0.5f) &&
IsNodeFusable(mul2_node, graph_viewer)) {
return true;
}
}
}
}
}
// Subgraph pattern 2
{
if (!IsNodeFusable(mul1_nodes[0], graph_viewer)) {
return false;
}
std::vector<const Node*> mul2_nodes;
for (auto i = mul1_nodes[0]->OutputNodesBegin(); i != mul1_nodes[0]->OutputNodesEnd(); ++i) {
mul2_nodes.push_back(&(*i));
}
if (mul2_nodes.size() != 1 && mul2_nodes[0]->OpType() != "Mul") {
return false;
}
//Check that Mul Op node_arg is 0.5f
bool is_mul2_input0 = mul1_nodes[0]->OutputDefs()[0]->Name() == mul2_nodes[0]->InputDefs()[0]->Name();
auto mul2_val = mul2_nodes[0]->InputDefs()[is_mul2_input0 ? 1 : 0];
if (mul2_val->Shape() != nullptr) {
if (!graph_viewer.IsConstantInitializer(mul2_val->Name(), true)) {
return false;
}
if (!IsScalar(mul2_val)) {
return false;
}
if (!IsInitilizedWithExpectedValue(graph_viewer, mul2_val, 0.5f)) {
return false;
}
}
}
return true;
}
} // namespace onnxruntime

13
onnxruntime/core/providers/dnnl/dnnl_node_capability.h
View file

@ -316,4 +316,17 @@ class DnnlSqueezeNodeCapability : public DnnlDefaultNodeCapability {
bool IsDimensionSupported(const Node* node, const GraphViewer& graph_viewer) const;
};
class DnnlErfNodeCapability : public DnnlDefaultNodeCapability {
public:
DnnlErfNodeCapability() : DnnlDefaultNodeCapability({type_float32}) {}
bool Supported(const Node* node, const GraphViewer& graph_viewer) const override;
private:
bool IsErfPartOfGelu(const Node* node, const GraphViewer& graph_viewer) const;
bool IsInitilizedWithExpectedValue(const GraphViewer& graph_viewer, const NodeArg* node_arg, float expected_value) const;
const Node* FirstParentByType(const Node& node, const std::string& parent_type) const;
bool IsNodeFusable(const Node* node, const GraphViewer& graph_viewer) const;
DnnlBinaryNodeCapability _binary;
};
} // namespace onnxruntime

4
onnxruntime/core/providers/dnnl/dnnl_op_manager.cc
View file

@ -10,11 +10,15 @@ DnnlOpManager::DnnlOpManager() {
dnnl_ops_map_.emplace(std::make_pair("Add", std::unique_ptr<DnnlNodeCapability>(new DnnlBinaryNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("AveragePool", std::unique_ptr<DnnlNodeCapability>(new DnnlPoolNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("BatchNormalization", std::unique_ptr<DnnlNodeCapability>(new DnnlBatchNormalizationNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("BiasGelu", std::unique_ptr<DnnlNodeCapability>(new DnnlDefaultNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("Conv", std::unique_ptr<DnnlNodeCapability>(new DnnlDefaultNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("Div", std::unique_ptr<DnnlNodeCapability>(new DnnlBinaryNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("DynamicQuantizeLinear", std::unique_ptr<DnnlNodeCapability>(new DnnlDynamicQuantizeLinearNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("Elu", std::unique_ptr<DnnlNodeCapability>(new DnnlElementwiseCapability())));
dnnl_ops_map_.emplace(std::make_pair("Erf", std::unique_ptr<DnnlNodeCapability>(new DnnlErfNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("Exp", std::unique_ptr<DnnlNodeCapability>(new DnnlElementwiseCapability())));
dnnl_ops_map_.emplace(std::make_pair("FastGelu", std::unique_ptr<DnnlNodeCapability>(new DnnlDefaultNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("Gelu", std::unique_ptr<DnnlNodeCapability>(new DnnlDefaultNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("Gemm", std::unique_ptr<DnnlNodeCapability>(new DnnlGemmNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("GlobalAveragePool", std::unique_ptr<DnnlNodeCapability>(new DnnlPoolNodeCapability())));
dnnl_ops_map_.emplace(std::make_pair("GlobalMaxPool", std::unique_ptr<DnnlNodeCapability>(new DnnlPoolNodeCapability())));

50
onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.cc
View file

@ -12,6 +12,11 @@ namespace ort_dnnl {
DnnlConv::DnnlConv() {}
void DnnlConv::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) {
bool has_relu = false;
if (node.OpType() == "ConvRelu") {
has_relu = true;
}
auto dnnl_engine = sp.GetEngine();
auto conv_src_mem = sp.GetMemory(node.Input(IN_X));
@ -82,25 +87,35 @@ void DnnlConv::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) {
auto dst_mem_dims = InferOutputShape(node, src_dims, weight_dims_original, kernel_shape, strides, dilations, padding);
dnnl::memory::desc dst_md = dnnl::memory::desc({dst_mem_dims}, node.Input(IN_X).Type(), dnnl::memory::format_tag::any);
#ifdef ENABLE_TRAINING
#ifdef ENABLE_TRAINING
auto prop_kind = dnnl::prop_kind::forward_training;
#else
auto prop_kind = dnnl::prop_kind::forward_inference;
#endif // ENABLE_TRAINING
dnnl::primitive_attr attr;
if (has_relu) {
const float ops_scale = 1.f;
const float ops_alpha = 0.f;
const float ops_beta = 0.f;
dnnl::post_ops ops;
ops.append_eltwise(ops_scale, dnnl::algorithm::eltwise_relu, ops_alpha, ops_beta);
attr.set_post_ops(ops);
}
dnnl::convolution_forward::primitive_desc conv_pd;
if (bias_exists) {
auto conv_desc = dnnl::convolution_forward::desc(
prop_kind, dnnl::algorithm::convolution_direct,
src_md, weight_md, bias_md, dst_md,
strides, dilations, padding_left, padding_right);
conv_pd = dnnl::convolution_forward::primitive_desc(conv_desc, dnnl_engine);
prop_kind, dnnl::algorithm::convolution_direct,
src_md, weight_md, bias_md, dst_md,
strides, dilations, padding_left, padding_right);
conv_pd = dnnl::convolution_forward::primitive_desc(conv_desc, attr, dnnl_engine);
} else {
auto conv_desc = dnnl::convolution_forward::desc(
prop_kind, dnnl::algorithm::convolution_direct,
src_md, weight_md, dst_md,
strides, dilations, padding_left, padding_right);
conv_pd = dnnl::convolution_forward::primitive_desc(conv_desc, dnnl_engine);
prop_kind, dnnl::algorithm::convolution_direct,
src_md, weight_md, dst_md,
strides, dilations, padding_left, padding_right);
conv_pd = dnnl::convolution_forward::primitive_desc(conv_desc, attr, dnnl_engine);
}
// If using GPU this will move the memory from the CPU to the GPU.
@ -151,7 +166,7 @@ std::vector<int64_t> DnnlConv::GetInferedPads(DnnlNode& node,
assert(src_dims.size() == shape + 2);
for (size_t i = 0; i < shape; ++i) {
if (ComputePad(src_dims[2 + i], strides[i], kernel_shape[i], (dilations[i] + 1), auto_pad, pad_head, pad_tail)) {
pads[i]= pad_head;
pads[i] = pad_head;
pads[shape + i] = pad_tail;
}
}
@ -304,8 +319,7 @@ dnnl::memory::dims DnnlConv::InferOutputShape(DnnlNode& node,
const std::vector<int64_t>& kernel_shape,
const dnnl::memory::dims& strides,
const dnnl::memory::dims& dilations,
const std::vector<int64_t>& pads)
{
const std::vector<int64_t>& pads) {
auto pad_type = GetAutoPad(node);
ConvShape shape = static_cast<ConvShape>(kernel_shape.size());
dnnl::memory::dims output_shape;
@ -323,12 +337,10 @@ dnnl::memory::dims DnnlConv::InferOutputShape(DnnlNode& node,
switch (pad_type) {
case onnxruntime::AutoPadType::NOTSET: {
output_shape.push_back(static_cast<int64_t>(static_cast<float>(src_dims[dim + 2] + pads[dim] + pads[dim + shape] - dkernel) / strides[dim] + 1));
}
break;
} break;
case onnxruntime::AutoPadType::VALID: {
output_shape.push_back((src_dims[dim + 2] - dkernel) / strides[dim] + 1);
}
break;
} break;
case onnxruntime::AutoPadType::SAME_UPPER: {
if (dilations[dim] != 0) {
LOGS_DEFAULT(ERROR) << "Dilation not supported for AutoPadType::SAME_UPPER or AutoPadType::SAME_LOWER.";
@ -337,14 +349,12 @@ dnnl::memory::dims DnnlConv::InferOutputShape(DnnlNode& node,
int64_t legacy_target_size = (src_dims[dim + 2] + strides[dim] - 1) / strides[dim];
int64_t pad_needed = (legacy_target_size - 1) * strides[dim] + kernel_shape[dim] - src_dims[dim + 2];
output_shape.push_back((src_dims[dim + 2] + pad_needed - dkernel) / strides[dim] + 1);
}
break;
} break;
case onnxruntime::AutoPadType::SAME_LOWER: {
int64_t legacy_target_size = (src_dims[dim + 2] + strides[dim] - 1) / strides[dim];
int64_t pad_needed = (legacy_target_size - 1) * strides[dim] + kernel_shape[dim] - src_dims[dim + 2];
output_shape.push_back((src_dims[dim + 2] + pad_needed - dkernel) / strides[dim] + 1);
}
break;
} break;
default:
break;
}

96
onnxruntime/core/providers/dnnl/subgraph/dnnl_gelu.cc Normal file
View file

@ -0,0 +1,96 @@
// Copyright(C) 2021 Intel Corporation
// Licensed under the MIT License
#include "dnnl_gelu.h"
#include "dnnl_subgraph.h"
#include "dnnl_subgraph_primitive.h"
namespace onnxruntime {
namespace ort_dnnl {
DnnlGelu::DnnlGelu() {}
void DnnlGelu::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) {
auto dnnl_engine = sp.GetEngine();
bool is_biased = node.Input(IN_BIAS).Exists();
dnnl::memory src_mem;
if (is_biased) {
src_mem = sp.GetMemoryInOrtFormat(node.Input(IN_X), dnnl_engine);
} else {
src_mem = sp.GetMemory(node.Input(IN_X));
}
auto gelu_src_mem = src_mem;
if (is_biased) {
auto bias_mem = sp.GetMemoryInOrtFormat(node.Input(IN_BIAS), dnnl_engine);
auto src0_ori_md = src_mem.get_desc();
auto src1_ori_md = bias_mem.get_desc();
auto src0_dims = src0_ori_md.dims();
auto src1_dims = src1_ori_md.dims();
if (src0_dims.size() != src1_dims.size()) {
while (src0_dims.size() < src1_dims.size()) {
src0_dims.insert(src0_dims.begin(), 1);
}
while (src0_dims.size() > src1_dims.size()) {
src1_dims.insert(src1_dims.begin(), 1);
}
}
auto src0_md = src0_ori_md.reshape(src0_dims);
auto src1_md = src1_ori_md.reshape(src1_dims);
auto output_shape = src0_dims;
for (size_t i = 0; i < output_shape.size(); i++) {
if (output_shape[i] == 1) {
output_shape[i] = src1_dims[i];
}
}
auto dst_md = dnnl::memory::desc(output_shape, node.Output(OUT_Y).Type(), dnnl::memory::format_tag::any);
auto binary_d = dnnl::binary::desc(dnnl::algorithm::binary_add, src0_md, src1_md, dst_md);
auto binary_pd = dnnl::binary::primitive_desc(binary_d, dnnl_engine);
auto binary_src0_mem = sp.GetMemoryAndReshape(node.Input(IN_X), binary_pd.src0_desc(), dnnl_engine);
auto binary_src1_mem = sp.GetMemoryAndReshape(node.Input(IN_BIAS), binary_pd.src1_desc(), dnnl_engine);
auto binary_dst_mem = dnnl::memory(binary_pd.dst_desc(), dnnl_engine);
auto binary_prim = dnnl::binary(binary_pd);
sp.AddPrimitive(binary_prim, {{DNNL_ARG_SRC_0, binary_src0_mem},
{DNNL_ARG_SRC_1, binary_src1_mem},
{DNNL_ARG_DST, binary_dst_mem}});
gelu_src_mem = binary_dst_mem;
}
dnnl::algorithm algo;
if (node.OpType() == "Gelu" || node.OpType() == "BiasGelu") {
algo = dnnl::algorithm::eltwise_gelu_erf;
} else if (node.OpType() == "FastGelu") {
algo = dnnl::algorithm::eltwise_gelu_tanh;
} else {
ORT_THROW("op type not supported");
}
auto gelu_desc = dnnl::eltwise_forward::desc(dnnl::prop_kind::forward_inference, algo, gelu_src_mem.get_desc());
auto gelu_pd = dnnl::eltwise_forward::primitive_desc(gelu_desc, dnnl_engine);
if(!is_biased) {
// If using GPU this will move the memory from the CPU to the GPU.
gelu_src_mem = sp.GetMemoryAndReshape(node.Input(IN_X), gelu_pd.src_desc(), dnnl_engine);
}
auto dst_mem = dnnl::memory(gelu_pd.dst_desc(), dnnl_engine);
auto gelu_op = dnnl::eltwise_forward(gelu_pd);
sp.AddPrimitive(gelu_op, {{DNNL_ARG_SRC, gelu_src_mem},
{DNNL_ARG_DST, dst_mem}});
sp.SetMemory(node.Output(OUT_Y), dst_mem);
}
} // namespace ort_dnnl
} // namespace onnxruntime

27
onnxruntime/core/providers/dnnl/subgraph/dnnl_gelu.h Normal file
View file

@ -0,0 +1,27 @@
// Copyright(C) 2021 Intel Corporation
// Licensed under the MIT License
#pragma once
#include "dnnl_subgraph.h"
#include "dnnl_subgraph_primitive.h"
namespace onnxruntime {
namespace ort_dnnl {
class DnnlGelu {
public:
enum InputTensors : int {
IN_X = 0,
IN_BIAS = 1
};
enum OutputTensors : int {
OUT_Y = 0
};
DnnlGelu();
void CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node);
};
} // namespace ort_dnnl
} // namespace onnxruntime

60
onnxruntime/core/providers/dnnl/subgraph/dnnl_matmul.cc
View file

@ -13,6 +13,13 @@ DnnlMatMul::DnnlMatMul() {}
void DnnlMatMul::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) {
auto eng = sp.GetEngine();
bool has_add = false;
if (node.OpType() == "MatMulAdd") {
has_add = true;
//if fused with add, need a third input
assert(node.Input(IN_BINARY).Exists());
}
auto src_dims = sp.GetMemory(node.Input(IN_A)).get_desc().dims();
auto weights_dims = sp.GetMemory(node.Input(IN_B)).get_desc().dims();
@ -37,20 +44,63 @@ void DnnlMatMul::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) {
}
}
/*
create a post op binary with possible unsqueezing in order to make sure onednn properly broadcast
current limitation
1. is no unsqueeze for matmul output as it is not exposed due to post op fusion
2. the third input has to be reordered to plain format (eg, no memory format propogation if the third input is internal to subgraph)
3. adding 1s to front (unsqueeze/expand) in logical dims would possibly fail if physcial layout is not plain format
*/
dnnl::primitive_attr attr;
if (has_add) {
dnnl::post_ops ops;
auto ori_binary_mem_desc = sp.GetMemory(node.Input(IN_BINARY).Name()).get_desc();
auto ori_binary_mem_dims = ori_binary_mem_desc.dims();
auto binary_mem_dims = ori_binary_mem_dims;
if (ori_binary_mem_dims.size() != output_shape.size()) {
if (ori_binary_mem_dims.size() > output_shape.size()) {
ORT_THROW("add fusion with matmul output broadcasting by unsqueezing is not supported");
}
//expand the third input (from the binary op) is possible
while (binary_mem_dims.size() < output_shape.size()) {
binary_mem_dims.insert(binary_mem_dims.begin(), 1);
}
}
//expand the dims by 1s (should always be possible)
//will throw exception if not possible
auto binary_mem_desc = ori_binary_mem_desc.reshape(binary_mem_dims);
//TODO: use format any to choose the best layout
ops.append_binary(dnnl::algorithm::binary_add, binary_mem_desc);
attr.set_post_ops(ops);
}
auto dst_md = dnnl::memory::desc(output_shape, node.Output(OUT_Y).Type(), dnnl::memory::format_tag::any);
auto matmul_d = dnnl::matmul::desc(src_md, weights_md, dst_md);
auto matmul_pd = dnnl::matmul::primitive_desc(matmul_d, eng);
auto matmul_pd = dnnl::matmul::primitive_desc(matmul_d, attr, eng);
auto matmul_src_mem = sp.GetMemoryAndReshape(node.Input(IN_A), matmul_pd.src_desc(), eng);
auto matmul_weights_mem = sp.GetMemoryAndReshape(node.Input(IN_B), matmul_pd.weights_desc(), eng);
auto matmul_dst_mem = dnnl::memory(matmul_pd.dst_desc(), eng);
auto matmul_prim = dnnl::matmul(matmul_pd);
sp.AddPrimitive(matmul_prim, {{DNNL_ARG_SRC, matmul_src_mem},
{DNNL_ARG_WEIGHTS, matmul_weights_mem},
{DNNL_ARG_DST, matmul_dst_mem}});
//a default memory map for matmul
std::unordered_map<int, dnnl::memory> mem_map({{DNNL_ARG_SRC, matmul_src_mem},
{DNNL_ARG_WEIGHTS, matmul_weights_mem},
{DNNL_ARG_DST, matmul_dst_mem}});
//add to memory map with extra third input if fused with add
if (has_add) {
dnnl::algorithm algo;
dnnl::memory::desc binary_mem_desc;
matmul_pd.get_primitive_attr().get_post_ops().get_params_binary(0, algo, binary_mem_desc);
assert(algo == dnnl::algorithm::binary_add);
auto binary_post_op_mem = sp.GetMemoryAndReshape(node.Input(IN_BINARY), binary_mem_desc, eng);
mem_map[DNNL_ARG_ATTR_MULTIPLE_POST_OP(0) | DNNL_ARG_SRC_1] = binary_post_op_mem;
}
sp.AddPrimitive(matmul_prim, mem_map);
sp.SetMemory(node.Output(OUT_Y), matmul_dst_mem);
}

3
onnxruntime/core/providers/dnnl/subgraph/dnnl_matmul.h
View file

@ -12,7 +12,8 @@ class DnnlMatMul {
public:
enum InputTensors : int {
IN_A = 0,
IN_B = 1
IN_B = 1,
IN_BINARY = 2 // the extra input due to matmulbinary fusion
};
enum OutputTensors : int {

256
onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph.cc
View file

@ -2,24 +2,35 @@
// Licensed under the MIT License
#include "dnnl_subgraph.h"
#include <queue>
namespace onnxruntime {
namespace ort_dnnl {
DnnlTensor DnnlNode::empty_tensor_ = DnnlTensor("");
DnnlTensor::DnnlTensor(const NodeArg* arg) {
arg_ = arg;
tensor_name_ = arg->Name();
if (!arg || !arg->Exists()) {
tensor_name_ = "";
} else {
tensor_name_ = arg->Name();
}
}
DnnlTensor::DnnlTensor(std::string name) {
tensor_name_ = name;
arg_ = nullptr;
}
std::string DnnlTensor::Name() const {
return tensor_name_;
}
dnnl::memory::dims DnnlTensor::Dim() {
dnnl::memory::dims DnnlTensor::Dim() const {
if (arg_ == nullptr) {
return dnnl::memory::dims();
}
auto shape_proto = arg_->Shape();
// a shape without any information
if (shape_proto == nullptr) {
@ -95,54 +106,80 @@ dnnl::memory::format_tag DnnlTensor::Format() {
return dnnl::memory::format_tag::any;
}
void DnnlTensor::SetProducer(const DnnlNodeArg& arg) {
producer_ = arg;
}
void DnnlTensor::ResetProducer() {
producer_ = DnnlNodeArg();
}
void DnnlTensor::AddConsumer(const DnnlNodeArg& arg) {
consumers_.push_back(arg);
}
void DnnlTensor::RemoveConsumer(const DnnlNodeArg& arg) {
consumers_.erase(std::remove(consumers_.begin(), consumers_.end(), arg), consumers_.end());
}
DnnlNode::DnnlNode(const Node* node) {
onnx_node_ = node;
name_ = node->Name();
op_type_ = node->OpType();
attr_->insert(node->GetAttributes());
}
std::string DnnlNode::Name() {
return onnx_node_->Name();
std::string& DnnlNode::Name() {
return name_;
}
std::string DnnlNode::OpType() {
return onnx_node_->OpType();
std::string& DnnlNode::OpType() {
return op_type_;
}
DnnlTensor DnnlNode::Input(int index) {
if (onnx_node_->InputDefs().size() <= (size_t)index) {
return DnnlTensor("");
std::vector<DnnlTensor*>& DnnlNode::Inputs() {
return inputs_;
}
std::vector<DnnlTensor*>& DnnlNode::Outputs() {
return outputs_;
}
size_t& DnnlNode::Index() {
return index_;
}
DnnlTensor& DnnlNode::Input(int index) {
if (inputs_.size() <= (size_t)index) {
return empty_tensor_;
}
if (!onnx_node_->InputDefs()[index]) {
return DnnlTensor("");
if (inputs_[index] && inputs_[index]->Exists()) {
return *inputs_[index];
}
if (onnx_node_->InputDefs()[index]->Exists()) {
auto def = onnx_node_->InputDefs()[index];
return DnnlTensor(def);
}
return DnnlTensor("");
return empty_tensor_;
}
size_t DnnlNode::InputCount() {
return onnx_node_->InputDefs().size();
return inputs_.size();
}
DnnlTensor DnnlNode::Output(int index) {
auto def = onnx_node_->OutputDefs()[index];
return DnnlTensor(def);
DnnlTensor& DnnlNode::Output(int index) {
return *outputs_[index];
}
size_t DnnlNode::OutputCount() {
return onnx_node_->OutputDefs().size();
return outputs_.size();
}
const NodeAttributes& DnnlNode::Attributes() {
return onnx_node_->GetAttributes();
NodeAttributes& DnnlNode::Attributes() {
return *attr_;
}
DnnlSubgraph::DnnlSubgraph(const GraphViewer& graph_viewer) : graph_viewer_(graph_viewer) {
Build();
is_dynamic_ = false;
for (auto& input : GetDnnlInputs()) {
if (input.IsDynamic()) {
for (auto input : GetDnnlInputs()) {
if (input->IsDynamic()) {
is_dynamic_ = true;
break;
}
@ -153,37 +190,184 @@ bool DnnlSubgraph::IsDynamic() {
return is_dynamic_;
}
std::vector<DnnlNode> DnnlSubgraph::GetDnnlNodes() {
return dnnl_nodes_;
void DnnlSubgraph::TopoSort() {
nodes_in_topological_order_.clear();
std::unordered_map<size_t, int> indegrees;
for (auto& node : dnnl_nodes_) {
if (node.get()) {
indegrees[node->Index()] = 0;
}
}
for (auto& e : dnnl_tensors_) {
auto tensor = e.second.get();
if (tensor->Exists() && tensor->GetProducer().GetNode()) {
for (auto edge : tensor->GetConsumers()) {
if (edge.GetNode()) {
indegrees[edge.GetNode()->Index()]++;
}
}
}
}
std::queue<DnnlNode*> queue;
for (auto e : indegrees) {
if (e.second == 0) {
queue.push(dnnl_nodes_[e.first].get());
}
}
//need to make sure all indegrees are computed before doing bfs
while (!queue.empty()) {
auto cur = queue.front();
queue.pop();
nodes_in_topological_order_.push_back(cur->Index());
for (auto output : cur->Outputs()) {
if (output && output->Exists()) {
for (auto edge : output->GetConsumers()) {
indegrees[edge.GetNode()->Index()] -= 1;
if (indegrees[edge.GetNode()->Index()] == 0) {
queue.push(edge.GetNode());
}
}
}
}
}
assert(indegrees.size() == nodes_in_topological_order_.size());
}
std::vector<DnnlTensor> DnnlSubgraph::GetDnnlInputs() {
std::vector<size_t> DnnlSubgraph::GetDnnlNodesInTopologicalOrder() {
TopoSort();
return nodes_in_topological_order_;
}
DnnlNode* DnnlSubgraph::GetDnnlNode(size_t node_index) {
return dnnl_nodes_[node_index].get();
}
DnnlTensor* DnnlSubgraph::GetDnnlTensor(const std::string& tensor_name){
if(dnnl_tensors_.count(tensor_name)){
return dnnl_tensors_[tensor_name].get();
}else{
return nullptr;
}
}
size_t DnnlSubgraph::GetMaxNodeIndex() {
return dnnl_nodes_.size();
}
std::vector<DnnlNode*> DnnlSubgraph::GetDnnlNodes() {
std::vector<DnnlNode*> result;
for (auto& node : dnnl_nodes_) {
if (node.get()) {
result.push_back(node.get());
}
}
return result;
}
std::vector<DnnlTensor*> DnnlSubgraph::GetDnnlInputs() {
return inputs_;
}
std::vector<DnnlTensor> DnnlSubgraph::GetDnnlOutputs() {
std::vector<DnnlTensor*> DnnlSubgraph::GetDnnlOutputs() {
return outputs_;
}
std::vector<DnnlTensor> DnnlSubgraph::GetDnnlInitializers() {
std::vector<DnnlTensor*> DnnlSubgraph::GetDnnlInitializers() {
return initializers_;
}
void DnnlSubgraph::Build() {
for (const auto* node_arg : graph_viewer_.GetInputsIncludingInitializers()) {
inputs_.push_back(DnnlTensor(node_arg));
void DnnlSubgraph::RemoveNode(size_t node_index) {
dnnl_nodes_[node_index].reset(nullptr);
}
void DnnlSubgraph::RemoveTensor(const std::string& tensor_name) {
inputs_.erase(std::remove_if(inputs_.begin(), inputs_.end(), [=](DnnlTensor* t){ return t->Name() == tensor_name;}), inputs_.end());
initializers_.erase(std::remove_if(initializers_.begin(), initializers_.end(), [=](DnnlTensor* t){ return t->Name() == tensor_name;}), initializers_.end());
dnnl_tensors_.erase(tensor_name);
}
void DnnlSubgraph::AddTensor(std::unique_ptr<DnnlTensor> new_tensor) {
if (!dnnl_tensors_.count(new_tensor->Name())) {
dnnl_tensors_.emplace(new_tensor->Name(), std::move(new_tensor));
} else {
ORT_THROW("tensor exists, modify or delete first before inseting");
}
}
void DnnlSubgraph::AddNode(std::unique_ptr<DnnlNode> new_node) {
auto index = dnnl_nodes_.size();
dnnl_nodes_.emplace_back(std::move(new_node));
dnnl_nodes_.back()->Index() = index;
}
bool DnnlSubgraph::GetInitializedTensor(const std::string& arg_name, const ONNX_NAMESPACE::TensorProto*& value) {
return graph_viewer_.GetInitializedTensor(arg_name, value);
}
bool DnnlSubgraph::IsConstantInitializer(const std::string& arg_name, bool check_outer_scope) {
return graph_viewer_.IsConstantInitializer(arg_name, check_outer_scope);
}
void DnnlSubgraph::Build() {
//establish nodes, tensors and nodeargs
const auto& node_indices = graph_viewer_.GetNodesInTopologicalOrder();
for (size_t i = 0; i < node_indices.size(); i++) {
const auto* node(graph_viewer_.GetNode(node_indices[i]));
dnnl_nodes_.push_back(DnnlNode(node));
AddNode(std::make_unique<DnnlNode>(node));
auto dnnl_node = dnnl_nodes_.back().get();
std::vector<DnnlTensor*> inputs;
size_t index = 0;
for (auto input : node->InputDefs()) {
if (input && input->Exists() && input->Name() != "") {
if (!dnnl_tensors_.count(input->Name())) {
dnnl_tensors_[input->Name()] = std::make_unique<DnnlTensor>(input);
}
dnnl_tensors_[input->Name()]->AddConsumer(DnnlNodeArg(dnnl_node, index, false));
inputs.push_back(dnnl_tensors_[input->Name()].get());
} else {
inputs.push_back(nullptr);
}
index++;
}
std::vector<DnnlTensor*> outputs;
index = 0;
for (auto output : node->OutputDefs()) {
if (output && output->Exists() && output->Name() != "") {
if (!dnnl_tensors_.count(output->Name())) {
dnnl_tensors_[output->Name()] = std::make_unique<DnnlTensor>(output);
}
dnnl_tensors_[output->Name()]->SetProducer(DnnlNodeArg(dnnl_node, index, true));
outputs.push_back(dnnl_tensors_[output->Name()].get());
} else {
outputs.push_back(nullptr);
}
index++;
}
dnnl_node->Inputs() = inputs;
dnnl_node->Outputs() = outputs;
}
//all tensors should have been established in graph
//establish inputs, outputs and initializers
//graph inputs including initializers and outputs can be deleted by graph transformation (eg, gelu fusion)
//delete unneeded inputs don't affect onnxruntime passing them as input data handle
//delete unneeded outputs will cause ep to output to fewer data handles then expected
for (const auto* node_arg : graph_viewer_.GetInputsIncludingInitializers()) {
inputs_.push_back(dnnl_tensors_[node_arg->Name()].get());
}
for (const auto* node_arg : graph_viewer_.GetOutputs()) {
outputs_.push_back(DnnlTensor(node_arg));
outputs_.push_back(dnnl_tensors_[node_arg->Name()].get());
}
for (auto& initializer : graph_viewer_.GetAllInitializedTensors()) {
auto& name = initializer.first;
initializers_.push_back(DnnlTensor(name));
initializers_.push_back(dnnl_tensors_[name].get());
}
}

92
onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph.h
View file

@ -6,65 +6,121 @@
#include <vector>
#include <string>
#include <map>
#include <limits>
#include "dnnl.hpp"
#include "core/providers/shared_library/provider_api.h"
namespace onnxruntime {
namespace ort_dnnl {
class DnnlNode;
class DnnlNodeArg {
public:
DnnlNodeArg(DnnlNode* node, size_t index, bool is_output)
: node_(node), index_(index), is_output_(is_output){};
DnnlNodeArg() = default;
DnnlNode* GetNode() { return node_; };
size_t GetIndex() { return index_; };
bool IsOutput() { return is_output_; };
bool Exists() { return node_ != nullptr; };
bool operator==(const DnnlNodeArg& e) const {
return node_ == e.node_ && index_ == e.index_ && is_output_ == e.is_output_;
};
private:
DnnlNode* node_ = nullptr;
size_t index_ = std::numeric_limits<size_t>::max();
bool is_output_ = false;
};
class DnnlTensor {
public:
DnnlTensor(const NodeArg* arg);
DnnlTensor(std::string name);
DnnlTensor() = default;
std::string Name() const;
dnnl::memory::dims Dim();
dnnl::memory::dims Dim() const;
dnnl::memory::data_type Type() const;
dnnl::memory::format_tag Format();
//check whether the tensor is dynamic, e.g. contains unspecified dimension
bool IsDynamic();
//check whether the tensor exsits for optional input output
bool Exists();
std::vector<DnnlNodeArg>& GetConsumers() { return consumers_; };
DnnlNodeArg& GetProducer() { return producer_; };
void SetProducer(const DnnlNodeArg& arg);
void ResetProducer();
void AddConsumer(const DnnlNodeArg& arg);
void RemoveConsumer(const DnnlNodeArg& arg);
private:
std::string tensor_name_;
const NodeArg* arg_;
bool is_dynamic_;
//a tensor can have no producer (input.initializer) or no consumer (output for subgraph)
DnnlNodeArg producer_;
std::vector<DnnlNodeArg> consumers_;
};
class DnnlNode {
public:
DnnlNode(const Node* node);
std::string Name();
std::string OpType();
DnnlTensor Input(int index);
DnnlNode() = default;
std::string& Name();
size_t& Index();
std::string& OpType();
DnnlTensor& Input(int index);
size_t InputCount();
DnnlTensor Output(int index);
DnnlTensor& Output(int index);
size_t OutputCount();
const NodeAttributes& Attributes();
NodeAttributes& Attributes();
std::vector<DnnlTensor*>& Inputs();
std::vector<DnnlTensor*>& Outputs();
private:
const Node* onnx_node_;
std::vector<DnnlTensor> inputs_;
std::vector<DnnlTensor> outputs_;
const Node* onnx_node_ = nullptr;
std::vector<DnnlTensor*> inputs_;
std::vector<DnnlTensor*> outputs_;
static DnnlTensor empty_tensor_;
std::string name_; // node can have empty/duplicate name, rely on index instead
std::string op_type_;
size_t index_ = std::numeric_limits<size_t>::max();
std::unique_ptr<NodeAttributes> attr_ = NodeAttributes::Create();
};
class DnnlSubgraph {
public:
DnnlSubgraph(const GraphViewer& graph_viewer);
std::vector<DnnlNode> GetDnnlNodes();
std::vector<DnnlTensor> GetDnnlInputs();
std::vector<DnnlTensor> GetDnnlOutputs();
std::vector<DnnlTensor> GetDnnlInitializers();
std::vector<DnnlNode*> GetDnnlNodes();
DnnlNode* GetDnnlNode(size_t node_index);
DnnlTensor* GetDnnlTensor(const std::string& tensor_name);
size_t GetMaxNodeIndex();
std::vector<size_t> GetDnnlNodesInTopologicalOrder();
std::vector<DnnlTensor*> GetDnnlInputs();
std::vector<DnnlTensor*> GetDnnlOutputs();
std::vector<DnnlTensor*> GetDnnlInitializers();
// build the subgraph IR
void Build();
//check whether the subgraph is dynamic
void TopoSort();
bool IsDynamic();
void AddNode(std::unique_ptr<DnnlNode> new_node);
void RemoveNode(size_t node_index);
void AddTensor(std::unique_ptr<DnnlTensor> new_tensor);
void RemoveTensor(const std::string& tensor_name);
bool GetInitializedTensor(const std::string& arg_name, const ONNX_NAMESPACE::TensorProto*& value);
bool IsConstantInitializer(const std::string& arg_name, bool check_outer_scope);
private:
std::vector<DnnlNode> dnnl_nodes_;
std::vector<DnnlTensor> inputs_;
std::vector<DnnlTensor> outputs_;
std::vector<DnnlTensor> initializers_;
//graph owns all nodes
std::vector<std::unique_ptr<DnnlNode>> dnnl_nodes_;
std::vector<size_t> nodes_in_topological_order_;
//graph owns all tensors
std::unordered_map<std::string, std::unique_ptr<DnnlTensor>> dnnl_tensors_;
std::vector<DnnlTensor*> inputs_;
std::vector<DnnlTensor*> outputs_; //output should never get deleted from graph transformation
std::vector<DnnlTensor*> initializers_;
const GraphViewer& graph_viewer_;
bool is_dynamic_;
};

30
onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc
View file

@ -8,6 +8,7 @@
#include "dnnl_conv.h"
#include "dnnl_dynamicquantizelinear.h"
#include "dnnl_elementwise.h"
#include "dnnl_gelu.h"
#include "dnnl_gemm.h"
#include "dnnl_lrn.h"
#include "dnnl_matmul.h"
@ -48,22 +49,29 @@ void DnnlSubgraphPrimitive::AddKernels() {
std::unordered_set<std::string> binary_ops = {"Add", "Div", "Mul", "Sub"};
std::unordered_set<std::string> elementwise_ops = {"Abs", "Elu", "Exp", "LeakyRelu", "Log", "Relu", "Round", "Sigmoid", "Softplus", "Sqrt", "Tanh"};
std::unordered_set<std::string> pool_ops = {"AveragePool", "GlobalAveragePool", "GlobalMaxPool", "MaxPool"};
for (auto& node : subgraph_->GetDnnlNodes()) {
auto indices = subgraph_->GetDnnlNodesInTopologicalOrder();
for (auto index : indices) {
auto& node = *(subgraph_->GetDnnlNode(index));
if (node.OpType() == "BatchNormalization") {
DnnlBatchNorm().CreatePrimitive(*this, node);
} else if (binary_ops.count(node.OpType())) {
DnnlBinary().CreatePrimitive(*this, node);
} else if (node.OpType() == "Conv") {
} else if (node.OpType() == "Conv" || node.OpType() == "ConvRelu") {
DnnlConv().CreatePrimitive(*this, node);
} else if (node.OpType() == "DynamicQuantizeLinear") {
DnnlDynamicQuantizeLinear().CreatePrimitive(*this, node);
} else if (elementwise_ops.count(node.OpType())) {
DnnlElementwise().CreatePrimitive(*this, node);
} else if (node.OpType() == "FastGelu"){
DnnlGelu().CreatePrimitive(*this, node);
} else if (node.OpType() == "Gelu" || node.OpType() == "BiasGelu") {
DnnlGelu().CreatePrimitive(*this, node);
} else if (node.OpType() == "Gemm") {
DnnlGemm().CreatePrimitive(*this, node);
} else if (node.OpType() == "LRN") {
DnnlLrn().CreatePrimitive(*this, node);
} else if (node.OpType() == "MatMul") {
} else if (node.OpType() == "MatMul" || node.OpType() == "MatMulAdd") {
DnnlMatMul().CreatePrimitive(*this, node);
} else if (node.OpType() == "MatMulInteger") {
DnnlMatMulInteger().CreatePrimitive(*this, node);
@ -161,9 +169,9 @@ void DnnlSubgraphPrimitive::Compile(const std::unordered_map<std::string, OnnxTe
//initializer should not be cleared upon recompile
//initializers_.clear();
for (auto& nodearg : subgraph_->GetDnnlInputs()) {
auto dnnl_tensor_name = nodearg.Name();
auto dnnl_data_type = nodearg.Type();
for (auto nodearg : subgraph_->GetDnnlInputs()) {
auto dnnl_tensor_name = nodearg->Name();
auto dnnl_data_type = nodearg->Type();
dnnl::memory::dims dnnl_dims = inputs.at(dnnl_tensor_name).tensor_info.shape;
if (dnnl_dims.size() == 0) {
dnnl_dims.push_back(1);
@ -257,8 +265,8 @@ dnnl::stream DnnlSubgraphPrimitive::GetStream() {
}
void DnnlSubgraphPrimitive::AddInitializers() {
for (auto& nodearg : subgraph_->GetDnnlInitializers()) {
auto dnnl_tensor_name = nodearg.Name();
for (auto nodearg : subgraph_->GetDnnlInitializers()) {
auto dnnl_tensor_name = nodearg->Name();
if (!Contains(initializers_, dnnl_tensor_name)) {
initializers_.insert(std::pair<std::string, std::vector<dnnl::memory> >(dnnl_tensor_name, std::vector<dnnl::memory>()));
}
@ -266,9 +274,9 @@ void DnnlSubgraphPrimitive::AddInitializers() {
}
void DnnlSubgraphPrimitive::AddOutputs() {
for (auto& tensor : subgraph_->GetDnnlOutputs()) {
auto dnnl_data_type = tensor.Type();
auto dnnl_tensor_name = tensor.Name();
for (auto tensor : subgraph_->GetDnnlOutputs()) {
auto dnnl_data_type = tensor->Type();
auto dnnl_tensor_name = tensor->Name();
auto engine = GetCPUEngine();
auto output_mem_dnnl = GetMemory(dnnl_tensor_name);
auto output_md = dnnl::memory::desc(output_mem_dnnl.get_desc().dims(), dnnl_data_type, GetDnnlFormat(output_mem_dnnl.get_desc().dims().size()));

750
onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_transformer.cc Normal file
View file

@ -0,0 +1,750 @@
// Copyright(C) 2021 Intel Corporation
// Licensed under the MIT License
#include "dnnl_subgraph_transformer.h"
#ifndef _USE_MATH_DEFINES
#define _USE_MATH_DEFINES
#endif
#include <cmath>
#include <iostream>
namespace onnxruntime {
namespace ort_dnnl {
//apply all transformation rules in order
void DnnlGraphTransformer::Apply(DnnlSubgraph& subgraph) {
ConvRelu(subgraph);
MatMulAdd(subgraph);
Gelu(subgraph);
FastGelu(subgraph);
}
//resolve a fusion by replacing old_indices nodes with a new_node
//unneeded tensors will be deleted, old news' edges will be cleared
//new_node will be set with new edges and inserted to subgraph
void DnnlGraphTransformer::ResolveFusion(DnnlSubgraph& subgraph, std::vector<size_t> old_indices, std::unique_ptr<DnnlNode> new_node) {
//the tensors to keep
std::unordered_set<std::string> keep_tensors;
//get keep tensors from new_node
//all tensors related to new_node needs to be kept
for (auto input : new_node->Inputs()) {
if (input && input->Exists()) {
keep_tensors.insert(input->Name());
}
}
for (auto output : new_node->Outputs()) {
if (output && output->Exists()) {
keep_tensors.insert(output->Name());
}
}
//find out tensors to remove, cleanup tensor consumers and producer
std::unordered_set<std::string> tensors_to_remove;
for (auto index : old_indices) {
auto cur_node = subgraph.GetDnnlNode(index);
{
int input_index = 0;
for (auto input : cur_node->Inputs()) {
if (input && input->Exists()) {
input->RemoveConsumer(DnnlNodeArg(cur_node, input_index, false));
if (!keep_tensors.count(input->Name())) {
tensors_to_remove.insert(input->Name());
}
}
input_index++;
}
}
for (auto output : cur_node->Outputs()) {
if (output && output->Exists()) {
output->ResetProducer();
if (!keep_tensors.count(output->Name())) {
tensors_to_remove.insert(output->Name());
}
}
}
}
//remove unused tensors
for (const auto& tensor_name : tensors_to_remove) {
auto tensor = subgraph.GetDnnlTensor(tensor_name);
if(tensor){
//has consumer and producer
if(tensor->GetConsumers().size() || tensor->GetProducer().Exists()){
continue;
}
else{
subgraph.RemoveTensor(tensor_name);
}
}
//subgraph.RemoveTensor(tensor_name);
}
//remove unused nodes
for (auto index : old_indices) {
subgraph.RemoveNode(index);
}
//reestablish producer and consumer for tensors related to new node
//such tensors should not get deleted
{
size_t input_index = 0;
for (auto input : new_node->Inputs()) {
if (input) {
input->AddConsumer(DnnlNodeArg(new_node.get(), input_index, false));
}
input_index++;
}
size_t output_index = 0;
for (auto output : new_node->Outputs()) {
if (output) {
output->SetProducer(DnnlNodeArg(new_node.get(), output_index, true));
}
output_index++;
}
}
//new node now has correct input output tensors as well as tensor connections
//subgraph now owns the new node
subgraph.AddNode(std::move(new_node));
}
//helper to determine whether a tensor acts as subgraph output
bool DnnlGraphTransformer::IsGraphOutput(DnnlSubgraph& subgraph, DnnlTensor& tensor) {
auto graph_outputs = subgraph.GetDnnlOutputs();
if (std::find(graph_outputs.cbegin(), graph_outputs.cend(), &tensor) != graph_outputs.cend()) {
return true;
}
return false;
}
//helper to determien whether
bool DnnlGraphTransformer::ProduceGraphOutput(DnnlSubgraph& subgraph, DnnlNode& node) {
auto graph_outputs = subgraph.GetDnnlOutputs();
for (auto output : node.Outputs()) {
if (output && output->Exists()) {
if (IsGraphOutput(subgraph, *output)) {
return true;
}
}
}
return false;
}
bool DnnlGraphTransformer::IsNodeFusable(DnnlSubgraph& subgraph, DnnlNode* node) const{
if (node == nullptr) {
return false;
}
//isSingleOutput(DnnlNode* node);
if (node->OutputCount() != 1) {
std::string s = "Invalid " + node->OpType() + " node";
ORT_THROW(s);
}
//isConsumedBySingleNode(DnnlNode* node);
if (node->Output(0).Exists() && node->Output(0).GetConsumers().size() != 1) {
return false;
}
//isOutputPartOfSubgraph(DnnlSubgraph& subgraph, DnnlNode* node);
auto graph_outputs = subgraph.GetDnnlOutputs();
if (std::find(graph_outputs.cbegin(), graph_outputs.cend(), &node->Output(0)) != graph_outputs.cend()) {
return false;
}
return true;
}
bool IsScalar(const DnnlTensor& input_arg) {
auto dim = input_arg.Dim();
auto dim_size = dim.size();
return dim_size == 0 || (dim_size == 1 && dim[0] == 1);
}
bool DnnlGraphTransformer::IsInitilizedWithExpectedValue(DnnlSubgraph& subgraph, DnnlTensor& input_arg, float expected_value) {
if (!IsScalar(input_arg)) {
return false;
}
const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr;
if (!subgraph.GetInitializedTensor(input_arg.Name(), tensor_proto)) {
return false;
}
if (tensor_proto == nullptr) {
return false;
}
if (!tensor_proto->has_raw_data()) {
return false;
}
const auto data_type = input_arg.Type();
if (data_type == dnnl::memory::data_type::f32) {
const float* val = reinterpret_cast<const float*>(tensor_proto->raw_data().data());
if (std::isnan(val[0]) || std::isinf(val[0])) {
if (std::isinf(val[0]) && std::isinf(expected_value) && (std::signbit(val[0]) == std::signbit(expected_value))) {
return true;
}
return false;
}
const float atol = 1e-8f;
const float rtol = 1e-5f;
float diff = std::abs(val[0] - expected_value);
if (diff > (atol + rtol * std::abs(expected_value))) {
return false;
}
} else {
// Not expected data types.
return false;
}
return true;
}
DnnlNode* FirstParentByType(DnnlNode* node, const std::string& parent_type) {
for (size_t i = 0; i < node->InputCount(); ++i) {
auto prev_node = node->Input(static_cast<int>(i)).GetProducer().GetNode();
if (prev_node != nullptr && prev_node->OpType() == parent_type) {
return prev_node;
}
}
return nullptr;
}
/*
This function fuses subgraph like the following into one Gelu node.
Subgraph pattern 1:
+-------Mul(0.5)---------------------+
| |
| v
[root] --> Div -----> Erf --> Add --> Mul ==>
(B=1.4142...) (1)
Subgraph pattern 2:
+------------------------------------+
| |
| v
[root] --> Div -----> Erf --> Add --> Mul -->Mul ==>
(B=1.4142...) (1) (0.5)
After Fusion:
[root]--> Gelu ==>
*/
void DnnlGraphTransformer::Gelu(DnnlSubgraph& subgraph) {
static int gelu_index = 0;
//traverse with max index as there will be empty nodes due to fusion
size_t max_index = subgraph.GetMaxNodeIndex();
for (size_t index = 0; index < max_index; index++) {
auto div_node = subgraph.GetDnnlNode(index);
std::vector<size_t> gelu_indices;
//----------------------------
if (div_node == nullptr || div_node->OpType() != "Div") {
continue;
}
// Check second input is sqrt(2)
// Some Bert models uses this approximation of SQRT2 in the Gelu function
float approximated_sqrt_two = 1.4142099618911743f;
if (!IsInitilizedWithExpectedValue(subgraph, div_node->Input(1), approximated_sqrt_two) &&
!IsInitilizedWithExpectedValue(subgraph, div_node->Input(1), static_cast<float>(M_SQRT2))) {
continue;
}
if (!IsNodeFusable(subgraph, div_node)) {
continue;
}
gelu_indices.push_back(div_node->Index());
//----------------------------
auto erf_node = div_node->Output(0).GetConsumers()[0].GetNode();
if (erf_node == nullptr || erf_node->OpType() != "Erf") {
continue;
}
if (!IsNodeFusable(subgraph, erf_node)) {
continue;
}
gelu_indices.push_back(erf_node->Index());
//----------------------------
auto add_node = erf_node->Output(0).GetConsumers()[0].GetNode();
if (add_node == nullptr || add_node->OpType() != "Add") {
continue;
}
bool is_add_input0 = add_node->Input(0).Name() == erf_node->Output(0).Name();
if (!IsInitilizedWithExpectedValue(subgraph, add_node->Input(is_add_input0 ? 1 : 0), 1.0f)) {
continue;
}
if (!IsNodeFusable(subgraph, add_node)) {
continue;
}
gelu_indices.push_back(add_node->Index());
//----------------------------
auto mul1_node = add_node->Output(0).GetConsumers()[0].GetNode();
if (mul1_node == nullptr || mul1_node->OpType() != "Mul") {
continue;
}
//if (!IsNodeFusable(subgraph, mul1_node)) {
// continue;
//}
gelu_indices.push_back(mul1_node->Index());
//----------------------------
// look for Mul(0.5) using pattern 1 shown above
bool is_pattern_1 = true;
auto mul2_node = FirstParentByType(mul1_node, "Mul");
if (mul2_node != nullptr) {
// the input Div and Mul2 should share at least one input.
bool is_mul2_input0 = div_node->Input(0).Name() == mul2_node->Input(0).Name();
bool is_mul2_input1 = div_node->Input(0).Name() == mul2_node->Input(1).Name();
if (!(is_mul2_input0 ^ is_mul2_input1)) {
is_pattern_1 = false;
}
if (is_pattern_1 && !IsInitilizedWithExpectedValue(subgraph, mul2_node->Input(is_mul2_input0 ? 1 : 0), 0.5f)) {
is_pattern_1 = false;
}
if (is_pattern_1 && !IsNodeFusable(subgraph, mul2_node)) {
is_pattern_1 = false;
}
} else {
is_pattern_1 = false;
}
// look for Mul(0.5) using pattern 2 shown above
if (!is_pattern_1) {
// We only need to check mul1_node IsNodeFusable for pattern 2
if (!IsNodeFusable(subgraph, mul1_node)) {
continue;
}
mul2_node = mul1_node->Output(0).GetConsumers()[0].GetNode();
if (mul2_node == nullptr || mul2_node->OpType() != "Mul") {
continue;
}
if (mul2_node->OutputCount() != 1) {
ORT_THROW("Invalid Mul node");
}
bool is_mul2_first_input = mul2_node->Input(0).Name() == mul1_node->Output(0).Name();
if (!IsInitilizedWithExpectedValue(subgraph, mul2_node->Input(is_mul2_first_input ? 1 : 0), 0.5f)) {
continue;
}
}
gelu_indices.push_back(mul2_node->Index());
//construct new node
auto new_node = std::make_unique<DnnlNode>();
new_node->Name() = div_node->Name() + "_Gelu_" + std::to_string(gelu_index++);
new_node->OpType() = "Gelu";
new_node->Inputs().push_back(div_node->Inputs()[0]);
if (is_pattern_1) {
for (auto def : mul1_node->Outputs()) {
new_node->Outputs().push_back(def);
}
} else {
for (auto def : mul2_node->Outputs()) {
new_node->Outputs().push_back(def);
}
}
// no attributes needed for Gelu if needed this can be updated
//new_node->Attributes().insert(div_node->Attributes());
//insert new node, remove original nodes, connect new edges
ResolveFusion(subgraph, {gelu_indices}, std::move(new_node));
if (debug_log_) {
LOGS_DEFAULT(ERROR) << "Gelu fusion found [" << gelu_index << "]";
}
}
}
/*
Rewrite graph fusing Gelu activation subgraph to a single Gelu node.
The formula corresponding to Gelu activation subgraph :
x * 0.5 * (1.0 + tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x))) or
x * 0.5 * (1.0 + tanh((sqrt(2 / pi) * (x + 0.044715 * pow(x, 3))))),
where x is the input.
*/
void DnnlGraphTransformer::FastGelu(DnnlSubgraph& subgraph) {
static int fastgelu_index = 0;
//traverse with max index as there will be empty nodes due to fusion
size_t max_index = subgraph.GetMaxNodeIndex();
for (size_t index = 0; index < max_index; index++) {
auto dnnl_node = subgraph.GetDnnlNode(index);
if (!FastGeluFirstFormula(subgraph, dnnl_node, fastgelu_index)) {
FastGeluSecondFormula(subgraph, dnnl_node, fastgelu_index);
}
}
}
/*
Rewrite graph fusing Gelu activation subgraph to a single Gelu node.
The formula corresponding to Gelu activation subgraph :
x * 0.5 * (1.0 + tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x)))
where x is the input.
*/
bool DnnlGraphTransformer::FastGeluFirstFormula(DnnlSubgraph& subgraph, DnnlNode* mul1_node, int& fastgelu_index) {
std::vector<size_t> gelu_indices;
//----------mul(0.44715)------------------
if (mul1_node == nullptr || mul1_node->OpType() != "Mul") {
return false;
}
int32_t mul1_input_index = -1;
const float mul_val = 0.044715f;
for (auto i = 0; i < 2; i++) {
if (IsInitilizedWithExpectedValue(subgraph, mul1_node->Input(i), mul_val)) {
mul1_input_index = i;
break;
}
}
if (mul1_input_index == -1) return false;
if (!IsNodeFusable(subgraph, mul1_node)) {
return false;
}
gelu_indices.push_back(mul1_node->Index());
//-------------Mul---------------
auto mul2_node = mul1_node->Output(0).GetConsumers()[0].GetNode();
if (mul2_node == nullptr || mul2_node->OpType() != "Mul") {
return false;
}
if (!IsNodeFusable(subgraph, mul2_node)) {
return false;
}
gelu_indices.push_back(mul2_node->Index());
//-------------Add(1.0)---------------
auto add1_node = mul2_node->Output(0).GetConsumers()[0].GetNode();
if (add1_node == nullptr || add1_node->OpType() != "Add") {
return false;
}
bool is_add_input0 = mul2_node->Output(0).Name() == add1_node->Input(0).Name();
if (!IsInitilizedWithExpectedValue(subgraph, add1_node->Input(is_add_input0 ? 1 : 0), 1.0f)) {
return false;
}
if (!IsNodeFusable(subgraph, add1_node)) {
return false;
}
gelu_indices.push_back(add1_node->Index());
//-------------Mul---------------
auto mul3_node = add1_node->Output(0).GetConsumers()[0].GetNode();
if (mul3_node == nullptr || mul3_node->OpType() != "Mul") {
return false;
}
if (!IsNodeFusable(subgraph, mul3_node)) {
return false;
}
gelu_indices.push_back(mul3_node->Index());
//-------------Mul(0.7978845834732056f)---------------
auto prev_mul4_node = FirstParentByType(mul3_node, "Mul");
if (prev_mul4_node == nullptr) {
return false;
}
int32_t mul4_input_index = -1;
const float mul4_val = 0.7978845834732056f;
for (auto i = 0; i < 2; i++) {
if (IsInitilizedWithExpectedValue(subgraph, prev_mul4_node->Input(i), mul4_val)) {
mul4_input_index = i;
break;
}
}
if (mul4_input_index == -1) return false;
if (!IsNodeFusable(subgraph, prev_mul4_node)) {
return false;
}
gelu_indices.push_back(prev_mul4_node->Index());
auto tanh_node = mul3_node->Output(0).GetConsumers()[0].GetNode();
int32_t x_input_index = (mul1_input_index == 0) ? 1 : 0;
if (FastGeluFormulaCommon(subgraph, mul1_node, x_input_index, tanh_node, gelu_indices, fastgelu_index)) {
if (debug_log_) {
LOGS_DEFAULT(ERROR) << "FastGelu fusion found [" << fastgelu_index << "] (first formula)";
}
return true;
}
return false;
}
/*
Rewrite graph fusing Gelu activation subgraph to a single Gelu node.
The formula corresponding to Gelu activation subgraph :
x * 0.5 * (1.0 + tanh((sqrt(2 / pi) * (x + 0.044715 * pow(x, 3))))),
where x is the input.
*/
void DnnlGraphTransformer::FastGeluSecondFormula(DnnlSubgraph& subgraph, DnnlNode* pow_node, int& fastgelu_index) {
std::vector<size_t> gelu_indices;
//---------Pow-------------------
if (pow_node == nullptr || pow_node->OpType() != "Pow") {
return;
}
auto pow_exponent = pow_node->Input(1);
if (!IsInitilizedWithExpectedValue(subgraph, pow_exponent, 3.0f)) {
return;
}
if (!IsNodeFusable(subgraph, pow_node)) {
return;
}
gelu_indices.push_back(pow_node->Index());
//----------Mul(0.044714998453855515f)------------------
auto mul1_node = pow_node->Output(0).GetConsumers()[0].GetNode();
if (mul1_node == nullptr || mul1_node->OpType() != "Mul") {
return;
}
float fastgelu_muliplyer = 0.044714998453855515f;
bool is_mul1_input0 = pow_node->Output(0).Name() == mul1_node->Input(0).Name();
if (!IsInitilizedWithExpectedValue(subgraph, mul1_node->Input(is_mul1_input0 ? 1 : 0), fastgelu_muliplyer)) {
return;
}
if (!IsNodeFusable(subgraph, mul1_node)) {
return;
}
gelu_indices.push_back(mul1_node->Index());
//----------Add------------------
auto add1_node = mul1_node->Output(0).GetConsumers()[0].GetNode();
if (add1_node == nullptr || add1_node->OpType() != "Add") {
return;
}
if (!IsNodeFusable(subgraph, add1_node)) {
return;
}
gelu_indices.push_back(add1_node->Index());
//----------Mul(sqrt(2/pi))------------------
auto mul2_node = add1_node->Output(0).GetConsumers()[0].GetNode();
if (mul2_node == nullptr || mul2_node->OpType() != "Mul") {
return;
}
// constant is sqrt(2/pi)
float fastgelu_sqrt_2_div_pi = 0.7978845834732056f;
bool is_mul2_input0 = add1_node->Output(0).Name() == mul2_node->Input(0).Name();
if (!IsInitilizedWithExpectedValue(subgraph, mul2_node->Input(is_mul2_input0 ? 1 : 0), fastgelu_sqrt_2_div_pi)) {
return;
}
if (!IsNodeFusable(subgraph, mul2_node)) {
return;
}
gelu_indices.push_back(mul2_node->Index());
//----------Tanh------------------
auto tanh_node = mul2_node->Output(0).GetConsumers()[0].GetNode();
// since the first node is pow the x_input_index is always 0
if (FastGeluFormulaCommon(subgraph, pow_node, 0, tanh_node, gelu_indices, fastgelu_index)) {
if (debug_log_) {
LOGS_DEFAULT(ERROR) << "FastGelu fusion found [" << fastgelu_index << "] (second formula)";
}
}
}
/*
Looks for the part of FastGelu that is common to both formulas if the pattern is found
return true otherwise return false.
i.e. x * 0.5 * (1.0 + tanh(...))
x * 0.5 * (1.0 + tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x))) or
x * 0.5 * (1.0 + tanh((sqrt(2 / pi) * (x + 0.044715 * pow(x, 3))))),
where x is the input.
*/
bool DnnlGraphTransformer::FastGeluFormulaCommon(DnnlSubgraph& subgraph, DnnlNode* gelu_start_node, int32_t x_input_index, DnnlNode* tanh_node, std::vector<size_t>& gelu_indices, int& fastgelu_index) {
//----------Tanh------------------
if (tanh_node == nullptr || tanh_node->OpType() != "Tanh") {
return false;
}
if (!IsNodeFusable(subgraph, tanh_node)) {
return false;
}
gelu_indices.push_back(tanh_node->Index());
//----------Add(1.0)------------------
auto add2_node = tanh_node->Output(0).GetConsumers()[0].GetNode();
if (add2_node == nullptr || add2_node->OpType() != "Add") {
return false;
}
bool is_add2_input0 = tanh_node->Output(0).Name() == add2_node->Input(0).Name();
if (!IsInitilizedWithExpectedValue(subgraph, add2_node->Input(is_add2_input0 ? 1 : 0), 1.0f)) {
return false;
}
if (!IsNodeFusable(subgraph, add2_node)) {
return false;
}
gelu_indices.push_back(add2_node->Index());
//----------Mul------------------
auto mul3_node = add2_node->Output(0).GetConsumers()[0].GetNode();
if (mul3_node == nullptr || mul3_node->OpType() != "Mul") {
return false;
}
if (mul3_node->OutputCount() != 1) {
ORT_THROW("Invalid Mul node");
}
gelu_indices.push_back(mul3_node->Index());
//---------Mul(0.5)---------------------
if (mul3_node->InputCount() != 2) {
return false;
}
auto prev_mul4_node = FirstParentByType(mul3_node, "Mul");
if (prev_mul4_node == nullptr) {
return false;
}
bool is_mul_input0 = gelu_start_node->Input(x_input_index).Name() == prev_mul4_node->Input(0).Name();
bool is_mul_input1 = gelu_start_node->Input(x_input_index).Name() == prev_mul4_node->Input(1).Name();
if (!(is_mul_input0 ^ is_mul_input1)) {
return false;
}
if (!IsInitilizedWithExpectedValue(subgraph, prev_mul4_node->Input(is_mul_input0 ? 1 : 0), 0.5f)) {
return false;
}
if (!IsNodeFusable(subgraph, prev_mul4_node)) {
return false;
}
gelu_indices.push_back(prev_mul4_node->Index());
//construct new node
auto new_node = std::make_unique<DnnlNode>();
new_node->Name() = "Dnnl_FastGelu_" + std::to_string(fastgelu_index++);
new_node->OpType() = "FastGelu";
new_node->Inputs().push_back(gelu_start_node->Inputs()[x_input_index]);
for (auto def : mul3_node->Outputs()) {
new_node->Outputs().push_back(def);
}
// No Attributes needed for FastGelu. If they are needed this can be added in.
//new_node->Attributes().insert(gelu_start_node->Attributes());
//insert new node, remove original nodes, connect new edges
ResolveFusion(subgraph, {gelu_indices}, std::move(new_node));
return true;
}
void DnnlGraphTransformer::ConvRelu(DnnlSubgraph& subgraph) {
//global index of convrelu
static int conv_relu_index = 0;
//traverse with max index as there will be empty nodes due to fusion
size_t max_index = subgraph.GetMaxNodeIndex();
for (size_t index = 0; index < max_index; index++) {
auto dnnl_node = subgraph.GetDnnlNode(index);
//look for conv relu pattern
if (dnnl_node == nullptr) {
continue;
}
if (dnnl_node->OpType() != "Conv") {
continue;
}
if (!IsNodeFusable(subgraph, dnnl_node)) {
continue;
}
auto next_dnnl_node = dnnl_node->Output(0).GetConsumers()[0].GetNode();
if (next_dnnl_node == nullptr) {
continue;
}
if (next_dnnl_node->OpType() != "Relu") {
continue;
}
//construct new node
auto new_node = std::make_unique<DnnlNode>();
new_node->Name() = dnnl_node->Name() + "_ConvRelu_" + std::to_string(conv_relu_index++);
new_node->OpType() = "ConvRelu";
for (auto def : dnnl_node->Inputs()) {
new_node->Inputs().push_back(def);
}
for (auto def : next_dnnl_node->Outputs()) {
new_node->Outputs().push_back(def);
}
new_node->Attributes().insert(dnnl_node->Attributes());
//insert new node, remove original nodes, connect new edges
if (debug_log_) {
LOGS_DEFAULT(ERROR) << "ConvRelu fusion of [" << dnnl_node->Name() << "] and [" << next_dnnl_node->Name() << "]";
}
ResolveFusion(subgraph, {dnnl_node->Index(), next_dnnl_node->Index()}, std::move(new_node));
}
}
void DnnlGraphTransformer::MatMulAdd(DnnlSubgraph& subgraph) {
static int fused_index = 0;
size_t max_index = subgraph.GetMaxNodeIndex();
for (size_t index = 0; index < max_index; index++) {
auto dnnl_node = subgraph.GetDnnlNode(index);
//look for conv relu pattern
if (dnnl_node == nullptr || dnnl_node->OpType() != "MatMul") {
continue;
}
auto matmul_node = dnnl_node;
if (!IsNodeFusable(subgraph, dnnl_node)) {
continue;
}
auto next_dnnl_node = matmul_node->Output(0).GetConsumers()[0].GetNode();
if (next_dnnl_node == nullptr) {
continue;
}
if (next_dnnl_node->OpType() != "Add") {
continue;
}
/*
add now has one of the input connecting to matmul's single output
different cases:
matmul input goes to the other add input
matmul output goes to the other add input (adding two identical tensor)
some other tensor goes to the other add input
*/
auto add_node = next_dnnl_node;
auto matmul_output_name = matmul_node->Output(0).Name();
auto add_inputs = add_node->Inputs();
//add is taking two inputs from the same matmul output
//not sure if onednn would support such post ops
if (add_inputs[0] == add_inputs[1]) {
continue;
}
auto fused_node_inputs = matmul_node->Inputs();
//the 3rd input to fused matmul
if (matmul_output_name == add_inputs[0]->Name()) {
fused_node_inputs.push_back(add_inputs[1]);
} else {
fused_node_inputs.push_back(add_inputs[0]);
}
auto fused_node_output = add_node->Outputs()[0];
auto fused_node_name = matmul_node->Name() + "_" + matmul_node->OpType() + "Add_" + std::to_string(fused_index++);
auto fused_node_type = matmul_node->OpType() + "Add";
//construct new node
auto fused_node = std::make_unique<DnnlNode>();
fused_node->Name() = fused_node_name;
fused_node->OpType() = fused_node_type;
fused_node->Inputs() = fused_node_inputs;
fused_node->Outputs() = {fused_node_output};
//no attribute for matmul and add
//insert new node, remove original nodes, connect new edges
if (debug_log_) {
LOGS_DEFAULT(ERROR) << "MatMulAdd fusion of [" << matmul_node->Name() << "] and [" << add_node->Name() << "]";
}
ResolveFusion(subgraph, {matmul_node->Index(), add_node->Index()}, std::move(fused_node));
}
}
} // namespace ort_dnnl
} // namespace onnxruntime

49
onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_transformer.h Normal file
View file

@ -0,0 +1,49 @@
// Copyright(C) 2021 Intel Corporation
// Licensed under the MIT License
#pragma once
#include "dnnl_subgraph.h"
namespace onnxruntime {
namespace ort_dnnl {
class DnnlGraphTransformer {
public:
void Apply(DnnlSubgraph& subgraph);
DnnlGraphTransformer() {
const std::string debug_log_env = onnxruntime::GetEnvironmentVar("ORT_DNNL_DEBUG_LOG");
if (!debug_log_env.empty()) {
debug_log_ = (std::stoi(debug_log_env) == 0 ? false : true);
}
}
private:
void Gelu(DnnlSubgraph& subgraph);
void FastGelu(DnnlSubgraph& subgraph);
bool FastGeluFirstFormula(DnnlSubgraph& subgraph, DnnlNode* node, int& fastgelu_index);
void FastGeluSecondFormula(DnnlSubgraph& subgraph, DnnlNode* node, int& fastgelu_index);
bool FastGeluFormulaCommon(DnnlSubgraph& subgraph, DnnlNode* gelu_start_node, int32_t x_input_index, DnnlNode* tanh_node, std::vector<size_t>& gelu_indices, int& fastgelu_index);
bool IsInitilizedWithExpectedValue(DnnlSubgraph& subgraph, DnnlTensor& input_arg, float expected_value);
void ConvRelu(DnnlSubgraph& subgraph);
void MatMulAdd(DnnlSubgraph& subgraph);
// This function checks a few things
// - the node in question has a single output
// - The output of the node is only consumed by a one other node
// - the output tensor from the node is going to another node within the subgraph
// If all of the above is true this will return true. It will return false otherwise.
//
// It is possible for a node to fail one or more of the checks above and still be fusable.
//
// The name of the function was chosen because this check is required for most of the node fusions
// found in the code.
//
// The last node in a fusion does not need to pass this check.
bool IsNodeFusable(DnnlSubgraph& subgraph, DnnlNode* node) const;
void ResolveFusion(DnnlSubgraph& subgraph, std::vector<size_t> old_indices, std::unique_ptr<DnnlNode> new_node);
bool ProduceGraphOutput(DnnlSubgraph& subgraph, DnnlNode& node);
bool IsGraphOutput(DnnlSubgraph& subgraph, DnnlTensor& tensor);
bool debug_log_ = false;
};
} // namespace ort_dnnl
} // namespace onnxruntime

12
onnxruntime/test/providers/cpu/activation/activation_op_test.h
View file

@ -67,6 +67,18 @@ inline void TestActivationOp(const char* szOp, const std::vector<std::vector<T>>
if (relu == 0) {
excluded_providers.insert(kNnapiExecutionProvider);
}
#endif
// Use relative error because of computation error for float::max
#if defined(USE_DNNL)
int gelu = strcmp(szOp, "Gelu");
if (gelu == 0) {
// OneDNN has a computation difference when computing FLT_MAX
// Expected: 3.4028234663852886e+38
// Actual: 3.4028232635611926e+38
// Since the numbers are large relative error is used instead of
// the default threshold which is a small value.
test.SetOutputRelErr("Y", .000001f);
}
#endif
test.Run(OpTester::ExpectResult::kExpectSuccess, "", excluded_providers);
}