diff --git a/onnxruntime/contrib_ops/contrib_kernels.cc b/onnxruntime/contrib_ops/contrib_kernels.cc index 76192ddcad..98aaef9c0f 100644 --- a/onnxruntime/contrib_ops/contrib_kernels.cc +++ b/onnxruntime/contrib_ops/contrib_kernels.cc @@ -9,6 +9,7 @@ namespace contrib { class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SampleOp); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ExpandDims); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedConv); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, AttnLSTM); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, string, Tokenizer); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, DequantizeLinear); @@ -28,6 +29,7 @@ void RegisterContribKernels(std::function fn) { // add more kernels here fn(BuildKernel()); + fn(BuildKernel()); fn(BuildKernel()); fn(BuildKernel()); fn(BuildKernel()); diff --git a/onnxruntime/core/graph/conv_activation_fusion.cc b/onnxruntime/core/graph/conv_activation_fusion.cc index 2d43d8c001..15b30e8fca 100644 --- a/onnxruntime/core/graph/conv_activation_fusion.cc +++ b/onnxruntime/core/graph/conv_activation_fusion.cc @@ -4,6 +4,7 @@ #include "core/graph/initializer.h" #include "core/graph/conv_activation_fusion.h" #include "core/graph/graph_utils.h" +#include using namespace onnx; using namespace ::onnxruntime::common; @@ -13,13 +14,31 @@ namespace { bool IsFusableActivation(const Node& node) { return utils::IsSupportedOptypeVersionAndDomain(node, "LeakyRelu", 6) || utils::IsSupportedOptypeVersionAndDomain(node, "Relu", 6) || utils::IsSupportedOptypeVersionAndDomain(node, "Sigmoid", 6) || utils::IsSupportedOptypeVersionAndDomain(node, "Tanh", 6); } + +void HandleActivationNodeEdges(Graph& g, const Node& act, Node& fused_conv) { + Node::EdgeSet output_edges; + for (auto it = act.OutputEdgesBegin(); it != act.OutputEdgesEnd(); ++it) { + output_edges.insert(*it); + } + + //remove output edge of activation + //connect fused_conv node and nodes after activation nodes + for (auto& output_edge : output_edges) { + NodeIndex dst_node_index = output_edge.GetNode().Index(); + int src_arg_index = output_edge.GetSrcArgIndex(); + int dst_arg_index = output_edge.GetDstArgIndex(); + g.RemoveEdge(act.Index(), dst_node_index, src_arg_index, dst_arg_index); + g.AddEdge(fused_conv.Index(), dst_node_index, 0, dst_arg_index); + } +} + } // namespace Status ConvActivationFusion::Apply(Graph& graph, bool& modified) const { GraphViewer graph_viewer(graph); const auto& order = graph_viewer.GetNodesInTopologicalOrder(); - std::vector removed_nodes; + std::deque removed_nodes; for (auto index : order) { auto node = graph.GetNode(index); if (!utils::IsSupportedOptypeVersionAndDomain(*node, "Conv", 1) || node->GetOutputEdgesCount() != 1) { @@ -32,7 +51,6 @@ Status ConvActivationFusion::Apply(Graph& graph, bool& modified) const { Node* conv_node = node; const Node& act_node = next_node; - std::vector input_args, output_args; Node& fused_conv = graph.AddNode(graph.GenerateNodeName("fused " + conv_node->Name()), "FusedConv", "fused Conv " + conv_node->Name() + "with activation " + act_node.OpType(), @@ -42,7 +60,7 @@ Status ConvActivationFusion::Apply(Graph& graph, bool& modified) const { "com.microsoft"); //Add a new attribute to specify the activation type - fused_conv.AddAttribute("activation", "string"); + fused_conv.AddAttribute("activation", act_node.OpType()); //Add optional attributes for activations if (act_node.OpType() == "LeakyRelu") { @@ -52,6 +70,8 @@ Status ConvActivationFusion::Apply(Graph& graph, bool& modified) const { } } + HandleActivationNodeEdges(graph, act_node, fused_conv); + // Replace the input of the node following activation node const NodeArg* act_output_def = act_node.OutputDefs()[0]; NodeArg* fused_conv_output_def = fused_conv.MutableOutputDefs()[0]; @@ -69,12 +89,12 @@ Status ConvActivationFusion::Apply(Graph& graph, bool& modified) const { } } - removed_nodes.push_back(act_node.Index()); - removed_nodes.push_back(conv_node->Index()); + removed_nodes.push_front(conv_node->Index()); + removed_nodes.push_front(act_node.Index()); } - for (auto i : removed_nodes) { - graph.RemoveNode(i); + for (auto node : removed_nodes) { + graph.RemoveNode(node); } if (!removed_nodes.empty()) {