From 1d7b74820b94799272d2516af62cf55fad47ca26 Mon Sep 17 00:00:00 2001 From: wychi Date: Thu, 7 May 2026 12:05:29 -0700 Subject: [PATCH 1/7] utlx: populate PluginInfo::tritonVersion (fixes plugin load on current Triton) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Triton commit 8497c845a (#9937) added isTritonAndPluginsVersionsMatch, which dereferences info->tritonVersion unconditionally at plugin load. uTLXPlugin's PluginInfo initializer omitted the field, leaving it nullptr — std::string construction throws and libtriton import dies with "basic_string::_M_construct null not valid" before any Python-side workaround can run. Fix by passing TRITON_VERSION (matches what support/Export.cpp already does). --- extensions/utlx/uTLXPlugin.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/extensions/utlx/uTLXPlugin.cpp b/extensions/utlx/uTLXPlugin.cpp index 5512d59..c9bf6c3 100644 --- a/extensions/utlx/uTLXPlugin.cpp +++ b/extensions/utlx/uTLXPlugin.cpp @@ -13,6 +13,7 @@ #include "triton/Dialect/TritonGPU/IR/Dialect.h" #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" #include "triton/Tools/PluginUtils.h" +#include "triton/Version.h" // TLX dialect headers #include "tlx/dialect/include/IR/Dialect.h" @@ -803,6 +804,11 @@ TRITON_PLUGIN_API plugin::PluginInfo *tritonGetPluginInfo() { 1, // numDialects ops, 48, // numOps + // Triton commit `8497c845a` (#9937) added `isTritonAndPluginsVersionsMatch` + // which dereferences `info->tritonVersion` unconditionally — leaving this + // field nullptr crashes libtriton's import with `basic_string::_M_construct + // null not valid` before any Python-side workaround can run. + TRITON_VERSION, }; return &info; } From 47debefa374fbc934aa42493e8bb472a77cac5d2 Mon Sep 17 00:00:00 2001 From: wychi Date: Thu, 7 May 2026 12:08:53 -0700 Subject: [PATCH 2/7] ci: add plugin-import smoke test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `make test` runs lit/FileCheck tests, which exercise MLIR passes but never load libutlx.so into Python. As a result, regressions in the plugin's static init path — fields read by libtriton's plugin loader (e.g. PluginInfo metadata) — slip through CI. The recent tritonVersion-nullptr regression rode through five Triton-pin bumps over a month before being noticed. Add a one-shot `python3 -c "import triton"` with TRITON_PLUGIN_PATHS pointed at the freshly-built libutlx.so. Cheap to run, requires no GPU, and would have caught the regression at the first pin bump that included the new version-check code in libtriton. --- .github/workflows/ci.yml | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c18fa1b..7b935e1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -112,6 +112,18 @@ jobs: TRITON_INSTALL_DIR="${{ steps.build-triton.outputs.triton_install_dir }}" EXTRA_CMAKE_ARGS="-DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache" + - name: Smoke test plugin import + # `make test` only runs lit/FileCheck tests, which exercise MLIR passes + # but never load the plugin into Python. Regressions in the plugin's + # static init (e.g. PluginInfo fields read by libtriton's loader) slip + # through unless we actually `import triton` with TRITON_PLUGIN_PATHS + # pointing at the freshly-built libutlx.so. + env: + TRITON_PLUGIN_PATHS: ${{ github.workspace }}/build/lib/libutlx.so + run: | + test -f "$TRITON_PLUGIN_PATHS" || { echo "libutlx.so missing at $TRITON_PLUGIN_PATHS"; exit 1; } + python3 -c "import triton; print('triton', triton.__version__, 'loaded with utlx plugin from', '$TRITON_PLUGIN_PATHS')" + - name: Run tests run: > make test LLVM_INSTALL_DIR="${{ steps.build-llvm.outputs.llvm_install_dir }}" From 44ba85b9c5728d15606338ab0143f7b7d749ff8f Mon Sep 17 00:00:00 2001 From: wychi Date: Thu, 7 May 2026 14:27:37 -0700 Subject: [PATCH 3/7] ci: convert plugin-import smoke test to pytest integration test Per code review on the prior commit: rather than a single inline shell command in CI hard-coded to libutlx.so, discover every plugin registered in this repo (every triton-ext.conf) and verify each loads. Wired into `make test` so the suite runs locally as well as in CI. testing/test_plugin_imports.py spawns a fresh interpreter per plugin and runs `import triton` with TRITON_PLUGIN_PATHS pointed at the .so. Each plugin runs in its own subprocess so failures isolate per plugin. This catches PluginInfo-level regressions on Triton builds that load plugins eagerly at import time -- the source-built pin used in CI does. Older release wheels load lazily and may pass even for a broken plugin; CI is the canonical environment for this signal. The Makefile splits `test` into `check-lit-tests` and `check-pytest-tests`; both run as part of `make test`. The standalone CI step is removed. --- .github/workflows/ci.yml | 12 ------ Makefile | 12 +++++- requirements.txt | 1 + testing/test_plugin_imports.py | 79 ++++++++++++++++++++++++++++++++++ 4 files changed, 91 insertions(+), 13 deletions(-) create mode 100644 testing/test_plugin_imports.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7b935e1..c18fa1b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -112,18 +112,6 @@ jobs: TRITON_INSTALL_DIR="${{ steps.build-triton.outputs.triton_install_dir }}" EXTRA_CMAKE_ARGS="-DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache" - - name: Smoke test plugin import - # `make test` only runs lit/FileCheck tests, which exercise MLIR passes - # but never load the plugin into Python. Regressions in the plugin's - # static init (e.g. PluginInfo fields read by libtriton's loader) slip - # through unless we actually `import triton` with TRITON_PLUGIN_PATHS - # pointing at the freshly-built libutlx.so. - env: - TRITON_PLUGIN_PATHS: ${{ github.workspace }}/build/lib/libutlx.so - run: | - test -f "$TRITON_PLUGIN_PATHS" || { echo "libutlx.so missing at $TRITON_PLUGIN_PATHS"; exit 1; } - python3 -c "import triton; print('triton', triton.__version__, 'loaded with utlx plugin from', '$TRITON_PLUGIN_PATHS')" - - name: Run tests run: > make test LLVM_INSTALL_DIR="${{ steps.build-llvm.outputs.llvm_install_dir }}" diff --git a/Makefile b/Makefile index 5d79d40..20b7876 100644 --- a/Makefile +++ b/Makefile @@ -20,9 +20,19 @@ build: configure cmake --build ${BUILD_DIR} .PHONY: test -test: +test: check-lit-tests check-pytest-tests + +.PHONY: check-lit-tests +check-lit-tests: ninja -C ${BUILD_DIR} check-lit-tests +# Plugin-import smoke tests: load each lib.so via TRITON_PLUGIN_PATHS and +# `import triton`. Catches regressions in the plugin's static-init path that +# lit/FileCheck tests can't see (e.g. PluginInfo fields read by libtriton's loader). +.PHONY: check-pytest-tests +check-pytest-tests: + TRITON_EXT_BUILD_DIR=${BUILD_DIR} python3 -m pytest testing/ + .PHONY: clean clean: rm -rf ${BUILD_DIR} diff --git a/requirements.txt b/requirements.txt index 9463836..0e49ab8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ lit~=18.1 ninja~=1.13 pre_commit~=4.6 +pytest~=8.3 requests~=2.33 diff --git a/testing/test_plugin_imports.py b/testing/test_plugin_imports.py new file mode 100644 index 0000000..9805000 --- /dev/null +++ b/testing/test_plugin_imports.py @@ -0,0 +1,79 @@ +"""Integration test: every plugin registered in this repo loads into Python. + +`make test` otherwise runs only lit/FileCheck tests, which exercise MLIR passes +but never load the plugin shared libraries into Python. That hid a regression +in the plugin's static init path (PluginInfo::tritonVersion left null) for five +Triton-pin bumps. + +For each `triton-ext.conf` in the source tree we resolve the corresponding +`lib.so` in the build dir and, in a fresh interpreter, run +`import triton` with `TRITON_PLUGIN_PATHS` pointed at the .so. Each plugin +runs in its own subprocess so a failure isolates to a single plugin. + +This catches PluginInfo-level regressions on Triton builds that load plugins +eagerly at import time — the source-built pin used in CI does. Older release +wheels load lazily and may pass even for a broken plugin; CI is the canonical +environment for this signal. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +BUILD_DIR = Path(os.environ.get("TRITON_EXT_BUILD_DIR", REPO_ROOT / "build")) +PLUGIN_LIB_DIR = BUILD_DIR / "lib" + + +def _discover_plugins() -> list[pytest.ParameterSet]: + plugins: list[pytest.ParameterSet] = [] + for conf in REPO_ROOT.rglob("triton-ext.conf"): + # Skip downloaded artifact trees (triton--..., llvm--...) and + # the build dir itself. + rel_parts = conf.relative_to(REPO_ROOT).parts + if rel_parts[0].startswith(("triton-", "llvm-", "build")): + continue + text = conf.read_text().strip() + if not text: + continue + # Format is `name;status[;hash]` (CMake list); we only need the name. + name = text.split(";", 1)[0].strip() + if not name: + continue + plugins.append(pytest.param(name, id=name)) + plugins.sort(key=lambda p: p.id) + return plugins + + +PLUGINS = _discover_plugins() + + +def test_plugins_discovered() -> None: + """Guard against silently testing nothing if discovery breaks.""" + assert PLUGINS, f"No triton-ext.conf files found under {REPO_ROOT}" + + +@pytest.mark.parametrize("name", PLUGINS) +def test_plugin_loads(name: str) -> None: + plugin_path = PLUGIN_LIB_DIR / f"lib{name}.so" + if not plugin_path.is_file(): + pytest.skip(f"Plugin not built at {plugin_path} (extension may be disabled)") + + env = {**os.environ, "TRITON_PLUGIN_PATHS": str(plugin_path)} + result = subprocess.run( + [sys.executable, "-c", "import triton"], + env=env, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, ( + f"Loading plugin {name} from {plugin_path} failed:\n" + f"--- stdout ---\n{result.stdout}\n" + f"--- stderr ---\n{result.stderr}" + ) From 6e9d25dbb0aa3a6f351cfe2ef7781ecb5a683ac1 Mon Sep 17 00:00:00 2001 From: wychi Date: Thu, 7 May 2026 14:51:43 -0700 Subject: [PATCH 4/7] utlx: drop doc comment from PluginInfo initializer Per review on the prior tritonVersion fix: the doc comment on the TRITON_VERSION field is overlong for what's effectively a one-liner historical note; the commit message and git blame already capture the context. --- extensions/utlx/uTLXPlugin.cpp | 4 ---- 1 file changed, 4 deletions(-) diff --git a/extensions/utlx/uTLXPlugin.cpp b/extensions/utlx/uTLXPlugin.cpp index c9bf6c3..72196bd 100644 --- a/extensions/utlx/uTLXPlugin.cpp +++ b/extensions/utlx/uTLXPlugin.cpp @@ -804,10 +804,6 @@ TRITON_PLUGIN_API plugin::PluginInfo *tritonGetPluginInfo() { 1, // numDialects ops, 48, // numOps - // Triton commit `8497c845a` (#9937) added `isTritonAndPluginsVersionsMatch` - // which dereferences `info->tritonVersion` unconditionally — leaving this - // field nullptr crashes libtriton's import with `basic_string::_M_construct - // null not valid` before any Python-side workaround can run. TRITON_VERSION, }; return &info; From 4772f44518c666fd74baf1bf16ca8e1ed04fca3a Mon Sep 17 00:00:00 2001 From: wychi Date: Thu, 7 May 2026 21:50:14 -0700 Subject: [PATCH 5/7] ci: apply yapf to test_plugin_imports.py Pre-commit's yapf hook prefers wrapping arguments with the closing paren on the same line as the last argument; reformat to match. --- testing/test_plugin_imports.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/testing/test_plugin_imports.py b/testing/test_plugin_imports.py index 9805000..ba3a10d 100644 --- a/testing/test_plugin_imports.py +++ b/testing/test_plugin_imports.py @@ -62,7 +62,8 @@ def test_plugins_discovered() -> None: def test_plugin_loads(name: str) -> None: plugin_path = PLUGIN_LIB_DIR / f"lib{name}.so" if not plugin_path.is_file(): - pytest.skip(f"Plugin not built at {plugin_path} (extension may be disabled)") + pytest.skip( + f"Plugin not built at {plugin_path} (extension may be disabled)") env = {**os.environ, "TRITON_PLUGIN_PATHS": str(plugin_path)} result = subprocess.run( @@ -75,5 +76,4 @@ def test_plugin_loads(name: str) -> None: assert result.returncode == 0, ( f"Loading plugin {name} from {plugin_path} failed:\n" f"--- stdout ---\n{result.stdout}\n" - f"--- stderr ---\n{result.stderr}" - ) + f"--- stderr ---\n{result.stderr}") From bc7cf40f30aa3245a4c89ab215835841fa055957 Mon Sep 17 00:00:00 2001 From: wychi Date: Thu, 7 May 2026 22:11:36 -0700 Subject: [PATCH 6/7] ci: source venv before running tests `make test` invokes `python3 -m pytest`, which needs the venv on PATH so that python3 resolves to the venv interpreter (where pytest is installed via requirements.txt). The earlier `Setup virtual environment` step propagates PATH via $GITHUB_ENV, but in practice python3 was resolving to the GitHub-hosted toolcache binary by the time `make test` ran -- so pytest came up as missing. Re-activating the venv right before invoking make is the robust spelling. (The previous inline `python3 -c "import triton"` smoke test happened to pass because Triton's `make dev-install` had landed in toolcache, so `import triton` worked even without the venv. Pytest is venv-only, so the same path resolution issue surfaces now.) --- .github/workflows/ci.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c18fa1b..a481c98 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -113,8 +113,9 @@ jobs: EXTRA_CMAKE_ARGS="-DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache" - name: Run tests - run: > - make test LLVM_INSTALL_DIR="${{ steps.build-llvm.outputs.llvm_install_dir }}" + run: | + source ~/.venv/bin/activate + make test LLVM_INSTALL_DIR="${{ steps.build-llvm.outputs.llvm_install_dir }}" \ TRITON_INSTALL_DIR="${{ steps.build-triton.outputs.triton_install_dir }}" - name: List files From bb214922613ee0789823d7a455a40dbc3f5a4700 Mon Sep 17 00:00:00 2001 From: wychi Date: Fri, 8 May 2026 17:08:41 -0700 Subject: [PATCH 7/7] ci: add plugin scenario tests using prebuilt triton from artifact Bundle python/triton/ in the triton artifact (suffix bumped to -v2 to invalidate cached v1 tarballs that lack the python tree). Add a dedicated "Run plugin tests" CI step that points PYTHONPATH at the artifact's python dir so pytest can `import triton` directly from the prebuilt tree -- no separate clone or pip install. The plugin test suite (testing/test_plugins.py, renamed from test_plugin_imports.py) auto-discovers plugins via triton-ext.conf and runs three categories per plugin: - test_plugin_loads[] -- static-init smoke (catches PluginInfo regressions) - test_plugin_compiles_kernel[] -- end-to-end JIT compile of a basic kernel through the plugin's pipeline - plugin-specific tests, e.g. test_utlx_registers_tlx_dsl, gated with @pytest.mark.skipif on the plugin .so existence Adding a plugin: drop a triton-ext.conf; both parametrized tests pick it up automatically. Exempting a plugin from a parametrized test: tag it via pytest.param(..., marks=pytest.mark.skip(...)) -- example dialect is exempted from the compile test (scaffolding-only Dialect::initialize() doesn't register StringAttr). Reverts the check-pytest-tests Makefile target added earlier on this branch -- the new CI step calls pytest directly. Run locally with: source ~/.venv-ci/bin/activate PYTHONPATH=~/oss/triton/python python -m pytest testing/test_plugins.py --- .github/actions/build-triton/action.yml | 9 +- .github/workflows/ci.yml | 11 +- Makefile | 12 +- testing/test_plugin_imports.py | 79 ---------- testing/test_plugins.py | 185 ++++++++++++++++++++++++ 5 files changed, 201 insertions(+), 95 deletions(-) delete mode 100644 testing/test_plugin_imports.py create mode 100644 testing/test_plugins.py diff --git a/.github/actions/build-triton/action.yml b/.github/actions/build-triton/action.yml index 58d2cc5..8f56446 100644 --- a/.github/actions/build-triton/action.yml +++ b/.github/actions/build-triton/action.yml @@ -35,7 +35,7 @@ runs: SHORT_TRITON_COMMIT_HASH="${TRITON_COMMIT_HASH:0:8}" SYSINFO=$(python ci/probe-sysinfo.py) - INSTALL_DIR="triton-${SHORT_TRITON_COMMIT_HASH}-${SYSINFO}" + INSTALL_DIR="triton-${SHORT_TRITON_COMMIT_HASH}-${SYSINFO}-v2" echo "Triton installation directory: ${INSTALL_DIR}" echo "triton_install_dir=${INSTALL_DIR}" >> ${GITHUB_ENV} echo "triton_install_dir=${INSTALL_DIR}" >> ${GITHUB_OUTPUT} @@ -93,8 +93,11 @@ runs: # Stripping debug symbols reduces the artifact size by an order of magnitude. run: | make install - find build/install -type f -executable | xargs strip --strip-debug - tar czf ../"${{ env.triton_install_dir }}.tar.gz" --transform="s|^build/install|${{ env.triton_install_dir }}|" build/install + find build/install python/triton -type f \( -executable -o -name "*.so" \) -print0 | xargs -0 -r strip --strip-debug || true + tar czf ../"${{ env.triton_install_dir }}.tar.gz" \ + --transform="s|^build/install|${{ env.triton_install_dir }}|" \ + --transform="s|^python|${{ env.triton_install_dir }}/python|" \ + build/install python/triton - name: Upload Triton artifact if: steps.check-artifact.outputs.exists == 'false' diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a481c98..fc704ae 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -113,10 +113,17 @@ jobs: EXTRA_CMAKE_ARGS="-DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache" - name: Run tests + run: > + make test LLVM_INSTALL_DIR="${{ steps.build-llvm.outputs.llvm_install_dir }}" + TRITON_INSTALL_DIR="${{ steps.build-triton.outputs.triton_install_dir }}" + + - name: Run plugin tests + env: + PYTHONPATH: ${{ github.workspace }}/${{ steps.build-triton.outputs.triton_install_dir }}/python + LD_LIBRARY_PATH: ${{ github.workspace }}/${{ steps.build-llvm.outputs.llvm_install_dir }}/lib:${{ github.workspace }}/${{ steps.build-triton.outputs.triton_install_dir }}/lib run: | source ~/.venv/bin/activate - make test LLVM_INSTALL_DIR="${{ steps.build-llvm.outputs.llvm_install_dir }}" \ - TRITON_INSTALL_DIR="${{ steps.build-triton.outputs.triton_install_dir }}" + python -m pytest testing/test_plugins.py - name: List files run: du -d 3 -h . diff --git a/Makefile b/Makefile index 20b7876..5d79d40 100644 --- a/Makefile +++ b/Makefile @@ -20,19 +20,9 @@ build: configure cmake --build ${BUILD_DIR} .PHONY: test -test: check-lit-tests check-pytest-tests - -.PHONY: check-lit-tests -check-lit-tests: +test: ninja -C ${BUILD_DIR} check-lit-tests -# Plugin-import smoke tests: load each lib.so via TRITON_PLUGIN_PATHS and -# `import triton`. Catches regressions in the plugin's static-init path that -# lit/FileCheck tests can't see (e.g. PluginInfo fields read by libtriton's loader). -.PHONY: check-pytest-tests -check-pytest-tests: - TRITON_EXT_BUILD_DIR=${BUILD_DIR} python3 -m pytest testing/ - .PHONY: clean clean: rm -rf ${BUILD_DIR} diff --git a/testing/test_plugin_imports.py b/testing/test_plugin_imports.py deleted file mode 100644 index ba3a10d..0000000 --- a/testing/test_plugin_imports.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Integration test: every plugin registered in this repo loads into Python. - -`make test` otherwise runs only lit/FileCheck tests, which exercise MLIR passes -but never load the plugin shared libraries into Python. That hid a regression -in the plugin's static init path (PluginInfo::tritonVersion left null) for five -Triton-pin bumps. - -For each `triton-ext.conf` in the source tree we resolve the corresponding -`lib.so` in the build dir and, in a fresh interpreter, run -`import triton` with `TRITON_PLUGIN_PATHS` pointed at the .so. Each plugin -runs in its own subprocess so a failure isolates to a single plugin. - -This catches PluginInfo-level regressions on Triton builds that load plugins -eagerly at import time — the source-built pin used in CI does. Older release -wheels load lazily and may pass even for a broken plugin; CI is the canonical -environment for this signal. -""" - -from __future__ import annotations - -import os -import subprocess -import sys -from pathlib import Path - -import pytest - -REPO_ROOT = Path(__file__).resolve().parent.parent -BUILD_DIR = Path(os.environ.get("TRITON_EXT_BUILD_DIR", REPO_ROOT / "build")) -PLUGIN_LIB_DIR = BUILD_DIR / "lib" - - -def _discover_plugins() -> list[pytest.ParameterSet]: - plugins: list[pytest.ParameterSet] = [] - for conf in REPO_ROOT.rglob("triton-ext.conf"): - # Skip downloaded artifact trees (triton--..., llvm--...) and - # the build dir itself. - rel_parts = conf.relative_to(REPO_ROOT).parts - if rel_parts[0].startswith(("triton-", "llvm-", "build")): - continue - text = conf.read_text().strip() - if not text: - continue - # Format is `name;status[;hash]` (CMake list); we only need the name. - name = text.split(";", 1)[0].strip() - if not name: - continue - plugins.append(pytest.param(name, id=name)) - plugins.sort(key=lambda p: p.id) - return plugins - - -PLUGINS = _discover_plugins() - - -def test_plugins_discovered() -> None: - """Guard against silently testing nothing if discovery breaks.""" - assert PLUGINS, f"No triton-ext.conf files found under {REPO_ROOT}" - - -@pytest.mark.parametrize("name", PLUGINS) -def test_plugin_loads(name: str) -> None: - plugin_path = PLUGIN_LIB_DIR / f"lib{name}.so" - if not plugin_path.is_file(): - pytest.skip( - f"Plugin not built at {plugin_path} (extension may be disabled)") - - env = {**os.environ, "TRITON_PLUGIN_PATHS": str(plugin_path)} - result = subprocess.run( - [sys.executable, "-c", "import triton"], - env=env, - capture_output=True, - text=True, - check=False, - ) - assert result.returncode == 0, ( - f"Loading plugin {name} from {plugin_path} failed:\n" - f"--- stdout ---\n{result.stdout}\n" - f"--- stderr ---\n{result.stderr}") diff --git a/testing/test_plugins.py b/testing/test_plugins.py new file mode 100644 index 0000000..db9d9ba --- /dev/null +++ b/testing/test_plugins.py @@ -0,0 +1,185 @@ +"""Plugin integration tests. + +Auto-discovers every plugin declared by a `triton-ext.conf` and exercises +its `lib.so` from a fresh Python interpreter with `TRITON_PLUGIN_PATHS` +set. Each plugin runs in its own subprocess so failures isolate cleanly. + +Tests: + - test_plugin_loads[] -- plugin static-init: `import triton` + succeeds with the .so loaded. + - test_plugin_compiles_kernel[] -- end-to-end: JIT-decorate and lower + a basic kernel through the plugin's + pipeline. + - test__ -- plugin-specific scenarios, gated + with `@pytest.mark.skipif` on the + plugin's .so existence. + +Adding a new plugin: drop a `triton-ext.conf`; both parametrized tests pick +it up. To exempt a plugin from a parametrized test, mark it at parametrize +time with `pytest.param(..., marks=pytest.mark.skip(...))` -- see +`_COMPILE_PLUGINS` for an example. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +BUILD_DIR = Path(os.environ.get("TRITON_EXT_BUILD_DIR", REPO_ROOT / "build")) +PLUGIN_LIB_DIR = BUILD_DIR / "lib" + + +def _discover_plugins() -> list[pytest.ParameterSet]: + plugins: list[pytest.ParameterSet] = [] + for conf in REPO_ROOT.rglob("triton-ext.conf"): + rel_parts = conf.relative_to(REPO_ROOT).parts + if rel_parts[0].startswith(("triton-", "llvm-", "build")): + continue + text = conf.read_text().strip() + if not text: + continue + # Format is `name;status[;hash]` (CMake list); we only need the name. + name = text.split(";", 1)[0].strip() + if not name: + continue + plugins.append(pytest.param(name, id=name)) + plugins.sort(key=lambda p: p.id) + return plugins + + +PLUGINS = _discover_plugins() + + +def _plugin_path(name: str) -> Path: + return PLUGIN_LIB_DIR / f"lib{name}.so" + + +def _run_with_plugin(plugin_path: Path, script: str) -> subprocess.CompletedProcess: + env = {**os.environ, "TRITON_PLUGIN_PATHS": str(plugin_path)} + return subprocess.run( + [sys.executable, "-c", textwrap.dedent(script)], + env=env, + capture_output=True, + text=True, + check=False, + ) + + +# --------------------------------------------------------------------------- +# Generic per-plugin tests (auto-discovered) +# --------------------------------------------------------------------------- + + +def test_plugins_discovered() -> None: + """Guard against silently testing nothing if discovery breaks.""" + assert PLUGINS, f"No triton-ext.conf files found under {REPO_ROOT}" + + +@pytest.mark.parametrize("name", PLUGINS) +def test_plugin_loads(name: str) -> None: + """Smoke: `import triton` succeeds with the plugin loaded.""" + path = _plugin_path(name) + if not path.is_file(): + pytest.skip(f"Plugin not built at {path} (extension may be disabled)") + result = _run_with_plugin(path, "import triton") + assert result.returncode == 0, ( + f"Loading plugin {name} from {path} failed:\n" + f"--- stdout ---\n{result.stdout}\n" + f"--- stderr ---\n{result.stderr}" + ) + + +# example dialect is scaffolding-only -- its Dialect::initialize() doesn't +# register StringAttr, so kernel compile aborts with an LLVM storage-uniquer +# error. Tag it as skip at parametrize time. +_COMPILE_PLUGINS = [ + pytest.param( + p.values[0], marks=pytest.mark.skip(reason="scaffolding-only dialect"), id=p.id + ) + if p.id == "example" + else p + for p in PLUGINS +] + + +@pytest.mark.parametrize("name", _COMPILE_PLUGINS) +def test_plugin_compiles_kernel(name: str) -> None: + """User scenario: with the plugin loaded, JIT-decorate and lower a basic kernel.""" + path = _plugin_path(name) + if not path.is_file(): + pytest.skip(f"Plugin not built at {path} (extension may be disabled)") + script = """ + import sys + import triton + import triton.language as tl + + @triton.jit + def kernel(in_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + tl.store(out_ptr + offs, tl.load(in_ptr + offs)) + + try: + target = triton.runtime.driver.active.get_current_target() + except Exception as e: + print(f"No target ({type(e).__name__}: {e}); skipping compile.") + sys.exit(0) + src = triton.compiler.ASTSource( + fn=kernel, + signature={"in_ptr": "*fp32", "out_ptr": "*fp32"}, + constexprs={"BLOCK": 128}, + ) + triton.compile(src, target=target) + """ + result = _run_with_plugin(path, script) + assert result.returncode == 0, ( + f"Plugin {name} broke kernel compile:\n" + f"--- stdout ---\n{result.stdout}\n" + f"--- stderr ---\n{result.stderr}" + ) + + +# --------------------------------------------------------------------------- +# Plugin-specific tests +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not _plugin_path("utlx").is_file(), reason="utlx plugin not built") +def test_utlx_registers_tlx_dsl() -> None: + """utlx registers `triton.language.extra.tlx` with local_alloc/view/store/load. + + The Python namespace is set up by `extensions/utlx/python/utlx_plugin/__init__.py` + when imported -- it inserts itself into `sys.modules` as + `triton.language.extra.tlx`. Loading the .so alone is not enough. + """ + plugin_path = _plugin_path("utlx") + utlx_python = REPO_ROOT / "extensions" / "utlx" / "python" + env = { + **os.environ, + "TRITON_PLUGIN_PATHS": str(plugin_path), + "PYTHONPATH": f"{utlx_python}{os.pathsep}{os.environ.get('PYTHONPATH', '')}", + } + script = """ + import triton # noqa: F401 + import utlx_plugin # noqa: F401 (registers triton.language.extra.tlx) + from triton.language.extra import tlx + for n in ("local_alloc", "local_view", "local_store", "local_load"): + assert hasattr(tlx, n), f"missing tlx.{n}" + """ + result = subprocess.run( + [sys.executable, "-c", textwrap.dedent(script)], + env=env, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, ( + f"utlx tlx-DSL check failed:\n" + f"--- stdout ---\n{result.stdout}\n" + f"--- stderr ---\n{result.stderr}" + )