Fix line numbers bug (#87247)

Fixes pytorch/torchdynamo#1462

Pull Request resolved: #87247
Approved by: https://github.com/anijain2305, https://github.com/jansel

torch/_dynamo/bytecode_analysis.py

@@ -59,6 +59,39 @@ def remove_pointless_jumps(instructions):
     return [inst for inst in instructions if id(inst) not in pointless_jumps]
 
 
+def propagate_line_nums(instructions):
+    """Ensure every instruction has line number set in case some are removed"""
+    cur_line_no = None
+
+    def populate_line_num(inst):
+        nonlocal cur_line_no
+        if inst.starts_line:
+            cur_line_no = inst.starts_line
+
+        inst.starts_line = cur_line_no
+
+    for inst in instructions:
+        populate_line_num(inst)
+
+
+def remove_extra_line_nums(instructions):
+    """Remove extra starts line properties before packing bytecode"""
+
+    cur_line_no = None
+
+    def remove_line_num(inst):
+        nonlocal cur_line_no
+        if inst.starts_line is None:
+            return
+        elif inst.starts_line == cur_line_no:
+            inst.starts_line = None
+        else:
+            cur_line_no = inst.starts_line
+
+    for inst in instructions:
+        remove_line_num(inst)
+
+
 @dataclasses.dataclass
 class ReadsWrites:
     reads: set

torch/_dynamo/bytecode_transformation.py

@@ -5,7 +5,11 @@
 import types
 from typing import Any, List, Optional
 
-from .bytecode_analysis import stacksize_analysis
+from .bytecode_analysis import (
+    propagate_line_nums,
+    remove_extra_line_nums,
+    stacksize_analysis,
+)
 
 
 @dataclasses.dataclass
@@ -332,6 +336,7 @@ def transform_code_object(code, transformations, safe=False):
     assert len(code_options["co_varnames"]) == code_options["co_nlocals"]
 
     instructions = cleaned_instructions(code, safe)
+    propagate_line_nums(instructions)
 
     transformations(instructions, code_options)
 
@@ -344,6 +349,7 @@ def transform_code_object(code, transformations, safe=False):
         # this pass might change offsets, if so we need to try again
         dirty = fix_extended_args(instructions)
 
+    remove_extra_line_nums(instructions)
     bytecode, lnotab = assemble(instructions, code_options["co_firstlineno"])
     if sys.version_info < (3, 10):
         code_options["co_lnotab"] = lnotab

torch/_dynamo/symbolic_convert.py

@@ -307,7 +307,7 @@ def step(self):
         else:
             self.instruction_pointer = None
             self.next_instruction = None
-        if inst.starts_line:
+        if inst.starts_line and self.lineno != inst.starts_line:
             self.lineno = inst.starts_line
             log.debug(f"TRACE starts_line {self.f_code.co_filename}:{self.lineno}")