{% 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/decomp_rule/decomp_vjp/details.h"
#include "paddle/fluid/primitive/vjp_interface/generated/generated_vjp.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}}::DecompVjp(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 == "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") {
    PADDLE_THROW(
        common::errors::Unimplemented("We don't support dynamic tensors "
                                        "attribute {{item.name}} for {{fwd_name}} decomposition "
                                        "for now. "));
  }
  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") {
    PADDLE_THROW(
        common::errors::Unimplemented("We don't support dynamic tensors "
                                        "attribute {{item.name}} for {{fwd_name}} decomposition "
                                        "for now. "));
  }
  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 {
      PADDLE_THROW(common::errors::Unimplemented("attr is not vector of {{temp_type}} "));
    }
  }
        {% 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 backward composite rule prepare";

  std::vector<std::vector<bool>> stop_gradients = ConstructStopGradient(op);

  std::vector<std::vector<paddle::Tensor>> tensor_res;
  for (auto arg : stop_gradients) {
    tensor_res.push_back(std::vector<paddle::Tensor>(arg.size()));
  }
  std::string op_name = "{{fwd_name}}";
  FLAGS_tensor_operants_mode = "static";
  VLOG(4) << "Call Pir Decomposed backward op {{fwd_name}}";


  {% for k in range(outputs|length) %}
  paddle::Tensor* {{outputs[k].name}} = !stop_gradients[{{k}}][0] ? &tensor_res[{{k}}][0] : nullptr;
  {% endfor %}

  {% for item in outputs %}
    {% do output_names.append(item.name) %}
    {% do output_types.append(item.mapped_type) %}
  {% endfor %}

  paddle::primitive::details::{{fwd_name}}<primitive::LazyTensor>(
      {{common.args(input_names, attr_names)}}, {{common.sequence('', '', ', ', output_names)}});
  std::vector<std::vector<pir::Value>> res(tensor_res.size());

  for (size_t i = 0; i < tensor_res.size(); ++i) {
    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())
                        ->value();
      }
    }
  }

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

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

}  // namespace dialect
}  // namespace paddle