refactor

Author: makslevental

mlir/lib/Bindings/Python/Traceback.cpp

@@ -440,22 +440,22 @@ void BuildTracebackSubmodule(nb::module_ &m) {
       },
       "Python wrapper around the Python C API function PyCode_Addr2Line");
 
+#if PY_VERSION_HEX >= 0x030b00f0
   type.attr("code_addr2location") = nb::cpp_function(
       [](nb::handle code, int lasti) {
         if (!PyCode_Check(code.ptr())) {
           throw std::runtime_error("code argument must be a code object");
         }
         int start_line, start_column, end_line, end_column;
-        // if (!PyCode_Addr2Location(reinterpret_cast<PyCodeObject
-        // *>(code.ptr()),
-        //                           lasti, &start_line, &start_column,
-        //                           &end_line, &end_column)) {
-        //   throw nb::python_error();
-        // }
-        throw nb::python_error();
+        if (!PyCode_Addr2Location(reinterpret_cast<PyCodeObject *>(code.ptr()),
+                                  lasti, &start_line, &start_column, &end_line,
+                                  &end_column)) {
+          throw nb::python_error();
+        }
         return nb::make_tuple(start_line, start_column, end_line, end_column);
       },
       "Python wrapper around the Python C API function PyCode_Addr2Location");
+#endif
 }
 } // namespace mlir::python
 

mlir/python/mlir/source_info_util.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import sys
 from collections.abc import Iterator
 import contextlib
 import dataclasses
@@ -25,24 +26,30 @@
 import threading
 import types
 from typing import NamedTuple
+from .ir import Location
 
 # import jax.version
 # from jax._src.lib import xla_client
 
 from . import traceback_util
+from .traceback_util import (
+    TracebackCaches,
+    Traceback,
+    _traceback_caches,
+    _traceback_in_locations_limit,
+    _include_full_tracebacks_in_locations,
+)
 
 traceback_util.register_exclusion(__file__)
 
-from ._mlir_libs._mlir import Traceback
-
 
 class Frame(NamedTuple):
     file_name: str
     function_name: str
     start_line: int
-    start_column: int
-    end_line: int
-    end_column: int
+    start_column: int | None
+    end_line: int | None
+    end_column: int | None
 
 
 _exclude_paths: list[str] = [
@@ -173,16 +180,29 @@ def is_user_filename(filename: str) -> bool:
 
 
 def raw_frame_to_frame(code: types.CodeType, lasti: int) -> Frame:
-    loc = Traceback.code_addr2location(code, lasti)
-    start_line, start_column, end_line, end_column = loc
-    return Frame(
-        file_name=code.co_filename,
-        function_name=code.co_qualname,
-        start_line=start_line,
-        start_column=start_column,
-        end_line=end_line,
-        end_column=end_column,
-    )
+    if sys.version_info.minor >= 11:
+        loc = Traceback.code_addr2location(code, lasti)
+        start_line, start_column, end_line, end_column = loc
+        frame = Frame(
+            file_name=code.co_filename,
+            function_name=code.co_qualname,
+            start_line=start_line,
+            start_column=start_column,
+            end_line=end_line,
+            end_column=end_column,
+        )
+    else:
+        start_line = Traceback.code_addr2line(code, lasti)
+        frame = Frame(
+            file_name=code.co_filename,
+            function_name=code.co_qualname,
+            start_line=start_line,
+            start_column=None,
+            end_line=None,
+            end_column=None,
+        )
+
+    return frame
 
 
 def user_frames(traceback: Traceback | None) -> Iterator[Frame]:
@@ -365,3 +385,100 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 
 transform_name_stack = TransformNameStackContextManager
+
+
+def get_canonical_source_file(file_name: str, caches: TracebackCaches) -> str:
+    canonical_file_name = caches.canonical_name_cache.get(file_name, None)
+    if canonical_file_name is not None:
+        return canonical_file_name
+
+    # pattern = config.hlo_source_file_canonicalization_regex.value
+    # if pattern:
+    #     file_name = re.sub(pattern, "", file_name)
+    caches.canonical_name_cache[file_name] = file_name
+    return file_name
+
+
+def _is_user_file(file_name: str) -> bool:
+    is_user = _traceback_caches.is_user_file_cache.get(file_name, None)
+    if is_user is not None:
+        return is_user
+    out = is_user_filename(file_name)
+    _traceback_caches.is_user_file_cache[file_name] = out
+    return out
+
+
+def _traceback_to_location(tb: Traceback) -> Location:
+    """Converts a full traceback to a callsite() MLIR location."""
+    loc = _traceback_caches.traceback_cache.get(tb, None)
+    if loc is not None:
+        return loc
+
+    frame_locs = []
+    frames_limit = _traceback_in_locations_limit
+    frames_limit = frames_limit if frames_limit >= 0 else 1000
+
+    codes, lastis = tb.raw_frames()
+    for i, code in enumerate(codes):
+        if not _is_user_file(code.co_filename):
+            continue
+
+        lasti = lastis[i]
+        code_lasti = code, lasti
+        loc = _traceback_caches.location_cache.get(code_lasti, None)
+        if loc is None:
+            frame = raw_frame_to_frame(code, lasti)
+            file_loc = Location.file(
+                get_canonical_source_file(frame.file_name, _traceback_caches),
+                frame.start_line,
+                frame.start_column,
+                frame.end_line,
+                frame.end_column,
+            )
+            loc = Location.name(frame.function_name, childLoc=file_loc)
+            _traceback_caches.location_cache[code_lasti] = loc
+        frame_locs.append(loc)
+        if len(frame_locs) >= frames_limit:
+            break
+
+    n = len(frame_locs)
+    if n == 0:
+        loc = Location.unknown()
+    elif n == 1:
+        loc = frame_locs[0]
+    else:
+        loc = Location.callsite(frame_locs[0], frame_locs[1:])
+    _traceback_caches.traceback_cache[tb] = loc
+    return loc
+
+
+def source_info_to_location(
+    primitive: None,
+    name_stack: NameStack,
+    traceback: Traceback | None,
+) -> Location:
+    if _include_full_tracebacks_in_locations:
+        if traceback is None:
+            loc = Location.unknown()
+        else:
+            loc = _traceback_to_location(traceback)
+    else:
+        frame = user_frame(traceback)
+        if frame is None:
+            loc = Location.unknown()
+        else:
+            loc = Location.file(
+                get_canonical_source_file(frame.file_name, _traceback_caches),
+                frame.start_line,
+                frame.start_column,
+            )
+    if primitive is None:
+        if name_stack.stack:
+            loc = Location.name(str(name_stack), childLoc=loc)
+    else:
+        eqn_str = (
+            f"{name_stack}/{primitive.name}" if name_stack.stack else primitive.name
+        )
+        loc = Location.name(eqn_str, childLoc=loc)
+        loc = Location.name(f"{primitive.name}:", childLoc=loc)
+    return loc

mlir/python/mlir/traceback_util.py

@@ -14,12 +14,15 @@
 
 from __future__ import annotations
 
+import dataclasses
 from collections.abc import Callable
 import functools
 import os
 import traceback
 import types
 from typing import Any, TypeVar, cast
+from ._mlir_libs._mlir import Traceback
+from .ir import Location
 
 
 C = TypeVar("C", bound=Callable[..., Any])
@@ -236,3 +239,22 @@ def reraise_with_filtered_traceback(*args, **kwargs):
                 del mode
 
     return cast(C, reraise_with_filtered_traceback)
+
+
+@dataclasses.dataclass
+class TracebackCaches:
+    traceback_cache: dict[Traceback, Location]
+    location_cache: dict[tuple[types.CodeType, int], Location]
+    canonical_name_cache: dict[str, str]
+    is_user_file_cache: dict[str, bool]
+
+    def __init__(self):
+        self.traceback_cache = {}
+        self.location_cache = {}
+        self.canonical_name_cache = {}
+        self.is_user_file_cache = {}
+
+
+_traceback_caches = TracebackCaches()
+_traceback_in_locations_limit = 100
+_include_full_tracebacks_in_locations = True

mlir/test/python/ir/line_info.py

@@ -0,0 +1,39 @@
+# RUN: %PYTHON %s | FileCheck %s
+import gc
+import traceback
+
+from mlir import source_info_util
+from mlir.source_info_util import _traceback_to_location
+from mlir import traceback_util
+from mlir.ir import Context
+
+# CHECK: hello
+print("hello")
+
+
+# traceback_util.register_exclusion(__file__)
+
+
+def run(f):
+    print("\nTEST:", f.__name__)
+    with Context() as ctx:
+        f()
+    gc.collect()
+    # assert Context._get_live_count() == 0
+    return f
+
+
+@run
+def foo():
+    def bar():
+        curr = source_info_util.current()
+        print(curr.name_stack)
+        print(curr.traceback)
+        traceback.print_tb(
+            traceback_util.filter_traceback(curr.traceback.as_python_traceback())
+        )
+
+        loc = _traceback_to_location(curr.traceback)
+        print(loc)
+
+    bar()