[SOT] Support user specified dynamic dims (#73509)

SigureMo · web-flow · commit 837c6a89177a · 2025-06-21T10:31:05.000+08:00
diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py
@@ -87,6 +87,8 @@
     "ENABLE_CINN_IN_DY2ST", True
 )
 
+DYNAMIC_DIMS_ATTR_NAME = "__sot_dynamic_dims"
+
 
 class Backend(Enum):
     CINN = auto()
@@ -1012,3 +1014,26 @@ def patch_method_guard(
         yield
     finally:
         restorer(instance)
+
+
+def extract_tensor_dynamic_dims(
+    tensor: paddle.Tensor,
+) -> tuple[int]:
+    """
+    Extract dynamic dimensions from a paddle.Tensor.
+    Returns a list of dynamic dimensions or None if no dynamic dimensions exist.
+    """
+    if not isinstance(tensor, paddle.Tensor):
+        raise TypeError(
+            f"Expected a paddle.Tensor, but got {type(tensor).__name__}"
+        )
+
+    if not hasattr(tensor, DYNAMIC_DIMS_ATTR_NAME):
+        return []
+
+    dynamic_dims = getattr(tensor, DYNAMIC_DIMS_ATTR_NAME)
+    if not isinstance(dynamic_dims, tuple):
+        raise TypeError(
+            f"Expected {DYNAMIC_DIMS_ATTR_NAME} to be a tuple, but got {type(dynamic_dims).__name__}"
+        )
+    return dynamic_dims
diff --git a/python/paddle/jit/marker.py b/python/paddle/jit/marker.py
@@ -15,23 +15,23 @@
 from __future__ import annotations
 
 import inspect
-from typing import (
-    Callable,
-    Protocol,
-    TypeVar,
-    overload,
-)
+from typing import TYPE_CHECKING, Callable, Protocol, TypeVar, overload
 
 from typing_extensions import (
     ParamSpec,
 )
 
 import paddle
+from paddle.jit.dy2static.utils import DYNAMIC_DIMS_ATTR_NAME
 
 from .dy2static.utils import (
     TransformOptions,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
 _RetT = TypeVar("_RetT")
 _InputT = ParamSpec("_InputT")
 
@@ -147,3 +147,32 @@ def force_dynamic(
     raise TypeError(
         f"Expected a callable or paddle.nn.Layer, but got {type(fn).__name__}."
     )
+
+
+def dynamic_dims(
+    tensor: paddle.Tensor,
+    dims: int | Sequence[int],
+) -> None:
+    """
+    Mark a tensor as having dynamic dimensions.
+    This is used to indicate that the tensor's shape may change dynamically
+    during execution, which is particularly useful in dynamic-to-static
+    conversion scenarios.
+    Args:
+        tensor (paddle.Tensor): The tensor to mark as dynamic.
+        dims (int | Sequence[int]): The dimensions to mark as dynamic.
+    """
+    if not isinstance(tensor, paddle.Tensor):
+        raise TypeError(
+            f"Expected a paddle.Tensor, but got {type(tensor).__name__}."
+        )
+
+    if isinstance(dims, int):
+        dims = (dims,)
+
+    if not isinstance(dims, (list, tuple)):
+        raise TypeError("Dimensions must be a list or tuple.")
+    if not all(isinstance(dim, int) for dim in dims):
+        raise TypeError("All dimensions must be integers.")
+
+    setattr(tensor, DYNAMIC_DIMS_ATTR_NAME, tuple(dims))
diff --git a/python/paddle/jit/sot/infer_meta.py b/python/paddle/jit/sot/infer_meta.py
@@ -36,6 +36,7 @@
     convert_to_dims_mapping,
 )
 from paddle.framework import use_pir_api
+from paddle.jit.dy2static.utils import extract_tensor_dynamic_dims
 from paddle.pir import is_fake_value
 from paddle.static import InputSpec
 from paddle.utils import flatten, is_sequence
@@ -157,6 +158,10 @@ def _handle_legacy_ir_amp_dtype(dtype):
             dtype = paddle.float32
         return dtype
 
+    @staticmethod
+    def mix_axes(axes1: list[int], axes2: list[int]) -> list[int]:
+        return sorted(set(axes1 + axes2))
+
     @staticmethod
     def from_tensor(
         tensor: paddle.Tensor, *, dynamic_axes: list[int] | None = None
@@ -171,9 +176,13 @@ def from_tensor(
         assert (
             -1 not in tensor.shape
         ), "Tensor shape should not contain -1, maybe you pass a Value to from_tensor"
+        user_specified_dynamic_axes = extract_tensor_dynamic_dims(tensor)
         dynamic_axes = dynamic_axes or []
+        dynamic_axes = MetaInfoOrNull.mix_axes(
+            dynamic_axes, list(user_specified_dynamic_axes)
+        )
         shape = [
-            SymbolicInt() if i in dynamic_axes else dim
+            SymbolicInt(dim) if i in dynamic_axes else dim
             for i, dim in enumerate(tensor.shape)
         ]
         if tensor.is_dist():
@@ -302,10 +311,19 @@ def shape_with_special_symbol(
         ]
 
     def with_dynamic_axes(self, name: str, dynamic_axes: list[int]) -> MetaInfo:
+        mixed_dynamic_axes = MetaInfoOrNull.mix_axes(
+            self.dynamic_axes, dynamic_axes
+        )
         # NOTE(SigureMo): Make sure create a new shape list with dynamic axes.
         # We will create a new shape list variable lazily in the future.
         shape = [
-            SymbolicInt(dim) if i in dynamic_axes else dim
+            (
+                SymbolicInt(dim)
+                if (
+                    i in mixed_dynamic_axes and not isinstance(dim, SymbolicInt)
+                )
+                else dim
+            )
             for i, dim in enumerate(self.shape)
         ]
         return MetaInfo(
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -1412,6 +1412,8 @@ def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
         if not ENV_SOT_ALLOW_DYNAMIC_SHAPE.get():
             return None
         if isinstance(value, SymbolicInt):
+            if value.is_backed():
+                return SymbolicVariable(value, graph, tracker)
             tensor_shape_source_result = (
                 SymbolicVariable.find_tensor_shape_source(tracker)
             )
diff --git a/test/sot/test_user_specified_dynamic_dims.py b/test/sot/test_user_specified_dynamic_dims.py
@@ -0,0 +1,120 @@
+# 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 test_case_base import (
+    TestCaseBase,
+    test_instruction_translator_cache_context,
+)
+
+import paddle
+from paddle.jit.sot.psdb import check_no_breakgraph
+
+
+@check_no_breakgraph
+def tensor_with_marked_dynamic_dims(tensor):
+    return tensor + 1
+
+
+@check_no_breakgraph
+def multiple_tensors_with_marked_dynamic_dims(tensor1, tensor2):
+    return tensor1.sum() + tensor2.sum()
+
+
+class TestUserSpecifiedDynamicDims(TestCaseBase):
+    def test_auto_inferred_dynamic_dims(self):
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            x1 = paddle.rand([5, 6], dtype='float32')
+            self.assert_results(tensor_with_marked_dynamic_dims, x1)
+            self.assertEqual(ctx.translate_count, 1)
+            x2 = paddle.rand([4, 6], dtype='float32')
+            self.assert_results(tensor_with_marked_dynamic_dims, x2)
+            self.assertEqual(ctx.translate_count, 2)
+            x3 = paddle.rand([3, 6], dtype='float32')
+            self.assert_results(tensor_with_marked_dynamic_dims, x3)
+            self.assertEqual(ctx.translate_count, 2)
+
+    def test_user_specified_dynamic_dims(self):
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            x1 = paddle.rand([5, 6], dtype='float32')
+            paddle.jit.marker.dynamic_dims(x1, [0])
+            self.assert_results(tensor_with_marked_dynamic_dims, x1)
+            self.assertEqual(ctx.translate_count, 1)
+            x2 = paddle.rand([4, 6], dtype='float32')
+            self.assert_results(tensor_with_marked_dynamic_dims, x2)
+            self.assertEqual(ctx.translate_count, 1)
+            x3 = paddle.rand([3, 6], dtype='float32')
+            self.assert_results(tensor_with_marked_dynamic_dims, x3)
+            self.assertEqual(ctx.translate_count, 1)
+
+    def test_user_specified_multiple_dynamic_dims(self):
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            x1 = paddle.rand([5, 6], dtype='float32')
+            paddle.jit.marker.dynamic_dims(x1, [0, 1])
+            self.assert_results(tensor_with_marked_dynamic_dims, x1)
+            self.assertEqual(ctx.translate_count, 1)
+            x2 = paddle.rand([4, 6], dtype='float32')
+            self.assert_results(tensor_with_marked_dynamic_dims, x2)
+            self.assertEqual(ctx.translate_count, 1)
+            x3 = paddle.rand([3, 7], dtype='float32')
+            self.assert_results(tensor_with_marked_dynamic_dims, x3)
+            self.assertEqual(ctx.translate_count, 1)
+
+    def test_multiple_tensors_with_marked_dynamic_dims(self):
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            x1 = paddle.rand([5, 6], dtype='float32')
+            x2 = paddle.rand([5, 6], dtype='float32')
+            paddle.jit.marker.dynamic_dims(x1, [0])
+            paddle.jit.marker.dynamic_dims(x2, [0])
+            self.assert_results(
+                multiple_tensors_with_marked_dynamic_dims, x1, x2
+            )
+            self.assertEqual(ctx.translate_count, 1)
+            x3 = paddle.rand([4, 6], dtype='float32')
+            x4 = paddle.rand([5, 6], dtype='float32')
+            self.assert_results(
+                multiple_tensors_with_marked_dynamic_dims, x3, x4
+            )
+            self.assertEqual(ctx.translate_count, 1)
+            x5 = paddle.rand([4, 6], dtype='float32')
+            x6 = paddle.rand([4, 6], dtype='float32')
+            self.assert_results(
+                multiple_tensors_with_marked_dynamic_dims, x5, x6
+            )
+            self.assertEqual(ctx.translate_count, 1)
+
+    def test_mix_auto_inferred_and_specified_dynamic_dims(self):
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            x1 = paddle.rand([5, 6], dtype='float32')
+            paddle.jit.marker.dynamic_dims(x1, [0])
+            self.assert_results(tensor_with_marked_dynamic_dims, x1)
+            self.assertEqual(ctx.translate_count, 1)
+            x2 = paddle.rand([4, 6], dtype='float32')
+            self.assert_results(tensor_with_marked_dynamic_dims, x2)
+            self.assertEqual(ctx.translate_count, 1)
+            x3 = paddle.rand([3, 7], dtype='float32')
+            self.assert_results(tensor_with_marked_dynamic_dims, x3)
+            self.assertEqual(ctx.translate_count, 2)
+
+
+if __name__ == '__main__':
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -1412,6 +1412,8 @@ def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):`
`1412`	`1412`	`if not ENV_SOT_ALLOW_DYNAMIC_SHAPE.get():`
`1413`	`1413`	`return None`
`1414`	`1414`	`if isinstance(value, SymbolicInt):`
	`1415`	`+ if value.is_backed():`
	`1416`	`+ return SymbolicVariable(value, graph, tracker)`
`1415`	`1417`	`tensor_shape_source_result = (`
`1416`	`1418`	`SymbolicVariable.find_tensor_shape_source(tracker)`
`1417`	`1419`	`)`