diff --git a/.github/scripts/select_triton_tests.py b/.github/scripts/select_triton_tests.py index 449cf561f05..2f5b5aad3f6 100755 --- a/.github/scripts/select_triton_tests.py +++ b/.github/scripts/select_triton_tests.py @@ -1,919 +1,424 @@ #!/usr/bin/env python +"""Select Triton unit tests to run for a PR, from its git diff. +.github/workflows/triton-test.yaml runs this as a dry run (--dry-run) while its +DRY_RUN switch is 'true': the selection is printed, and CI runs what it ran +before test selection existed. The ci:triton-355 PR label runs the full suite +(--all). Any error falls back to the full suite. +""" -# Imports -# ------------------------------------------------------------------------------ - -# Python standard library. import argparse import ast -import functools -import logging import os import re -import shlex import subprocess import sys -import time +from collections import Counter +from functools import cache from pathlib import Path -# Third party libraries. -import networkx as nx - -logger = logging.getLogger(__name__) - -# Small utility functions. -# ------------------------------------------------------------------------------ - - -def log_file_list(log_level: int, files: list[Path]) -> None: - if not logging.getLogger().isEnabledFor(log_level): - return - for f in files: - logger.log(log_level, "* %s", f) - - -# Structure of Triton source files. -# ------------------------------------------------------------------------------ - - -def check_dir(p: Path) -> Path: - if not p.exists(): - logger.critical("Required directory [%s] doesn't exist.", p) - sys.exit(1) - if not p.is_dir(): - logger.critical("Required directory [%s] isn't a directory.", p) - sys.exit(1) - return p - - -@functools.cache -def root_dir() -> Path: - return check_dir(Path(__file__).parent.parent.parent) - - -@functools.cache -def triton_op_dir() -> Path: - return check_dir(root_dir() / "aiter" / "ops" / "triton") - - -@functools.cache -def triton_config_dir() -> Path: - return check_dir(triton_op_dir() / "configs") - - -@functools.cache -def triton_gemm_config_dir() -> Path: - return check_dir(triton_config_dir() / "gemm") - - -def list_files(dir: Path, suffix: str = "") -> set[Path]: - return {p.relative_to(root_dir()) for p in dir.glob(f"**/*{suffix}") if p.is_file()} - - -def list_triton_op_files() -> set[Path]: - files = list_files(triton_op_dir(), suffix=".py") - logger.debug("Found %d Triton operator source files.", len(files)) - return files - - -def list_triton_kernel_files(kernel_dir: Path) -> set[Path]: - files = list_files(kernel_dir, suffix=".py") - logger.debug("Found %d Triton kernel source files.", len(files)) - return files - - -def list_triton_config_files() -> set[Path]: - files = list_files(triton_config_dir(), suffix=".json") - logger.debug("Found %d Triton kernel config files.", len(files)) - return files - - -def list_triton_test_files(test_dir: Path) -> set[Path]: - files = list_files(test_dir, suffix=".py") - logger.debug("Found %d Triton test source files.", len(files)) - return files - - -def list_triton_bench_files(bench_dir: Path) -> set[Path]: - files = list_files(bench_dir, suffix=".py") - logger.debug("Found %d Triton benchmark source files.", len(files)) - return files - - -def list_triton_source_files() -> ( - tuple[set[Path], list[Path], list[Path], list[Path], list[Path]] -): - kernels_dir = check_dir(triton_op_dir() / "_triton_kernels") - op_test_dir = check_dir(root_dir() / "op_tests") - test_dir = check_dir(op_test_dir / "triton_tests") - bench_dir = check_dir(op_test_dir / "op_benchmarks" / "triton") - op_files = list_triton_op_files() - kernel_files = list_triton_kernel_files(kernels_dir) - config_files = list_triton_config_files() - test_files = list_triton_test_files(test_dir) - bench_files = list_triton_bench_files(bench_dir) - all_files = op_files | kernel_files | config_files | test_files | bench_files - return ( - all_files, - sorted(kernel_files), - sorted(config_files), - sorted(test_files), - sorted(bench_files), - ) - - -# Matching of kernel config files. -# ------------------------------------------------------------------------------ - - -DEVICES: frozenset[str] = frozenset(["gfx942", "gfx950"]) - -# Single letter placeholders for JSON config file template strings. You can add -# more letters if necessary. -MNK_PLACEHOLDERS: str = "MNK" -NUM_PATTERN: str = r"\d+" +SRC = "aiter/ops/triton/" +KERNELS = SRC + "_triton_kernels/" +GLUON_KERNELS = SRC + "_gluon_kernels/" +CONFIGS = SRC + "configs/" +TESTS = "op_tests/triton_tests/" +BENCH = "op_tests/op_benchmarks/triton/" +ROOT = Path(__file__).resolve().parents[2] +IMPORT_ROOTS = ("aiter.ops.triton", "op_tests.triton_tests") +CI_FILES = { + ".github/scripts/build_aiter_triton.sh", + ".github/scripts/download_triton_wheel.sh", + ".github/scripts/install_triton.sh", + ".github/scripts/select_triton_tests.py", + ".github/scripts/split_tests.sh", + ".github/scripts/verify_triton_pin.py", + ".github/scripts/ensure_zstd.sh", + ".github/requirements/triton-test.txt", + ".github/workflows/triton-test.yaml", + ".github/workflows/prepare-triton-wheel.yaml", + ".github/workflows/ci-config.yaml", +} + +# A change here can affect anything: run the full suite. +GLOBAL_PREFIXES = ( + ".github/", + SRC + "utils/", + KERNELS + "common/", + GLUON_KERNELS + "common/", + TESTS + "utils/", +) -MOE_DTYPES: tuple[str, ...] = ( - "DEFAULT", - "FP8_W8A8", - "INT8_W8A16", - "INT8_W8A8", - "INT4_W4A16", - "MX_FP4", +# Directories under the source tree that are not op categories. +NON_CATEGORY_DIRS = {"utils", "configs", "_triton_kernels", "_gluon_kernels"} + +# Config categories whose source/test folders use a different layout. +CONFIG_CATEGORIES = { + "attention": {"attention", "chunk_delta_attn"}, + "mhc": {"fusions"}, +} + +# triton-test.yaml's `paths:` filter before test selection. A dry run runs the +# full suite exactly when this matches: when the workflow used to run at all. +LEGACY_PATHS = ( + "aiter/ops/triton/**", + "op_tests/triton_tests/**", + "op_tests/op_benchmarks/triton/**", + ".github/scripts/build_aiter_triton.sh", + ".github/scripts/download_triton_wheel.sh", + ".github/scripts/install_triton.sh", + ".github/scripts/select_triton_tests.py", + ".github/scripts/split_tests.sh", + ".github/scripts/verify_triton_pin.py", + ".github/requirements/triton-test.txt", + "!**/*.md", + "!docs/**", + "!LICENSE", + "!.gitignore", + "!.github/workflows/**", + "!.github/scripts/sglang_downstream.py", + ".github/workflows/triton-test.yaml", + ".github/workflows/prepare-triton-wheel.yaml", + ".github/workflows/ci-config.yaml", ) -def expand_mnk(json_string: str, config_files: list[Path]) -> list[str]: - """ - Expands template strings containing M/N/K placeholders by matching them - against actual config file paths. - - :param json_string: Template string that references kernel config JSON files. - May contain M/N/K placeholders to be expanded. - :param config_files: All Triton kernel config JSON files available in the - filesystem. - :return: List of template strings with expanded M/N/K placeholders or the - a list containing just the input template string as-is if there - are no M/N/K placeholders to be expanded. - """ - # Early exit if no placeholders present. - if not any(f"{{{p}}}" in json_string for p in MNK_PLACEHOLDERS): - logger.debug("No M/N/K placeholders in [%s].", json_string) - return [json_string] - # Strip `f"` prefix and `"` suffix, then escape special regex characters. For - # example, `f"gfx950-GEMM-N={N}-K={K}.json"` becomes `gfx950-GEMM-N=\{N\}-K=\{K\}.json`. - pattern_str = re.escape(json_string[2:-1]) - # Replace each escaped placeholder with a named capture regex group. - for p in MNK_PLACEHOLDERS: # `p` will be `M`, then `N`, then `K`. - # For instance, if `p` is `M` then `\{M\}` will be replaced by `(?P\d+)`. - # Following the example: - # 1. there's no M, pattern is not changed. - # 2. `gfx950-GEMM-N=\{N\}-K=\{K\}.json` becomes `gfx950-GEMM-N=(?P\d+)-K=\{K\}.json` - # 3. `gfx950-GEMM-N=(?P\d+)-K=\{K\}.json` becomes `gfx950-GEMM-N=(?P\d+)-K=(?P\d+).json` - pattern_str = pattern_str.replace( - rf"\{{{p}\}}", - rf"(?P<{p}>{NUM_PATTERN})", - ) - # Compile regex with anchors to match entire path. - pattern = re.compile(f"^{pattern_str}$") - logger.debug("M/N/K regex is [%s].", pattern.pattern) - # Return f-string representations of matching config paths. - return [f'f"{path}"' for c in config_files if pattern.match(path := c.as_posix())] - - -def expand_moe_dtypes(json_strings: list[str]) -> list[str]: - expanded_moe_dtypes: list[str] = [] - for s in json_strings: - if r"MOE-{dtype_str}" in s: - expanded_moe_dtypes.extend( - s.replace(r"{dtype_str}", dtype) for dtype in MOE_DTYPES - ) - else: - expanded_moe_dtypes.append(s) - return expanded_moe_dtypes - - -def expand_interpolations(json_string: str, config_files: list[Path]) -> list[str]: - if not (json_string.startswith(("f'", 'f"'))): - return [json_string] - # Replace config path placeholder - if r"{AITER_TRITON_CONFIGS_PATH}" in json_string: - json_string = json_string.replace( - r"{AITER_TRITON_CONFIGS_PATH}", - str(triton_config_dir().relative_to(root_dir()).as_posix()), +def legacy_trigger(diff): + """Whether LEGACY_PATHS matches a changed file, with GitHub's rules: `**` + crosses directories, `*` does not, and the last matching pattern wins.""" + glob = {"**/": "(?:.*/)?", "**": ".*", "*": "[^/]*", "?": "[^/]"} + rules = [ + ( + not pattern.startswith("!"), + re.compile( + "".join( + glob.get(token, re.escape(token)) + for token in re.split(r"(\*\*/|\*\*|\*|\?)", pattern.lstrip("!")) + ) + ), ) - logger.debug("Resolved {AITER_TRITON_CONFIGS_PATH}: [%s]", json_string) - # Expand device variants - expanded = [json_string] - if r"{dev}" in json_string: - expanded = [s.replace(r"{dev}", dev) for s in expanded for dev in DEVICES] - logger.debug("Resolved {dev}: %s", expanded) - # Expand GEMM M-N-K patterns - expanded = [ - expanded_mnk for s in expanded for expanded_mnk in expand_mnk(s, config_files) + for pattern in LEGACY_PATHS ] - # Expand MOE data type variants - expanded = expand_moe_dtypes(expanded) - # Clean up f-string delimiters if no more interpolation needed - expanded = [s[2:-1] if not any(c in s for c in "{}") else s for s in expanded] - return expanded - - -def resolve_path(json_string: str) -> Path | None: - p = Path(json_string) - # Try absolute path - if p.is_absolute() and p.exists() and p.is_file(): - return p.relative_to(root_dir()) - # Try relative to root - p = root_dir() / json_string - if p.exists() and p.is_file(): - return p.relative_to(root_dir()) - return None + def included(path): + keep = False + for include, regex in rules: + if regex.fullmatch(path): + keep = include + return keep -def resolve_json_strings( - json_strings: list[str], config_files: list[Path] -) -> list[Path]: - # Expand all interpolations first - expanded_strings = [ - expanded - for json_string in json_strings - for expanded in expand_interpolations(json_string, config_files) - ] - # Sort and deduplicate - expanded_strings = sorted(set(expanded_strings)) - # Resolve to actual paths - resolved: list[Path] = [] - unresolved: list[str] = [] - for json_string in expanded_strings: - if resolved_path := resolve_path(json_string): - resolved.append(resolved_path) - else: - unresolved.append(json_string) - # Log results - if resolved: - logger.debug("Resolved JSON config files:") - log_file_list(logging.DEBUG, resolved) - log_level = logging.DEBUG - if unresolved and logging.getLogger().isEnabledFor(log_level): - logger.log(log_level, "Unresolved JSON strings:") - for s in unresolved: - logger.log(log_level, "* %s", s) - return resolved - - -def resolve_gemm_config_names(gemm_config_names: list[str]) -> list[Path]: - gemm_configs = [ - p.relative_to(root_dir()) - for dev in DEVICES - for gemm_config_name in gemm_config_names - for p in triton_gemm_config_dir().glob(f"{dev}-{gemm_config_name}*.json") - if p.is_file() - ] - if gemm_configs: - logger.debug("Resolved JSON GEMM config files:") - log_file_list(logging.DEBUG, gemm_configs) - return gemm_configs - - -def resolve_configs( - json_strings: list[str], - gemm_config_names: list[str], - config_files: list[Path], -) -> list[Path]: - return sorted( - set(resolve_json_strings(json_strings, config_files)).union( - resolve_gemm_config_names(gemm_config_names) - ) - ) + return any(included(path) for path in diff) -# Git commands. -# ------------------------------------------------------------------------------ +def list_files(base, pattern): + return sorted(p.relative_to(ROOT).as_posix() for p in (ROOT / base).rglob(pattern)) -def git(args: str, check: bool = True) -> subprocess.CompletedProcess: - try: - return subprocess.run( - ["git"] + shlex.split(args), - capture_output=True, - text=True, - check=check, - cwd=root_dir(), # always run git commands from repo root - ) - except FileNotFoundError: - logger.critical("Git not found.") - sys.exit(1) - except subprocess.CalledProcessError: - logger.critical("Malformed Git command: [git %s].", args) - sys.exit(1) - - -def git_current_branch() -> str: - return git("rev-parse --abbrev-ref HEAD").stdout.rstrip() - - -def git_check_branch(branch: str) -> None: - if git(f"rev-parse --verify --quiet {branch}", check=False).returncode != 0: - logger.critical("Branch [%s] doesn't exist.", branch) - sys.exit(1) - - -def git_filename_diff(source_branch: str, target_branch: str) -> set[Path]: - diff_output = git( - f"diff --name-only {target_branch} {source_branch}" - ).stdout.rstrip() - if not diff_output: - return set() - files = set() - for diff_p in diff_output.splitlines(): - abs_path = root_dir() / diff_p - if abs_path.exists() and abs_path.is_file(): - # Add path relative to root_dir() - this will match paths from `list_files` function. - files.add(abs_path.relative_to(root_dir())) - logger.debug( - "There %s %d file%s in the diff from [%s] to [%s].", - "is" if len(files) == 1 else "are", - len(files), - "" if len(files) == 1 else "s", - source_branch, - target_branch, - ) - return files +def category_of(path): + """Op category a source or test file belongs to, or None.""" + for base in (KERNELS, GLUON_KERNELS, SRC, TESTS): + if not path.startswith(base): + continue + parts = path[len(base) :].split("/") + # Gluon is a backend, not an op: _gluon_kernels///... + if base == GLUON_KERNELS and parts[0].startswith("gfx"): + parts = parts[1:] + if len(parts) < 2 or parts[0] in NON_CATEGORY_DIRS: + return None + return parts[0] + return None -def get_filename_diff(source_branch: str | None, target_branch: str) -> set[Path]: - if source_branch is None: - source_branch = git_current_branch() - logger.info( - "Source branch wasn't provided, using current branch [%s] as source branch.", - source_branch, - ) - else: - git_check_branch(source_branch) - - git_check_branch(target_branch) - - if target_branch == source_branch: - logger.error("Source and target branches must be different.") - sys.exit(1) - - return git_filename_diff(source_branch, target_branch) - - -# Source file parsing. -# ------------------------------------------------------------------------------ - - -class Visitor(ast.NodeVisitor): - imports_to_ignore: frozenset[str] = frozenset( - [ - "einops", - "iris", - "jax", - "jinja2", - "matplotlib", - "mori", - "numpy", - "packaging", - "pandas", - "prettytable", - "psutil", - "pybind11", - "pytest", - "setuptools", - "torch", - "triton", - "vllm", - "yaml", - "zmq", - ] - ) - json_strings_to_ignore: frozenset[str] = frozenset( - [ - "./utils/model_configs.json", - ".json", - "empty_kernel.json", - "f'{kernel_name}.json'", - "utils/model_configs.json", - ] - ) +@cache +def compatibility_modules(): + """Read legacy module aliases without importing GPU code.""" + path = ROOT / (SRC + "__init__.py") + if path.is_file(): + for node in ast.parse(path.read_text(encoding="utf-8")).body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == "_BACKWARD_COMPAT_MAP" + for target in node.targets + ): + return { + f"aiter.ops.triton.{old}": f"aiter.ops.triton.{new}" + for old, new in ast.literal_eval(node.value).items() + } + return {} + + +@cache +def resolve_module(dotted, gone): + """Resolve local modules, retaining edges to deleted files.""" + if not any( + dotted == root or dotted.startswith(root + ".") for root in IMPORT_ROOTS + ): + return None + rel = compatibility_modules().get(dotted, dotted).replace(".", "/") + for cand in (rel + ".py", rel + "/__init__.py"): + if cand in gone or (ROOT / cand).is_file(): + return cand + return None - @classmethod - def is_import_of_interest(cls, import_: str) -> bool: - return ( - import_ not in sys.stdlib_module_names - and import_ not in cls.imports_to_ignore - and not any( - import_.startswith(module + ".") for module in cls.imports_to_ignore + +def scan_imports(path, gone, dynamic): + """Collect local imports and flag loaders with unknown runtime targets.""" + found = set() + tree = ast.parse((ROOT / path).read_text(encoding="utf-8")) + # Resolve constants such as PUBLIC_WRAPPER or a local legacy import path. + writes = Counter( + node.id + for node in ast.walk(tree) + if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store) + ) + writes.update(node.arg for node in ast.walk(tree) if isinstance(node, ast.arg)) + constants = { + target.id: node.value.value + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and isinstance(node.value, ast.Constant) + and isinstance(node.value.value, str) + for target in node.targets + if isinstance(target, ast.Name) and writes[target.id] == 1 + } + loaders = { + name: name + for name in ("import_module", "__import__", "spec_from_file_location") + } + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + loaders.update( + (a.asname or a.name, loaders[a.name]) + for a in node.names + if a.name in loaders ) - and not any( - import_.startswith(module + ".") for module in sys.stdlib_module_names + for node in ast.walk(tree): + if isinstance(node, ast.Import): + names = [a.name for a in node.names] + elif isinstance(node, ast.ImportFrom): + module = node.module or "" + if node.level: + package = path.split("/")[:-1] + if node.level > len(package): + raise RuntimeError(f"{path}: relative import escapes its package") + package = package[: len(package) - node.level + 1] + module = ".".join(package + ([module] if module else [])) + # Imported names can themselves be submodules. + names = [f"{module}.{a.name}" for a in node.names] + if not all(name in compatibility_modules() for name in names): + names.append(module) + elif isinstance(node, ast.Call): + name = getattr(node.func, "id", getattr(node.func, "attr", None)) + if name not in loaders: + continue + arg = node.args[0] if node.args else None + target = ( + arg.value + if isinstance(arg, ast.Constant) + else constants.get(arg.id) if isinstance(arg, ast.Name) else None ) - ) - - @classmethod - def is_json_string_of_interest(cls, json_string: str) -> bool: - return json_string not in cls.json_strings_to_ignore - - def __init__(self, source_file: Path) -> None: - # Remove extension from source file, and split directories into module parts. - self.source_file: Path = source_file - self.source_module_parts: tuple[str, ...] = self.source_file.with_suffix( - "" - ).parts - self.dependencies: set[Path] = set() - self.json_strings: set[str] = set() - self.gemm_config_names: set[str] = set() - - def add_dependency(self, import_: str, suppress_warning: bool = False) -> None: - if not import_ or not self.__class__.is_import_of_interest(import_): - return - import_py_file = import_.replace(".", os.sep) + ".py" - # Check if it's a module file (e.g., `foo/bar.py`). - p = root_dir() / import_py_file - if p.exists() and p.is_file(): - self.dependencies.add(p.relative_to(root_dir())) - return - # Check if it's a package directory with `__init__.py`. - p_dir = (root_dir() / import_py_file).with_suffix("") - p_init = p_dir / "__init__.py" - if p_init.exists() and p_init.is_file(): - # Add dependency as the `__init__.py` file of the package. - self.dependencies.add(p_init.relative_to(root_dir())) - return - # Check for relative imports within the same package. - p = root_dir() / self.source_file.parent / import_py_file - if p.exists() and p.is_file(): - self.dependencies.add(p.relative_to(root_dir())) - return - # Check for relative package imports. - p_dir = (root_dir() / self.source_file.parent / import_py_file).with_suffix("") - p_init = p_dir / "__init__.py" - if p_init.exists() and p_init.is_file(): - self.dependencies.add(p_init.relative_to(root_dir())) - return - if not suppress_warning: - # Check if it's a directory without `__init__.py`` (namespace package or external) - p_dir = (root_dir() / import_py_file).with_suffix("") - if p_dir.exists() and p_dir.is_dir(): - logger.debug( - "Directory [%s] exists but has no '__init__.py', skipping dependency [%s] of [%s].", - p_dir.relative_to(root_dir()), - import_, - self.source_file, - ) - else: - logger.warning( - "Unable to find [%s] dependency of [%s] on filesystem.", - import_, - self.source_file, + if ( + loaders[name] == "spec_from_file_location" + or not isinstance(target, str) + or target.startswith(".") + or ( + target.startswith(IMPORT_ROOTS) + and resolve_module(target, gone) is None ) - - def add_json_string(self, json_string: str) -> None: - if not json_string or not self.__class__.is_json_string_of_interest( - json_string - ): - return - self.json_strings.add(json_string) - - def visit_Import(self, node: ast.Import) -> None: - for alias in node.names: - self.add_dependency(alias.name) - self.generic_visit(node) - - def visit_ImportFrom(self, node: ast.ImportFrom) -> None: - module_name = node.module if node.module else "" - if node.level > 0: - # Resolve absolute import from relative import. - full_module_name = ".".join(self.source_module_parts[: -node.level]) - if module_name: - full_module_name += "." + module_name - else: - full_module_name = module_name - self.add_dependency(full_module_name) - # For "from X import Y", also try to add Y as a potential module - # (in case Y is a submodule rather than a name defined in X). - if full_module_name and node.names[0].name != "*": - for alias in node.names: - potential_submodule = f"{full_module_name}.{alias.name}" - self.add_dependency(potential_submodule, suppress_warning=True) - self.generic_visit(node) - - def visit_Constant(self, node: ast.Constant) -> None: - # Plain string literals. - if isinstance(node.value, str) and node.value.lower().endswith(".json"): - self.add_json_string(node.value) - - def visit_JoinedStr(self, node: ast.JoinedStr) -> None: - # f-strings. - parts = [] - for value in node.values: - if isinstance(value, ast.Constant) and isinstance(value.value, str): - parts.append(value.value) - elif isinstance(value, ast.FormattedValue): - # Unparse the inner expression to a readable form. - expr_str = ast.unparse(value.value) - parts.append(f"{{{expr_str}}}") - joined = "".join(parts) - if joined.lower().endswith(".json"): - self.add_json_string(f"f{joined!r}") - - def visit_Call(self, node: ast.Call) -> None: - func = node.func - is_get_gemm_config = ( - isinstance(func, ast.Name) - and func.id == "get_gemm_config" - or isinstance(func, ast.Attribute) - and func.attr == "get_gemm_config" - ) - if is_get_gemm_config and node.args: - gemm_config_name = node.args[0] - if isinstance(gemm_config_name, ast.Constant) and isinstance( - gemm_config_name.value, str ): - self.gemm_config_names.add(gemm_config_name.value) - self.generic_visit(node) - - -def parse_source_file(source_file: Path) -> tuple[list[Path], list[str], list[str]]: - try: - logger.debug("Parsing source file [%s]...", str(source_file)) - source = (root_dir() / source_file).read_text(encoding="utf-8") - tree = ast.parse(source, filename=str(source_file)) - except Exception: - logger.exception("Skipping source file [%s].", source_file) - return [], [], [] - - visitor = Visitor(source_file) - visitor.visit(tree) - - dependecies = sorted(visitor.dependencies) - if not dependecies: - logger.debug("No dependecies of interest in [%s].", source_file) - else: - logger.debug("Dependecies of interest in [%s]:", source_file) - log_file_list(logging.DEBUG, dependecies) - - json_strings = sorted(visitor.json_strings) - if not json_strings: - logger.debug("No JSON strings in [%s].", source_file) - else: - logger.debug("JSON strings in [%s]: %s", source_file, str(json_strings)) - - gemm_config_names = sorted(visitor.gemm_config_names) - if not gemm_config_names: - logger.debug("No GEMM config names in [%s].", source_file) - else: - logger.debug( - "GEMM config names in [%s]: %s", source_file, str(gemm_config_names) - ) - - return dependecies, json_strings, gemm_config_names - - -def parse_source_file_recursively( - graph: nx.DiGraph, - source_file: Path, - config_files: list[Path], - visited: set[Path], - deps_to_ignore: set[Path] | None = None, -) -> None: - if deps_to_ignore is None: - deps_to_ignore = set() - stack = [source_file] - - while stack: - current = stack.pop() - if current in visited: - continue - - dependencies, json_strings, gemm_config_names = parse_source_file(current) - configs = resolve_configs(json_strings, gemm_config_names, config_files) - - # Add current node to the graph. - current_str = str(current) - graph.add_node(current_str) - logger.debug("Added graph node [%s].", current_str) - - # Add dependencies of current node, and respective edges, to the graph. - for d in dependencies: - d_str = str(d) - graph.add_node(d_str) - logger.debug("Added graph node [%s].", d_str) - graph.add_edge(d_str, current_str) - logger.debug("Added graph edge [%s]->[%s].", d_str, current_str) - - # Add configs of current node, and respective edges, to the graph. - for c in configs: - c_str = str(c) - graph.add_node(c_str) - logger.debug("Added graph node [%s].", c_str) - graph.add_edge(c_str, current_str) - logger.debug("Added graph edge [%s]->[%s].", c_str, current_str) - - stack.extend( - d - for d in dependencies - if (rd := root_dir() / d).is_file() and rd not in deps_to_ignore - ) - visited.add(current) - - -# Dependency graph. -# ------------------------------------------------------------------------------ - - -def tag_node(graph: nx.DiGraph, file: Path, tag: str) -> None: - file_str = str(file) - if file_str in graph.nodes: - graph.nodes[file_str]["type"] = tag - logger.debug("Tagged file [%s] as a '%s' in the graph.", file_str, tag) - else: - logger.warning( - "Couldn't find file [%s] in the graph, unable to tag it as '%s'.", - file_str, - tag, - ) - - -def add_files_to_dependency_graph( - graph: nx.DiGraph, - files: list[Path], - file_type: str, - config_files: list[Path], - visited: set[Path], - deps_to_ignore: set[Path] | None = None, -) -> None: - if deps_to_ignore is None: - deps_to_ignore = set() - for f in files: - parse_source_file_recursively( - graph, f, config_files, visited, deps_to_ignore=deps_to_ignore - ) - tag_node(graph, f, file_type) - - -def build_dependency_graph( - kernel_files: list[Path], - config_files: list[Path], - test_files: list[Path], - bench_files: list[Path], -) -> nx.DiGraph: - graph: nx.DiGraph = nx.DiGraph() - visited: set[Path] = set() - # These source files have dependencies that are hard to track and can be safely ignored for - # Triton test selection purposes: - deps_to_ignore: set[Path] = { - root_dir() / "aiter" / "jit" / "core.py", - root_dir() / "aiter" / "dist" / "utils.py", - root_dir() / "csrc" / "cpp_itfs" / "hsaco_launcher.py", - } - assert all( - d.is_file() and d.exists() for d in deps_to_ignore - ), "All ignored source file dependencies must exist in the filesystem." - # Add files that tests depends on. - add_files_to_dependency_graph( - graph, test_files, "test", config_files, visited, deps_to_ignore=deps_to_ignore - ) - # Add files that benchmarks depends on. - add_files_to_dependency_graph( - graph, - bench_files, - "bench", - config_files, - visited, - deps_to_ignore=deps_to_ignore, - ) - logger.debug( - "Built dependency graph of Triton source files with %d nodes and %d edges.", - graph.number_of_nodes(), - graph.number_of_edges(), - ) - # Tag kernel files. - for kernel_file in kernel_files: - if kernel_file.name != "__init__.py": - tag_node(graph, kernel_file, "kernel") - # Tag config files. - for config_file in config_files: - tag_node(graph, config_file, "config") - return graph - - -def find_tests_to_run(graph: nx.DiGraph, diff_inter_triton: list[Path]) -> list[Path]: - # Filter out unreachable files, i.e. Triton source files that are in the diff content but - # aren't in the dependency graph. - reachable_diff_inter_triton: list[Path] = [] - for p in diff_inter_triton: - p_str = str(p) - if p_str in graph: - reachable_diff_inter_triton.append(p) + dynamic.add(path) + continue + names = [target] else: - logger.warning( - "Triton source file [%s] isn't in the dependency graph, it's unreachable.", - p_str, - ) - if not reachable_diff_inter_triton: - logger.warning("There are no reachable tests from the Triton diff.") - logger.warning( - "Please check test selection script, there might be a bug in it." - ) - logger.warning( - "Please check Triton code base, there may be some filesystem organizations that aren't taken into account." - ) - return [] - - # Figure out if all reachable files are benchmarks. If this is the case, then we should use - # other graph traversal approach. - is_all_bench: bool = all( - graph.nodes[str(p)].get("type") == "bench" for p in reachable_diff_inter_triton + continue + found.update(filter(None, (resolve_module(n, gone) for n in names))) + return found + + +def dependents(paths, importers): + """Follow imports backwards from changed files to their consumers.""" + seen = set(paths) + frontier = list(seen) + while frontier: + for path in importers.get(frontier.pop(), ()): + if path not in seen: + seen.add(path) + frontier.append(path) + return seen + + +def changed_files(merge_ref): + # Compare the checked-out PR merge to its base parent; retain both rename paths. + refs = [merge_ref + "^1", merge_ref] + cmd = ["git", "diff", "--name-only", "--no-renames", "-z", *refs] + out = subprocess.run(cmd, cwd=ROOT, capture_output=True, text=True, check=True) + return [path for path in out.stdout.split("\0") if path] + + +def select(diff): + """Select affected test files, raising when a subset cannot be determined.""" + diff = [p for p in diff if not p.endswith(".md") and Path(p).name != ".gitkeep"] + if not any( + p.startswith((SRC, TESTS, "aiter/aot/triton/", "aiter/utility/triton/")) + or p in CI_FILES + for p in diff + ): + return [], ["no Triton/Gluon changes"] + for path in diff: + if path.startswith(GLOBAL_PREFIXES): + raise RuntimeError(f"{path}: shared machinery/CI changed") + if not path.startswith((SRC, TESTS, BENCH)) and ( + path.startswith(("aiter/", "op_tests/", "requirements")) + or path + in {"setup.py", "setup.cfg", "pyproject.toml", "pytest.ini", "conftest.py"} + ): + raise RuntimeError(f"{path}: shared code outside the import graph") + diff = [p for p in diff if p.startswith((SRC, TESTS))] + + tests = set(list_files(TESTS, "test_*.py")) + sources = list_files(SRC, "*.py") + gone = frozenset( + f + for f in diff + if f.endswith(".py") and f.startswith((SRC, TESTS)) and not (ROOT / f).is_file() ) - - # Traverse the dependency graph, searching from tests that are reachable from the diff content. - tests_to_run: set[Path] = set() - for p in reachable_diff_inter_triton: - p_str = str(p) - logger.debug("Searching for tests related to Triton source file [%s]...", p_str) - if not is_all_bench: - # Forward traversal for non-benchmarks. - reachable_files = nx.descendants(graph, p_str) | {p_str} - else: - # Backward traversal for benchmarks, filtering only dependencies on kernel files. After - # that, perform a forward traversal from kernels. This strategy isn't perfect but it's - # an attempt to take benchmark files into account. It may fail if a benchmark utility is - # changed in isolation. - reachable_files = { - kernel_descendant - for p_ancestor in nx.ancestors(graph, p_str) - if graph.nodes[p_ancestor].get("type") == "kernel" - for kernel_descendant in nx.descendants(graph, p_ancestor) + all_sources = set(sources) | {f for f in gone if f.startswith(SRC)} + by_name = {} + for source in all_sources: + by_name.setdefault(Path(source).stem, set()).add(source) + + dynamic, importers = set(), {} + for path in sources + list_files(TESTS, "*.py"): + dependencies = scan_imports(path, gone, dynamic) + if path in tests: + # Some tests select their wrapper by name instead of importing it. + subject = Path(path).stem.removeprefix("test_") + for name in (subject, subject.removeprefix("compile_")): + dependencies.update(by_name.get(name, ())) + for dependency in dependencies: + importers.setdefault(dependency, set()).add(path) + + def affected(paths): + return tests & dependents(paths, importers) + + def folder_tests(path): + category = category_of(path) + hits = {test for test in tests if category and category_of(test) == category} + if not hits: + raise RuntimeError(f"{path}: no mapped tests or op test folder") + return hits + + selected = set() + reasons = [] + for path in diff: + if Path(path).name == "__init__.py": + raise RuntimeError(f"{path}: package initializer changed") + + if path.startswith(CONFIGS): + # configs/////.json + parts = path[len(CONFIGS) :].split("/") + if not ( + path.endswith(".json") + and len(parts) >= 5 + and parts[1] in ("triton", "gluon") + ): + raise RuntimeError(f"{path}: unknown config layout") + op, family = parts[2:4] + categories = CONFIG_CATEGORIES.get(op, {op}) + hits = {test for test in tests if category_of(test) in categories} + if not hits and not any( + category_of(source) in categories or Path(source).stem == op + for source in sources + ): + raise RuntimeError(f"{path}: unknown config op '{op}'") + # Family names need not match filenames; include the whole op and dependents. + consumers = { + source + for source in sources + if category_of(source) in categories + or Path(source).stem in (op, family) } - logger.debug( - "There %s %d file%s reachable from [%s].", - "is" if len(reachable_files) == 1 else "are", - len(reachable_files), - "" if len(reachable_files) == 1 else "s", - p_str, - ) - - test_files = { - Path(f) for f in reachable_files if graph.nodes[f].get("type") == "test" - } - if test_files: - logger.debug( - "There %s %d test%s reachable from [%s].", - "is" if len(test_files) == 1 else "are", - len(test_files), - "" if len(test_files) == 1 else "s", - p_str, - ) - tests_to_run.update(test_files) + hits.update(affected(consumers)) + if not hits: + raise RuntimeError(f"{path}: config maps to no tests") + reason = f"'{op}' folder and dependents" + elif path.startswith(TESTS): + hits = affected([path]) + if Path(path).name.startswith("test_") and path.endswith(".py"): + reason = "changed test and importers" + else: + hits.update(folder_tests(path)) + reason = "test helper folder and importers" else: - logger.warning( - "Couldn't find test files related to [%s] Triton source.", p_str - ) - - if tests_to_run: - sorted_tests_to_run = sorted(tests_to_run) - logger.info( - "There %s %d test%s reachable from the Triton diff:", - "is" if len(sorted_tests_to_run) == 1 else "are", - len(sorted_tests_to_run), - "" if len(sorted_tests_to_run) == 1 else "s", - ) - log_file_list(logging.INFO, sorted_tests_to_run) - return sorted_tests_to_run - else: - logger.warning("Couldn't find any test file related to Triton diff.") - logger.warning( - "Please check test selection script, there might be a bug in it." - ) - logger.warning("Please check Triton code base, there may be untested kernels.") - return [] - - -# Writing to GitHub environment file. -# ------------------------------------------------------------------------------ - - -def write_env_file(env_var: str, env_file: str, tests_to_run: list[Path]) -> None: - if env_var is None or not (env_var := env_var.strip()): - logger.info( - "Environment variable is absent, environment file won't be written." - ) - return - if env_file is None or not (env_file := env_file.strip()): - logger.info("Environment file is absent, it won't be written.") - return - if not tests_to_run: - logger.warning( - "List of tests to run is empty, enviroment file won't be written." - ) - return - tests_to_run_joined_str = " ".join(str(t) for t in tests_to_run) - env_file_data = f"{env_var}={tests_to_run_joined_str}" - logger.debug("Writing [%s] to [%s]...", env_file_data, env_file) - try: - with open(env_file, "a") as env_file_fd: - env_file_fd.write(env_file_data + "\n") - logger.info("Wrote tests to run to [%s] environment file.", env_file) - except OSError: - logger.exception("I/O error while writing to [%s] environment file.", env_file) - logger.info("The entire Triton test suite will be executed.") - - -# Command line interface parsing. -# ------------------------------------------------------------------------------ - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="select which Triton tests to run based on git diff" + if not path.endswith(".py"): + raise RuntimeError(f"{path}: non-Python source changed") + hits = affected([path]) + reason = "paired and importing tests" + if not hits: + hits = folder_tests(path) + reason = "unmapped source; whole op folder" + selected.update(hits) + reasons.append(f"{path}: {reason} ({len(hits)} tests)") + + if not selected: + raise RuntimeError("relevant files changed but no tests were selected") + # Only unresolved runtime imports need conservative extra coverage. + extra = affected(dynamic) - selected + selected.update(extra) + if extra: + reasons.append("unresolved dynamic imports: " + ", ".join(sorted(extra))) + return sorted(selected), reasons + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) - parser.add_argument( - "-s", "--source", help="source branch, defaults to current branch" + ap.add_argument( + "--merge-ref", + help="PR merge ref; diff is taken against its first parent. Without it " + "the full suite runs", ) - parser.add_argument( - "-t", "--target", default="main", help="target branch, defaults to main" + ap.add_argument("--all", action="store_true", help="select the full suite") + ap.add_argument( + "--dry-run", + action="store_true", + help="print the selection, but run what CI ran before test selection", ) - parser.add_argument( - "-v", - "--env-var", - default="TRITON_TEST", - help="environment variable to store which tests to run, defaults to TRITON_TEST", - ) - parser.add_argument( - "-f", - "--env-file", - help="environment file to write, won't write anything if absent", - ) - parser.add_argument( - "-l", - "--log-level", - type=str.lower, - choices=["critical", "error", "warning", "info", "debug", "off"], - default="info", - help="log level to enable (default: info)", - ) - args = parser.parse_args() - args.log_level = { - "critical": logging.CRITICAL, - "error": logging.ERROR, - "warning": logging.WARNING, - "info": logging.INFO, - "debug": logging.DEBUG, - "off": logging.CRITICAL + 1000, - }[args.log_level] - return args - - -# Main logic and script entry point. -# ------------------------------------------------------------------------------ - - -def main_logic(args: argparse.Namespace) -> None: - diff_files = get_filename_diff(args.source, args.target) - all_files, kernel_files, config_files, test_files, bench_files = ( - list_triton_source_files() - ) - diff_inter_triton = diff_files & all_files - del diff_files, all_files - - if not diff_inter_triton: - logger.info( - "There are no Triton source files in diff, there's no need to run Triton tests." - ) - return - - logger.info( - "There %s %d Triton source file%s in the diff:", - "is" if len(diff_inter_triton) == 1 else "are", - len(diff_inter_triton), - "" if len(diff_inter_triton) == 1 else "s", - ) - - sorted_diff_inter_triton = sorted(diff_inter_triton) - del diff_inter_triton - log_file_list(logging.INFO, sorted_diff_inter_triton) - - graph = build_dependency_graph(kernel_files, config_files, test_files, bench_files) - tests_to_run = find_tests_to_run(graph, sorted_diff_inter_triton) - - if not tests_to_run: - return - - write_env_file(args.env_var, args.env_file, tests_to_run) - - -def main() -> None: - start_timestamp = time.perf_counter() - args = parse_args() - logging.basicConfig(format="%(levelname)s|%(message)s", level=args.log_level) - main_logic(args) - end_timestamp = time.perf_counter() - elapsed_time_s = end_timestamp - start_timestamp - logger.info("Finished, execution took %.2f seconds.", elapsed_time_s) + ap.add_argument("--output", default="selected_triton_tests.list") + args = ap.parse_args() + everything = list_files(TESTS, "test_*.py") + needs_mi300x = True + diff = None + tests, reasons = everything, ["full suite requested"] + try: + if args.merge_ref: + diff = changed_files(args.merge_ref) + needs_mi300x = any( + path.startswith((CONFIGS + "gfx942/", GLUON_KERNELS + "gfx942/")) + for path in diff + ) + if not args.all: + tests, reasons = select(diff) + except Exception as why: # noqa: BLE001 -- any failure falls open to a full run + tests, reasons = everything, [f"FULL SUITE: {why}"] + report = f"Triton tests: {len(tests)} files\n" + report += "".join(f"- {reason}\n" for reason in reasons) + if args.dry_run and not args.all: + # Print the selection, but keep the tests CI ran before selection. + report = "DRY RUN, selection not applied.\n" + report + if len(tests) < len(everything): + report += "".join(f" {test}\n" for test in tests) + tests = everything if diff is None or legacy_trigger(diff) else [] + report += f"Running {len(tests)} files, as CI did before test selection.\n" + Path(args.output).write_text("".join(t + "\n" for t in tests), encoding="utf-8") + if output := os.environ.get("GITHUB_OUTPUT"): + with open(output, "a", encoding="utf-8") as fh: + fh.write( + f"selected_count={len(tests)}\nneeds_mi300x={str(needs_mi300x).lower()}\n" + ) + print(report, file=sys.stderr) + if summary := os.environ.get("GITHUB_STEP_SUMMARY"): + with open(summary, "a", encoding="utf-8") as fh: + fh.write(report + "\n") if __name__ == "__main__": diff --git a/.github/scripts/split_tests.sh b/.github/scripts/split_tests.sh index 9bb427f58bd..1174577a8a3 100755 --- a/.github/scripts/split_tests.sh +++ b/.github/scripts/split_tests.sh @@ -8,20 +8,23 @@ # Parameters: # --shards N number of shards (required) # --test-type TYPE test type, default aiter +# --select-file F only shard test files listed in F (one per line) # --dry-run only output allocation plan, do not execute # -v Pytest's -v option, no effect -# Exit code: always 0 +# Exit code: 0 on success, 1 for invalid arguments or selected paths set -euo pipefail SHARDS=0 TEST_TYPE="aiter" DRY_RUN=0 +SELECT_FILE="" while [[ $# -gt 0 ]]; do case "$1" in --shards) SHARDS="$2"; shift 2 ;; --test-type) TEST_TYPE="$2"; shift 2 ;; + --select-file) SELECT_FILE="$2"; shift 2 ;; --dry-run) DRY_RUN=1; shift ;; -v|--verbose) shift ;; # compatibility, ignore *) @@ -67,6 +70,42 @@ if [[ ${#ALL_FILES[@]} -eq 0 ]]; then exit 1 fi +# Apply the optional selection before sharding. +if [[ -n "$SELECT_FILE" ]]; then + if [[ ! -f "$SELECT_FILE" ]]; then + echo "Selection file not found: $SELECT_FILE" >&2 + exit 1 + fi + declare -A SELECTED=() + while IFS= read -r line || [[ -n "$line" ]]; do + [[ -n "$line" ]] && SELECTED["$line"]=1 + done < "$SELECT_FILE" + FILTERED=() + for f in "${ALL_FILES[@]}"; do + if [[ -n "${SELECTED[$f]:-}" ]]; then + FILTERED+=("$f") + unset "SELECTED[$f]" + fi + done + # Never silently drop a path the selector asked to run. + if [[ ${#SELECTED[@]} -gt 0 ]]; then + echo "Selection lists paths that are not test files under ${TEST_DIR}:" >&2 + for f in "${!SELECTED[@]}"; do echo " ${f}" >&2; done + exit 1 + fi + echo "Test selection: ${#FILTERED[@]} of ${#ALL_FILES[@]} test files selected." + if [[ ${#FILTERED[@]} -eq 0 ]]; then + echo "Selection is empty — writing ${SHARDS} empty shard lists." + if [[ $DRY_RUN -eq 0 ]]; then + for ((s=0; s < SHARDS; s++)); do + : > "${TEST_TYPE}_shard_${s}.list" + done + fi + exit 0 + fi + ALL_FILES=("${FILTERED[@]}") +fi + # ------------------------------ # FILE_TIMES (seconds), unknown files default 15 # ------------------------------ diff --git a/.github/workflows/pr-title-tags.yaml b/.github/workflows/pr-title-tags.yaml index c3fb24185d1..d68aba27b9f 100644 --- a/.github/workflows/pr-title-tags.yaml +++ b/.github/workflows/pr-title-tags.yaml @@ -12,7 +12,8 @@ name: PR Title Tags & Labels # component namespace ([fix], [Perf], [MLA], [gfx1250], ...) are never # touched. The same rule applies to labels: only the component labels listed # in COMPONENT_LABELS are ever added or removed — human labels (ci:*, ...) -# are left alone, while hand-toggling a managed component label is corrected +# are left alone, except gfx942 config changes add ci:triton-300x. +# Hand-toggling a managed component label is corrected # on the spot. To keep a PR's title untouched, add the 'no-auto-title' label — # it takes effect immediately, and removing it re-syncs the title on the spot. # Component labels are synced regardless of the opt-out: the AITER-triton @@ -50,7 +51,8 @@ jobs: script: | // The component labels this workflow owns (name -> color). Only // these are ever added or removed; anything else on the PR - // (ci:*, no-auto-title, ...) belongs to humans. + // (ci:*, no-auto-title, ...) belongs to humans, except the + // add-only ci:triton-300x label for gfx942 config changes below. const COMPONENT_LABELS = { 'Triton/Gluon': '1d76db', 'ASM': '5319e7', @@ -86,12 +88,17 @@ jobs: const optedOut = pr.labels.some(l => l.name === 'no-auto-title'); // ---- changed files (paginated) ---- - const files = (await github.paginate(github.rest.pulls.listFiles, { + const changedFiles = await github.paginate(github.rest.pulls.listFiles, { owner: context.repo.owner, repo: context.repo.repo, pull_number: pr.number, per_page: 100, - })).map(f => f.filename); + }); + if (changedFiles.length !== pr.changed_files) { + throw new Error(`Incomplete PR file list: ${changedFiles.length}/${pr.changed_files}`); + } + const files = changedFiles.map(f => f.filename); + const changedPaths = changedFiles.flatMap(f => [f.filename, f.previous_filename].filter(Boolean)); // ---- path predicates ---- const isDoc = f => f.startsWith('docs/') || f.endsWith('.md'); @@ -218,3 +225,13 @@ jobs: } else { core.info('Labels already in sync.'); } + + if (changedPaths.some(f => f.startsWith('aiter/ops/triton/configs/gfx942/') && !isDoc(f)) + && !pr.labels.some(l => l.name === 'ci:triton-300x')) { + await github.rest.issues.addLabels({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pr.number, + labels: ['ci:triton-300x'], + }); + } diff --git a/.github/workflows/pr-welcome-comment.yaml b/.github/workflows/pr-welcome-comment.yaml index bcf243fa202..9823410d1ee 100644 --- a/.github/workflows/pr-welcome-comment.yaml +++ b/.github/workflows/pr-welcome-comment.yaml @@ -31,7 +31,8 @@ jobs: | Label | Tests | |-------|-------| | \`ci:gfx1250-ffm-triton\` | Run the five-shard gfx1250 FFM Triton test suite | - | \`ci:triton-300x\` | Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X | + | \`ci:triton-300x\` | Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X | + | \`ci:triton-355\` | Run the full Triton test suite on MI35X, not only the tests the change affects | | \`multigpu\` | Aiter multi-GPU tests on the 8-GPU runner | | \`ci:sglang\` | SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy | | \`ci:atom\` | ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B | diff --git a/.github/workflows/triton-test.yaml b/.github/workflows/triton-test.yaml index 604e05ce468..c586ee10d18 100644 --- a/.github/workflows/triton-test.yaml +++ b/.github/workflows/triton-test.yaml @@ -1,49 +1,39 @@ name: Triton Test +# Test selection: .github/scripts/select_triton_tests.py picks the unit tests a +# PR affects. +# - It runs as a dry run (DRY_RUN in the "Select Triton tests from PR diff" +# step): the selection is printed, and CI runs what it ran before test +# selection existed. DRY_RUN: 'false' applies the selection; setting it +# back to 'true' is the revert for every PR. +# - PR label ci:triton-355: that PR runs the full suite. +# Pushes to main always run the full suite, on MI35X and MI300X. + on: push: branches: [main] pull_request: - # labeled: re-run when adding ci:triton-300x to start extra MI300X jobs - types: [opened, synchronize, reopened, ready_for_review, labeled] + # Selection skips unrelated changes; the result check is reported for every PR. + types: [opened, synchronize, reopened, ready_for_review, labeled, unlabeled] branches: [main] - paths: - - "aiter/ops/triton/**" - - "op_tests/triton_tests/**" - - "op_tests/op_benchmarks/triton/**" - - ".github/scripts/build_aiter_triton.sh" - - ".github/scripts/download_triton_wheel.sh" - - ".github/scripts/install_triton.sh" - - ".github/scripts/select_triton_tests.py" - - ".github/scripts/split_tests.sh" - - ".github/scripts/verify_triton_pin.py" - - ".github/requirements/triton-test.txt" - # Docs-only changes (README.md, configs/CLAUDE.md, conv/DESIGN.md, ...) - # don't affect test results — skip the test suites when a PR touches - # nothing else. - - "!**/*.md" - - '!docs/**' - - '!LICENSE' - - '!.gitignore' - - '!.github/workflows/**' - - '!.github/scripts/sglang_downstream.py' - # Re-include Triton workflow files after the broad workflow exclusion. - - ".github/workflows/triton-test.yaml" - - ".github/workflows/prepare-triton-wheel.yaml" - - ".github/workflows/ci-config.yaml" workflow_dispatch: concurrency: - group: ${{ github.workflow }}-${{ github.ref }} + # Labels other than ci:triton-300x and ci:triton-355 change nothing here: such + # a run gets a group of its own, so it cannot cancel the run in progress, and + # all its jobs skip. + group: ${{ github.workflow }}-${{ github.ref }}${{ (github.event.action == 'labeled' || github.event.action == 'unlabeled') && !contains(fromJSON('["ci:triton-300x", "ci:triton-355"]'), github.event.label.name) && format('-{0}', github.run_id) || '' }} cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} jobs: ci_config: - if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} + needs: [split_triton_tests] + if: ${{ needs.split_triton_tests.outputs.selected_count != '0' }} uses: ./.github/workflows/ci-config.yaml check-signal: - if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} + needs: [split_triton_tests] + if: ${{ needs.split_triton_tests.outputs.selected_count != '0' }} runs-on: ubuntu-latest steps: - name: Checkout code @@ -57,17 +47,46 @@ jobs: # Step 1: split triton tests into 8 shards, output triton_shard_0.list ... triton_shard_7.list split_triton_tests: - if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} + # Not for drafts, nor for labels other than ci:triton-300x and ci:triton-355. + if: ${{ !github.event.pull_request.draft && (github.event.action != 'labeled' && github.event.action != 'unlabeled' || contains(fromJSON('["ci:triton-300x", "ci:triton-355"]'), github.event.label.name)) }} runs-on: ubuntu-latest - needs: [check-signal] outputs: shard_count: 8 + shards: ${{ steps.split.outputs.shards }} + selected_count: ${{ steps.select.outputs.selected_count }} + # main always; a PR for gfx942-specific changes or the ci:triton-300x label. + needs_mi300x: ${{ github.ref == 'refs/heads/main' || steps.select.outputs.needs_mi300x == 'true' || contains(github.event.pull_request.labels.*.name, 'ci:triton-300x') }} steps: - name: Checkout code uses: actions/checkout@v4 + with: + # Compare the event's merge commit with its base parent. + fetch-depth: 2 + + - name: Select Triton tests from PR diff + id: select + env: + # 'true': print the selection, but run what CI ran before test + # selection existed. 'false': run only the selected tests. + DRY_RUN: "true" + FULL_SUITE: ${{ github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'ci:triton-355') }} + run: | + set -euo pipefail + args=(--merge-ref HEAD --output selected_triton_tests.list) + if [ "$FULL_SUITE" = true ]; then + args+=(--all) + elif [ "$DRY_RUN" = true ]; then + args+=(--dry-run) + fi + python3 .github/scripts/select_triton_tests.py "${args[@]}" - name: Split Triton Tests (8 shards) - run: ./.github/scripts/split_tests.sh --shards 8 --test-type triton + id: split + run: | + ./.github/scripts/split_tests.sh --shards 8 --test-type triton --select-file selected_triton_tests.list + # Only shards that got tests start a GPU job. + shards=$(for s in $(seq 0 7); do grep -q '[^[:space:]]' "triton_shard_${s}.list" && echo "$s"; done | paste -sd, -) + echo "shards=[${shards}]" >> "$GITHUB_OUTPUT" - name: Upload test shard lists as artifact uses: actions/upload-artifact@v4 @@ -77,8 +96,8 @@ jobs: retention-days: 7 prepare_triton_wheel: - if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} - needs: [check-signal, ci_config] + if: ${{ needs.split_triton_tests.outputs.selected_count != '0' }} + needs: [split_triton_tests, check-signal, ci_config] uses: ./.github/workflows/prepare-triton-wheel.yaml with: docker-image: ${{ needs.ci_config.outputs.pytorch_py312_image }} @@ -87,7 +106,7 @@ jobs: # Step 2: MI35X matrix jobs triton: - if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} + if: ${{ needs.split_triton_tests.outputs.selected_count != '0' }} name: Triton Tests (MI35X) / Shard ${{ matrix.shard }} runs-on: linux-aiter-mi35x-1 timeout-minutes: 60 @@ -95,7 +114,7 @@ jobs: strategy: fail-fast: false matrix: - shard: [0, 1, 2, 3, 4, 5, 6, 7] + shard: ${{ fromJSON(needs.split_triton_tests.outputs.shards) }} env: DOCKER_IMAGE: ${{ needs.ci_config.outputs.pytorch_py312_image }} TRITON_WHEEL_ARTIFACT_NAME: ${{ needs.ci_config.outputs.triton_wheel_artifact_name }} @@ -248,35 +267,17 @@ jobs: echo "Setting up Aiter and Triton..." docker exec -w /workspace triton_test ./.github/scripts/build_aiter_triton.sh - # TODO: Uncomment [docker exec -w /workspace triton_test cat "${ENV_FILE}" >> "${GITHUB_ENV}"] - # command to enable test selection. - # - name: Triton Test Selection Script - # # main branch should always run the full test suite. - # if: ${{ github.ref != 'refs/heads/main' }} - # run: | - # set -ex - # git fetch --no-tags origin --depth=1 \ - # "pull/${{ github.event.pull_request.number }}/merge:pr-merge" - # git fetch --no-tags origin --depth=1 \ - # "refs/heads/${{ github.event.pull_request.base.ref }}:target" - # ENV_FILE=$(docker exec -w /workspace triton_test \ - # mktemp /workspace/github_env.XXXXXXXXXX.tmp) - # docker exec -w /workspace triton_test \ - # git config --global --add safe.directory /workspace - # docker exec -w /workspace triton_test \ - # python .github/scripts/select_triton_tests.py \ - # --source pr-merge --target target \ - # --env-var TRITON_TEST --env-file "${ENV_FILE}" - # # docker exec -w /workspace triton_test cat "${ENV_FILE}" >> "${GITHUB_ENV}" - # docker exec -w /workspace triton_test rm "${ENV_FILE}" - - name: Triton Tests run: | set -exo pipefail echo "Running Triton Tests..." mkdir -p test-reports docker exec -w /workspace triton_test mkdir -p test-reports - # MI35X: skip the MHA-PE tests that fail an accuracy check on gfx950 with the coming Triton release. + if [ -z "$TRITON_TEST" ]; then + echo "No tests selected for this shard." + echo '' > test-reports/triton.xml + exit 0 + fi # Unit tests must not autotune: surface any tuning and fail the job if it happened. docker exec -e TRITON_PRINT_AUTOTUNING=1 -w /workspace triton_test pytest -v -s ${TRITON_TEST} \ --junitxml=test-reports/triton.xml 2>&1 | tee test-reports/pytest.log @@ -310,9 +311,9 @@ jobs: run: | docker rm -f triton_test || true - # Step 2b: MI300X matrix jobs (opt-in via ci:triton-300x on PRs, always on main) + # MI300X: see needs_mi300x above. triton-mi300x: - if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') && (github.ref == 'refs/heads/main' || (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'ci:triton-300x'))) }} + if: ${{ needs.split_triton_tests.outputs.selected_count != '0' && needs.split_triton_tests.outputs.needs_mi300x == 'true' }} name: Triton Tests (MI300X) / Shard ${{ matrix.shard }} runs-on: linux-aiter-oci-mi300x-1 timeout-minutes: 60 @@ -320,7 +321,7 @@ jobs: strategy: fail-fast: false matrix: - shard: [0, 1, 2, 3, 4, 5, 6, 7] + shard: ${{ fromJSON(needs.split_triton_tests.outputs.shards) }} env: DOCKER_IMAGE: ${{ needs.ci_config.outputs.pytorch_py312_image }} TRITON_WHEEL_ARTIFACT_NAME: ${{ needs.ci_config.outputs.triton_wheel_artifact_name }} @@ -427,6 +428,11 @@ jobs: echo "Running Triton Tests..." mkdir -p test-reports docker exec -w /workspace triton_test mkdir -p test-reports + if [ -z "$TRITON_TEST" ]; then + echo "No tests selected for this shard." + echo '' > test-reports/triton.xml + exit 0 + fi # Unit tests must not autotune: surface any tuning and fail the job if it happened. docker exec -e TRITON_PRINT_AUTOTUNING=1 -w /workspace triton_test pytest -v -s ${TRITON_TEST} \ --junitxml=test-reports/triton.xml 2>&1 | tee test-reports/pytest.log @@ -460,65 +466,34 @@ jobs: run: | docker rm -f triton_test || true - triton-mi300x-test-finish: - if: ${{ (!github.event.pull_request || github.event.pull_request.draft == false) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') && (github.ref == 'refs/heads/main' || (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'ci:triton-300x'))) }} - name: Triton MI300X Test Results - runs-on: ubuntu-latest - needs: [triton-mi300x] - steps: - - name: Download all MI300X test reports - uses: actions/download-artifact@v4 - timeout-minutes: 15 - with: - pattern: triton-test-mi300x-shard-* - path: . - - - name: Check Triton MI300X Test Results - run: | - set -ex - echo "Checking Triton MI300X Test Results..." - all_passed=true - for shard in {0..7}; do - if [ ! -f triton-test-mi300x-shard-${shard}/triton.xml ]; then - echo "MI300X test report for shard ${shard} not found." - all_passed=false - break - fi - done - if [ "$all_passed" = true ]; then - echo "All MI300X tests passed." - else - echo "MI300X test failures or errors detected." - exit 1 - fi - triton-test-finish: - if: ${{ (!github.event.pull_request || !github.event.pull_request.draft) && (github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'ci:triton-300x') }} name: Triton Test Results + # Every run but an ignored label's, which skips split_triton_tests; drafts + # report here too, as passing with nothing run. + if: ${{ always() && (needs.split_triton_tests.result != 'skipped' || github.event.pull_request.draft == true) }} runs-on: ubuntu-latest - needs: [triton] + needs: [split_triton_tests, triton, triton-mi300x] steps: - - name: Download all test reports - uses: actions/download-artifact@v4 - timeout-minutes: 15 - with: - pattern: triton-test-shard-* - path: . - - name: Check Triton Test Results + - name: Check required Triton jobs + env: + PR_DRAFT: ${{ github.event.pull_request.draft }} + SELECTION_RESULT: ${{ needs.split_triton_tests.result }} + SELECTED_COUNT: ${{ needs.split_triton_tests.outputs.selected_count }} + NEEDS_MI300X: ${{ needs.split_triton_tests.outputs.needs_mi300x }} + MI35X_RESULT: ${{ needs.triton.result }} + MI300X_RESULT: ${{ needs.triton-mi300x.result }} run: | - set -ex - echo "Checking Triton Test Results..." - all_passed=true - for shard in {0..7}; do - if [ ! -f triton-test-shard-${shard}/triton.xml ]; then - echo "Test report for shard ${shard} not found." - all_passed=false - break - fi - done - if [ "$all_passed" = true ]; then - echo "All tests passed." - else - echo "Test failures or errors detected." - exit 1 + set -euo pipefail + if [ "$PR_DRAFT" = true ]; then + echo "Draft PR; tests are deferred until ready for review." + exit 0 + fi + test "$SELECTION_RESULT" = success + if [ "$SELECTED_COUNT" = 0 ]; then + echo "No Triton/Gluon tests required for this change." + exit 0 + fi + test "$MI35X_RESULT" = success + if [ "$NEEDS_MI300X" = true ]; then + test "$MI300X_RESULT" = success fi