diff --git a/onnxruntime/core/optimizer/bias_softmax_fusion.cc b/onnxruntime/core/optimizer/bias_softmax_fusion.cc old mode 100644 new mode 100755 index f44c2169fd..0ca2e8b13f --- a/onnxruntime/core/optimizer/bias_softmax_fusion.cc +++ b/onnxruntime/core/optimizer/bias_softmax_fusion.cc @@ -59,9 +59,23 @@ bool TryBiasSoftmaxSubgraphMatch(Graph& graph, Node& start, Node*& add, Node*& s return false; } + // BiasSoftmax supports only float/float16/double - see ./onnxruntime/core/graph/contrib_ops/contrib_defs.cc + auto type_allowed = [](NodeArg* input) { + auto data_type = input->TypeAsProto()->tensor_type().elem_type(); + if (data_type != ONNX_NAMESPACE::TensorProto_DataType_DOUBLE && + data_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 && + data_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { + return false; + } + return true; + }; + if (!type_allowed(input1) || !type_allowed(input2)) { + return false; + } + // check add is only consumed by softmax with matching exec provider Node& softmax_node = *graph.GetNode(add_node.OutputNodesBegin()->Index()); - if (!graph_utils::IsSupportedOptypeVersionAndDomain(softmax_node, "Softmax", {1, 11}) || + if (!graph_utils::IsSupportedOptypeVersionAndDomain(softmax_node, "Softmax", {1, 11, 13}) || softmax_node.GetExecutionProviderType() != add_node.GetExecutionProviderType()) { return false; } @@ -107,11 +121,12 @@ bool TrySelectInputAndBiasWithAlignment( // confirm all dimensions starting from softmax axis match for input and mask bool singlebatch_shape_matches = true; - int axis = 1; + // default axis = -1 if opset >= 13 + int axis = graph_utils::MatchesOpSinceVersion(softmax_node, {1, 11}) ? 1 : -1; auto& softmax_attr = softmax_node.GetAttributes(); if (softmax_attr.find("axis") != softmax_attr.end()) { auto& axis_attr = softmax_attr.at("axis"); - axis = utils::HasInt(axis_attr) ? (int)axis_attr.i() : 1; + axis = utils::HasInt(axis_attr) ? (int)axis_attr.i() : axis; } int N1 = input1->Shape()->dim_size(); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc old mode 100644 new mode 100755 index aa85e9029f..f24f317b18 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -3514,12 +3514,9 @@ struct BiasSoftmaxFusionTester { } } - void TestFusionOccurs(int expected_broadcast_axis) { + void TestFusionOccurs(int expected_broadcast_axis, int expected_softmax_axis) { ASSERT_STATUS_OK(model_load_); - int expected_softmax_axis = 1; - GetAxis("Softmax", "axis", &expected_softmax_axis); - ASSERT_STATUS_OK(graph_transformation_mgr_.ApplyTransformers(p_model_->MainGraph(), TransformerLevel::Level2, *logger_)); std::map op_to_count = CountOpsInGraph(p_model_->MainGraph()); @@ -3556,25 +3553,37 @@ TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_GpuOnly) { TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_Simple_Rocm) { auto model_uri = MODEL_FOLDER "fusion/bias_softmax_fusion_simple.onnx"; BiasSoftmaxFusionTester tester(model_uri, logger_.get(), kRocmExecutionProvider); - tester.TestFusionOccurs(1); + tester.TestFusionOccurs(1, 1); } TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_Simple_Cuda) { auto model_uri = MODEL_FOLDER "fusion/bias_softmax_fusion_simple.onnx"; BiasSoftmaxFusionTester tester(model_uri, logger_.get()); - tester.TestFusionOccurs(1); + tester.TestFusionOccurs(1, 1); +} + +TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_Simple_Opset13_DefaultAxis) { + auto model_uri = MODEL_FOLDER "fusion/bias_softmax_fusion_simple_no_axis_opset13.onnx"; + BiasSoftmaxFusionTester tester(model_uri, logger_.get()); + tester.TestFusionOccurs(1, 1); +} + +TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_BFloat16_Input) { + auto model_uri = MODEL_FOLDER "fusion/bias_softmax_fusion_bfloat16.onnx"; + BiasSoftmaxFusionTester tester(model_uri, logger_.get()); + tester.TestNoFusionOccurs(); } TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_MiddleOnes) { auto model_uri = MODEL_FOLDER "fusion/bias_softmax_fusion_middleones.onnx"; BiasSoftmaxFusionTester tester(model_uri, logger_.get()); - tester.TestFusionOccurs(3); + tester.TestFusionOccurs(3, 6); } TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_ReversedInputs) { auto model_uri = MODEL_FOLDER "fusion/bias_softmax_fusion_middleones_reversed.onnx"; BiasSoftmaxFusionTester tester(model_uri, logger_.get()); - tester.TestFusionOccurs(3); + tester.TestFusionOccurs(3, 6); } TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_BadAxis) { @@ -3586,19 +3595,19 @@ TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_BadAxis) { TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_AllLeadingOnes) { auto model_uri = MODEL_FOLDER "fusion/bias_softmax_fusion_allleadingones.onnx"; BiasSoftmaxFusionTester tester(model_uri, logger_.get()); - tester.TestFusionOccurs(0); + tester.TestFusionOccurs(0, 6); } TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_SomeLeadingOnes) { auto model_uri = MODEL_FOLDER "fusion/bias_softmax_fusion_someleadingones.onnx"; BiasSoftmaxFusionTester tester(model_uri, logger_.get()); - tester.TestFusionOccurs(0); + tester.TestFusionOccurs(0, 6); } TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_NoLeadingOnes) { auto model_uri = MODEL_FOLDER "fusion/bias_softmax_fusion_noleadingones.onnx"; BiasSoftmaxFusionTester tester(model_uri, logger_.get()); - tester.TestFusionOccurs(0); + tester.TestFusionOccurs(0, 6); } static void TestBiasDropoutFusion(const PathString& file_path, const logging::Logger& logger, const int add_count = 0) { diff --git a/onnxruntime/test/testdata/transform/fusion/bias_softmax_fusion_bfloat16.onnx b/onnxruntime/test/testdata/transform/fusion/bias_softmax_fusion_bfloat16.onnx new file mode 100644 index 0000000000..04b53417ec Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/bias_softmax_fusion_bfloat16.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/bias_softmax_fusion_simple_no_axis_opset13.onnx b/onnxruntime/test/testdata/transform/fusion/bias_softmax_fusion_simple_no_axis_opset13.onnx new file mode 100644 index 0000000000..a0a8905dcd Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/bias_softmax_fusion_simple_no_axis_opset13.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/bias_softmax_gen.py b/onnxruntime/test/testdata/transform/fusion/bias_softmax_gen.py old mode 100644 new mode 100755 index 26002476c6..9c5312db77 --- a/onnxruntime/test/testdata/transform/fusion/bias_softmax_gen.py +++ b/onnxruntime/test/testdata/transform/fusion/bias_softmax_gen.py @@ -1,11 +1,59 @@ import onnx -from onnx import TensorProto, helper +from onnx import OperatorSetIdProto, TensorProto, helper add = helper.make_node("Add", ["input", "bias"], ["add_out"], "add") reverseadd = helper.make_node("Add", ["bias", "input"], ["add_out"], "add") softmax1 = helper.make_node("Softmax", ["add_out"], ["output"], "softmax", axis=1) softmax3 = helper.make_node("Softmax", ["add_out"], ["output"], "softmax", axis=3) softmax6 = helper.make_node("Softmax", ["add_out"], ["output"], "softmax", axis=6) +softmax_no_axis = helper.make_node("Softmax", ["add_out"], ["output"], "softmax") + +onnxdomain = OperatorSetIdProto() +onnxdomain.version = 13 +# The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification. +onnxdomain.domain = "" +msdomain = OperatorSetIdProto() +msdomain.version = 1 +msdomain.domain = "com.microsoft" +opsets = [onnxdomain, msdomain] + +onnx.save( + helper.make_model( + helper.make_graph( + [add, softmax_no_axis], + "Add_Softmax_Fusion", + [ + helper.make_tensor_value_info("input", TensorProto.FLOAT, ["d_1", "d_2"]), + helper.make_tensor_value_info("bias", TensorProto.FLOAT, ["d_1", "d_2"]), + ], + [ + helper.make_tensor_value_info("output", TensorProto.FLOAT, ["d_1", "d_2"]), + ], + [], + ), + opset_imports=opsets, + ), + r"bias_softmax_fusion_simple_no_axis_opset13.onnx", +) + +onnx.save( + helper.make_model( + helper.make_graph( + [add, softmax1], + "Add_Softmax_Fusion", + [ + helper.make_tensor_value_info("input", TensorProto.BFLOAT16, ["d_1", "d_2"]), + helper.make_tensor_value_info("bias", TensorProto.BFLOAT16, ["d_1", "d_2"]), + ], + [ + helper.make_tensor_value_info("output", TensorProto.BFLOAT16, ["d_1", "d_2"]), + ], + [], + ), + opset_imports=opsets, + ), + r"bias_softmax_fusion_bfloat16.onnx", +) onnx.save( helper.make_model( @@ -25,6 +73,7 @@ onnx.save( r"bias_softmax_fusion_simple.onnx", ) + onnx.save( helper.make_model( helper.make_graph(