[SOT] Add fine gate side effects changed check for dict like proxy data (#73549)

Author: SigureMo

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

@@ -148,6 +148,9 @@ def version(self):
     def has_changed(self):
         return self.version != 0
 
+    def check_changed(self, key: Any) -> bool:
+        raise NotImplementedError
+
     def rollback(self, version: int):
         assert version <= self.version
         self.records[:] = self.records[:version]
@@ -182,6 +185,17 @@ def __init__(self, data: Any, getter: DataGetter):
     def clear_read_cache(self):
         self.read_cache.clear()
 
+    def check_changed(self, key: Any) -> bool:
+        if not self.has_changed:
+            return False
+        for mutation in self.records:
+            if (
+                isinstance(mutation, (MutationNew, MutationDel, MutationSet))
+                and mutation.key == key
+            ):
+                return True
+        return False
+
     def get(self, key: Any):
         # TODO(SigureMo): Optimize performance of this.
         write_cache = self.reproduce(self.version)
@@ -240,6 +254,9 @@ def __init__(self, data: Any, getter: DataGetter):
     def clear_read_cache(self):
         self.read_cache[:] = []
 
+    def check_changed(self, key: Any) -> bool:
+        return self.has_changed
+
     @property
     def length(self):
         return len(self.reproduce())

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

@@ -2475,7 +2475,7 @@ def proxy_getter(self, proxy: MutableDictLikeData, key: Any):
         return VariableFactory.from_value(
             proxy.original_data[key],
             self.graph,
-            tracker=GetAttrTracker(self, key, changed=proxy.has_changed),
+            tracker=GetAttrTracker(self, key, changed=proxy.check_changed(key)),
         )
 
     def get_py_type(self):

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

@@ -199,7 +199,7 @@ def proxy_getter(self, proxy: MutableListLikeData, key: Any):
         return VariableFactory.from_value(
             proxy.original_data[key],
             self.graph,
-            tracker=GetItemTracker(self, key, changed=proxy.has_changed),
+            tracker=GetItemTracker(self, key, changed=proxy.check_changed(key)),
         )
 
     def get_py_value(self, allow_tensor=False):
@@ -249,7 +249,7 @@ def getitem(self, key):
                 items[key],
                 self.graph,
                 tracker=GetItemTracker(
-                    self, key, changed=self.proxy.has_changed
+                    self, key, changed=self.proxy.check_changed(key)
                 ),
             )
         else:
@@ -870,7 +870,7 @@ def proxy_getter(self, proxy: MutableDictLikeData, key: Any):
         return VariableFactory.from_value(
             proxy.original_data[key],
             self.graph,
-            tracker=GetItemTracker(self, key, changed=proxy.has_changed),
+            tracker=GetItemTracker(self, key, changed=proxy.check_changed(key)),
         )
 
     def get_py_value(self, allow_tensor=False):

test/sot/test_side_effects.py

@@ -181,6 +181,24 @@ def slice_list_after_change(l):
     return sum
 
 
+class ReadBufferAfterChanged(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.buffer1 = paddle.to_tensor(1)
+        self.buffer2 = paddle.to_tensor(2)
+
+    def forward(self, x):
+        self.buffer1 += 1
+        return x + self.buffer1 + self.buffer2
+
+    def __eq__(self, other):
+        if not isinstance(other, ReadBufferAfterChanged):
+            return False
+        return paddle.equal(self.buffer1, other.buffer1) and paddle.equal(
+            self.buffer2, other.buffer2
+        )
+
+
 class TestDictSideEffect(TestCaseBase):
     def test_dict_setitem(self):
         self.assert_results_with_side_effects(
@@ -246,6 +264,7 @@ def test_list_insert(self):
     def test_list_remove(self):
         self.assert_results_with_side_effects(list_remove, [1, 1, 1])
         self.assert_results_with_side_effects(list_remove, [0, 1, 2])
+        # TODO(DrRyanHuang): change this to ValueError
         with self.assertRaises(InnerError):
             symbolic_translate(list_remove)([0, 2, 4])
 
@@ -329,5 +348,16 @@ def test_attr_set_breakgraph(self):
         self.attr_check(object_attr_breakgraph, ["x"], CustomObject, 1000)
 
 
+class TestReadBufferAfterChanged(TestCaseBase):
+    def test_read_buffer_after_change(self):
+        layer = ReadBufferAfterChanged()
+        x = paddle.randn([1, 2, 3])
+        self.assert_results_with_side_effects(
+            layer.__class__.forward,
+            layer,
+            x,
+        )
+
+
 if __name__ == "__main__":
     unittest.main()