diff --git a/onnxruntime/core/optimizer/div_mul_fusion.cc b/onnxruntime/core/optimizer/div_mul_fusion.cc new file mode 100644 index 0000000000..4f191b1b63 --- /dev/null +++ b/onnxruntime/core/optimizer/div_mul_fusion.cc @@ -0,0 +1,101 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/div_mul_fusion.h" + +#include "core/graph/graph_utils.h" +#include "core/optimizer/initializer.h" +#include "core/optimizer/utils.h" + +using namespace ONNX_NAMESPACE; +using namespace onnxruntime::common; +namespace onnxruntime { + +/** +Transform that fuses two Div -> Mul nodes to a single Div node +when the first input to Div is 1. +1 / x1 * x2 -> x2 / x1 + */ +bool DivMulFusion::SatisfyCondition(const Graph& graph, const Node& node, const logging::Logger&) const { + if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "Div", {7, 13}) || + node.GetOutputEdgesCount() != 1) { + return false; + } + + const auto& next_node = *node.OutputNodesBegin(); + if (!graph_utils::IsSupportedOptypeVersionAndDomain(next_node, "Mul", {7, 13}) || + // Make sure the two nodes do not span execution providers. + next_node.GetExecutionProviderType() != node.GetExecutionProviderType()) { + return false; + } + + // Check that the appropriate input to the Div node is a constant. + if (!graph_utils::NodeArgIsConstant(graph, *node.InputDefs()[0])) { + return false; + } + + const auto* initializer = graph_utils::GetConstantInitializer(graph, node.InputDefs()[0]->Name()); + if (!initializer) { + return false; + } + + int32_t data_type = initializer->data_type(); + Initializer div_A(*initializer, graph.ModelPath()); + if (div_A.size() > 1) { + return false; + } + switch (data_type) { + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + if (*div_A.data() != 1.f) { + return false; + } + break; + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: + if (math::halfToFloat(div_A.data()->val) != 1.f) { + return false; + } + break; + case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE: + if (*div_A.data() != static_cast(1.f)) { + return false; + } + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT32: + if (*div_A.data() != static_cast(1)) { + return false; + } + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + if (*div_A.data() != static_cast(1)) { + return false; + } + break; + default: + return false; + } + + if (!graph.GetNodeOutputsInGraphOutputs(node).empty()) { + return false; + } + + return true; +} + +Status DivMulFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_effect, const logging::Logger&) const { + auto& div_node = node; + auto& mul_node = *graph.GetNode(div_node.OutputNodesBegin()->Index()); // get mutable next node + const auto& div_output = div_node.OutputDefs(); + auto& mul_inputs = mul_node.MutableInputDefs(); + + //get other input of mul + auto& mul_other_input = mul_inputs[0] == div_output[0] ? mul_inputs[1] : mul_inputs[0]; + + graph_utils::ReplaceNodeInput(div_node, 0, *mul_other_input); + // move the output definition and edges from the mul_node to the div_node and delete the mul_node + graph_utils::FinalizeNodeFusion(graph, div_node, mul_node); + + rule_effect = RewriteRuleEffect::kModifiedRestOfGraph; + + return Status::OK(); +} +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/div_mul_fusion.h b/onnxruntime/core/optimizer/div_mul_fusion.h new file mode 100644 index 0000000000..809f6f7aef --- /dev/null +++ b/onnxruntime/core/optimizer/div_mul_fusion.h @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/rewrite_rule.h" + +namespace onnxruntime { +/** +@Class DivMulFusion + +Rewrite rule that fuses two Div -> Mul nodes to a single Div node +when the first input to Div is 1. +1 / x1 * x2 -> x2 / x1 + +*/ +class DivMulFusion : public RewriteRule { + public: + DivMulFusion() noexcept : RewriteRule("DivMulFusion") {} + + std::vector TargetOpTypes() const noexcept override { + return {"Div"}; + } + + private: + bool SatisfyCondition(const Graph& graph, const Node& node, const logging::Logger& logger) const override; + + Status Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_effect, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 14de60cfe9..3b7556ac2f 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -14,6 +14,7 @@ #include "core/optimizer/conv_add_fusion.h" #include "core/optimizer/conv_bn_fusion.h" #include "core/optimizer/conv_mul_fusion.h" +#include "core/optimizer/div_mul_fusion.h" #include "core/optimizer/dropout_elimination.h" #include "core/optimizer/dynamic_quantize_matmul_fusion.h" #include "core/optimizer/embed_layer_norm_fusion.h" @@ -63,6 +64,7 @@ std::vector> GenerateRewriteRules( rules.push_back(onnxruntime::make_unique()); rules.push_back(onnxruntime::make_unique()); rules.push_back(onnxruntime::make_unique()); + rules.push_back(onnxruntime::make_unique()); rules.push_back(onnxruntime::make_unique()); rules.push_back(onnxruntime::make_unique()); rules.push_back(onnxruntime::make_unique()); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 941fae0b76..0e7b6bd220 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -27,6 +27,7 @@ #include "core/optimizer/conv_add_fusion.h" #include "core/optimizer/conv_bn_fusion.h" #include "core/optimizer/conv_mul_fusion.h" +#include "core/optimizer/div_mul_fusion.h" #include "core/optimizer/dropout_elimination.h" #include "core/optimizer/dynamic_quantize_matmul_fusion.h" #include "core/optimizer/embed_layer_norm_fusion.h" @@ -527,6 +528,26 @@ TEST_F(GraphTransformationTests, FuseConvBNMulAddUnsqueeze) { } } +TEST_F(GraphTransformationTests, DivMulFusion) { + auto model_uri = MODEL_FOLDER "fusion/div_mul.onnx"; + std::shared_ptr model; + ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_)); + Graph& graph = model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["Div"] == 5); + ASSERT_TRUE(op_to_count["Mul"] == 5); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + auto rule_transformer_L1 = onnxruntime::make_unique("RuleTransformer1"); + rule_transformer_L1->Register(onnxruntime::make_unique()); + graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["Div"] == 5); + ASSERT_TRUE(op_to_count["Mul"] == 2); +} + #if defined(USE_CUDA) && !defined(DISABLE_CONTRIB_OPS) TEST_F(GraphTransformationTests, FuseCudaConvAddRelu) { auto model_uri = MODEL_FOLDER "fusion/conv_add_relu.onnx"; diff --git a/onnxruntime/test/testdata/transform/fusion/div_mul.onnx b/onnxruntime/test/testdata/transform/fusion/div_mul.onnx new file mode 100644 index 0000000000..a37dc14b97 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/div_mul.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/div_mul.py b/onnxruntime/test/testdata/transform/fusion/div_mul.py new file mode 100644 index 0000000000..db05d5b84a --- /dev/null +++ b/onnxruntime/test/testdata/transform/fusion/div_mul.py @@ -0,0 +1,69 @@ +import onnx +from onnx import helper +from onnx import TensorProto, OperatorSetIdProto +from enum import Enum + +opsets = [] +onnxdomain = OperatorSetIdProto() +onnxdomain.version = 12 +onnxdomain.domain = "" # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification. +opsets.append(onnxdomain) + +msdomain = OperatorSetIdProto() +msdomain.version = 1 +msdomain.domain = 'com.microsoft' + +opsets.append(msdomain) +kwargs={} +kwargs['opset_imports'] = opsets + +def GenerateModel(model_name): + nodes = [ # subgraph + # float + helper.make_node("Div", ["float_1", "A"], ["div_1"], "div_1"), + helper.make_node("Mul", ["div_1", "B"], ["mul_1"], "mul_1"), + helper.make_node("Cast", ["mul_1"], ["cast_1"], "cast_1", to=10), + # float_16 + helper.make_node("Div", ["float16_1", "cast_1"], ["div_2"], "div_2"), + helper.make_node("Mul", ["C", "div_2"], ["mul_2"], "mul_2"), + helper.make_node("Cast", ["mul_2"], ["cast_2"], "cast_2", to=7), + # int64 + helper.make_node("Div", ["int64_1", "cast_2"], ["div_3"], "div_3"), + helper.make_node("Mul", ["D", "div_3"], ["mul_3"], "mul_3"), + helper.make_node("Identity", ["mul_3"], ["Y"], "output"), + # div has >1 consumers + helper.make_node("Div", ["float_1", "A"], ["div_4"], "div_4"), + helper.make_node("Mul", ["div_4", "B"], ["mul_4"], "mul_4"), + # div is graph output + helper.make_node("Div", ["float_1", "div_4"], ["div_5"], "div_5"), + helper.make_node("Mul", ["div_5", "B"], ["mul_5"], "mul_5"), + ] + + inputs = [ # inputs + helper.make_tensor_value_info('A', TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info('B', TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info('C', TensorProto.FLOAT16, ['M', 'K']), + helper.make_tensor_value_info('D', TensorProto.INT64, ['M', 'K']), + ] + + initializers = [ + helper.make_tensor('float_1', TensorProto.FLOAT, [1], [1.0]), + helper.make_tensor('float16_1', TensorProto.FLOAT16, [1], [15360]), # 15360 is the fp16 representation of 1.f + helper.make_tensor('int64_1', TensorProto.INT64, [1], [1]), + ] + + graph = helper.make_graph( + nodes, + "DivMul", #name + inputs, + [ # outputs + helper.make_tensor_value_info('Y', TensorProto.INT64, ['M', 'K']), + helper.make_tensor_value_info('div_5', TensorProto.FLOAT, ['M', 'K']), + ], + initializers) + + model = helper.make_model(graph, **kwargs) + onnx.save(model, model_name) + +if __name__ == "__main__": + GenerateModel('div_mul.onnx') \ No newline at end of file diff --git a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc index f6c6ec3f67..535a2eca74 100644 --- a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc +++ b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc @@ -15,6 +15,7 @@ #include "core/optimizer/conv_add_fusion.h" #include "core/optimizer/conv_bn_fusion.h" #include "core/optimizer/conv_mul_fusion.h" +#include "core/optimizer/div_mul_fusion.h" #include "core/optimizer/dropout_elimination.h" #include "core/optimizer/embed_layer_norm_fusion.h" #include "core/optimizer/expand_elimination.h" @@ -73,6 +74,7 @@ std::vector> GeneratePreTrainingTransformers( rule_transformer->Register(make_unique()); rule_transformer->Register(make_unique()); rule_transformer->Register(make_unique()); + rule_transformer->Register(make_unique()); rule_transformer->Register(make_unique()); rule_transformer->Register(make_unique()); rule_transformer->Register(make_unique());