[JIT]Layer supports eager dygraph mode and Polish Function interface (#44283)

* [JIT]Layer support eager dygraph mode and polish Function interface

* remove usless code

* fix #define

paddle/fluid/jit/base_function.h

@@ -14,23 +14,23 @@
 
 #pragma once
 
-#include <ostream>
-#include <string>
-
-#include "paddle/phi/common/place.h"
-
-#include "paddle/fluid/framework/variable.h"
+#include "paddle/phi/api/include/tensor.h"
+#include "paddle/phi/core/dense_tensor.h"
 
 namespace paddle {
 namespace jit {
 
-using Variable = paddle::framework::Variable;
+using Tensor = paddle::experimental::Tensor;
+using DenseTensor = phi::DenseTensor;
+
 class BaseFunction {
  public:
-  virtual std::vector<Variable> operator()(
-      const std::vector<Variable> &inputs) = 0;
+  virtual std::vector<DenseTensor> operator()(
+      const std::vector<DenseTensor> &inputs) = 0;
+
+  virtual std::vector<Tensor> operator()(const std::vector<Tensor> &inputs) = 0;
+
   virtual ~BaseFunction() {}
-  // virtual void SetPalce(const phi::Place &place);
 };
 
 }  // namespace jit

paddle/fluid/jit/compilation_unit.cc

@@ -24,7 +24,7 @@ std::shared_ptr<BaseFunction> CompilationUnit::Function(
   PADDLE_ENFORCE_EQ(
       function_map_.count(name),
       1,
-      platform::errors::InvalidArgument(
+      phi::errors::InvalidArgument(
           "Funciton name %s is not exist in function_map_.", name));
   return function_map_.at(name);
 }

paddle/fluid/jit/executor_function.h

@@ -42,17 +42,21 @@ class ExecutorFunction : public BaseFunction {
 
   ~ExecutorFunction() noexcept {}
 
-  std::vector<Variable> operator()(const std::vector<Variable> &inputs) {
-    utils::ShareInputsIntoScope(info_->InputArgNames(), inputs, &scope_);
+  std::vector<Tensor> operator()(const std::vector<Tensor> &inputs) {
+    auto dense_tensors = utils::ToDenseTensors(inputs);
+    return utils::ToTensors(this->operator()(dense_tensors));
+  }
+
+  std::vector<DenseTensor> operator()(const std::vector<DenseTensor> &inputs) {
+    utils::ShareIntoScope(info_->InputArgNames(), inputs, &scope_);
     inner_exe_.Run(info_->ProgramDesc(),
                    &scope_,
                    /*blockID=*/0,
                    false,
                    true,
                    info_->OutputArgNames());
-    VLOG(6) << framework::GenScopeTreeDebugInfo(&scope_);
-    std::vector<Variable> res;
-    utils::FetchVarsByNames(info_->OutputArgNames(), scope_, &res);
+    std::vector<DenseTensor> res;
+    utils::FetchOuts(info_->OutputArgNames(), scope_, &res);
     return res;
   }
 

paddle/fluid/jit/function_utils.cc

@@ -21,36 +21,50 @@
 namespace paddle {
 namespace jit {
 namespace utils {
-void FetchVarsByNames(const std::vector<std::string> &names,
-                      const framework::Scope &scope,
-                      std::vector<Variable> *outs) {
-  for (auto &out_name : names) {
+
+std::vector<DenseTensor> ToDenseTensors(const std::vector<Tensor> &tensors) {
+  std::vector<DenseTensor> ret;
+  for (auto &t : tensors) {
+    ret.emplace_back(*std::dynamic_pointer_cast<phi::DenseTensor>(t.impl()));
+  }
+  return ret;
+}
+
+std::vector<Tensor> ToTensors(const std::vector<DenseTensor> &tensors) {
+  std::vector<Tensor> ret;
+  for (auto &t : tensors) {
+    ret.emplace_back(std::make_shared<DenseTensor>(t));
+  }
+  return ret;
+}
+
+void FetchOuts(const std::vector<std::string> &names,
+               const framework::Scope &scope,
+               std::vector<DenseTensor> *outs) {
+  outs->reserve(names.size());
+  for (size_t i = 0; i < names.size(); ++i) {
+    auto &out_name = names[i];
     VLOG(3) << "fetch out: " << out_name;
     auto *var = scope.FindVar(out_name);
     auto &src_tensor = var->Get<DenseTensor>();
-    Variable v;
-    auto *p = v.GetMutable<DenseTensor>();
-    *p = src_tensor;
-    outs->emplace_back(v);
+    outs->emplace_back(src_tensor);
   }
 }
 
-void ShareInputsIntoScope(const std::vector<std::string> &ordered_input_names,
-                          const std::vector<Variable> &vars,
-                          framework::Scope *scope) {
-  VLOG(3) << "vars size: " << vars.size();
+void ShareIntoScope(const std::vector<std::string> &ordered_input_names,
+                    const std::vector<DenseTensor> &tensors,
+                    framework::Scope *scope) {
+  VLOG(3) << "tensors size: " << tensors.size();
   PADDLE_ENFORCE_EQ(
-      vars.size(),
+      tensors.size(),
       ordered_input_names.size(),
       platform::errors::InvalidArgument(
-          "vars.size() should be equal to ordered_input_names.size()."));
-
-  for (size_t i = 0; i < vars.size(); i++) {
+          "tensors.size() should be equal to ordered_input_names.size()."));
+  for (size_t i = 0; i < tensors.size(); ++i) {
     VLOG(3) << "share into scope: " << ordered_input_names[i];
-    auto &dense_tensor = vars[i].Get<DenseTensor>();
     auto *var = scope->Var(ordered_input_names[i]);
     auto *dst_tensor = var->GetMutable<DenseTensor>();
-    *dst_tensor = dense_tensor;
+    *dst_tensor = tensors[i];
   }
 }
 

paddle/fluid/jit/function_utils.h

@@ -20,6 +20,7 @@
 
 #include "paddle/fluid/framework/scope.h"
 #include "paddle/fluid/framework/variable.h"
+#include "paddle/phi/api/include/tensor.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/dense_tensor.h"
 
@@ -30,15 +31,20 @@ namespace jit {
 using Variable = paddle::framework::Variable;
 using Name2VariableMap = std::unordered_map<std::string, Variable>;
 using DenseTensor = phi::DenseTensor;
+using Tensor = paddle::experimental::Tensor;
+
 namespace utils {
 
-void FetchVarsByNames(const std::vector<std::string> &names,
-                      const framework::Scope &scope,
-                      std::vector<Variable> *outs);
+std::vector<DenseTensor> ToDenseTensors(const std::vector<Tensor> &tensors);
+std::vector<Tensor> ToTensors(const std::vector<DenseTensor> &tensors);
 
-void ShareInputsIntoScope(const std::vector<std::string> &ordered_input_names,
-                          const std::vector<Variable> &vars,
-                          framework::Scope *scope);
+void FetchOuts(const std::vector<std::string> &names,
+               const framework::Scope &scope,
+               std::vector<DenseTensor> *outs);
+
+void ShareIntoScope(const std::vector<std::string> &ordered_input_names,
+                    const std::vector<DenseTensor> &vars,
+                    framework::Scope *scope);
 
 void ShareParamsIntoScope(const std::vector<std::string> &param_names,
                           const Name2VariableMap &params_dict,

paddle/fluid/jit/layer.cc

@@ -16,9 +16,6 @@
 
 namespace paddle {
 namespace jit {
-// TODO(dev): Make vector<string>, num_slot as in argument
-// Layer(const std::shared_ptr<ClassType>& type) : obj_(type, /*num_slot*/ 0U)
-// {}
 Layer::Layer(const std::vector<std::shared_ptr<FunctionInfo>>& infos,
              const Name2VariableMap& params_dict,
              const phi::Place& place)
@@ -30,7 +27,13 @@ std::shared_ptr<BaseFunction> Layer::Function(const std::string& name) const {
   return unit_.Function(name);
 }
 
-std::vector<Variable> Layer::forward(const std::vector<Variable>& inputs) {
+std::vector<Tensor> Layer::forward(const std::vector<Tensor>& inputs) {
+  auto func = Function("forward");
+  return (*func)(inputs);
+}
+
+std::vector<DenseTensor> Layer::forward(
+    const std::vector<DenseTensor>& inputs) {
   auto func = Function("forward");
   return (*func)(inputs);
 }

paddle/fluid/jit/layer.h

@@ -32,9 +32,6 @@ using Name2VariableMap = std::unordered_map<std::string, Variable>;
 
 class Layer {
  public:
-  // TODO(dev): Make vector<string>, num_slot as in argument
-  // Layer(const std::shared_ptr<ClassType>& type) : obj_(type, /*num_slot*/ 0U)
-  // {}
   Layer(const std::vector<std::shared_ptr<FunctionInfo>>& infos,
         const Name2VariableMap& params_dict,
         const phi::Place& place);
@@ -43,7 +40,9 @@ class Layer {
 
   Variable Attribute(const std::string& name) const;
 
-  std::vector<Variable> forward(const std::vector<Variable>& inputs);
+  std::vector<Tensor> forward(const std::vector<Tensor>& inputs);
+
+  std::vector<DenseTensor> forward(const std::vector<DenseTensor>& inputs);
 
   void to(const phi::Place& place);
 
@@ -55,7 +54,6 @@ class Layer {
   const Name2FunctionMap& FunctionMap() const;
 
  private:
-  // internal::Object obj_;
   Name2VariableMap params_dict_;
   Name2VariableMap attrs_dict_;
   CompilationUnit unit_;

paddle/fluid/jit/layer_test.cc

@@ -52,17 +52,16 @@ namespace paddle {
 namespace jit {
 using DenseTensor = phi::DenseTensor;
 
-std::vector<Variable> PrepareInputs(const phi::Place& place) {
+std::vector<DenseTensor> PrepareInputs(const phi::Place& place) {
   platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
   auto& dev_ctx = *pool.Get(place);
 
-  Variable v;
-  auto* dense_tensor = v.GetMutable<DenseTensor>();
-  dense_tensor->Resize(phi::make_ddim({2, 4}));
-  dense_tensor->mutable_data<float>(place);
-  phi::funcs::set_constant(dev_ctx, dense_tensor, 2.);
+  DenseTensor t;
+  t.Resize(phi::make_ddim({2, 4}));
+  t.mutable_data<float>(place);
+  phi::funcs::set_constant(dev_ctx, &t, 2.);
 
-  return {v};
+  return {t};
 }
 
 TEST(CpuLayerTest, Construct) {
@@ -72,16 +71,12 @@ TEST(CpuLayerTest, Construct) {
   auto inputs = PrepareInputs(place);
 
   auto outs = layer.forward(inputs);
-  auto out_vars = outs[0];
-  auto out_dense_tensor = out_vars.Get<DenseTensor>();
-  auto out_data = out_dense_tensor.data<float>();
+  auto out_data = outs[0].data<float>();
   EXPECT_NEAR(out_data[0], 0.02194316, 1e-6);
 
   auto func = layer.Function("infer");
   outs = (*func)(inputs);
-  out_vars = outs[0];
-  out_dense_tensor = out_vars.Get<DenseTensor>();
-  out_data = out_dense_tensor.data<float>();
+  out_data = outs[0].data<float>();
   EXPECT_NEAR(out_data[0], 1.41562390, 1e-6);
 }
 
@@ -98,8 +93,7 @@ TEST(GpuLayerTest, Construct) {
   auto inputs = PrepareInputs(place);
 
   auto outs = layer.forward(inputs);
-  auto out_vars = outs[0];
-  auto out_dense_tensor = out_vars.Get<DenseTensor>();
+  auto out_dense_tensor = outs[0];
   phi::Copy(
       *dev_ctx_gpu, out_dense_tensor, phi::CPUPlace(), true, &cpu_dense_tensor);
   auto out_data = cpu_dense_tensor.data<float>();
@@ -107,8 +101,7 @@ TEST(GpuLayerTest, Construct) {
 
   auto func = layer.Function("infer");
   outs = (*func)(inputs);
-  out_vars = outs[0];
-  out_dense_tensor = out_vars.Get<DenseTensor>();
+  out_dense_tensor = outs[0];
   phi::Copy(
       *dev_ctx_gpu, out_dense_tensor, phi::CPUPlace(), true, &cpu_dense_tensor);
   out_data = cpu_dense_tensor.data<float>();

paddle/fluid/jit/pe_function.h

@@ -43,24 +43,29 @@ class PEFunction : public BaseFunction {
 
   ~PEFunction() noexcept {}
 
-  std::vector<Variable> operator()(const std::vector<Variable> &inputs) {
-    // bool is_test = true;
+  std::vector<Tensor> operator()(const std::vector<Tensor> &inputs) {
+    auto dense_tensors = utils::ToDenseTensors(inputs);
+    return utils::ToTensors(this->operator()(dense_tensors));
+  }
+
+  std::vector<DenseTensor> operator()(const std::vector<DenseTensor> &inputs) {
     std::string prog_string;
     std::hash<std::string> string_hash;
     auto &program_desc = info_->ProgramDesc();
+    // TODO(dev): Serialize is very slow.
     const_cast<framework::ProgramDesc *>(&program_desc)
         ->Proto()
         ->SerializePartialToString(&prog_string);
-    // program_desc.Proto()->SerializePartialToString(&prog_string);
+
     int64_t program_id = static_cast<int64_t>(string_hash(prog_string));
     const framework::BlockDesc &global_block = program_desc.Block(0);
     int64_t start_op_index = 0;
     int64_t end_op_index = static_cast<int64_t>(global_block.OpSize());
 
-    utils::ShareInputsIntoScope(info_->InputArgNames(), inputs, &scope_);
+    utils::ShareIntoScope(info_->InputArgNames(), inputs, &scope_);
     std::vector<std::string> input_var_names = info_->InputArgNames();
     std::vector<std::string> output_var_names = info_->OutputArgNames();
-    std::vector<std::string> dout_var_names;
+
     if (end_op_index > start_op_index) {
       auto cache_info = framework::GetExecutorInfoFromCache(program_desc,
                                                             place_,
@@ -78,9 +83,7 @@ class PEFunction : public BaseFunction {
         skip_eager_delete_vars.insert(skip_eager_delete_vars.end(),
                                       output_var_names.begin(),
                                       output_var_names.end());
-        skip_eager_delete_vars.insert(skip_eager_delete_vars.end(),
-                                      dout_var_names.begin(),
-                                      dout_var_names.end());
+
         framework::details::ParseSafeEagerDeletionSkipVars(
             program_desc,
             end_op_index,
@@ -89,9 +92,8 @@ class PEFunction : public BaseFunction {
       }
       parallel_executor->RunWithoutFetch(skip_eager_delete_vars);
     }
-    VLOG(6) << framework::GenScopeTreeDebugInfo(&scope_);
-    std::vector<Variable> res;
-    utils::FetchVarsByNames(info_->OutputArgNames(), scope_, &res);
+    std::vector<DenseTensor> res;
+    utils::FetchOuts(info_->OutputArgNames(), scope_, &res);
     return res;
   }
 

paddle/fluid/pybind/eager_functions.cc

@@ -357,6 +357,19 @@ static std::vector<paddle::any> CastAttrsToTragetType(
   return res;
 }
 
+static PyObject* eager_api_jit_function_call(PyObject* self,
+                                             PyObject* args,
+                                             PyObject* kwargs) {
+  EAGER_TRY
+  std::shared_ptr<jit::BaseFunction> function =
+      CastPyArg2BaseFunction(PyTuple_GET_ITEM(args, 0), 0);
+  std::vector<paddle::experimental::Tensor> ins =
+      CastPyArg2VectorOfTensor(PyTuple_GET_ITEM(args, 1), 1);
+  std::vector<paddle::experimental::Tensor> outs = (*function)(ins);
+  return ToPyObject(outs);
+  EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 static PyObject* eager_api_run_costum_op(PyObject* self,
                                          PyObject* args,
                                          PyObject* kwargs) {
@@ -911,6 +924,10 @@ PyMethodDef variable_functions[] = {
      (PyCFunction)(void (*)(void))eager_api_read_next_tensor_list,
      METH_VARARGS | METH_KEYWORDS,
      NULL},
+    {"jit_function_call",
+     (PyCFunction)(void (*)(void))eager_api_jit_function_call,
+     METH_VARARGS | METH_KEYWORDS,
+     NULL},
     /**sparse functions**/
     {"sparse_coo_tensor",
      (PyCFunction)(void (*)(void))eager_api_sparse_coo_tensor,

paddle/fluid/pybind/eager_utils.cc

@@ -51,6 +51,7 @@ extern PyTypeObject* g_customplace_pytype;
 extern PyTypeObject* g_framework_tensor_pytype;
 extern PyTypeObject* g_framework_lodtensorarray_pytype;
 extern PyTypeObject* g_custom_op_kernel_ctx_pytype;
+extern PyTypeObject* g_executor_function_pytype;
 
 int TensorDtype2NumpyDtype(phi::DataType dtype) {
   switch (dtype) {
@@ -227,6 +228,21 @@ std::shared_ptr<imperative::VarBase> CastPyArg2VarBase(PyObject* obj,
   return py::cast<std::shared_ptr<imperative::VarBase>>(obj);
 }
 
+std::shared_ptr<jit::BaseFunction> CastPyArg2BaseFunction(PyObject* obj,
+                                                          ssize_t arg_pos) {
+  if (PyObject_IsInstance(
+          obj, reinterpret_cast<PyObject*>(g_executor_function_pytype))) {
+    return ::pybind11::handle(obj)
+        .cast<std::shared_ptr<jit::ExecutorFunction>>();
+  } else {
+    PADDLE_THROW(platform::errors::InvalidArgument(
+        "argument (position %d) must be "
+        "BaseFunction, but got %s",
+        arg_pos + 1,
+        reinterpret_cast<PyTypeObject*>(obj->ob_type)->tp_name));
+  }
+}
+
 std::vector<paddle::experimental::Tensor> CastPyArg2VectorOfTensor(
     PyObject* obj, ssize_t arg_pos) {
   std::vector<paddle::experimental::Tensor> result;

paddle/fluid/pybind/eager_utils.h

@@ -19,6 +19,7 @@ typedef SSIZE_T ssize_t;
 
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/jit/executor_function.h"
 #include "paddle/fluid/platform/place.h"
 #include "paddle/phi/common/backend.h"
 #include "paddle/phi/common/data_type.h"
@@ -72,6 +73,8 @@ framework::proto::VarType::Type CastPyArg2ProtoType(PyObject* obj,
 std::unordered_map<std::wstring, int> CastPyArg2Vocab(PyObject* obj,
                                                       ssize_t arg_pos);
 std::vector<std::string> CastPyArg2Strings(PyObject* obj, ssize_t arg_pos);
+std::shared_ptr<jit::BaseFunction> CastPyArg2BaseFunction(PyObject* obj,
+                                                          ssize_t arg_pos);
 
 PyObject* ToPyObject(int value);
 PyObject* ToPyObject(uint32_t value);

paddle/fluid/pybind/jit.cc

@@ -28,39 +28,21 @@ namespace py = pybind11;
 namespace paddle {
 namespace pybind {
 
+PyTypeObject *g_executor_function_pytype = nullptr;
 using Variable = paddle::framework::Variable;
 
 void BindJit(pybind11::module *m) {
   py::class_<jit::Layer>(*m, "Layer", R"DOC(Layer Class.)DOC")
-      .def("function_dict", &jit::Layer::FunctionMap);
-
-  py::class_<jit::ExecutorFunction, std::shared_ptr<jit::ExecutorFunction>>(
-      *m, "ExectorFunction", R"DOC(ExectorFunction Class.)DOC")
-      .def("__call__",
-           [](jit::ExecutorFunction &self,
-              const std::vector<std::shared_ptr<imperative::VarBase>>
-                  &tensor_inputs) {
-             std::vector<Variable> var_inputs;
-             for (auto &tensor : tensor_inputs) {
-               var_inputs.emplace_back(tensor->Var());
-             }
-             auto var_outputs = self(var_inputs);
-
-             std::vector<std::shared_ptr<imperative::VarBase>> tensor_outputs;
-             auto output_names = self.Info()->OutputArgNames();
-             for (size_t i = 0; i < var_outputs.size(); ++i) {
-               auto var = var_outputs[i];
-               std::string name = output_names[i];
-               imperative::VariableWrapper var_wrapper(name, var);
-               auto shared_wrapper =
-                   std::make_shared<imperative::VariableWrapper>(var_wrapper);
-               auto shared_varbase =
-                   std::make_shared<imperative::VarBase>(shared_wrapper);
-               tensor_outputs.emplace_back(shared_varbase);
-             }
-             return tensor_outputs;
-           })
-      .def("info", &jit::ExecutorFunction::Info);
+      .def("function_dict",
+           &jit::Layer::FunctionMap,
+           py::return_value_policy::reference);
+
+  py::class_<jit::ExecutorFunction, std::shared_ptr<jit::ExecutorFunction>>
+      executor_function(
+          *m, "ExectorFunction", R"DOC(ExectorFunction Class.)DOC");
+  g_executor_function_pytype =
+      reinterpret_cast<PyTypeObject *>(executor_function.ptr());
+  executor_function.def("info", &jit::ExecutorFunction::Info);
 
   py::class_<jit::FunctionInfo, std::shared_ptr<jit::FunctionInfo>>(
       *m, "FunctionInfo", R"DOC(FunctionInfo Class.)DOC")

python/paddle/fluid/tests/unittests/test_jit_layer.py

@@ -22,7 +22,6 @@ from paddle.fluid.framework import _enable_legacy_dygraph
 from paddle.jit.layer import Layer
 from paddle.fluid.dygraph.dygraph_to_static.program_translator import ProgramTranslator
 
-_enable_legacy_dygraph()
 paddle.seed(1)
 
 

python/paddle/jit/layer.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from paddle.fluid import core
 from paddle.fluid.core import Load
 
 
@@ -39,7 +40,7 @@ class Function():
         self.info = FunctionInfo(function.info())
 
     def __call__(self, *args):
-        return self.function(args)
+        return core.eager.jit_function_call(self.function, args)
 
 
 class FunctionInfo():