[FIX] [SANITIZER] Report correct line numbers for OOB errors in loops (#251)

Author: mark14wu

tests/end_to_end/test_sanitizer.py

@@ -131,9 +131,12 @@ def _check_range_satisfiable(
         access_addr: Z3Expr,
         expr_constraints: Optional[BoolRef],
         symbolic_expr: SymbolicExpr,
+        source_location: Optional[tuple[str, int, str]] = None,
     ) -> None:
         self.check_inside_loop.append((access_addr, expr_constraints))
-        super()._check_range_satisfiable(access_addr, expr_constraints, symbolic_expr)
+        super()._check_range_satisfiable(
+            access_addr, expr_constraints, symbolic_expr, source_location
+        )
 
     def register_for_loop_callback(self) -> ForLoopCallbacks:
         callbacks = super().register_for_loop_callback()
@@ -296,3 +299,64 @@ def test_loop_deferred_checks_simplify():
     loop_deferred_check_simplify_kernel[(2,)](out)
 
     assert load_index_checker.observed_offsets == []
+
+
+# ======== Line Number Reporting Tests ===========
+
+
+# Create a dedicated sanitizer for line number tests
+line_number_checker: SymbolicSanitizer = SymbolicSanitizer(abort_on_error=False)
+
+
+@triton_viz.trace(client=line_number_checker)
+@triton.jit
+def oob_in_loop_kernel(ptr, N: tl.constexpr, BLOCK_SIZE: tl.constexpr):
+    """Kernel where OOB occurs inside a loop at the tl.load line."""
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    for _ in range(2):
+        # This line should be reported in the traceback when OOB occurs
+        val = tl.load(ptr + offsets + 1000)  # OOB access due to +1000 offset
+        tl.store(ptr + offsets, val)
+
+
+def test_loop_oob_reports_correct_line_number():
+    """
+    Test that sanitizer reports the correct line number for OOB errors in loops.
+
+    Previously, the sanitizer would report the function definition line instead
+    of the actual tl.load/tl.store line that caused the error. This was because
+    traceback was captured at loop exit time rather than when the memory
+    operation was executed.
+    """
+    line_number_checker.records.clear()
+
+    data = torch.zeros((16,), dtype=torch.float32)
+
+    oob_in_loop_kernel[(1,)](data, N=16, BLOCK_SIZE=16)
+
+    # Verify that OOB was detected
+    assert len(line_number_checker.records) > 0, "Expected OOB to be detected"
+
+    record = line_number_checker.records[0]
+    assert len(record.user_code_tracebacks) > 0, "Expected traceback info"
+
+    tb_info = record.user_code_tracebacks[0]
+
+    # The error should point to the tl.load line, not the function definition
+    assert "tl.load" in tb_info.line_of_code, (
+        f"Expected traceback to point to tl.load line, "
+        f"but got: {tb_info.line_of_code!r}"
+    )
+
+    # Verify the line contains the OOB offset
+    assert "+1000" in tb_info.line_of_code or "1000" in tb_info.line_of_code, (
+        f"Expected line to contain the OOB offset, "
+        f"but got: {tb_info.line_of_code!r}"
+    )
+
+    # Verify function name
+    assert tb_info.func_name == "oob_in_loop_kernel", (
+        f"Expected func_name to be 'oob_in_loop_kernel', "
+        f"but got: {tb_info.func_name!r}"
+    )

triton_viz/clients/sanitizer/report.py

@@ -1,5 +1,8 @@
+import linecache
+import sys
 import traceback
-from typing import Optional, TYPE_CHECKING
+from pathlib import Path
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -15,6 +18,58 @@
     from .sanitizer import SymbolicExpr
 
 
+# Paths that identify framework code (not user code)
+_FRAMEWORK_PATHS = [
+    "triton_viz/core/",
+    "triton_viz/clients/",
+    "triton/runtime/",
+    "triton/language/",
+    "site-packages/triton/",
+]
+
+
+def _get_user_code_location() -> tuple[str, int, str] | None:
+    from types import FrameType
+
+    frame: FrameType | None = sys._getframe()
+
+    while frame is not None:
+        filename = Path(frame.f_code.co_filename).as_posix()
+
+        # Skip Python internals
+        if filename.startswith("<"):
+            frame = frame.f_back
+            continue
+
+        # Check if this is user code (not in framework paths)
+        is_framework = any(path in filename for path in _FRAMEWORK_PATHS)
+
+        # User code: not in framework paths, OR in examples directory
+        if not is_framework or "examples/" in filename:
+            return (
+                frame.f_code.co_filename,
+                frame.f_lineno,
+                frame.f_code.co_name,
+            )
+
+        frame = frame.f_back
+
+    return None
+
+
+def _location_to_traceback_info(
+    source_location: tuple[str, int, str],
+) -> TracebackInfo:
+    filename, lineno, func_name = source_location
+    line_of_code = linecache.getline(filename, lineno).rstrip()
+    return TracebackInfo(
+        filename=filename,
+        lineno=lineno,
+        func_name=func_name,
+        line_of_code=line_of_code,
+    )
+
+
 def print_oob_record(oob_record: OutOfBoundsRecord, max_display=10):
     """
     Print detailed logs for a given OOB record.
@@ -95,7 +150,7 @@ def print_oob_record(oob_record: OutOfBoundsRecord, max_display=10):
 
 
 def print_oob_record_pdb_style(
-    oob_record: OutOfBoundsRecord, symbolic_expr: Optional["SymbolicExpr"] = None
+    oob_record: OutOfBoundsRecord, symbolic_expr: "SymbolicExpr | None" = None
 ):
     """
     Print a comprehensive diagnostic report for OOB errors in PDB-style format.
@@ -104,7 +159,7 @@ def print_oob_record_pdb_style(
     ----------
     oob_record : OutOfBoundsRecord
         The record containing information about out-of-bounds accesses.
-    symbolic_expr : Optional[SymbolicExpr]
+    symbolic_expr : SymbolicExpr | None
         The symbolic expression tree that led to the OOB access.
     """
     from pathlib import Path

triton_viz/clients/sanitizer/sanitizer.py

@@ -84,7 +84,13 @@
     check_inner_stride_equal_to_one,
 )
 from .data import OutOfBoundsRecordZ3
-from .report import _get_traceback_info, print_oob_record, print_oob_record_pdb_style
+from .report import (
+    _get_traceback_info,
+    _get_user_code_location,
+    _location_to_traceback_info,
+    print_oob_record,
+    print_oob_record_pdb_style,
+)
 from ...core.config import config as cfg
 
 
@@ -163,6 +169,9 @@ class PendingCheck:
     symbolic_expr: "SymbolicExpr"
     addr_expr: Z3Expr
     constraints: ConstraintConjunction
+    # Lightweight source location: (filename, lineno, func_name)
+    # Captured immediately to preserve accurate line info for deferred checks
+    source_location: Optional[tuple[str, int, str]] = None
 
 
 @dataclass
@@ -1707,6 +1716,7 @@ def _check_range_satisfiable(
         access_addr: Z3Expr,
         expr_constraints: ConstraintConjunction,
         symbolic_expr: SymbolicExpr,
+        source_location: Optional[tuple[str, int, str]] = None,
     ) -> None:
         # Use push/pop on persistent solver
         solver = self.solver
@@ -1739,8 +1749,10 @@ def _check_single_addr(addr_expr: Z3Expr) -> None:
                 else:
                     op_type = Load
 
-                # Report with symbolic expression
-                self._report(op_type, tensor, violation_addr, symbolic_expr)
+                # Report with symbolic expression and source location
+                self._report(
+                    op_type, tensor, violation_addr, symbolic_expr, source_location
+                )
             solver.pop()
 
         if isinstance(access_addr, list):
@@ -1767,11 +1779,16 @@ def _handle_access_check(self, expr: SymbolicExpr) -> None:
         signature = _make_signature(z3_addr, z3_constraints)
         pending_idx = ctx.signature_cache.get(signature)
         if pending_idx is None:
+            # Capture source location now while we're still in the user's tl.load/tl.store call.
+            # This is a lightweight operation that only traverses frame objects.
+            # The actual source line will be read later only if an error is detected.
+            source_location = _get_user_code_location()
             ctx.signature_cache[signature] = len(ctx.pending_checks)
             pending_check = PendingCheck(
                 symbolic_expr=expr,
                 addr_expr=z3_addr,
                 constraints=z3_constraints,
+                source_location=source_location,
             )
             ctx.pending_checks.append(pending_check)
         else:
@@ -1784,8 +1801,15 @@ def _report(
         tensor: Tensor,
         violation_address: int,
         symbolic_expr: Optional[SymbolicExpr] = None,
+        source_location: Optional[tuple[str, int, str]] = None,
     ) -> None:
-        traceback_info = _get_traceback_info()
+        # Use pre-captured location if available (for deferred checks in loops),
+        # otherwise capture it now (for immediate checks outside loops)
+        if source_location is not None:
+            traceback_info = [_location_to_traceback_info(source_location)]
+        else:
+            traceback_info = _get_traceback_info()
+
         tensor_name = self._get_tensor_name(tensor)
         oob_record = OutOfBoundsRecordZ3(
             op_type=op_type,
@@ -2285,7 +2309,7 @@ def loop_hook_after(lineno: int) -> None:
             if not self.loop_stack or self.loop_stack[-1].lineno != lineno:
                 return
             ctx = self.loop_stack.pop()
-            # execute pending checks
+            # Execute pending checks that were deferred during loop execution
             solver = self.solver
             addr_sym = self.addr_sym
             assert solver is not None
@@ -2302,6 +2326,9 @@ def loop_hook_after(lineno: int) -> None:
                 addr_expr = pending_check.addr_expr
                 expr_constraints = pending_check.constraints
                 symbolic_expr = pending_check.symbolic_expr
+                # Use the source location captured when the check was created,
+                # not the current location (which would be the loop exit point)
+                source_location = pending_check.source_location
 
                 if cfg.verbose:
                     print(
@@ -2315,6 +2342,7 @@ def loop_hook_after(lineno: int) -> None:
                     addr_expr,
                     expr_constraints,
                     symbolic_expr,
+                    source_location,
                 )
             if ctx.pending_checks:
                 solver.pop()