【PIR API adaptor No.14-15】 Migrate nn.initializer.Assign/Bilinear into pir (#60114)

MarioLulab · web-flow · commit d433ad512546 · 2023-12-19T11:10:31.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -102,12 +102,12 @@
     'print',
     'sequence_mask',
     'number_count',
+    'assign_value',
 ]
 
 NO_NEED_GEN_STATIC_ONLY_APIS = [
     'add_n_',
     'add_n_with_kernel',
-    'assign_value',
     'c_allgather',
     'c_allreduce_max',
     'c_allreduce_sum',
diff --git a/python/paddle/nn/initializer/Bilinear.py b/python/paddle/nn/initializer/Bilinear.py
@@ -14,10 +14,14 @@
 
 import numpy as np
 
-from paddle import _C_ops
+from paddle import _C_ops, pir
 
 from ...base import core, framework, unique_name
-from ...base.framework import _current_expected_place, in_dygraph_mode
+from ...base.framework import (
+    _current_expected_place,
+    in_dygraph_mode,
+    in_pir_mode,
+)
 from .initializer import Initializer
 
 __all__ = []
@@ -87,11 +91,13 @@ def forward(self, var, block=None):
         """
         block = self._check_block(block)
 
-        if not isinstance(var, framework.Variable):
-            raise ValueError("var must be framework.Variable.")
+        if not isinstance(var, (framework.Variable, pir.core.ParameterMeta)):
+            raise ValueError(
+                "var must be framework.Variable or pir.core.ParameterMeta."
+            )
 
-        if not isinstance(block, framework.Block):
-            raise ValueError("block must be framework.Block.")
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise ValueError("block must be framework.Block or pir.Block.")
 
         shape = var.shape
         if len(shape) != 4:
@@ -127,11 +133,18 @@ def forward(self, var, block=None):
                 type=core.VarDesc.VarType.LOD_TENSOR,
                 persistable=False,
             )
+        elif var.dtype in [
+            core.DataType.FLOAT16,
+            core.DataType.BFLOAT16,
+            core.DataType.FLOAT64,
+        ]:
+            out_dtype = core.DataType.FLOAT32
+            out_var = var
         else:
             out_dtype = var.dtype
             out_var = var
 
-        if out_dtype == core.VarDesc.VarType.FP32:
+        if out_dtype in (core.VarDesc.VarType.FP32, core.DataType.FLOAT32):
             value_name = "fp32_values"
             values = [float(v) for v in weight.flat]
         else:
@@ -158,6 +171,20 @@ def forward(self, var, block=None):
             else:
                 out_var._share_underline_tensor_to(var)
             return None
+        elif in_pir_mode():
+            out_var = _C_ops.assign_value(
+                list(shape),
+                out_dtype,
+                values,
+                _current_expected_place(),
+            )
+            if var.dtype in [
+                core.DataType.FLOAT16,
+                core.DataType.BFLOAT16,
+                core.DataType.FLOAT64,
+            ]:
+                out_var = _C_ops.cast(out_var, var.dtype)
+            return out_var
         else:
             op = block.append_op(
                 type='assign_value',
diff --git a/python/paddle/nn/initializer/assign.py b/python/paddle/nn/initializer/assign.py
@@ -16,7 +16,11 @@
 
 from ...base import core, framework, unique_name
 from ...base.data_feeder import check_type
-from ...base.framework import _current_expected_place, in_dygraph_mode
+from ...base.framework import (
+    _current_expected_place,
+    in_dygraph_mode,
+    in_pir_mode,
+)
 from .initializer import Initializer
 
 __all__ = []
@@ -54,8 +58,10 @@ def forward(self, var, block=None):
         """
         block = self._check_block(block)
 
-        assert isinstance(var, framework.Variable)
-        assert isinstance(block, framework.Block)
+        assert isinstance(
+            var, (framework.Variable, paddle.pir.core.ParameterMeta)
+        )
+        assert isinstance(block, (framework.Block, paddle.pir.Block))
 
         # to be compatible of fp16 initalizers
         if var.dtype in [core.VarDesc.VarType.FP16, core.VarDesc.VarType.BF16]:
@@ -70,23 +76,29 @@ def forward(self, var, block=None):
                 type=core.VarDesc.VarType.LOD_TENSOR,
                 persistable=False,
             )
+        elif var.dtype in [core.DataType.FLOAT16, core.DataType.BFLOAT16]:
+            out_var = var
+            out_dtype = core.DataType.FLOAT32
+            np_value = self._value.astype("float32")
         else:
             out_var = var
             out_dtype = var.dtype
             np_value = self._value
 
-        if out_dtype == core.VarDesc.VarType.FP32:
+        if out_dtype in (core.VarDesc.VarType.FP32, core.DataType.FLOAT32):
             value_name = "fp32_values"
             values = [float(v) for v in np_value.flat]
-        elif out_dtype == core.VarDesc.VarType.FP64:
+        elif out_dtype in (core.VarDesc.VarType.FP64, core.DataType.FLOAT64):
             value_name = "fp64_values"
             values = [float(v) for v in np_value.flat]
-        elif out_dtype == core.VarDesc.VarType.INT32:
+        elif out_dtype in (core.VarDesc.VarType.INT32, core.DataType.INT32):
             value_name = "int32_values"
             values = [int(v) for v in np_value.flat]
-        elif (
-            out_dtype == core.VarDesc.VarType.INT8
-            or out_dtype == core.VarDesc.VarType.UINT8
+        elif out_dtype in (
+            core.VarDesc.VarType.INT8,
+            core.VarDesc.VarType.UINT8,
+            core.DataType.INT8,
+            core.DataType.UINT8,
         ):
             value_name = "int8_values"
             values = [int(v) for v in np_value.flat]
@@ -115,6 +127,16 @@ def forward(self, var, block=None):
             else:
                 out_var._share_underline_tensor_to(var)
             return None
+        elif in_pir_mode():
+            out_var = _C_ops.assign_value(
+                list(self._value.shape),
+                out_dtype,
+                values,
+                _current_expected_place(),
+            )
+            if var.dtype in [core.DataType.FLOAT16, core.DataType.BFLOAT16]:
+                out_var = _C_ops.cast(out_var, var.dtype)
+            return out_var
         else:
             op = block.append_op(
                 type='assign_value',
diff --git a/test/legacy_test/test_initializer.py b/test/legacy_test/test_initializer.py
