From 62c0abac80a35b5b1c0c5648d2d4ee5bb9dfb080 Mon Sep 17 00:00:00 2001 From: Weilong Wu Date: Tue, 18 Oct 2022 10:24:01 +0800 Subject: [PATCH] [Eager, Performance optimization] support pow( ** operator) to sink to Cpp layer (#47077) --- paddle/fluid/pybind/eager_math_op_patch.cc | 180 +++++++++++++++++++ python/paddle/fluid/dygraph/math_op_patch.py | 6 +- 2 files changed, 182 insertions(+), 4 deletions(-) diff --git a/paddle/fluid/pybind/eager_math_op_patch.cc b/paddle/fluid/pybind/eager_math_op_patch.cc index 4cd76f4a10..521a8f9f74 100644 --- a/paddle/fluid/pybind/eager_math_op_patch.cc +++ b/paddle/fluid/pybind/eager_math_op_patch.cc @@ -143,6 +143,8 @@ paddle::experimental::Tensor CallScalarFuction( ret = scale_ad_func(self_tensor, phi::Scalar(other), 0.0, true); } else if (op_type == "div") { ret = scale_ad_func(self_tensor, phi::Scalar(1.0 / other), 0.0, true); + } else if (op_type == "pow") { + ret = pow_ad_func(self_tensor, other); } return ret; @@ -1385,6 +1387,176 @@ static PyObject* tensor__floordiv__method(TensorObject* self, EAGER_CATCH_AND_THROW_RETURN_NULL } +static PyObject* tensor__pow__method(TensorObject* self, + PyObject* args, + PyObject* kwargs) { + paddle::platform::RecordEvent pythonc_record_event( + "pow pybind_patch_func", + paddle::platform::TracerEventType::UserDefined, + 1); + + EAGER_TRY + VLOG(6) << "Running Eager tensor__pow__method"; + + // Set Device ID + auto place = egr::Controller::Instance().GetExpectedPlace(); + SetDevice(place); + + paddle::experimental::Tensor ret; + paddle::experimental::Tensor self_tensor = self->tensor; + + PyObject* other_obj = PyTuple_GET_ITEM(args, 0); + + // 1. scalar exists cases + if (PyFloat_Check(other_obj) || PyCheckInteger(other_obj) || + IsNumpyType(other_obj)) { + double other = 0.0; + if (PyFloat_Check(other_obj)) { + other = CastPyArg2Double(other_obj, "__pow__", 0); + if (_supported_int_dtype_.find(self_tensor.dtype()) != + _supported_int_dtype_.end()) { + eager_gil_scoped_release guard; + self_tensor = cast_ad_func(self_tensor, DataType::FLOAT32); + } + } else if (PyCheckInteger(other_obj) || IsNumpyType(other_obj)) { + other = CastPyArg2Double(other_obj, "__pow__", 0); + } + { + eager_gil_scoped_release guard; + ret = CallScalarFuction(self_tensor, other, "pow"); + } + return ToPyObject(ret); + } + + // 2. create or get tensor for other_obj + paddle::experimental::Tensor other_tensor; + if (!PyCheckTensor(other_obj)) { + paddle::experimental::Scalar value = + CastPyArg2Scalar(other_obj, "__pow__", 0); + if (PyComplex_Check(other_obj)) { + eager_gil_scoped_release guard; + other_tensor = + full_ad_func({1}, value, DataType::COMPLEX64, self_tensor.place()); + } else { + eager_gil_scoped_release guard; + other_tensor = + full_ad_func({1}, value, self_tensor.dtype(), self_tensor.place()); + } + } else { + other_tensor = CastPyArg2Tensor(other_obj, 0); + } + + // 3. promote types or unify right var type to left var + phi::DataType lhs_dtype = self_tensor.dtype(); + phi::DataType rhs_dtype = other_tensor.dtype(); + if (lhs_dtype != rhs_dtype) { + VLOG(6) << "The dtype of left and right Tensor are not the same, left " + "dtype is " + << lhs_dtype << ", but right dtype is " << rhs_dtype + << ", the right dtype will convert to " << lhs_dtype; + eager_gil_scoped_release guard; + other_tensor = cast_ad_func(other_tensor, lhs_dtype); + } + + // 4. calculation + VLOG(6) << "Calling elementwise_pow_ad_func in tensor__pow__method"; + { + eager_gil_scoped_release guard; + ret = elementwise_pow_ad_func(self_tensor, other_tensor); + } + + return ToPyObject(ret); + EAGER_CATCH_AND_THROW_RETURN_NULL +} + +static PyObject* tensor__rpow__method(TensorObject* self, + PyObject* args, + PyObject* kwargs) { + paddle::platform::RecordEvent pythonc_record_event( + "__rpow__ pybind_patch_func", + paddle::platform::TracerEventType::UserDefined, + 1); + + EAGER_TRY + VLOG(6) << "Running Eager tensor__rpow__method"; + + // Set Device ID + auto place = egr::Controller::Instance().GetExpectedPlace(); + SetDevice(place); + + paddle::experimental::Tensor ret; + paddle::experimental::Tensor self_tensor = self->tensor; + + PyObject* other_obj = PyTuple_GET_ITEM(args, 0); + + // 1. scalar exists cases or not + // there is no scalar case for rpow, but alse need to cast self_tensor in + // need. + double other_double = 0.0; + bool has_other_double = false; + if (PyFloat_Check(other_obj) || PyCheckInteger(other_obj) || + IsNumpyType(other_obj)) { + if (PyFloat_Check(other_obj)) { + other_double = CastPyArg2Double(other_obj, "__rpow__", 0); + has_other_double = true; + if (_supported_int_dtype_.find(self_tensor.dtype()) != + _supported_int_dtype_.end()) { + eager_gil_scoped_release guard; + self_tensor = cast_ad_func(self_tensor, DataType::FLOAT32); + } + } else if (PyCheckInteger(other_obj) || IsNumpyType(other_obj)) { + other_double = CastPyArg2Double(other_obj, "__rpow__", 0); + has_other_double = true; + } + } + + // 2. create or get tensor for other_obj + paddle::experimental::Tensor other_tensor; + if (has_other_double) { + eager_gil_scoped_release guard; + other_tensor = full_ad_func(self_tensor.shape(), + phi::Scalar(other_double), + self_tensor.dtype(), + self_tensor.place()); + } else if (!PyCheckTensor(other_obj)) { + paddle::experimental::Scalar value = + CastPyArg2Scalar(other_obj, "__rpow__", 0); + if (PyComplex_Check(other_obj)) { + eager_gil_scoped_release guard; + other_tensor = + full_ad_func({1}, value, DataType::COMPLEX64, self_tensor.place()); + } else { + eager_gil_scoped_release guard; + other_tensor = full_ad_func( + self_tensor.shape(), value, self_tensor.dtype(), self_tensor.place()); + } + } else { + other_tensor = CastPyArg2Tensor(other_obj, 0); + } + + // 3. promote types or unify right var type to left var + phi::DataType lhs_dtype = self_tensor.dtype(); + phi::DataType rhs_dtype = other_tensor.dtype(); + if (lhs_dtype != rhs_dtype) { + VLOG(6) << "The dtype of left and right Tensor are not the same, left " + "dtype is " + << lhs_dtype << ", but right dtype is " << rhs_dtype + << ", the right dtype will convert to " << lhs_dtype; + eager_gil_scoped_release guard; + other_tensor = cast_ad_func(other_tensor, lhs_dtype); + } + + // 4. calculation + VLOG(6) << "Calling elementwise_pow_ad_func in tensor__rpow__method"; + { + eager_gil_scoped_release guard; + ret = elementwise_pow_ad_func(other_tensor, self_tensor); + } + + return ToPyObject(ret); + EAGER_CATCH_AND_THROW_RETURN_NULL +} + PyMethodDef math_op_patch_methods[] = { {"__add__", (PyCFunction)(void (*)(void))tensor__add__method, @@ -1430,6 +1602,14 @@ PyMethodDef math_op_patch_methods[] = { (PyCFunction)(void (*)(void))tensor__floordiv__method, METH_VARARGS | METH_KEYWORDS, NULL}, + {"__pow__", + (PyCFunction)(void (*)(void))tensor__pow__method, + METH_VARARGS | METH_KEYWORDS, + NULL}, + {"__rpow__", + (PyCFunction)(void (*)(void))tensor__rpow__method, + METH_VARARGS | METH_KEYWORDS, + NULL}, {"__mod__", (PyCFunction)(void (*)(void))tensor__mod__method, METH_VARARGS | METH_KEYWORDS, diff --git a/python/paddle/fluid/dygraph/math_op_patch.py b/python/paddle/fluid/dygraph/math_op_patch.py index f754fb93c4..c5514b55ef 100644 --- a/python/paddle/fluid/dygraph/math_op_patch.py +++ b/python/paddle/fluid/dygraph/math_op_patch.py @@ -388,10 +388,6 @@ def monkey_patch_math_varbase(): ('ndim', _ndim_), ('size', _size_), ('T', _T_), - ('__pow__', _binary_creator_('__pow__', 'pow', False, _C_ops.pow, - True)), - ('__rpow__', _binary_creator_('__rpow__', 'elementwise_pow', True, - None)), # for logical compare ('__eq__', _binary_creator_('__eq__', 'equal', False, None, True)), ('__ne__', _binary_creator_('__ne__', 'not_equal', False, None, True)), @@ -416,6 +412,8 @@ def monkey_patch_math_varbase(): '__lt__', '__le__', '__floordiv__', + '__pow__', + '__rpow__', ] global _already_patch_varbase -- GitLab