[FRONTEND] Restore error traceback filtering (#5731)

The `filter_traceback` call was commented out during the tuple PR. This
just restores it and adds a check in the relevant tests.

Author: peterbell10

python/test/unit/language/test_compile_errors.py

@@ -10,6 +10,11 @@
 from triton._internal_testing import is_interpreter, is_cuda, is_hip, is_hip_mi300, is_hip_mi200
 
 
+def format_exception(type, value, tb):
+    list_msg = traceback.format_exception(type, value, tb, chain=False)
+    return "\n".join(list_msg)
+
+
 def test_err_undefined_variable():
 
     @triton.jit
@@ -20,7 +25,9 @@ def kernel():
         triton.compile(triton.compiler.ASTSource(fn=kernel, signature={}, constexprs={}))
 
     try:
-        assert "is not defined" in str(e.value), "error should mention the undefined variable"
+        err_msg = format_exception(e.type, value=e.value, tb=e.tb)
+        assert "is not defined" in err_msg, "error should mention the undefined variable"
+        assert "code_generator.py" not in err_msg
     except AssertionError as assertion_err:
         raise assertion_err from e.value
 
@@ -35,7 +42,9 @@ def kernel():
         triton.compile(triton.compiler.ASTSource(fn=kernel, signature={}, constexprs={}))
 
     try:
-        assert "at 2:4:" in str(e.value), "error should point to the 0"
+        err_msg = format_exception(e.type, value=e.value, tb=e.tb)
+        assert "at 2:4:" in err_msg, "error should point to the 0"
+        assert "code_generator.py" not in err_msg
     except AssertionError as assertion_err:
         raise assertion_err from e.value
 
@@ -52,8 +61,11 @@ def kernel():
     try:
         assert isinstance(e.value, CompileTimeAssertionFailure)
         assert e.value.__cause__ is None
-        assert "at 2:4:" in str(e.value), "error should point to the static_assert call"
-        assert "<source unavailable>" not in str(e.value)
+        err_msg = format_exception(e.type, value=e.value, tb=e.tb)
+        print(err_msg)
+        assert "at 2:4:" in err_msg, "error should point to the static_assert call"
+        assert "<source unavailable>" not in err_msg
+        assert "code_generator.py" not in err_msg
     except AssertionError as assertion_err:
         raise assertion_err from e.value
 
@@ -70,8 +82,10 @@ def kernel():
 
     try:
         assert e.value.__cause__ is None
-        assert "at 2:4:" in str(e.value), "error should point to the `not`"
-        assert "<source unavailable>" not in str(e.value)
+        err_msg = format_exception(e.type, value=e.value, tb=e.tb)
+        assert "at 2:4:" in err_msg, "error should point to the `not`"
+        assert "<source unavailable>" not in err_msg
+        assert "code_generator.py" not in err_msg
     except AssertionError as assertion_err:
         raise assertion_err from e.value
 
@@ -86,8 +100,10 @@ def kernel():
         triton.compile(triton.compiler.ASTSource(fn=kernel, signature={}, constexprs={}))
 
     try:
-        assert "at 2:4:" in str(e.value), "error should point to the 1.0"
-        assert "<source unavailable>" not in str(e.value)
+        err_msg = format_exception(e.type, value=e.value, tb=e.tb)
+        assert "at 2:4:" in err_msg, "error should point to the 1.0"
+        assert "<source unavailable>" not in err_msg
+        assert "code_generator.py" not in err_msg
     except AssertionError as assertion_err:
         raise assertion_err from e.value
 
@@ -110,13 +126,16 @@ def kernel():
         triton.compile(triton.compiler.ASTSource(fn=kernel, signature={}, constexprs={}))
 
     try:
-        inner = e.value.__cause__
-        outer = e.value
-        assert "at 2:4:" in str(inner), "error should point to xyz"
-        assert "<source unavailable>" not in str(inner)
-
-        assert "at 3:4" in str(outer), "error should point to the nested_call"
-        assert "<source unavailable>" not in str(outer)
+        inner_exc = e.value.__cause__
+        inner = format_exception(inner_exc.__class__, inner_exc, inner_exc.__traceback__)
+        assert "at 2:4:" in inner, "error should point to xyz"
+        assert "<source unavailable>" not in inner
+        assert "code_generator.py" not in inner
+
+        outer = format_exception(e.type, value=e.value, tb=e.tb)
+        assert "at 3:4" in outer, "error should point to the nested_call"
+        assert "<source unavailable>" not in outer
+        assert "code_generator.py" not in outer
     except AssertionError as assertion_err:
         raise assertion_err from e.value
 
@@ -133,13 +152,15 @@ def kernel():
         triton.compile(triton.compiler.ASTSource(fn=kernel, signature={}, constexprs={}))
 
     try:
-        inner = e.value.__cause__
-        outer = e.value
-        assert f"{os.sep}core.py" in '\n'.join(traceback.format_tb(
-            inner.__traceback__)), "error should point inside core.py"
-
-        assert "at 2:4:" in str(outer), "error should point to expand_dims call"
-        assert "<source unavailable>" not in str(outer)
+        inner_exc = e.value.__cause__
+        inner = format_exception(inner_exc.__class__, inner_exc, inner_exc.__traceback__)
+        assert f"{os.sep}core.py" in inner, "error should point inside core.py"
+        assert "code_generator.py" not in inner
+
+        outer = format_exception(e.type, value=e.value, tb=e.tb)
+        assert "at 2:4:" in outer, "error should point to expand_dims call"
+        assert "<source unavailable>" not in outer
+        assert "code_generator.py" not in outer
     except AssertionError as assertion_err:
         raise assertion_err from e.value
 

python/triton/compiler/compiler.py

@@ -271,11 +271,11 @@ def compile(src, target=None, options=None):
 
     codegen_fns = backend.get_codegen_implementation(options)
     module_map = backend.get_module_map()
-    # try:
-    module = src.make_ir(options, codegen_fns, module_map, context)
-    # except Exception as e:
-    #     filter_traceback(e)
-    #     raise
+    try:
+        module = src.make_ir(options, codegen_fns, module_map, context)
+    except Exception as e:
+        filter_traceback(e)
+        raise
     use_ir_loc = os.environ.get("USE_IR_LOC", None)
     for ext, compile_ir in list(stages.items())[first_stage:]:
         next_module = compile_ir(module, metadata)