[FRONTEND] Add TRITON_DISABLE_PYTHON_STACKTRACE envvar (#4130)

Used to disable stacktrace handler registration in the python module.

fixes #4129

The stacktrace handler disrupts normal signal propagation, so we introduce
an environment variable to disable registration at import time.

include/triton/Tools/Sys/GetEnv.hpp

@@ -33,7 +33,10 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
 };
 
 inline const std::set<std::string> CACHE_NEUTRAL_ENV_VARS = {
+    // clang-format off
     "TRITON_REPRODUCER_PATH",
+    "TRITON_DISABLE_PYTHON_STACKTRACE"
+    // clang-format on
 };
 
 namespace tools {

python/src/llvm.cc

@@ -17,6 +17,7 @@
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
 #include "llvm/Support/CodeGen.h"
+#include "llvm/Support/Signals.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/IPO/AlwaysInliner.h"
@@ -420,3 +421,9 @@ void init_triton_llvm(py::module &&m) {
     }
   });
 }
+
+void init_triton_stacktrace_hook(pybind11::module &m) {
+  if (!mlir::triton::tools::getBoolEnv("TRITON_DISABLE_PYTHON_STACKTRACE")) {
+    llvm::sys::PrintStackTraceOnErrorSignal("triton_python");
+  }
+}

python/src/main.cc

@@ -40,11 +40,12 @@ void init_triton_ir(pybind11::module &&m);
 void init_triton_llvm(pybind11::module &&m);
 void init_triton_interpreter(pybind11::module &&m);
 void init_triton_passes(pybind11::module &&m);
+void init_triton_stacktrace_hook(pybind11::module &m);
 FOR_EACH_P(DECLARE_BACKEND, TRITON_BACKENDS_TUPLE)
 
 PYBIND11_MODULE(libtriton, m) {
   m.doc() = "Python bindings to the C++ Triton API";
-  llvm::sys::PrintStackTraceOnErrorSignal("triton_python");
+  init_triton_stacktrace_hook(m);
   init_triton_env_vars(m);
   init_triton_ir(m.def_submodule("ir"));
   init_triton_passes(m.def_submodule("passes"));