PaddlePaddle
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
Copy file name to clipboard
+7-3Lines changed: 7 additions & 3 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
Copy file name to clipboard
+7-3Lines changed: 7 additions & 3 deletions
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/function_graph.py
Copy file name to clipboard
+4-2Lines changed: 4 additions & 2 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/function_graph.py
Copy file name to clipboard
+4-2Lines changed: 4 additions & 2 deletions
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
Copy file name to clipboard
+9-27Lines changed: 9 additions & 27 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
Copy file name to clipboard
+9-27Lines changed: 9 additions & 27 deletions
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
Copy file name to clipboard
+33-8Lines changed: 33 additions & 8 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
Copy file name to clipboard
+33-8Lines changed: 33 additions & 8 deletions
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/variables/base.py
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎python/paddle/jit/sot/opcode_translator/executor/variables/base.py
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
Copy file name to clipboard
+6-6Lines changed: 6 additions & 6 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
Copy file name to clipboard
+6-6Lines changed: 6 additions & 6 deletions
@@ -25,7 +25,9 @@
 from ...utils import (
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
     BreakGraphError,
+    CompileCountInfo,
     FallbackError,
+    InfoCollector,
     InnerError,
     Singleton,
     is_strict_mode,
@@ -148,9 +150,9 @@ def lookup(
                     )
             except Exception as e:
                 log(2, f"[Cache]: Guard function error: {e}\n")
-                log(
+                log_do(
                     2,
-                    f"[Cache]: Guard string is: {getattr(guard_fn, 'expr', 'None')}\n",
+                    self.analyse_guard_error(guard_fn, frame),
                 )
 
                 continue
@@ -205,8 +207,9 @@ def inner():
                     result = guard(frame)
                 except Exception as e:
                     print(
-                        f"[Cache]: skip checking {guard_str}\n         because error occurred {e}"
+                        f"[Cache]: Error occurred when checking guard {guard_str}: {e}"
                     )
+                    return
                 if result is False:
                     print(f"[Cache]: missed at {guard_str}")
                     return
@@ -231,6 +234,7 @@ def start_translate(
     simulator = OpcodeExecutor(frame, **kwargs)
     try:
         simulator.check_code_simulatable()
+        InfoCollector().attach(CompileCountInfo, frame.f_code)
         with sot_simulation_mode_guard(True):
             new_custom_code, guard_fn = simulator.transform()
         if not simulator._graph.need_cache:
 
@@ -369,8 +369,9 @@ def _restore_origin_opcode(self, stack_vars, store_var_info, instr_idx):
         # here is not update changed values, it just give names to stack vars
         # and want keep same interface as _build_compile_fn_with_name_store
         for var in stack_vars[::-1]:
-            if not store_var_info[var.id]:
+            if not store_var_info.get(var.id, []):
                 name = name_gen.next()
+                store_var_info.setdefault(var.id, [])
                 store_var_info[var.id].append(name)
                 self.pycode_gen.gen_store_fast(name)
             else:
@@ -398,8 +399,9 @@ def _build_compile_fn_with_name_store(
         name_gen = NameGenerator("___graph_fn_saved_")
 
         for var in to_store_vars[::-1]:
-            if not store_var_info[var.id]:
+            if not store_var_info.get(var.id, []):
                 name = name_gen.next()
+                store_var_info.setdefault(var.id, [])
                 store_var_info[var.id].append(name)
                 self.pycode_gen.gen_store_fast(name)
             else:
 
@@ -770,22 +770,6 @@ def BINARY_SLICE(self, instr: Instruction):
 
     def binary_subscr_operation(self, key, container, opname):
         assert isinstance(key, VariableBase)
-        # TODO(xiongkun): getitem / getattr support key and attr as variable.
-        if isinstance(key, TensorVariable) and isinstance(
-            container, TensorVariable
-        ):
-            # NOTE(xiongkun): tensor[tensor] should support.
-            output = self._graph.call_tensor_method(
-                "__getitem__", container, key
-            )
-            self.stack.push(output)
-            return
-
-        if isinstance(key, TensorVariable):
-            raise BreakGraphError(
-                f"Key is a TensorVariable in {opname}, {container}[{key}]"
-            )
-
         result = BuiltinVariable(
             operator.getitem, self._graph, DanglingTracker()
         )(container, key)
@@ -989,20 +973,15 @@ def STORE_SLICE(self, instr: Instruction):
 
     def store_subscr_operation(self, key, container, value, opname):
         assert isinstance(key, VariableBase)
-        self._graph.add_global_guarded_variable(key)
-        if isinstance(key, TensorVariable):
-            raise BreakGraphError(
-                f"Key is a TensorVariable in {opname}, {container}[{key}] = {value}"
-            )
-        # TODO(xiongkun): support tensor[tensor] = tensor, dy2static is not the same with dygraph.
-        container[key.get_py_value()] = value
+        BuiltinVariable(operator.setitem, self._graph, DanglingTracker())(
+            container, key, value
+        )
         value.debug_name = f"{container.debug_name}[{key.debug_name}]"
 
     def DELETE_SUBSCR(self, instr: Instruction):
         key = self.stack.pop()
         container = self.stack.pop()
         assert isinstance(key, VariableBase)
-        self._graph.add_global_guarded_variable(key)
         BuiltinVariable(operator.delitem, self._graph, DanglingTracker())(
             container, key
         )
@@ -2079,7 +2058,7 @@ def get_compute_fn_and_update_changed_vars(
             _var = self.get_var(name, allow_undefined=True)
             if _var is SotUndefinedVar():
                 continue
-            if _var not in stack:
+            if _var not in store_vars:
                 store_vars.append(_var)
             store_var_info.setdefault(_var.id, [])
             store_var_info[_var.id].append(name)
@@ -2499,9 +2478,12 @@ def create_after_loop_fn():
             OrderedSet(loop_body_inputs[:-1]) | after_loop_read_names
         )
         extra_store_vars = (
-            [iterator]
+            [
+                item
+                for item in iterator.flatten_items()
+                if isinstance(item, (TensorVariable, SymbolicVariable))
+            ]
             if isinstance(iterator, IterVariable)
-            and isinstance(iterator.hold, TensorVariable)
             else []
         )
         var_loader = self.get_compute_fn_and_update_changed_vars(
 
@@ -572,9 +572,11 @@ def dispatch_list_ne(lhs: ListVariable, rhs: ListVariable):
 # stop
 Dispatcher.register(
     range,
-    ("ConstantVariable",),
+    ("ConstantVariable | TensorVariable",),
     lambda stop: RangeVariable(
-        range(stop.get_py_value()),
+        ConstantVariable.wrap_literal(0, stop.graph),
+        stop,
+        ConstantVariable.wrap_literal(1, stop.graph),
         graph=stop.graph,
         tracker=DummyTracker([stop]),
     ),
@@ -583,19 +585,27 @@ def dispatch_list_ne(lhs: ListVariable, rhs: ListVariable):
 # start, stop
 Dispatcher.register(
     range,
-    ("ConstantVariable", "ConstantVariable"),
+    ("ConstantVariable | TensorVariable", "ConstantVariable | TensorVariable"),
     lambda start, stop: RangeVariable(
-        range(start.get_py_value(), stop.get_py_value()),
+        start,
+        stop,
+        ConstantVariable.wrap_literal(1, stop.graph),
         graph=stop.graph,
         tracker=DummyTracker([start, stop]),
     ),
 )
 # start, stop, step
 Dispatcher.register(
     range,
-    ("ConstantVariable", "ConstantVariable", "ConstantVariable"),
+    (
+        "ConstantVariable | TensorVariable",
+        "ConstantVariable | TensorVariable",
+        "ConstantVariable | TensorVariable",
+    ),
     lambda start, stop, step: RangeVariable(
-        range(start.get_py_value(), stop.get_py_value(), step.get_py_value()),
+        start,
+        stop,
+        step,
         graph=stop.graph,
         tracker=DummyTracker([start, stop, step]),
     ),
@@ -792,14 +802,29 @@ def str_endswith(
 )
 
 # setitem
+Dispatcher.register(
+    operator.setitem,
+    (
+        "TensorVariable",
+        "Any",
+        "VariableBase",
+    ),
+    lambda var, key, value: var.setitem(
+        VariableFactory.from_value(
+            key, graph=var.graph, tracker=ConstTracker(key)
+        ),
+        value,
+    ),
+)
+
 Dispatcher.register(
     operator.setitem,
     (
         "VariableBase",
         "int | str | ConstantVariable | TensorVariable",
         "int | str | ConstantVariable | TensorVariable",
     ),
-    lambda var, key, value: var.setitem(key.get_py_value(), value),
+    lambda var, key, value: var.setitem(add_guard(key).get_py_value(), value),
 )
 
 # delitem
@@ -817,7 +842,7 @@ def str_endswith(
         "VariableBase",
         "ConstantVariable",
     ),
-    lambda var, key: var.delitem(key.get_py_value()),
+    lambda var, key: var.delitem(add_guard(key).get_py_value()),
 )
 
 
 
@@ -430,8 +430,9 @@ def flatten_items(self) -> list[VariableBase]:
             list[VariableBase]: Flattened items of a container variable.
         """
         from .container import ContainerVariable
+        from .iter import IterVariable
 
-        if not isinstance(self, ContainerVariable):
+        if not isinstance(self, (ContainerVariable, IterVariable)):
             return [self]
         flattened_items = []
         for item in self.get_items():
 
@@ -527,15 +527,10 @@ def getitem(self, key):
         return self.graph.call_tensor_method("__getitem__", self, key)
 
     def setitem(self, key, value):
-        self.graph.add_global_guarded_variable(value)
-
-        key_var = VariableFactory.from_value(
-            key, self.graph, tracker=ConstTracker(key)
-        )
         new_tensor = self.graph.call_paddle_api(
             paddle.static.setitem,
             self,
-            key_var,
+            key,
             value,
         )
 
@@ -807,6 +802,7 @@ def disable_symbolic(var: VariableBase):
 
         disable_symbolic(self)
         self.graph.need_cache = False
+        log(3, f"Fallback {self} to ConstantVariable")
         return ConstantVariable(
             self.get_py_value(), self.graph, DummyTracker([self])
         )
@@ -815,6 +811,10 @@ def get_py_value(self, allow_tensor: bool = False) -> bool | int | float:
         if ENV_SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE.get():
             raise BreakGraphError("get_py_value from SymbolicVariable")
         self.need_guard_value = True
+        log(
+            3,
+            f"get_py_value from SymbolicVariable {self} caused value need guard",
+        )
         if isinstance(self.value, SymbolicValue):
             assert isinstance(
                 self.tracker, SymbolicOperationTracker