{"owner":"triton-lang","repo":"triton","hasSkills":true,"totalSkillsCount":2,"totalTokensCount":4172,"categories":["plugin-manifest","root-instruction"],"hasMcp":false,"mcpConfig":null,"found":["examples/plugins/README.md","AGENTS.md"],"skills":{"examples/plugins/README.md":"# Triton TTIR and TTGIR Out of Tree Plugin Passes\n\n## Overview\nTriton’s existing pass pipelines are assembled in the various extended compiler.py files that live in Triton’s backends. Currently when we want to insert\npasses either for downstream optimizations, custom ops, or instrumentation it is required for the compiler.py file itself to be modified and all of Triton to be\nrecompiled.\n\nIn order to allow for more downstream configurability we have implemented a custom MLIR level (TTIR and TTGIR) pass plugin and configuration system that allows for either\noverriding the compiler.py pipeline entirely or inserting passes and custom ops through a compiler pipeline hook. Example use cases include:\n- Custom ops and lowering passes\n- Custom optimization passes\n- Instrumentation and analysis passes\n- Specialized per kernel passes (e.g. kernel/model specific warp specialization)\n\nCustom passes/ops are implemented as a shared library that is loaded by Triton at JIT compile/runtime. The plugins can be implement entirely out of tree or in the Triton source tree as\nlong as the libtriton.so is linked to the plugin and the Triton include passes are used to build the plugin.\n\n## Example 1: Developing a custom pass and running triton-opt to inspect the modified IR\n``` bash\nexport TRITON_EXT_ENABLED=1;  make dev-install-llvm\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so triton-opt -tritongpu-plugin test/Plugins/test-plugin.mlir\n```\n``` MLIR\nmodule attributes {\"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:80\"} {\n  tt.func @foo() {\n    tt.return\n  }\n}\n```\n\nAfter the out of tree pass runs, becomes:\n``` MLIR\nmodule attributes {\"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:80\"} {\n  tt.func @bar() {\n    tt.return\n  }\n}\n```\nFunction \"foo\" is renamed to \"bar\" by the out of tree pass.\n\n## Example 2: Inserting a new pass into the compiler pipeline\nLet's take the following toy kernel example:\n``` python\nimport torch\nimport os\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n@triton.jit\ndef kernel(BLOCK_SIZE: tl.constexpr):\n    return\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n\nRunning as is will produce the expected output of printing the TTGIR of the kernel:\n``` bash\npython test.py\n```\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nRunning same code but loading the plugin library also produces the same results since, while the plugin pass has been loaded and registered with the\npass manager it is not inserted into the compiler pass pipeline:\n\n``` bash\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so python test.py\n```\n\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nFinally, if we both load the plugin at runtime and insert the pass pipeline hook into the kernel code:\n\n``` python\nimport torch\nimport os\nimport hashlib\nimport pathlib\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n@triton.jit\ndef kernel(BLOCK_SIZE: tl.constexpr):\n    return\n\n#These two methods must be implemented by the plugin\ndef get_key():\n    return pathlib.Path(__file__).read_text()\ndef get_hash():\n    return hashlib.sha256(get_key().encode('utf-8')).hexdigest()\n\ndef inspect_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    # If the hook is called with no arguments we assume were just after the key and hash and don't want to\n    # actually execute the pipeline yet.\n    # This no argument early return must be implemented.\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n\n    def make_ttir_wrapper(mod, metadata, opt, capability):\n        mod = self.make_ttir(mod, metadata, opt, capability)\n        pm = ir.pass_manager(mod.context)\n        pm.enable_debug()\n        passes.plugin.add_plugin(pm)\n        pm.run(mod, 'make_ttir_plugin')\n        return mod\n\n    stages[\"ttir\"] = lambda src, metadata: make_ttir_wrapper(src, metadata, options, capability)\n\n    return get_key(), get_hash()\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n\n    if \"TRITON_PLUGIN_PATHS\" in os.environ:\n      knobs.runtime.add_stages_inspection_hook = inspect_stages_hook\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n\n    # Unset the hook to go back to the standard pipeline\n    knobs.runtime.add_stages_inspection_hook = None\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n\n``` bash\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so python test.py\n```\n\nShows the pass ran and modified the kernel name but only after the hook is set. Any kernels before the hook or after the hook is unset are left unchanged.\n\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @foo() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nThe hook, as defined, in the example will insert the pass at the end of the make_ttir pipeline but it's placement in the Triton pipeline is arbitrary.\nThis functionality can be toggled on and off by just commenting out this line in kernel code (or setting to None):\nknobs.runtime.add_stages_inspection_hook = inspect_stages_hook\nwithout needing any core compiler changes or rebuilding Triton.\n\n## Example 3: Inserting a new pass into the compiler pipeline at an arbitrary point.\n\nExample 2 added a new pass to the end of the ttgir \"stage\". However the plugin pass's location is arbitrary and can be dynamically inserted anywhere in the pipeline. Replacing the inspect_stages_hook function from example 2 instead with:\n\n```python\ndef inspect_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    module_name = 'dynamic_module'\n    spec = importlib.util.spec_from_loader(module_name, loader=None)\n    module = importlib.util.module_from_spec(spec)\n    sys.modules[module_name] = module\n    stage_src = textwrap.dedent(inspect.getsource(self.make_ttir))\n    stage_src = 'from triton._C.libtriton import ir, passes, llvm, amd, nvidia\\n' + stage_src\n    # Inject plugin pass right after loop unroll in the dynamically loaded stage source\n    stage_src = stage_src.replace(\n        \"passes.ttir.add_loop_unroll(pm)\",\n        \"passes.ttir.add_loop_unroll(pm)\\n    passes.plugin.add_plugin(pm)\"\n    )\n    exec(stage_src, module.__dict__)\n    make_lambda = lambda f: lambda src, metadata: f(src, metadata, options, capability)\n    stages[\"ttir\"] = make_lambda(module.make_ttir)\n    return get_key(), get_hash()\n```\ndirects the new pass's placement based on other surrounding passes. Knowing which passes are in the pipeline a priori can be challenging, therefore in the next example we show how to dump and inspect the entire pipeline that is run for a particular kernel to allow for precise placement of specialized out of tree passes even if the upstream pass pipeline structure changes.\n\n## Example 4: Fully customizing the compiler pipeline with pass and op insertions at arbitrary locations\n\nHere we now run two kernels one with the full standard Triton pipeline and one with fully customized pipeline entirely from within\nkernel code with modifying any core Triton compiler code or recompiling. We run the kernel with a hook to output the standard pipeline, modify\nthe compiler.py file to insert our out of tree pass before add_loop_unroll pass (although there is no restriction of where it can be inserted),\nthen run the second kernel with a different pipeline. This modification can, as before, be seen in the kernel function name modification by the\ninserted pass.\n\n``` python\nimport torch\nimport os\nimport hashlib\nimport pathlib\nimport sys\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\nimport inspect\nfrom importlib.util import module_from_spec, spec_from_file_location\n\nfrom triton.backends.compiler import Language\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n\n@triton.jit\ndef kernel1(BLOCK_SIZE: tl.constexpr):\n    return\n@triton.jit\ndef kernel2(BLOCK_SIZE: tl.constexpr):\n    return\n\ndef get_key():\n    return pathlib.Path(__file__).read_text()\ndef get_hash():\n    return hashlib.sha256(get_key().encode('utf-8')).hexdigest()\n\ndef dump_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    source_code = \"# This is generated from Triton compiler.py\"\n    source_code = (\n        source_code\n        + \"\\n\"\n        + \"from triton._C.libtriton import ir, passes, llvm, amd, nvidia\"\n    )\n    source_code = source_code + \"\\n\" + \"class GPUOverrideBackend:\"\n    source_code = source_code + \"\\n\" + inspect.getsource(self.make_ttir)\n    source_code = source_code + \"\\n\" + inspect.getsource(self.make_ttgir)\n\n    with open(\"compiler_override.py\", \"w\") as file:\n        file.write(source_code)\n    return get_key(), get_hash()\ndef override_stages(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    if language != Language.TRITON:\n        return\n    full_name = \"compiler_override.py\"\n\n    print(f\"\\nOverriding compile pass stages with file {full_name}\")\n    module_name = \"triton_override_compiler_stages\"\n    spec = (\n        spec_from_file_location(module_name, full_name)\n        if os.path.isfile(full_name)\n        else None\n    )\n    if not spec:\n        return\n\n    module = module_from_spec(spec)\n    sys.modules[module_name] = module\n    spec.loader.exec_module(module)\n    if not hasattr(module, \"GPUOverrideBackend\"):\n        return\n    module = getattr(module, \"GPUOverrideBackend\")\n\n    has_func = lambda mod, name: hasattr(mod, name) and callable(getattr(mod, name))\n    make_lambda = lambda f: lambda src, metadata: f(src, metadata, options, capability)\n    if has_func(module, \"make_ttir\"):\n        stages[\"ttir\"] = make_lambda(module.make_ttir)\n    if has_func(module, \"make_ttgir\"):\n        stages[\"ttgir\"] = make_lambda(module.make_ttgir)\n    return get_key(), get_hash()\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    knobs.runtime.add_stages_inspection_hook = dump_stages_hook\n    h = kernel1[grid](BLOCK_SIZE=1024)\n    filename = \"compiler_override.py\"\n\n    with open(filename, \"r\") as infile:\n        file_str = infile.readlines()\n\n    with open(filename, \"w\") as outfile:\n        for line in file_str:\n            if \"add_loop_unroll\" in line:\n                outfile.write(\"\\n        passes.plugin.add_plugin(pm)\\n\")\n            outfile.write(line)\n    if \"TRITON_PLUGIN_PATHS\" in os.environ:\n      knobs.runtime.add_stages_inspection_hook = override_stages\n    h = kernel2[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n","AGENTS.md":"# Working on Triton\n\n## Block Programming Model\n- Triton and Gluon programs operate on logical blocks, not independent hardware threads. Scalar values, scalar control flow, and memory descriptor origins are uniform within a logical block or warp-specialized execution region. Per-element variation must be represented by tensors with valid layouts. Membar inserts the required intra-CTA barriers so shared-memory operations obey this block-level model from the user's perspective.\n- Do not introduce hidden lane- or warp-varying scalars through inline assembly, hardware IDs, or other backdoors. Such programs violate the programming model and must not be used as reproducers or as justification for compiler guards. Hardware IDs used by compiler lowerings to implement layouts are different from source-level scalar values.\n- Block-uniform does not mean loop-invariant or identical across warp-specialized regions. Preserve physical alias, execution-region, and iteration-boundary checks when reasoning about synchronization, as well as the documented waits for asynchronous operations.\n- Regression tests must be expressible as valid Triton/Gluon block programs. Targeted PTX delays may expose an existing race, but must not change values, memory effects, synchronization, or thread participation, or introduce divergence.\n\n## Descriptor Memory Effects\n- Synchronization footprints conservatively include storage reserved by zero bases in a descriptor's allocation layout. True subviews restrict the logical coordinates while retaining that allocation padding. This applies to all shared-memory and tensor-memory accesses, including FP4 padding.\n- Respect physical allocation boundaries: a zero TMEM row-16 basis leaves the other half of each warp's rows unallocated. Operations must not access beyond a true subview; verifiers must reject instruction shapes that require doing so.\n\n## Build and Testing Guidelines\n- Before running tests for native/compiler changes, run `make` in the triton directory to rebuild triton. DO NOT RUN `make` if you only changed Python code or code in `python/triton_kernels`.\n- For compiler changes, add tests in `python/test/` (pytest) or test (lit). Keep GPU-only tests in `python/test/unit/` or `python/test/gluon/`, name them `test_<feature>_<condition>`, and avoid creating new test files unless requested.\n- Run pytest with `-s --tb=short`. Run a single test with `pytest file.py::test_name`.\n- The build dir is given by `BUILD_DIR := $(shell PYTHONPATH=\"./python\" python3 -c 'from build_helpers import get_cmake_dir; print(get_cmake_dir())')`\n- Run lit from the build dir:  `cd BUILD_DIR; ninja triton-opt; lit -v test/<path>.mlir` (example: `lit -v test/TritonNvidiaGPU/tmem_layouts.mlir`).\n- Lit tests can be run locally (no GPU required).\n- Compiler crashes sometimes print an MLIR reproducer (external_resources / mlir_reproducer). Save the full MLIR + {-# ... #-} metadata to `/tmp/<file>.mlir`, then run `triton-opt /tmp/<file>.mlir --run-reproducer` to reproduce locally.\n\n## C++ Guidelines\n- In C++, never put side-effecting code in `assert`. Assertions may be compiled out, so perform mutations and other required computation before the assertion and assert only the resulting condition. This guideline does not apply to Python `assert` statements.\n\n## Lowering Guidelines\n- Triton IR uses fixed-width integer types, not MLIR's `index` type. Do not report missing `index` support as a Triton bug.\n- Lowerings must inspect only the operation they lower. Do not inspect other operations or follow loop-carried values; perform cross-operation reasoning in a separate analysis or transformation pass.\n"},"files":{"examples/plugins/README.md":"# Triton TTIR and TTGIR Out of Tree Plugin Passes\n\n## Overview\nTriton’s existing pass pipelines are assembled in the various extended compiler.py files that live in Triton’s backends. Currently when we want to insert\npasses either for downstream optimizations, custom ops, or instrumentation it is required for the compiler.py file itself to be modified and all of Triton to be\nrecompiled.\n\nIn order to allow for more downstream configurability we have implemented a custom MLIR level (TTIR and TTGIR) pass plugin and configuration system that allows for either\noverriding the compiler.py pipeline entirely or inserting passes and custom ops through a compiler pipeline hook. Example use cases include:\n- Custom ops and lowering passes\n- Custom optimization passes\n- Instrumentation and analysis passes\n- Specialized per kernel passes (e.g. kernel/model specific warp specialization)\n\nCustom passes/ops are implemented as a shared library that is loaded by Triton at JIT compile/runtime. The plugins can be implement entirely out of tree or in the Triton source tree as\nlong as the libtriton.so is linked to the plugin and the Triton include passes are used to build the plugin.\n\n## Example 1: Developing a custom pass and running triton-opt to inspect the modified IR\n``` bash\nexport TRITON_EXT_ENABLED=1;  make dev-install-llvm\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so triton-opt -tritongpu-plugin test/Plugins/test-plugin.mlir\n```\n``` MLIR\nmodule attributes {\"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:80\"} {\n  tt.func @foo() {\n    tt.return\n  }\n}\n```\n\nAfter the out of tree pass runs, becomes:\n``` MLIR\nmodule attributes {\"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:80\"} {\n  tt.func @bar() {\n    tt.return\n  }\n}\n```\nFunction \"foo\" is renamed to \"bar\" by the out of tree pass.\n\n## Example 2: Inserting a new pass into the compiler pipeline\nLet's take the following toy kernel example:\n``` python\nimport torch\nimport os\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n@triton.jit\ndef kernel(BLOCK_SIZE: tl.constexpr):\n    return\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n\nRunning as is will produce the expected output of printing the TTGIR of the kernel:\n``` bash\npython test.py\n```\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nRunning same code but loading the plugin library also produces the same results since, while the plugin pass has been loaded and registered with the\npass manager it is not inserted into the compiler pass pipeline:\n\n``` bash\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so python test.py\n```\n\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nFinally, if we both load the plugin at runtime and insert the pass pipeline hook into the kernel code:\n\n``` python\nimport torch\nimport os\nimport hashlib\nimport pathlib\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n@triton.jit\ndef kernel(BLOCK_SIZE: tl.constexpr):\n    return\n\n#These two methods must be implemented by the plugin\ndef get_key():\n    return pathlib.Path(__file__).read_text()\ndef get_hash():\n    return hashlib.sha256(get_key().encode('utf-8')).hexdigest()\n\ndef inspect_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    # If the hook is called with no arguments we assume were just after the key and hash and don't want to\n    # actually execute the pipeline yet.\n    # This no argument early return must be implemented.\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n\n    def make_ttir_wrapper(mod, metadata, opt, capability):\n        mod = self.make_ttir(mod, metadata, opt, capability)\n        pm = ir.pass_manager(mod.context)\n        pm.enable_debug()\n        passes.plugin.add_plugin(pm)\n        pm.run(mod, 'make_ttir_plugin')\n        return mod\n\n    stages[\"ttir\"] = lambda src, metadata: make_ttir_wrapper(src, metadata, options, capability)\n\n    return get_key(), get_hash()\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n\n    if \"TRITON_PLUGIN_PATHS\" in os.environ:\n      knobs.runtime.add_stages_inspection_hook = inspect_stages_hook\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n\n    # Unset the hook to go back to the standard pipeline\n    knobs.runtime.add_stages_inspection_hook = None\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n\n``` bash\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so python test.py\n```\n\nShows the pass ran and modified the kernel name but only after the hook is set. Any kernels before the hook or after the hook is unset are left unchanged.\n\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @foo() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nThe hook, as defined, in the example will insert the pass at the end of the make_ttir pipeline but it's placement in the Triton pipeline is arbitrary.\nThis functionality can be toggled on and off by just commenting out this line in kernel code (or setting to None):\nknobs.runtime.add_stages_inspection_hook = inspect_stages_hook\nwithout needing any core compiler changes or rebuilding Triton.\n\n## Example 3: Inserting a new pass into the compiler pipeline at an arbitrary point.\n\nExample 2 added a new pass to the end of the ttgir \"stage\". However the plugin pass's location is arbitrary and can be dynamically inserted anywhere in the pipeline. Replacing the inspect_stages_hook function from example 2 instead with:\n\n```python\ndef inspect_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    module_name = 'dynamic_module'\n    spec = importlib.util.spec_from_loader(module_name, loader=None)\n    module = importlib.util.module_from_spec(spec)\n    sys.modules[module_name] = module\n    stage_src = textwrap.dedent(inspect.getsource(self.make_ttir))\n    stage_src = 'from triton._C.libtriton import ir, passes, llvm, amd, nvidia\\n' + stage_src\n    # Inject plugin pass right after loop unroll in the dynamically loaded stage source\n    stage_src = stage_src.replace(\n        \"passes.ttir.add_loop_unroll(pm)\",\n        \"passes.ttir.add_loop_unroll(pm)\\n    passes.plugin.add_plugin(pm)\"\n    )\n    exec(stage_src, module.__dict__)\n    make_lambda = lambda f: lambda src, metadata: f(src, metadata, options, capability)\n    stages[\"ttir\"] = make_lambda(module.make_ttir)\n    return get_key(), get_hash()\n```\ndirects the new pass's placement based on other surrounding passes. Knowing which passes are in the pipeline a priori can be challenging, therefore in the next example we show how to dump and inspect the entire pipeline that is run for a particular kernel to allow for precise placement of specialized out of tree passes even if the upstream pass pipeline structure changes.\n\n## Example 4: Fully customizing the compiler pipeline with pass and op insertions at arbitrary locations\n\nHere we now run two kernels one with the full standard Triton pipeline and one with fully customized pipeline entirely from within\nkernel code with modifying any core Triton compiler code or recompiling. We run the kernel with a hook to output the standard pipeline, modify\nthe compiler.py file to insert our out of tree pass before add_loop_unroll pass (although there is no restriction of where it can be inserted),\nthen run the second kernel with a different pipeline. This modification can, as before, be seen in the kernel function name modification by the\ninserted pass.\n\n``` python\nimport torch\nimport os\nimport hashlib\nimport pathlib\nimport sys\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\nimport inspect\nfrom importlib.util import module_from_spec, spec_from_file_location\n\nfrom triton.backends.compiler import Language\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n\n@triton.jit\ndef kernel1(BLOCK_SIZE: tl.constexpr):\n    return\n@triton.jit\ndef kernel2(BLOCK_SIZE: tl.constexpr):\n    return\n\ndef get_key():\n    return pathlib.Path(__file__).read_text()\ndef get_hash():\n    return hashlib.sha256(get_key().encode('utf-8')).hexdigest()\n\ndef dump_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    source_code = \"# This is generated from Triton compiler.py\"\n    source_code = (\n        source_code\n        + \"\\n\"\n        + \"from triton._C.libtriton import ir, passes, llvm, amd, nvidia\"\n    )\n    source_code = source_code + \"\\n\" + \"class GPUOverrideBackend:\"\n    source_code = source_code + \"\\n\" + inspect.getsource(self.make_ttir)\n    source_code = source_code + \"\\n\" + inspect.getsource(self.make_ttgir)\n\n    with open(\"compiler_override.py\", \"w\") as file:\n        file.write(source_code)\n    return get_key(), get_hash()\ndef override_stages(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    if language != Language.TRITON:\n        return\n    full_name = \"compiler_override.py\"\n\n    print(f\"\\nOverriding compile pass stages with file {full_name}\")\n    module_name = \"triton_override_compiler_stages\"\n    spec = (\n        spec_from_file_location(module_name, full_name)\n        if os.path.isfile(full_name)\n        else None\n    )\n    if not spec:\n        return\n\n    module = module_from_spec(spec)\n    sys.modules[module_name] = module\n    spec.loader.exec_module(module)\n    if not hasattr(module, \"GPUOverrideBackend\"):\n        return\n    module = getattr(module, \"GPUOverrideBackend\")\n\n    has_func = lambda mod, name: hasattr(mod, name) and callable(getattr(mod, name))\n    make_lambda = lambda f: lambda src, metadata: f(src, metadata, options, capability)\n    if has_func(module, \"make_ttir\"):\n        stages[\"ttir\"] = make_lambda(module.make_ttir)\n    if has_func(module, \"make_ttgir\"):\n        stages[\"ttgir\"] = make_lambda(module.make_ttgir)\n    return get_key(), get_hash()\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    knobs.runtime.add_stages_inspection_hook = dump_stages_hook\n    h = kernel1[grid](BLOCK_SIZE=1024)\n    filename = \"compiler_override.py\"\n\n    with open(filename, \"r\") as infile:\n        file_str = infile.readlines()\n\n    with open(filename, \"w\") as outfile:\n        for line in file_str:\n            if \"add_loop_unroll\" in line:\n                outfile.write(\"\\n        passes.plugin.add_plugin(pm)\\n\")\n            outfile.write(line)\n    if \"TRITON_PLUGIN_PATHS\" in os.environ:\n      knobs.runtime.add_stages_inspection_hook = override_stages\n    h = kernel2[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n","AGENTS.md":"# Working on Triton\n\n## Block Programming Model\n- Triton and Gluon programs operate on logical blocks, not independent hardware threads. Scalar values, scalar control flow, and memory descriptor origins are uniform within a logical block or warp-specialized execution region. Per-element variation must be represented by tensors with valid layouts. Membar inserts the required intra-CTA barriers so shared-memory operations obey this block-level model from the user's perspective.\n- Do not introduce hidden lane- or warp-varying scalars through inline assembly, hardware IDs, or other backdoors. Such programs violate the programming model and must not be used as reproducers or as justification for compiler guards. Hardware IDs used by compiler lowerings to implement layouts are different from source-level scalar values.\n- Block-uniform does not mean loop-invariant or identical across warp-specialized regions. Preserve physical alias, execution-region, and iteration-boundary checks when reasoning about synchronization, as well as the documented waits for asynchronous operations.\n- Regression tests must be expressible as valid Triton/Gluon block programs. Targeted PTX delays may expose an existing race, but must not change values, memory effects, synchronization, or thread participation, or introduce divergence.\n\n## Descriptor Memory Effects\n- Synchronization footprints conservatively include storage reserved by zero bases in a descriptor's allocation layout. True subviews restrict the logical coordinates while retaining that allocation padding. This applies to all shared-memory and tensor-memory accesses, including FP4 padding.\n- Respect physical allocation boundaries: a zero TMEM row-16 basis leaves the other half of each warp's rows unallocated. Operations must not access beyond a true subview; verifiers must reject instruction shapes that require doing so.\n\n## Build and Testing Guidelines\n- Before running tests for native/compiler changes, run `make` in the triton directory to rebuild triton. DO NOT RUN `make` if you only changed Python code or code in `python/triton_kernels`.\n- For compiler changes, add tests in `python/test/` (pytest) or test (lit). Keep GPU-only tests in `python/test/unit/` or `python/test/gluon/`, name them `test_<feature>_<condition>`, and avoid creating new test files unless requested.\n- Run pytest with `-s --tb=short`. Run a single test with `pytest file.py::test_name`.\n- The build dir is given by `BUILD_DIR := $(shell PYTHONPATH=\"./python\" python3 -c 'from build_helpers import get_cmake_dir; print(get_cmake_dir())')`\n- Run lit from the build dir:  `cd BUILD_DIR; ninja triton-opt; lit -v test/<path>.mlir` (example: `lit -v test/TritonNvidiaGPU/tmem_layouts.mlir`).\n- Lit tests can be run locally (no GPU required).\n- Compiler crashes sometimes print an MLIR reproducer (external_resources / mlir_reproducer). Save the full MLIR + {-# ... #-} metadata to `/tmp/<file>.mlir`, then run `triton-opt /tmp/<file>.mlir --run-reproducer` to reproduce locally.\n\n## C++ Guidelines\n- In C++, never put side-effecting code in `assert`. Assertions may be compiled out, so perform mutations and other required computation before the assertion and assert only the resulting condition. This guideline does not apply to Python `assert` statements.\n\n## Lowering Guidelines\n- Triton IR uses fixed-width integer types, not MLIR's `index` type. Do not report missing `index` support as a Triton bug.\n- Lowerings must inspect only the operation they lower. Do not inspect other operations or follow loop-carried values; perform cross-operation reasoning in a separate analysis or transformation pass.\n"},"items":[{"name":"AGENTS.md","path":"AGENTS.md","rawUrl":"https://raw.githubusercontent.com/triton-lang/triton/HEAD/AGENTS.md","title":"AI Agent Protocol & Instructions","category":"root-instruction","format":"markdown","content":"# Working on Triton\n\n## Block Programming Model\n- Triton and Gluon programs operate on logical blocks, not independent hardware threads. Scalar values, scalar control flow, and memory descriptor origins are uniform within a logical block or warp-specialized execution region. Per-element variation must be represented by tensors with valid layouts. Membar inserts the required intra-CTA barriers so shared-memory operations obey this block-level model from the user's perspective.\n- Do not introduce hidden lane- or warp-varying scalars through inline assembly, hardware IDs, or other backdoors. Such programs violate the programming model and must not be used as reproducers or as justification for compiler guards. Hardware IDs used by compiler lowerings to implement layouts are different from source-level scalar values.\n- Block-uniform does not mean loop-invariant or identical across warp-specialized regions. Preserve physical alias, execution-region, and iteration-boundary checks when reasoning about synchronization, as well as the documented waits for asynchronous operations.\n- Regression tests must be expressible as valid Triton/Gluon block programs. Targeted PTX delays may expose an existing race, but must not change values, memory effects, synchronization, or thread participation, or introduce divergence.\n\n## Descriptor Memory Effects\n- Synchronization footprints conservatively include storage reserved by zero bases in a descriptor's allocation layout. True subviews restrict the logical coordinates while retaining that allocation padding. This applies to all shared-memory and tensor-memory accesses, including FP4 padding.\n- Respect physical allocation boundaries: a zero TMEM row-16 basis leaves the other half of each warp's rows unallocated. Operations must not access beyond a true subview; verifiers must reject instruction shapes that require doing so.\n\n## Build and Testing Guidelines\n- Before running tests for native/compiler changes, run `make` in the triton directory to rebuild triton. DO NOT RUN `make` if you only changed Python code or code in `python/triton_kernels`.\n- For compiler changes, add tests in `python/test/` (pytest) or test (lit). Keep GPU-only tests in `python/test/unit/` or `python/test/gluon/`, name them `test_<feature>_<condition>`, and avoid creating new test files unless requested.\n- Run pytest with `-s --tb=short`. Run a single test with `pytest file.py::test_name`.\n- The build dir is given by `BUILD_DIR := $(shell PYTHONPATH=\"./python\" python3 -c 'from build_helpers import get_cmake_dir; print(get_cmake_dir())')`\n- Run lit from the build dir:  `cd BUILD_DIR; ninja triton-opt; lit -v test/<path>.mlir` (example: `lit -v test/TritonNvidiaGPU/tmem_layouts.mlir`).\n- Lit tests can be run locally (no GPU required).\n- Compiler crashes sometimes print an MLIR reproducer (external_resources / mlir_reproducer). Save the full MLIR + {-# ... #-} metadata to `/tmp/<file>.mlir`, then run `triton-opt /tmp/<file>.mlir --run-reproducer` to reproduce locally.\n\n## C++ Guidelines\n- In C++, never put side-effecting code in `assert`. Assertions may be compiled out, so perform mutations and other required computation before the assertion and assert only the resulting condition. This guideline does not apply to Python `assert` statements.\n\n## Lowering Guidelines\n- Triton IR uses fixed-width integer types, not MLIR's `index` type. Do not report missing `index` support as a Triton bug.\n- Lowerings must inspect only the operation they lower. Do not inspect other operations or follow loop-carried values; perform cross-operation reasoning in a separate analysis or transformation pass.\n","isInternal":false,"tokens":776,"sizeBytes":3646},{"name":"README.md","path":"examples/plugins/README.md","rawUrl":"https://raw.githubusercontent.com/triton-lang/triton/HEAD/examples/plugins/README.md","title":"plugins Documentation","category":"plugin-manifest","format":"markdown","content":"# Triton TTIR and TTGIR Out of Tree Plugin Passes\n\n## Overview\nTriton’s existing pass pipelines are assembled in the various extended compiler.py files that live in Triton’s backends. Currently when we want to insert\npasses either for downstream optimizations, custom ops, or instrumentation it is required for the compiler.py file itself to be modified and all of Triton to be\nrecompiled.\n\nIn order to allow for more downstream configurability we have implemented a custom MLIR level (TTIR and TTGIR) pass plugin and configuration system that allows for either\noverriding the compiler.py pipeline entirely or inserting passes and custom ops through a compiler pipeline hook. Example use cases include:\n- Custom ops and lowering passes\n- Custom optimization passes\n- Instrumentation and analysis passes\n- Specialized per kernel passes (e.g. kernel/model specific warp specialization)\n\nCustom passes/ops are implemented as a shared library that is loaded by Triton at JIT compile/runtime. The plugins can be implement entirely out of tree or in the Triton source tree as\nlong as the libtriton.so is linked to the plugin and the Triton include passes are used to build the plugin.\n\n## Example 1: Developing a custom pass and running triton-opt to inspect the modified IR\n``` bash\nexport TRITON_EXT_ENABLED=1;  make dev-install-llvm\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so triton-opt -tritongpu-plugin test/Plugins/test-plugin.mlir\n```\n``` MLIR\nmodule attributes {\"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:80\"} {\n  tt.func @foo() {\n    tt.return\n  }\n}\n```\n\nAfter the out of tree pass runs, becomes:\n``` MLIR\nmodule attributes {\"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:80\"} {\n  tt.func @bar() {\n    tt.return\n  }\n}\n```\nFunction \"foo\" is renamed to \"bar\" by the out of tree pass.\n\n## Example 2: Inserting a new pass into the compiler pipeline\nLet's take the following toy kernel example:\n``` python\nimport torch\nimport os\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n@triton.jit\ndef kernel(BLOCK_SIZE: tl.constexpr):\n    return\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n\nRunning as is will produce the expected output of printing the TTGIR of the kernel:\n``` bash\npython test.py\n```\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nRunning same code but loading the plugin library also produces the same results since, while the plugin pass has been loaded and registered with the\npass manager it is not inserted into the compiler pass pipeline:\n\n``` bash\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so python test.py\n```\n\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nFinally, if we both load the plugin at runtime and insert the pass pipeline hook into the kernel code:\n\n``` python\nimport torch\nimport os\nimport hashlib\nimport pathlib\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n@triton.jit\ndef kernel(BLOCK_SIZE: tl.constexpr):\n    return\n\n#These two methods must be implemented by the plugin\ndef get_key():\n    return pathlib.Path(__file__).read_text()\ndef get_hash():\n    return hashlib.sha256(get_key().encode('utf-8')).hexdigest()\n\ndef inspect_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    # If the hook is called with no arguments we assume were just after the key and hash and don't want to\n    # actually execute the pipeline yet.\n    # This no argument early return must be implemented.\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n\n    def make_ttir_wrapper(mod, metadata, opt, capability):\n        mod = self.make_ttir(mod, metadata, opt, capability)\n        pm = ir.pass_manager(mod.context)\n        pm.enable_debug()\n        passes.plugin.add_plugin(pm)\n        pm.run(mod, 'make_ttir_plugin')\n        return mod\n\n    stages[\"ttir\"] = lambda src, metadata: make_ttir_wrapper(src, metadata, options, capability)\n\n    return get_key(), get_hash()\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n\n    if \"TRITON_PLUGIN_PATHS\" in os.environ:\n      knobs.runtime.add_stages_inspection_hook = inspect_stages_hook\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n\n    # Unset the hook to go back to the standard pipeline\n    knobs.runtime.add_stages_inspection_hook = None\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n\n``` bash\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so python test.py\n```\n\nShows the pass ran and modified the kernel name but only after the hook is set. Any kernels before the hook or after the hook is unset are left unchanged.\n\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @foo() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nThe hook, as defined, in the example will insert the pass at the end of the make_ttir pipeline but it's placement in the Triton pipeline is arbitrary.\nThis functionality can be toggled on and off by just commenting out this line in kernel code (or setting to None):\nknobs.runtime.add_stages_inspection_hook = inspect_stages_hook\nwithout needing any core compiler changes or rebuilding Triton.\n\n## Example 3: Inserting a new pass into the compiler pipeline at an arbitrary point.\n\nExample 2 added a new pass to the end of the ttgir \"stage\". However the plugin pass's location is arbitrary and can be dynamically inserted anywhere in the pipeline. Replacing the inspect_stages_hook function from example 2 instead with:\n\n```python\ndef inspect_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    module_name = 'dynamic_module'\n    spec = importlib.util.spec_from_loader(module_name, loader=None)\n    module = importlib.util.module_from_spec(spec)\n    sys.modules[module_name] = module\n    stage_src = textwrap.dedent(inspect.getsource(self.make_ttir))\n    stage_src = 'from triton._C.libtriton import ir, passes, llvm, amd, nvidia\\n' + stage_src\n    # Inject plugin pass right after loop unroll in the dynamically loaded stage source\n    stage_src = stage_src.replace(\n        \"passes.ttir.add_loop_unroll(pm)\",\n        \"passes.ttir.add_loop_unroll(pm)\\n    passes.plugin.add_plugin(pm)\"\n    )\n    exec(stage_src, module.__dict__)\n    make_lambda = lambda f: lambda src, metadata: f(src, metadata, options, capability)\n    stages[\"ttir\"] = make_lambda(module.make_ttir)\n    return get_key(), get_hash()\n```\ndirects the new pass's placement based on other surrounding passes. Knowing which passes are in the pipeline a priori can be challenging, therefore in the next example we show how to dump and inspect the entire pipeline that is run for a particular kernel to allow for precise placement of specialized out of tree passes even if the upstream pass pipeline structure changes.\n\n## Example 4: Fully customizing the compiler pipeline with pass and op insertions at arbitrary locations\n\nHere we now run two kernels one with the full standard Triton pipeline and one with fully customized pipeline entirely from within\nkernel code with modifying any core Triton compiler code or recompiling. We run the kernel with a hook to output the standard pipeline, modify\nthe compiler.py file to insert our out of tree pass before add_loop_unroll pass (although there is no restriction of where it can be inserted),\nthen run the second kernel with a different pipeline. This modification can, as before, be seen in the kernel function name modification by the\ninserted pass.\n\n``` python\nimport torch\nimport os\nimport hashlib\nimport pathlib\nimport sys\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\nimport inspect\nfrom importlib.util import module_from_spec, spec_from_file_location\n\nfrom triton.backends.compiler import Language\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n\n@triton.jit\ndef kernel1(BLOCK_SIZE: tl.constexpr):\n    return\n@triton.jit\ndef kernel2(BLOCK_SIZE: tl.constexpr):\n    return\n\ndef get_key():\n    return pathlib.Path(__file__).read_text()\ndef get_hash():\n    return hashlib.sha256(get_key().encode('utf-8')).hexdigest()\n\ndef dump_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    source_code = \"# This is generated from Triton compiler.py\"\n    source_code = (\n        source_code\n        + \"\\n\"\n        + \"from triton._C.libtriton import ir, passes, llvm, amd, nvidia\"\n    )\n    source_code = source_code + \"\\n\" + \"class GPUOverrideBackend:\"\n    source_code = source_code + \"\\n\" + inspect.getsource(self.make_ttir)\n    source_code = source_code + \"\\n\" + inspect.getsource(self.make_ttgir)\n\n    with open(\"compiler_override.py\", \"w\") as file:\n        file.write(source_code)\n    return get_key(), get_hash()\ndef override_stages(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    if language != Language.TRITON:\n        return\n    full_name = \"compiler_override.py\"\n\n    print(f\"\\nOverriding compile pass stages with file {full_name}\")\n    module_name = \"triton_override_compiler_stages\"\n    spec = (\n        spec_from_file_location(module_name, full_name)\n        if os.path.isfile(full_name)\n        else None\n    )\n    if not spec:\n        return\n\n    module = module_from_spec(spec)\n    sys.modules[module_name] = module\n    spec.loader.exec_module(module)\n    if not hasattr(module, \"GPUOverrideBackend\"):\n        return\n    module = getattr(module, \"GPUOverrideBackend\")\n\n    has_func = lambda mod, name: hasattr(mod, name) and callable(getattr(mod, name))\n    make_lambda = lambda f: lambda src, metadata: f(src, metadata, options, capability)\n    if has_func(module, \"make_ttir\"):\n        stages[\"ttir\"] = make_lambda(module.make_ttir)\n    if has_func(module, \"make_ttgir\"):\n        stages[\"ttgir\"] = make_lambda(module.make_ttgir)\n    return get_key(), get_hash()\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    knobs.runtime.add_stages_inspection_hook = dump_stages_hook\n    h = kernel1[grid](BLOCK_SIZE=1024)\n    filename = \"compiler_override.py\"\n\n    with open(filename, \"r\") as infile:\n        file_str = infile.readlines()\n\n    with open(filename, \"w\") as outfile:\n        for line in file_str:\n            if \"add_loop_unroll\" in line:\n                outfile.write(\"\\n        passes.plugin.add_plugin(pm)\\n\")\n            outfile.write(line)\n    if \"TRITON_PLUGIN_PATHS\" in os.environ:\n      knobs.runtime.add_stages_inspection_hook = override_stages\n    h = kernel2[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n","isInternal":false,"tokens":3396,"sizeBytes":13193}],"systemPromptSnippet":"<yakaai_skills repo=\"triton-lang/triton\">\n<!-- File: AGENTS.md (Tokens: ~776 | Category: root-instruction) -->\n# Working on Triton\n\n## Block Programming Model\n- Triton and Gluon programs operate on logical blocks, not independent hardware threads. Scalar values, scalar control flow, and memory descriptor origins are uniform within a logical block or warp-specialized execution region. Per-element variation must be represented by tensors with valid layouts. Membar inserts the required intra-CTA barriers so shared-memory operations obey this block-level model from the user's perspective.\n- Do not introduce hidden lane- or warp-varying scalars through inline assembly, hardware IDs, or other backdoors. Such programs violate the programming model and must not be used as reproducers or as justification for compiler guards. Hardware IDs used by compiler lowerings to implement layouts are different from source-level scalar values.\n- Block-uniform does not mean loop-invariant or identical across warp-specialized regions. Preserve physical alias, execution-region, and iteration-boundary checks when reasoning about synchronization, as well as the documented waits for asynchronous operations.\n- Regression tests must be expressible as valid Triton/Gluon block programs. Targeted PTX delays may expose an existing race, but must not change values, memory effects, synchronization, or thread participation, or introduce divergence.\n\n## Descriptor Memory Effects\n- Synchronization footprints conservatively include storage reserved by zero bases in a descriptor's allocation layout. True subviews restrict the logical coordinates while retaining that allocation padding. This applies to all shared-memory and tensor-memory accesses, including FP4 padding.\n- Respect physical allocation boundaries: a zero TMEM row-16 basis leaves the other half of each warp's rows unallocated. Operations must not access beyond a true subview; verifiers must reject instruction shapes that require doing so.\n\n## Build and Testing Guidelines\n- Before running tests for native/compiler changes, run `make` in the triton directory to rebuild triton. DO NOT RUN `make` if you only changed Python code or code in `python/triton_kernels`.\n- For compiler changes, add tests in `python/test/` (pytest) or test (lit). Keep GPU-only tests in `python/test/unit/` or `python/test/gluon/`, name them `test_<feature>_<condition>`, and avoid creating new test files unless requested.\n- Run pytest with `-s --tb=short`. Run a single test with `pytest file.py::test_name`.\n- The build dir is given by `BUILD_DIR := $(shell PYTHONPATH=\"./python\" python3 -c 'from build_helpers import get_cmake_dir; print(get_cmake_dir())')`\n- Run lit from the build dir:  `cd BUILD_DIR; ninja triton-opt; lit -v test/<path>.mlir` (example: `lit -v test/TritonNvidiaGPU/tmem_layouts.mlir`).\n- Lit tests can be run locally (no GPU required).\n- Compiler crashes sometimes print an MLIR reproducer (external_resources / mlir_reproducer). Save the full MLIR + {-# ... #-} metadata to `/tmp/<file>.mlir`, then run `triton-opt /tmp/<file>.mlir --run-reproducer` to reproduce locally.\n\n## C++ Guidelines\n- In C++, never put side-effecting code in `assert`. Assertions may be compiled out, so perform mutations and other required computation before the assertion and assert only the resulting condition. This guideline does not apply to Python `assert` statements.\n\n## Lowering Guidelines\n- Triton IR uses fixed-width integer types, not MLIR's `index` type. Do not report missing `index` support as a Triton bug.\n- Lowerings must inspect only the operation they lower. Do not inspect other operations or follow loop-carried values; perform cross-operation reasoning in a separate analysis or transformation pass.\n\n<!-- File: examples/plugins/README.md (Tokens: ~3396 | Category: plugin-manifest) -->\n# Triton TTIR and TTGIR Out of Tree Plugin Passes\n\n## Overview\nTriton’s existing pass pipelines are assembled in the various extended compiler.py files that live in Triton’s backends. Currently when we want to insert\npasses either for downstream optimizations, custom ops, or instrumentation it is required for the compiler.py file itself to be modified and all of Triton to be\nrecompiled.\n\nIn order to allow for more downstream configurability we have implemented a custom MLIR level (TTIR and TTGIR) pass plugin and configuration system that allows for either\noverriding the compiler.py pipeline entirely or inserting passes and custom ops through a compiler pipeline hook. Example use cases include:\n- Custom ops and lowering passes\n- Custom optimization passes\n- Instrumentation and analysis passes\n- Specialized per kernel passes (e.g. kernel/model specific warp specialization)\n\nCustom passes/ops are implemented as a shared library that is loaded by Triton at JIT compile/runtime. The plugins can be implement entirely out of tree or in the Triton source tree as\nlong as the libtriton.so is linked to the plugin and the Triton include passes are used to build the plugin.\n\n## Example 1: Developing a custom pass and running triton-opt to inspect the modified IR\n``` bash\nexport TRITON_EXT_ENABLED=1;  make dev-install-llvm\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so triton-opt -tritongpu-plugin test/Plugins/test-plugin.mlir\n```\n``` MLIR\nmodule attributes {\"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:80\"} {\n  tt.func @foo() {\n    tt.return\n  }\n}\n```\n\nAfter the out of tree pass runs, becomes:\n``` MLIR\nmodule attributes {\"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:80\"} {\n  tt.func @bar() {\n    tt.return\n  }\n}\n```\nFunction \"foo\" is renamed to \"bar\" by the out of tree pass.\n\n## Example 2: Inserting a new pass into the compiler pipeline\nLet's take the following toy kernel example:\n``` python\nimport torch\nimport os\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n@triton.jit\ndef kernel(BLOCK_SIZE: tl.constexpr):\n    return\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n\nRunning as is will produce the expected output of printing the TTGIR of the kernel:\n``` bash\npython test.py\n```\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nRunning same code but loading the plugin library also produces the same results since, while the plugin pass has been loaded and registered with the\npass manager it is not inserted into the compiler pass pipeline:\n\n``` bash\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so python test.py\n```\n\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nFinally, if we both load the plugin at runtime and insert the pass pipeline hook into the kernel code:\n\n``` python\nimport torch\nimport os\nimport hashlib\nimport pathlib\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n@triton.jit\ndef kernel(BLOCK_SIZE: tl.constexpr):\n    return\n\n#These two methods must be implemented by the plugin\ndef get_key():\n    return pathlib.Path(__file__).read_text()\ndef get_hash():\n    return hashlib.sha256(get_key().encode('utf-8')).hexdigest()\n\ndef inspect_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    # If the hook is called with no arguments we assume were just after the key and hash and don't want to\n    # actually execute the pipeline yet.\n    # This no argument early return must be implemented.\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n\n    def make_ttir_wrapper(mod, metadata, opt, capability):\n        mod = self.make_ttir(mod, metadata, opt, capability)\n        pm = ir.pass_manager(mod.context)\n        pm.enable_debug()\n        passes.plugin.add_plugin(pm)\n        pm.run(mod, 'make_ttir_plugin')\n        return mod\n\n    stages[\"ttir\"] = lambda src, metadata: make_ttir_wrapper(src, metadata, options, capability)\n\n    return get_key(), get_hash()\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n\n    if \"TRITON_PLUGIN_PATHS\" in os.environ:\n      knobs.runtime.add_stages_inspection_hook = inspect_stages_hook\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n\n    # Unset the hook to go back to the standard pipeline\n    knobs.runtime.add_stages_inspection_hook = None\n    h = kernel[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n\n``` bash\nTRITON_PLUGIN_PATHS=/home/triton/python/triton/plugins/libTritonPluginsTestLib.so python test.py\n```\n\nShows the pass ran and modified the kernel name but only after the hook is set. Any kernels before the hook or after the hook is unset are left unchanged.\n\n``` MLIR\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @foo() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n\nmodule attributes {\"ttg.num-ctas\" = 1 : i32, \"ttg.num-warps\" = 4 : i32, ttg.target = \"cuda:90\", \"ttg.threads-per-warp\" = 32 : i32} {\n  tt.func public @kernel() attributes {noinline = false} {\n    tt.return loc(#loc1)\n  } loc(#loc)\n} loc(#loc)\n#loc = loc(\"/home/triton/test.py\":13:0)\n#loc1 = loc(\"/home/triton/test.py\":14:4)\n```\n\nThe hook, as defined, in the example will insert the pass at the end of the make_ttir pipeline but it's placement in the Triton pipeline is arbitrary.\nThis functionality can be toggled on and off by just commenting out this line in kernel code (or setting to None):\nknobs.runtime.add_stages_inspection_hook = inspect_stages_hook\nwithout needing any core compiler changes or rebuilding Triton.\n\n## Example 3: Inserting a new pass into the compiler pipeline at an arbitrary point.\n\nExample 2 added a new pass to the end of the ttgir \"stage\". However the plugin pass's location is arbitrary and can be dynamically inserted anywhere in the pipeline. Replacing the inspect_stages_hook function from example 2 instead with:\n\n```python\ndef inspect_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    module_name = 'dynamic_module'\n    spec = importlib.util.spec_from_loader(module_name, loader=None)\n    module = importlib.util.module_from_spec(spec)\n    sys.modules[module_name] = module\n    stage_src = textwrap.dedent(inspect.getsource(self.make_ttir))\n    stage_src = 'from triton._C.libtriton import ir, passes, llvm, amd, nvidia\\n' + stage_src\n    # Inject plugin pass right after loop unroll in the dynamically loaded stage source\n    stage_src = stage_src.replace(\n        \"passes.ttir.add_loop_unroll(pm)\",\n        \"passes.ttir.add_loop_unroll(pm)\\n    passes.plugin.add_plugin(pm)\"\n    )\n    exec(stage_src, module.__dict__)\n    make_lambda = lambda f: lambda src, metadata: f(src, metadata, options, capability)\n    stages[\"ttir\"] = make_lambda(module.make_ttir)\n    return get_key(), get_hash()\n```\ndirects the new pass's placement based on other surrounding passes. Knowing which passes are in the pipeline a priori can be challenging, therefore in the next example we show how to dump and inspect the entire pipeline that is run for a particular kernel to allow for precise placement of specialized out of tree passes even if the upstream pass pipeline structure changes.\n\n## Example 4: Fully customizing the compiler pipeline with pass and op insertions at arbitrary locations\n\nHere we now run two kernels one with the full standard Triton pipeline and one with fully customized pipeline entirely from within\nkernel code with modifying any core Triton compiler code or recompiling. We run the kernel with a hook to output the standard pipeline, modify\nthe compiler.py file to insert our out of tree pass before add_loop_unroll pass (although there is no restriction of where it can be inserted),\nthen run the second kernel with a different pipeline. This modification can, as before, be seen in the kernel function name modification by the\ninserted pass.\n\n``` python\nimport torch\nimport os\nimport hashlib\nimport pathlib\nimport sys\n\nimport triton\nimport triton.language as tl\nfrom triton._C.libtriton import ir, passes\nfrom triton import knobs\nimport inspect\nfrom importlib.util import module_from_spec, spec_from_file_location\n\nfrom triton.backends.compiler import Language\n\nDEVICE = triton.runtime.driver.active.get_active_torch_device()\n\n\n@triton.jit\ndef kernel1(BLOCK_SIZE: tl.constexpr):\n    return\n@triton.jit\ndef kernel2(BLOCK_SIZE: tl.constexpr):\n    return\n\ndef get_key():\n    return pathlib.Path(__file__).read_text()\ndef get_hash():\n    return hashlib.sha256(get_key().encode('utf-8')).hexdigest()\n\ndef dump_stages_hook(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    source_code = \"# This is generated from Triton compiler.py\"\n    source_code = (\n        source_code\n        + \"\\n\"\n        + \"from triton._C.libtriton import ir, passes, llvm, amd, nvidia\"\n    )\n    source_code = source_code + \"\\n\" + \"class GPUOverrideBackend:\"\n    source_code = source_code + \"\\n\" + inspect.getsource(self.make_ttir)\n    source_code = source_code + \"\\n\" + inspect.getsource(self.make_ttgir)\n\n    with open(\"compiler_override.py\", \"w\") as file:\n        file.write(source_code)\n    return get_key(), get_hash()\ndef override_stages(self=None, stages=None, options=None, language=None, capability=None):\n    if all(arg is None for arg in (stages, options, language, capability)):\n        return get_key(), get_hash()\n    if language != Language.TRITON:\n        return\n    full_name = \"compiler_override.py\"\n\n    print(f\"\\nOverriding compile pass stages with file {full_name}\")\n    module_name = \"triton_override_compiler_stages\"\n    spec = (\n        spec_from_file_location(module_name, full_name)\n        if os.path.isfile(full_name)\n        else None\n    )\n    if not spec:\n        return\n\n    module = module_from_spec(spec)\n    sys.modules[module_name] = module\n    spec.loader.exec_module(module)\n    if not hasattr(module, \"GPUOverrideBackend\"):\n        return\n    module = getattr(module, \"GPUOverrideBackend\")\n\n    has_func = lambda mod, name: hasattr(mod, name) and callable(getattr(mod, name))\n    make_lambda = lambda f: lambda src, metadata: f(src, metadata, options, capability)\n    if has_func(module, \"make_ttir\"):\n        stages[\"ttir\"] = make_lambda(module.make_ttir)\n    if has_func(module, \"make_ttgir\"):\n        stages[\"ttgir\"] = make_lambda(module.make_ttgir)\n    return get_key(), get_hash()\n\nif __name__ == '__main__':\n\n    size = 98432\n    x = torch.rand(size, device=DEVICE)\n    output = torch.empty_like(x)\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n\n    knobs.runtime.add_stages_inspection_hook = dump_stages_hook\n    h = kernel1[grid](BLOCK_SIZE=1024)\n    filename = \"compiler_override.py\"\n\n    with open(filename, \"r\") as infile:\n        file_str = infile.readlines()\n\n    with open(filename, \"w\") as outfile:\n        for line in file_str:\n            if \"add_loop_unroll\" in line:\n                outfile.write(\"\\n        passes.plugin.add_plugin(pm)\\n\")\n            outfile.write(line)\n    if \"TRITON_PLUGIN_PATHS\" in os.environ:\n      knobs.runtime.add_stages_inspection_hook = override_stages\n    h = kernel2[grid](BLOCK_SIZE=1024)\n    print(h.asm[\"ttgir\"])\n```\n</yakaai_skills>"}