diff --git a/depyf/decompiler.py b/depyf/decompiler.py index 906d463a..a13ce81e 100644 --- a/depyf/decompiler.py +++ b/depyf/decompiler.py @@ -300,7 +300,11 @@ def STORE_ATTR(self, inst: Instruction): def DELETE_SUBSCR(self, inst: Instruction): index = self.state.stack.pop() x = self.state.stack.pop() - self.state.source_code += f"del {x}[{index}]\n" + if not f"{x}[{index}]" in self.state.stack: + # if f"{x}[{index}]" already exists somewhere in the stack, + # we don't need to delete it, otherwise that reference will be invalidated. + # TODO: this is a hack, we should use a more robust way to track the references. + self.state.source_code += f"del {x}[{index}]\n" def generic_delete(self, inst: Instruction): self.state.source_code += f"del {inst.argval}\n" @@ -1042,9 +1046,8 @@ def LIST_TO_TUPLE(self, inst: Instruction): self.state.stack.append(f"tuple({item})") def LIST_EXTEND(self, inst: Instruction): - assert inst.argval == 1, "Only tested for argval==1" values = self.state.stack.pop() - temp = self.replace_mutable_tos_with_temp() + temp = self.replace_mutable_tos_with_temp(pos=inst.argval) self.state.source_code += f"{temp}.extend({values})\n" def LIST_APPEND(self, inst: Instruction): @@ -1150,11 +1153,12 @@ def get_temp_name(self): Decompiler.temp_count += 1 return f"{self.temp_prefix}{Decompiler.temp_count}" - def replace_mutable_tos_with_temp(self): - ans = self.state.stack.pop() + def replace_mutable_tos_with_temp(self, pos: int = 1): + # pos is basically argval of the instruction + ans = self.state.stack[-pos] temp_name = self.get_temp_name() self.state.source_code += f"{temp_name} = {ans}\n" - self.state.stack.append(temp_name) + self.state.stack[-pos] = temp_name return temp_name @staticmethod @@ -1181,8 +1185,13 @@ def decompile( source_code = self.state.source_code # the header might have invalid function name in torchdynamo. only # optimize the function body. - source_code = remove_some_temp( - source_code, self.temp_prefix, indentation) + if os.environ.get("DEPYF_REMOVE_TEMP", "1") == "1": + # remove temp can be dangerous if some code in between contains + # mutation, e.g. as described in https://github.com/thuml/depyf/issues/85 . + # doing a full mutation analysis in python is too complex, so we remove temp by default, + # and only disable it when users set the environment variable DEPYF_REMOVE_TEMP to 0. + source_code = remove_some_temp( + source_code, self.temp_prefix, indentation) header = get_function_signature(self.code, overwite_fn_name) # we cannot rely on `co_names`. For example, `from math import sqrt` will make `math` and `sqrt` in `co_names`. global_names = set(inst.argval for inst in dis.get_instructions(self.code) if inst.opname == "STORE_GLOBAL") diff --git a/tests/test_pytorch/test_pytorch.py b/tests/test_pytorch/test_pytorch.py index 96773a6f..61141e85 100644 --- a/tests/test_pytorch/test_pytorch.py +++ b/tests/test_pytorch/test_pytorch.py @@ -2,8 +2,7 @@ def toy_function(a, b): x = a / (torch.abs(a) + 1) - if b.sum() < 0: - b = b * -1 + b = b * -1 + b.sum() return x * b class ToyModule(torch.nn.Module): @@ -12,8 +11,7 @@ def __init__(self): def forward(self, a, b): x = a / (torch.abs(a) + 1) - if b.sum() < 0: - b = b * -1 + b = b * -1 + b.sum() return x * b