[PIR]Migrate fill_constant into pir (#57020)

0x45f · web-flow · commit 99831e0dc489 · 2023-09-16T15:30:42.000+08:00
* Call _C_ops.full in new ir

* Fix shape attr

* Fix yaml and infermeta

* Refine code

* Use full with tensor op

* Refine code

* Remove breakpoint

* Fix full

* Call stack when use full_with_tensor

* Refine code

* Refine code

* Refine code
diff --git a/paddle/fluid/pir/dialect/op_generator/python_c_gen.py b/paddle/fluid/pir/dialect/op_generator/python_c_gen.py
@@ -212,6 +212,8 @@
     "std::vector<double>": "FLOAT64",
 }
 
+MANUAL_STATIC_OP_FUNCTION_LIST = ['full']
+
 
 class PythonCCodeGen(CodeGen):
     def __init__(self) -> None:
@@ -414,6 +416,12 @@ def _gen_one_impl(self, op_info, op_name):
         ret = re.sub(r' +\n', '', ret)
         return ret
 
+    def _need_skip(self, op_info, op_name):
+        return (
+            super()._need_skip(op_info, op_name)
+            or op_name in MANUAL_STATIC_OP_FUNCTION_LIST
+        )
+
     def _gen_cpp_file(self, op_info_items, namespaces, cpp_file_path):
         impl_str = ''
         for op_info in op_info_items:
diff --git a/paddle/fluid/pybind/manual_static_op_function.h b/paddle/fluid/pybind/manual_static_op_function.h
@@ -15,6 +15,7 @@
 #pragma once
 
 #include "paddle/fluid/pir/dialect/operator/ir/manual_api.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
 #include "paddle/fluid/pybind/eager_utils.h"
 #include "paddle/fluid/pybind/exception.h"
 #include "paddle/fluid/pybind/op_function_common.h"
@@ -73,6 +74,63 @@ static PyObject *static_api_set_parameter(PyObject *self,
   }
 }
 
+PyObject *static_api_full(PyObject *self, PyObject *args, PyObject *kwargs) {
+  try {
+    VLOG(6) << "Add full op into program";
+    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);
+
+    // Parse Attributes
+    PyObject *shape_obj = PyTuple_GET_ITEM(args, 0);
+    PyObject *value_obj = PyTuple_GET_ITEM(args, 1);
+    PyObject *dtype_obj = PyTuple_GET_ITEM(args, 2);
+    PyObject *place_obj = PyTuple_GET_ITEM(args, 3);
+
+    phi::DataType dtype = CastPyArg2DataTypeDirectly(dtype_obj, "full", 2);
+    Place place = CastPyArg2Place(place_obj, "full", 3);
+
+    if (!PyObject_CheckIROpResult(shape_obj) &&
+        !PyObject_CheckIRVectorOfOpResult(shape_obj) &&
+        !PyObject_CheckIROpResult(value_obj)) {
+      std::vector<int64_t> shape = CastPyArg2Longs(shape_obj, "full", 0);
+      float value = CastPyArg2Float(value_obj, "full", 1);
+      auto static_api_out = paddle::dialect::full(shape, value, dtype, place);
+      return ToPyObject(static_api_out);
+    } else {
+      pir::OpResult shape;
+      pir::OpResult value;
+
+      if (PyObject_CheckIROpResult(shape_obj)) {
+        shape = CastPyArg2OpResult(shape_obj, "full", 0);
+      } else if (PyObject_CheckIRVectorOfOpResult(shape_obj)) {
+        std::vector<pir::OpResult> shape_tmp =
+            CastPyArg2VectorOfOpResult(shape_obj, "full", 0);
+        shape = paddle::dialect::stack(shape_tmp, 0);
+      } else {
+        std::vector<int64_t> shape_tmp = CastPyArg2Longs(shape_obj, "full", 0);
+        shape = paddle::dialect::full_int_array(
+            shape_tmp, phi::DataType::INT64, phi::CPUPlace());
+      }
+
+      if (PyObject_CheckIROpResult(value_obj)) {
+        value = CastPyArg2OpResult(value_obj, "full", 1);
+      } else {
+        float value_tmp = CastPyArg2Float(value_obj, "full", 1);
+        value = paddle::dialect::full(std::vector<int64_t>{1},
+                                      value_tmp,
+                                      phi::DataType::FLOAT32,
+                                      phi::CPUPlace());
+      }
+
+      auto static_api_out =
+          paddle::dialect::full_with_tensor(shape, value, dtype);
+      return ToPyObject(static_api_out);
+    }
+  } catch (...) {
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  }
+}
+
 static PyMethodDef ManualOpsAPI[] = {
     {"set_parameter",
      (PyCFunction)(void (*)(void))static_api_set_parameter,
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -4398,7 +4398,7 @@ void MaskedMultiheadAttentionInferMeta(const MetaTensor& x,
 void FullWithTensorInferMeta(const MetaTensor& shape,
                              DataType dtype,
                              MetaTensor* out) {
-  out->set_dims(make_ddim({-1}));
+  out->set_dims(make_ddim(std::vector<int64_t>(shape.numel(), -1)));
   out->set_dtype(dtype);
 }
 
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -40,6 +40,7 @@
     convert_np_dtype_to_dtype_,
     core,
     in_dynamic_mode,
+    in_dynamic_or_pir_mode,
     in_pir_mode,
 )
 
@@ -878,42 +879,29 @@ def full_like(x, fill_value, dtype=None, name=None):
 
 
 def fill_constant(shape, dtype, value, force_cpu=False, out=None, name=None):
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         place = _current_expected_place()
         if force_cpu:
             place = core.CPUPlace()
         if isinstance(shape, (list, tuple)):
             shape = paddle.utils.convert_shape_to_list(shape)
 
-        if not isinstance(dtype, core.VarDesc.VarType):
+        if not isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
             dtype = convert_np_dtype_to_dtype_(dtype)
 
         if out is None:
-            out = _C_ops.full(shape, float(value), dtype, place)
+            value = float(value) if in_dynamic_mode() else value
+            out = _C_ops.full(shape, value, dtype, place)
             out.stop_gradient = True
             return out
 
         if out is not None:
+            value = float(value) if in_dynamic_mode() else value
             # final state mode is support out is not None.
-            _C_ops.full_(out, shape, float(value), dtype, place)
+            _C_ops.full_(out, shape, value, dtype, place)
             out.stop_gradient = True
             return out
     else:
-        if paddle.ir.core._use_new_ir_api():
-            # Below code will be removed after we can generate IR api automatically
-            place = _current_expected_place()
-            if force_cpu:
-                place = core.CPUPlace()
-            if isinstance(shape, (list, tuple)):
-                shape = paddle.utils.convert_shape_to_list(shape)
-
-            if not isinstance(dtype, core.DataType):
-                dtype = paddle.ir.core.convert_np_dtype_to_dtype_(dtype)
-
-            if out is None:
-                out = paddle._ir_ops.full(shape, float(value), dtype, place)
-                out.stop_gradient = True
-                return out
         attrs = {'force_cpu': force_cpu}
         dtype = convert_dtype(dtype)
         if not isinstance(value, Variable):
diff --git a/test/legacy_test/test_fill_constant_op.py b/test/legacy_test/test_fill_constant_op.py
@@ -16,7 +16,7 @@
 
 import numpy as np
 from op import Operator
-from op_test import OpTest, convert_float_to_uint16
+from op_test import OpTest, convert_float_to_uint16, paddle_static_guard
 
 import paddle
 from paddle import base
@@ -44,7 +44,7 @@ def setUp(self):
         self.outputs = {'Out': np.full(self.shape, self.value)}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
     def init_dtype(self):
         self.dtype = np.float64
@@ -115,7 +115,7 @@ def setUp(self):
 
     def test_check_output(self):
         place = core.CUDAPlace(0)
-        self.check_output_with_place(place)
+        self.check_output_with_place(place, check_new_ir=True)
 
 
 class TestFillConstantOpWithSelectedRows(unittest.TestCase):
@@ -168,7 +168,7 @@ def init_data(self):
         self.value = 3.8
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
 
 class TestFillConstantOp2_ShapeTensorList(OpTest):
@@ -192,7 +192,7 @@ def init_data(self):
         self.infer_shape = [-1, -1]
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
 
 class TestFillConstantOp3_ShapeTensorList(TestFillConstantOp1_ShapeTensorList):
@@ -226,7 +226,7 @@ def init_data(self):
         self.value = 3.8
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
 
 # Situation 4: value is a tensor
@@ -250,7 +250,7 @@ def init_data(self):
         self.dtype = np.float32
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
 
 # Situation 5: value is a tensor
@@ -274,7 +274,7 @@ def init_data(self):
         self.dtype = np.int32
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
 
 # Test python API
@@ -417,7 +417,7 @@ def test_ninf(self):
 
 class TestFillConstantOpError(unittest.TestCase):
     def test_errors(self):
-        with program_guard(Program(), Program()):
+        with paddle_static_guard(), program_guard(Program(), Program()):
             # for ci coverage
             x1 = paddle.static.data(name='x1', shape=[-1, 1], dtype="int16")
             self.assertRaises(
@@ -512,7 +512,9 @@ def init_data(self):
 
     def test_check_output(self):
         # no dynamic graph test for mkldnn
-        self.check_output_with_place(core.CPUPlace(), check_dygraph=False)
+        self.check_output_with_place(
+            core.CPUPlace(), check_dygraph=False, check_new_ir=False
+        )
 
 
 if __name__ == "__main__":

Original file line number	Diff line number	Diff line change
`@@ -4398,7 +4398,7 @@ void MaskedMultiheadAttentionInferMeta(const MetaTensor& x,`
`4398`	`4398`	`void FullWithTensorInferMeta(const MetaTensor& shape,`
`4399`	`4399`	`DataType dtype,`
`4400`	`4400`	`MetaTensor* out) {`
`4401`		`- out->set_dims(make_ddim({-1}));`
	`4401`	`+ out->set_dims(make_ddim(std::vector<int64_t>(shape.numel(), -1)));`
`4402`	`4402`	`out->set_dtype(dtype);`
`4403`	`4403`	`}`
`4404`	`4404`