[SOT] fix bug in llm stable diffusion (#62257)

feifei-111 · web-flow · commit 98f48ba29477 · 2024-03-02T10:16:32.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -88,6 +88,7 @@
     TensorVariable,
     TupleVariable,
     UserDefinedFunctionVariable,
+    UserDefinedGeneratorFunctionVariable,
     VariableBase,
     VariableFactory,
 )
@@ -1318,11 +1319,21 @@ def g(z=x):
             default_args,
             closure,
         )
-        self.stack.push(
-            UserDefinedFunctionVariable(
-                new_fn, self._graph, DummyTracker(related_list)
+        # new_fn is created for which is binded with Variables
+        # so new_fn.__module__ is a ConstantVariable
+        # can not use VariableFactory.from_value
+        if inspect.isgeneratorfunction(new_fn):
+            self.stack.push(
+                UserDefinedGeneratorFunctionVariable(
+                    new_fn, self._graph, DummyTracker(related_list)
+                )
+            )
+        else:
+            self.stack.push(
+                UserDefinedFunctionVariable(
+                    new_fn, self._graph, DummyTracker(related_list)
+                )
             )
-        )
 
     def GET_ITER(self, instr: Instruction):
         source_obj = self.stack.pop()
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py b/python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py
@@ -44,7 +44,7 @@
     PaddleApiVariable,
     PaddleLayerVariable,
     UserDefinedFunctionVariable,
-    UserDefinedGeneratorVariable,
+    UserDefinedGeneratorFunctionVariable,
     UserDefinedLayerVariable,
 )
 from .container import (  # noqa: F401
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -681,9 +681,9 @@ def main_info(self) -> dict[str, Any]:
         }
 
 
-class UserDefinedGeneratorVariable(FunctionVariable):
+class UserDefinedGeneratorFunctionVariable(FunctionVariable):
     """
-    UserDefinedGeneratorVariable is a subclass of FunctionVariable used to wrap a user-defined generator.
+    UserDefinedGeneratorFunctionVariable is a subclass of FunctionVariable used to wrap a user-defined generator.
     Args:
         fn (Callable[..., Any]): The user-defined generator to be wrapped.
         graph(FunctionGraph): The FunctionGraph object that this variable is associated with.
@@ -711,7 +711,7 @@ def main_info(self) -> dict[str, Any]:
     )
     def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
         if inspect.isgeneratorfunction(value):
-            return UserDefinedGeneratorVariable(value, graph, tracker)
+            return UserDefinedGeneratorFunctionVariable(value, graph, tracker)
         return None
 
 
diff --git a/python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_analysis.py b/python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_analysis.py
@@ -23,21 +23,19 @@
     ALL_JUMP,
     HAS_FREE,
     HAS_LOCAL,
-    RETURN,
     UNCONDITIONAL_JUMP,
 )
 
 
 @dataclasses.dataclass
-class State:
+class NameRecorder:
     reads: OrderedSet[str]
     writes: OrderedSet[str]
-    visited: OrderedSet[int]
 
     def __or__(self, other):
         reads = self.reads | other.reads
         writes = self.writes | other.writes
-        return State(reads, writes, OrderedSet())
+        return NameRecorder(reads, writes)
 
 
 def is_read_opcode(opname):
@@ -90,46 +88,70 @@ def analysis_used_names(
     Returns:
         State: The analysis result.
     """
-    root_state = State(OrderedSet(), OrderedSet(), OrderedSet())
-
-    def fork(state: State, start: int, jump: bool, jump_target: int) -> State:
+    name_recorder = NameRecorder(OrderedSet(), OrderedSet())
+
+    # start idx and writes names can decide the analysis result below
+    # so, just check the pair of (idx, writes), to skip repeat simulation
+    # (writes can decide if a name should be add to reads)
+    # one idx can has multi writes for whom is not subset with each other
+    # if A is subset of B, we just record A, simulate A might add more reads
+    visited_states = {}
+
+    def check_and_update_visited_states(idx, writes):
+        writes = set(writes)
+
+        if idx in visited_states:
+            history = visited_states[idx]
+            for record in history:
+                if record.issubset(writes):
+                    return True
+                elif writes.issubset(record):
+                    history.remove(record)
+                    history.append(writes)
+                    return False
+        else:
+            visited_states[idx] = [writes]
+
+        return False
+
+    def fork(
+        name_recorder: NameRecorder, start: int, jump: bool, jump_target: int
+    ) -> NameRecorder:
         new_start = start + 1 if not jump else jump_target
-        new_state = State(
-            OrderedSet(state.reads),
-            OrderedSet(state.writes),
-            OrderedSet(state.visited),
+        new_state = NameRecorder(
+            OrderedSet(name_recorder.reads),
+            OrderedSet(name_recorder.writes),
         )
         return walk(new_state, new_start)
 
-    def walk(state: State, start: int) -> State:
+    def walk(name_recorder: NameRecorder, start: int) -> NameRecorder:
         end = len(instructions) if stop_instr_idx is None else stop_instr_idx
         for i in range(start, end):
-            if i in state.visited:
-                return state
-            state.visited.add(i)
+            if check_and_update_visited_states(i, name_recorder.writes):
+                return name_recorder
 
             instr = instructions[i]
             if instr.opname in HAS_LOCAL | HAS_FREE:
                 if is_read_opcode(instr.opname) and instr.argval not in (
-                    state.writes
+                    name_recorder.writes
                 ):
-                    state.reads.add(instr.argval)
+                    name_recorder.reads.add(instr.argval)
                 elif is_write_opcode(instr.opname):
-                    state.writes.add(instr.argval)
+                    name_recorder.writes.add(instr.argval)
             elif instr.opname in ALL_JUMP:
                 assert instr.jump_to is not None
                 target_idx = instructions.index(instr.jump_to)
                 # Fork to two branches, jump or not
-                jump_branch = fork(state, i, True, target_idx)
+                jump_branch = fork(name_recorder, i, True, target_idx)
                 not_jump_branch = (
-                    fork(state, i, False, target_idx)
+                    fork(name_recorder, i, False, target_idx)
                     if instr.opname not in UNCONDITIONAL_JUMP
-                    else State(OrderedSet(), OrderedSet(), OrderedSet())
+                    else NameRecorder(OrderedSet(), OrderedSet())
                 )
                 return jump_branch | not_jump_branch
-            elif instr.opname in RETURN:
-                return state
-        return state
+            elif instr.opname == "RETURN_VALUE":
+                return name_recorder
+        return name_recorder
 
-    state = walk(root_state, current_instr_idx)
-    return state.reads, state.writes
+    name_recorder = walk(name_recorder, current_instr_idx)
+    return name_recorder.reads, name_recorder.writes
diff --git a/python/paddle/jit/sot/utils/paddle_api_config.py b/python/paddle/jit/sot/utils/paddle_api_config.py
@@ -82,7 +82,6 @@ def get_paddle_api():
 # considered as paddle module？
 paddle_api_module_prefix = {
     "paddle.nn.functional",
-    "paddle.nn.layer.activation",
 }
 
 break_graph_set = set()
diff --git a/test/sot/test_break_graph.py b/test/sot/test_break_graph.py
@@ -185,5 +185,20 @@ def test_break_graph_in_layer(self):
         self.assert_results(net.forward, x)
 
 
+def dummy(*args):
+    return None
+
+
+def break_graph_call_generator_function(x):
+    return dummy(y for y in x)
+
+
+class TestBreakGraphCallGeneratorFunction(TestCaseBase):
+    def test_break_graph_when_call_generator_function(self):
+        x = paddle.rand([1], dtype=paddle.float32)
+        y = paddle.rand([1], dtype=paddle.float32)
+        self.assert_results(break_graph_call_generator_function, [x, y])
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@`
`44`	`44`	`PaddleApiVariable,`
`45`	`45`	`PaddleLayerVariable,`
`46`	`46`	`UserDefinedFunctionVariable,`
`47`		`- UserDefinedGeneratorVariable,`
	`47`	`+ UserDefinedGeneratorFunctionVariable,`
`48`	`48`	`UserDefinedLayerVariable,`
`49`	`49`	`)`
`50`	`50`	`from .container import ( # noqa: F401`
Original file line number	Diff line number	Diff line change
`@@ -82,7 +82,6 @@ def get_paddle_api():`
`82`	`82`	`# considered as paddle module？`
`83`	`83`	`paddle_api_module_prefix = {`
`84`	`84`	`"paddle.nn.functional",`
`85`		`- "paddle.nn.layer.activation",`
`86`	`85`	`}`
`87`	`86`
`88`	`87`	`break_graph_set = set()`