[SOT] move_gpu_pinned_to_gpu (#60395)

feifei-111 · web-flow · commit 33cb1be658a5 · 2024-01-09T11:27:35.000+08:00
diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py
@@ -58,6 +58,7 @@
     NO_SHAPE_VAR_TYPE,
     ast_to_func,
     backend_guard,
+    cuda_pinned_tensors_move_to_excepted_place,
     func_to_source_code,
     input_specs_compatible,
     is_paddle_func,
@@ -714,6 +715,8 @@ def _perform_call(self, *args, **kwargs):
         from ..sot import symbolic_translate
 
         args, kwargs = self._function_spec.unified_args_and_kwargs(args, kwargs)
+        cuda_pinned_tensors_move_to_excepted_place(args)
+
         (
             input_args_with_spec,
             input_kwargs_with_spec,
diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py
@@ -35,7 +35,8 @@
 from paddle.base.data_feeder import convert_dtype
 from paddle.base.layer_helper import LayerHelper
 from paddle.base.wrapped_decorator import signature_safe_contextmanager
-from paddle.utils import gast
+from paddle.framework import CUDAPinnedPlace
+from paddle.utils import flatten, gast
 
 from .ast_utils import ast_to_source_code
 from .utils_helper import (  # noqa: F401
@@ -1279,3 +1280,19 @@ def tensor_name_guard(tensors, names):
     finally:
         for t, name in zip(tensors, origin_names):
             t.name = name
+
+
+def cuda_pinned_tensors_move_to_excepted_place(inputs):
+    if paddle.is_compiled_with_cuda():
+        expected_place = framework._current_expected_place()
+        cuda_pinned_place = CUDAPinnedPlace()
+
+        for value in flatten(inputs):
+            if (
+                isinstance(value, core.eager.Tensor)
+                and value.stop_gradient
+                and value.place._equals(cuda_pinned_place)
+            ):
+                var = value._copy_to(expected_place, True)
+                var.stop_gradient = True
+                var._share_buffer_to(value)
diff --git a/test/dygraph_to_static/test_move_cuda_pinned_tensor.py b/test/dygraph_to_static/test_move_cuda_pinned_tensor.py
@@ -0,0 +1,61 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import unittest
+
+import paddle
+from paddle.jit.dy2static.utils import (
+    cuda_pinned_tensors_move_to_excepted_place,
+)
+
+
+class TestCopyCudaPinnedTensors(unittest.TestCase):
+    def test_copy_cuda_pinned_tensors(self):
+        if paddle.device.is_compiled_with_cuda():
+
+            @paddle.jit.to_static
+            def f(x):
+                return x
+
+            cuda_pinned_place = paddle.CUDAPinnedPlace()
+            x = paddle.to_tensor(
+                [1, 2, 3], place=paddle.CUDAPinnedPlace(), stop_gradient=True
+            )
+
+            cuda_pinned_tensors_move_to_excepted_place(x)
+            assert not x.place._equals(cuda_pinned_place)
+
+            y = {
+                "a": [
+                    paddle.to_tensor(
+                        1, place=paddle.CUDAPinnedPlace(), stop_gradient=True
+                    ),
+                    paddle.to_tensor(
+                        2, place=paddle.CUDAPinnedPlace(), stop_gradient=True
+                    ),
+                ],
+                "b": {
+                    "c": paddle.to_tensor(
+                        3, place=paddle.CUDAPinnedPlace(), stop_gradient=True
+                    )
+                },
+            }
+            cuda_pinned_tensors_move_to_excepted_place(y)
+            for var in paddle.utils.flatten(y):
+                if isinstance(var, paddle.Tensor):
+                    assert not var.place._equals(cuda_pinned_place)
+
+
+if __name__ == '__main__':
+    unittest.main()
