diff --git a/aiter/aot/flydsl/README.md b/aiter/aot/flydsl/README.md index 02352b38eb8..7a78ed19a7e 100644 --- a/aiter/aot/flydsl/README.md +++ b/aiter/aot/flydsl/README.md @@ -98,14 +98,18 @@ python -m aiter.aot.flydsl.chunk_gdn_h --csv /path/to/tuned.csv | `AITER_FLYDSL_AOT_TIMEOUT` | Per-kernel wall-clock cap (seconds). A worker stuck *alive* past this is killed (and retried); `0` disables. | `1200` | | `AITER_FLYDSL_AOT_MAX_RETRIES` | Retries for a worker that **died abnormally** (OOM-kill / segfault / timeout-kill). A clean compile error is never retried. `0` disables. | `2` | | `AITER_CONFIGS` | Resolves the default CSV lookup path (same as the runtime JIT) | repo built-in | -| `ARCH` / `GPU_ARCHS` | Selects which jobs to build, not what arch a job compiles *for* (that comes from the CSV's `cu_num`). `conv.py` applies it inside `parse_csv`, so both `python -m` and the `setup.py` path (`run_aot`) honour it. `gemm.py` still filters in `main()` only, so `run_aot` builds all of its archs. | auto-detect | - -> **About the compile target arch.** The arch each kernel is actually compiled -> for is derived per-job from the CSV's `cu_num` column (`cu_num_to_arch(...)`) -> and applied internally via `FLYDSL_GPU_ARCH`. That internal var is overwritten -> for every job, so setting `ARCH` / `GPU_ARCHS` / `FLYDSL_GPU_ARCH` in your shell -> does **not** change what gets built. To cross-compile, edit the `cu_num` -> column in the CSV. +| `AITER_GPU_TARGETS` / `ARCH` / `GPU_ARCHS` | Filter **GEMM** AOT jobs by target. The first one set wins, in that order. Unset builds every job. | unset | + +> **Compile target.** Each job is compiled for the arch in its `gfx` field, or +> from the CSV `cu_num` via `cu_num_to_arch(...)` when `gfx` is empty. AOT +> applies that through `FLYDSL_GPU_ARCH` and overwrites it per job, so setting +> `FLYDSL_GPU_ARCH` in the shell does not change what gets built. + +`AITER_GPU_TARGETS=gfx950:128` is the exact `(gfx, cu_num)` form on both +paths. The other two differ by path. Running `python -m aiter.aot.flydsl.gemm` +directly, `ARCH` selects every GEMM job of that arch and `GPU_ARCHS` does the +same (`CU_NUM` is ignored); packaging (`setup.py`) reads no `ARCH` and pairs +`GPU_ARCHS` with `CU_NUM` into exact SKU pairs instead. Example: diff --git a/aiter/aot/flydsl/common.py b/aiter/aot/flydsl/common.py index 8018f28302f..21376fe0d77 100755 --- a/aiter/aot/flydsl/common.py +++ b/aiter/aot/flydsl/common.py @@ -133,6 +133,18 @@ def override_env(var_name: str, value: str | None) -> Iterator[None]: os.environ[var_name] = prev +def _filter_collected_aot_jobs( + kind: OpKind, jobs: list[dict[str, Any]] +) -> list[dict[str, Any]]: + # GEMM only: parse_csv for the others keep just one of gfx/cu_num, + # so a (gfx, cu_num) filter would drop all of them. + if kind is OpKind.GEMM: + from .gemm import filter_jobs_for_build_targets + + return filter_jobs_for_build_targets(jobs) + return jobs + + def _collect_aot_jobs_for(kind: OpKind) -> list[dict[str, Any]]: """Load DEFAULT_CSVS + parse_csv for the named kind and return its job list. Note: importing .gemm / .moe / .chunk_gdn_h here also @@ -161,7 +173,8 @@ def _collect_aot_jobs_for(kind: OpKind) -> list[dict[str, Any]]: from .chunk_gdn_h import DEFAULT_CSVS, parse_csv else: raise ValueError(f"unknown FlyDSL AOT kind: {kind!r}") - return collect_aot_jobs(DEFAULT_CSVS, parse_csv) + jobs = collect_aot_jobs(DEFAULT_CSVS, parse_csv) + return _filter_collected_aot_jobs(kind, jobs) def _compile_one_config_for(kind: OpKind) -> Callable[..., dict[str, Any]]: diff --git a/aiter/aot/flydsl/gemm.py b/aiter/aot/flydsl/gemm.py index 3da3e63b80e..3d4bb796d42 100755 --- a/aiter/aot/flydsl/gemm.py +++ b/aiter/aot/flydsl/gemm.py @@ -31,7 +31,12 @@ Environment variables: FLYDSL_RUNTIME_CACHE_DIR Cache directory (default: ~/.flydsl/cache) - GPU_ARCHS / ARCH Target GPU architecture information for logging. + AITER_GPU_TARGETS ';'/','-separated 'gfx' or 'gfx:cu_num' build + targets; missing pairs retain runtime fallback. + GPU_ARCHS Arch list. Packaging resolves it with CU_NUM into + exact targets; the CLI keeps every job of that arch. + CU_NUM Compute-unit count paired with GPU_ARCHS. + ARCH Bare-arch filter on the CLI; packaging ignores it. """ from __future__ import annotations @@ -57,6 +62,11 @@ run_jobs_parallel, ) from aiter.jit.core import AITER_CONFIGS +from aiter.jit.utils.build_targets import ( + KNOWN_GFX, + _parse_gpu_targets_env, +) +from aiter.jit.utils.chip_info import get_build_targets, get_gfx_runtime from aiter.ops.flydsl.batched_gemm_a8w8_gfx950 import ( BMM_MFMA_NAME_PREFIX, compile_bmm_a8w8_mxfp8_gfx950, @@ -213,7 +223,7 @@ def _bmm_mfma_job( "m": 1, "n": n, "k": k, - "cu_num": 0, # the gfx column names the arch + "cu_num": int(row.get("cu_num") or 0), "gfx": row.get("gfx", "").strip(), "b": b, # A table from before the column holds only 128x128 rows (as at runtime). @@ -253,7 +263,7 @@ def parse_csv(csv_path: str): m = int(row["M"]) n = int(row["N"]) k = int(row["K"]) - cu_num = int(row.get("cu_num", "0")) + cu_num = int(row.get("cu_num") or 0) gfx = row.get("gfx", "").strip() if row.get("kernelId", "").strip() == "bmm": # the batched GEMM, B = 1 @@ -994,6 +1004,106 @@ def job_arch(cu_num: int = 0, gfx: str = "") -> str: return gfx or cu_num_to_arch(cu_num, default=GEMM_AOT_ARCH_DEFAULT) +def active_target_env(*, cli: bool = True) -> tuple[str, str] | None: + """First configured target variable; packaging ignores the CLI-only ARCH.""" + if cli: + names = ("AITER_GPU_TARGETS", "ARCH", "GPU_ARCHS") + else: + names = ("AITER_GPU_TARGETS", "GPU_ARCHS") + for var_name in names: + value = (os.environ.get(var_name) or "").strip() + if value: + return var_name, value + return None + + +def _archs_from_env(value: str) -> set[str]: + """Arch names in an ARCH/GPU_ARCHS value. Raises ValueError on a bad one.""" + archs = {a.strip().lower() for a in re.split(r"[;,]", value) if a.strip()} + if not archs: + raise ValueError("contains no valid architecture names") + if "native" in archs and len(archs) != 1: + raise ValueError("'native' cannot be combined with other targets") + unknown = archs - KNOWN_GFX - {"native"} + if unknown: + raise ValueError(f"contains unknown target(s): {sorted(unknown)}") + return archs + + +def _requested_targets( + var_name: str, value: str, *, arch_wide: bool +) -> set[tuple[str, int | None]]: + """Resolve (gfx, CU count) targets; None means every CU count of an arch.""" + if var_name == "AITER_GPU_TARGETS": + return set(_parse_gpu_targets_env()) + + archs = _archs_from_env(value) + if arch_wide: + if archs == {"native"}: + archs = {get_gfx_runtime()} + return {(gfx, None) for gfx in archs} + + # Normalize separators/case before resolving packaging targets. The shared + # resolver handles CU_NUM, native detection, and live-device CU counts. + with override_env("GPU_ARCHS", ";".join(sorted(archs))): + return set(get_build_targets()) + + +def filter_jobs_for_build_targets( + jobs: list[dict], *, arch_wide: bool = False +) -> list[dict]: + """Select jobs for AITER_GPU_TARGETS, or for GPU_ARCHS with CU_NUM. + + The CLI passes arch_wide=True: ARCH takes priority over GPU_ARCHS and both + select every CU count of the requested architectures. AITER_GPU_TARGETS + always selects exact targets and takes priority over either variable. + """ + active = active_target_env(cli=arch_wide) + if active is None: + return jobs + var_name, value = active + try: + targets = _requested_targets(var_name, value, arch_wide=arch_wide) + except ValueError as e: + raise RuntimeError(f"{var_name} {e}.") from None + + selected = [] + covered = set() + legacy_cu_nums = set() + for job in jobs: + gfx = str(job.get("gfx") or "").strip().lower() + cu_num = int(job.get("cu_num") or 0) + if not gfx: + gfx = job_arch(cu_num) + legacy_cu_nums.add(cu_num) + + # A job with no recorded CU count (0) matches every target of its arch. + matches = { + (target_gfx, target_cu) + for target_gfx, target_cu in targets + if gfx == target_gfx and (target_cu is None or cu_num in (0, target_cu)) + } + if matches: + selected.append(job) + covered.update(matches) + + for cu_num in sorted(legacy_cu_nums): + print( + f" [WARN] FlyDSL GEMM row has no gfx; inferring {job_arch(cu_num)} " + f"from legacy cu_num={cu_num}. Re-run the tuner to record gfx." + ) + missing = [ + gfx if cu is None else f"{gfx}:{cu}" for gfx, cu in sorted(targets - covered) + ] + if missing: + print( + f" [WARN] The configured GEMM CSVs have no FlyDSL GEMM jobs for " + f"{', '.join(missing)}; those targets will use their existing " + "runtime fallback." + ) + return selected + + def compile_one_config( kernel_name: str, kind: str, @@ -1086,17 +1196,17 @@ def main(): cache_dir = os.path.expanduser( os.environ.get("FLYDSL_RUNTIME_CACHE_DIR", "~/.flydsl/cache") ) - arch = os.environ.get("ARCH") or os.environ.get("GPU_ARCHS") - all_jobs = collect_aot_jobs(csv_paths, parse_csv) - if arch: - # GPU_ARCHS may be a ';'- or ','-separated list (e.g. "gfx942;gfx950"). - arch_set = {a.strip() for a in re.split(r"[;,]", arch) if a.strip()} - n_before = len(all_jobs) - all_jobs = [ - j for j in all_jobs if job_arch(j["cu_num"], j.get("gfx", "")) in arch_set - ] - print(f"[aiter] ARCH={arch}: {len(all_jobs)}/{n_before} jobs match") + n_before = len(all_jobs) + try: + all_jobs = filter_jobs_for_build_targets(all_jobs, arch_wide=True) + except RuntimeError as e: + print(f"Error: {e}") + sys.exit(1) + active = active_target_env() + target_spec = active[1] if active else "" + if target_spec: + print(f"[aiter] targets={target_spec}: {len(all_jobs)}/{n_before} jobs match") print("=" * 72) print("FlyDSL GEMM AOT Pre-compilation") @@ -1107,7 +1217,7 @@ def main(): print(f" {kind} jobs:{' ' * max(1, 15 - len(kind))}{count}") print(f" Total jobs: {len(all_jobs)}") print(f" Cache dir: {cache_dir}") - print(f" Target arch: {arch or '(all archs found in CSVs)'}") + print(f" Target arch: {target_spec or '(all archs found in CSVs)'}") print("=" * 72) total_t0 = time.time() diff --git a/op_tests/test_gemm_codegen.py b/op_tests/test_gemm_codegen.py index 84bc31d54f2..0f5f17e7cf8 100644 --- a/op_tests/test_gemm_codegen.py +++ b/op_tests/test_gemm_codegen.py @@ -25,11 +25,11 @@ """ import contextlib -import inspect import os import sys import tempfile import textwrap +from unittest import mock # Ensure the repo-local aiter is imported, not any system/site-packages install. _REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -156,50 +156,8 @@ def test_get_build_targets(): "gfx942" in GFX_CU_NUM_MAP and "gfx950" in GFX_CU_NUM_MAP, ) - # 1.8 AITER_GPU_TARGETS: two CU counts of one arch, which GPU_ARCHS - # plus a single global CU_NUM cannot express. + # 1.8 Live GPU fallback — requires torch and a GPU; skipped otherwise del os.environ["GPU_ARCHS"] - targets_env = f"{TARGET_B[0]}:{TARGET_B[1]};{TARGET_D[0]}:{TARGET_D[1]}" - os.environ["AITER_GPU_TARGETS"] = targets_env - t = get_build_targets_env() - _check( - f"AITER_GPU_TARGETS={targets_env} → two targets same gfx", - t == [TARGET_B, TARGET_D], - str(t), - ) - - # 1.9 Exact duplicates are removed without changing caller order. - os.environ["AITER_GPU_TARGETS"] = ( - f"{TARGET_D[0]}:{TARGET_D[1]},{TARGET_B[0]}:{TARGET_B[1]};" - f"{TARGET_D[0]}:{TARGET_D[1]}" - ) - t = get_build_targets_env() - _check( - "AITER_GPU_TARGETS preserves first-occurrence order", - t == [TARGET_D, TARGET_B], - str(t), - ) - - # 1.10 Bare entry falls back to the GFX_CU_NUM_MAP default - os.environ["AITER_GPU_TARGETS"] = TARGET_B[0] - t = get_build_targets_env() - _check( - f"AITER_GPU_TARGETS={TARGET_B[0]} → [{TARGET_B}]", t == [TARGET_B], str(t) - ) - - # 1.11 Wins over a conflicting GPU_ARCHS + CU_NUM - os.environ["GPU_ARCHS"] = TARGET_A[0] - os.environ["CU_NUM"] = str(TARGET_C[1]) - os.environ["AITER_GPU_TARGETS"] = f"{TARGET_D[0]}:{TARGET_D[1]}" - t = get_build_targets_env() - _check( - "AITER_GPU_TARGETS wins over GPU_ARCHS + CU_NUM", t == [TARGET_D], str(t) - ) - del os.environ["GPU_ARCHS"] - del os.environ["CU_NUM"] - del os.environ["AITER_GPU_TARGETS"] - - # 1.12 Live GPU fallback — requires torch and a GPU; skipped otherwise try: from aiter.jit.utils.chip_info import get_build_targets @@ -989,51 +947,150 @@ def __init__(self, name): del os.environ["CU_NUM"] +def test_explicit_build_targets(): + cases = [ + # (gpu_target_string, expected_targets) + ("gfx950:128,gfx950:256;gfx950:128", [TARGET_D, TARGET_B]), + ("gfx950", [TARGET_B]), + ] + for gpu_target_string, expected_targets in cases: + with _target_env( + AITER_GPU_TARGETS=gpu_target_string, GPU_ARCHS="gfx942", CU_NUM="80" + ): + targets = get_build_targets_env() + + _check( + f"{gpu_target_string}: explicit targets override GPU_ARCHS and CU_NUM", + targets == expected_targets, + str(targets), + ) + + def test_unmatched_targets(): - _section("2b. unmatched_targets — which build targets have no tuned rows") + df = pd.DataFrame([TARGET_A, TARGET_B], columns=["gfx", "cu_num"]) + cases = [ + (df, [TARGET_A, TARGET_B], []), + (df, [TARGET_A, TARGET_D], ["gfx950:128"]), + # A library with no tuned rows must report every requested target. + (df.iloc[:0], [TARGET_A, TARGET_B], ["gfx942:304", "gfx950:256"]), + ] + for rows, targets, expected in cases: + missing = unmatched_targets(rows, targets) + _check( + f"unmatched targets for {targets} with {len(rows)} tuned rows", + missing == expected, + str(missing), + ) - import pandas as pd - df = pd.DataFrame( - { - "gfx": [TARGET_A[0], TARGET_B[0]], - "cu_num": [TARGET_A[1], TARGET_B[1]], - "libtype": ["ck", "ck"], - } - ) +def test_flydsl_aot_target_filter(): + try: + import aiter.aot.flydsl.common as aot_common + from aiter.aot.flydsl.common import OpKind + from aiter.aot.flydsl.gemm import filter_jobs_for_build_targets + from aiter.jit.utils import chip_info + except ModuleNotFoundError as e: + if e.name not in ("torch", "flydsl"): + raise + print(f" SKIP FlyDSL AOT tests require torch and FlyDSL ({e})") + return + + jobs = [ + {"kernel_name": "k128", "gfx": "gfx950", "cu_num": 128}, + {"kernel_name": "k256", "gfx": "gfx950", "cu_num": 256}, + # Same CU as k128, so gfx must participate in the filter. + {"kernel_name": "other_arch", "gfx": "gfx942", "cu_num": 128}, + ] + cases = [ + # env, arch_wide, expected + ({"AITER_GPU_TARGETS": "gfx950:256;gfx950:128"}, False, ["k128", "k256"]), + ({"AITER_GPU_TARGETS": "gfx950:128;gfx950:64"}, False, ["k128"]), + ({"GPU_ARCHS": "gfx950", "CU_NUM": "128"}, False, ["k128"]), + ({"GPU_ARCHS": "gfx950", "CU_NUM": "128"}, True, ["k128", "k256"]), + ({"GPU_ARCHS": "gfx942,gfx950"}, False, ["k256"]), + ] + for env, arch_wide, expected in cases: + with ( + _target_env(**env), + mock.patch.object(aot_common, "collect_aot_jobs", return_value=jobs), + mock.patch.object( + chip_info, "get_gfx_runtime", side_effect=RuntimeError("No GPU") + ), + ): + if arch_wide: + selected = filter_jobs_for_build_targets(jobs, arch_wide=True) + else: + selected = aot_common._collect_aot_jobs_for(OpKind.GEMM) + + names = [job["kernel_name"] for job in selected] + mode = "CLI" if arch_wide else "packaging" + _check(f"FlyDSL {mode}: {env}", names == expected, str(names)) + + +def test_flydsl_aot_legacy_rows(): + try: + from aiter.aot.flydsl.gemm import filter_jobs_for_build_targets + except ModuleNotFoundError as e: + if e.name not in ("torch", "flydsl"): + raise + print(f" SKIP FlyDSL AOT tests require torch and FlyDSL ({e})") + return + + jobs = [ + {"kernel_name": "legacy", "gfx": "", "cu_num": 128}, + {"kernel_name": "no_cu", "gfx": "gfx950", "cu_num": 0}, + {"kernel_name": "other_arch", "gfx": "gfx942", "cu_num": 0}, + ] + with _target_env(AITER_GPU_TARGETS="gfx950:128;gfx950:256"): + selected = filter_jobs_for_build_targets(jobs) _check( - "every target tuned -> nothing reported", - unmatched_targets(df, [TARGET_A, TARGET_B]) == [], - ) - _check( - "untuned target is named as 'gfx:cu_num'", - unmatched_targets(df, [TARGET_A, TARGET_D]) == [f"{TARGET_D[0]}:{TARGET_D[1]}"], - str(unmatched_targets(df, [TARGET_A, TARGET_D])), - ) - # The frame reaching the check is empty when the CSV has no rows of that - # libtype at all -- the case the old len(tune_df) guard swallowed. - empty = df[df["libtype"] == "cktile"] - _check( - "wholly missing libtype reports every target", - unmatched_targets(empty, [TARGET_A, TARGET_B]) - == [f"{TARGET_A[0]}:{TARGET_A[1]}", f"{TARGET_B[0]}:{TARGET_B[1]}"], - str(unmatched_targets(empty, [TARGET_A, TARGET_B])), - ) - _check( - "filter_tune_df takes only the frame and the targets", - list(inspect.signature(filter_tune_df).parameters) == ["tune_df", "targets"], - str(inspect.signature(filter_tune_df)), + "legacy gfx inference and missing CU counts preserve matching rows once", + selected == jobs[:2], + str(selected), ) +def test_flydsl_aot_invalid_archs(): + try: + from aiter.aot.flydsl.gemm import filter_jobs_for_build_targets + except ModuleNotFoundError as e: + if e.name not in ("torch", "flydsl"): + raise + print(f" SKIP FlyDSL AOT tests require torch and FlyDSL ({e})") + return + + for spec in (" ; ", "gfx9999", "native;gfx942"): + with _target_env(GPU_ARCHS=spec): + raised = False + try: + filter_jobs_for_build_targets([]) + except RuntimeError: + raised = True + + _check(f"FlyDSL rejects GPU_ARCHS={spec!r}", raised) + + +@contextlib.contextmanager +def _target_env(**values): + with mock.patch.dict(os.environ): + for name in ("AITER_GPU_TARGETS", "GPU_ARCHS", "CU_NUM", "ARCH"): + os.environ.pop(name, None) + os.environ.update(values) + yield + + # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- if __name__ == "__main__": test_get_build_targets() + test_explicit_build_targets() test_unmatched_targets() + test_flydsl_aot_target_filter() + test_flydsl_aot_legacy_rows() + test_flydsl_aot_invalid_archs() test_gen_instances_filter( csv_path=REPRO_CSV, target_a=TARGET_C,