【New IR] Concat python api and vjp (#56316)

xiaoguoguo626807 · cyber-pioneer · Charles-hit · web-flow · commit 1da20aa8bf98 · 2023-08-21T10:56:31.000+08:00
* support ir api form prim

* convert vector of int to intarray

* support ir api for prim

* Add more gen api

* concat python api to concat_grad

* fix gen conflict

* support vjp prim mode in new ir

* remove useless code

* add vjp autogen v1.0

* add test for prim

* resolve type conflict

* modify utils

* remove useless code

* add split op and modify some bug of vectorType

* fix conflict

* add concat python test

---------

Co-authored-by: cyber-pioneer &lt;chenzhuo@tju.edu.cn&gt;
Co-authored-by: Charles-hit &lt;wanghao107@baidu.com&gt;
Co-authored-by: 0x45f &lt;wangzhen45@baidu.com&gt;
Co-authored-by: chenzhiyang &lt;1792266893@qq.com&gt;
Co-authored-by: Chen Zhiyang &lt;chenzhiyang99@126.com&gt;
diff --git a/paddle/fluid/ir/dialect/op_generator/api_gen.py b/paddle/fluid/ir/dialect/op_generator/api_gen.py
@@ -29,6 +29,7 @@
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/common/scalar.h"
+#include "paddle/fluid/ir/dialect/pd_manual_api.h"
 
 {body}
 
diff --git a/paddle/fluid/ir/dialect/op_generator/vjp_interface_gen_op_list.py b/paddle/fluid/ir/dialect/op_generator/vjp_interface_gen_op_list.py
@@ -21,5 +21,13 @@
 # TODO(wanghao107)
 # remove this file and support Vjp methods
 # code gen.
-vjp_interface_declare_gen_op_list = ["tanh", "mean", "divide", "sum", "add"]
+
+vjp_interface_declare_gen_op_list = [
+    "tanh",
+    "mean",
+    "divide",
+    "sum",
+    "add",
+    "concat",
+]
 vjp_interface_implementation_gen_op_list = ["tanh", "mean", "divide", "add"]
diff --git a/paddle/fluid/ir/dialect/pd_manual_api.cc b/paddle/fluid/ir/dialect/pd_manual_api.cc
@@ -0,0 +1,36 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/ir/dialect/pd_manual_api.h"
+#include "paddle/fluid/ir/dialect/pd_dialect.h"
+#include "paddle/fluid/ir/dialect/pd_op.h"
+#include "paddle/ir/core/builder.h"
+#include "paddle/ir/core/builtin_op.h"
+
+namespace paddle {
+namespace dialect {
+std::vector<ir::OpResult> concat_grad(std::vector<ir::OpResult> x,
+                                      ir::OpResult out_grad,
+                                      ir::OpResult axis) {
+  auto combine_op =
+      APIBuilder::Instance().GetBuilder()->Build<ir::CombineOp>(x);
+  paddle::dialect::ConcatGradOp concat_grad_op =
+      APIBuilder::Instance().GetBuilder()->Build<paddle::dialect::ConcatGradOp>(
+          combine_op.out(), out_grad, axis);
+  auto split_op = APIBuilder::Instance().GetBuilder()->Build<ir::SplitOp>(
+      concat_grad_op.result(0));
+  return split_op.outputs();
+}
+}  // namespace dialect
+}  // namespace paddle
diff --git a/paddle/fluid/ir/dialect/pd_manual_api.h b/paddle/fluid/ir/dialect/pd_manual_api.h
@@ -0,0 +1,30 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <vector>
+
+#include "paddle/ir/core/value.h"
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/common/place.h"
+
+namespace paddle {
+namespace dialect {
+
+std::vector<ir::OpResult> concat_grad(std::vector<ir::OpResult> x,
+                                      ir::OpResult out_grad,
+                                      ir::OpResult axis);
+}  // namespace dialect
+}  // namespace paddle
diff --git a/paddle/fluid/ir/dialect/pd_op_vjp_manual.cc b/paddle/fluid/ir/dialect/pd_op_vjp_manual.cc
@@ -16,6 +16,7 @@
 #include "paddle/fluid/ir/dialect/pd_op.h"
 #include "paddle/fluid/primitive/rule/vjp/vjp.h"
 #include "paddle/fluid/primitive/type/lazy_tensor.h"
+#include "paddle/ir/core/builtin_op.h"
 #include "paddle/ir/core/op_base.h"
 #include "paddle/phi/common/int_array.h"
 
@@ -24,6 +25,42 @@
 
 namespace paddle {
 namespace dialect {
+using IntArray = paddle::experimental::IntArray;
+
+std::vector<std::vector<ir::OpResult>> ConcatOp::Vjp(
+    ir::Operation* op,
+    const std::vector<std::vector<ir::OpResult>>& out_grads,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  ConcatOp op_obj = op->dyn_cast<ConcatOp>();
+  ir::CombineOp combine_op_obj =
+      op_obj.x().GetDefiningOp()->dyn_cast<ir::CombineOp>();
+  std::vector<Tensor> x;
+  for (size_t idx = 0; idx < combine_op_obj.inputs().size(); idx++) {
+    x.emplace_back(
+        std::make_shared<primitive::LazyTensor>(combine_op_obj.inputs()[idx]));
+  }
+
+  Tensor out_grad(std::make_shared<primitive::LazyTensor>(out_grads[0][0]));
+  Tensor axis(std::make_shared<primitive::LazyTensor>(op_obj.axis()));
+
+  std::vector<std::vector<Tensor>> tensor_res =
+      primitive::concat_vjp(x, out_grad, axis, stop_gradients);
+  std::vector<std::vector<ir::OpResult>> res(tensor_res.size(),
+                                             std::vector<ir::OpResult>());
+  for (uint64_t i = 0; i < tensor_res.size(); i++) {
+    res[i].resize(tensor_res[i].size());
+    for (uint64_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>();
+      }
+    }
+  }
+  return res;
+}
+
 std::vector<std::vector<ir::OpResult>> SumOp::Vjp(
     ir::Operation* op,
     const std::vector<std::vector<ir::OpResult>>& out_grads,
diff --git a/paddle/fluid/primitive/backend/static_backend.cc b/paddle/fluid/primitive/backend/static_backend.cc
@@ -242,6 +242,38 @@ Tensor sum_grad<LazyTensor>(const Tensor& x,
       x_res, out_grad_res, axis.GetData(), keepdim, reduce_all);
   return Tensor(std::make_shared<LazyTensor>(op_res));
 }
+
+template <>
+std::vector<Tensor> concat_grad<LazyTensor>(const std::vector<Tensor>& x,
+                                            const Tensor& out_grad,
+                                            const Tensor& axis) {
+  std::vector<ir::OpResult> x_res;
+  for (uint64_t idx = 0; idx < x.size(); idx++) {
+    x_res.emplace_back(std::static_pointer_cast<LazyTensor>(x[idx].impl())
+                           ->getValue()
+                           .dyn_cast<ir::OpResult>());
+  }
+
+  ir::OpResult out_grad_res =
+      std::static_pointer_cast<LazyTensor>(out_grad.impl())
+          ->getValue()
+          .dyn_cast<ir::OpResult>();
+
+  ir::OpResult axis_res = std::static_pointer_cast<LazyTensor>(axis.impl())
+                              ->getValue()
+                              .dyn_cast<ir::OpResult>();
+
+  std::vector<ir::OpResult> op_res =
+      paddle::dialect::concat_grad(x_res, out_grad_res, axis_res);
+
+  std::vector<Tensor> op_result;
+  for (uint64_t idx = 0; idx < op_res.size(); idx++) {
+    op_result.emplace_back(
+        std::make_shared<primitive::LazyTensor>(op_res[idx]));
+  }
+  return op_result;
+}
+
 }  // namespace backend
 }  // namespace primitive
 }  // namespace paddle
diff --git a/paddle/fluid/primitive/backend/static_backend.h b/paddle/fluid/primitive/backend/static_backend.h
@@ -37,6 +37,11 @@ Tensor mean_grad(const Tensor& x,
                  bool keepdim = false,
                  bool reduce_all = false);
 
+template <typename T>
+std::vector<Tensor> concat_grad(const std::vector<Tensor>& x,
+                                const Tensor& out_grad,
+                                const Tensor& axis);
+
 template <typename T>
 std::tuple<Tensor, Tensor> add_grad(const Tensor& x,
                                     const Tensor& y,
diff --git a/paddle/fluid/primitive/rule/vjp/vjp.cc b/paddle/fluid/primitive/rule/vjp/vjp.cc
@@ -148,6 +148,28 @@ std::vector<std::vector<paddle::Tensor>> add_vjp(
   return vjp_res;
 }
 
+std::vector<std::vector<paddle::Tensor>> concat_vjp(
+    const std::vector<Tensor>& x,
+    const Tensor& out_grad,
+    const Tensor& axis,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  std::vector<std::vector<paddle::Tensor>> vjp_res(2, std::vector<Tensor>());
+  // get concat_grad res.
+  std::vector<Tensor> op_res =
+      backend::concat_grad<primitive::LazyTensor>(x, out_grad, axis);
+
+  // construct vjp result by op result and stop_gradients info
+  vjp_res[0].resize(op_res.size());
+  for (uint64_t idx = 0; idx < op_res.size(); idx++) {
+    if (!stop_gradients[0][idx]) {
+      vjp_res[0][idx] = op_res[idx];
+    }
+  }
+  // vjp_res[1] is axis's grad which is attribute (no grad).
+  vjp_res[1].resize(1);
+  return vjp_res;
+}
+
 std::vector<std::vector<paddle::Tensor>> divide_vjp(
     const Tensor& x,
     const Tensor& y,
diff --git a/paddle/fluid/primitive/rule/vjp/vjp.h b/paddle/fluid/primitive/rule/vjp/vjp.h
@@ -38,6 +38,12 @@ std::vector<std::vector<paddle::Tensor>> mean_vjp(
     bool reduce_all,
     const std::vector<std::vector<bool>>& stop_gradients);
 
+std::vector<std::vector<paddle::Tensor>> concat_vjp(
+    const std::vector<Tensor>& x,
+    const Tensor& out_grad,
+    const Tensor& axis,
+    const std::vector<std::vector<bool>>& stop_gradients);
+
 std::vector<std::vector<paddle::Tensor>> add_vjp(
     const Tensor& x,
     const Tensor& y,
diff --git a/paddle/fluid/pybind/ir.cc b/paddle/fluid/pybind/ir.cc
@@ -249,6 +249,7 @@ void BindValue(py::module *m) {
       .def("get_defining_op",
            &Value::GetDefiningOp,
            return_value_policy::reference)
+      .def("first_use", &Value::first_use, return_value_policy::reference)
       .def("__eq__", &Value::operator==)
       .def("__eq__",
            [](Value &self, OpResult &other) {
@@ -272,9 +273,11 @@ void BindOpOperand(py::module *m) {
   op_operand
       .def("source",
            [](OpOperand &self) { return self.source().dyn_cast<OpResult>(); })
-      .def("set_source", [](OpOperand &self, const OpResult &result) {
-        self.set_source(result);
-      });
+      .def("set_source",
+           [](OpOperand &self, const OpResult &result) {
+             self.set_source(result);
+           })
+      .def("owner", &OpOperand::owner, return_value_policy::reference);
 }
 
 bool GetStopGradient(const OpResult &self) {
@@ -331,6 +334,7 @@ void BindOpResult(py::module *m) {
       .def("get_defining_op",
            &OpResult::GetDefiningOp,
            return_value_policy::reference)
+      .def("first_use", &OpResult::first_use, return_value_policy::reference)
       .def("use_empty", &OpResult::use_empty)
       .def("type", &OpResult::type)
       .def_property(
diff --git a/paddle/fluid/pybind/ops_api.cc b/paddle/fluid/pybind/ops_api.cc
@@ -40,7 +40,11 @@ static PyObject *divide(PyObject *self, PyObject *args, PyObject *kwargs) {
   return static_api_divide(self, args, kwargs);
 }
 
-static PyMethodDef OpsAPI[] = {{"add_n",  // NOLINT
+static PyObject *concat(PyObject *self, PyObject *args, PyObject *kwargs) {
+  return static_api_concat(self, args, kwargs);
+}
+
+static PyMethodDef OpsAPI[] = {{"add_n",
                                 (PyCFunction)(void (*)(void))add_n,
                                 METH_VARARGS | METH_KEYWORDS,
                                 "C++ interface function for add_n."},
@@ -56,6 +60,10 @@ static PyMethodDef OpsAPI[] = {{"add_n",  // NOLINT
                                 (PyCFunction)(void (*)(void))divide,
                                 METH_VARARGS | METH_KEYWORDS,
                                 "C++ interface function for divide."},
+                               {"concat",
+                                (PyCFunction)(void (*)(void))concat,
+                                METH_VARARGS | METH_KEYWORDS,
+                                "C++ interface function for concat."},
                                {"full",
                                 (PyCFunction)(void (*)(void))full,
                                 METH_VARARGS | METH_KEYWORDS,
diff --git a/paddle/fluid/pybind/static_op_function.cc b/paddle/fluid/pybind/static_op_function.cc
@@ -109,6 +109,26 @@ PyObject *static_api_divide(PyObject *self, PyObject *args, PyObject *kwargs) {
   }
 }
 
+PyObject *static_api_concat(PyObject *self, PyObject *args, PyObject *kwargs) {
+  try {
+    VLOG(6) << "Add concat op into program";
+    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);
+    // Get OpResult from args
+    PyObject *x_obj = PyTuple_GET_ITEM(args, 0);
+    auto x = CastPyArg2VectorOfOpResult("concat", x_obj, 0);
+
+    PyObject *axis_obj = PyTuple_GET_ITEM(args, 1);
+    paddle::experimental::Scalar axis = CastPyArg2Scalar(axis_obj, "concat", 1);
+
+    // Call ir static api
+    auto out = paddle::dialect::concat(x, axis.to<float>());
+    return ToPyObject(out);
+  } catch (...) {
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  }
+}
+
 PyObject *static_api_full(PyObject *self, PyObject *args, PyObject *kwargs) {
   try {
     VLOG(6) << "Add full op into program";
diff --git a/paddle/fluid/pybind/static_op_function.h b/paddle/fluid/pybind/static_op_function.h
@@ -28,6 +28,7 @@ PyObject *static_api_add_n(PyObject *self, PyObject *args, PyObject *kwargs);
 PyObject *static_api_mean(PyObject *self, PyObject *args, PyObject *kwargs);
 PyObject *static_api_sum(PyObject *self, PyObject *args, PyObject *kwargs);
 PyObject *static_api_divide(PyObject *self, PyObject *args, PyObject *kwargs);
+PyObject *static_api_concat(PyObject *self, PyObject *args, PyObject *kwargs);
 PyObject *static_api_full(PyObject *self, PyObject *args, PyObject *kwargs);
 
 }  // namespace pybind
diff --git a/python/paddle/autograd/backward.py b/python/paddle/autograd/backward.py
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
diff --git a/test/cpp/prim/test_vjp.cc b/test/cpp/prim/test_vjp.cc
diff --git a/test/ir/new_ir/test_build_op.py b/test/ir/new_ir/test_build_op.py
diff --git a/test/ir/new_ir/test_ir_backward.py b/test/ir/new_ir/test_ir_backward.py
