chore: import upstream snapshot with attribution

paddle/fluid/pybind/eager_custom_python_api.h

@@ -0,0 +1,172 @@
+// Copyright (c) 2022 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
+//
+//     http://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 <iostream>
+
+#include "paddle/fluid/eager/to_static/run_program_func.h"
+#include "paddle/fluid/eager/utils.h"
+#include "paddle/phi/core/enforce.h"
+
+using egr::ConvertAllInputsToDistTensor;
+using egr::InputsContainDistTensor;
+
+namespace paddle {
+namespace pybind {
+
+static PyObject *eager_api_linear(PyObject *self,
+                                  PyObject *args,
+                                  PyObject *kwargs) {
+  PyThreadState *tstate = nullptr;
+  try {
+    auto &x = GetTensorFromArgs("linear", "X", args, 0, false);
+    auto &weight = GetTensorFromArgs("linear", "weight", args, 1, false);
+    auto &bias = GetTensorFromArgs("linear", "Bias", args, 2, true);
+
+    tstate = PyEval_SaveThread();
+    SetPythonStack();
+    if (bias.is_dist_tensor() || bias.has_allocation()) {
+      const phi::distributed::ProcessMesh *mesh = nullptr;
+      if (InputsContainDistTensor(&mesh, x, weight, bias)) {
+        ConvertAllInputsToDistTensor(mesh, x, weight, bias);
+      }
+
+      auto mm_out = matmul_ad_func(x, weight, false, false);
+      auto out = add_ad_func(mm_out, bias);
+      PyEval_RestoreThread(tstate);
+      tstate = nullptr;
+      return ToPyObject(out);
+    } else {
+      const phi::distributed::ProcessMesh *mesh = nullptr;
+      if (InputsContainDistTensor(&mesh, x, weight)) {
+        ConvertAllInputsToDistTensor(mesh, x, weight);
+      }
+
+      auto mm_out = matmul_ad_func(x, weight, false, false);
+      PyEval_RestoreThread(tstate);
+      tstate = nullptr;
+      return ToPyObject(mm_out);
+    }
+  } catch (paddle::platform::EnforceNotMet &exception) {
+    if (tstate) {
+      PyEval_RestoreThread(tstate);
+    }
+    std::ostringstream sout;
+    sout << exception.what();
+    sout << "  [operator < linear > error]";
+    exception.set_error_str(sout.str());
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  } catch (...) {
+    if (tstate) {
+      PyEval_RestoreThread(tstate);
+    }
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  }
+}
+
+static PyObject *eager_api_run_program(PyObject *self,
+                                       PyObject *args,
+                                       PyObject *kwargs) {
+  PyThreadState *tstate = nullptr;
+  try {
+    auto X_info = GetPyArgumentInfo("run_program", "X", args, 0, true);
+    TensorListBufferAllocator X_allocator(X_info.second);
+    auto &X = GetTensorListFromArgsWithBuffer("run_program",
+                                              "X",
+                                              0,
+                                              nullptr,
+                                              X_info.first,
+                                              X_info.second,
+                                              X_allocator);
+
+    auto Params_info =
+        GetPyArgumentInfo("run_program", "Params", args, 1, true);
+    TensorListBufferAllocator Params_allocator(Params_info.second);
+    auto &Params = GetTensorListFromArgsWithBuffer("run_program",
+                                                   "Params",
+                                                   0,
+                                                   nullptr,
+                                                   Params_info.first,
+                                                   Params_info.second,
+                                                   Params_allocator);
+
+    auto OutScope =
+        GetScopePtrListFromArgs("run_program", "OutScope", args, 2, false);
+    const phi::distributed::ProcessMesh *mesh = nullptr;
+    if (InputsContainDistTensor(&mesh, X, Params)) {
+      X = GetTensorListFromArgsWithBuffer("run_program",
+                                          "X",
+                                          0,
+                                          nullptr,
+                                          X_info.first,
+                                          X_info.second,
+                                          X_allocator);
+      Params = GetTensorListFromArgsWithBuffer("run_program",
+                                               "Params",
+                                               0,
+                                               nullptr,
+                                               Params_info.first,
+                                               Params_info.second,
+                                               Params_allocator);
+    }
+    VLOG(6) << "Start PIR GetProgramAttributesMapPtrFromPyArgs";
+    auto prog_attrs_ptr =
+        GetProgramAttributesMapPtrFromPyArgs("run_program", args, 3);
+    VLOG(6) << "Finish PIR GetProgramAttributesMapPtrFromPyArgs";
+
+    VLOG(6) << "Start PIR ConstructCudaGraphAttrMapForRunProgram";
+    paddle::framework::AttributeMap cuda_graph_attrs;
+    ConstructCudaGraphAttrMapForRunProgram(
+        "run_program", args, 4, cuda_graph_attrs);
+    VLOG(6) << "Finish PIR ConstructCudaGraphAttrMapForRunProgram";
+    tstate = PyEval_SaveThread();
+    auto out = egr::to_static::run_program_ad_func(
+        X, Params, OutScope, *prog_attrs_ptr, cuda_graph_attrs);
+    PyEval_RestoreThread(tstate);
+    tstate = nullptr;
+    return ToPyObject(out);
+  } catch (paddle::platform::EnforceNotMet &exception) {
+    if (tstate) {
+      PyEval_RestoreThread(tstate);
+    }
+    std::ostringstream sout;
+    sout << exception.what();
+    sout << "  [operator < run_program > error]";
+    exception.set_error_str(sout.str());
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  } catch (...) {
+    if (tstate) {
+      PyEval_RestoreThread(tstate);
+    }
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  }
+}
+
+static PyMethodDef CustomEagerFinalStateMethods[] = {
+    {"linear",
+     (PyCFunction)(void (*)(void))eager_api_linear,
+     METH_VARARGS | METH_KEYWORDS,
+     "C++ interface function for linear."},
+    {"run_program",
+     (PyCFunction)(void (*)(void))eager_api_run_program,
+     METH_VARARGS | METH_KEYWORDS,
+     "C++ interface function for run_program in dygraph."},
+    {nullptr, nullptr, 0, nullptr}};
+
+}  // namespace pybind
+}  // namespace paddle