Optimizer - add missing supported version for BiasSoftmaxFusion (#11616)

* add missing version

* opset check

* fix format

* reject fusion if type not allowed

* per comments

* trigger new build

Co-authored-by: Ethan Tao <ettao@microsoft.com@orttrainingdev7.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>

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();

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<std::string, int> 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) {

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(