support optional and invoke (#57084)

Charles-hit · web-flow · commit 430a6572daa2 · 2023-09-09T03:19:24.000+08:00
diff --git a/paddle/fluid/ir/dialect/op_generator/op_interface_gen.py b/paddle/fluid/ir/dialect/op_generator/op_interface_gen.py
@@ -62,7 +62,7 @@
         res[i].resize(tensor_res[i].size());
         for (size_t j = 0; j < tensor_res[i].size(); ++j) {
             if(tensor_res[i][j].defined()){
-                res[i][j] = std::static_pointer_cast<primitive::LazyTensor>(tensor_res[i][j].impl())->getValue().dyn_cast<ir::OpResult>();
+                res[i][j] = std::static_pointer_cast<primitive::LazyTensor>(tensor_res[i][j].impl())->value().dyn_cast<ir::OpResult>();
             }
         }
     }"""
diff --git a/paddle/fluid/ir/dialect/paddle_dialect/ir/pd_manual_op_vjp.cc b/paddle/fluid/ir/dialect/paddle_dialect/ir/pd_manual_op_vjp.cc
@@ -45,7 +45,7 @@ std::vector<std::vector<ir::OpResult>> SumOp::Vjp(
   if (tensor_res[0][0].defined()) {
     res[0][0] =
         std::static_pointer_cast<primitive::LazyTensor>(tensor_res[0][0].impl())
-            ->getValue()
+            ->value()
             .dyn_cast<ir::OpResult>();
   }
   return res;
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -16,7 +16,6 @@
 import hashlib
 import pathlib
 import sys
-from typing import Dict, List
 
 import jinja2
 import yaml
@@ -28,6 +27,7 @@
 )
 import filters as op_gen_filters
 import tests_utils as op_gen_tests
+from parse_utils import to_named_dict
 
 # import from paddle/fluid/ir/dialect/op_generator/api_gen.py
 sys.path.append(
@@ -62,6 +62,9 @@
     'slice_grad',
     'transpose_grad',
     'dropout_grad',
+    'cast_grad',
+    'slice_double_grad',
+    'layer_norm_grad',
 ]
 VJP_COMPS = ['divide_grad', 'sum_grad', 'gelu_grad']
 BACKENDS = [
@@ -130,6 +133,8 @@
     'scatter',
     'scatter_nd_add',
     'dropout_grad',
+    'slice',
+    'layer_norm_grad',
 ]
 
 
@@ -219,21 +224,6 @@ def save(content: str, path: pathlib.Path):
             print(f"Generate source file {path}")
 
 
-def to_compat_dict(items: List[Dict]) -> Dict[str, Dict]:
-    compat_dict = {}
-    for item in items:
-        name = item["op"]
-        compat_dict[name] = item
-    return compat_dict
-
-
-def to_apis_dict(apis):
-    apis_dict = {}
-    for api in apis:
-        apis_dict[api['name']] = api
-    return apis_dict
-
-
 def get_inplace_api(apis):
     inplace_apis = []
     for api in apis:
@@ -271,7 +261,7 @@ def extend_compat_info(apis, compats):
                 attr['typename']
             ) or op_gen_tests.is_intarray(attr['typename']):
                 attr["support_tensor"] = False
-    apis_dict = to_apis_dict(apis)
+    apis_dict = to_named_dict(apis)
     for compat_item in compats:
         fwd_op_name = compat_item["op"]
         if fwd_op_name not in apis_dict:
@@ -322,6 +312,31 @@ def extend_compat_info(apis, compats):
     return apis
 
 
+def process_backward_invoke_info(apis):
+    apis_dict = to_named_dict(apis)
+    for api in apis:
+        if api['is_fwd']:
+            continue
+        if 'invoke' in api and api['invoke']['func'] in apis_dict:
+            args = api['invoke']['args'].split(',')
+            args = [arg.strip() for arg in args]
+            attrs_dict = to_named_dict(api['attrs'])
+            inputs_dict = to_named_dict(api['inputs'])
+            arg_inputs = []
+            arg_attrs = []
+            for arg in args:
+                if arg in inputs_dict:
+                    arg_inputs.append(arg)
+                elif arg in attrs_dict and attrs_dict[arg].get(
+                    "support_tensor", False
+                ):
+                    arg_inputs.append(arg + '_')
+                else:
+                    arg_attrs.append(arg)
+            args = arg_inputs + arg_attrs
+            api['invoke']['args'] = ', '.join(args)
+
+
 def gen(
     prim_path: pathlib.Path,
     fwd_path: pathlib.Path,
@@ -369,6 +384,7 @@ def gen(
     ]
     apis = extend_compat_info(apis, compats)
     apis = apis + get_inplace_api(apis)
+    process_backward_invoke_info(apis)
     render(
         templates_dir,
         destination_dir,
diff --git a/paddle/fluid/primitive/codegen/templates/backend/generated/generated_backend.h.j2 b/paddle/fluid/primitive/codegen/templates/backend/generated/generated_backend.h.j2
@@ -7,6 +7,7 @@
 #include <vector>
 
 #include "paddle/phi/api/include/tensor.h"
+#include "paddle/utils/optional.h"
 
 
 namespace paddle {
diff --git a/paddle/fluid/primitive/codegen/templates/backend/generated/generated_static_backend.cc.j2 b/paddle/fluid/primitive/codegen/templates/backend/generated/generated_static_backend.cc.j2
@@ -17,22 +17,69 @@ template <>
 {{common.ret(outputs)}} {{name}}<LazyTensor>({{common.params(inputs, attrs, mutable_attribute_as_inputs, False)}})
 {%- endmacro -%}
 
-{% macro body(name, inputs, outputs, attrs, mutable_attribute_as_inputs=False) %}
-  {%- set output_names = [] -%}
-  {%- for o in outputs -%} {%- do output_names.append(o.name) -%} {%-endfor-%}
+{%- macro prepare_ir_api_inputs(inputs)-%}
   {%- for input in inputs -%}
-    {% if input.typename=='Tensor[]' %}
-  std::vector<ir::OpResult> {{input.name}}_res({{input.name}}.size());
-  std::transform({{input.name}}.begin(), {{input.name}}.end(), {{input.name}}_res.begin(), [](const Tensor& t) {
-    return std::static_pointer_cast<LazyTensor>(t.impl())->getValue().dyn_cast<ir::OpResult>();
+    {% if input.typename=='Tensor[]' and not input.optional %}
+std::vector<ir::OpResult> {{input.name}}_res({{input.name}}.size());
+std::transform({{input.name}}.begin(), {{input.name}}.end(), {{input.name}}_res.begin(), [](const Tensor& t) {
+  return std::static_pointer_cast<LazyTensor>(t.impl())->value().dyn_cast<ir::OpResult>();
+});
+    {% elif input.typename=='Tensor[]' and input.optional %}
+std::vector<ir::OpResult> {{input.name}}_res({{input.name}}.size());
+if({{input.name}}) {
+  std::transform({{input.name}}.get().begin(), {{input.name}}.get().end(), {{input.name}}_res.begin(), [](const Tensor& t) {
+    return std::static_pointer_cast<LazyTensor>(t.impl())->value().dyn_cast<ir::OpResult>();
   });
+}
+    {% elif input.typename=='Tensor' and not input.optional %}
+ir::OpResult {{input.name}}_res = std::static_pointer_cast<LazyTensor>({{input.name}}.impl())->value().dyn_cast<ir::OpResult>();
     {% else %}
-  ir::OpResult {{input.name}}_res = std::static_pointer_cast<LazyTensor>({{input.name}}.impl())->getValue().dyn_cast<ir::OpResult>();
+ir::OpResult {{input.name}}_res;
+if({{input.name}}) {
+  {{input.name}}_res = std::static_pointer_cast<LazyTensor>({{input.name}}.get().impl())->value().dyn_cast<ir::OpResult>();
+}
     {% endif %}
   {% endfor %}
-  {%- for attr in attrs -%}
+{%- endmacro -%}
+
+{%- macro get_static_backend_outputs(outputs)-%}
+  {%- if outputs|length == 1 -%}
+    {%- if outputs[0].typename == 'Tensor' -%}
+Tensor {{outputs[0].name}}(std::make_shared<LazyTensor>(op_res));
+return {{outputs[0].name}};
+    {%- elif outputs[0].typename == 'Tensor[]' -%}
+std::vector<Tensor> {{outputs[0].name}}(op_res.size());
+std::transform(op_res.begin(), op_res.end(), {{outputs[0].name}}.begin(), [](const ir::OpResult& res) {
+return Tensor(std::make_shared<LazyTensor>(res));
+  });
+return {{outputs[0].name}};
+    {%- else -%} {#- render nothing -#}
+    {%- endif -%}
+  {%- elif outputs|length > 1 -%}
+    {%- for i in range(outputs|length) %}
+auto op_res_{{i}} = std::get<{{i}}>(op_res);
+      {% if outputs[i].typename == 'Tensor' %}
+Tensor {{outputs[i].name}}(std::make_shared<LazyTensor>(op_res_{{i}}));
+      {% elif outputs[i].typename == 'Tensor[]' %}
+std::vector<Tensor> {{outputs[i].name}}(op_res_{{i}}.size());
+std::transform(op_res_{{i}}.begin(), op_res_{{i}}.end(), {{outputs[i].name}}.begin(), [](const ir::OpResult& res) {
+return Tensor(std::make_shared<LazyTensor>(res));
+  });
+      {% else %} {#- render nothing -#}
+      {% endif %}
+    {% endfor -%}
+return std::make_tuple({%- for i in range(outputs|length) -%}{{outputs[i].name}}{%- if i!=outputs|length - 1 -%}, {% endif -%}{%- endfor -%});
+  {%- else -%} {#- render nothing -#}
+  {%- endif -%}
+{%- endmacro -%}
+
+{% macro body(name, inputs, outputs, attrs, mutable_attribute_as_inputs=False) %}
+  {%- set output_names = [] -%}
+  {%- for o in outputs -%} {%- do output_names.append(o.name) -%} {%-endfor-%}
+{{prepare_ir_api_inputs(inputs)}}
+  {%- for attr in attrs %}
     {% if mutable_attribute_as_inputs and attr is mutable_attribute %}
-  ir::OpResult {{attr.name}}_res = std::static_pointer_cast<LazyTensor>({{attr.name~'_'}}.impl())->getValue().dyn_cast<ir::OpResult>();
+ir::OpResult {{attr.name}}_res = std::static_pointer_cast<LazyTensor>({{attr.name~'_'}}.impl())->value().dyn_cast<ir::OpResult>();
     {% endif %}
   {% endfor %}
   {%- set input_names = [] -%}
@@ -52,48 +99,25 @@ template <>
       {%- do attr_names.append(common.phi2ir_attr(i)) -%} 
     {%- endif -%}
   {% endfor %}
-  auto op_res = paddle::dialect::{{name}}({{common.args(input_names, attr_names)}});
-  {% if outputs|length == 1 %}
-    {% if outputs[0].typename == 'Tensor' %}
-  Tensor {{outputs[0].name}}(std::make_shared<LazyTensor>(op_res));
-  return {{outputs[0].name}};
-    {% elif outputs[0].typename == 'Tensor[]' %}
-  std::vector<Tensor> {{outputs[0].name}}(op_res.size());
-  std::transform(op_res.begin(), op_res.end(), {{outputs[0].name}}.begin(), [](const ir::OpResult& res) {
-    return Tensor(std::make_shared<LazyTensor>(res));
-  });
-  return {{outputs[0].name}};
-    {% else %} {#- render nothing -#}
-    {% endif %}
-  {% elif outputs|length > 1 %}
-    {% for i in range(outputs|length) %}
-  auto op_res_{{i}} = std::get<{{i}}>(op_res);
-      {% if outputs[i].typename == 'Tensor' %}
-  Tensor {{outputs[i].name}}(std::make_shared<LazyTensor>(op_res_{{i}}));
-      {% elif outputs[i].typename == 'Tensor[]' %}
-  std::vector<Tensor> {{outputs[i].name}}(op_res_{{i}}.size());
-  std::transform(op_res_{{i}}.begin(), op_res_{{i}}.end(), {{outputs[i].name}}.begin(), [](const ir::OpResult& res) {
-    return Tensor(std::make_shared<LazyTensor>(res));
-  });
-      {% else %} {#- render nothing -#}
-      {% endif %}
-    {% endfor %}
-  return std::make_tuple({% for i in range(outputs|length) %}{{outputs[i].name}}{%- if i!=outputs|length - 1 -%}, {% endif %}{% endfor %});
-  {% else %} {#- render nothing -#}
-  {% endif %}
-{% endmacro %}
+auto op_res = paddle::dialect::{{name}}({{common.args(input_names, attr_names)}});
+{{get_static_backend_outputs(outputs)}}
+{%- endmacro %}
 
 
 {% for api in apis %}
 {% if api.name in backend_white_list %}
   {% set api_outputs = api.outputs | trip_intermediate %}
 {{sig(api.name, api.inputs, api_outputs, api.attrs)}} {
+  {% filter indent(2, True) %}
 {{body(api.name, api.inputs, api_outputs, api.attrs)}} 
+  {% endfilter %}
 }
 
   {% if api.attrs is exist_mutable_attribute %}
 {{sig(api.name, api.inputs, api_outputs, api.attrs, True)}} {
+  {% filter indent(2, True) %}
 {{body(api.name, api.inputs, api_outputs, api.attrs, True)}} 
+  {% endfilter %}
 }
 
   {% endif %}
diff --git a/paddle/fluid/primitive/codegen/templates/rule/vjp/generated/generated_vjp.cc.j2 b/paddle/fluid/primitive/codegen/templates/rule/vjp/generated/generated_vjp.cc.j2
@@ -11,6 +11,7 @@
 #include "paddle/fluid/primitive/utils/utils.h"
 #include "paddle/ir/core/operation.h"
 #include "paddle/phi/core/flags.h"
+#include "paddle/utils/optional.h"
 
 PHI_DECLARE_string(tensor_operants_mode);
 
@@ -33,14 +34,14 @@ if (paddle::prim::StaticCompositeContext::Instance().IsBwdPrimEnabled()) {
 }
   {% else %}
 {{body_unprim(api)}}
-  {% endif %}
+  {%- endif %}
 return vjp_res;
-{% endmacro %}
+{%- endmacro -%}
 
 {% macro get_mutable_attribute(attrs, api_name) %}
   {% for i in attrs %}
     {%- if i is mutable_attribute -%}
-auto* {{i.name}}_define_op = std::static_pointer_cast<primitive::LazyTensor>({{i.name~'_'}}.impl())->getValue().dyn_cast<ir::OpResult>().GetDefiningOp();
+auto* {{i.name}}_define_op = std::static_pointer_cast<primitive::LazyTensor>({{i.name~'_'}}.impl())->value().dyn_cast<ir::OpResult>().GetDefiningOp();
       {% if i.typename is scalar %}
 if({{i.name}}_define_op->name() != "pd.full") {
   PADDLE_THROW(platform::errors::Unimplemented(
@@ -62,6 +63,7 @@ auto {{i.name}} = {{i.name}}_define_op->attribute("value").dyn_cast<paddle::dial
 
 {% macro body_unprim(api) %}
   {%- set input_names=[] -%}
+  {%- for api in apis -%} {%- do api_map.update({api.name: api}) -%} {%- endfor -%}
   {%- for i in api.inputs -%} {%- do input_names.append(i.name) -%} {%- endfor -%}
   {%- set attr_names=[] -%}
   {%- for i in api.attrs -%} 
@@ -71,29 +73,29 @@ auto {{i.name}} = {{i.name}}_define_op->attribute("value").dyn_cast<paddle::dial
       {%- do attr_names.append(i.name) -%} 
     {%- endif -%}
   {%- endfor %}
+  {% if 'invoke' in api and api.invoke.func in api_map %}
+auto op_res = backend::{{api.invoke.func}}<LazyTensor>({{api.invoke.args}});
+  {% else %}
 auto op_res = backend::{{api.name}}<LazyTensor>({{common.args(input_names, attr_names)}});
+  {% endif %}
   {% set outputs = api.outputs|trip_intermediate %} {#- ignore intermediate output -#}
   {% if outputs|length > 1 %}
     {% for i in range(outputs|length) %}
-auto out{{i}} = std::get<{{i}}>(op_res);
       {% if outputs[i].typename=='Tensor' %}
-vjp_res[{{i}}][0] = !stop_gradients[{{i}}][0] ? out{{i}} : vjp_res[{{i}}][0];
+vjp_res[{{i}}][0] = std::get<{{i}}>(op_res);
       {% else %}
-for (size_t i=0; i< stop_gradients[{{i}}].size(); i++ ) {
-    vjp_res[{{i}}][i] =  !stop_gradients[{{i}}][i] ? out{{i}}[i] : vjp_res[{{i}}][i];
-}
+vjp_res[{{i}}] = std::get<{{i}}>(op_res);
       {% endif %}
     {% endfor %}
   {% elif outputs|length == 1 %}
     {% if outputs[0].typename=='Tensor' %}
-vjp_res[0][0] = !stop_gradients[0][0] ? op_res : vjp_res[0][0];
+vjp_res[0][0] = op_res;
     {% else %}
-for (size_t i=0; i< stop_gradients[0].size(); i++ ) {
-  vjp_res[0][i] =  !stop_gradients[0][i] ? op_res[i] : vjp_res[0][i];
-}
+vjp_res[0] = op_res;
     {% endif %}
   {% else %} {#- render nothing -#}
   {% endif %}
+vjp_res = ConstructVjpResultByStopGradients(vjp_res, stop_gradients);
 {% endmacro %}
 
 {% macro body_prim(api) %}
@@ -120,7 +122,7 @@ details::{{api.composite.func_name}}<LazyTensor>({{api.composite.func_args}});
 {{sig(api.name, backward_api.name, backward_api.inputs, backward_api.attrs, backward_api.outputs)}} {
     {% filter indent(2, True) %}
 {{body(backward_api)}}
-    {% endfilter %}
+    {% endfilter -%}
 }
 
   {% endif %}
diff --git a/paddle/fluid/primitive/type/lazy_tensor.h b/paddle/fluid/primitive/type/lazy_tensor.h
@@ -41,12 +41,14 @@ class LazyTensor : public phi::ExtendedTensor,
         value_.type().dyn_cast<paddle::dialect::DenseTensorType>().dtype());
   }
 
-  ir::Value getValue() const { return value_; }
+  ir::Value value() const { return value_; }
 
   const phi::Place& place() const override { return place_; }
 
   bool initialized() const override { return value_.impl() != nullptr; }
 
+  void set_empty_type() { value_.set_type(ir::Type()); }
+
  private:
   ir::Value value_;
   mutable phi::DDim dims_;
diff --git a/paddle/fluid/primitive/utils/static_utils.cc b/paddle/fluid/primitive/utils/static_utils.cc
@@ -21,5 +21,23 @@ void set_output<LazyTensor>(const paddle::Tensor& x_tmp, paddle::Tensor* x) {
   x->set_impl(x_tmp.impl());
 }
 
+std::vector<std::vector<Tensor>> ConstructVjpResultByStopGradients(
+    const std::vector<std::vector<Tensor>>& outputs,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  std::vector<std::vector<Tensor>> vjp_results(outputs.size());
+  for (size_t i = 0; i < outputs.size(); ++i) {
+    vjp_results[i].reserve(outputs[i].size());
+    for (size_t j = 0; j < outputs[i].size(); ++j) {
+      if (stop_gradients[i][j]) {
+        // Use Tensor's impl is nullptr to indicate it has no gradient
+        vjp_results[i].emplace_back(Tensor());
+      } else {
+        vjp_results[i].emplace_back(outputs[i][j]);
+      }
+    }
+  }
+  return vjp_results;
+}
+
 }  // namespace primitive
 }  // namespace paddle
diff --git a/paddle/fluid/primitive/utils/utils.h b/paddle/fluid/primitive/utils/utils.h
diff --git a/test/ir/new_ir/test_ir_backward.py b/test/ir/new_ir/test_ir_backward.py

Original file line number	Diff line number	Diff line change
`@@ -62,7 +62,7 @@`
`62`	`62`	`res[i].resize(tensor_res[i].size());`
`63`	`63`	`for (size_t j = 0; j < tensor_res[i].size(); ++j) {`
`64`	`64`	`if(tensor_res[i][j].defined()){`
`65`		`- res[i][j] = std::static_pointer_cast<primitive::LazyTensor>(tensor_res[i][j].impl())->getValue().dyn_cast<ir::OpResult>();`
	`65`	`+ res[i][j] = std::static_pointer_cast<primitive::LazyTensor>(tensor_res[i][j].impl())->value().dyn_cast<ir::OpResult>();`
`66`	`66`	`}`
`67`	`67`	`}`
`68`	`68`	`}"""`
Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ std::vector<std::vector<ir::OpResult>> SumOp::Vjp(`
`45`	`45`	`if (tensor_res[0][0].defined()) {`
`46`	`46`	`res[0][0] =`
`47`	`47`	`std::static_pointer_cast<primitive::LazyTensor>(tensor_res[0][0].impl())`
`48`		`- ->getValue()`
	`48`	`+ ->value()`
`49`	`49`	`.dyn_cast<ir::OpResult>();`
`50`	`50`	`}`
`51`	`51`	`return res;`