[PRIM][PIR]Support optional and more vjp gen (#57693)

Charles-hit · 0x45f · web-flow · commit 09321420ce60 · 2023-09-25T21:22:41.000+08:00
* Support optional input and output for pir api

* Polish gen code

* Fix py3 ci compile error

* Fix code

* Fix error

* Fix windows ci error

* support optional in vjp and add more ops vjp gen

---------

Co-authored-by: 0x45f &lt;wangzhen45@baidu.com&gt;
diff --git a/paddle/fluid/operators/generator/filters.py b/paddle/fluid/operators/generator/filters.py
@@ -24,6 +24,7 @@
     input_types_map,
     opmaker_attr_types_map,
     optional_input_types_map,
+    optional_output_type_map,
     output_type_map,
     phi_attr_types_map,
     sr_output_types_map,
@@ -154,7 +155,9 @@ def delete_last_underline(op_name):
 
 
 # ------------------------------ output  ----------------------------------
-def to_paddle_output_type(s):
+def to_paddle_output_type(s, optional=False):
+    if optional:
+        return optional_output_type_map[s]
     return output_type_map[s]
 
 
diff --git a/paddle/fluid/operators/generator/type_mapping.py b/paddle/fluid/operators/generator/type_mapping.py
@@ -85,6 +85,11 @@
     'SelectedRows': 'SelectedRows',
 }
 
+optional_output_type_map = {
+    'Tensor': 'const paddle::optional<Tensor>&',
+    'Tensor[]': 'const paddle::optional<std::vector<Tensor>>&',
+}
+
 # ------------------------------ phi attr ------------------------------
 phi_attr_types_map = attr_types_map.copy()
 phi_attr_types_map.update(
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -78,8 +78,60 @@
     'unsqueeze_grad',
     'poisson_grad',
     'gumbel_softmax_grad',
-    'squeeze_grad',
-    'unsqueeze_grad',
+    'conv2d_grad',
+    'depthwise_conv2d_grad',
+    'sqrt_grad',
+    'flatten_grad',
+    'relu_grad',
+    'abs_grad',
+    'log_grad',
+    'clip_grad',
+    'ceil_grad',
+    'frobenius_norm_grad',
+    'p_norm_grad',
+    'maximum_grad',
+    'argsort_grad',
+    'min_grad',
+    'batch_norm_grad',
+    'max_pool2d_with_index_grad',
+    'pool2d_grad',
+    'minimum_grad',
+    'prod_grad',
+    'round_grad',
+    'sin_grad',
+    'cos_grad',
+    'dot_grad',
+    'floor_grad',
+    'topk_grad',
+    'square_grad',
+    'gather_grad',
+    'label_smooth_grad',
+    'cross_entropy_with_softmax_grad',
+    'mean_all_grad',
+    'cumsum_grad',
+    'linear_interp_grad',
+    'bilinear_interp_grad',
+    'trilinear_interp_grad',
+    'nearest_interp_grad',
+    'bicubic_interp_grad',
+    'assign_grad',
+    'assign_out__grad',
+    'real_grad',
+    'flip_grad',
+    'softmax_grad',
+    'expand_grad',
+    'conv2d_transpose_grad',
+    'depthwise_conv2d_transpose_grad',
+    'sigmoid_grad',
+    'pad_grad',
+    'pad3d_grad',
+    'einsum_grad',
+    'leaky_relu_grad',
+    'log10_grad',
+    'conv3d_grad',
+    'solve_grad',
+    'diag_grad',
+    'trace_grad',
 ]
 
 
@@ -183,6 +235,58 @@
     'stack_grad',
     'squeeze_grad',
     'unsqueeze_grad',
+    'conv2d_grad',
+    'depthwise_conv2d_grad',
+    'sqrt_grad',
+    'flatten_grad',
+    'relu_grad',
+    'abs_grad',
+    'log_grad',
+    'clip_grad',
+    'ceil_grad',
+    'frobenius_norm_grad',
+    'p_norm_grad',
+    'maximum_grad',
+    'argsort_grad',
+    'min_grad',
+    'batch_norm_grad',
+    'max_pool2d_with_index_grad',
+    'pool2d_grad',
+    'minimum_grad',
+    'prod_grad',
+    'round_grad',
+    'sin_grad',
+    'cos_grad',
+    'dot_grad',
+    'floor_grad',
+    'topk_grad',
+    'square_grad',
+    'gather_grad',
+    'label_smooth_grad',
+    'cross_entropy_with_softmax_grad',
+    'mean_all_grad',
+    'cumsum_grad',
+    'linear_interp_grad',
+    'bilinear_interp_grad',
+    'trilinear_interp_grad',
+    'nearest_interp_grad',
+    'bicubic_interp_grad',
+    'assign_out__grad',
+    'real_grad',
+    'softmax_grad',
+    'conv2d_transpose_grad',
+    'depthwise_conv2d_transpose_grad',
+    'sigmoid_grad',
+    'pad_grad',
+    'pad3d_grad',
+    'einsum_grad',
+    'leaky_relu_grad',
+    'log10_grad',
+    'conv3d_grad',
+    'solve_grad',
+    'diag_grad',
+    'trace_grad',
+    'flip',
 ]
 
 
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
@@ -21,11 +21,17 @@ using DataType = phi::DataType;
 
 {% for api in apis %}
   {%- if api.name in backend_white_list -%}
+  {% set inplace_map = {} %}
+  {% if 'inplace' in api and api.inplace != None %}
+    {% for source, target in api.inplace.items() %}
+      {% do inplace_map.update({source: target}) %}
+    {% endfor %}
+  {% endif %}
   {% if api.attrs is exist_mutable_attribute %}
-{{common.sig(api.name, api.inputs, api.outputs|trip_intermediate , api.attrs, True, True)}};
+{{common.sig(api.name, api.inputs, api.outputs|trip_intermediate , api.attrs, inplace_map, True, True)}};
 
   {% endif %}
-{{common.sig(api.name, api.inputs, api.outputs|trip_intermediate , api.attrs, False, True)}};
+{{common.sig(api.name, api.inputs, api.outputs|trip_intermediate , api.attrs, inplace_map, False, True)}};
 
   {% endif %}
 {% endfor %}
diff --git a/paddle/fluid/primitive/codegen/templates/backend/generated/generated_eager_backend.cc.j2 b/paddle/fluid/primitive/codegen/templates/backend/generated/generated_eager_backend.cc.j2
@@ -16,9 +16,9 @@ namespace backend {
   {{common.sequence('', '', ', ', attrs)}}
 {%- endmacro -%}
 
-{%- macro sig(name, inputs, attrs, outputs) -%}
+{%- macro sig(name, inputs, attrs, outputs, inplace_map) -%}
 template <>
-{{common.ret(outputs)}} {{name}}<Tensor>({{common.params(inputs, attrs, False)}})
+{{common.ret(outputs, inplace_map)}} {{name}}<Tensor>({{common.params(inputs, attrs, False)}})
 {%- endmacro -%}
 
 {% macro body(name, inputs, attrs, outputs) %}
@@ -35,7 +35,13 @@ return ::{{name}}_ad_func({{common.args(input_names, attr_names)}});
 
 {% for api in apis %}
   {%- if api.is_prim and api.name in backend_white_list -%}
-{{sig(api.name, api.inputs, api.attrs, api.outputs | trip_intermediate)}} {
+  {% set inplace_map = {} %}
+  {% if 'inplace' in api and api.inplace != None %}
+    {% for source, target in api.inplace.items() %}
+      {% do inplace_map.update({source: target}) %}
+    {% endfor %}
+  {% endif %}
+{{sig(api.name, api.inputs, api.attrs, api.outputs | trip_intermediate, inplace_map)}} {
 {{body(api.name, api.inputs, api.attrs, api.outputs | trip_intermediate)}} 
 }
 
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
@@ -12,9 +12,9 @@ namespace backend {
 
 using LazyTensor = paddle::primitive::LazyTensor;
 
-{%- macro sig(name, inputs, outputs, attrs, mutable_attribute_as_inputs=False) -%}
+{%- macro sig(name, inputs, outputs, attrs, inplace_map, mutable_attribute_as_inputs=False) -%}
 template <>
-{{common.ret(outputs)}} {{name}}<LazyTensor>({{common.params(inputs, attrs, mutable_attribute_as_inputs, False)}})
+{{common.ret(outputs, inplace_map)}} {{name}}<LazyTensor>({{common.params(inputs, attrs, mutable_attribute_as_inputs, False)}})
 {%- endmacro -%}
 
 {%- macro prepare_ir_api_inputs(inputs)-%}
@@ -25,46 +25,80 @@ std::transform({{input.name}}.begin(), {{input.name}}.end(), {{input.name}}_res.
   return std::static_pointer_cast<LazyTensor>(t.impl())->value();
 });
     {% elif input.typename=='Tensor[]' and input.optional %}
-std::vector<pir::Value> {{input.name}}_res({{input.name}}.size());
+paddle::optional<std::vector<pir::Value>> {{input.name}}_res;
 if({{input.name}}) {
-  std::transform({{input.name}}.get().begin(), {{input.name}}.get().end(), {{input.name}}_res.begin(), [](const Tensor& t) {
+  std::vector<pir::Value> {{input.name}}_res_inner({{input.name}}.get().size());
+  std::transform({{input.name}}.get().begin(), {{input.name}}.get().end(), {{input.name}}_res_inner.begin(), [](const Tensor& t) {
     return std::static_pointer_cast<LazyTensor>(t.impl())->value();
   });
+  {{input.name}}_res = paddle::make_optional<std::vector<pir::Value>>({{input.name}}_res_inner);
 }
     {% elif input.typename=='Tensor' and not input.optional %}
 pir::Value {{input.name}}_res = std::static_pointer_cast<LazyTensor>({{input.name}}.impl())->value();
     {% else %}
-pir::Value {{input.name}}_res;
+paddle::optional<pir::Value> {{input.name}}_res;
 if({{input.name}}) {
-  {{input.name}}_res = std::static_pointer_cast<LazyTensor>({{input.name}}.get().impl())->value();
+  pir::Value {{input.name}}_res_inner;
+  {{input.name}}_res_inner = std::static_pointer_cast<LazyTensor>({{input.name}}.get().impl())->value();
+  {{input.name}}_res = paddle::make_optional<pir::Value>({{input.name}}_res_inner);
 }
     {% endif %}
   {% endfor %}
 {%- endmacro -%}
 
 {%- macro get_static_backend_outputs(outputs)-%}
   {%- if outputs|length == 1 -%}
-    {%- if outputs[0].typename == 'Tensor' -%}
+    {%- if outputs[0].typename == 'Tensor' and not outputs[0].optional-%}
 Tensor {{outputs[0].name}}(std::make_shared<LazyTensor>(op_res));
 return {{outputs[0].name}};
-    {%- elif outputs[0].typename == 'Tensor[]' -%}
+    {%- elif outputs[0].typename == 'Tensor' and outputs[0].optional -%}
+paddle::optional<Tensor> {{outputs[0].name}};
+if(op_res){
+  {{outputs[0].name}} = paddle::make_optional<Tensor>(Tensor(std::make_shared<LazyTensor>(op_res.get()));
+}
+return {{outputs[0].name}};
+    {%- elif outputs[0].typename == 'Tensor[]' and not outputs[0].optional -%}
 std::vector<Tensor> {{outputs[0].name}}(op_res.size());
 std::transform(op_res.begin(), op_res.end(), {{outputs[0].name}}.begin(), [](const pir::OpResult& res) {
 return Tensor(std::make_shared<LazyTensor>(res));
   });
+return {{outputs[0].name}};
+    {%- elif outputs[0].typename == 'Tensor[]' and outputs[0].optional -%}
+paddle::optional<std::vector<Tensor>> {{outputs[0].name}};
+if({{op_res}}) {
+  std::vector<pir::Value> {{outputs[0].name}}_inner(op_res.get().size());
+  std::transform(op_res.get().begin(), op_res.get().end(), {{outputs[0].name}}_inner.begin(), [](const pir::OpResult& res) {
+    return Tensor(std::make_shared<LazyTensor>(res));
+  });
+  {{outputs[0].name}} = paddle::make_optional<std::vector<Tensor>>({{outputs[0].name}}_inner);
+}
 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' %}
+      {% if outputs[i].typename == 'Tensor' and not outputs[i].optional %}
 Tensor {{outputs[i].name}}(std::make_shared<LazyTensor>(op_res_{{i}}));
-      {% elif outputs[i].typename == 'Tensor[]' %}
+      {% elif outputs[i].typename == 'Tensor' and  outputs[i].optional %}
+paddle::optional<Tensor> {{outputs[i].name}};
+if(op_res_{{i}}){
+  {{outputs[i].name}} = paddle::make_optional<Tensor>(Tensor(std::make_shared<LazyTensor>(op_res_{{i}}.get())));
+}
+      {% elif outputs[i].typename == 'Tensor[]' and not outputs[i].optional %}
 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 pir::OpResult& res) {
 return Tensor(std::make_shared<LazyTensor>(res));
   });
+      {% elif outputs[i].typename == 'Tensor[]' and outputs[i].optional %}
+paddle::optional<std::vector<Tensor>> {{outputs[i].name}};
+if(op_res_{{i}}){
+  std::vector<Tensor> {{outputs[i].name}}_inner(op_res_{{i}}.get().size());
+  std::transform(op_res_{{i}}.get().begin(), op_res_{{i}}.get().end(), {{outputs[i].name}}_inner.begin(), [](const pir::OpResult& res) {
+    return Tensor(std::make_shared<LazyTensor>(res));
+  });
+  {{outputs[i].name}} = paddle::make_optional<std::vector<Tensor>>({{outputs[i].name}}_inner);
+}
       {% else %} {#- render nothing -#}
       {% endif %}
     {% endfor -%}
@@ -107,14 +141,20 @@ auto op_res = paddle::dialect::{{name}}({{common.args(input_names, attr_names)}}
 {% 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)}} {
+  {% set inplace_map = {} %}
+  {% if 'inplace' in api and api.inplace != None %}
+    {% for source, target in api.inplace.items() %}
+      {% do inplace_map.update({source: target}) %}
+    {% endfor %}
+  {% endif %}
+{{sig(api.name, api.inputs, api_outputs, api.attrs, inplace_map)}} {
   {% 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)}} {
+{{sig(api.name, api.inputs, api_outputs, api.attrs, inplace_map, True)}} {
   {% filter indent(2, True) %}
 {{body(api.name, api.inputs, api_outputs, api.attrs, True)}} 
   {% endfilter %}
diff --git a/paddle/fluid/primitive/codegen/templates/common.j2 b/paddle/fluid/primitive/codegen/templates/common.j2
@@ -1,6 +1,6 @@
-{%- macro sig(name, inputs, outputs, attrs, mutable_attribute_as_inputs=False, default=False) -%}
+{%- macro sig(name, inputs, outputs, attrs, inplace_map, mutable_attribute_as_inputs=False, default=False) -%}
 template <typename T>
-{{ret(outputs)}} {{name}}({{params(inputs, attrs, mutable_attribute_as_inputs, default)}})
+{{ret(outputs, inplace_map)}} {{name}}({{params(inputs, attrs, mutable_attribute_as_inputs, default)}})
 {%- endmacro %}
 
 
@@ -40,9 +40,9 @@ template <typename T>
 {%- endmacro -%}
 
 
-{%- macro ret(outputs) -%}
+{%- macro ret(outputs, inplace_map) -%}
   {%- set names = [] -%}
-  {%- for i in outputs -%} {%- do names.append(i.typename|to_paddle_output_type) -%} {%- endfor -%}
+  {%- for i in outputs -%} {%- do names.append(i.typename|to_paddle_output_type(i.name in inplace_map and i.optional)) -%} {%- endfor -%}
   {%- if names|length > 1 -%} 
 std::tuple<{{sequence('', '', ', ', names)}}>
   {%- else -%}
diff --git a/paddle/fluid/primitive/codegen/templates/primitive/primitive.h.j2 b/paddle/fluid/primitive/codegen/templates/primitive/primitive.h.j2
@@ -18,7 +18,13 @@ using IntArray = paddle::experimental::IntArray;
   {%- for i in api.inputs -%} {%- do input_names.append(i.name) -%} {%- endfor -%}
   {%- set attr_names = [] -%}
   {%- for i in api.attrs -%} {%- do attr_names.append(i.name) -%} {% endfor %}
-{{common.sig(api.name, api.inputs, api.outputs | trip_intermediate, api.attrs, False, True)}} {
+  {% set inplace_map = {} %}
+  {% if 'inplace' in api and api.inplace != None %}
+    {% for source, target in api.inplace.items() %}
+      {% do inplace_map.update({source: target}) %}
+    {% endfor %}
+  {% endif %}
+{{common.sig(api.name, api.inputs, api.outputs | trip_intermediate, api.attrs, inplace_map, False, True)}} {
     return backend::{{api.name}}<T>({{common.args(input_names, attr_names)}});
 }
 
