[triton][beta] [Cherry-pick] '[BACKEND] Add hook for configurable/overridable compiler pass pipeline (#8137)' (#1014)

Summary:
Pull Request resolved: #1014

This is a cherry-pick of an upstream PR: triton-lang/triton#8137

Upstream commit message:
```
> [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 <puyan@puyan.org>
```

***Do not remove the following line from this commit***
Reactor Cherry-pick Revision: 905b3d1
 ---

This diff was generated by running:
```
buck run fbcode//triton/tools/reactor:reactor -- cherrypick --num-commits 1
```

Reviewed By: dshi7

Differential Revision: D94678547

fbshipit-source-id: a2af93c1274f56130c21b2b91621b7fea3763121

Author: agron911

README.md

@@ -133,6 +133,15 @@ qk_storage_alias.set_buffer_overlap(
 )
 ```
 
+**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
+```
 
 
 ### Remote buffer operations

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

@@ -455,6 +455,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
@@ -475,6 +481,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

@@ -501,6 +501,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

@@ -617,6 +617,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):