Add an option to remove debug info from the module (#6837)

Sometimes having minimal module dump without location info clobbering
the IR is usefull.

Author: pawelszczerbuk

README.md

@@ -212,6 +212,7 @@ See [`python/triton/knobs.py`](python/triton/knobs.py) for the full list of conf
 - `TRITON_OVERRIDE_DIR` specifies the directory from which to load the IR/ptx/amdgcn files when `TRITON_KERNEL_OVERRIDE` is set to 1.
 - `TRITON_F32_DEFAULT` sets the default input precision of `tl.dot` when using 32-bit floats, which can be either `ieee`, `tf32`, or `tf32x3`.
 - `TRITON_FRONT_END_DEBUGGING=1` disables exception wrapping when an error occurs in the compiler frontend, allowing the full stack trace to be seen.
+- `TRITON_STRIP_DEBUG_INFO` removes all debug information from the module, including location information
 
 N.B. Some of these environment variables don't have a knob in `knobs.py`-- those are only relevant to the C++ layer(s), hence they don't exist in the python layer.
 

python/src/passes.cc

@@ -25,6 +25,7 @@ void init_triton_analysis(py::module &&m) {
 
 void init_triton_passes_common(py::module &&m) {
   using namespace mlir;
+  ADD_PASS_WRAPPER_0("add_strip_debug_info", createStripDebugInfoPass);
   ADD_PASS_WRAPPER_0("add_sccp", createSCCPPass);
   ADD_PASS_WRAPPER_0("add_symbol_dce", createSymbolDCEPass);
   ADD_PASS_WRAPPER_0("add_inliner", createInlinerPass);

python/triton/knobs.py

@@ -358,6 +358,7 @@ class compilation_knobs(base_knobs):
     disable_line_info: env_bool = env_bool("TRITON_DISABLE_LINE_INFO")
     front_end_debugging: env_bool = env_bool("TRITON_FRONT_END_DEBUGGING")
     allow_non_constexpr_globals: env_bool = env_bool("TRITON_ALLOW_NON_CONSTEXPR_GLOBALS")
+    strip_debug_info: env_bool = env_bool("TRITON_STRIP_DEBUG_INFO")
     listener: Union[CompilationListener, None] = None
 
 

third_party/amd/backend/compiler.py

@@ -195,6 +195,8 @@ def path_to_rocm_lld():
     def make_ttir(mod, metadata, options):
         pm = ir.pass_manager(mod.context)
         pm.enable_debug()
+        if knobs.compilation.strip_debug_info:
+            passes.common.add_strip_debug_info(pm)
         passes.common.add_inliner(pm)
         passes.ttir.add_rewrite_tensor_pointer(pm)
         passes.ttir.add_rewrite_tensor_descriptor_to_pointer(pm)

third_party/nvidia/backend/compiler.py

@@ -207,6 +207,8 @@ def load_dialects(self, ctx):
     def make_ttir(mod, metadata, opt, capability):
         pm = ir.pass_manager(mod.context)
         pm.enable_debug()
+        if knobs.compilation.strip_debug_info:
+            passes.common.add_strip_debug_info(pm)
         passes.common.add_inliner(pm)
         passes.ttir.add_rewrite_tensor_pointer(pm)
         if capability // 10 < 9: