[SideEffect] add side effects support for ObjectVariable

sot/opcode_translator/executor/function_graph.py

@@ -34,6 +34,8 @@
     GlobalDelSideEffectRestorer,
     GlobalSetSideEffectRestorer,
     ListSideEffectRestorer,
+    ObjDelSideEffectRestorer,
+    ObjSetSideEffectRestorer,
     SideEffectRestorer,
     SideEffects,
 )
@@ -43,6 +45,7 @@
     GlobalVariable,
     ListVariable,
     NullVariable,
+    ObjectVariable,
     PaddleLayerVariable,
     TensorVariable,
     VariableBase,
@@ -555,7 +558,7 @@ def _find_tensor_outputs(
                         var, TensorVariable
                     ):
                         output_tensors.add(var)
-                    if isinstance(var, GlobalVariable):
+                    if isinstance(var, (GlobalVariable, ObjectVariable)):
                         for record in var.proxy.records:
                             if (
                                 isinstance(record, (MutationSet, MutationNew))
@@ -641,6 +644,24 @@ def restore_side_effects(self, variables: list[VariableBase]):
                                 GlobalDelSideEffectRestorer(record.key)
                             )
                 # TODO: support attribute restore
+                elif isinstance(var, ObjectVariable):
+                    for record in var.proxy.records[::-1]:
+                        if isinstance(record, (MutationSet, MutationNew)):
+                            restorers.append(
+                                ObjSetSideEffectRestorer(
+                                    var,
+                                    record.key,
+                                    record.value,
+                                )
+                            )
+                        elif isinstance(record, MutationDel):
+                            restorers.append(
+                                ObjDelSideEffectRestorer(
+                                    var,
+                                    record.key,
+                                )
+                            )
+
         for restorer in restorers:
             restorer.pre_gen(self.pycode_gen)
         for restorer in restorers[::-1]:

sot/opcode_translator/executor/opcode_executor.py

@@ -969,7 +969,7 @@ def LOAD_METHOD(self, instr: Instruction):
             self.stack.push(NullVariable())
             self.stack.push(method)
 
-    def STORE_ATTR(self, instr):
+    def STORE_ATTR(self, instr: Instruction):
         obj = self.stack.pop()
         val = self.stack.pop()
         key = self._code.co_names[instr.arg]
@@ -985,9 +985,11 @@ def STORE_ATTR(self, instr):
                 val,
             )
         else:
-            raise BreakGraphError(
-                f"STORE_ATTR don't support {type(obj)}.{key}={val}"
-            )
+            obj.setattr(key, val)
+
+    def DELETE_ATTR(self, instr: Instruction):
+        obj = self.stack.pop()
+        obj.delattr(instr.argval)
 
     def STORE_DEREF(self, instr):
         namemap = self._code.co_cellvars + self._code.co_freevars

sot/opcode_translator/executor/pycode_generator.py

@@ -721,6 +721,18 @@ def gen_load_attr(self, name: str):
         idx = self._code_options["co_names"].index(name)
         self._add_instr("LOAD_ATTR", arg=idx, argval=name)
 
+    def gen_store_attr(self, name: str):
+        if name not in self._code_options["co_names"]:
+            self._code_options["co_names"].append(name)
+        idx = self._code_options["co_names"].index(name)
+        self._add_instr("STORE_ATTR", arg=idx, argval=name)
+
+    def gen_delete_attr(self, name: str):
+        if name not in self._code_options["co_names"]:
+            self._code_options["co_names"].append(name)
+        idx = self._code_options["co_names"].index(name)
+        self._add_instr("DELETE_ATTR", arg=idx, argval=name)
+
     def gen_load_method(self, name: str):
         if name not in self._code_options["co_names"]:
             self._code_options["co_names"].append(name)

sot/opcode_translator/executor/side_effects.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
 
 from .mutable_data import MutableData
-from .variables import VariableBase
+from .variables import ObjectVariable, VariableBase
 
 if TYPE_CHECKING:
     from .mutable_data import DataGetter
@@ -180,3 +180,59 @@ def pre_gen(self, codegen: PyCodeGen):
 
     def post_gen(self, codegen: PyCodeGen):
         codegen.gen_delete_global(self.name)
+
+
+class ObjSetSideEffectRestorer(SideEffectRestorer):
+    """
+    class CustomObject:
+        def __init__(self):
+            self.x = 0
+
+    def attr_set(cus_obj):
+        cus_obj.x = 2
+
+    cus_obj = CustomObject()
+    attr_set(cus_obj)
+    """
+
+    def __init__(self, obj: ObjectVariable, name: str, var: VariableBase):
+        super().__init__()
+        self.obj = obj
+        self.name = name
+        self.var = var
+
+    def pre_gen(self, codegen: PyCodeGen):
+        self.var.reconstruct(codegen)
+
+    def post_gen(self, codegen: PyCodeGen):
+        # TODO(gouzil): Find this name, doc: cus_obj
+        codegen.gen_load_fast(self.obj.tracker.name)
+        codegen.gen_store_attr(self.name)
+
+
+class ObjDelSideEffectRestorer(SideEffectRestorer):
+    """
+    class CustomObject:
+        def __init__(self):
+            self.x = 0
+
+    def attr_set(cus_obj):
+        del cus_obj.x
+
+    cus_obj = CustomObject()
+    attr_set(cus_obj)
+    """
+
+    def __init__(self, obj: ObjectVariable, name: str):
+        super().__init__()
+        self.obj = obj
+        self.name = name
+
+    def pre_gen(self, codegen: PyCodeGen):
+        # do nothing
+        ...
+
+    def post_gen(self, codegen: PyCodeGen):
+        # TODO(gouzil): Find this name, doc: cus_obj
+        codegen.gen_load_fast(self.obj.tracker.name)
+        codegen.gen_delete_attr(self.name)

sot/opcode_translator/executor/variable_dispatch.py

@@ -51,7 +51,7 @@ def inner(*args, **kwargs):
 # slice
 Dispatcher.register(
     slice,
-    ("VariableBase"),
+    ("VariableBase",),
     lambda stop: SliceVariable(
         slice(stop),
         graph=stop.graph,

sot/opcode_translator/executor/variables/base.py

@@ -490,6 +490,20 @@ def getattr(self, name: str, default=None):
             )
         return result
 
+    def setattr(self, key: str, value):
+        from .basic import ConstantVariable
+
+        self.attr_proxy.set(key, value)
+        self.graph.side_effects.record_proxy_variable(self)
+        return ConstantVariable.wrap_literal(None, self.graph)
+
+    def delattr(self, key: str):
+        from .basic import ConstantVariable
+
+        self.attr_proxy.delete(key)
+        self.graph.side_effects.record_proxy_variable(self)
+        return ConstantVariable.wrap_literal(None, self.graph)
+
     def __setitem__(self, key, value):
         return self.setitem(key, value)
 

sot/opcode_translator/executor/variables/basic.py

@@ -510,6 +510,12 @@ def getattr(self, name: str, default=None):
         else:
             raise HasNoAttributeError(f"Unknown Tensor attribute: {name}")
 
+    def setattr(self, key, val):
+        raise BreakGraphError("Don't support TensorVariable setattr")
+
+    def delattr(self, key):
+        raise BreakGraphError("Don't support TensorVariable delattr")
+
     @VariableFactory.register_from_value()
     def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
         if isinstance(value, (paddle.Tensor, MetaInfo)):
@@ -527,11 +533,12 @@ class ObjectVariable(VariableBase):
         tracker(Tracker): The Tracker object that tracks the information of this variable.
     """
 
+    make_stringify_guard = object_equal_stringify_guard
+
     def __init__(self, obj, graph, tracker):
         super().__init__(graph, tracker)
         self.value = obj
-
-    make_stringify_guard = object_equal_stringify_guard
+        self.proxy = self.attr_proxy
 
     @property
     def main_info(self) -> dict[str, Any]:
@@ -619,6 +626,12 @@ def _reconstruct(self, codegen: PyCodeGen):
         else:
             super()._reconstruct(codegen)
 
+    def setattr(self, key, val):
+        raise BreakGraphError("Don't support SliceVariable setattr")
+
+    def delattr(self, key):
+        raise BreakGraphError("Don't support SliceVariable delattr")
+
     @VariableFactory.register_from_value()
     def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
         if isinstance(value, slice):

sot/opcode_translator/executor/variables/callable.py

@@ -414,6 +414,7 @@ def __init__(
         self, layer: paddle.nn.Layer, graph: FunctionGraph, tracker: Tracker
     ):
         super().__init__(layer, graph, tracker)
+        self.proxy = self.attr_proxy
 
     def call_function(self, /, *args, **kwargs):
         fn_var = UserDefinedFunctionVariable(
@@ -634,6 +635,7 @@ class ClassVariable(CallableVariable):
     def __init__(self, class_: type, graph: FunctionGraph, tracker: Tracker):
         super().__init__(graph, tracker)
         self.value = class_
+        self.proxy = self.attr_proxy
 
     def get_py_value(self, allow_tensor=False):
         return self.value

tests/test_side_effects.py

@@ -126,11 +126,15 @@ def __init__(self):
         self.x = 0
 
 
-def object_attr(cus_obj, t):
+def object_attr_set(cus_obj, t):
     """object side effect."""
     t = t + 1
     cus_obj.x = t
-    return t, cus_obj
+    return t, cus_obj.x
+
+
+def object_attr_del(cus_obj):
+    del cus_obj.x
 
 
 def slice_list_after_change(l):
@@ -252,5 +256,25 @@ def test_slice_list_after_change(self):
         )
 
 
+class TestATTRSideEffect(TestCaseBase):
+    def attr_check(self, func, attr_keys: list[str], obj, *inputs):
+        cus_obj1 = obj()
+        cus_obj2 = obj()
+        sym_output = symbolic_translate(func)(cus_obj1, *inputs)
+        paddle_output = func(cus_obj2, *inputs)
+        for key in attr_keys:
+            self.assert_nest_match(
+                getattr(cus_obj1, key, "Key does not exist"),
+                getattr(cus_obj2, key, "Key does not exist"),
+            )
+        self.assert_nest_match(sym_output, paddle_output)
+
+    def test_attr_set(self):
+        self.attr_check(object_attr_set, ["x"], CustomObject, 5)
+
+    def test_attr_del(self):
+        self.attr_check(object_attr_del, ["x"], CustomObject)
+
+
 if __name__ == "__main__":
     unittest.main()