[Extension Operants] Extension supports tensor operants

Author: jiahy0825

paddle/fluid/pybind/eager_functions.cc

@@ -499,15 +499,18 @@ static PyObject* eager_api_jit_function_call(PyObject* self,
   EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
-static PyObject* eager_api_init_eager_and_static_tensor_operants(
-    PyObject* self, PyObject* args, PyObject* kwargs) {
+static PyObject* eager_api_init_tensor_operants(PyObject* self,
+                                                PyObject* args,
+                                                PyObject* kwargs) {
   EAGER_TRY
 
   paddle::OperantsManager::Instance().eager_operants.reset(
       new paddle::prim::EagerTensorOperants());
   paddle::OperantsManager::Instance().static_operants.reset(
       new paddle::prim::StaticTensorOperants());
-  VLOG(4) << "Initialize eager and static tensor operants successfully";
+  paddle::OperantsManager::Instance().phi_operants.reset(
+      new paddle::operants::PhiTensorOperants());
+  VLOG(4) << "Initialize tensor operants successfully";
 
   RETURN_PY_NONE
   EAGER_CATCH_AND_THROW_RETURN_NULL
@@ -1123,9 +1126,8 @@ PyMethodDef variable_functions[] = {
      (PyCFunction)(void (*)(void))eager_api_run_custom_op,
      METH_VARARGS | METH_KEYWORDS,
      NULL},
-    {"_init_eager_and_static_tensor_operants",
-     (PyCFunction)(void (*)(
-         void))eager_api_init_eager_and_static_tensor_operants,
+    {"_init_tensor_operants",
+     (PyCFunction)(void (*)(void))eager_api_init_tensor_operants,
      METH_VARARGS | METH_KEYWORDS,
      NULL},
     {"tensor_copy",

paddle/utils/CMakeLists.txt

@@ -17,5 +17,5 @@ if(NOT ((NOT WITH_PYTHON) AND ON_INFER))
   cc_library(
     pybind_util
     SRCS pybind.cc
-    DEPS phi_tensor_raw)
+    DEPS phi_tensor_raw flags)
 endif()

paddle/utils/pybind.cc

@@ -13,8 +13,11 @@
 // limitations under the License.
 
 #include "paddle/utils/pybind.h"
+
+#include "gflags/gflags.h"
 #include "paddle/phi/core/enforce.h"
 
+DECLARE_string(tensor_operants_mode);
 namespace paddle {
 namespace pybind {
 
@@ -66,5 +69,7 @@ PyObject* ToPyObject(const paddle::experimental::Tensor& value,
   return obj;
 }
 
+void SwitchTensorOperantsMode() { FLAGS_tensor_operants_mode = "phi"; }
+
 }  // namespace pybind
 }  // namespace paddle

paddle/utils/pybind.h

@@ -46,6 +46,9 @@ paddle::experimental::Tensor CastPyArg2Tensor(PyObject* obj, ssize_t arg_pos);
 PyObject* ToPyObject(const paddle::experimental::Tensor& value,
                      bool return_py_none_if_not_initialize = false);
 
+// Internal use only, switch tensor_operants_mode to phi
+void SwitchTensorOperantsMode();
+
 }  // namespace pybind
 }  // namespace paddle
 
@@ -59,6 +62,7 @@ struct type_caster<paddle::experimental::Tensor> {
                        _("paddle::experimental::Tensor"));
 
   bool load(handle src, bool) {
+    paddle::pybind::SwitchTensorOperantsMode();
     PyObject* obj = src.ptr();
     if (paddle::pybind::PyCheckTensor(obj)) {
       value = paddle::pybind::CastPyArg2Tensor(obj, 0);

python/paddle/fluid/__init__.py

@@ -231,7 +231,7 @@ def remove_flag_if_exists(name):
         core.init_glog(sys.argv[0])
     # don't init_p2p when in unittest to save time.
     core.init_devices()
-    core.eager._init_eager_and_static_tensor_operants()
+    core.eager._init_tensor_operants()
     core.init_default_kernel_signatures()
     core.init_memory_method()
 

python/paddle/fluid/tests/cpp_extension/custom_extension.cc

@@ -21,7 +21,7 @@
 paddle::Tensor custom_sub(paddle::Tensor x, paddle::Tensor y);
 
 paddle::Tensor custom_add(const paddle::Tensor& x, const paddle::Tensor& y) {
-  return paddle::add(paddle::exp(x), paddle::exp(y));
+  return x.exp() + y.exp();
 }
 
 paddle::Tensor nullable_tensor(bool return_none = false) {