[SOT] Explicit access layer attribute from _parameters/_sub_layers/_buffers to avoid __getattr__ overhead and move dataclass utilities to utils.py (#73469)

SigureMo · web-flow · commit 2d2065a25de5 · 2025-06-20T10:06:12.000+08:00
diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
@@ -127,6 +127,7 @@
     linalg as linalg,
     signal as signal,
     tensor as tensor,
+    utils as utils,
 )
 from .autograd import (
     enable_grad,
diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py
@@ -16,6 +16,7 @@
 
 import atexit
 import builtins
+import dataclasses
 import functools
 import importlib.util
 import inspect
@@ -287,13 +288,30 @@ def type_name(v):
     return type(v).__name__
 
 
-def _is_dataclass_instance(obj):
+def is_dataclass_instance(obj):
     """Check if the object is an instance of a dataclass.
     Refer to https://docs.python.org/3/library/dataclasses.html#dataclasses.is_dataclass
     """
     return is_dataclass(obj) and not isinstance(obj, type)
 
 
+def is_dataclass_type(obj):
+    return is_dataclass(obj) and isinstance(obj, type)
+
+
+def dataclass_as_dict(obj):
+    return {f.name: getattr(obj, f.name) for f in dataclasses.fields(obj)}
+
+
+def dataclass_from_dict(dataclass_type: type[Any], data: dict[str, Any]):
+    # NOTE(SigureMo): Create dataclass without __post_init__,
+    # because __post_init__ has been run in simulation
+    instance = dataclass_type.__new__(dataclass_type, **data)
+    for fd in dataclasses.fields(dataclass_type):
+        setattr(instance, fd.name, data[fd.name])
+    return instance
+
+
 def make_hashable(x, error_msg=None):
     """
     Makes input `x` hashable.
@@ -303,7 +321,7 @@ def make_hashable(x, error_msg=None):
     if isinstance(x, (tuple, list, set)):
         return tuple(map(make_hashable, x))
 
-    if _is_dataclass_instance(x):
+    if is_dataclass_instance(x):
         return (
             type(x).__name__,
             *map(
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/base.py b/python/paddle/jit/sot/opcode_translator/executor/variables/base.py
@@ -23,6 +23,9 @@
 from typing import TYPE_CHECKING, Any, Callable, Optional
 
 import paddle
+from paddle.jit.dy2static.utils import (
+    dataclass_from_dict,
+)
 
 from ....profiler import event_register
 from ....utils import (
@@ -157,7 +160,7 @@ def _map_slice_variable(variable: VariableBase | object):
     def _map_dataclass_variable(variable: VariableBase | object):
         if not isinstance(variable, DataClassInstanceVariable):
             return variable
-        new_dataclass = DataClassInstanceVariable._dataclass_from_dict(
+        new_dataclass = dataclass_from_dict(
             variable.get_py_type(),
             {
                 fd.name: map_func(variable.getattr(fd.name))
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -18,7 +18,6 @@
 import operator
 import sys
 import types
-from dataclasses import is_dataclass
 from enum import Enum
 from functools import cached_property, reduce
 from typing import TYPE_CHECKING, Any
@@ -28,6 +27,11 @@
 import paddle
 from paddle._typing import unreached
 from paddle.framework import core
+from paddle.jit.dy2static.utils import (
+    dataclass_as_dict,
+    dataclass_from_dict,
+    is_dataclass_instance,
+)
 from paddle.jit.sot.opcode_translator.executor.pycode_generator import PyCodeGen
 from paddle.pir.core import _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE
 
@@ -453,7 +457,7 @@ def __init__(
             and not self.meta.is_null()
         ):
             dynamic_axes = self.analyse_dynamic_axes(tracker)
-        self.var_name = TensorVariable.var_name_generator.next()
+        self.var_name = self.var_name_generator.next()
         self.graph.side_effects.record_mutable_variable(self)
         self.meta = self.meta.with_dynamic_axes(self.var_name, dynamic_axes)
         self.origin_meta = self.meta
@@ -1463,6 +1467,8 @@ def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
 
 
 class ParameterVariable(TensorVariable):
+    var_name_generator = NameGenerator("param_")
+
     def __init__(
         self,
         meta: MetaInfoOrNull,
@@ -2428,18 +2434,9 @@ def __init__(
             id_getter=lambda _: data_id or id(data_dict),
         )
 
-    @staticmethod
-    def _dataclass_from_dict(dataclass_type: type[Any], data: dict[str, Any]):
-        # NOTE(SigureMo): Create dataclass without __post_init__,
-        # because __post_init__ has been run in simulation
-        instance = dataclass_type.__new__(dataclass_type, **data)
-        for fd in dataclasses.fields(dataclass_type):
-            setattr(instance, fd.name, data[fd.name])
-        return instance
-
     def _reconstruct(self, codegen: PyCodeGen) -> None:
         codegen.gen_load_object(
-            DataClassInstanceVariable._dataclass_from_dict,
+            dataclass_from_dict,
             "___dataclass_from_dict",
         )
         self.getattr("__class__").reconstruct(codegen)
@@ -2483,7 +2480,7 @@ def get_py_type(self):
         return self.class_var.get_py_value()
 
     def get_py_value(self, allow_tensor=False):
-        return DataClassInstanceVariable._dataclass_from_dict(
+        return dataclass_from_dict(
             self.get_py_type(),
             {
                 key: value.get_py_value(allow_tensor)
@@ -2565,18 +2562,14 @@ def flatten_inner_vars(self) -> list[VariableBase]:
             for fd in dataclasses.fields(self.get_py_type())
         ]
 
-    @classmethod
-    def _custom_asdict(cls, obj):
-        return {f.name: getattr(obj, f.name) for f in dataclasses.fields(obj)}
-
     @VariableFactory.register_from_value()
     def from_value(value: object, graph: FunctionGraph, tracker: Tracker):
-        if is_dataclass(value) and not isinstance(value, type):
+        if is_dataclass_instance(value):
             class_var = VariableFactory.from_value(
                 type(value), graph, DanglingTracker()
             )
             var = DataClassInstanceVariable(
-                DataClassInstanceVariable._custom_asdict(value),
+                dataclass_as_dict(value),
                 class_var,
                 id(value),
                 graph=graph,
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -23,7 +23,7 @@
 import random
 import sys
 import types
-from dataclasses import fields, is_dataclass
+from dataclasses import fields
 from functools import partial, reduce
 from typing import (
     TYPE_CHECKING,
@@ -36,7 +36,7 @@
     _DecoratorContextManager,
     in_sot_simulation_mode,
 )
-from paddle.jit.dy2static.utils import TransformOptions
+from paddle.jit.dy2static.utils import TransformOptions, is_dataclass_type
 
 from .... import psdb
 from ....profiler import EventGuard
@@ -659,6 +659,47 @@ def __init__(
         super().__init__(graph, tracker)
         self.value = layer
 
+    def getattr(self, name: str, default=None):
+        # TODO(SigureMo): Use more common logic to handle this case,
+        # e.g. call Layer's __getattr__
+        layer = self.value
+        if (
+            '_parameters' in layer.__dict__
+            and name in layer.__dict__["_parameters"]
+        ):
+            return self.getattr("_parameters").getitem(
+                ConstantVariable.wrap_literal(name, self.graph)
+            )
+        if (
+            '_sub_layers' in layer.__dict__
+            and name in layer.__dict__["_sub_layers"]
+        ):
+            out = self.getattr("_sub_layers").getitem(
+                ConstantVariable.wrap_literal(name, self.graph)
+            )
+            return out
+        if '_buffers' in layer.__dict__ and name in layer.__dict__["_buffers"]:
+            return self.getattr("_buffers").getitem(
+                ConstantVariable.wrap_literal(name, self.graph)
+            )
+        return super().getattr(name, default)
+
+    def setattr(self, name: str, value):
+        layer = self.value
+        if (
+            '_parameters' in layer.__dict__
+            and name in layer.__dict__["_parameters"]
+        ):
+            return self.getattr("_parameters").setitem(name, value)
+        if (
+            '_sub_layers' in layer.__dict__
+            and name in layer.__dict__["_sub_layers"]
+        ):
+            return self.getattr("_sub_layers").setitem(name, value)
+        if '_buffers' in layer.__dict__ and name in layer.__dict__["_buffers"]:
+            return self.getattr("_buffers").setitem(name, value)
+        return super().setattr(name, value)
+
     def get_py_value(self, allow_tensor=False):
         return self.value
 
@@ -1311,7 +1352,7 @@ def create_dataclass_init_signature(
 
     @VariableFactory.register_from_value(successor="ClassVariable")
     def from_value(value: object, graph: FunctionGraph, tracker: Tracker):
-        if is_dataclass(value) and isinstance(value, type):
+        if is_dataclass_type(value):
             var = DataClassVariable(value, graph=graph, tracker=tracker)
             return var
         return None
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -23,14 +23,19 @@
 import weakref
 from collections import OrderedDict
 from contextlib import contextmanager
+from dataclasses import is_dataclass
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Callable, TypeVar
 from weakref import WeakValueDictionary
 
 import numpy as np
 
 import paddle
-from paddle.jit.dy2static.utils import TransformOptions
+from paddle.jit.dy2static.utils import (
+    TransformOptions,
+    dataclass_as_dict,
+    dataclass_from_dict,
+)
 from paddle.utils import flatten, map_structure
 
 from .envs import (
@@ -281,13 +286,21 @@ def map_if_extend(structure, pred, true_fn, false_fn):
     def wrapped_pred(x):
         if isinstance(x, slice):
             return True
+        if is_dataclass(x) and not isinstance(x, type):
+            return True
         return pred(x)
 
     def wrapped_true_fn(x):
         if isinstance(x, (slice)):
             l = [x.start, x.stop, x.step]
             l = map_if_extend(l, pred, true_fn, false_fn)
             return slice(*l)
+
+        if is_dataclass(x) and not isinstance(x, type):
+            dt_dict = dataclass_as_dict(x)
+            dt_dict = map_if_extend(dt_dict, pred, true_fn, false_fn)
+            return dataclass_from_dict(type(x), dt_dict)
+
         return true_fn(x)
 
     return map_if(

Original file line number	Diff line number	Diff line change
`@@ -127,6 +127,7 @@`
`127`	`127`	`linalg as linalg,`
`128`	`128`	`signal as signal,`
`129`	`129`	`tensor as tensor,`
	`130`	`+ utils as utils,`
`130`	`131`	`)`
`131`	`132`	`from .autograd import (`
`132`	`133`	`enable_grad,`