[BACKEND] Add hook for configurable/overridable compiler pass pipeline (#8137)

Triton’s existing pass pipelines are explicitly defined in the various
extended compiler.py files that live in Triton’s backends. Currently
when we require insertion of passes either for instrumentation or for
the addition of downstream optimization and custom lowering it is
required for the compiler.py file itself to be modified.

In order to allow for more downstream configurability and as a first
step toward more custom MLIR level pass plugins, we add a hook into the
compiler stages to allow for a more configurable pass manager system
setup.

Using Python inspection routines coupled with the hook allows for more
fine grained control of things like enabling/disabling passes for
specific kernels with eventually being able to load and insert
completely out of tree ops/passes in arbitrary places in the stages
pipeline.

Co-authored with @plotfi

---------

Co-authored-by: Puyan Lotfi <[email protected]>

Author: CRobeck

README.md

@@ -262,6 +262,15 @@ export TRITON_OVERRIDE_DIR=<override_dir>
 # Step 4: Run the kernel again to see the overridden result
 ```
 
+**Compiler Pipeline Inspection Steps**
+To introspect the pipeline `add_stages`, before running your kernels, simply set
+the add_stages_inspection_hook like so:
+
+```python
+def inspect_stages(_self, stages, options, language, capability):
+    # inspect or modify add_stages here
+triton.knobs.runtime.add_stages_inspection_hook = inspect_stages
+```
 
 # Changelog
 

python/test/backend/test_stages_inspection.py

@@ -0,0 +1,53 @@
+import triton
+from triton import knobs
+
+import os
+import pathlib
+
+
+def test_inspection(monkeypatch, tmp_path: pathlib.Path):
+    stage_name = 'make_ttgir'
+    curr_repro_path = tmp_path / ("repro_prefix." + stage_name + ".repro.mlir")
+    repro_path = tmp_path / "repro_prefix"
+
+    monkeypatch.setenv("TRITON_ALWAYS_COMPILE", "1")
+    monkeypatch.setenv("TRITON_REPRODUCER_PATH", str(repro_path))
+
+    inspect_stages_hook_called = False
+    make_ttgir_wrapper_called = False
+
+    def inspect_stages_hook(self, stages, options, language, capability):
+        nonlocal inspect_stages_hook_called
+        inspect_stages_hook_called = True
+
+        def make_ttgir_wrapper(src, metadata, options, capability):
+            nonlocal make_ttgir_wrapper_called
+            make_ttgir_wrapper_called = True
+            return self.make_ttgir(src, metadata, options, capability)
+
+        stages["ttgir"] = lambda src, metadata: make_ttgir_wrapper(src, metadata, options, capability)
+
+    @triton.jit
+    def k1():
+        return
+
+    @triton.jit
+    def k2():
+        return
+
+    # Run once to get the clean/golden repro dump
+    k1[(1, )]()
+    assert not inspect_stages_hook_called and not make_ttgir_wrapper_called
+    assert os.path.exists(curr_repro_path)
+    golden_repro = curr_repro_path.read_text()
+    curr_repro_path.unlink()
+
+    # Setup hook and call again, check if hooks got called
+    knobs.runtime.add_stages_inspection_hook = inspect_stages_hook
+    k2[(1, )]()
+    assert inspect_stages_hook_called and make_ttgir_wrapper_called
+    assert os.path.exists(curr_repro_path)
+    hook_repro = curr_repro_path.read_text()
+
+    # Check that repros match
+    assert golden_repro.replace('k1', 'dummy') == hook_repro.replace('k2', 'dummy')

python/triton/knobs.py

@@ -447,6 +447,12 @@ def __call__(self, *, key: str, repr: str, fn: JitFunctionInfo, compile: JITHook
         ...
 
 
+class PipelineStagesHook(Protocol):
+
+    def __call__(self, stages, options, language, capability):
+        ...
+
+
 class runtime_knobs(base_knobs):
     interpret: env_bool = env_bool("TRITON_INTERPRET")
     # debug is on critical path for kernel launches
@@ -465,6 +471,9 @@ class runtime_knobs(base_knobs):
     # jit_cache_hook will always be called before compilation and jit_post_compile_hook after.
     jit_post_compile_hook: Optional[JITHook] = None
 
+    # Hook for inspecting compiler pipeline stages
+    add_stages_inspection_hook: Optional[PipelineStagesHook] = None
+
 
 class language_knobs(base_knobs):
     fp32_default: env_opt_str = env_opt_str("TRITON_F32_DEFAULT")

third_party/amd/backend/compiler.py

@@ -456,6 +456,8 @@ def add_stages(self, stages, options, language):
         stages["llir"] = lambda src, metadata: self.make_llir(src, metadata, options)
         stages["amdgcn"] = lambda src, metadata: self.make_amdgcn(src, metadata, options)
         stages["hsaco"] = lambda src, metadata: self.make_hsaco(src, metadata, options)
+        if knobs.runtime.add_stages_inspection_hook is not None:
+            knobs.runtime.add_stages_inspection_hook(self, stages, options, language, None)
 
     @functools.lru_cache()
     def hash(self):

third_party/nvidia/backend/compiler.py

@@ -519,6 +519,8 @@ def add_stages(self, stages, options, language):
         stages["llir"] = lambda src, metadata: self.make_llir(src, metadata, options, capability)
         stages["ptx"] = lambda src, metadata: self.make_ptx(src, metadata, options, self.target.arch)
         stages["cubin"] = lambda src, metadata: self.make_cubin(src, metadata, options, self.target.arch)
+        if knobs.runtime.add_stages_inspection_hook is not None:
+            knobs.runtime.add_stages_inspection_hook(self, stages, options, language, capability)
 
     @functools.lru_cache()
     def hash(self):