[FFI][PYTHON] Improve the traceback generation in python

python/tvm/ffi/cython/base.pxi

@@ -24,18 +24,6 @@ from cpython cimport PyErr_CheckSignals, PyGILState_Ensure, PyGILState_Release,
 from cpython cimport pycapsule, PyCapsule_Destructor
 from cpython cimport PyErr_SetNone
 
-# Cython internal for custom traceback creation
-cdef extern from *:
-    """
-    static void __Pyx_AddTraceback(const char *funcname, int c_line, int py_line, const char *filename);
-    static void __Pyx_AddTraceback_(const char *funcname, int c_line, int py_line, const char *filename) {
-      __Pyx_AddTraceback(funcname, c_line, py_line, filename);
-    }
-    """
-    # __Pyx_AddTraceback is a Cython internal function to add custom traceback to exception
-    # We declare __Pyx_AddTraceback_ and redirect to it to avoid mixed C/C++ calling
-    # This is a nice hack to enable us to create customized traceback refers to c++ code
-    void __Pyx_AddTraceback_(const char *funcname, int c_line, int py_line, const char *filename)
 
 # Cython binding for TVM FFI C API
 cdef extern from "tvm/ffi/c_api.h":

python/tvm/ffi/cython/error.pxi

@@ -26,43 +26,6 @@ _WITH_APPEND_TRACEBACK = None
 _TRACEBACK_TO_STR = None
 
 
-cdef inline int _raise_not_implemented_with_extra_frame(
-    const char* func_name,
-    int lineno,
-    const char* file_path
-) except -1:
-    """Helper util, raise internal """
-    PyErr_SetNone(NotImplementedError)
-    __Pyx_AddTraceback_(func_name, 0, lineno, file_path)
-    return -1
-
-
-def _append_traceback_frame(tb, filename, lineno, func):
-    """Append tracebacks to frame.
-
-    Parameters
-    ----------
-    tb : types.TracebackType
-        The traceback to append to.
-
-    filename : str
-        The filename of the traceback
-    lineno : int
-        The line number of the traceback
-    func : str
-        The function name of the traceback
-    """
-    try:
-        _raise_not_implemented_with_extra_frame(
-            c_str(func), lineno, c_str(filename)
-        )
-    except NotImplementedError as e:
-        dummy_tb = e.__traceback__
-    # skip the first frame which is the dummy frame
-    new_frame = dummy_tb.tb_next
-    return types.TracebackType(tb, new_frame.tb_frame, new_frame.tb_lasti, new_frame.tb_lineno)
-
-
 cdef class Error(Object):
     """Base class for all FFI errors, usually they are attached to errors
 

python/tvm/ffi/error.py

@@ -17,6 +17,9 @@
 # pylint: disable=invalid-name
 """Error handling."""
 import re
+import types
+import sys
+import ast
 from . import core
 
 
@@ -48,11 +51,76 @@ def _parse_traceback(traceback):
     return result
 
 
+class TracebackManager:
+    """
+    Helper to manage traceback generation
+    """
+
+    def __init__(self):
+        self._code_cache = {}
+
+    def _get_cached_code_object(self, filename, lineno, func):
+        # Hack to create a code object that points to the correct
+        # line number and function name
+        key = (filename, lineno, func)
+        # cache the code object to avoid re-creating it
+        if key in self._code_cache:
+            return self._code_cache[key]
+        # Parse to AST and zero out column info
+        # since column info are not accurate in original trace
+        tree = ast.parse("_getframe()", filename=filename, mode="eval")
+        for node in ast.walk(tree):
+            if hasattr(node, "col_offset"):
+                node.col_offset = 0
+            if hasattr(node, "end_col_offset"):
+                node.end_col_offset = 0
+        # call into get frame, bt changes the context
+        code_object = compile(tree, filename, "eval")
+        # replace the function name and line number
+        code_object = code_object.replace(co_name=func, co_firstlineno=lineno)
+        self._code_cache[key] = code_object
+        return code_object
+
+    def _create_frame(self, filename, lineno, func):
+        """Create a frame object from the filename, lineno, and func"""
+        code_object = self._get_cached_code_object(filename, lineno, func)
+        # call into get frame, but changes the context so the code
+        # points to the correct frame
+        context = {"_getframe": sys._getframe}
+        # pylint: disable=eval-used
+        return eval(code_object, context, context)
+
+    def append_traceback(self, tb, filename, lineno, func):
+        """Append a traceback to the given traceback
+
+        Parameters
+        ----------
+        tb : types.TracebackType
+            The traceback to append to.
+        filename : str
+            The filename of the traceback
+        lineno : int
+            The line number of the traceback
+        func : str
+            The function name of the traceback
+
+        Returns
+        -------
+        new_tb : types.TracebackType
+            The new traceback with the appended frame.
+        """
+        frame = self._create_frame(filename, lineno, func)
+        return types.TracebackType(tb, frame, frame.f_lasti, lineno)
+
+
+_TRACEBACK_MANAGER = TracebackManager()
+
+
 def _with_append_traceback(py_error, traceback):
     """Append the traceback to the py_error and return it"""
     tb = py_error.__traceback__
     for filename, lineno, func in reversed(_parse_traceback(traceback)):
-        tb = core._append_traceback_frame(tb, filename, lineno, func)
+        tb = _TRACEBACK_MANAGER.append_traceback(tb, filename, lineno, func)
     return py_error.with_traceback(tb)