Support DataClassInstanceVariable

Author: 0x3878f

python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py

@@ -671,6 +671,7 @@ def step(self, instr: Instruction):
             assert opname != "CALL", "CALL should fused with PRECALL"
         with EventGuard(f"{opname}", event_level=2):
             try:
+                # breakpoint()
                 return getattr(self, opname)(instr)  # run single step.
             except SotCapturedException as e:
                 self.handle_exception(e)

python/paddle/jit/sot/opcode_translator/executor/side_effects.py

@@ -17,10 +17,15 @@
 from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from typing_extensions import TypeAlias
+
     from .mutable_data import DataGetter, MutableData
     from .pycode_generator import PyCodeGen
     from .variables import VariableBase
 
+    IdGetter: TypeAlias = Callable[[Any], int]
     MutableDataT = TypeVar("MutableDataT", bound=MutableData)
 
 
@@ -51,8 +56,9 @@ def get_proxy(
         proxy_type: type[MutableDataT],
         data: Any,
         getter: DataGetter,
+        id_getter: IdGetter = id,
     ) -> MutableDataT:
-        data_id = id(data)
+        data_id = id_getter(data)
         if data_id not in self.data_id_to_proxy:
             self.data_id_to_proxy[data_id] = proxy_type(data, getter)
         return self.data_id_to_proxy[data_id]  # type: ignore

python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py

@@ -18,6 +18,7 @@
 import inspect
 import math
 import operator
+from dataclasses import fields
 from functools import partial, reduce
 from typing import TYPE_CHECKING
 
@@ -67,6 +68,7 @@
     CallableVariable,
     ConstantVariable,
     ContainerVariable,
+    DataClassInstanceVariable,
     DictVariable,
     EnumerateVariable,
     EnumVariable,
@@ -1556,6 +1558,31 @@ def exception_variable_equal(left: ExceptionVariable, right: ExceptionVariable):
     lambda left, right: exception_variable_equal(left, right),
 )
 
+
+@Dispatcher.register_decorator(operator.eq)
+def dataclass_instance_eq(
+    lhs: DataClassInstanceVariable, rhs: DataClassInstanceVariable
+):
+    if lhs.get_py_type() != rhs.get_py_type():
+        return ConstantVariable(False, lhs.graph, DummyTracker([lhs, rhs]))
+
+    return ConstantVariable(
+        all(
+            Dispatcher.call(
+                operator.eq, lhs.getattr(field.name), rhs.getattr(field.name)
+            )
+            for field in fields(lhs.get_py_type())
+        ),
+        lhs.graph,
+        DummyTracker([lhs, rhs]),
+    )
+
+
+@Dispatcher.register_decorator(operator.ne)
+def dataclass_instance_ne(lhs: TupleVariable, rhs: TupleVariable):
+    return Dispatcher.call(operator.eq, lhs, rhs).bool_not()
+
+
 Dispatcher.register(
     operator.eq,
     ("EnumVariable", "EnumVariable"),

python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py

@@ -21,6 +21,7 @@
 from .basic import (  # noqa: F401
     CellVariable,
     ConstantVariable,
+    DataClassInstanceVariable,
     DataVariable,
     DygraphTracerVariable,
     EnumVariable,

python/paddle/jit/sot/opcode_translator/executor/variables/basic.py

@@ -17,6 +17,7 @@
 import operator
 import sys
 import types
+from dataclasses import dataclass, is_dataclass
 from enum import Enum
 from functools import cached_property, reduce
 from typing import TYPE_CHECKING, Any
@@ -2396,3 +2397,133 @@ def get_enum_class_id(cls, enum_class: type[Enum]):
             id = len(same_name_enums)
             same_name_enums.append(enum_class)
         return id
+
+
+class DataClassInstanceVariable(VariableBase):
+    known_dataclasses = {}
+
+    def __init__(
+        self,
+        value,
+        graph: FunctionGraph,
+        tracker: Tracker,
+    ):
+        super().__init__(graph=graph, tracker=tracker)
+        self.dataclass_type = type(value)
+        var_dict = value.__dict__
+        var_dict.update(
+            {'__post_init__': getattr(value, '__post_init__', None)}
+        )
+        self.proxy = self.graph.side_effects.get_proxy(
+            MutableDictLikeData,
+            var_dict,
+            self.proxy_getter,
+            id_getter=lambda _: id(value),
+        )
+
+    @cached_property
+    def attr_proxy(self):
+        return self.proxy
+
+    def getattr(self, name: str, default=None):
+        return self.proxy.get(name)
+
+    def setattr(self, key: str, value):
+        self.proxy.set(key, value)
+        self.graph.side_effects.record_proxy_variable(self)
+        return ConstantVariable.wrap_literal(None, self.graph)
+
+    def proxy_getter(self, proxy: MutableDictLikeData, key: Any):
+        if key not in proxy.original_data:
+            return MutableDictLikeData.Empty()
+        return VariableFactory.from_value(
+            proxy.original_data[key],
+            self.graph,
+            tracker=GetAttrTracker(self, key, changed=proxy.has_changed),
+        )
+
+    def get_py_type(self):
+        return self.dataclass_type
+
+    def get_py_value(self, allow_tensor=False):
+        return self.dataclass_type(
+            **{
+                key: value.get_py_value(allow_tensor)
+                for key, value in self.proxy.get_all().items()
+            }
+        )
+
+    @check_faster_guard
+    def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
+        expr_node = self.tracker.guard_tree_expr_node()
+        type_guard = paddle.framework.core.GuardNode(
+            paddle.framework.core.TypeMatchGuard(self.get_py_type()),
+            [expr_node],
+        )
+        guard_variables = filter(
+            lambda var: not isinstance(var, MutableDictLikeData.Empty),
+            self.proxy.reproduce(0).values(),
+        )
+        return reduce(
+            operator.add,
+            [[type_guard]]
+            + [
+                item.make_faster_guard()
+                for item in guard_variables
+                if item.tracker.need_guard()
+            ],
+        )
+
+    @check_guard
+    def make_stringified_guard(self) -> list[StringifiedExpression]:
+        data_class = self.get_py_type()
+        class_name = data_class.__name__
+        data_class_id = DataClassInstanceVariable.get_class_id(data_class)
+        extern_var_name = f"__{class_name}_{data_class_id}"
+        frame_value_tracer = self.tracker.trace_value_from_frame()
+        type_guard = FasterStringifiedExpression(
+            f"isinstance({{}}, {extern_var_name})",
+            paddle.framework.core.InstanceCheckGuard(self.get_py_type()),
+            [frame_value_tracer],
+            union_free_vars(
+                frame_value_tracer.free_vars,
+                {f"{extern_var_name}": self.get_py_type()},
+            ),
+        )
+        guard_variables = filter(
+            lambda var: not isinstance(var, MutableDictLikeData.Empty),
+            self.proxy.reproduce(0).values(),
+        )
+        return reduce(
+            operator.add,
+            [[type_guard]]
+            + [
+                item.make_stringified_guard()
+                for item in guard_variables
+                if item.tracker.need_guard()
+            ],
+        )
+
+    @VariableFactory.register_from_value()
+    def from_value(value: dataclass, graph: FunctionGraph, tracker: Tracker):
+        if is_dataclass(value) and not isinstance(value, type):
+            var = DataClassInstanceVariable(value, graph=graph, tracker=tracker)
+            return var
+        return None
+
+    @classmethod
+    def get_class_id(cls, data_class: type[dataclass]):
+        class_name = data_class.__name__
+        DataClassInstanceVariable.known_dataclasses.setdefault(class_name, [])
+        same_name_dataclasses = DataClassInstanceVariable.known_dataclasses[
+            class_name
+        ]
+        id = 0
+        for i, cls in enumerate(same_name_dataclasses):
+            if data_class == cls:
+                id = i
+                break
+        else:
+            id = len(same_name_dataclasses)
+            same_name_dataclasses.append(data_class)
+        return id

test/sot/test_dataclass.py

@@ -12,12 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import unittest
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from enum import IntEnum
+from typing import Callable
 
-from test_case_base import TestCaseBase
+from test_case_base import (
+    TestCaseBase,
+    test_instruction_translator_cache_context,
+)
 
 import paddle
+from paddle.jit.sot.psdb import check_no_breakgraph
 from paddle.jit.sot.utils import strict_mode_guard
 
 
@@ -54,5 +62,101 @@ def test_dtype_reconstruct_with_post_init(self):
         self.assert_results(return_dataclass_with_post_init, x)
 
 
+class DataType(IntEnum):
+    FLOAT32 = 1
+    FLOAT64 = 2
+    INT32 = 3
+    INT64 = 4
+
+
+@dataclass
+class DataMeta:
+    x: paddle.Tensor
+    y: paddle.Tensor = None
+    z: DataType = DataType.FLOAT32
+    m: list[list[paddle.Tensor]] = field(default_factory=list)
+    n: int = 0
+    f: Callable[[DataMeta], list] = None
+
+    def __post_init__(self):
+        self.x += 1
+
+
+@check_no_breakgraph
+def is_eq(data: DataMeta, data2: DataMeta):
+    return data == data2
+
+
+@check_no_breakgraph
+def get_attr(data: DataMeta):
+    return data.x + data.y
+
+
+@check_no_breakgraph
+def set_attr(data: DataMeta):
+    ori_x = data.x
+    data.x = data.x + data.n
+    res = data.x
+    data.x = ori_x
+    return res
+
+
+@check_no_breakgraph
+def callable_attr(data: DataMeta):
+    return data.f(data)
+
+
+class TestDataClassInstance(TestCaseBase):
+    def test_guard(self):
+        d1 = Data(x=paddle.randn([1]))
+        dm1 = DataMeta(x=paddle.randn([1]))
+        dm2 = DataMeta(x=paddle.randn([1]))
+        dm3 = DataMeta(x=paddle.zeros([1]))
+        dm4 = DataMeta(x=paddle.randn([1]), z=DataType.INT32)
+        dm5 = DataMeta(x=paddle.randn([1]), n=1)
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            self.assert_results(is_eq, dm1, dm2)
+            self.assertEqual(ctx.translate_count, 1)
+            self.assert_results(is_eq, dm1, dm2)
+            self.assertEqual(ctx.translate_count, 1)
+            self.assert_results(is_eq, dm1, d1)
+            self.assertEqual(ctx.translate_count, 2)
+            self.assert_results(is_eq, dm1, dm3)
+            self.assertEqual(ctx.translate_count, 2)
+            self.assert_results(is_eq, dm1, dm4)
+            self.assertEqual(ctx.translate_count, 3)
+            self.assert_results(is_eq, dm1, dm5)
+            self.assertEqual(ctx.translate_count, 4)
+
+    def test_get_attr(self):
+        dm = DataMeta(x=paddle.randn([1, 2]), y=paddle.randn([1]))
+        self.assert_results(get_attr, dm)
+
+    def test_set_attr(self):
+        dm = DataMeta(x=paddle.ones([1, 2]), n=2)
+        self.assert_results(set_attr, dm)
+
+    def test_callable_attr(self):
+
+        def process_func(data: DataMeta):
+            return data.x.shape
+
+        dm = DataMeta(x=paddle.randn([1, 2]), f=process_func)
+        self.assert_results(callable_attr, dm)
+
+    def test_eq(self):
+        dm1 = DataMeta(x=paddle.randn([1]))
+        dm2 = DataMeta(x=paddle.randn([1]))
+        dm3 = DataMeta(x=paddle.zeros([1]))
+        dm4 = DataMeta(x=paddle.randn([1]), z=DataType.INT32)
+        self.assert_results(is_eq, dm1, dm2)
+        self.assert_results(is_eq, dm1, dm3)
+        self.assert_results(is_eq, dm1, dm4)
+        # TODO(wangmingkai): operator.eq with args UserDefinedFunctionVariable
+        # dm5 = DataMeta(x= paddle.randn([1]), f=lambda _: [])
+        # self.assert_results(is_eq, dm1, dm5)
+
+
 if __name__ == "__main__":
     unittest.main()