diff --git a/3rdparty/cutlass_fpA_intB_gemm b/3rdparty/cutlass_fpA_intB_gemm index b71d94a4ccd6..0697a5113661 160000 --- a/3rdparty/cutlass_fpA_intB_gemm +++ b/3rdparty/cutlass_fpA_intB_gemm @@ -1 +1 @@ -Subproject commit b71d94a4ccd6573c9cbd4056c9ce660f110d33f0 +Subproject commit 0697a511366194fc305649da0746308439fd7a75 diff --git a/python/tvm/ffi/cython/base.pxi b/python/tvm/ffi/cython/base.pxi index 50831be462ad..40dfffe3d230 100644 --- a/python/tvm/ffi/cython/base.pxi +++ b/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": diff --git a/python/tvm/ffi/cython/error.pxi b/python/tvm/ffi/cython/error.pxi index 8da630873ede..968860390a3c 100644 --- a/python/tvm/ffi/cython/error.pxi +++ b/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 diff --git a/python/tvm/ffi/error.py b/python/tvm/ffi/error.py index 7c97ff43b357..a7714cb58ffd 100644 --- a/python/tvm/ffi/error.py +++ b/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)