[PIR] fix logspace and matmul_triple_grad (#67430)

wanghuancoder · web-flow · commit ff9e2d83d828 · 2024-08-16T10:11:12.000+08:00
* fix logspace and matmul_triple_grad
diff --git a/paddle/phi/ops/yaml/inconsistent/static_backward.yaml b/paddle/phi/ops/yaml/inconsistent/static_backward.yaml
@@ -244,6 +244,7 @@
     func : matmul_double_grad
   composite : matmul_double_grad(x, y, grad_out, grad_x_grad, grad_y_grad, transpose_x=false, transpose_y=false)
   optional : grad_x_grad, grad_y_grad
+  backward : matmul_triple_grad
 
 - backward_op : matmul_grad
   forward : matmul (Tensor x, Tensor y, bool transpose_x=false, bool transpose_y=false) -> Tensor(out)
@@ -258,6 +259,16 @@
   composite: matmul_grad(x, y, out_grad, transpose_x, transpose_y, x_grad, y_grad)
   backward : matmul_double_grad
 
+- backward_op : matmul_triple_grad
+  forward : matmul_double_grad (Tensor x, Tensor y, Tensor grad_out, Tensor grad_grad_x, Tensor grad_grad_y, bool transpose_x=false, bool transpose_y=false) -> Tensor(grad_x), Tensor(grad_y), Tensor(grad_grad_out)
+  args : (Tensor x, Tensor y, Tensor grad_out, Tensor grad_grad_x, Tensor grad_grad_y, Tensor grad_x_grad, Tensor grad_y_grad, Tensor grad_grad_out_grad, bool transpose_x=false, bool transpose_y=false)
+  output : Tensor(x_grad), Tensor(y_grad), Tensor(grad_out_grad), Tensor(grad_grad_x_grad), Tensor(grad_grad_y_grad)
+  infer_meta :
+    func : GeneralQuinaryGradInferMeta
+    param : [x, y, grad_out, grad_grad_x, grad_grad_y]
+  kernel :
+    func : matmul_triple_grad
+
 - backward_op : matmul_with_flatten_grad
   forward : matmul_with_flatten (Tensor x, Tensor y, int x_num_col_dims=1, int y_num_col_dims=1) -> Tensor(out)
   args : (Tensor x, Tensor y, Tensor out_grad, int x_num_col_dims=1, int y_num_col_dims=1)
diff --git a/python/paddle/nn/layer/layers.py b/python/paddle/nn/layer/layers.py
@@ -2224,19 +2224,33 @@ def _set_var(var, ndarray):
                 t.set(ndarray, place)
 
             try:
-                executor = Executor(_get_device())._default_executor
                 # restore parameter states
-                core._create_loaded_parameter(
-                    [param for param, state in matched_param_state],
-                    global_scope(),
-                    executor,
-                )
+                if in_pir_mode():
+                    executor = Executor(
+                        paddle.base.framework._current_expected_place_()
+                    )._default_executor
+                    paddle.base.libpaddle.pir.create_loaded_parameter(
+                        [param for param, state in matched_param_state],
+                        global_scope(),
+                        executor,
+                    )
+                else:
+                    executor = Executor(_get_device())._default_executor
+                    core._create_loaded_parameter(
+                        [param for param, state in matched_param_state],
+                        global_scope(),
+                        executor,
+                    )
                 for param, state in matched_param_state:
                     _set_var(param, state)
             except ValueError as e:
                 raise ValueError(
                     "This error might happens in dy2static, while calling 'set_state_dict' dynamically in 'forward', which is not supported. If you only need call 'set_state_dict' once, move it to '__init__'."
                 )
+            except TypeError as e:
+                raise ValueError(
+                    "This error might happens in dy2static, while calling 'set_state_dict' dynamically in 'forward', which is not supported. If you only need call 'set_state_dict' once, move it to '__init__'."
+                )
 
         return missing_keys, unexpected_keys
 
diff --git a/python/paddle/sparse/nn/layer/norm.py b/python/paddle/sparse/nn/layer/norm.py
@@ -18,9 +18,9 @@
 from typing import TYPE_CHECKING, Literal
 
 import paddle
-from paddle import _C_ops, in_dynamic_mode
+from paddle import _C_ops
 from paddle.base.layer_helper import LayerHelper
-from paddle.framework import no_grad
+from paddle.framework import in_dynamic_or_pir_mode, no_grad
 from paddle.nn.layer.norm import _BatchNormBase
 
 if TYPE_CHECKING:
@@ -150,7 +150,7 @@ def forward(self, input: Tensor) -> Tensor:
 
         data_format = 'NCHW' if self._data_format[1] == 'C' else 'NHWC'
 
-        if in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             batch_norm_out, _, _, _, _, _ = _C_ops.sparse_batch_norm_(
                 input,
                 self._mean,
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -556,7 +556,42 @@ def logspace(
     if not isinstance(base, (Variable, paddle.pir.Value)):
         with device_guard("cpu"):
             tensor_base = fill_constant([1], dtype, base)
-    if in_dynamic_or_pir_mode():
+    if in_dynamic_mode():
+        return _C_ops.logspace(
+            tensor_start,
+            tensor_stop,
+            tensor_num,
+            tensor_base,
+            dtype,
+            _current_expected_place(),
+        )
+    elif in_pir_mode():
+        start_dtype = convert_dtype(tensor_start.dtype)
+        stop_dtype = convert_dtype(tensor_stop.dtype)
+        base_dtype = convert_dtype(tensor_base.dtype)
+        out_dtype = convert_dtype(dtype)
+        if (
+            (
+                stop_dtype == "float64"
+                or start_dtype == "float64"
+                or base_dtype == "float64"
+            )
+            and out_dtype in ["float32", "int32"]
+        ) or (
+            (
+                stop_dtype == "int64"
+                or start_dtype == "int64"
+                or base_dtype == "int64"
+            )
+            and out_dtype == "int32"
+        ):
+            raise ValueError(
+                f"The dtype of start/stop/base is {start_dtype}/{stop_dtype}/{base_dtype} but the attr(dtype) of logspace is {dtype}, "
+                "which may cause data type overflows. Please reset attr(dtype) of logspace."
+            )
+        if isinstance(num, paddle.pir.Value):
+            check_dtype(num.dtype, 'num', ['int32'], 'logspace')
+
         return _C_ops.logspace(
             tensor_start,
             tensor_stop,
diff --git a/test/deprecated/legacy_test/CMakeLists.txt b/test/deprecated/legacy_test/CMakeLists.txt
@@ -694,7 +694,6 @@ set(STATIC_BUILD_TESTS
     test_momentum_op_deprecated
     test_nce_deprecated
     test_sparse_conv_op
-    test_sparse_norm_op
     test_tensor_array_to_tensor_deprecated
     test_unique
     test_one_hot_v2_op)
diff --git a/test/deprecated/legacy_test/test_learning_rate_scheduler_deprecated.py b/test/deprecated/legacy_test/test_learning_rate_scheduler_deprecated.py
@@ -616,4 +616,5 @@ def run_places(lr, start_lr, end_lr):
 
 
 if __name__ == '__main__':
+    paddle.enable_static()
     unittest.main()
diff --git a/test/dygraph_to_static/test_high_order_net.py b/test/dygraph_to_static/test_high_order_net.py
@@ -26,10 +26,12 @@
 class HighOrderNet(paddle.nn.Layer):
     def __init__(self):
         super().__init__()
-        self.linear = paddle.nn.Linear(3, 4, bias_attr=False)
+        self.bilinear = paddle.nn.Bilinear(
+            in1_features=5, in2_features=4, out_features=1000
+        )
 
-    def forward(self, x):
-        y = self.linear(x)
+    def forward(self, x, y):
+        y = self.bilinear(x, y)
         z = paddle.pow(y, 2)
         x_grad = paddle.grad(z, x, create_graph=True)[0]
         x_grad_grad = paddle.grad(x_grad, x, create_graph=True)[0]
@@ -43,14 +45,16 @@ def test_backward_has_no_grad_error(self):
         net = HighOrderNet()
         static_net = paddle.jit.to_static(net, full_graph=True)
 
-        x = paddle.to_tensor([[1, 1, 1], [1, 1, 1]], 'float32')
+        x = layer1 = paddle.rand((5, 5)).astype('float32')
         x.stop_gradient = False
+        y = layer1 = paddle.rand((5, 4)).astype('float32')
+        y.stop_gradient = False
 
         with self.assertRaisesRegex(
             ValueError,
-            "op 'pd_op.matmul_double_grad' has no grad op, consider enable prim to decompose it.",
+            "op 'pd_op.bilinear_grad' has no grad op, consider enable prim to decompose it.",
         ):
-            x_grad_grad = static_net(x)
+            x_grad_grad = static_net(x, y)
             x_grad_grad.backward()
 
 
diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt
@@ -1103,6 +1103,7 @@ set(STATIC_BUILD_TESTS
     test_random_routing_op
     test_searchsorted_op
     test_sparse_momentum_op
+    test_sparse_norm_op
     test_sgd_op_bf16
     test_softmax_mask_fuse_upper_triangle_op
     test_sparse_pooling_op
diff --git a/test/legacy_test/test_logspace.py b/test/legacy_test/test_logspace.py
@@ -192,7 +192,8 @@ def test_name(self):
             out = paddle.logspace(
                 0, 10, 5, 2, dtype='float32', name='logspace_res'
             )
-            assert 'logspace_res' in out.name
+            if not paddle.framework.use_pir_api():
+                assert 'logspace_res' in out.name
 
     def test_imperative(self):
         paddle.disable_static()
diff --git a/test/legacy_test/test_nn_matmul_v2_grad.py b/test/legacy_test/test_nn_matmul_v2_grad.py
diff --git a/test/legacy_test/test_parameter.py b/test/legacy_test/test_parameter.py
@@ -28,26 +28,29 @@
 
 class ParameterChecks(unittest.TestCase):
     def test_parameter(self):
-        shape = [784, 100]
-        val = 1.0625
-        b = main_program.global_block()
-        param = b.create_parameter(
-            name='fc.w',
-            shape=shape,
-            dtype='float32',
-            initializer=paddle.nn.initializer.Constant(val),
-        )
-        self.assertIsNotNone(param)
-        self.assertEqual('fc.w', param.name)
-        self.assertEqual((784, 100), param.shape)
-        self.assertEqual(paddle.float32, param.dtype)
-        self.assertEqual(0, param.block.idx)
-        exe = Executor(paddle.CPUPlace())
-        p = exe.run(main_program, fetch_list=[param])[0]
-        np.testing.assert_array_equal(p, np.ones(shape) * val)
-
-        zero_dim_param = b.create_parameter(name='x', shape=[], dtype='float32')
-        self.assertEqual(zero_dim_param.shape, ())
+        with paddle.pir_utils.OldIrGuard():
+            shape = [784, 100]
+            val = 1.0625
+            b = main_program.global_block()
+            param = b.create_parameter(
+                name='fc.w',
+                shape=shape,
+                dtype='float32',
+                initializer=paddle.nn.initializer.Constant(val),
+            )
+            self.assertIsNotNone(param)
+            self.assertEqual('fc.w', param.name)
+            self.assertEqual((784, 100), param.shape)
+            self.assertEqual(paddle.float32, param.dtype)
+            self.assertEqual(0, param.block.idx)
+            exe = Executor(paddle.CPUPlace())
+            p = exe.run(main_program, fetch_list=[param])[0]
+            np.testing.assert_array_equal(p, np.ones(shape) * val)
+
+            zero_dim_param = b.create_parameter(
+                name='x', shape=[], dtype='float32'
+            )
+            self.assertEqual(zero_dim_param.shape, ())
 
     def test_parambase(self):
         with guard():
diff --git a/test/legacy_test/test_sparse_norm_op.py b/test/legacy_test/test_sparse_norm_op.py
@@ -172,8 +172,9 @@ class TestStatic(unittest.TestCase):
     @compare_legacy_with_pt
     def test(self):
         paddle.enable_static()
-        main_program = paddle.static.Program()
-        with paddle.static.program_guard(main_program):
+        main_program = paddle.base.Program()
+        startup_program = paddle.base.Program()
+        with paddle.base.program_guard(main_program, startup_program):
             indices = paddle.static.data(
                 name='indices', shape=[4, 4], dtype='int32'
             )
@@ -189,6 +190,7 @@ def test(self):
             out = sp_y.to_dense()
 
             exe = paddle.static.Executor()
+            exe.run(startup_program)
             indices_data = [
                 [0, 0, 0, 0],
                 [0, 0, 0, 0],
@@ -198,19 +200,11 @@ def test(self):
             values_data = np.array([[1.0], [2.0], [3.0], [4.0]]).astype(
                 'float32'
             )
-            bias_data = np.array([1.0]).astype('float32')
-            weight_data = np.array([2.0]).astype('float32')
-            mean_data = np.array([1.0]).astype('float32')
-            variance_data = np.array([2.0]).astype('float32')
-
             fetch = exe.run(
+                main_program,
                 feed={
                     'indices': indices_data,
                     'values': values_data,
-                    sparse_batch_norm.bias.name: bias_data,
-                    sparse_batch_norm.weight.name: weight_data,
-                    sparse_batch_norm._mean.name: mean_data,
-                    sparse_batch_norm._variance.name: variance_data,
                 },
                 fetch_list=[out],
                 return_numpy=True,
@@ -219,9 +213,9 @@ def test(self):
                 [
                     [
                         [
-                            [[0.0], [-1.6832708], [0.0], [0.1055764]],
-                            [[0.0], [0.0], [1.8944236], [0.0]],
-                            [[0.0], [0.0], [0.0], [3.683271]],
+                            [[0.0], [-1.3416353], [0.0], [-0.44721174]],
+                            [[0.0], [0.0], [0.44721198], [0.0]],
+                            [[0.0], [0.0], [0.0], [1.3416355]],
                         ]
                     ]
                 ]
diff --git a/tools/parallel_UT_rule.py b/tools/parallel_UT_rule.py
@@ -1145,7 +1145,7 @@
     'test_normal',
     'test_tensor_scalar_type_promotion_static',
     'test_trt_group_norm_op',
-    'test_learning_rate_scheduler',
+    'test_learning_rate_scheduler_deprecated',
     'test_numel_op',
     'test_adaptive_max_pool3d',
     'test_sequential',
@@ -2663,7 +2663,7 @@
     'test_label_smooth_op',
     'test_logsumexp',
     'test_log_softmax',
-    'test_learning_rate_scheduler',
+    'test_learning_rate_scheduler_deprecated',
     'test_linspace',
     'test_linear_interp_op',
     'test_lamb_op',
diff --git a/tools/static_mode_white_list.py b/tools/static_mode_white_list.py
@@ -285,7 +285,7 @@
     'test_layer_norm_mkldnn_op',
     'test_layer_norm_bf16_mkldnn_op',
     'test_layer_norm_op_v2',
-    'test_learning_rate_scheduler',
+    'test_learning_rate_scheduler_deprecated',
     'test_linear_interp_op',
     'test_linear_interp_v2_op',
     'test_linspace',

Original file line number	Diff line number	Diff line change
`@@ -616,4 +616,5 @@ def run_places(lr, start_lr, end_lr):`
`616`	`616`
`617`	`617`
`618`	`618`	`if __name__ == '__main__':`
	`619`	`+ paddle.enable_static()`
`619`	`620`	`unittest.main()`
Original file line number	Diff line number	Diff line change
`@@ -192,7 +192,8 @@ def test_name(self):`
`192`	`192`	`out = paddle.logspace(`
`193`	`193`	`0, 10, 5, 2, dtype='float32', name='logspace_res'`
`194`	`194`	`)`
`195`		`- assert 'logspace_res' in out.name`
	`195`	`+ if not paddle.framework.use_pir_api():`
	`196`	`+ assert 'logspace_res' in out.name`
`196`	`197`
`197`	`198`	`def test_imperative(self):`
`198`	`199`	`paddle.disable_static()`