paddlepaddle--paddle/paddle/fluid/primitive/codegen/templates/decomp/generated_decomp.j2

{% import "common.j2" as common %}
// Auto Generated by decomp_gen.py, DO NOT EDIT!

#include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
#include "paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h"
#include "paddle/fluid/primitive/base/lazy_tensor.h"
#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/common/int_array.h"
#include "paddle/pir/include/core/builtin_op.h"
#include "paddle/pir/include/core/op_base.h"

namespace paddle {
namespace dialect {
using IntArray = paddle::experimental::IntArray;
{% macro sig(fwd_name, class_name, inputs, attrs, outputs) %}

{% set input_names=[] %}
{% set attr_names=[] %}
{% set output_names=[] %}
{% set output_types=[] %}

std::vector<std::vector<pir::Value>> {{class_name}}::Decomp(pir::Operation* op) {
  VLOG(4) << "Decomp call {{fwd_name}}'s decomp interface begin";

  {{class_name}} op_obj = op->dyn_cast<{{class_name}}>();
  (void)op_obj;

  FLAGS_tensor_operants_mode = "static";

  VLOG(6) << "Decomp Prepare inputs of {{fwd_name}}";

    {% for item in inputs -%}
      {% do input_names.append(item.name) %}
      {% if item.typename == "Tensor" %}  {#- Tensor or Tensor[] #}
        {% if item.optional %}
  paddle::optional<Tensor> {{item.name}};
  if (!IsEmptyValue(op_obj.{{item.name}}())){
      {{item.name}} = paddle::make_optional<Tensor>(Tensor(std::make_shared<primitive::LazyTensor>(op_obj.{{item.name}}())));
  }
        {% else %}
  {{item.typename}} {{item.name}}(std::make_shared<primitive::LazyTensor>(op_obj.{{item.name}}()));
        {% endif %}
      {% elif item.typename == "Tensor[]" %}
        {% if item.optional %}

  paddle::optional<std::vector<Tensor>> {{item.name}};
  if (!IsEmptyValue(op_obj.{{item.name}}())){
      pir::CombineOp combine_op_obj =
          op_obj.{{item.name}}().defining_op()->dyn_cast<pir::CombineOp>();
      std::vector<Tensor> optional_{{item.name}};
      for (size_t idx = 0; idx < combine_op_obj.inputs().size(); idx++) {
          optional_{{item.name}}.emplace_back(
              std::make_shared<primitive::LazyTensor>(combine_op_obj.inputs()[idx]));
      }
      {{item.name}} = paddle::make_optional<std::vector<Tensor>>(optional_{{item.name}});
  }

        {% else %}
  pir::CombineOp combine_op_obj_{{item.name}} =
    op_obj.{{item.name}}().defining_op()->dyn_cast<pir::CombineOp>();
  std::vector<Tensor> {{item.name}};
  for (size_t idx = 0; idx < combine_op_obj_{{item.name}}.inputs().size(); idx++) {
      {{item.name}}.emplace_back(
          std::make_shared<primitive::LazyTensor>(combine_op_obj_{{item.name}}.inputs()[idx]));
  }
        {% endif %}
      {% endif %}
    {% endfor %}

  VLOG(6) << "Decomp prepare attributes of {{fwd_name}}";
    {% if attrs %}
      {% for item in attrs %}
      {% do attr_names.append(item.name) %}
      {% if item.typename.startswith("Scalar") and item.support_tensor %}

  Tensor {{item.name}}_(std::make_shared<primitive::LazyTensor>(op_obj.{{item.name}}()));

  auto* {{item.name}}_define_op =
      std::static_pointer_cast<primitive::LazyTensor>({{item.name}}_.impl())
          ->value()
          .defining_op();
  if ({{item.name}}_define_op->name() != "pd_op.full") {
    return {};
  }
  Scalar {{item.name}} = {{item.name}}_define_op->attribute("value").dyn_cast<paddle::dialect::ScalarAttribute>().data();

      {% elif item.typename == "IntArray" and item.support_tensor %}

  Tensor {{item.name}}_(std::make_shared<primitive::LazyTensor>(op_obj.{{item.name}}()));

  auto* {{item.name}}_define_op =
      std::static_pointer_cast<primitive::LazyTensor>({{item.name}}_.impl())
          ->value()
          .defining_op();
  if ({{item.name}}_define_op->name() != "pd_op.full_int_array") {
    return {};
  }
  IntArray {{item.name}} = phi::IntArray(
      paddle::dialect::GetInt64Vector({{item.name}}_define_op->attribute("value")));

      {% else %}
        {% if item.mapped_type[0] == "pir::StrAttribute" %}
  {{item.mapped_type[1]}} {{item.name}} = op->attribute("{{item.name}}").dyn_cast<{{item.mapped_type[0]}}>().AsString();
        {% elif "[]" in item.typename %}
  auto array_list = op->attribute("{{item.name}}").dyn_cast<pir::ArrayAttribute>().AsVector();
  {% set temp_type= item.mapped_type[0]|replace('pir::ArrayAttribute<', '')|replace('>', '')%}
  {{item.mapped_type[1]|replace('const ', '')|replace('&', '')}} {{item.name}};
  if (array_list.size() > 0) {
    if (array_list[0].isa<{{temp_type}}>()) {
      for (size_t i = 0; i < array_list.size(); ++i) {
        {{item.name}}.push_back(
            array_list[i].dyn_cast<{{temp_type}}>().data());
      }

    } else {
      return {};
    }
  }
        {% else %}
  {{item.mapped_type[1]}} {{item.name}} = op->attribute("{{item.name}}").dyn_cast<{{item.mapped_type[0]}}>().data();
        {% endif %}
      {% endif %}
      {% endfor %}
    {% endif %}

  VLOG(6) << "Decomp call {{fwd_name}}'s forward composite rule prepare";

  auto org_res = op->results();
  std::vector<std::vector<pir::Value>> res(org_res.size());

  VLOG(6) << "Decomp call {{fwd_name}}'s forward composite rule begin";
  {% if outputs|length == 1 %}
    {% if outputs[0].typename == "Tensor[]" %}
  std::vector<Tensor> op_res = paddle::primitive::details::{{fwd_name}}_decomp<primitive::LazyTensor>({{common.args(input_names, attr_names)}});
    {% else %}
  Tensor op_res = paddle::primitive::details::{{fwd_name}}_decomp<primitive::LazyTensor>({{common.args(input_names, attr_names)}});
    {% endif %}
  VLOG(6) << "Decomp call {{fwd_name}}'s forward composite rule end";
    {% if outputs[0].typename == "Tensor[]" %}
  for (size_t idx = 0; idx < op_res.size(); idx++) {
      res[0].push_back(
          std::static_pointer_cast<primitive::LazyTensor>(op_res[idx].impl())
              ->value());
  }
    {% else %}
  res[0].push_back(
    std::static_pointer_cast<primitive::LazyTensor>(op_res.impl())
        ->value());
    {% endif %}
  {% else %}
    {% for item in outputs %}
      {% do output_names.append(item.name) %}
      {% do output_types.append(item.mapped_type) %}
    {% endfor %}
  std::tuple<{{common.sequence('', '', ', ', output_types)}}> op_res = paddle::primitive::details::{{fwd_name}}_decomp<primitive::LazyTensor>(
          {{common.args(input_names, attr_names)}});
  VLOG(6) << "Decomp call {{fwd_name}}'s forward composite rule end";
  {% for k in range(outputs|length) %}
    {% if outputs[k].intermediate and fwd_name in decomp_ops_list_contain_unused_output %}
  pir::Value {{outputs[k].name}};
  res[{{k}}].push_back({{outputs[k].name}});
    {% else %}
  res[{{k}}].push_back(std::static_pointer_cast<primitive::LazyTensor>(std::get<{{k}}>(op_res).impl())->value());
    {% endif %}
  {% endfor %}
  {% endif %}

  VLOG(4) << "Decomp call {{fwd_name}}'s decomp interface end";
  return res;
}
{% endmacro %}

{% for api in apis -%}
  {% if api.name in decomp_white_list %}
    {{sig(api.name, api.class_name, api.inputs, api.attrs, api.outputs)}}
  {% else %} {# render nothing #}
  {% endif %}
{% endfor %}

}  // namespace dialect
}  // namespace paddle