[PIR][Dy2St] Convert Tensor to Value in PIR bind API (#59761)

SigureMo · web-flow · commit 1266da13814d · 2023-12-15T10:20:52.000+08:00
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -1916,7 +1916,7 @@ paddle::Tensor CreateTensorFromVarDesc(
   return tensor;
 }
 
-PyObject* GetEmpytyTensorsWithVarDesc(PyObject* self, PyObject* args) {
+PyObject* GetEmptyTensorsWithVarDesc(PyObject* self, PyObject* args) {
   std::vector<paddle::Tensor> result;
   std::unordered_map<std::string, paddle::Tensor> out_tensor_map;
 
@@ -2097,9 +2097,24 @@ paddle::experimental::Scalar CastNumpy2Scalar(PyObject* obj,
   }
 }
 
+PyObject* CastPyArg2ValuePreHook(PyObject* obj) {
+  PyObject* hook = static_op_arg_pre_cast_hook_get();
+  if (hook == Py_None) {
+    return obj;
+  }
+  Py_INCREF(obj);
+  PyObject* result = PyObject_CallFunction(hook, "O", obj);
+  PADDLE_ENFORCE(result,
+                 paddle::platform::errors::Fatal(
+                     "Call static_op_arg_pre_cast_hook failed."));
+  Py_DECREF(obj);
+  return result;
+}
+
 pir::Value CastPyArg2Value(PyObject* obj,
                            const std::string& op_type,
                            size_t arg_pos) {
+  obj = CastPyArg2ValuePreHook(obj);
   if (PyObject_TypeCheck(obj, g_ir_opresult_pytype)) {
     return ::pybind11::handle(obj).cast<pir::OpResult>();
   } else if (PyObject_TypeCheck(obj, g_ir_value_pytype)) {
@@ -2133,6 +2148,7 @@ std::vector<pir::Value> CastPyArg2VectorOfValue(PyObject* obj,
     PyObject* item = nullptr;
     for (Py_ssize_t i = 0; i < len; i++) {
       item = PyList_GetItem(obj, i);
+      item = CastPyArg2ValuePreHook(item);
       if (PyObject_TypeCheck(item, g_ir_opresult_pytype)) {
         value_list.emplace_back(::pybind11::handle(item).cast<pir::OpResult>());
       } else if (item == Py_None) {
@@ -2152,6 +2168,7 @@ std::vector<pir::Value> CastPyArg2VectorOfValue(PyObject* obj,
     PyObject* item = nullptr;
     for (Py_ssize_t i = 0; i < len; i++) {
       item = PyTuple_GetItem(obj, i);
+      item = CastPyArg2ValuePreHook(item);
       if (PyObject_TypeCheck(item, g_ir_opresult_pytype)) {
         value_list.emplace_back(::pybind11::handle(item).cast<pir::OpResult>());
       } else if (item == Py_None) {
@@ -2581,6 +2598,50 @@ void* UnPackHook::operator()(void* packed_value, void* other) {
   return reinterpret_cast<void*>(ret);
 }
 
+/* ------------------ for SetStaticOpArgPreCastHook ----------------------- */
+
+static Py_tss_t static_op_arg_pre_cast_hook_key = {0, 0};
+
+inline static PyObject* static_op_arg_pre_cast_hook_get() {
+  void* result = PyThread_tss_get(&static_op_arg_pre_cast_hook_key);
+  if (result == NULL) {
+    return Py_None;
+  } else {
+    return reinterpret_cast<PyObject*>(result);
+  }
+}
+
+inline static void static_op_arg_pre_cast_hook_set(PyObject* obj) {
+  PyThread_tss_set(&static_op_arg_pre_cast_hook_key, obj);
+}
+
+static PyObject* set_static_op_arg_pre_cast_hook(PyObject* new_callback,
+                                                 PyThreadState* tstate) {
+  PyObject* old_callback = static_op_arg_pre_cast_hook_get();
+  Py_INCREF(new_callback);
+  static_op_arg_pre_cast_hook_set(new_callback);
+
+  return old_callback;
+}
+
+PyObject* SetStaticOpArgPreCastHook(PyObject* dummy, PyObject* callback) {
+  if (callback != Py_None && !PyCallable_Check(callback)) {
+    VLOG(7) << "callback is not a callable or none, invalid arguments.";
+    Py_INCREF(Py_None);
+    return Py_None;
+  }
+  return set_static_op_arg_pre_cast_hook(callback, PyThreadState_GET());
+}
+
+PyMODINIT_FUNC PyInit__static_op_arg_pre_cast_hook() {
+  auto result = PyThread_tss_create(&static_op_arg_pre_cast_hook_key);
+  VLOG(7) << "Set PyThread_tss_create return: " << result;
+
+  Py_INCREF(Py_None);
+  static_op_arg_pre_cast_hook_set(Py_None);
+  return NULL;
+}
+
 /* ------------------ for auto parallel ----------------------- */
 
 void DistTensorTypeParser::operator()(const Tensor& x) {
@@ -2671,16 +2732,21 @@ void DistTensorConverter::operator()(paddle::optional<std::vector<Tensor>>* x) {
 
 static PyMethodDef EagerUtilMethods[] = {
     {"create_empty_tensors_with_var_descs",
-     (PyCFunction)(void (*)(void))GetEmpytyTensorsWithVarDesc,
+     (PyCFunction)(void (*)(void))GetEmptyTensorsWithVarDesc,
      METH_VARARGS,
-     "GetEmpytyTensorsWithVarDesc"},
+     "GetEmptyTensorsWithVarDesc"},
     {"create_empty_tensors_with_op_results",
      (PyCFunction)(void (*)(void))GetEmpytyTensorsWithOpResult,
      METH_VARARGS,
      "GetEmpytyTensorsWithOpResult."},
+    {"set_static_op_arg_pre_cast_hook",
+     (PyCFunction)SetStaticOpArgPreCastHook,
+     METH_O,
+     "Set hook for pre cast a static OP argument."},
     {nullptr, nullptr, 0, nullptr}};
 
 void BindEagerUtils(PyObject* module) {
+  PyInit__static_op_arg_pre_cast_hook();
   if (PyModule_AddFunctions(module, EagerUtilMethods) < 0) {
     PADDLE_THROW(platform::errors::Fatal(
         "Init Paddle error in BindEagerUtils(PyModule_AddFunctions)."));
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
@@ -418,6 +418,19 @@ class eager_gil_scoped_release {
   PyThreadState* tstate{nullptr};
 };
 
+/* ------------------ for SetStaticOpArgPreCastHook ----------------------- */
+
+inline static PyObject* static_op_arg_pre_cast_hook_get();
+
+inline static void static_op_arg_pre_cast_hook_set(PyObject* obj);
+
+static PyObject* set_static_op_arg_pre_cast_hook(PyObject* new_callback,
+                                                 PyThreadState* tstate);
+
+PyObject* SetStaticOpArgPreCastHook(PyObject* callback);
+
+PyMODINIT_FUNC PyInit__static_op_arg_pre_cast_hook();
+
 /* ------------------ for auto parallel ----------------------- */
 using paddle::experimental::detail::ArgsIterator;
 
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1181,7 +1181,9 @@ SplitedResult SplitForwardBackward(
             dtype,
             place);
     counter += 1;
-    backward_value_map[v] = op->results()[0].Value::impl();
+    pir::Value target = op->results()[0].Value::impl();
+    target.set_type(v.type());
+    backward_value_map[v] = target;
   };
 
   auto create_output_fn_forward = [&ctx,
diff --git a/python/paddle/base/dygraph/base.py b/python/paddle/base/dygraph/base.py
@@ -121,7 +121,7 @@ def _convert_into_variable(tensor):
     Convert Tensor into Variable.
     """
     if paddle.framework.use_pir_api():
-        return paddle.pir.core._convert_into_opresult(tensor)
+        return paddle.pir.core._convert_into_value(tensor)
     if isinstance(tensor, core.eager.Tensor):
         # Check whether has been created before.
         new_var = tensor.block._find_var_recursive(tensor.name)
diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py
@@ -33,6 +33,7 @@
 from paddle.framework import in_dynamic_mode, use_pir_api
 from paddle.nn.layer import layers
 from paddle.pir import Value
+from paddle.pir.core import _convert_into_value, static_op_arg_cast_guard
 from paddle.utils import flatten, gast
 
 from . import error, logging_utils
@@ -1202,7 +1203,9 @@ def pir_from_func_spec(
         # framework.default_startup_program().random_seed
         # ) }}}
         with ir_static.program_guard(main_program, startup_program):
-            with _to_static_mode_guard_(is_to_static=True):
+            with _to_static_mode_guard_(
+                is_to_static=True
+            ), static_op_arg_cast_guard(_convert_into_value):
                 # 1. Adds `paddle.static.data` layers for input if needed
                 static_inputs = func_spec.pir_to_static_inputs_with_spec(
                     input_spec, main_program
diff --git a/python/paddle/pir/core.py b/python/paddle/pir/core.py
@@ -26,6 +26,7 @@
 
 from .._pir_ops import parameter, set_parameter
 from ..base import unique_name
+from ..base.core import set_static_op_arg_pre_cast_hook
 from ..base.wrapped_decorator import signature_safe_contextmanager
 
 vartype_to_datatype = {
@@ -304,27 +305,30 @@ def create_parameter(
     return param
 
 
-def _convert_into_opresult(tensor):
+def _convert_into_value(tensor):
     """
     Convert Tensor into OpResult.
     """
     import paddle
-    from paddle.base import core, framework
     from paddle.jit.pir_dy2static.parameter_recorder import (
         _global_parameter_recorder,
     )
 
-    if isinstance(tensor, core.eager.Tensor):
-        # Check whether has been created before.
-        new_var = tensor.block._find_var_recursive(tensor.name)
-        is_persistable = True
-        if new_var is not None:
-            assert isinstance(new_var, framework.Variable)
-        else:
-            new_var = _global_parameter_recorder.get(
-                paddle.pir.core.default_main_program(), tensor
-            )
-        # add param into parameter recorder to collect all the params used in this program.
-        return new_var
-    else:
-        return tensor
+    if isinstance(tensor, paddle.Tensor):
+        return _global_parameter_recorder.get(
+            paddle.pir.core.default_main_program(), tensor
+        )
+    return tensor
+
+
+@signature_safe_contextmanager
+def static_op_arg_cast_guard(hook):
+    """
+    Set a hook function to cast the arguments of static op.
+    """
+
+    original_callback = set_static_op_arg_pre_cast_hook(hook)
+    try:
+        yield
+    finally:
+        set_static_op_arg_pre_cast_hook(original_callback)
diff --git a/test/dygraph_to_static/test_set_static_op_arg_pre_cast_hook.py b/test/dygraph_to_static/test_set_static_op_arg_pre_cast_hook.py
@@ -0,0 +1,77 @@
+# Copyright (c) 2020 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.
+
+import unittest
+
+import numpy as np
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    enable_to_static_guard,
+    static_guard,
+    test_ast_only,
+    test_pir_only,
+)
+
+import paddle
+from paddle.pir.core import _convert_into_value, static_op_arg_cast_guard
+
+
+class TestSetStaticOpArgPreCastHook(Dy2StTestBase):
+    @test_ast_only
+    @test_pir_only
+    def test_set_static_op_arg_pre_cast_hook(self):
+        eager_tensor = paddle.rand((10, 10), 'float32')
+
+        with static_guard():
+            with self.assertRaisesRegex(
+                ValueError,
+                r"abs\(\): argument \(position 1\) must be OpResult, but got Tensor",
+            ):
+                paddle.abs(eager_tensor)
+
+            with static_op_arg_cast_guard(_convert_into_value):
+                paddle.abs(eager_tensor)
+
+
+class NetWithEagerTensor(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.extra_inputs = []
+
+    # Disable the _jst.Ld to avoid convert the Eager Tensor to Value
+    # at transform time.
+    @paddle.jit.not_to_static
+    def forward_impl(self, x):
+        return paddle.concat([x] + self.extra_inputs, axis=0)
+
+    def forward(self, x):
+        return self.forward_impl(x)
+
+
+class TestSetStaticOpArgPreCastHookWithEagerTensor(Dy2StTestBase):
+    @test_ast_only
+    @test_pir_only
+    def test_net_with_eager_tensor(self):
+        net = NetWithEagerTensor()
+        net.extra_inputs.append(paddle.rand((10, 10), 'float32'))
+        net = paddle.jit.to_static(net, full_graph=True)
+        x = paddle.rand((10, 10), 'float32')
+        with enable_to_static_guard(False):
+            dygraph_out = net(x)
+        static_out = net(x)
+        np.testing.assert_allclose(dygraph_out.numpy(), static_out.numpy())
+
+
+if __name__ == '__main__':
+    unittest.main()
