update get/set parameter (#57539)

YuanRisheng · web-flow · commit 00bd3aa99f33 · 2023-09-21T10:29:31.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/api_builder.cc b/paddle/fluid/pir/dialect/operator/ir/api_builder.cc
@@ -48,5 +48,16 @@ void APIBuilder::ResetInsertionPointToEnd() {
   builder_->SetInsertionPointToEnd(builder_->block());
 }
 
+pir::Parameter* APIBuilder::GetParameter(const std::string& name) const {
+  pir::Program* program = builder_->block()->GetParentOp()->GetParentProgram();
+  return program->GetParameter(name);
+}
+
+void APIBuilder::SetParameter(const std::string& name,
+                              std::unique_ptr<pir::Parameter>&& parameter) {
+  pir::Program* program = builder_->block()->GetParentOp()->GetParentProgram();
+  program->SetParameter(name, std::move(parameter));
+}
+
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/pir/dialect/operator/ir/api_builder.h b/paddle/fluid/pir/dialect/operator/ir/api_builder.h
@@ -17,6 +17,7 @@
 
 #include "paddle/pir/core/builder.h"
 #include "paddle/pir/core/macros.h"
+#include "paddle/pir/core/parameter.h"
 #include "paddle/pir/core/program.h"
 
 namespace paddle {
@@ -40,6 +41,11 @@ class APIBuilder {
 
   void ResetInsertionPointToEnd();
 
+  pir::Parameter* GetParameter(const std::string& name) const;
+
+  void SetParameter(const std::string& name,
+                    std::unique_ptr<pir::Parameter>&& parameter);
+
   std::shared_ptr<pir::Builder> GetBuilder() { return builder_; }
 
  private:
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_api.cc b/paddle/fluid/pir/dialect/operator/ir/manual_api.cc
@@ -18,7 +18,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/pir/core/builtin_op.h"
-
+#include "paddle/pir/core/parameter.h"
 namespace paddle {
 namespace dialect {
 
@@ -46,25 +46,18 @@ pir::OpResult zeros_like(pir::Value x,
   return paddle::dialect::full_like(x, 0, dtype, place);
 }
 
-pir::OpResult get_parameter(const std::string& name,
-                            phi::DataType dtype,
-                            const std::vector<int64_t>& shape) {
-  phi::LoD lod;
-  size_t offset{0};
-  pir::Type out_dense_tensor_type = paddle::dialect::DenseTensorType::get(
-      pir::IrContext::Instance(),
-      TransToIrDataType(dtype),
-      phi::DDim(shape.data(), shape.size()),
-      phi::DataLayout::UNDEFINED,
-      lod,
-      offset);
+pir::OpResult get_parameter(const std::string& name) {
+  pir::Parameter* param = APIBuilder::Instance().GetParameter(name);
   pir::GetParameterOp get_parameter_op =
       APIBuilder::Instance().GetBuilder()->Build<pir::GetParameterOp>(
-          name, out_dense_tensor_type);
+          name, param->type());
   return get_parameter_op.result(0);
 }
 
 void set_parameter(pir::Value parameter, const std::string& name) {
+  std::unique_ptr<pir::Parameter> param(
+      new pir::Parameter(nullptr, 0, parameter.type()));
+  APIBuilder::Instance().SetParameter(name, std::move(param));
   APIBuilder::Instance().GetBuilder()->Build<pir::SetParameterOp>(parameter,
                                                                   name);
 }
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_api.h b/paddle/fluid/pir/dialect/operator/ir/manual_api.h
@@ -32,9 +32,7 @@ pir::OpResult zeros_like(pir::Value x,
                          phi::DataType dtype = phi::DataType::UNDEFINED,
                          const Place& place = {});
 
-pir::OpResult get_parameter(const std::string& name,
-                            phi::DataType dtype,
-                            const std::vector<int64_t>& shape);
+pir::OpResult get_parameter(const std::string& name);
 
 void set_parameter(pir::Value parameter, const std::string& name);
 
diff --git a/paddle/fluid/pybind/ir.cc b/paddle/fluid/pybind/ir.cc
@@ -153,6 +153,11 @@ void BindProgram(py::module *m) {
            [](const std::shared_ptr<Program> &self) {
              return self->parameters_num();
            })
+      .def("move_parameters_from",
+           [](const std::shared_ptr<Program> &self,
+              const std::shared_ptr<Program> &other) {
+             self->set_parameters(std::move(other->parameters()));
+           })
       .def(
           "global_block",
           [](std::shared_ptr<Program> self) { return self->block(); },
@@ -375,9 +380,19 @@ void BindOpOperand(py::module *m) {
 bool GetOpResultBoolAttr(const OpResult &self, const std::string &attr_name) {
   auto *defining_op = self.owner();
   if (defining_op->HasAttribute(attr_name)) {
+    PADDLE_ENFORCE(
+        defining_op->attribute(attr_name).isa<pir::ArrayAttribute>(),
+        paddle::platform::errors::InvalidArgument(
+            "%s: Callstack attributes of %s is not ArrayAttribute type",
+            attr_name));
     auto attrs = defining_op->attribute(attr_name)
                      .dyn_cast<pir::ArrayAttribute>()
                      .AsVector();
+    PADDLE_ENFORCE(attrs[self.index()].isa<pir::BoolAttribute>(),
+                   paddle::platform::errors::InvalidArgument(
+                       "The index %d in %s is not BoolAttribute type",
+                       self.index(),
+                       attr_name));
     return attrs[self.index()].dyn_cast<pir::BoolAttribute>().data();
   } else {
     return true;
diff --git a/paddle/fluid/pybind/manual_static_op_function.h b/paddle/fluid/pybind/manual_static_op_function.h
@@ -35,13 +35,8 @@ static PyObject *static_api_get_parameter(PyObject *self,
     // Parse Attributes
     PyObject *name_obj = PyTuple_GET_ITEM(args, 0);
     std::string name = CastPyArg2String(name_obj, "name", 0);
-    PyObject *dtype_obj = PyTuple_GET_ITEM(args, 1);
-    phi::DataType dtype = CastPyArg2DataTypeDirectly(dtype_obj, "dtype", 1);
-    PyObject *shape_obj = PyTuple_GET_ITEM(args, 2);
-    phi::IntArray shape = CastPyArg2IntArray(shape_obj, "shape", 2);
     // Call ir static api
-    auto static_api_out =
-        paddle::dialect::get_parameter(name, dtype, shape.GetData());
+    auto static_api_out = paddle::dialect::get_parameter(name);
 
     return ToPyObject(static_api_out);
   } catch (...) {
diff --git a/python/paddle/base/data_feeder.py b/python/paddle/base/data_feeder.py
@@ -47,7 +47,7 @@
 _PADDLE_NEW_IR_DTYPE_2_NUMPY_DTYPE = {
     core.DataType.BOOL: 'bool',
     core.DataType.FLOAT16: 'float16',
-    core.DataType.UINT16: 'uint16',
+    core.DataType.BFLOAT16: 'uint16',
     core.DataType.FLOAT32: 'float32',
     core.DataType.FLOAT64: 'float64',
     core.DataType.INT8: 'int8',
diff --git a/python/paddle/base/executor.py b/python/paddle/base/executor.py
@@ -515,11 +515,12 @@ def _add_pir_fetch_ops(program, fetch_list, fetch_var_name):
     if not has_fetch_operations(
         global_block, fetch_list, fetch_var_name, fetch_op
     ):
-        for i, fetch_input in enumerate(fetch_list):
-            assert isinstance(
-                fetch_input, OpResult
-            ), "Wrong type for fetch_list[%s]: %s" % (i, type(fetch_input))
-            paddle._ir_ops.fetch(fetch_input, fetch_var_name + str(i), i)
+        with paddle.static.program_guard(program):
+            for i, fetch_input in enumerate(fetch_list):
+                assert isinstance(
+                    fetch_input, OpResult
+                ), "Wrong type for fetch_list[%s]: %s" % (i, type(fetch_input))
+                paddle._ir_ops.fetch(fetch_input, fetch_var_name + str(i), i)
 
 
 def _merge_tensors(tensor, micro_batch_num):
diff --git a/python/paddle/ir/core.py b/python/paddle/ir/core.py
@@ -251,6 +251,12 @@ def program_guard(main_program, startup_program=None):
             switch_startup_program(startup_program)
 
 
+class ParameterMeta:
+    def __init__(self, shape, dtype):
+        self.shape = shape
+        self.dtype = dtype
+
+
 def create_parameter(
     dtype,
     shape,
@@ -266,19 +272,21 @@ def create_parameter(
     op_result_name = unique_name.generate('parameter')
     startup_program = default_startup_program()
     main_program = default_main_program()
-
-    with program_guard(default_main_program()):
-        param = get_parameter(op_result_name, dtype, shape)
-        trainable = kwargs.get('trainable', True)
-        param.stop_gradient = not trainable
-        param.is_persistable = True
+    parameter_meta = ParameterMeta(shape, dtype)
 
     with program_guard(startup_program):
         initializer = kwargs['initializer']
         init_result = initializer(
-            param, param.get_defining_op().get_parent_block()
+            parameter_meta, startup_program.global_block()
         )
         init_result.is_persistable = True
         set_parameter(init_result, op_result_name)
 
+    main_program.move_parameters_from(startup_program)
+    with program_guard(default_main_program()):
+        param = get_parameter(op_result_name, dtype, shape)
+        trainable = kwargs.get('trainable', True)
+        param.stop_gradient = not trainable
+        param.is_persistable = True
+
     return param
diff --git a/python/paddle/nn/initializer/constant.py b/python/paddle/nn/initializer/constant.py
@@ -58,7 +58,12 @@ def forward(self, var, block=None):
 
         assert isinstance(
             var,
-            (framework.Variable, framework.EagerParamBase, paddle.ir.OpResult),
+            (
+                framework.Variable,
+                framework.EagerParamBase,
+                paddle.ir.OpResult,
+                paddle.ir.core.ParameterMeta,
+            ),
         )
         assert isinstance(block, (framework.Block, paddle.ir.Block))
 
diff --git a/python/paddle/nn/initializer/xavier.py b/python/paddle/nn/initializer/xavier.py
@@ -88,12 +88,13 @@ def forward(self, var, block=None):
 
         block = self._check_block(block)
         assert isinstance(block, (framework.Block, paddle.ir.Block))
-        check_variable_and_dtype(
-            var,
-            "Out",
-            ["uint16", "float16", "float32", "float64"],
-            "xavier_init",
-        )
+        if not isinstance(var, paddle.ir.core.ParameterMeta):
+            check_variable_and_dtype(
+                var,
+                "Out",
+                ["uint16", "float16", "float32", "float64"],
+                "xavier_init",
+            )
 
         f_in, f_out = self._compute_fans(var)
 
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -889,8 +889,6 @@ def divide(x, y, name=None):
     """
     if in_dynamic_or_pir_mode():
         return _C_ops.divide(x, y)
-    elif in_pir_mode():
-        return paddle._ir_ops.divide(x, y)
     else:
         return _elementwise_op(LayerHelper('elementwise_div', **locals()))
 
diff --git a/test/ir/new_ir/test_build_model.py b/test/ir/new_ir/test_build_model.py
@@ -31,12 +31,12 @@ def test_basic_network(self):
             exe = paddle.static.Executor()
             x_feed = np.ones([4, 4], dtype=np.float32) * 10
             y_feed = np.ones([4, 4], dtype=np.float32) * 2
-            (sum_value,) = exe.run(
-                main_program,
-                feed={'x': x_feed, 'y': y_feed},
-                fetch_list=[sum_out],
-            )
-            self.assertEqual(sum_value, 5 * 4 * 4)
+        (sum_value,) = exe.run(
+            main_program,
+            feed={'x': x_feed, 'y': y_feed},
+            fetch_list=[sum_out],
+        )
+        self.assertEqual(sum_value, 5 * 4 * 4)
 
         main_program = paddle.static.Program()
         with paddle.static.program_guard(main_program):
