test_stride (#67452)

ooooo-create · web-flow · commit ba37d2796495 · 2024-08-22T11:06:51.000+08:00
* add stride valueerror in pir

* refine

* refine

* fix

* refine and check after dy2st

* refine

* fix

* fix import

* refine

* apply review

* update

* fix
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -103,6 +103,27 @@
     "atan2_grad": ['X1', 'X2'],
 }
 
+stride_ops = [
+    "pd_op.slice",
+    "pd_op.strided_slice",
+    "pd_op.index_select",
+    "pd_op.split",
+    "pd_op.unsqueeze",
+    "pd_op.unsqueeze2",
+    "pd_op.squeeze",
+    "pd_op.squeeze2",
+    "pd_op.transpose",
+    "pd_op.transpose2",
+    "pd_op.unbind",
+    "pd_op.diagonal",
+    "pd_op.flatten",
+    "pd_op.imag",
+    "pd_op.real",
+    "pd_op.reshape",
+    "pd_op.reshape2",
+    "pd_op.as_real",
+]
+
 
 def _global_flags():
     return _global_flags_
diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py
@@ -16,6 +16,7 @@
 
 import collections
 import inspect
+import os
 import threading
 import warnings
 import weakref
@@ -99,6 +100,54 @@ def lock_func(*args, **kwargs):
     return lock_func
 
 
+def show_op_callstack(op):
+    op_callstack = op.callstack
+    index = op_callstack.index("    outputs = static_func(*inputs)")
+    op_callstack_result = '\n'.join(op_callstack[index + 1 :])
+    raise ValueError(
+        f'In transformed code:\n\n{op_callstack_result}\n\nSorry about what\'s happened. In to_static mode, {op.name()}\'s output variable is a viewed Tensor in dygraph. This will result in inconsistent calculation behavior between dynamic and static graphs. You must find the location of the strided ops be called, and call paddle.assign() before inplace input.If you certainly make sure it\'s safe, you can set env stride_in_no_check_dy2st_diff to 1.'
+    )
+
+
+def check_view_api_used_by_inplace(program: paddle.pir.Program) -> None:
+    """
+    check viewed value used by inplace op in pir mode.
+
+    Two scenarios will raise ValueError:
+        # one
+        a = transpose(b)
+        a.add_(c)
+        # two
+        a = transpose(b)
+        b.add_(c)
+    """
+    all_vars_list = program.list_vars()
+    for value in all_vars_list:
+        if len(value.all_used_ops()) == 0:
+            return
+        uesd_by_stride_ops = []
+        for op in value.all_used_ops()[::-1]:
+            inplace_info = paddle.core.pir.get_op_inplace_info(op)
+            if op.name() in framework.stride_ops and op.operand_source(
+                0
+            ).is_same(value):
+                uesd_by_stride_ops.append(op)
+            if op.name().endswith("_") and any(
+                op.operand_source(index).is_same(value)
+                for index in inplace_info.keys()
+            ):
+                if value.get_defining_op().name() in framework.stride_ops:
+                    show_op_callstack(op)
+                if len(uesd_by_stride_ops) == 0:
+                    continue
+                if (
+                    op.name() == "pd_op.set_value_"
+                    or op.name() == "pd_op.set_value_with_tensor_"
+                ):
+                    continue
+                show_op_callstack(op)
+
+
 class FunctionCache:
     """
     Caches the transformed functions to avoid redundant conversions of the same function.
@@ -1260,6 +1309,8 @@ def pir_from_func_spec(
                         outputs = [outputs]
 
         main_program = update_op_callstack_with_origin_info(main_program)
+        if not os.environ.get("stride_in_no_check_dy2st_diff", "0") == "1":
+            check_view_api_used_by_inplace(main_program)
 
         return ConcreteProgram(
             inputs=static_inputs,
diff --git a/python/paddle/utils/inplace_utils.py b/python/paddle/utils/inplace_utils.py
@@ -11,19 +11,41 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
 
+import os
 import warnings
-from typing import Callable, TypeVar
+from typing import TYPE_CHECKING, Callable, TypeVar
 
 from typing_extensions import ParamSpec
 
 import paddle  # noqa: F401
+from paddle.base.framework import stride_ops
 from paddle.base.wrapped_decorator import wrap_decorator
-from paddle.framework import in_dynamic_mode
+from paddle.framework import in_dynamic_mode, in_pir_mode
 
 _InputT = ParamSpec("_InputT")
 _RetT = TypeVar("_RetT")
 
+if TYPE_CHECKING:
+    from paddle.pir import Value
+
+
+def check_view_value(value: Value) -> bool:
+    # check if the value is a view tensor
+    if value.get_defining_op().name() in stride_ops:
+        # TODO(ooooo-create): The `x = stride_op(x)` shouldn't return True.
+        return True
+    all_used_ops = value.all_used_ops()
+    if len(all_used_ops) == 0:
+        return False
+    for op in all_used_ops:
+        if op.name() in stride_ops and op.operand_source(0).is_same(value):
+            # TODO(ooooo-create): The `y = stride_op(x).clone()` and `y = stride_op(x) + op` should also return False.
+            # Now is True.
+            return True
+    return False
+
 
 # NOTE(pangyoki): The Inplace APIs with underline(`_`) is only valid for the method of calling `_C_ops`
 # in dygraph mode. If static graph mode is used, the inplace mechanism will not be used, and the static method
@@ -42,11 +64,23 @@ def __impl__(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
             from ..base.dygraph.base import in_to_static_mode
 
             if in_to_static_mode():
-                for arg in args:
-                    if hasattr(arg, "is_view_var") and arg.is_view_var:
+                stride_in_no_check_dy2st_diff = os.environ.get(
+                    "stride_in_no_check_dy2st_diff", "0"
+                )
+                if in_pir_mode():
+                    if (
+                        stride_in_no_check_dy2st_diff != '1'
+                        and check_view_value(args[0])
+                    ):
                         raise ValueError(
-                            f'Sorry about what\'s happened. In to_static mode, {func.__name__}\'s output variable {arg.name} is a viewed Tensor in dygraph. This will result in inconsistent calculation behavior between dynamic and static graphs. You must find the location of the strided API be called, and call {arg.name} = {arg.name}.assign().'
+                            f'Sorry about what\'s happened. In to_static mode, {func.__name__}\'s output variable is a viewed Tensor in dygraph. This will result in inconsistent calculation behavior between dynamic and static graphs. You must find the location of the strided API be called, and call paddle.assign() before inplace input.'
                         )
+                else:
+                    for arg in args:
+                        if hasattr(arg, "is_view_var") and arg.is_view_var:
+                            raise ValueError(
+                                f'Sorry about what\'s happened. In to_static mode, {func.__name__}\'s output variable {arg.name} is a viewed Tensor in dygraph. This will result in inconsistent calculation behavior between dynamic and static graphs. You must find the location of the strided API be called, and call {arg.name} = paddle.assign({arg.name}).'
+                            )
 
             origin_func = f"{func.__module__}.{origin_api_name}"
             return eval(origin_func)(*args, **kwargs)
diff --git a/test/legacy_test/test_stride.py b/test/legacy_test/test_stride.py
@@ -17,6 +17,7 @@
 import numpy as np
 
 import paddle
+from paddle.pir_utils import DygraphPirGuard, test_with_dygraph_pir
 
 
 def ref_view_as_real(x):
@@ -751,27 +752,95 @@ def test_stride_gpu(self):
 
 
 class TestToStaticCheck(unittest.TestCase):
+    @test_with_dygraph_pir
     def test_error(self):
         @paddle.jit.to_static(full_graph=True)
-        def func():
+        def func1():
             x_np = np.random.random(size=[2, 3, 4]).astype('float32')
             x = paddle.to_tensor(x_np)
             y = paddle.transpose(x, perm=[1, 0, 2])
-            x.add_(x)
+            z = paddle.ones([3, 2, 4])
+            y.add_(z)
 
-        if not paddle.framework.use_pir_api():
-            self.assertRaises(ValueError, func)
+        self.assertRaises(ValueError, func1)
 
+        @paddle.jit.to_static(full_graph=True)
+        def func2():
+            x_np = np.random.random(size=[2, 3, 4]).astype('float32')
+            x = paddle.to_tensor(x_np)
+            y = paddle.transpose(x, perm=[1, 0, 2])
+            z = paddle.ones([2, 3, 4])
+            x.add_(z)
+
+        self.assertRaises(ValueError, func2)
+
+    def test_error_with_program(self):
+        with DygraphPirGuard():
+
+            @paddle.jit.to_static(full_graph=True)
+            def func1():
+                x = paddle.ones((4, 3)) * 2
+                y = paddle.transpose(x, [1, 0])
+                z = paddle.ones((3,))
+                paddle.tensor.manipulation.fill_diagonal_tensor_(y, z)
+
+            self.assertRaises(ValueError, func1)
+
+            @paddle.jit.to_static(full_graph=True)
+            def func2():
+                x = paddle.ones((4, 3)) * 2
+                y = paddle.transpose(x, [1, 0])
+                z = paddle.ones((3,))
+                paddle.tensor.manipulation.fill_diagonal_tensor_(x, z)
+
+            self.assertRaises(ValueError, func2)
+
+    def test_no_error_with_program(self):
+        with DygraphPirGuard():
+
+            @paddle.jit.to_static(full_graph=True)
+            def func1():
+                x = paddle.ones((4, 3)) * 2
+                y = paddle.transpose(x, [1, 0])
+                yy = paddle.assign(y)
+                z = paddle.ones((3,))
+                paddle.tensor.manipulation.fill_diagonal_tensor_(yy, z)
+
+            func1()
+
+            @paddle.jit.to_static(full_graph=True)
+            def func2():
+                x = paddle.ones((4, 3)) * 2
+                y = paddle.transpose(x, [1, 0])
+                xx = paddle.assign(x)
+                z = paddle.ones((3,))
+                paddle.tensor.manipulation.fill_diagonal_tensor_(xx, z)
+
+            func2()
+
+    @test_with_dygraph_pir
     def test_no_error(self):
         @paddle.jit.to_static(full_graph=True)
-        def func():
+        def func1():
             x_np = np.random.random(size=[2, 3, 4]).astype('float32')
             x = paddle.to_tensor(x_np)
+            z = paddle.ones([3, 2, 4])
+            y = paddle.transpose(x, perm=[1, 0, 2])
+            yy = paddle.assign(y)
+            yy.add_(z)
+
+        func1()
+
+        @paddle.jit.to_static(full_graph=True)
+        def func2():
+            x_np = np.random.random(size=[2, 3, 4]).astype('float32')
+            x = paddle.to_tensor(x_np)
+            y = paddle.transpose(x, perm=[1, 0, 2])
             xx = paddle.assign(x)
-            y = paddle.transpose(xx, perm=[1, 0, 2])
-            x.add_(x)
+            z = paddle.ones([2, 3, 4])
+            xx.add_(z)
 
-        func()
+        func2()
 
 
 if __name__ == '__main__':
