Support 'complex promote' in yaml (#50611)

* support 'complex promote' in yaml

* change the compplex_promote

* change 'kron' in math.py

* change 'kron' comment in python

* change kron comment in python

* change kron comment in python

Author: heavyrain-lzy

paddle/fluid/operators/generator/generate_op.py

@@ -293,6 +293,11 @@ def add_grad_op_compat_name(grad_op_item, args_name_map):
         if new_op_name != op_name:
             forward_op_item['op_name'] = op_name
 
+        # add complex promote infomation
+        if "complex_promote" in op_args:
+            forward_op_item["complex_promote"] = op_args["complex_promote"]
+            if has_backward:
+                backward_op_item["complex_promote"] = op_args["complex_promote"]
         scalar_configs = None
         int_array_configs = None
         if 'scalar' in op_args:

paddle/fluid/operators/generator/templates/operator_utils.c.j2

@@ -279,30 +279,52 @@ phi::KernelKey GetExpectedKernelType(
     data_type = framework::OperatorWithKernel::IndicateVarDataType(ctx, {{data_type_args[1] | to_opmaker_name}});
   }
   {% endif %}
+{% elif "complex_promote" in op and "forward" not in op%}
+  {% set inputs = op["complex_promote"]%}
+  auto data_type =
+        OperatorWithKernel::IndicateOrPromoteVarDataTypes(ctx, "{{inputs[0]}}", "{{inputs[1]}}");
 {% endif %}
   return phi::KernelKey(data_type, ctx.GetPlace());
 }
-{% endmacro %}
+{% endmacro -%}
+
+{% macro get_kernel_for_var(op) %}
+{% set skip_args = none %}
+{% if op["data_transform"] is not none%}
+  {% if "skip_transform" in op["data_transform"] %}
+    {% set skip_args = op["data_transform"]["skip_transform"] %}
+  {% elif "support_trans_dtype" in op["data_transform"] %}
+    {% set skip_args = op["data_transform"]["support_trans_dtype"] %}
+  {% endif %}
+{% endif %}
+{% set var_name = "var_name" -%}
 
-{% macro get_kernel_for_var(op) %} {# only for data_transform #}
-{% set skip_args = op["data_transform"]["skip_transform"] %}
-{% set var_name = "var_name" %}
-{% set skip_args_len = skip_args | length %}
 phi::KernelKey GetKernelTypeForVar(
     const std::string& {{var_name}},
     const phi::DenseTensor& tensor,
     const phi::KernelKey& expected_kernel_type) const override {
-    
+{%if skip_args is not none%}{# deal data_transform #}
+  {% set skip_args_len = skip_args | length %}
       if (
         {%- for skip_arg in skip_args -%}
           var_name == "{{ skip_arg }}"
           {%- if skip_args_len != 1 and loop.index != skip_args_len %} || {% endif -%}
         {%- endfor -%}
       ){
+    {% if "skip_transform" in op["data_transform"] %}
         return phi::KernelKey(phi::Backend::ALL_BACKEND,
                             expected_kernel_type.layout(),
                             expected_kernel_type.dtype());
+    {% elif "support_trans_dtype" in op["data_transform"] %}
+       return phi::KernelKey(tensor.place(), tensor.layout(), tensor.dtype());
+    {% endif %}
+      }
+{% else %}{# deal complex_promote #}
+      if (framework::IsComplexType(expected_kernel_type.dtype())) {
+      // only promote inputs’s types when contains complex input
+        return phi::KernelKey(tensor.place(), tensor.layout(), tensor.dtype());
       }
+{% endif %}
       else{
           return phi::KernelKey(
             tensor.place(), tensor.layout(), expected_kernel_type.dtype());
@@ -317,20 +339,23 @@ class {{op["op_name"] | to_pascal_case}}Op : public framework::OperatorWithKerne
   using framework::OperatorWithKernel::OperatorWithKernel;
   {# ----------- get expected kernel type function -------------------------- #}
   {% set kernel = op["kernel"] %}
-  {% if kernel["data_type"] is not none %}
+  {% if kernel["data_type"] is not none or "complex_promote" in op or "data_transform" in op%}
  protected:
-    {% filter indent(2, True)%}
+    {% if kernel["data_type"] is not none or "complex_promote" in op %}
+      {% filter indent(2, True)%}
 {{get_expected_kernel(op)}}
-    {% endfilter %}
-    {%- if "data_transform" in op and op["data_transform"] is not none -%}
-      {%- if "skip_transform" in op["data_transform"] -%}
-        {% filter indent(2, True) %}
+      {% endfilter %}
+    {% endif %}
+  {% endif %}
+  {%- if "data_transform" in op and op["data_transform"] is not none -%}
+      {% filter indent(2, True) %}
+{{get_kernel_for_var(op)}}
+      {% endfilter %}
+  {%- elif "complex_promote" in op and op["complex_promote"] is not none -%}
+      {% filter indent(2, True) %}
 {{get_kernel_for_var(op)}}
         {% endfilter %}
       {%- endif %}
-    {%- endif -%}
-{# TODO(lizhiyu): add the 'support_trans_dtype' #}
-  {% endif %}
 };
 
 DECLARE_INFER_SHAPE_FUNCTOR({{op["op_name"]}}, {{op["op_name"] | to_pascal_case}}InferShapeFunctor,

paddle/fluid/operators/kron_op.cc

@@ -663,6 +663,17 @@
   kernel :
     func : inverse_grad
 
+- backward_op : kron_grad
+  forward : kron (Tensor x, Tensor y) -> Tensor(out)
+  args : (Tensor x, Tensor y, Tensor out_grad)
+  output : Tensor(x_grad), Tensor(y_grad)
+  infer_meta :
+    func : GeneralBinaryGradInferMeta
+    param : [x, y]
+  kernel :
+    func : kron_grad
+    data_type : out_grad
+
 - backward_op : kthvalue_grad
   forward : kthvalue(Tensor x, int k, int axis, bool keepdim) -> Tensor(out), Tensor(indices)
   args : (Tensor x, Tensor indices, Tensor out_grad, int k, int axis, bool keepdim)

paddle/phi/api/yaml/backward.yaml

@@ -636,17 +636,6 @@
     func : kldiv_loss_grad
   no_need_buffer : x
 
-- backward_op : kron_grad
-  forward : kron (Tensor x, Tensor y) -> Tensor(out)
-  args : (Tensor x, Tensor y, Tensor out_grad)
-  output : Tensor(x_grad), Tensor(y_grad)
-  infer_meta :
-    func : GeneralBinaryGradInferMeta
-    param : [x, y]
-  kernel :
-    func : kron_grad
-    data_type : out_grad
-
 - backward_op : layer_norm_grad
   forward : layer_norm (Tensor x, Tensor scale, Tensor bias, float epsilon, int begin_norm_axis) -> Tensor(out), Tensor(mean), Tensor(variance)
   args : (Tensor x,  Tensor scale, Tensor bias, Tensor mean, Tensor variance, Tensor out_grad, float epsilon, int begin_norm_axis)

paddle/phi/api/yaml/legacy_backward.yaml

@@ -921,15 +921,6 @@
     data_type : x
   backward : kldiv_loss_grad
 
-- op : kron
-  args : (Tensor x, Tensor y)
-  output : Tensor
-  infer_meta :
-    func : KronInferMeta
-  kernel :
-    func : kron
-  backward : kron_grad
-
 - op : lamb_
   args : (Tensor param, Tensor grad, Tensor learning_rate, Tensor moment1, Tensor moment2, Tensor beta1_pow, Tensor beta2_pow, Tensor master_param, Tensor skip_update, float weight_decay, float beta1, float beta2, float epsilon, bool multi_precision)
   output : Tensor(param_out), Tensor(moment1_out), Tensor(moment2_out), Tensor(beta1_pow_out), Tensor(beta2_pow_out), Tensor(master_param_outs)

paddle/phi/api/yaml/legacy_ops.yaml

@@ -828,6 +828,14 @@
   outputs :
     out : Out
 
+- op : kron
+  backward : kron_grad
+  inputs :
+    {x : X, y : Y}
+  outputs :
+    {out : Out}
+  complex_promote : [X, Y]
+
 - op : kthvalue
   inputs :
     x : X

paddle/phi/api/yaml/op_compat.yaml

@@ -653,6 +653,15 @@
     func : isnan {dense -> dense},
            isnan_sr {selected_rows -> selected_rows}
 
+- op : kron
+  args : (Tensor x, Tensor y)
+  output : Tensor
+  infer_meta :
+    func : KronInferMeta
+  kernel :
+    func : kron
+  backward : kron_grad
+
 - op : kthvalue
   args : (Tensor x, int k = 1, int axis = -1, bool keepdim = false)
   output : Tensor(out), Tensor(indices)