graph: interface: op: softmax supports mixed data types

Author: TaoLv

src/graph/interface/op_def.hpp

@@ -1032,12 +1032,15 @@ DNNL_GRAPH_OP_SCHEMA(SoftMax, 1,
         op_schema_t()
                 .set_num_inputs(1)
                 .set_num_outputs(1)
-                .set_input(0, "src", "T")
-                .set_output(0, "dst", "T")
+                .set_input(0, "src", "T1")
+                .set_output(0, "dst", "T2")
                 .set_attr(op_attr::axis, false, attribute_kind::i, (int64_t)1)
                 .set_type_constraints(
-                        "T", {data_type::f32, data_type::bf16, data_type::f16})
-                .set_shape_inference_function(infer_identity_output_shape))
+                        "T1", {data_type::f32, data_type::bf16, data_type::f16})
+                .set_type_constraints(
+                        "T2", {data_type::f32, data_type::bf16, data_type::f16})
+                .set_shape_inference_function(infer_identity_output_shape)
+                .set_op_def_constraint_function(check_softmax_dtype))
 
 DNNL_GRAPH_OP_SCHEMA(SoftMaxBackward, 1,
         op_schema_t()

src/graph/interface/op_def_constraint.cpp

@@ -105,6 +105,25 @@ bool check_matmul_dtype(const op_t *mm) {
     return true;
 }
 
+// For SoftMax, if the src is f32, dst can be xf16. Otherwise, src and dst
+// should have the same data type.
+bool check_softmax_dtype(const op_t *n) {
+    const auto &inputs = n->get_input_values();
+    const auto &outputs = n->get_output_values();
+
+    const logical_tensor_t &src = inputs[0]->get_logical_tensor();
+    const logical_tensor_t &dst = outputs[0]->get_logical_tensor();
+    if (src.data_type != dst.data_type) {
+        if (src.data_type != data_type::f32) {
+            VCHECK_SHAPE_INFER(false, "%s, %s src + %s dst is not supported",
+                    op_t::kind2str(n->get_kind()).c_str(),
+                    dnnl_dt2str(src.data_type), dnnl_dt2str(dst.data_type));
+        }
+    }
+
+    return true;
+}
+
 // check function for data_type of LayerNorm and GroupNorm.
 // only when data is bf16, gamma/beta/mean/var can be bf16.
 // If data is bf16, gamma/beta/mean/var can be f32 or bf16.

src/graph/interface/op_def_constraint.hpp

@@ -30,6 +30,8 @@ bool check_bn_data_type(const op_t *n);
 
 bool check_matmul_dtype(const op_t *n);
 
+bool check_softmax_dtype(const op_t *n);
+
 bool check_ln_gn_data_type(const op_t *n);
 
 bool check_typecast_data_type(const op_t *n);