[SOT] Non-break support for paddle.get_device (#72004)

SigureMo · Copilot · web-flow · commit 85fe49823796 · 2025-04-01T23:31:40.000+08:00
---------

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/python/paddle/device/__init__.py b/python/paddle/device/__init__.py
@@ -342,7 +342,6 @@ def get_device() -> str:
     elif isinstance(place, core.IPUPlace):
         num_devices = core.get_ipu_device_count()
         device = f"ipus:{{0-{num_devices - 1}}}"
-        device = f"ipus:{{0-{num_devices - 1}}}"
     elif isinstance(place, core.CustomPlace):
         device_id = place.get_device_id()
         device_type = place.get_device_type()
diff --git a/python/paddle/jit/sot/opcode_translator/executor/dispatch_functions.py b/python/paddle/jit/sot/opcode_translator/executor/dispatch_functions.py
@@ -59,3 +59,11 @@ def tensor_dim(x):
 
 def generator_send(x):
     pass
+
+
+def place_get_device_id():
+    pass
+
+
+def place_get_device_type():
+    pass
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py b/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
@@ -47,6 +47,8 @@
     operator_is_none,
     operator_is_not_none,
     operator_not_in,
+    place_get_device_id,
+    place_get_device_type,
     tensor_dim,
 )
 from .dispatcher import Dispatcher, optional
@@ -1586,3 +1588,15 @@ def dispatch_all(var: ContainerVariable | IterVariable):
             ufunc,
         ),
     )
+
+# place
+Dispatcher.register(
+    place_get_device_id,
+    ("PlaceVariable",),
+    lambda var: var.get_device_id(),
+)
+Dispatcher.register(
+    place_get_device_type,
+    ("PlaceVariable",),
+    lambda var: var.get_device_type(),
+)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py b/python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py
@@ -32,6 +32,7 @@
     NumpyVariable,
     ObjectVariable,
     ParameterVariable,
+    PlaceVariable,
     SliceVariable,
     SuperVariable,
     SymbolicVariable,
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -23,6 +23,7 @@
 import numpy as np
 
 import paddle
+from paddle._typing import unreached
 from paddle.framework import core
 
 from ....infer_meta import (
@@ -54,7 +55,11 @@
     InnerError,
     UnsupportedPaddleAPIBreak,
 )
-from ..dispatch_functions import tensor_dim
+from ..dispatch_functions import (
+    place_get_device_id,
+    place_get_device_type,
+    tensor_dim,
+)
 from ..guard import (
     FasterStringifiedExpression,
     StringifiedExpression,
@@ -1428,7 +1433,7 @@ def make_stringified_guard(self) -> None:
 
     @VariableFactory.register_from_value()
     def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
-        if isinstance(value, (np.ndarray)):
+        if isinstance(value, np.ndarray):
             return NumpyArrayVariable(value, graph, tracker)
         return None
 
@@ -1483,7 +1488,7 @@ def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
 class NumpyBoolVariable(NumpyNumberVariable):
     @VariableFactory.register_from_value()
     def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
-        if isinstance(value, (np.bool_)):
+        if isinstance(value, np.bool_):
             return NumpyBoolVariable(value, graph, tracker)
         return None
 
@@ -1519,6 +1524,54 @@ def make_stringified_guard(self) -> list[StringifiedExpression]:
         ]
 
 
+class PlaceVariable(ObjectVariable):
+    def __init__(self, obj, graph, tracker):
+        super().__init__(obj, graph, tracker)
+
+    def getattr(self, name: str, default=None):
+        if default is not None:
+            raise FallbackError(
+                "default argument for getattr is not implemented"
+            )
+        if name not in ["get_device_id", "get_device_type"]:
+            return super().getattr(name, default)
+        from .callable import BuiltinVariable
+
+        if name == "get_device_id":
+            return BuiltinVariable(
+                place_get_device_id, self.graph, DanglingTracker()
+            ).bind_dangling_fn(self, name)
+        elif name == "get_device_type":
+            return BuiltinVariable(
+                place_get_device_type, self.graph, DanglingTracker()
+            ).bind_dangling_fn(self, name)
+        unreached()
+
+    def get_device_id(self):
+        return VariableFactory.from_value(
+            self.value.get_device_id(), self.graph, DummyTracker([self])
+        )
+
+    def get_device_type(self):
+        return VariableFactory.from_value(
+            self.value.get_device_type(), self.graph, DummyTracker([self])
+        )
+
+    @VariableFactory.register_from_value()
+    def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
+        if paddle.is_compiled_with_cuda() and isinstance(
+            value, (paddle.CUDAPlace, paddle.CUDAPinnedPlace)
+        ):
+            return PlaceVariable(value, graph, tracker)
+        if paddle.is_compiled_with_xpu() and isinstance(
+            value, (paddle.XPUPlace, paddle.XPUPinnedPlace)
+        ):
+            return PlaceVariable(value, graph, tracker)
+        if isinstance(value, paddle.CustomPlace):
+            return PlaceVariable(value, graph, tracker)
+        return None
+
+
 class NullVariable(VariableBase):
     """
     NullVariable is a subclass of VariableBase used to represent a placeholder variable that has no value or reference associated with it.
diff --git a/test/sot/test_sot_place.py b/test/sot/test_sot_place.py
@@ -0,0 +1,104 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+import unittest
+from contextlib import contextmanager
+
+from test_case_base import (
+    TestCaseBase,
+    test_instruction_translator_cache_context,
+)
+
+import paddle
+from paddle.jit.sot.psdb import check_no_breakgraph
+
+
+@contextmanager
+def device_guard(place: str):
+    original_place = paddle.get_device()
+    try:
+        paddle.set_device(place)
+        yield
+    finally:
+        paddle.set_device(original_place)
+
+
+@check_no_breakgraph
+def run_diff_logic_by_check_expected_place(x: paddle.Tensor):
+    expected_place_str = paddle.get_device()
+    if "cpu" in expected_place_str:
+        return x + 1
+    elif "gpu" in expected_place_str:
+        return x + 2
+    elif "xpu" in expected_place_str:
+        return x + 3
+    elif "npu" in expected_place_str:
+        return x + 4
+    return x
+
+
+class TestCheckExpectedPlace(TestCaseBase):
+    def test_check_cpu(self):
+        x = paddle.to_tensor(0.0)
+        with device_guard("cpu"):
+            self.assert_results(run_diff_logic_by_check_expected_place, x.cpu())
+
+    @unittest.skipUnless(
+        paddle.is_compiled_with_cuda(),
+        "This test case needs to be compiled with CUDA",
+    )
+    def test_check_gpu(self):
+        x = paddle.to_tensor(0.0)
+        with device_guard("gpu"):
+            self.assert_results(
+                run_diff_logic_by_check_expected_place, x.cuda()
+            )
+
+    @unittest.skipUnless(
+        paddle.is_compiled_with_xpu(),
+        "This test case needs to be compiled with XPU",
+    )
+    def test_check_xpu(self):
+        x = paddle.to_tensor(0.0)
+        with device_guard("xpu"):
+            self.assert_results(
+                run_diff_logic_by_check_expected_place, x.to("xpu")
+            )
+
+
+class TestExpectedPlaceGuard(TestCaseBase):
+    @unittest.skipUnless(
+        paddle.is_compiled_with_cuda(),
+        "This test case needs to be compiled with cuda",
+    )
+    def test_expected_place_guard(self):
+        x = paddle.to_tensor(0.0)
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            with device_guard("cpu"):
+                self.assert_results(
+                    run_diff_logic_by_check_expected_place, x.cpu()
+                )
+            self.assertEqual(ctx.translate_count, 1)
+            with device_guard("gpu"):
+                self.assert_results(
+                    run_diff_logic_by_check_expected_place, x.cuda()
+                )
+            self.assertEqual(ctx.translate_count, 2)
+
+
+if __name__ == "__main__":
+    unittest.main()
