diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 00000000..c0251fff --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,48 @@ +name: CI + +on: + pull_request: + push: + branches: + - main + +permissions: + contents: read + +jobs: + python-syntax: + name: Python syntax (${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: + - "3.10" + - "3.13" + + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Compile Python sources + run: python -m compileall -q . + + cpu-tests: + name: CPU tests + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + - uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + - name: Install development dependencies + run: uv sync --extra dev + - name: Run CPU test suite + run: uv run pytest -m "not gpu" diff --git a/.gitignore b/.gitignore index 345fb412..b4c13cb5 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,4 @@ traces/ # Workspace (runtime artifacts) workspace/ +uv.lock diff --git a/CHANGELOG.md b/CHANGELOG.md index 12cf02e3..f57746ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,119 @@ # Changelog +## Unreleased (downstream) + +### Universal workload and discovery foundation + +- Added versioned, metadata-only FastVideo workload manifests under + `workloads/` (Wan 2.1 T2V 1.3B 480p and LTX 480p), a resume-safe + native-versus-optimized launcher bridge, structured generation results with + full-frame parity enforcement, and the `workload.py` CLI + (`validate`, `show`, `run-ab`, `validate-result`) +- Added the discovery layer under `autokernel/discovery/`: metadata-only + discovery report schema with stable graph fingerprints, fail-closed + pure-tensor safety checks, model-independent CPU FX region capture, + profiler-export ingestion, profiler-to-region timing correlation, and + Amdahl-style impact ranking with a configurable end-to-end floor through + the `discovery.py` CLI (`validate`, `rank`, `ingest-profiler`) +- Profiled Wan 2.1 T2V 1.3B and LTX-2 distilled T2V generations on GB200 + through the model-agnostic path and ingested both into validated, ranked + discovery reports + +### Model optimization campaigns + +- Added a versioned, metadata-only campaign contract with strict validation, + impact ranking, legacy orchestration-plan generation, and trusted + starter-kernel preparation through `campaign.py` +- Added a one-command, time-bounded and resumable overnight campaign runner + with per-target benchmark instructions, durable logs, terminal receipts, and + a consolidated morning report + +### MotionKernel identity + +- Renamed the downstream distribution to MotionKernel and reset its independent + package version to `0.1.0` +- Repositioned the project around verified GPU kernel optimization for video + generation models, with FastVideo as the first target integration +- Preserved the `autokernel` Python import namespace as a temporary + compatibility boundary and retained the upstream MIT license and attribution + +### Custom operation registry + +- Added the `autokernel` package with `autokernel/specs/`: a typed `KernelSpec` + that owns one operation's reference, deterministic inputs, sizes, dtypes, + tolerances, edge cases, FLOP/byte accounting, profiler shape aliases and + starter kernels +- Added `KernelRegistry` with deterministic ordering, duplicate detection and + per-command isolation (`create_builtin_registry()`); registry discovery imports + no `torch` and initializes no GPU, so it works on CPU-only machines +- Migrated all nine built-in operations (`matmul`, `softmax`, `layernorm`, + `flash_attention`, `fused_mlp`, `cross_entropy`, `rotary_embedding`, + `rmsnorm`, `reduce`) to specifications; `bench.py` and `extract.py` now read + metadata only from those specs +- Removed the duplicated metadata maps from `extract.py` (`SHAPE_KEYS`, + `SHAPE_ALIAS_MAP`, `TOLERANCES_MAP`, `FLOPS_FN_SRC`, `BYTES_FN_SRC`, + `SPEEDUP_ESTIMATES`, hard-coded default shapes). FLOP/byte accounting is a + serializable expression tree instead of stored Python source strings +- Added `--spec LOCATOR` and `--spec-override` to `bench.py` and `extract.py`. + Precedence is `--spec`, then `--kernel`, then `kernel.py::KERNEL_TYPE`, so + existing invocations are unchanged. `--help` never imports an external spec +- Added `examples/custom_ops/add.py` (external spec) and + `examples/custom_ops/add_kernel.py` (its starter kernel) +- Added a CPU test suite (`uv run pytest -m "not gpu"`) that freezes the + built-in metadata against its pre-refactor values, plus a `dev` extra and a + CI job to run it. GPU tests are marked `gpu` +- `bench.py` keeps a deprecated `KERNEL_CONFIGS` view derived from the registry + for out-of-tree callers +- `Tolerance` rejects negative, NaN, and infinite `atol`/`rtol` values, so a + malformed specification cannot silently disable the correctness gate +- All declared built-in, edge-case and default-shape dimensions must be + positive integers; empty starter-kernel mappings are explicitly supported + for benchmark-only external specifications + +### Generalized verification + +- Added deterministic comparison for tensor, tuple, list, dictionary, + named-tuple and nested output trees, including exact metadata comparison and + per-leaf NaN, infinity and error diagnostics +- Added versioned production shape corpora, append and corpus-only benchmark + modes, and weighted aggregates that remain separated by dtype +- Added optional `BackwardSpec` gradient verification with deterministic + upstream gradients and per-input diagnostics +- Added optional `CompileSpec` verification and `--check-compile`; candidates + compile with full-graph mode by default, run at least twice, reuse one + compiled callable for dynamic shapes and compare through the normal output + tree gate outside performance timing +- Added `FORWARD_CORRECTNESS`, `BACKWARD_CORRECTNESS` and + `COMPILE_CORRECTNESS` console verdicts +- Added schema-versioned, atomic JSON results under + `workspace/bench_result.json`, configurable with `--result-json` +- Added `examples/custom_ops/affine.py`, its candidate and a metadata-only + shape corpus as a structured-output, backward and compile fixture +- Made the float32 matmul starter request IEEE dot inputs instead of Triton's + TF32 default, kept strict BF16 LayerNorm parity through an explicit PyTorch + fallback pending a Welford Triton implementation, and made the affine + fixture's residual rounding stable under Inductor fusion +- Kept the top-level `profile.py` CLI compatible with the standard-library + `profile` API so importing `cProfile` and initializing `torch.compile` from + the repository root no longer fails + +### Wan kernel fusion + +- Added the first production video-DiT operation specification: Wan's + post-self-attention gated residual update plus FP32 affine LayerNorm +- Added specifications, production corpora, and Triton starters for Wan's + modulated pre-attention LayerNorm and post-MLP gated residual, completing the + first three-target Wan optimization pack +- Validated both new starters across their complete production corpora on + GB200, with all correctness stages passing and weighted isolated speedups of + 10.449x and 10.628x respectively +- Added a metadata-only shape corpus covering Wan 2.1 1.3B and 14B at common + 480p token counts, including four-way sequence-parallel layouts +- Added a structured-output Triton baseline that returns both the normalized + activation and updated residual stream in the model dtype +- Validated the full production corpus on GB200 with all correctness stages + passing and a weighted 8.638x operator speedup over eager PyTorch + ## v1.3.0 -- 2026-03-13 ### AMD ROCm GPU Support (PR #3 by @andyluo7) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 00000000..036edb4a --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,53 @@ +# Contributing to MotionKernel + +## Development workflow + +1. Start from an up-to-date `main`. +2. Create a focused branch. +3. Keep framework changes separate from generated kernel experiments. +4. Run CPU validation before pushing. +5. Run the relevant GPU correctness and performance suites before promoting a + kernel. + +Do not run autonomous experiments in a checkout containing unrelated or +uncommitted work. Use a disposable clone or Git worktree so an experiment can +be abandoned without affecting development state. + +## Validation levels + +### CPU baseline + +The baseline check requires no GPU and verifies that tracked Python sources +compile: + +```bash +python -m compileall -q . +``` + +### GPU correctness + +GPU changes must run the relevant benchmark correctness stages across their +declared shapes, dtypes, layouts, and edge cases. Multi-output or training +operations must also validate every returned tensor and requested gradient. + +### Performance + +Performance claims must include: + +- GPU model and compute capability; +- PyTorch, Triton, CUDA, and driver versions; +- input shapes, dtypes, and layouts; +- warmup and measurement methodology; +- median latency and variance; and +- the exact baseline being compared. + +Generated candidates are experimental until their correctness and performance +results are reproducible. + +## Git safety + +- Never force-push shared branches. +- Never push changes to the `upstream` remote. +- Never use a destructive reset outside an isolated experiment branch or + disposable worktree. +- Preserve the MIT license and upstream attribution. diff --git a/DOWNSTREAM.md b/DOWNSTREAM.md new file mode 100644 index 00000000..7fb8545e --- /dev/null +++ b/DOWNSTREAM.md @@ -0,0 +1,60 @@ +# MotionKernel provenance + +MotionKernel is an independently maintained, MIT-licensed downstream fork of +[RightNow-AI/AutoKernel](https://github.com/RightNow-AI/autokernel). Its focus +is GPU kernel discovery, optimization, verification, and packaging for video +generation models. + +## Provenance + +- Upstream repository: `https://github.com/RightNow-AI/autokernel` +- Initial downstream base: `7843582` (`test hf kernels export`) +- License: MIT +- Original copyright: Copyright (c) 2026 RightNow AI + +The upstream `LICENSE` file is preserved. Source files substantially derived +from upstream remain covered by that notice. + +## MotionKernel direction + +MotionKernel is intended to become a video-first, framework-agnostic platform +for discovering, testing, tuning, and exporting production GPU kernels. Its +initial work focuses on: + +- external custom-operation specifications; +- multi-output, backward, determinism, and compile verification; +- production shape corpora captured from real models; +- modulated normalization, gated-residual, attention, and layout fusion; +- architecture-aware tuning and reproducible experiment records; and +- clean export into runtime kernel packages for FastVideo, Diffusers, and other + PyTorch video runtimes. + +The optimization platform and shipped runtime kernels are separate products: +the platform searches and validates candidates, while downstream applications +consume only promoted kernel implementations. + +The initial model families are Wan, LTX-Video, Cosmos, and Kandinsky. Listing a +model as a target does not imply complete support: support is earned through a +published integration, representative workload corpus, correctness results, +and an end-to-end benchmark. + +## Compatibility identity + +The distribution is named `motionkernel`. The Python import namespace remains +`autokernel` temporarily so existing specifications, scripts, and downstream +users do not break during the project transition. A future namespace migration +will include a compatibility release and explicit upgrade instructions. + +## Upstream relationship + +Useful upstream changes can be incorporated without making upstream a release +dependency: + +```bash +git fetch upstream +git switch main +git merge upstream/main +``` + +MotionKernel features do not require upstream approval. Improvements may still +be offered upstream when doing so benefits both projects. diff --git a/README.md b/README.md index 02d766bb..f3b29dfe 100644 --- a/README.md +++ b/README.md @@ -1,21 +1,63 @@ -# AutoKernel - -[![Discord](https://img.shields.io/badge/Discord-Join%20us-5865F2?logo=discord&logoColor=white)](https://discord.gg/UfEyc72t) - -**Autoresearch for GPU kernels.** Give it any PyTorch model, go to sleep, wake up to optimized Triton or CUDA C++ kernels. - -![AutoKernel Progress](progress.png) - -Inspired by [@karpathy/autoresearch](https://github.com/karpathy/autoresearch) -- which demonstrated autonomous AI agents for LLM training research. AutoKernel applies the same philosophy to GPU kernel optimization: agent modifies one file, runs a fixed evaluation, keeps or reverts, repeats forever. +# MotionKernel + +**Verified GPU kernel optimization for video generation models.** + +> [!NOTE] +> MotionKernel is an independently maintained, MIT-licensed fork of +> [RightNow-AI/AutoKernel](https://github.com/RightNow-AI/autokernel), focused +> on GPU kernel optimization for video diffusion transformers, video VAEs, and +> production video-generation workloads. It preserves the upstream license and +> attribution. See +> [DOWNSTREAM.md](DOWNSTREAM.md) for provenance and [ROADMAP.md](ROADMAP.md) +> for the video-first project plan. + +MotionKernel profiles real model executions, captures production tensor shapes, +develops and tunes Triton or CUDA C++ kernels, and verifies numerical +correctness, gradients, `torch.compile` compatibility, performance, and +end-to-end integration behavior. + +The first target integration is +[FastVideo](https://github.com/hao-ai-lab/FastVideo), beginning with Wan and +expanding to LTX-Video, Cosmos, and Kandinsky. The optimization platform remains +framework-agnostic: promoted kernels can be consumed by FastVideo, Diffusers, +or other PyTorch video runtimes without requiring the research harness at +inference time. + +## Project Status + +The reusable kernel specification registry, production shape corpora, +structured-output comparison, backward verification, `torch.compile` +verification, and reproducible JSON result artifacts are implemented. + +The first video-specific pack covers three Wan boundaries: modulated +pre-attention LayerNorm, post-attention gated residual plus LayerNorm, and the +post-MLP gated residual. All three have been validated across their +production shape corpora on an NVIDIA GB200 (see +[docs/WAN_KERNEL_RESULTS.md](docs/WAN_KERNEL_RESULTS.md)). These are isolated +operator results; complete model packs still require end-to-end benchmark +publication before support is claimed. + +A model-independent discovery foundation is also in place: declarative +FastVideo workload manifests, a resumable native-versus-optimized launcher +bridge, metadata-only profiler ingestion, CPU FX graph capture with stable +fingerprints, and impact ranking with an end-to-end floor. Graph-derived +spec generation and generic artifact dispatch are the next stages. + +MotionKernel currently retains the `autokernel` Python import namespace for +compatibility with the upstream project. The import namespace will only move +after a documented migration path exists. ## How It Works -Give AutoKernel any PyTorch model. It will: +Give MotionKernel a PyTorch model or an external operation specification. It +will: 1. **Profile** the model to find which GPU kernels are bottlenecks -2. **Extract** each bottleneck as a standalone Triton or CUDA C++ kernel -3. **Optimize** each kernel autonomously (edit, benchmark, keep/revert -- forever) -4. **Verify** end-to-end correctness and report the total speedup +2. **Capture** representative shapes, dtypes, layouts, and environment metadata +3. **Extract** each bottleneck as a standalone Triton or CUDA C++ kernel +4. **Optimize** candidates through an iterative edit, benchmark, and keep/revert loop +5. **Verify** outputs, optional gradients, compilation, and end-to-end behavior +6. **Promote** reproducible kernels into runtime integration packages The agent reads `program.md` -- the "research org code" -- which contains comprehensive instructions for autonomous operation. It edits `kernel.py` one kernel at a time, runs `bench.py` (fixed benchmark with 5-stage correctness checks + roofline analysis), and either keeps or reverts the change. The orchestrator decides when to move to the next kernel using Amdahl's law. @@ -30,8 +72,8 @@ Each experiment takes ~90 seconds. That's ~40 experiments/hour, ~320 overnight, curl -LsSf https://astral.sh/uv/install.sh | sh # Clone and setup -git clone https://github.com/RightNow-AI/autokernel.git -cd autokernel +git clone https://github.com/aryan5v/motionkernel.git +cd motionkernel uv sync # One-time setup: test data + baselines @@ -58,7 +100,7 @@ Read program.md and let's kick off a new experiment. Start with setup. The agent will: 1. Profile your model and present the optimization plan -2. Create a branch (e.g., `autokernel/mar10-llama7b`) +2. Create a branch (e.g., `motionkernel/wan-gated-residual`) 3. Optimize each bottleneck kernel in priority order 4. Verify end-to-end correctness and report total speedup @@ -98,9 +140,208 @@ Any PyTorch ──> Rank kernels ──> Generate baseline ──> Optimiz Each has a PyTorch reference in `reference.py`, a starter Triton kernel in `kernels/`, and a starter CUDA C++ kernel in `kernels/cuda/`. +Every built-in operation is described by one `KernelSpec` in +`autokernel/specs/builtins.py`. That +specification is the single source of truth for sizes, dtypes, tolerances, edge cases, +FLOP/byte accounting, profiler shape aliases and starter kernels -- `bench.py` and +`extract.py` read it instead of carrying their own per-operation tables. + +## Custom Operations + +Any operation can be added from outside the repository, without editing `bench.py`, +`extract.py`, `reference.py` or any central map. Write a `KernelSpec` and export it: + +```python +# my_ops/gelu_tanh.py +from autokernel.specs import DT_BYTES, EdgeCase, KernelSpec, Tolerance, resolve_torch_dtype, size + + +def gelu_tanh_ref(x): + import torch + return 0.5 * x * (1 + torch.tanh(0.7978845608 * (x + 0.044715 * x ** 3))) + + +def gen_inputs(size_map, dtype, device, seed=42): + import torch + torch.manual_seed(seed) + rows, cols = size_map["rows"], size_map["cols"] + return {"x": torch.randn(rows, cols, device=device, dtype=resolve_torch_dtype(dtype))} + + +SPEC = KernelSpec( + name="gelu_tanh", + reference_fn=gelu_tanh_ref, + input_generator=gen_inputs, + sizes={ + "small": {"rows": 256, "cols": 512}, + "medium": {"rows": 1024, "cols": 1024}, + "large": {"rows": 4096, "cols": 4096}, + }, + dtypes=("float16", "bfloat16", "float32"), + tolerances={ + "float16": Tolerance(atol=1e-3, rtol=1e-3), + "bfloat16": Tolerance(atol=2e-3, rtol=2e-3), + "float32": Tolerance(atol=1e-5, rtol=1e-5), + }, + flops_fn=8 * size("rows") * size("cols"), + bytes_fn=2 * size("rows") * size("cols") * DT_BYTES, + edge_cases=(EdgeCase(name="edge_1023", size={"rows": 1023, "cols": 1023}),), + shape_keys=("rows", "cols"), + starter_kernels={"triton": "my_ops/gelu_tanh_kernel.py"}, +) +``` + +Then point the existing commands at it with `--spec LOCATOR`, where a locator is +`path/to/spec.py:ATTRIBUTE` or `package.module:ATTRIBUTE`: + +```bash +# benchmark a candidate kernel.py against the external spec +cp examples/custom_ops/add_kernel.py kernel.py +uv run bench.py --spec examples/custom_ops/add.py:SPEC --quick + +# generate a starter kernel file for it under workspace/ +uv run extract.py --spec examples/custom_ops/add.py:SPEC --top 1 +``` + +Operation selection precedence is `--spec`, then `--kernel`, then `kernel.py::KERNEL_TYPE`, +so existing invocations keep working unchanged. A spec whose name collides with a built-in +is rejected unless `--spec-override` is passed. `ATTRIBUTE` may be a `KernelSpec` or a +zero-argument callable returning one. + +Requirements the harness validates before allocating anything on the GPU: an +identifier-like name, `small`/`medium`/`large` sizes, canonical dtype names +(`float16`, `bfloat16`, `float32`), a tolerance for every declared dtype, size keys that +match `shape_keys`, positive integer dimensions in every built-in, edge and default +shape, and starter-kernel files that exist. `starter_kernels={}` is valid for a +benchmark-only specification; extraction skips a backend whose starter is not declared. + +A complete, runnable example lives in `examples/custom_ops/add.py` (spec) and +`examples/custom_ops/add_kernel.py` (starter kernel). + +Note that loading a spec executes the Python file you point at, exactly like running +`python that_file.py`. Only pass locators you trust. + +## Model Optimization Campaigns + +FastVideo and other runtimes can export a versioned campaign containing only +operation identities, tensor shape/layout signatures, call counts, aggregate +timings, and environment identity. Validate and rank a campaign without loading +PyTorch or executing any referenced Python: + +```bash +uv run campaign.py validate /path/to/campaign.json +uv run campaign.py rank /path/to/campaign.json +uv run campaign.py plan /path/to/campaign.json +``` + +Once the campaign and its spec locators have been reviewed, prepare all ranked +starter kernels and the existing orchestration state in one step: + +```bash +uv run campaign.py prepare /path/to/campaign.json --trust-specs +uv run orchestrate.py plan +``` + +Preparation writes `workspace/optimization_plan.json`, one candidate kernel per +ranked target, and `workspace/campaign_receipt.json`. The explicit trust flag is +required because a Python spec locator executes code. Continue with +`program.md` for the autonomous experiment loop; every candidate still passes +the fixed correctness gates in `bench.py` before a result can be kept. + +For an unattended, resumable run, preparation and the agent loop are one +command: + +```bash +uv run campaign.py run /path/to/wan-campaign.json --trust-specs --budget-hours 10 +``` + +Like `prepare`, `run` refuses to load a campaign's Python spec locators +without the explicit `--trust-specs` flag, because loading a spec executes the +Python file it points at. Use `--dry-run` to inspect +`workspace/overnight_prompt.md` without launching an agent, and `--resume` +after an interrupted run. A non-`completed` terminal status is reported as +`CAMPAIGN_RUN: FAIL` with a non-zero exit code. By default the runner invokes +the Codex CLI; `--agent-command` supports trusted alternatives with `{repo}` +and `{prompt_file}` placeholders. The next morning, inspect +`workspace/morning_report.md`, the terminal receipt, agent log, and verified +`kernel___optimized.py` artifacts in the same directory. + +## Workloads and Discovery + +The universal optimization path starts from a declarative workload manifest +instead of model-specific scripts. A manifest in `workloads/` describes one +reproducible FastVideo generation benchmark: model identifier, task and +prompt reference, resolution, frame count, inference steps, seed, dtype, +warmup and measured repetitions, and the output-parity policy. Canonical +manifests exist for Wan 2.1 T2V 1.3B 480p and LTX 480p. + +```bash +# validate and inspect a manifest +uv run workload.py validate workloads/wan_t2v_1.3b_480p.yaml +uv run workload.py show workloads/wan_t2v_1.3b_480p.yaml + +# run a resumable native-versus-optimized A/B through a FastVideo checkout +uv run workload.py run-ab --fastvideo-checkout /path/to/FastVideo \ + --workload workloads/wan_t2v_1.3b_480p.yaml --output workspace/wan_ab + +# validate a structured generation result +uv run workload.py validate-result workspace/wan_ab/native_result.json +``` + +Discovery reports are metadata-only records of where a profiled generation +actually spends time: profiler operator rows, captured FX graph regions with +stable fingerprints, graph breaks, and unsupported operations. They never +contain prompts, weights, activations, tensor values, or model outputs. + +```bash +# convert a FastVideo profiler export into a discovery report +uv run discovery.py ingest-profiler workspace/profiler_export.json \ + --output workspace/discovery_report.json + +# validate and rank candidate regions by optimistic end-to-end impact +uv run discovery.py validate workspace/discovery_report.json +uv run discovery.py rank workspace/discovery_report.json --impact-floor 0.005 +``` + +Ranking uses measured production frequency and an Amdahl-style ceiling: a +candidate is only worth searching when its optimistic end-to-end improvement +clears the impact floor (0.5% by default). Regions with mutation, collectives, +data-dependent control flow, or unknown aliasing are rejected fail-closed +before they can enter the search pipeline. + +## Generalized Verification + +The harness compares complete output trees, including nested tensors and metadata. An +external spec may also declare `BackwardSpec` and `CompileSpec` policies. The structured +affine example exercises both: + +```bash +cp examples/custom_ops/affine_kernel.py kernel.py + +# forward output-tree comparison plus production-shape corpus +uv run bench.py --spec examples/custom_ops/affine.py:SPEC \ + --shape-corpus examples/custom_ops/affine_corpus.json --quick + +# optional gradient and full-graph compile gates +uv run bench.py --spec examples/custom_ops/affine.py:SPEC \ + --check-backward --check-compile --quick +``` + +Compile verification calls `torch.compile` with `fullgraph=True` by default and runs +before performance timing. A dynamic `CompileSpec` exercises two declared shapes through +the same compiled callable. If `torch.compile` is unavailable, the result is +`UNSUPPORTED`, never `PASS`. + +Every normal benchmark run atomically writes a schema-versioned JSON record to +`workspace/bench_result.json`; override it with `--result-json PATH`. It includes +forward leaf errors, optional gradient and compile results, shape-corpus identity, +environment metadata and performance results. Stable console verdicts are +`FORWARD_CORRECTNESS`, `BACKWARD_CORRECTNESS` and `COMPILE_CORRECTNESS`. + ## Example Models -Self-contained model definitions ship with AutoKernel (no `transformers` library needed): +Self-contained model definitions inherited by MotionKernel require no +`transformers` library: | Model | File | Params | Usage | |-------|------|--------|-------| @@ -119,9 +360,9 @@ uv run profile.py --module transformers --class-name AutoModelForCausalLM \ ## KernelBench Integration -AutoKernel integrates with [KernelBench](https://github.com/ScalingIntelligence/KernelBench), +MotionKernel retains integration with [KernelBench](https://github.com/ScalingIntelligence/KernelBench), the standard benchmark for evaluating AI-generated GPU kernels (250+ problems across 4 difficulty -levels). While most KernelBench evaluations use one-shot LLM generation, AutoKernel runs +levels). While most KernelBench evaluations use one-shot LLM generation, MotionKernel runs **50-300+ iterative refinement experiments per problem** -- systematically exploring the optimization space instead of guessing. @@ -175,7 +416,7 @@ kernels upload . --repo_id your-username/my_matmul ## Project Structure ``` -autokernel/ +motionkernel/ kernel.py the file the agent modifies (one kernel at a time) program.md agent instructions -- the "research org code" @@ -183,6 +424,19 @@ autokernel/ reference.py PyTorch reference implementations (ground truth) prepare.py one-time setup: test data, baselines + autokernel/specs/ KernelSpec types, registry, external spec loader, + built-in operation metadata, input generators + autokernel/campaign/ campaign contract, ranking, overnight runner + autokernel/workload/ workload manifest schema, FastVideo launcher bridge, + structured generation results and parity checks + autokernel/discovery/ discovery report schema, FX capture, profiler + ingestion, timing correlation, impact ranking + + campaign.py validate, rank, prepare, and run campaigns + workload.py validate and A/B-run FastVideo workload manifests + discovery.py validate, ingest, and rank discovery reports + workloads/ canonical workload manifests (Wan, LTX) + profile.py profile any PyTorch model, rank kernels by GPU time extract.py extract bottleneck kernels into workspace/ orchestrate.py multi-kernel scheduler (Amdahl's law) @@ -194,6 +448,8 @@ autokernel/ kernels/cuda/ starter CUDA C++ kernels (9 types, tensor core accelerated) kernelbench/ KernelBench integration (bridge, eval harness, scorer) models/ self-contained model definitions (GPT-2, LLaMA, BERT) + examples/custom_ops/ external KernelSpec example + its starter kernel + tests/ CPU test suite (uv run pytest -m "not gpu") workspace/ runtime artifacts (gitignored) ``` @@ -228,11 +484,18 @@ Every experiment is logged to `results.tsv` (tab-separated): ## Credits -This project is **autoresearch for GPU kernels** -- directly inspired by Andrej Karpathy's [autoresearch](https://github.com/karpathy/autoresearch), the original experiment in autonomous AI research agents for LLM training. Karpathy showed that an AI agent can run hundreds of experiments overnight, methodically exploring a search space and logging every result. AutoKernel applies that same loop -- agent edits one file, runs a fixed evaluation, keeps or reverts -- to the domain of GPU kernel optimization with Triton and native CUDA C++. +MotionKernel builds on AutoKernel's **autoresearch for GPU kernels** approach, +which was directly inspired by Andrej Karpathy's +[autoresearch](https://github.com/karpathy/autoresearch). MotionKernel retains +the iterative agent loop while extending the platform toward production video +workloads, explicit operation specifications, representative shape corpora, +and stronger verification. -**KernelBench** integration is based on the work of Simon Guo, Sean Resta, et al. at Stanford's Scaling Intelligence Lab. Their paper ["KernelBench: Can LLMs Write GPU Kernels?"](https://arxiv.org/abs/2502.10517) (2025) established the standard benchmark for evaluating AI-generated GPU kernels. AutoKernel extends this by applying iterative optimization (300+ experiments per problem) instead of one-shot generation. KernelBench dataset and evaluation protocol: [ScalingIntelligence/KernelBench](https://github.com/ScalingIntelligence/KernelBench). +**KernelBench** integration is based on the work of Simon Guo, Sean Resta, et al. at Stanford's Scaling Intelligence Lab. Their paper ["KernelBench: Can LLMs Write GPU Kernels?"](https://arxiv.org/abs/2502.10517) (2025) established the standard benchmark for evaluating AI-generated GPU kernels. The inherited AutoKernel integration applies iterative optimization instead of one-shot generation. KernelBench dataset and evaluation protocol: [ScalingIntelligence/KernelBench](https://github.com/ScalingIntelligence/KernelBench). -Built by [RightNow AI](https://www.rightnowai.co). For enterprise GPU optimization, check out [RightNow Enterprise](https://www.rightnowai.co/forge). +MotionKernel is independently maintained. The original AutoKernel project was +built by [RightNow AI](https://www.rightnowai.co); see +[DOWNSTREAM.md](DOWNSTREAM.md) for the exact fork provenance. ## Changelog @@ -260,4 +523,5 @@ See [CHANGELOG.md](CHANGELOG.md) for full details. ## License -MIT +MIT. MotionKernel retains the original AutoKernel copyright and permission +notice. See [LICENSE](LICENSE) and [DOWNSTREAM.md](DOWNSTREAM.md). diff --git a/ROADMAP.md b/ROADMAP.md new file mode 100644 index 00000000..2f5024ec --- /dev/null +++ b/ROADMAP.md @@ -0,0 +1,82 @@ +# MotionKernel Roadmap + +MotionKernel is the video-kernel optimization layer, not a competing +video-generation runtime. It discovers and validates kernels that frameworks +such as FastVideo and Diffusers can consume. + +Detailed execution instructions for the first two milestones are available in +[docs/WEEK_1_2_AGENT_BRIEF.md](docs/WEEK_1_2_AGENT_BRIEF.md). + +## Milestone 0: standalone foundation + +- Preserve upstream provenance and MIT attribution. +- Maintain separate `origin` and `upstream` remotes. +- Establish the MotionKernel identity while retaining a compatible + `autokernel` import namespace. +- Establish safe contribution and experiment practices. +- Add lightweight CPU-only validation for every change. + +## Milestone 1: custom-operation platform + +- Introduce a `KernelSpec` registry. +- Load operation specifications without editing the core benchmark. +- Move existing built-in operations onto the same registry. +- Define stable interfaces for reference functions, inputs, cases, tolerances, + output comparison, performance metrics, and integration hooks. + +Exit criterion: a new single-output operation can be optimized from an external +specification without changes to `bench.py`, `extract.py`, or `reference.py`. + +## Milestone 2: production verification + +- Compare tensor, tuple, and nested outputs. +- Add optional backward and gradient verification. +- Add deterministic execution checks. +- Add `torch.compile` full-graph compatibility checks. +- Support model-specific replacement adapters. +- Record GPU, software, shape, and benchmark metadata with every result. + +Exit criterion: a custom multi-output operation can be validated in isolation +and inside a model, including backward execution when requested. + +## Milestone 3: video and diffusion transformer kernels + +- Add modulated LayerNorm and RMSNorm. +- Add gated residual updates. +- Add combined gated-residual, normalization, and modulation. +- Cover affine and non-affine variants. +- Cover batch, frame, token, and spatial broadcast layouts. +- Tune FP16 and BF16 paths using FP32 accumulation. + +Exit criterion: promoted kernels pass correctness and compile gates and show a +meaningful speedup on production shape distributions. + +## Milestone 4: model adoption + +- Integrate and benchmark Wan through FastVideo. +- Integrate and benchmark LTX-Video. +- Integrate and benchmark Cosmos. +- Integrate and benchmark Kandinsky. +- Validate single-GPU and sequence-parallel execution. + +Runtime integrations will use exported kernels with native PyTorch fallbacks; +they will not require the optimization platform at inference time. + +## Milestone 5: continuous kernel research + +- Run parallel searches across GPU workers. +- Maintain architecture-specific tuning records. +- Track performance regressions between revisions. +- Promote candidates through experimental, validated, and production stages. +- Expand into attention, MLP, quantization, and communication-aware fusion. + +## Long-term: VideoKernelBench and ecosystem adoption + +- Publish reproducible operator and end-to-end video workload benchmarks. +- Track latency, throughput, memory, compile time, numerical accuracy, and + output-quality regressions across hardware generations. +- Maintain FastVideo and Diffusers adapters with stable replacement APIs. +- Grow shared and model-specific kernel packs without hard-coding model logic + into the optimizer core. +- Add training, lower-precision, multi-GPU, and additional hardware backends + only after the inference path is reliable. diff --git a/autokernel/__init__.py b/autokernel/__init__.py new file mode 100644 index 00000000..73bcd271 --- /dev/null +++ b/autokernel/__init__.py @@ -0,0 +1,17 @@ +"""AutoKernel downstream platform package. + +This package holds the reusable, importable core of the project. The +command-line entry points (``bench.py``, ``extract.py``, ``profile.py``, +``verify.py``) live at the repository root and import from here. + +Nothing in this package may initialize a GPU at import time: registry +discovery and specification inspection must work on a CPU-only machine. + +Derived from the upstream AutoKernel project (MIT License). See LICENSE. +""" + +from __future__ import annotations + +__all__ = ["__version__"] + +__version__ = "1.0.0" diff --git a/autokernel/campaign/__init__.py b/autokernel/campaign/__init__.py new file mode 100644 index 00000000..200c0bef --- /dev/null +++ b/autokernel/campaign/__init__.py @@ -0,0 +1,35 @@ +"""Versioned optimization campaigns produced by model runtimes.""" + +from .runner import ( + build_overnight_prompt, + parse_agent_command, + run_campaign, +) +from .types import ( + CAMPAIGN_SCHEMA_VERSION, + CampaignError, + CampaignTarget, + OptimizationCampaign, + ShapeObservation, + TensorSignature, + load_campaign, + prepare_campaign, + rank_targets, + write_optimization_plan, +) + +__all__ = [ + "CAMPAIGN_SCHEMA_VERSION", + "CampaignError", + "CampaignTarget", + "OptimizationCampaign", + "ShapeObservation", + "TensorSignature", + "build_overnight_prompt", + "load_campaign", + "parse_agent_command", + "prepare_campaign", + "rank_targets", + "run_campaign", + "write_optimization_plan", +] diff --git a/autokernel/campaign/runner.py b/autokernel/campaign/runner.py new file mode 100644 index 00000000..42c51b17 --- /dev/null +++ b/autokernel/campaign/runner.py @@ -0,0 +1,415 @@ +"""Resumable unattended execution for prepared optimization campaigns.""" + +from __future__ import annotations + +import json +import os +import shlex +import shutil +import subprocess +import sys +import tempfile +from collections.abc import Sequence +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from .types import ( + CampaignError, + OptimizationCampaign, + prepare_campaign, + rank_targets, +) + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _write_json_atomic(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as handle: + json.dump(payload, handle, indent=2) + handle.write("\n") + temporary = Path(handle.name) + temporary.replace(path) + + +def _write_text_atomic(path: Path, value: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as handle: + handle.write(value) + temporary = Path(handle.name) + temporary.replace(path) + + +def _corpus_for(operation: str, repo_root: Path) -> str | None: + candidate = repo_root / "models" / f"{operation}_corpus.json" + if candidate.is_file(): + return str(candidate.relative_to(repo_root)) + return None + + +def _reset_runtime_files(workspace: Path) -> None: + for name in ( + "orchestration_state.json", + "aggregate_report.md", + "overnight_agent.log", + "agent_last_message.md", + "morning_report.md", + ): + path = workspace / name + if path.is_file(): + path.unlink() + + +def build_overnight_prompt( + campaign: OptimizationCampaign, + *, + repo_root: Path, + budget_hours: float, +) -> str: + """Build the campaign-specific instructions layered over ``program.md``.""" + rows = [] + for rank, target in enumerate(rank_targets(campaign), start=1): + corpus = _corpus_for(target.operation, repo_root) + bench = f"uv run bench.py --spec {target.spec_locator}" + if corpus is not None: + bench += f" --shape-corpus {corpus}" + rows.append( + "\n".join( + [ + f"{rank}. {target.name}", + f" operation: {target.operation}", + f" candidate: workspace/kernel_{target.operation}_{rank}.py", + f" spec: {target.spec_locator}", + f" benchmark: {bench}", + ( + f" estimated model impact: " + f"{target.impact_pct(campaign.total_profiled_device_time_us):.2f}%" + ), + ] + ) + ) + targets = "\n\n".join(rows) + return f"""\ +You are running a prepared MotionKernel optimization campaign unattended. +Read program.md for the optimization playbook, correctness rules, experiment +discipline, crash recovery, and move-on criteria. This prompt overrides its +interactive Phase A: profiling, target selection, and approval are complete. +Do not ask the user questions. + +Workload: {campaign.workload['workload_id']} +Model: {campaign.workload['model_id']} +Budget: {budget_hours:.2f} hours + +Ranked targets: + +{targets} + +For each target in rank order: +1. Use `uv run orchestrate.py next` to confirm the active target. +2. Copy its prepared candidate to kernel.py. +3. Run the exact target-specific benchmark command shown above. Always retain + its `--spec` and `--shape-corpus` arguments for every experiment. +4. Record the baseline and every kept/reverted experiment with orchestrate.py. +5. Never weaken correctness tolerances or edit references, specs, corpora, + bench.py, verification code, or orchestration code. +6. Save the best passing implementation as + `workspace/kernel___optimized.py` before moving on. +7. Continue until every target is done/plateaued or the budget is nearly + exhausted. Leave enough time to run `uv run orchestrate.py report`. + +Write a concise final summary to the agent output requested by the runner. +""" + + +def _default_codex_command( + repo_root: Path, + prompt: str, + last_message: Path, +) -> list[str]: + executable = shutil.which("codex") + if executable is None: + raise CampaignError( + "Codex CLI was not found; install it or pass --agent-command" + ) + return [ + executable, + "exec", + "-C", + str(repo_root), + "-s", + "workspace-write", + "--output-last-message", + str(last_message), + prompt, + ] + + +def parse_agent_command( + value: str, + *, + repo_root: Path, + prompt_path: Path, +) -> list[str]: + """Parse a shell-like command without invoking a shell.""" + command = shlex.split(value) + if not command: + raise CampaignError("--agent-command must not be empty") + replacements = { + "{repo}": str(repo_root), + "{prompt_file}": str(prompt_path), + } + return [ + argument.replace("{repo}", replacements["{repo}"]).replace( + "{prompt_file}", replacements["{prompt_file}"] + ) + for argument in command + ] + + +def _write_morning_report( + workspace: Path, + receipt: dict[str, Any], + *, + report_stdout: str, +) -> Path: + aggregate = workspace / "aggregate_report.md" + agent_summary = workspace / "agent_last_message.md" + lines = [ + "# MotionKernel Overnight Campaign", + "", + f"Status: **{receipt['status']}**", + f"Workload: `{receipt['workload_id']}`", + f"Started: {receipt['started_at']}", + f"Finished: {receipt['finished_at']}", + f"Budget: {receipt['budget_hours']:.2f} hours", + "", + "## Targets", + "", + ] + for target in receipt["targets"]: + lines.append( + f"- {target['rank']}. `{target['operation']}` " + f"({target['name']})" + ) + if aggregate.is_file(): + lines.extend(["", aggregate.read_text(encoding="utf-8")]) + elif report_stdout.strip(): + lines.extend(["", "## Orchestrator output", "", "```text"]) + lines.extend([report_stdout.strip(), "```"]) + if agent_summary.is_file() and agent_summary.stat().st_size: + lines.extend( + [ + "", + "## Agent summary", + "", + agent_summary.read_text(encoding="utf-8").strip(), + ] + ) + path = workspace / "morning_report.md" + _write_text_atomic(path, "\n".join(lines).rstrip() + "\n") + return path + + +def _campaign_progress(workspace: Path) -> dict[str, Any]: + state_path = workspace / "orchestration_state.json" + if not state_path.is_file(): + return { + "complete": False, + "completed_targets": 0, + "total_targets": 0, + } + try: + state = json.loads(state_path.read_text(encoding="utf-8")) + kernels = state["kernels"] + complete = [ + kernel + for kernel in kernels + if kernel.get("status") in ("done", "skipped") + ] + except (json.JSONDecodeError, KeyError, TypeError): + return { + "complete": False, + "completed_targets": 0, + "total_targets": 0, + } + return { + "complete": bool(kernels) and len(complete) == len(kernels), + "completed_targets": len(complete), + "total_targets": len(kernels), + } + + +def run_campaign( + campaign: OptimizationCampaign, + *, + repo_root: str | Path, + budget_hours: float = 10.0, + resume: bool = False, + dry_run: bool = False, + trust_specs: bool = False, + agent_command: Sequence[str] | None = None, + timeout_seconds: float | None = None, +) -> dict[str, Any]: + """Prepare and run a campaign, preserving logs and a terminal receipt.""" + root = Path(repo_root).resolve() + workspace = root / "workspace" + if not root.joinpath("program.md").is_file(): + raise CampaignError(f"MotionKernel repository not found at {root}") + if budget_hours <= 0: + raise CampaignError("budget_hours must be greater than zero") + + prepared_receipt = workspace / "campaign_receipt.json" + fresh_run = not resume or not prepared_receipt.is_file() + if fresh_run: + _reset_runtime_files(workspace) + prepare_campaign( + campaign, + workspace, + trust_specs=trust_specs, + spec_root=root, + ) + + prompt = build_overnight_prompt( + campaign, repo_root=root, budget_hours=budget_hours + ) + prompt_path = workspace / "overnight_prompt.md" + _write_text_atomic(prompt_path, prompt) + receipt_path = workspace / "overnight_receipt.json" + last_message = workspace / "agent_last_message.md" + log_path = workspace / "overnight_agent.log" + + receipt: dict[str, Any] = { + "schema_version": 1, + "campaign_source": campaign.source, + "workload_id": campaign.workload["workload_id"], + "started_at": _utc_now(), + "budget_hours": budget_hours, + "resume": resume, + "status": "prepared" if dry_run else "running", + "prompt": str(prompt_path), + "log": str(log_path), + "targets": [ + { + "rank": rank, + "name": target.name, + "operation": target.operation, + } + for rank, target in enumerate(rank_targets(campaign), start=1) + ], + } + _write_json_atomic(receipt_path, receipt) + if dry_run: + receipt["finished_at"] = _utc_now() + _write_json_atomic(receipt_path, receipt) + return receipt + + command = ( + list(agent_command) + if agent_command is not None + else _default_codex_command(root, prompt, last_message) + ) + timeout = ( + timeout_seconds + if timeout_seconds is not None + else budget_hours * 60 * 60 + ) + env = os.environ.copy() + env["AUTOKERNEL_CAMPAIGN"] = campaign.source + env["AUTOKERNEL_BUDGET_HOURS"] = str(budget_hours) + timed_out = False + returncode: int | None = None + launch_error: str | None = None + with log_path.open("a", encoding="utf-8") as log: + try: + process = subprocess.Popen( + command, + cwd=root, + env=env, + stdout=log, + stderr=subprocess.STDOUT, + text=True, + ) + except OSError as exc: + launch_error = str(exc) + log.write(f"agent launch failed: {exc}\n") + else: + try: + returncode = process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + timed_out = True + process.terminate() + try: + returncode = process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + returncode = process.wait() + + if launch_error is not None: + receipt.update( + { + "finished_at": _utc_now(), + "status": "agent_launch_failed", + "error": launch_error, + "agent_returncode": None, + "timed_out": False, + } + ) + morning_report = _write_morning_report( + workspace, receipt, report_stdout="" + ) + receipt["morning_report"] = str(morning_report) + _write_json_atomic(receipt_path, receipt) + return receipt + + report = subprocess.run( + [sys.executable, "orchestrate.py", "report"], + cwd=root, + check=False, + capture_output=True, + text=True, + ) + progress = _campaign_progress(workspace) + if timed_out: + status = "budget_exhausted" + elif returncode != 0: + status = "agent_failed" + elif progress["complete"]: + status = "completed" + else: + status = "incomplete" + receipt.update( + { + "finished_at": _utc_now(), + "status": status, + "agent_returncode": returncode, + "timed_out": timed_out, + "report_returncode": report.returncode, + "aggregate_report": str(workspace / "aggregate_report.md"), + "progress": progress, + } + ) + morning_report = _write_morning_report( + workspace, + receipt, + report_stdout=report.stdout, + ) + receipt["morning_report"] = str(morning_report) + _write_json_atomic(receipt_path, receipt) + return receipt diff --git a/autokernel/campaign/types.py b/autokernel/campaign/types.py new file mode 100644 index 00000000..83e9ebe0 --- /dev/null +++ b/autokernel/campaign/types.py @@ -0,0 +1,720 @@ +"""Portable FastVideo-to-AutoKernel optimization campaign contract. + +The campaign contains metadata only: operation identities, tensor +shape/layout signatures, call counts, timing aggregates, and workload/runtime +identity. It must never contain tensor values, prompts, model weights, or +credentials. +""" + +from __future__ import annotations + +import json +import math +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +CAMPAIGN_SCHEMA_VERSION = 1 + +_TOP_LEVEL_FIELDS = { + "schema_version", + "producer", + "workload", + "environment", + "total_profiled_device_time_us", + "targets", +} +_TARGET_FIELDS = { + "name", + "operation", + "kind", + "spec_locator", + "total_device_time_us", + "self_device_time_us", + "calls", + "requires_backward", + "observations", + "attributes", +} +_OBSERVATION_FIELDS = { + "name", + "count", + "total_device_time_us", + "inputs", + "tags", +} +_TENSOR_FIELDS = { + "name", + "shape", + "stride", + "dtype", + "device_type", + "requires_grad", +} +_TARGET_KINDS = {"operator", "fusion", "graph_fragment"} +_OPERATION_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +_FORBIDDEN_METADATA_KEYS = { + "credential", + "credentials", + "data", + "password", + "prompt", + "secret", + "secrets", + "tensor_values", + "token", + "values", + "weights", +} + + +class CampaignError(ValueError): + """Raised when a campaign is malformed or unsafe to execute.""" + + +def _fail(source: object, location: str, message: str) -> CampaignError: + return CampaignError(f"optimization campaign {source!r}: {location}: {message}") + + +def _mapping( + value: Any, + source: object, + location: str, + *, + non_empty: bool = False, +) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or (non_empty and not value): + qualifier = "non-empty " if non_empty else "" + raise _fail(source, location, f"must be a {qualifier}object") + for key in value: + if not isinstance(key, str) or not key: + raise _fail(source, location, "keys must be non-empty strings") + return value + + +def _unknown_fields( + raw: Mapping[str, Any], + allowed: set[str], + source: object, + location: str, +) -> None: + unknown = sorted(set(raw) - allowed) + if unknown: + raise _fail(source, location, f"unknown field(s) {unknown}") + + +def _text(value: Any, source: object, location: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise _fail(source, location, "must be a non-empty string") + return value + + +def _finite_non_negative(value: Any, source: object, location: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise _fail(source, location, "must be a finite non-negative number") + normalized = float(value) + if not math.isfinite(normalized) or normalized < 0: + raise _fail(source, location, "must be a finite non-negative number") + return normalized + + +def _positive_int(value: Any, source: object, location: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise _fail(source, location, "must be a positive integer") + return value + + +def _metadata_value(value: Any, source: object, location: str) -> Any: + """Validate JSON metadata while excluding common content/secret fields.""" + if value is None or isinstance(value, (str, bool)): + return value + if isinstance(value, (int, float)) and not isinstance(value, bool): + if not math.isfinite(float(value)): + raise _fail(source, location, "numbers must be finite") + return value + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [ + _metadata_value(item, source, f"{location}[{index}]") + for index, item in enumerate(value) + ] + if isinstance(value, Mapping): + result = {} + for key, item in _mapping(value, source, location).items(): + if key.lower() in _FORBIDDEN_METADATA_KEYS: + raise _fail( + source, + f"{location}.{key}", + "content or secret fields are forbidden", + ) + result[key] = _metadata_value(item, source, f"{location}.{key}") + return result + raise _fail(source, location, "must contain JSON metadata only") + + +@dataclass(frozen=True) +class TensorSignature: + """One tensor's layout metadata without any tensor contents.""" + + name: str + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: str + device_type: str + requires_grad: bool = False + + @classmethod + def from_dict( + cls, + raw_value: Any, + *, + source: object, + location: str, + ) -> "TensorSignature": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _TENSOR_FIELDS, source, location) + name = _text(raw.get("name"), source, f"{location}.name") + dtype = _text(raw.get("dtype"), source, f"{location}.dtype") + device_type = _text( + raw.get("device_type"), source, f"{location}.device_type" + ) + + shape_raw = raw.get("shape") + if not isinstance(shape_raw, Sequence) or isinstance( + shape_raw, (str, bytes) + ): + raise _fail(source, f"{location}.shape", "must be a list of dimensions") + shape: list[int] = [] + for index, dimension in enumerate(shape_raw): + if ( + isinstance(dimension, bool) + or not isinstance(dimension, int) + or dimension < 0 + ): + raise _fail( + source, + f"{location}.shape[{index}]", + "must be a non-negative integer", + ) + shape.append(dimension) + + stride_raw = raw.get("stride") + if not isinstance(stride_raw, Sequence) or isinstance( + stride_raw, (str, bytes) + ): + raise _fail(source, f"{location}.stride", "must be a list of strides") + stride: list[int] = [] + for index, value in enumerate(stride_raw): + if isinstance(value, bool) or not isinstance(value, int): + raise _fail( + source, + f"{location}.stride[{index}]", + "must be an integer", + ) + stride.append(value) + if len(stride) != len(shape): + raise _fail( + source, + f"{location}.stride", + "must have the same length as shape", + ) + + requires_grad = raw.get("requires_grad", False) + if not isinstance(requires_grad, bool): + raise _fail(source, f"{location}.requires_grad", "must be a bool") + return cls( + name=name, + shape=tuple(shape), + stride=tuple(stride), + dtype=dtype, + device_type=device_type, + requires_grad=requires_grad, + ) + + def as_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "shape": list(self.shape), + "stride": list(self.stride), + "dtype": self.dtype, + "device_type": self.device_type, + "requires_grad": self.requires_grad, + } + + +@dataclass(frozen=True) +class ShapeObservation: + """One repeated input signature observed during a model run.""" + + name: str + count: int + total_device_time_us: float + inputs: tuple[TensorSignature, ...] + tags: tuple[str, ...] = () + + @classmethod + def from_dict( + cls, + raw_value: Any, + *, + source: object, + location: str, + ) -> "ShapeObservation": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _OBSERVATION_FIELDS, source, location) + name = _text(raw.get("name"), source, f"{location}.name") + count = _positive_int(raw.get("count"), source, f"{location}.count") + total = _finite_non_negative( + raw.get("total_device_time_us", 0), + source, + f"{location}.total_device_time_us", + ) + inputs_raw = raw.get("inputs") + if ( + not isinstance(inputs_raw, Sequence) + or isinstance(inputs_raw, (str, bytes)) + or not inputs_raw + ): + raise _fail(source, f"{location}.inputs", "must be a non-empty list") + inputs = tuple( + TensorSignature.from_dict( + item, + source=source, + location=f"{location}.inputs[{index}]", + ) + for index, item in enumerate(inputs_raw) + ) + names = [item.name for item in inputs] + if len(names) != len(set(names)): + raise _fail(source, f"{location}.inputs", "contains duplicate names") + + tags_raw = raw.get("tags", []) + if not isinstance(tags_raw, Sequence) or isinstance( + tags_raw, (str, bytes) + ): + raise _fail(source, f"{location}.tags", "must be a list of strings") + tags = tuple( + _text(tag, source, f"{location}.tags[{index}]") + for index, tag in enumerate(tags_raw) + ) + if len(tags) != len(set(tags)): + raise _fail(source, f"{location}.tags", "contains duplicates") + return cls( + name=name, + count=count, + total_device_time_us=total, + inputs=inputs, + tags=tags, + ) + + def as_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "count": self.count, + "total_device_time_us": self.total_device_time_us, + "inputs": [item.as_dict() for item in self.inputs], + "tags": list(self.tags), + } + + +@dataclass(frozen=True) +class CampaignTarget: + """One optimization opportunity ranked by model-level device time.""" + + name: str + operation: str + kind: str + total_device_time_us: float + self_device_time_us: float + calls: int + requires_backward: bool + observations: tuple[ShapeObservation, ...] + spec_locator: str | None = None + attributes: Mapping[str, Any] | None = None + + @classmethod + def from_dict( + cls, + raw_value: Any, + *, + source: object, + location: str, + ) -> "CampaignTarget": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _TARGET_FIELDS, source, location) + name = _text(raw.get("name"), source, f"{location}.name") + operation = _text( + raw.get("operation"), source, f"{location}.operation" + ) + if not _OPERATION_PATTERN.fullmatch(operation): + raise _fail( + source, + f"{location}.operation", + "must be a safe identifier containing letters, digits, and underscores", + ) + kind = _text(raw.get("kind"), source, f"{location}.kind") + if kind not in _TARGET_KINDS: + raise _fail( + source, + f"{location}.kind", + f"must be one of {sorted(_TARGET_KINDS)}", + ) + total = _finite_non_negative( + raw.get("total_device_time_us"), + source, + f"{location}.total_device_time_us", + ) + self_time = _finite_non_negative( + raw.get("self_device_time_us"), + source, + f"{location}.self_device_time_us", + ) + if self_time > total: + raise _fail( + source, + f"{location}.self_device_time_us", + "cannot exceed total_device_time_us", + ) + calls = _positive_int(raw.get("calls"), source, f"{location}.calls") + requires_backward = raw.get("requires_backward", False) + if not isinstance(requires_backward, bool): + raise _fail( + source, f"{location}.requires_backward", "must be a bool" + ) + spec_locator_raw = raw.get("spec_locator") + spec_locator = ( + None + if spec_locator_raw is None + else _text( + spec_locator_raw, source, f"{location}.spec_locator" + ) + ) + observations_raw = raw.get("observations") + if ( + not isinstance(observations_raw, Sequence) + or isinstance(observations_raw, (str, bytes)) + or not observations_raw + ): + raise _fail( + source, f"{location}.observations", "must be a non-empty list" + ) + observations = tuple( + ShapeObservation.from_dict( + item, + source=source, + location=f"{location}.observations[{index}]", + ) + for index, item in enumerate(observations_raw) + ) + observation_calls = sum(item.count for item in observations) + if observation_calls != calls: + raise _fail( + source, + f"{location}.observations", + f"counts sum to {observation_calls}, expected calls={calls}", + ) + attributes_raw = raw.get("attributes", {}) + attributes = _metadata_value( + _mapping(attributes_raw, source, f"{location}.attributes"), + source, + f"{location}.attributes", + ) + return cls( + name=name, + operation=operation, + kind=kind, + total_device_time_us=total, + self_device_time_us=self_time, + calls=calls, + requires_backward=requires_backward, + observations=observations, + spec_locator=spec_locator, + attributes=attributes, + ) + + def impact_pct(self, total_profiled_device_time_us: float) -> float: + if total_profiled_device_time_us <= 0: + return 0.0 + return 100.0 * self.total_device_time_us / total_profiled_device_time_us + + def as_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "operation": self.operation, + "kind": self.kind, + "spec_locator": self.spec_locator, + "total_device_time_us": self.total_device_time_us, + "self_device_time_us": self.self_device_time_us, + "calls": self.calls, + "requires_backward": self.requires_backward, + "observations": [item.as_dict() for item in self.observations], + "attributes": dict(self.attributes or {}), + } + + +@dataclass(frozen=True) +class OptimizationCampaign: + """A validated, versioned request to optimize one model workload.""" + + producer: Mapping[str, Any] + workload: Mapping[str, Any] + environment: Mapping[str, Any] + total_profiled_device_time_us: float + targets: tuple[CampaignTarget, ...] + source: str + schema_version: int = CAMPAIGN_SCHEMA_VERSION + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object = "" + ) -> "OptimizationCampaign": + raw = _mapping(raw_value, source, "top level", non_empty=True) + _unknown_fields(raw, _TOP_LEVEL_FIELDS, source, "top level") + version = raw.get("schema_version") + if isinstance(version, bool) or not isinstance(version, int): + raise _fail(source, "schema_version", "must be an integer") + if version != CAMPAIGN_SCHEMA_VERSION: + raise _fail( + source, + "schema_version", + f"unsupported version {version}; expected {CAMPAIGN_SCHEMA_VERSION}", + ) + producer = dict( + _mapping(raw.get("producer"), source, "producer", non_empty=True) + ) + workload = dict( + _mapping(raw.get("workload"), source, "workload", non_empty=True) + ) + environment = dict( + _mapping( + raw.get("environment"), + source, + "environment", + non_empty=True, + ) + ) + for location, mapping, required in ( + ("producer", producer, ("name", "version")), + ( + "workload", + workload, + ("workload_id", "model_id", "task", "variant_id"), + ), + ( + "environment", + environment, + ("hardware_profile_id", "software_profile_id"), + ), + ): + for field in required: + _text(mapping.get(field), source, f"{location}.{field}") + + total = _finite_non_negative( + raw.get("total_profiled_device_time_us"), + source, + "total_profiled_device_time_us", + ) + if total <= 0: + raise _fail( + source, + "total_profiled_device_time_us", + "must be greater than zero", + ) + targets_raw = raw.get("targets") + if ( + not isinstance(targets_raw, Sequence) + or isinstance(targets_raw, (str, bytes)) + or not targets_raw + ): + raise _fail(source, "targets", "must be a non-empty list") + targets = tuple( + CampaignTarget.from_dict( + item, source=source, location=f"targets[{index}]" + ) + for index, item in enumerate(targets_raw) + ) + names = [target.name for target in targets] + if len(names) != len(set(names)): + raise _fail(source, "targets", "contains duplicate target names") + return cls( + producer=producer, + workload=workload, + environment=environment, + total_profiled_device_time_us=total, + targets=targets, + source=str(source), + schema_version=version, + ) + + def as_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "producer": dict(self.producer), + "workload": dict(self.workload), + "environment": dict(self.environment), + "total_profiled_device_time_us": self.total_profiled_device_time_us, + "targets": [target.as_dict() for target in self.targets], + } + + +def load_campaign(path: str | Path) -> OptimizationCampaign: + """Load and validate a campaign without importing torch or touching a GPU.""" + source = str(path) + file_path = Path(path) + if not file_path.is_file(): + raise _fail(source, "file", "not found") + try: + raw = json.loads(file_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise _fail(source, "JSON", f"invalid JSON: {exc}") from exc + return OptimizationCampaign.from_dict(raw, source=source) + + +def rank_targets(campaign: OptimizationCampaign) -> tuple[CampaignTarget, ...]: + """Rank by end-to-end device-time impact with stable name tie-breaking.""" + return tuple( + sorted( + campaign.targets, + key=lambda target: ( + -target.total_device_time_us, + -target.self_device_time_us, + target.name, + ), + ) + ) + + +def write_optimization_plan( + campaign: OptimizationCampaign, + path: str | Path, + *, + workspace_dir: str | Path = "workspace", +) -> dict[str, Any]: + """Write the legacy orchestrator plan derived from a campaign.""" + kernels = [] + workspace = Path(workspace_dir) + for rank, target in enumerate(rank_targets(campaign), start=1): + kernels.append( + { + "rank": rank, + "file": str( + workspace / f"kernel_{target.operation}_{rank}.py" + ), + "op_type": target.operation, + "target_name": target.name, + "spec_locator": target.spec_locator, + "pct_total": target.impact_pct( + campaign.total_profiled_device_time_us + ), + "calls": target.calls, + "requires_backward": target.requires_backward, + } + ) + plan = { + "schema_version": 1, + "campaign_source": campaign.source, + "workload": dict(campaign.workload), + "environment": dict(campaign.environment), + "kernels_to_optimize": kernels, + } + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(output.suffix + ".tmp") + temporary.write_text(json.dumps(plan, indent=2) + "\n", encoding="utf-8") + temporary.replace(output) + return plan + + +def prepare_campaign( + campaign: OptimizationCampaign, + output_dir: str | Path = "workspace", + *, + trust_specs: bool = False, + spec_root: str | Path | None = None, +) -> dict[str, Any]: + """Materialize trusted starter kernels and an orchestration receipt.""" + if not trust_specs: + raise CampaignError( + "campaign preparation loads Python spec locators; " + "pass trust_specs=True only for campaigns and specs you trust" + ) + + from autokernel.specs import SpecLoadError, load_spec + + workspace = Path(output_dir) + workspace.mkdir(parents=True, exist_ok=True) + plan = write_optimization_plan( + campaign, + workspace / "optimization_plan.json", + workspace_dir=workspace, + ) + prepared = [] + ranked = rank_targets(campaign) + for target, entry in zip(ranked, plan["kernels_to_optimize"]): + if target.spec_locator is None: + raise _fail( + campaign.source, + f"target {target.name!r}.spec_locator", + "is required for campaign preparation", + ) + locator = target.spec_locator + if spec_root is not None: + module, separator, attribute = locator.rpartition(":") + module_path = Path(module) + rooted = Path(spec_root).resolve() / module_path + if separator and not module_path.is_absolute() and rooted.is_file(): + locator = f"{rooted}:{attribute}" + try: + spec = load_spec(locator) + except SpecLoadError as exc: + raise _fail( + campaign.source, + f"target {target.name!r}.spec_locator", + str(exc), + ) from exc + if spec.name != target.operation: + raise _fail( + campaign.source, + f"target {target.name!r}.operation", + f"{target.operation!r} does not match spec name {spec.name!r}", + ) + starter = spec.starter_kernel("triton") + if starter is None: + raise _fail( + campaign.source, + f"target {target.name!r}", + "spec does not declare a Triton starter kernel", + ) + destination = Path(entry["file"]) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(destination.suffix + ".tmp") + temporary.write_text( + starter.read_text(encoding="utf-8"), encoding="utf-8" + ) + temporary.replace(destination) + prepared.append( + { + "rank": entry["rank"], + "target_name": target.name, + "operation": target.operation, + "candidate": str(destination), + "spec_locator": target.spec_locator, + "status": "prepared", + } + ) + + receipt = { + "schema_version": 1, + "status": "prepared", + "campaign_source": campaign.source, + "optimization_plan": str(workspace / "optimization_plan.json"), + "targets": prepared, + "next_command": "Follow program.md with orchestrate.py next and bench.py", + } + receipt_path = workspace / "campaign_receipt.json" + temporary = receipt_path.with_suffix(".json.tmp") + temporary.write_text( + json.dumps(receipt, indent=2) + "\n", encoding="utf-8" + ) + temporary.replace(receipt_path) + return receipt diff --git a/autokernel/discovery/__init__.py b/autokernel/discovery/__init__.py new file mode 100644 index 00000000..b3a294fb --- /dev/null +++ b/autokernel/discovery/__init__.py @@ -0,0 +1,79 @@ +"""Universal profiling and graph-region discovery contracts.""" + +from .correlation import ( + correlate_discovery_report, + correlate_profiler_to_regions, +) +from .fingerprint import fingerprint_payload, graph_fingerprint +from .fx_capture import ( + CaptureResult, + RegionCaptureSession, + capture_callable_region, + capture_model_regions, + capture_module_region, +) +from .profiler_export import load_profiler_export, profiler_export_to_report +from .profiler_parse import parse_key_averages_rows +from .ranking import ( + DEFAULT_IMPACT_FLOOR, + DEFAULT_PROMOTION_TARGET, + RankedCandidate, + classify_pattern_family, + optimistic_e2e_improvement, + rank_operators, + rank_regions, +) +from .safety import ( + ALLOWED_ATEN_OPS, + is_region_safe, + normalize_op_name, + reject_region, +) +from .types import ( + DISCOVERY_SCHEMA_VERSION, + DiscoveryError, + DiscoveryReport, + GraphBreakRecord, + GraphRegion, + OperatorHotspot, + TensorMeta, + UnsupportedOpRecord, + load_discovery_report, + write_discovery_report, +) + +__all__ = [ + "ALLOWED_ATEN_OPS", + "DEFAULT_IMPACT_FLOOR", + "DEFAULT_PROMOTION_TARGET", + "DISCOVERY_SCHEMA_VERSION", + "CaptureResult", + "DiscoveryError", + "DiscoveryReport", + "GraphBreakRecord", + "GraphRegion", + "OperatorHotspot", + "RankedCandidate", + "RegionCaptureSession", + "TensorMeta", + "UnsupportedOpRecord", + "capture_callable_region", + "capture_model_regions", + "capture_module_region", + "classify_pattern_family", + "correlate_discovery_report", + "correlate_profiler_to_regions", + "fingerprint_payload", + "graph_fingerprint", + "is_region_safe", + "load_discovery_report", + "load_profiler_export", + "normalize_op_name", + "optimistic_e2e_improvement", + "parse_key_averages_rows", + "profiler_export_to_report", + "rank_operators", + "rank_regions", + "reject_region", + "write_discovery_report", +] diff --git a/autokernel/discovery/correlation.py b/autokernel/discovery/correlation.py new file mode 100644 index 00000000..c4984efd --- /dev/null +++ b/autokernel/discovery/correlation.py @@ -0,0 +1,399 @@ +"""Offline correlation between profiler rows and captured FX regions. + +This module implements the logic to match profiler operator rows with FX graph +regions, aggregate timing data, and populate DiscoveryReport with correlated +metrics without double-counting nested scopes. +""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, Mapping, Sequence + +from .profiler_parse import parse_key_averages_rows +from .ranking import optimistic_e2e_improvement +from .safety import normalize_op_name, reject_region +from .types import ( + DiscoveryReport, + GraphRegion, + OperatorHotspot, +) + + +@dataclass(frozen=True) +class ScopeMatch: + """Result of matching a profiler row to a region.""" + + profiler_row: OperatorHotspot + region: GraphRegion + confidence: float + match_reason: str + + +@dataclass +class RegionAccumulator: + """Aggregate data for regions with the same fingerprint.""" + + base_region: GraphRegion + total_cuda_time_us: float = 0.0 + total_self_cuda_time_us: float = 0.0 + total_calls: int = 0 + shape_frequencies: dict[str, int] = field(default_factory=dict) + matched_profiler_rows: list[OperatorHotspot] = field(default_factory=list) + unmatched_scopes: list[str] = field(default_factory=list) + capture_failures: list[str] = field(default_factory=list) + + +def _match_scope_to_region( + profiler_row: OperatorHotspot, + region: GraphRegion, +) -> ScopeMatch | None: + """Attempt to match a profiler row to a region based on scope heuristics. + + Matching strategies (in order of precedence): + 1. Exact parent_module match + 2. Op key appears in region operations + 3. Name similarity between profiler row and region + + Returns None if no reasonable match is found. + """ + # Shape-specific record_function ranges exported by FastVideo carry the + # captured region name verbatim. Prefer this identity over broader module + # or operation heuristics. + if profiler_row.name == region.name: + return ScopeMatch( + profiler_row=profiler_row, + region=region, + confidence=1.0, + match_reason="exact_region_name", + ) + + # Strategy 1: Exact parent_module match + if profiler_row.parent_module and region.parent_module: + if profiler_row.parent_module == region.parent_module: + return ScopeMatch( + profiler_row=profiler_row, + region=region, + confidence=0.9, + match_reason="exact_parent_module", + ) + # Check for hierarchical relationship (e.g., "blocks.0.attn" matches "blocks.0") + if (profiler_row.parent_module.startswith(region.parent_module + ".") or + region.parent_module.startswith(profiler_row.parent_module + ".")): + return ScopeMatch( + profiler_row=profiler_row, + region=region, + confidence=0.7, + match_reason="hierarchical_parent_module", + ) + + # Strategy 2: Op key appears in region operations. Region operations are + # normalized during FX capture (overload suffixes stripped), while + # profiler rows keep overload-qualified names, so normalize both sides. + normalized_op_key = normalize_op_name(profiler_row.op_key) + if normalized_op_key in { + normalize_op_name(op) for op in region.operations + }: + return ScopeMatch( + profiler_row=profiler_row, + region=region, + confidence=0.6, + match_reason="op_key_in_operations", + ) + + # Strategy 3: Name similarity + profiler_name_lower = profiler_row.name.lower() + region_name_lower = region.name.lower() + if (profiler_name_lower in region_name_lower or + region_name_lower in profiler_name_lower): + return ScopeMatch( + profiler_row=profiler_row, + region=region, + confidence=0.5, + match_reason="name_similarity", + ) + + return None + + +def _calculate_exclusive_time( + profiler_row: OperatorHotspot, + parent_time_us: float, +) -> tuple[float, float]: + """Calculate exclusive CUDA time avoiding double-counting. + + Returns (cuda_time_us, self_cuda_time_us) where: + - cuda_time_us is the total time (inclusive) + - self_cuda_time_us is exclusive time, capped at parent_time_us if nested + """ + # Use the profiler's self_cuda_time_us as the exclusive time + exclusive = profiler_row.self_cuda_time_us + + # If this is nested within a parent region, ensure we don't exceed parent time + if parent_time_us > 0 and exclusive > parent_time_us: + # This can happen if profiler attribution is inconsistent + # Conservative approach: cap at parent time + exclusive = parent_time_us + + return profiler_row.cuda_time_us, exclusive + + +def _aggregate_region_timing( + region: GraphRegion, + profiler_rows: Sequence[OperatorHotspot], +) -> tuple[float, float, int]: + """Aggregate timing data for a region from matched profiler rows. + + Returns (total_cuda_time_us, total_self_cuda_time_us, total_calls). + Uses exclusive time to avoid double-counting nested scopes. + """ + if not profiler_rows: + return region.cuda_time_us, region.self_cuda_time_us, region.calls + + # A shape-specific record_function range names the region itself. PyTorch + # reports its useful device attribution as inclusive CUDA time while its + # self CUDA time is normally zero (the range launches no kernel directly). + # Treat that inclusive duration as the region's attributed duration. This + # is safe per candidate; callers must not sum nested candidate shares. + attributed_self = [ + row.cuda_time_us if row.name == region.name else row.self_cuda_time_us + for row in profiler_rows + ] + + total_cuda = sum(row.cuda_time_us for row in profiler_rows) + total_self_cuda = sum(attributed_self) + total_calls = sum(row.calls for row in profiler_rows) + + # Fall back to region's own timing if no profiler data + if total_self_cuda == 0: + total_self_cuda = region.self_cuda_time_us + if total_cuda == 0: + total_cuda = region.cuda_time_us + if total_calls == 0: + total_calls = region.calls + + return total_cuda, total_self_cuda, total_calls + + +def _deduplicate_regions_by_fingerprint( + regions: Sequence[GraphRegion], +) -> dict[str, RegionAccumulator]: + """Group equivalent regions by their stable graph fingerprint.""" + accumulators: dict[str, RegionAccumulator] = {} + + for region in regions: + fingerprint = region.fingerprint + if fingerprint not in accumulators: + accumulators[fingerprint] = RegionAccumulator(base_region=region) + + # Merge shape frequencies + if region.shape_frequency: + for shape_key, count in region.shape_frequency.items(): + accumulators[fingerprint].shape_frequencies[shape_key] = ( + accumulators[fingerprint].shape_frequencies.get(shape_key, 0) + count + ) + + return accumulators + + +def correlate_profiler_to_regions( + profiler_rows: Sequence[OperatorHotspot], + fx_regions: Sequence[GraphRegion], + *, + total_cuda_time_us: float, +) -> tuple[tuple[GraphRegion, ...], tuple[OperatorHotspot, ...]]: + """Correlate profiler rows with FX regions and populate timing data. + + This is the main entry point for offline correlation. It: + 1. Matches profiler rows to regions based on scope heuristics + 2. Calculates exclusive CUDA time without double-counting nested scopes + 3. Deduplicates equivalent regions using stable graph fingerprints + 4. Aggregates timing, call counts, and shape frequencies + 5. Computes confidence and rejection reasons + 6. Returns populated GraphRegion tuples plus the unmatched rows + + Args: + profiler_rows: OperatorHotspot rows from the profiler + fx_regions: GraphRegion instances from FX capture + total_cuda_time_us: Total end-to-end CUDA time for percentage calculations + + Returns: + A ``(regions, unmatched_rows)`` pair: populated GraphRegion + instances, and the profiler rows that matched no captured region. + Unmatched rows are reported separately because they carry no input + metadata and therefore cannot form a valid, serializable region. + """ + # Step 1: Group regions by fingerprint for deduplication + fingerprint_groups = _deduplicate_regions_by_fingerprint(fx_regions) + + # Step 2: Match profiler rows to regions + matched_regions: dict[str, list[ScopeMatch]] = defaultdict(list) + unmatched_rows: list[OperatorHotspot] = [] + + for row in profiler_rows: + candidates: list[ScopeMatch] = [] + + # Try to match against each region group (use first region as representative) + for accumulator in fingerprint_groups.values(): + match = _match_scope_to_region(row, accumulator.base_region) + if match: + candidates.append(match) + + best_match: ScopeMatch | None = None + if candidates: + best_confidence = max(match.confidence for match in candidates) + strongest = [ + match for match in candidates + if match.confidence == best_confidence + ] + # Never assign an aggregate operator row arbitrarily when the same + # op occurs in several captured regions. A shape-specific scope + # range or other unique match is required for trustworthy timing. + if len(strongest) == 1: + best_match = strongest[0] + + if best_match is not None: + matched_regions[best_match.region.fingerprint].append(best_match) + else: + unmatched_rows.append(row) + + # Step 3: Aggregate timing data and build final regions + final_regions: list[GraphRegion] = [] + + for fingerprint, accumulator in fingerprint_groups.items(): + matches = matched_regions.get(fingerprint, []) + matched_rows = [m.profiler_row for m in matches] + + # Aggregate timing from matched profiler rows + total_cuda, total_self_cuda, total_calls = _aggregate_region_timing( + accumulator.base_region, + matched_rows, + ) + + # Merge shape frequencies + merged_shape_freq = dict(accumulator.shape_frequencies) + for match in matches: + if match.profiler_row.input_shapes: + # Create a shape key from input shapes + shape_key = "|".join( + "x".join(str(d) for d in shape) + for shape in match.profiler_row.input_shapes + ) + merged_shape_freq[shape_key] = ( + merged_shape_freq.get(shape_key, 0) + match.profiler_row.calls + ) + + # Calculate confidence + safety_reasons = tuple(reject_region(accumulator.base_region.operations)) + confidence = 0.4 # Base confidence + if not safety_reasons: + confidence += 0.3 # Safety bonus + if matches: + confidence += 0.2 # Profiler match bonus + if total_calls >= 8: + confidence += 0.1 # Call count bonus + if merged_shape_freq: + confidence += 0.1 # Shape frequency bonus + confidence = min(1.0, confidence) + + # Calculate percentage of end-to-end time + e2e_share = 0.0 + if total_cuda_time_us > 0: + e2e_share = 100.0 * total_self_cuda / total_cuda_time_us + + # Estimate maximum end-to-end improvement + estimated_improvement = optimistic_e2e_improvement(e2e_share / 100.0) + + # Build rejection reasons + rejection_reasons = list(safety_reasons) + if not matches: + rejection_reasons.append("no_profiler_match") + if accumulator.base_region.rejection_reasons: + rejection_reasons.extend(accumulator.base_region.rejection_reasons) + + # Create the populated region + populated_region = GraphRegion.build( + name=accumulator.base_region.name, + operations=accumulator.base_region.operations, + inputs=accumulator.base_region.inputs, + outputs=accumulator.base_region.outputs, + dependencies=accumulator.base_region.dependencies, + parent_module=accumulator.base_region.parent_module, + safe_constants=accumulator.base_region.safe_constants, + pattern_family=accumulator.base_region.pattern_family, + rejection_reasons=tuple(dict.fromkeys(rejection_reasons)), + calls=total_calls, + shape_frequency=merged_shape_freq or None, + cuda_time_us=total_cuda, + self_cuda_time_us=total_self_cuda, + attributes={ + **(accumulator.base_region.attributes or {}), + "e2e_share_pct": round(e2e_share, 4), + "estimated_max_e2e_improvement": round(estimated_improvement, 4), + "confidence": round(confidence, 4), + "matched_profiler_rows": len(matched_rows), + }, + ) + + final_regions.append(populated_region) + + return tuple(final_regions), tuple(unmatched_rows) + + +def correlate_discovery_report( + profiler_export_rows: Sequence[Mapping[str, Any]], + fx_discovery_report: DiscoveryReport, +) -> DiscoveryReport: + """Correlate profiler data with an existing FX discovery report. + + This function takes profiler export rows and a discovery report containing + FX-captured regions, correlates them, and returns a new discovery report with + populated timing data. + + Args: + profiler_export_rows: Raw profiler export rows (list of dicts) + fx_discovery_report: Discovery report from FX capture (CPU-only) + + Returns: + New DiscoveryReport with correlated timing data + """ + # Parse profiler rows + profiler_operators = parse_key_averages_rows(profiler_export_rows) + + # Correlate profiler rows with FX regions + searchable_regions, unmatched_rows = correlate_profiler_to_regions( + profiler_operators, + fx_discovery_report.regions, + total_cuda_time_us=fx_discovery_report.total_cuda_time_us, + ) + unsupported = [ + item.as_dict() for item in fx_discovery_report.unsupported + ] + if unmatched_rows: + unsupported.append( + { + "op_name": "profiler::unmatched", + "reason": ( + f"{len(unmatched_rows)} profiler row(s) did not match a " + "captured FX region" + ), + "count": len(unmatched_rows), + "scope": "profiler_correlation", + } + ) + + # Create new discovery report with populated regions + return DiscoveryReport.from_dict( + { + "schema_version": fx_discovery_report.schema_version, + "producer": dict(fx_discovery_report.producer), + "workload": dict(fx_discovery_report.workload), + "environment": dict(fx_discovery_report.environment), + "total_cuda_time_us": fx_discovery_report.total_cuda_time_us, + "operators": [op.as_dict() for op in profiler_operators], + "regions": [region.as_dict() for region in searchable_regions], + "graph_breaks": [item.as_dict() for item in fx_discovery_report.graph_breaks], + "unsupported": unsupported, + } + ) diff --git a/autokernel/discovery/fingerprint.py b/autokernel/discovery/fingerprint.py new file mode 100644 index 00000000..0c7be43a --- /dev/null +++ b/autokernel/discovery/fingerprint.py @@ -0,0 +1,60 @@ +"""Stable fingerprints for captured graph regions and operator sequences. + +Fingerprints are derived only from operation names, tensor signatures, and +safe constants. They must be identical across repeated runs of the same +region and must never incorporate tensor values, prompts, or weights. +""" + +from __future__ import annotations + +import hashlib +import json +import math +from typing import Any, Mapping, Sequence + + +def _canonical(value: Any) -> Any: + if value is None or isinstance(value, (bool, int, float, str)): + if isinstance(value, float): + if not math.isfinite(value): + raise ValueError("fingerprint values must be finite") + return float(value) + return value + if isinstance(value, Mapping): + return {str(k): _canonical(v) for k, v in sorted(value.items(), key=lambda kv: str(kv[0]))} + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [_canonical(item) for item in value] + raise ValueError( + f"unsupported fingerprint value type: {type(value).__name__}" + ) + + +def fingerprint_payload(payload: Mapping[str, Any], *, length: int = 32) -> str: + """Hash a JSON-canonical payload into a short hex fingerprint.""" + encoded = json.dumps( + _canonical(dict(payload)), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ) + digest = hashlib.sha256(encoded.encode("utf-8")).hexdigest() + return digest[:length] + + +def graph_fingerprint( + *, + operations: Sequence[str], + input_signatures: Sequence[Mapping[str, Any]], + output_signatures: Sequence[Mapping[str, Any]] | None = None, + safe_constants: Mapping[str, Any] | None = None, + parent_module: str | None = None, +) -> str: + """Fingerprint a pure tensor region from metadata only.""" + payload = { + "operations": list(operations), + "inputs": list(input_signatures), + "outputs": list(output_signatures or ()), + "constants": dict(safe_constants or {}), + "parent_module": parent_module or "", + } + return fingerprint_payload(payload) diff --git a/autokernel/discovery/fx_capture.py b/autokernel/discovery/fx_capture.py new file mode 100644 index 00000000..8f0b60e8 --- /dev/null +++ b/autokernel/discovery/fx_capture.py @@ -0,0 +1,703 @@ +"""Model-independent Dynamo/FX region capture (metadata only). + +Captures repeated module calls as executable FX tensor subgraphs for discovery +and ranking. Records only: + +- ordered operations and dependencies +- input/output tensor signatures (shape, stride, dtype, device, requires_grad) +- safe scalar constants needed for semantics +- parent module scope, call counts, shape frequency +- graph breaks and unsupported operations + +Never serializes tensor values, weights, prompts, model outputs, credentials, +or arbitrary Python source. Fail closed on mutation, collectives, +data-dependent control flow, unknown aliasing, and unsupported custom ops. + +CUDA timings remain optional and are filled by a separate profiler path. +""" + +from __future__ import annotations + +import re +from collections import defaultdict +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any, Callable, Iterator, Mapping, Sequence + +from .ranking import classify_pattern_family +from .safety import is_region_safe, normalize_op_name, reject_region +from .types import ( + DiscoveryReport, + GraphBreakRecord, + GraphRegion, + TensorMeta, + UnsupportedOpRecord, +) + +_FORBIDDEN_SERIALIZED_KEYS = frozenset( + { + "activations", + "credential", + "credentials", + "data", + "password", + "prompt", + "secret", + "secrets", + "source", + "source_code", + "tensor_values", + "token", + "values", + "weights", + } +) + +_SAFE_SCALAR_TYPES = (bool, int, float, str) + + +@dataclass(frozen=True) +class CaptureResult: + """Result of attempting to capture one module or callable.""" + + region: GraphRegion | None + graph_breaks: tuple[GraphBreakRecord, ...] + unsupported: tuple[UnsupportedOpRecord, ...] + operations: tuple[str, ...] + + +@dataclass +class _RegionAccumulator: + """Mutable aggregation for repeated captures of the same fingerprint.""" + + region: GraphRegion + calls: int = 0 + shape_keys: dict[str, int] = field(default_factory=dict) + breaks: list[GraphBreakRecord] = field(default_factory=list) + unsupported: list[UnsupportedOpRecord] = field(default_factory=list) + + +def _tensor_meta_from_example(name: str, tensor: Any) -> TensorMeta: + """Layout metadata only — never reads or stores tensor values.""" + shape = tuple(int(x) for x in tensor.shape) + if hasattr(tensor, "stride"): + stride = tuple(int(x) for x in tensor.stride()) + else: + stride_list: list[int] = [] + running = 1 + for dim in reversed(shape): + stride_list.append(running) + running *= max(dim, 1) + stride = tuple(reversed(stride_list)) + dtype = str(tensor.dtype).replace("torch.", "") + device_type = "cpu" + if hasattr(tensor, "device") and hasattr(tensor.device, "type"): + device_type = tensor.device.type + requires_grad = bool(getattr(tensor, "requires_grad", False)) + return TensorMeta( + name=name, + shape=shape, + stride=stride, + dtype=dtype, + device_type=device_type, + requires_grad=requires_grad, + ) + + +def _shape_frequency_key(inputs: Sequence[TensorMeta]) -> str: + parts = [] + for item in inputs: + shape = "x".join(str(d) for d in item.shape) + parts.append(f"{item.name}:{shape}:{item.dtype}") + return "|".join(parts) + + +def _function_to_op_key(target: Any) -> str: + text = str(target) + if "aten::" in text: + return normalize_op_name( + "aten::" + text.split("aten::", 1)[1].split(".")[0].split("(")[0] + ) + if "aten." in text: + return normalize_op_name("aten::" + text.split("aten.", 1)[1].split(".")[0]) + name = getattr(target, "__name__", None) or text + module = getattr(target, "__module__", "") or "" + if "torch" in module or name in { + "add", + "mul", + "sub", + "div", + "silu", + "gelu", + "relu", + "sigmoid", + "layer_norm", + "softmax", + "tanh", + "rsqrt", + "sqrt", + "pow", + "mean", + "var", + "neg", + "exp", + "clone", + "contiguous", + "view", + "reshape", + "permute", + "transpose", + "unsqueeze", + "squeeze", + "cat", + "stack", + "expand", + "type_as", + }: + return normalize_op_name(f"aten::{name}") + return normalize_op_name(str(name)) + + +def _is_safe_constant_value(value: Any) -> bool: + if isinstance(value, _SAFE_SCALAR_TYPES): + if isinstance(value, float): + import math + + return math.isfinite(value) + if isinstance(value, str): + # Short enum-like tags only — never free-form source/prompts. + return 0 < len(value) <= 64 and "\n" not in value + return True + if isinstance(value, (list, tuple)): + if len(value) > 32: + return False + return all(_is_safe_constant_value(item) for item in value) + return False + + +def _sanitize_safe_constants(raw: Mapping[str, Any]) -> dict[str, Any]: + """Keep only finite scalars / short tags; drop anything else.""" + result: dict[str, Any] = {} + for key, value in raw.items(): + if not isinstance(key, str) or not key: + continue + if key.lower() in _FORBIDDEN_SERIALIZED_KEYS: + continue + if _is_safe_constant_value(value): + result[key] = value + return result + + +def _extract_graph_structure(graph: Any) -> tuple[list[str], list[str], dict[str, Any], list[str]]: + """Return (operations, dependencies, safe_constants, structural_rejections).""" + operations: list[str] = [] + dependencies: list[str] = [] + safe_constants: dict[str, Any] = {} + structural: list[str] = [] + node_index: dict[str, int] = {} + + for node in graph.nodes: + if node.op in {"placeholder", "output"}: + continue + if node.op == "get_attr": + # Parameters / buffers are weights — never export values. + # Only allow pure Python attribute scalars if present on the module + # via string target name recording (no value read here). + target = str(node.target) + if any( + part in target.lower() + for part in ("weight", "bias", "embed", "param", "buffer") + ): + structural.append(f"get_attr:{target}: weights/parameters forbidden") + else: + # Record attribute *name* only as a dependency tag, not its value. + safe_constants.setdefault(f"attr:{target}", True) + continue + if node.op == "call_function": + op_key = _function_to_op_key(node.target) + idx = len(operations) + node_index[node.name] = idx + operations.append(op_key) + for arg in node.args: + arg_name = getattr(arg, "name", None) + if arg_name in node_index: + dependencies.append(f"{node_index[arg_name]}->{idx}") + for key, value in (node.kwargs or {}).items(): + if _is_safe_constant_value(value): + safe_constants[f"{node.name}.{key}"] = value + elif value is not None and not hasattr(value, "op"): + # Non-safe non-node kwarg — fail closed (unknown constant). + structural.append( + f"{op_key}: unsafe or non-scalar constant {key!r}" + ) + elif node.op == "call_method": + op_key = normalize_op_name(f"aten::{node.target}") + idx = len(operations) + node_index[node.name] = idx + operations.append(op_key) + for arg in node.args: + arg_name = getattr(arg, "name", None) + if arg_name in node_index: + dependencies.append(f"{node_index[arg_name]}->{idx}") + elif node.op == "call_module": + # Nested modules collapse to a single opaque node — fail closed for + # leaf fusion search unless expanded. + op_key = normalize_op_name(f"module::{node.target}") + idx = len(operations) + node_index[node.name] = idx + operations.append(op_key) + structural.append(f"{op_key}: nested module not expanded") + else: + structural.append(f"unknown_fx_op:{node.op}") + + # Unknown aliasing: multiple outputs writing overlapping views without + # explicit clone is hard to prove; flag in-place method names already + # covered by reject_region. Detect star-deps fan-in of getitem/scatter-like. + for op in operations: + if "scatter" in op or "index_put" in op or "copy_" in op: + structural.append(f"{op}: potential aliasing mutation") + + return operations, dependencies, safe_constants, structural + + +def _trace_module(module: Any, *, tracer: str, example_inputs: Sequence[Any]) -> Any: + """Return an FX GraphModule or raise.""" + import torch + import torch.fx as fx + + if tracer == "symbolic": + return fx.symbolic_trace(module) + if tracer == "dynamo": + try: + exported = torch._dynamo.export(module)(*example_inputs) + # dynamo.export returns (gm, guards) on some versions or an object + if isinstance(exported, tuple): + return exported[0] + graph_module = getattr(exported, "graph_module", None) + if graph_module is not None: + return graph_module + return exported + except Exception: + # Fall back to symbolic when Dynamo cannot export the module. + return fx.symbolic_trace(module) + raise ValueError(f"unsupported tracer {tracer!r}; use 'symbolic' or 'dynamo'") + + +def _assert_region_metadata_only(region: GraphRegion) -> None: + """Hard privacy check on the produced region dict.""" + payload = region.as_dict() + + def walk(obj: Any, path: str) -> None: + if isinstance(obj, Mapping): + for key, value in obj.items(): + lower = str(key).lower() + if lower in _FORBIDDEN_SERIALIZED_KEYS: + raise RuntimeError( + f"forbidden key {key!r} at {path} in captured region" + ) + walk(value, f"{path}.{key}") + elif isinstance(obj, Sequence) and not isinstance(obj, (str, bytes)): + for index, item in enumerate(obj): + walk(item, f"{path}[{index}]") + + walk(payload, "region") + + +def capture_module_region( + module: Any, + example_inputs: Sequence[Any], + *, + name: str, + parent_module: str | None = None, + tracer: str = "symbolic", + calls: int = 1, + shape_frequency: Mapping[str, int] | None = None, +) -> CaptureResult: + """Trace a module and build a metadata-only GraphRegion. + + ``example_inputs`` are used only for tracing shapes/dtypes and one forward + for output signatures — values are never written into the region. + """ + import torch + + breaks: list[GraphBreakRecord] = [] + unsupported: list[UnsupportedOpRecord] = [] + operations: list[str] = [] + dependencies: list[str] = [] + safe_constants: dict[str, Any] = {} + + try: + traced = _trace_module( + module, tracer=tracer, example_inputs=example_inputs + ) + graph = getattr(traced, "graph", None) + if graph is None: + raise RuntimeError("trace result has no FX graph") + operations, dependencies, safe_constants, structural = _extract_graph_structure( + graph + ) + for reason in structural: + breaks.append( + GraphBreakRecord(scope=name, reason=reason, count=1) + ) + except Exception as exc: # noqa: BLE001 - capture failures are data + breaks.append( + GraphBreakRecord( + scope=name, + reason=f"fx_trace_failed: {type(exc).__name__}: {exc}", + count=1, + ) + ) + return CaptureResult(None, tuple(breaks), tuple(unsupported), ()) + + if not operations: + breaks.append( + GraphBreakRecord(scope=name, reason="empty_graph", count=1) + ) + return CaptureResult(None, tuple(breaks), tuple(unsupported), ()) + + safe_constants = _sanitize_safe_constants(safe_constants) + rejection = list(reject_region(operations)) + # Nested module / structural reasons already in breaks; fold into rejection. + for item in breaks: + if item.reason not in rejection and not item.reason.startswith( + "fx_trace_failed" + ): + rejection.append(item.reason) + + for reason in rejection: + if ( + "unsupported custom" in reason + or "not in pure-tensor" in reason + or "nested module" in reason + ): + op_name = reason.rsplit(": ", 1)[0] + unsupported.append( + UnsupportedOpRecord( + op_name=op_name, reason=reason, count=1, scope=name + ) + ) + + inputs = tuple( + _tensor_meta_from_example(f"input_{i}", tensor) + for i, tensor in enumerate(example_inputs) + ) + outputs: tuple[TensorMeta, ...] = () + try: + with torch.no_grad(): + out = module(*example_inputs) + if isinstance(out, torch.Tensor): + outputs = (_tensor_meta_from_example("output_0", out),) + elif isinstance(out, (tuple, list)): + outputs = tuple( + _tensor_meta_from_example(f"output_{i}", item) + for i, item in enumerate(out) + if isinstance(item, torch.Tensor) + ) + except Exception as exc: # noqa: BLE001 + breaks.append( + GraphBreakRecord( + scope=name, + reason=f"output_meta_failed: {type(exc).__name__}: {exc}", + count=1, + ) + ) + + freq = dict(shape_frequency) if shape_frequency else {} + if not freq: + freq[_shape_frequency_key(inputs)] = max(1, calls) + + family = classify_pattern_family(operations) + region = GraphRegion.build( + name=name, + operations=operations, + inputs=inputs, + outputs=outputs, + dependencies=tuple(dependencies), + parent_module=parent_module, + safe_constants=safe_constants or None, + pattern_family=family, + rejection_reasons=tuple(dict.fromkeys(rejection)), + calls=max(1, calls), + shape_frequency=freq, + cuda_time_us=0.0, + self_cuda_time_us=0.0, + ) + _assert_region_metadata_only(region) + return CaptureResult( + region=region, + graph_breaks=tuple(breaks), + unsupported=tuple(unsupported), + operations=tuple(operations), + ) + + +def capture_callable_region( + fn: Callable[..., Any], + example_inputs: Sequence[Any], + *, + name: str, + tracer: str = "symbolic", +) -> CaptureResult: + """Wrap a pure function as an nn.Module and capture it.""" + import torch.nn as nn + + arity = len(example_inputs) + if arity == 1: + + class _Wrapper1(nn.Module): + def forward(self, x): # type: ignore[no-untyped-def] + return fn(x) + + wrapper: nn.Module = _Wrapper1() + elif arity == 2: + + class _Wrapper2(nn.Module): + def forward(self, x, y): # type: ignore[no-untyped-def] + return fn(x, y) + + wrapper = _Wrapper2() + elif arity == 3: + + class _Wrapper3(nn.Module): + def forward(self, x, y, z): # type: ignore[no-untyped-def] + return fn(x, y, z) + + wrapper = _Wrapper3() + else: + return CaptureResult( + None, + ( + GraphBreakRecord( + scope=name, + reason=f"callable arity {arity} not supported for FX wrap", + count=1, + ), + ), + (), + (), + ) + + return capture_module_region( + wrapper, + example_inputs, + name=name, + parent_module=None, + tracer=tracer, + ) + + +class RegionCaptureSession: + """Hook selected modules, capture FX regions on each forward, aggregate. + + Designed for repeated DiT block / leaf-module calls without capturing the + entire generation pipeline as one graph. + """ + + def __init__( + self, + *, + tracer: str = "symbolic", + name_prefix: str = "region", + ) -> None: + self.tracer = tracer + self.name_prefix = name_prefix + self._accumulators: dict[str, _RegionAccumulator] = {} + self._graph_breaks: list[GraphBreakRecord] = [] + self._unsupported: list[UnsupportedOpRecord] = [] + self._hooks: list[Any] = [] + self._call_counters: dict[str, int] = defaultdict(int) + # Prevent re-entry when capture_module_region runs symbolic_trace / + # a meta forward on the same hooked module. + self._capturing: bool = False + + def register_module( + self, + module: Any, + *, + scope: str, + ) -> None: + """Install a forward hook that captures each invocation.""" + + def _hook(_mod: Any, inputs: tuple[Any, ...], _output: Any) -> None: + if self._capturing: + return + self._call_counters[scope] += 1 + # Only tensor args are used for signatures. + tensor_inputs = tuple( + arg for arg in inputs if hasattr(arg, "shape") and hasattr(arg, "dtype") + ) + if not tensor_inputs: + self._graph_breaks.append( + GraphBreakRecord( + scope=scope, + reason="no_tensor_inputs", + count=1, + ) + ) + return + capture_name = re.sub( + r"[^A-Za-z0-9._-]", "_", f"{self.name_prefix}.{scope}" + ) + if not re.match(r"^[A-Za-z0-9]", capture_name): + capture_name = f"r.{capture_name}" + # GraphRegion names must match _NAME_PATTERN (128 chars max), + # so truncate after any prefix is applied. + capture_name = capture_name[:128] + self._capturing = True + try: + result = capture_module_region( + _mod, + tensor_inputs, + name=capture_name, + parent_module=scope, + tracer=self.tracer, + ) + finally: + self._capturing = False + self._graph_breaks.extend(result.graph_breaks) + self._unsupported.extend(result.unsupported) + if result.region is None: + return + key = result.region.fingerprint + shape_key = _shape_frequency_key(result.region.inputs) + if key not in self._accumulators: + self._accumulators[key] = _RegionAccumulator(region=result.region) + acc = self._accumulators[key] + acc.calls += 1 + acc.shape_keys[shape_key] = acc.shape_keys.get(shape_key, 0) + 1 + acc.breaks.extend(result.graph_breaks) + acc.unsupported.extend(result.unsupported) + + handle = module.register_forward_hook(_hook) + self._hooks.append(handle) + + def register_named_children( + self, + root: Any, + *, + predicate: Callable[[str, Any], bool] | None = None, + ) -> int: + """Register hooks on ``named_modules`` matching predicate (default: leaves).""" + count = 0 + for name, child in root.named_modules(): + if name == "": + continue + if predicate is not None: + if not predicate(name, child): + continue + else: + # Default: leaf modules with no children + if any(True for _ in child.children()): + continue + self.register_module(child, scope=name or "root") + count += 1 + return count + + def close(self) -> None: + for handle in self._hooks: + handle.remove() + self._hooks.clear() + + def __enter__(self) -> "RegionCaptureSession": + return self + + def __exit__(self, *exc: Any) -> None: + self.close() + + def regions(self) -> tuple[GraphRegion, ...]: + """Materialize aggregated GraphRegion records.""" + regions: list[GraphRegion] = [] + for acc in self._accumulators.values(): + base = acc.region + regions.append( + GraphRegion.build( + name=base.name, + operations=base.operations, + inputs=base.inputs, + outputs=base.outputs, + dependencies=base.dependencies, + parent_module=base.parent_module, + safe_constants=base.safe_constants, + pattern_family=base.pattern_family, + rejection_reasons=base.rejection_reasons, + calls=max(1, acc.calls), + shape_frequency=dict(acc.shape_keys) or base.shape_frequency, + cuda_time_us=base.cuda_time_us, + self_cuda_time_us=base.self_cuda_time_us, + attributes=base.attributes, + ) + ) + return tuple(regions) + + def graph_breaks(self) -> tuple[GraphBreakRecord, ...]: + # Coalesce identical break reasons. + counts: dict[tuple[str, str, str | None], int] = defaultdict(int) + for item in self._graph_breaks: + counts[(item.scope, item.reason, item.op_name)] += item.count + return tuple( + GraphBreakRecord( + scope=scope, reason=reason, op_name=op_name, count=count + ) + for (scope, reason, op_name), count in sorted(counts.items()) + ) + + def unsupported(self) -> tuple[UnsupportedOpRecord, ...]: + counts: dict[tuple[str, str, str | None], int] = defaultdict(int) + for item in self._unsupported: + counts[(item.op_name, item.reason, item.scope)] += item.count + return tuple( + UnsupportedOpRecord( + op_name=op_name, reason=reason, count=count, scope=scope + ) + for (op_name, reason, scope), count in sorted(counts.items()) + ) + + def to_discovery_report( + self, + *, + workload: Mapping[str, Any], + environment: Mapping[str, Any] | None = None, + producer: Mapping[str, Any] | None = None, + total_cuda_time_us: float = 0.0, + ) -> DiscoveryReport: + """Build a DiscoveryReport from aggregated FX captures (CPU-safe).""" + return DiscoveryReport.from_dict( + { + "schema_version": 1, + "producer": dict( + producer + or {"name": "motionkernel.fx_capture", "version": "1"} + ), + "workload": dict(workload), + "environment": dict( + environment + or { + "hardware_profile_id": "cpu", + "software_profile_id": "fx-capture", + } + ), + "total_cuda_time_us": float(total_cuda_time_us), + "operators": [], + "regions": [region.as_dict() for region in self.regions()], + "graph_breaks": [item.as_dict() for item in self.graph_breaks()], + "unsupported": [item.as_dict() for item in self.unsupported()], + } + ) + + +@contextmanager +def capture_model_regions( + root: Any, + *, + tracer: str = "symbolic", + predicate: Callable[[str, Any], bool] | None = None, + name_prefix: str = "region", +) -> Iterator[RegionCaptureSession]: + """Context manager: hook modules, run caller forward, yield session.""" + session = RegionCaptureSession(tracer=tracer, name_prefix=name_prefix) + try: + session.register_named_children(root, predicate=predicate) + yield session + finally: + session.close() diff --git a/autokernel/discovery/profiler_export.py b/autokernel/discovery/profiler_export.py new file mode 100644 index 00000000..8a7eeced --- /dev/null +++ b/autokernel/discovery/profiler_export.py @@ -0,0 +1,175 @@ +"""Load FastVideo metadata-only torch.profiler exports.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from .profiler_parse import parse_key_averages_rows +from .types import ( + DISCOVERY_SCHEMA_VERSION, + DiscoveryError, + DiscoveryReport, +) + +_EXPORT_FIELDS = { + "schema_version", + "producer", + "workload", + "environment", + "total_cuda_time_us", + "rows", + # Optional FX capture block, present only when the producer ran graph + # capture alongside the timing profile. Older exports omit all four. + "capture", + "regions", + "graph_breaks", + "unsupported", +} + +# Capture-block format the loader understands. Independent of the export's +# schema_version so a capture change does not invalidate timing-only readers. +SUPPORTED_CAPTURE_SCHEMA_VERSION = 1 + + +def _fail(source: object, location: str, message: str) -> DiscoveryError: + return DiscoveryError(f"profiler export {source!r}: {location}: {message}") + + +def _mapping( + value: Any, + *, + source: object, + location: str, + non_empty: bool = False, +) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or (non_empty and not value): + qualifier = "non-empty " if non_empty else "" + raise _fail(source, location, f"must be a {qualifier}object") + return value + + +def profiler_export_to_report( + raw_value: Any, + *, + source: object = "", +) -> DiscoveryReport: + """Validate a portable profiler export and create a discovery report.""" + raw = _mapping( + raw_value, + source=source, + location="top level", + non_empty=True, + ) + unknown = sorted(set(raw) - _EXPORT_FIELDS) + if unknown: + raise _fail(source, "top level", f"unknown field(s) {unknown}") + if raw.get("schema_version") != DISCOVERY_SCHEMA_VERSION: + raise _fail( + source, + "schema_version", + f"expected {DISCOVERY_SCHEMA_VERSION}", + ) + rows = raw.get("rows") + if not isinstance(rows, Sequence) or isinstance(rows, (str, bytes)): + raise _fail(source, "rows", "must be a list") + + try: + operators = parse_key_averages_rows(rows) + except (TypeError, ValueError) as exc: + raise _fail(source, "rows", str(exc)) from exc + + total = raw.get("total_cuda_time_us") + if total is None: + total = sum(max(item.self_cuda_time_us, 0.0) for item in operators) + + capture = raw.get("capture") + if capture is not None: + capture = _mapping( + capture, + source=source, + location="capture", + non_empty=True, + ) + version = capture.get("capture_schema_version") + if isinstance(version, bool) or not isinstance(version, int): + raise _fail(source, "capture.capture_schema_version", "must be an integer") + if version != SUPPORTED_CAPTURE_SCHEMA_VERSION: + raise _fail( + source, + "capture.capture_schema_version", + f"unsupported version {version}; " + f"expected {SUPPORTED_CAPTURE_SCHEMA_VERSION}", + ) + else: + # Captured payloads are versioned through the capture mapping; refuse + # capture data that arrives without its version declaration. + for name in ("regions", "graph_breaks", "unsupported"): + if raw.get(name): + raise _fail( + source, + "capture", + f"required when {name!r} is present", + ) + + def _list(name: str) -> list[Any]: + value = raw.get(name, []) + if value is None: + return [] + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise _fail(source, name, "must be a list") + return list(value) + + payload = { + "schema_version": DISCOVERY_SCHEMA_VERSION, + "producer": dict( + _mapping( + raw.get("producer"), + source=source, + location="producer", + non_empty=True, + ) + ), + "workload": dict( + _mapping( + raw.get("workload"), + source=source, + location="workload", + non_empty=True, + ) + ), + "environment": dict( + _mapping( + raw.get("environment"), + source=source, + location="environment", + non_empty=True, + ) + ), + "total_cuda_time_us": total, + "operators": [ + { + **item.as_dict(), + "op_key": item.op_key, + } + for item in operators + ], + "regions": _list("regions"), + "graph_breaks": _list("graph_breaks"), + "unsupported": _list("unsupported"), + } + return DiscoveryReport.from_dict(payload, source=source) + + +def load_profiler_export(path: str | Path) -> DiscoveryReport: + """Load a FastVideo profiler JSON artifact.""" + input_path = Path(path) + if not input_path.is_file(): + raise _fail(str(input_path), "file", "not found") + try: + raw = json.loads(input_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise _fail(str(input_path), "JSON", f"invalid JSON: {exc}") from exc + return profiler_export_to_report(raw, source=str(input_path)) diff --git a/autokernel/discovery/profiler_parse.py b/autokernel/discovery/profiler_parse.py new file mode 100644 index 00000000..9f213c98 --- /dev/null +++ b/autokernel/discovery/profiler_parse.py @@ -0,0 +1,113 @@ +"""Parse torch.profiler key-average style tables into OperatorHotspot rows. + +Accepts already-exported JSON lists (no live profiler required). Real GPU +collection is a separate step; this module only normalizes metadata. +""" + +from __future__ import annotations + +import re +from typing import Any, Mapping, Sequence + +from .types import OperatorHotspot + +_UNSAFE_OP_CHARACTERS = re.compile(r"[^A-Za-z0-9_./:-]") + + +def canonical_op_key(name: str) -> str: + normalized = _UNSAFE_OP_CHARACTERS.sub("_", name).strip("_") + return (normalized or "unknown_operator")[:256] + + +def _num(value: Any, default: float = 0.0) -> float: + if value is None: + return default + try: + return float(value) + except (TypeError, ValueError): + return default + + +def parse_key_averages_rows( + rows: Sequence[Mapping[str, Any]], + *, + source: str = "torch_profiler", +) -> tuple[OperatorHotspot, ...]: + """Convert list-of-dicts profiler exports into OperatorHotspot tuples. + + Recognized keys (flexible aliases): + - name / key / op_name + - cuda_time_total / cuda_time_us / device_time_total + - self_cuda_time_total / self_cuda_time_us + - cpu_time_total / cpu_time_us + - count / calls + """ + hotspots: list[OperatorHotspot] = [] + for index, row in enumerate(rows): + if not isinstance(row, Mapping): + raise TypeError(f"rows[{index}] must be a mapping") + name = ( + row.get("name") + or row.get("key") + or row.get("op_name") + or row.get("operator") + ) + if not name: + raise ValueError(f"rows[{index}] missing operator name") + name = str(name) + cuda = _num( + row.get("cuda_time_us", + row.get("cuda_time_total", + row.get("device_time_total", + row.get("cuda_time", 0)))), + ) + # Profiler often reports times in microseconds already; if a *_ms key + # is present, convert. + if "cuda_time_ms" in row: + cuda = _num(row["cuda_time_ms"]) * 1000.0 + self_cuda = _num( + row.get( + "self_cuda_time_us", + row.get("self_cuda_time_total", row.get("self_cuda_time", cuda)), + ) + ) + if "self_cuda_time_ms" in row: + self_cuda = _num(row["self_cuda_time_ms"]) * 1000.0 + cpu = _num(row.get("cpu_time_us", row.get("cpu_time_total", 0))) + if "cpu_time_ms" in row: + cpu = _num(row["cpu_time_ms"]) * 1000.0 + calls = row.get("calls", row.get("count", 1)) + try: + calls_i = int(calls) + except (TypeError, ValueError) as exc: + raise ValueError(f"rows[{index}].calls invalid") from exc + if calls_i <= 0: + continue + shapes_raw = row.get("input_shapes") or row.get("shapes") or [] + shapes: list[tuple[int, ...]] = [] + if isinstance(shapes_raw, Sequence) and not isinstance( + shapes_raw, (str, bytes) + ): + for shape in shapes_raw: + if isinstance(shape, Sequence) and not isinstance( + shape, (str, bytes) + ): + shapes.append(tuple(int(x) for x in shape)) + hotspots.append( + OperatorHotspot( + name=name, + op_key=canonical_op_key(name), + calls=calls_i, + cuda_time_us=cuda, + self_cuda_time_us=min(self_cuda, cuda) if cuda else self_cuda, + cpu_time_us=cpu, + input_shapes=tuple(shapes), + parent_module=( + str(row["parent_module"]) + if row.get("parent_module") is not None + else None + ), + source=source, + ) + ) + return tuple(hotspots) diff --git a/autokernel/discovery/ranking.py b/autokernel/discovery/ranking.py new file mode 100644 index 00000000..e5bd253b --- /dev/null +++ b/autokernel/discovery/ranking.py @@ -0,0 +1,181 @@ +"""Candidate impact ranking and Amdahl-style end-to-end ceilings. + +Uses measured CUDA time share and an optimistic reducible fraction to decide +whether a region is worth searching. Defaults match the universal plan: + +- impact floor: do not search when optimistic e2e improvement is below 0.5% +- promotion preference: candidates that can plausibly exceed 1% e2e gain +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + +from .safety import reject_region +from .types import GraphRegion, OperatorHotspot + + +DEFAULT_IMPACT_FLOOR = 0.005 # 0.5% optimistic end-to-end +DEFAULT_PROMOTION_TARGET = 0.01 # 1% +DEFAULT_REDUCIBLE_FRACTION = 0.9 # optimistic upper bound on kernel speedup + + +@dataclass(frozen=True) +class RankedCandidate: + """One discovery region with model-level value estimates.""" + + region: GraphRegion + share_of_e2e: float + estimated_reducible_fraction: float + estimated_max_e2e_improvement: float + confidence: float + search_worthy: bool + meets_promotion_target: bool + rejection_reasons: tuple[str, ...] + pattern_family: str | None = None + + def as_dict(self) -> dict: + return { + "name": self.region.name, + "fingerprint": self.region.fingerprint, + "share_of_e2e": self.share_of_e2e, + "estimated_reducible_fraction": self.estimated_reducible_fraction, + "estimated_max_e2e_improvement": self.estimated_max_e2e_improvement, + "confidence": self.confidence, + "search_worthy": self.search_worthy, + "meets_promotion_target": self.meets_promotion_target, + "rejection_reasons": list(self.rejection_reasons), + "pattern_family": self.pattern_family or self.region.pattern_family, + "calls": self.region.calls, + "cuda_time_us": self.region.cuda_time_us, + } + + +def e2e_share(cuda_time_us: float, total_cuda_time_us: float) -> float: + if total_cuda_time_us <= 0: + return 0.0 + return max(0.0, float(cuda_time_us) / float(total_cuda_time_us)) + + +def optimistic_e2e_improvement( + share: float, + *, + reducible_fraction: float = DEFAULT_REDUCIBLE_FRACTION, +) -> float: + """Amdahl-style upper bound: share * reducible_fraction. + + If a region is 2% of e2e and we can remove 90% of its time, the model + improves by at most ~1.8%. + """ + if reducible_fraction < 0.0 or reducible_fraction > 1.0: + raise ValueError("reducible_fraction must be in [0, 1]") + return max(0.0, share * reducible_fraction) + + +def classify_pattern_family(operations: Sequence[str]) -> str: + """Coarse family label for reports; not a correctness claim.""" + ops = " ".join(operations).lower() + if any(tok in ops for tok in ("layer_norm", "rms_norm", "native_layer_norm")): + if any(tok in ops for tok in ("mul", "add")): + return "residual_gate_norm" + return "normalization" + if any(tok in ops for tok in ("silu", "gelu", "relu", "sigmoid")): + return "activation_epilogue" + if any(tok in ops for tok in ("permute", "transpose", "contiguous", "clone", "to")): + if all( + any(x in o for x in ("permute", "transpose", "contiguous", "clone", "to", "view", "reshape")) + for o in operations + ): + return "layout_cast_copy" + if any(tok in ops for tok in ("mul", "add", "sub", "div")): + return "elementwise_chain" + return "unknown" + + +def rank_regions( + regions: Sequence[GraphRegion], + *, + total_cuda_time_us: float, + impact_floor: float = DEFAULT_IMPACT_FLOOR, + promotion_target: float = DEFAULT_PROMOTION_TARGET, + reducible_fraction: float = DEFAULT_REDUCIBLE_FRACTION, +) -> tuple[RankedCandidate, ...]: + """Rank graph regions by optimistic end-to-end value.""" + ranked: list[RankedCandidate] = [] + for region in regions: + share = e2e_share(region.self_cuda_time_us, total_cuda_time_us) + safety_reasons = tuple(reject_region(region.operations)) + if region.rejection_reasons: + safety_reasons = tuple( + dict.fromkeys([*region.rejection_reasons, *safety_reasons]) + ) + improvement = optimistic_e2e_improvement( + share, reducible_fraction=reducible_fraction + ) + safe = not safety_reasons + # Confidence: higher when more calls and pure allowlist. + confidence = 0.4 + if safe: + confidence += 0.3 + if region.calls >= 8: + confidence += 0.2 + if region.shape_frequency: + confidence += 0.1 + confidence = min(1.0, confidence) + + search_worthy = safe and improvement >= impact_floor + reasons = list(safety_reasons) + if safe and improvement < impact_floor: + reasons.append( + f"below_impact_floor: optimistic e2e {improvement:.4f} " + f"< floor {impact_floor:.4f}" + ) + search_worthy = False + + family = region.pattern_family or classify_pattern_family( + region.operations + ) + ranked.append( + RankedCandidate( + region=region, + share_of_e2e=share, + estimated_reducible_fraction=reducible_fraction, + estimated_max_e2e_improvement=improvement, + confidence=confidence, + search_worthy=search_worthy, + meets_promotion_target=improvement >= promotion_target, + rejection_reasons=tuple(reasons), + pattern_family=family, + ) + ) + + ranked.sort( + key=lambda item: ( + -item.estimated_max_e2e_improvement, + -item.share_of_e2e, + -item.confidence, + item.region.name, + ) + ) + return tuple(ranked) + + +def rank_operators( + operators: Sequence[OperatorHotspot], + *, + total_cuda_time_us: float, +) -> tuple[tuple[OperatorHotspot, float], ...]: + """Return operators sorted by self CUDA e2e share (descending). + + ``cuda_time_us`` is inclusive in torch profiler exports, so ranking with it + counts the same device work once for every enclosing record_function or + custom-op scope. ``total_cuda_time_us`` is the sum of self CUDA time and + operator shares must use the same accounting basis. + """ + rows = [ + (op, e2e_share(op.self_cuda_time_us, total_cuda_time_us)) + for op in operators + ] + rows.sort(key=lambda item: (-item[1], -item[0].calls, item[0].name)) + return tuple(rows) diff --git a/autokernel/discovery/safety.py b/autokernel/discovery/safety.py new file mode 100644 index 00000000..fc686865 --- /dev/null +++ b/autokernel/discovery/safety.py @@ -0,0 +1,127 @@ +"""Allowlist and rejection checks for captured graph regions. + +Fail closed: mutation, data-dependent Python control flow, collectives, +unknown aliasing, and unsupported custom ops are rejected before search. +""" + +from __future__ import annotations + +from typing import Iterable, Sequence + +# Minimal pure ATen/elementwise subset useful for early LTX candidates. +# Expand from real profile evidence, not from a desire to support all of PyTorch. +ALLOWED_ATEN_OPS: frozenset[str] = frozenset( + { + "aten::add", + "aten::mul", + "aten::sub", + "aten::div", + "aten::neg", + "aten::exp", + "aten::silu", + "aten::gelu", + "aten::relu", + "aten::sigmoid", + "aten::tanh", + "aten::rsqrt", + "aten::sqrt", + "aten::pow", + "aten::mean", + "aten::var", + "aten::layer_norm", + "aten::rms_norm", + "aten::native_layer_norm", + "aten::to", + "aten::clone", + "aten::contiguous", + "aten::view", + "aten::reshape", + "aten::permute", + "aten::transpose", + "aten::unsqueeze", + "aten::squeeze", + "aten::cat", + "aten::stack", + "aten::expand", + "aten::broadcast_to", + "aten::type_as", + } +) + +REJECT_SUBSTRINGS: tuple[tuple[str, str], ...] = ( + ("c10d", "collective communication"), + ("all_reduce", "collective communication"), + ("all_gather", "collective communication"), + ("reduce_scatter", "collective communication"), + ("barrier", "collective communication"), + ("aten::item", "data-dependent host sync"), + ("aten::nonzero", "data-dependent indexing"), + ("aten::unique", "data-dependent indexing"), + ("aten::argsort", "data-dependent ordering"), + ("aten::index_put", "mutation / complex indexing"), + ("aten::scatter", "mutation / scatter"), + ("aten::sort", "data-dependent ordering"), + ("aten::randint", "rng / nondeterminism boundary"), + ("aten::rand", "rng / nondeterminism boundary"), + ("aten::copy_", "mutation / aliasing write"), + ("unknown aliasing", "unknown aliasing"), + ("prims::", "unsupported prims op"), +) + + +def normalize_op_name(op_name: str) -> str: + name = op_name.strip() + if name.startswith("aten.") and not name.startswith("aten::"): + name = "aten::" + name[len("aten.") :] + # Strip schema overloads: aten::add.Tensor -> aten::add + if "." in name and name.startswith("aten::"): + base, _sep, _rest = name.partition(".") + # Keep dtype/device markers that use a single segment after :: only when + # they are overload names (add.Tensor). Nested module names stay intact. + if _rest and "." not in base: + name = base + return name + + +def reject_region( + operations: Sequence[str], + *, + allowlist: Iterable[str] | None = None, +) -> list[str]: + """Return rejection reasons for a candidate op sequence (empty if safe).""" + allowed = frozenset(allowlist) if allowlist is not None else ALLOWED_ATEN_OPS + reasons: list[str] = [] + if not operations: + return ["empty operation sequence"] + + for op in operations: + normalized = normalize_op_name(op) + if normalized.startswith("aten::") and normalized.endswith("_"): + reasons.append(f"{normalized}: in-place mutation") + continue + lower = normalized.lower() + for token, reason in REJECT_SUBSTRINGS: + if token in lower: + reasons.append(f"{normalized}: {reason}") + break + else: + if normalized not in allowed and not normalized.startswith("aten::"): + reasons.append(f"{normalized}: unsupported custom operator") + elif normalized not in allowed: + reasons.append(f"{normalized}: not in pure-tensor allowlist") + # Deduplicate while preserving order. + seen: set[str] = set() + unique: list[str] = [] + for reason in reasons: + if reason not in seen: + seen.add(reason) + unique.append(reason) + return unique + + +def is_region_safe( + operations: Sequence[str], + *, + allowlist: Iterable[str] | None = None, +) -> bool: + return not reject_region(operations, allowlist=allowlist) diff --git a/autokernel/discovery/types.py b/autokernel/discovery/types.py new file mode 100644 index 00000000..badc6568 --- /dev/null +++ b/autokernel/discovery/types.py @@ -0,0 +1,852 @@ +"""Metadata-only schemas for universal profiling and graph capture. + +These records are the Workstream 2 foundation: torch.profiler hotspots and +Dynamo/FX graph regions without requiring model-specific annotations. +Content, secrets, weights, activations, and prompts are forbidden. +""" + +from __future__ import annotations + +import json +import math +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +from .fingerprint import graph_fingerprint + +DISCOVERY_SCHEMA_VERSION = 1 + +_TOP_LEVEL_FIELDS = { + "schema_version", + "producer", + "workload", + "environment", + "total_cuda_time_us", + "operators", + "regions", + "graph_breaks", + "unsupported", +} +_OPERATOR_FIELDS = { + "name", + "op_key", + "calls", + "cuda_time_us", + "self_cuda_time_us", + "cpu_time_us", + "input_shapes", + "parent_module", + "source", + "attributes", +} +_REGION_FIELDS = { + "name", + "fingerprint", + "operations", + "dependencies", + "inputs", + "outputs", + "safe_constants", + "shape_frequency", + "cuda_time_us", + "self_cuda_time_us", + "calls", + "parent_module", + "pattern_family", + "rejection_reasons", + "attributes", +} +_TENSOR_FIELDS = { + "name", + "shape", + "stride", + "dtype", + "device_type", + "requires_grad", +} +_GRAPH_BREAK_FIELDS = { + "scope", + "reason", + "op_name", + "count", +} +_UNSUPPORTED_FIELDS = { + "op_name", + "reason", + "count", + "scope", +} + +_FORBIDDEN_KEYS = { + "credential", + "credentials", + "data", + "password", + "prompt", + "secret", + "secrets", + "tensor_values", + "token", + "values", + "weights", + "activations", +} +_OP_KEY_PATTERN = re.compile(r"^[A-Za-z0-9_./:-]{1,256}$") +_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") + + +class DiscoveryError(ValueError): + """Raised when a discovery/profile capture is malformed or unsafe.""" + + +def _fail(source: object, location: str, message: str) -> DiscoveryError: + return DiscoveryError(f"discovery report {source!r}: {location}: {message}") + + +def _mapping( + value: Any, + source: object, + location: str, + *, + non_empty: bool = False, +) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or (non_empty and not value): + qualifier = "non-empty " if non_empty else "" + raise _fail(source, location, f"must be a {qualifier}object") + for key in value: + if not isinstance(key, str) or not key: + raise _fail(source, location, "keys must be non-empty strings") + if key.lower() in _FORBIDDEN_KEYS: + raise _fail( + source, + f"{location}.{key}", + "content or secret fields are forbidden", + ) + return value + + +def _unknown_fields( + raw: Mapping[str, Any], + allowed: set[str], + source: object, + location: str, +) -> None: + unknown = sorted(set(raw) - allowed) + if unknown: + raise _fail(source, location, f"unknown field(s) {unknown}") + + +def _text(value: Any, source: object, location: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise _fail(source, location, "must be a non-empty string") + return value.strip() + + +def _optional_text(value: Any, source: object, location: str) -> str | None: + if value is None: + return None + return _text(value, source, location) + + +def _positive_int(value: Any, source: object, location: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise _fail(source, location, "must be a positive integer") + return value + + +def _non_negative_int(value: Any, source: object, location: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise _fail(source, location, "must be a non-negative integer") + return value + + +def _finite_non_negative(value: Any, source: object, location: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise _fail(source, location, "must be a finite non-negative number") + number = float(value) + if not math.isfinite(number) or number < 0: + raise _fail(source, location, "must be a finite non-negative number") + return number + + +def _metadata_value(value: Any, source: object, location: str) -> Any: + if value is None or isinstance(value, (str, bool)): + return value + if isinstance(value, (int, float)) and not isinstance(value, bool): + if not math.isfinite(float(value)): + raise _fail(source, location, "numbers must be finite") + return value + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [ + _metadata_value(item, source, f"{location}[{index}]") + for index, item in enumerate(value) + ] + if isinstance(value, Mapping): + result = {} + for key, item in _mapping(value, source, location).items(): + result[key] = _metadata_value(item, source, f"{location}.{key}") + return result + raise _fail(source, location, "must contain JSON metadata only") + + +@dataclass(frozen=True) +class TensorMeta: + """Layout metadata for one tensor without any values.""" + + name: str + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: str + device_type: str + requires_grad: bool = False + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "TensorMeta": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _TENSOR_FIELDS, source, location) + shape_raw = raw.get("shape") + stride_raw = raw.get("stride") + if not isinstance(shape_raw, Sequence) or isinstance(shape_raw, (str, bytes)): + raise _fail(source, f"{location}.shape", "must be a list") + if not isinstance(stride_raw, Sequence) or isinstance( + stride_raw, (str, bytes) + ): + raise _fail(source, f"{location}.stride", "must be a list") + shape = tuple( + _non_negative_int(dim, source, f"{location}.shape[{i}]") + for i, dim in enumerate(shape_raw) + ) + stride = [] + for i, value in enumerate(stride_raw): + if isinstance(value, bool) or not isinstance(value, int): + raise _fail(source, f"{location}.stride[{i}]", "must be an integer") + stride.append(value) + if len(stride) != len(shape): + raise _fail( + source, + f"{location}.stride", + "must have the same length as shape", + ) + requires_grad = raw.get("requires_grad", False) + if not isinstance(requires_grad, bool): + raise _fail(source, f"{location}.requires_grad", "must be a bool") + return cls( + name=_text(raw.get("name"), source, f"{location}.name"), + shape=shape, + stride=tuple(stride), + dtype=_text(raw.get("dtype"), source, f"{location}.dtype"), + device_type=_text( + raw.get("device_type"), source, f"{location}.device_type" + ), + requires_grad=requires_grad, + ) + + def as_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "shape": list(self.shape), + "stride": list(self.stride), + "dtype": self.dtype, + "device_type": self.device_type, + "requires_grad": self.requires_grad, + } + + def signature_dict(self) -> dict[str, Any]: + """Shape/dtype/layout only — used for fingerprinting.""" + return { + "shape": list(self.shape), + "stride": list(self.stride), + "dtype": self.dtype, + "device_type": self.device_type, + "requires_grad": self.requires_grad, + } + + +@dataclass(frozen=True) +class OperatorHotspot: + """One torch.profiler-attributed operator or ATen op aggregate.""" + + name: str + op_key: str + calls: int + cuda_time_us: float + self_cuda_time_us: float + cpu_time_us: float = 0.0 + input_shapes: tuple[tuple[int, ...], ...] = () + parent_module: str | None = None + source: str = "torch_profiler" + attributes: Mapping[str, Any] | None = None + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "OperatorHotspot": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _OPERATOR_FIELDS, source, location) + op_key = _text(raw.get("op_key"), source, f"{location}.op_key") + if not _OP_KEY_PATTERN.fullmatch(op_key): + raise _fail(source, f"{location}.op_key", "invalid op key") + shapes_raw = raw.get("input_shapes", []) + if not isinstance(shapes_raw, Sequence) or isinstance( + shapes_raw, (str, bytes) + ): + raise _fail(source, f"{location}.input_shapes", "must be a list") + shapes: list[tuple[int, ...]] = [] + for index, shape in enumerate(shapes_raw): + if not isinstance(shape, Sequence) or isinstance(shape, (str, bytes)): + raise _fail( + source, + f"{location}.input_shapes[{index}]", + "must be a list of dimensions", + ) + shapes.append( + tuple( + _non_negative_int( + dim, + source, + f"{location}.input_shapes[{index}][{dim_i}]", + ) + for dim_i, dim in enumerate(shape) + ) + ) + attributes = _metadata_value( + raw.get("attributes", {}), + source, + f"{location}.attributes", + ) + return cls( + name=_text(raw.get("name"), source, f"{location}.name"), + op_key=op_key, + calls=_positive_int(raw.get("calls"), source, f"{location}.calls"), + cuda_time_us=_finite_non_negative( + raw.get("cuda_time_us"), source, f"{location}.cuda_time_us" + ), + self_cuda_time_us=_finite_non_negative( + raw.get("self_cuda_time_us"), + source, + f"{location}.self_cuda_time_us", + ), + cpu_time_us=_finite_non_negative( + raw.get("cpu_time_us", 0), source, f"{location}.cpu_time_us" + ), + input_shapes=tuple(shapes), + parent_module=_optional_text( + raw.get("parent_module"), source, f"{location}.parent_module" + ), + source=_text( + raw.get("source", "torch_profiler"), + source, + f"{location}.source", + ), + attributes=attributes if attributes else None, + ) + + def impact_pct(self, total_cuda_time_us: float) -> float: + """Return this operator's exclusive share of measured CUDA time.""" + if total_cuda_time_us <= 0: + return 0.0 + return 100.0 * self.self_cuda_time_us / total_cuda_time_us + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "name": self.name, + "op_key": self.op_key, + "calls": self.calls, + "cuda_time_us": self.cuda_time_us, + "self_cuda_time_us": self.self_cuda_time_us, + "cpu_time_us": self.cpu_time_us, + "input_shapes": [list(shape) for shape in self.input_shapes], + "source": self.source, + } + if self.parent_module is not None: + payload["parent_module"] = self.parent_module + if self.attributes: + payload["attributes"] = dict(self.attributes) + return payload + + +@dataclass(frozen=True) +class GraphBreakRecord: + """A recorded Dynamo/FX graph break or capture limitation.""" + + scope: str + reason: str + op_name: str | None = None + count: int = 1 + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "GraphBreakRecord": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _GRAPH_BREAK_FIELDS, source, location) + return cls( + scope=_text(raw.get("scope"), source, f"{location}.scope"), + reason=_text(raw.get("reason"), source, f"{location}.reason"), + op_name=_optional_text( + raw.get("op_name"), source, f"{location}.op_name" + ), + count=_positive_int( + raw.get("count", 1), source, f"{location}.count" + ), + ) + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "scope": self.scope, + "reason": self.reason, + "count": self.count, + } + if self.op_name is not None: + payload["op_name"] = self.op_name + return payload + + +@dataclass(frozen=True) +class UnsupportedOpRecord: + """An observed op that cannot yet enter the search pipeline.""" + + op_name: str + reason: str + count: int = 1 + scope: str | None = None + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "UnsupportedOpRecord": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _UNSUPPORTED_FIELDS, source, location) + return cls( + op_name=_text(raw.get("op_name"), source, f"{location}.op_name"), + reason=_text(raw.get("reason"), source, f"{location}.reason"), + count=_positive_int( + raw.get("count", 1), source, f"{location}.count" + ), + scope=_optional_text( + raw.get("scope"), source, f"{location}.scope" + ), + ) + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "op_name": self.op_name, + "reason": self.reason, + "count": self.count, + } + if self.scope is not None: + payload["scope"] = self.scope + return payload + + +@dataclass(frozen=True) +class GraphRegion: + """One captured pure-tensor subgraph candidate (metadata only).""" + + name: str + fingerprint: str + operations: tuple[str, ...] + inputs: tuple[TensorMeta, ...] + outputs: tuple[TensorMeta, ...] = () + dependencies: tuple[str, ...] = () + safe_constants: Mapping[str, Any] | None = None + shape_frequency: Mapping[str, int] | None = None + cuda_time_us: float = 0.0 + self_cuda_time_us: float = 0.0 + calls: int = 1 + parent_module: str | None = None + pattern_family: str | None = None + rejection_reasons: tuple[str, ...] = () + attributes: Mapping[str, Any] | None = None + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "GraphRegion": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _REGION_FIELDS, source, location) + name = _text(raw.get("name"), source, f"{location}.name") + if not _NAME_PATTERN.fullmatch(name): + raise _fail(source, f"{location}.name", "invalid region name") + + ops_raw = raw.get("operations") + if ( + not isinstance(ops_raw, Sequence) + or isinstance(ops_raw, (str, bytes)) + or not ops_raw + ): + raise _fail( + source, f"{location}.operations", "must be a non-empty list" + ) + operations = tuple( + _text(op, source, f"{location}.operations[{i}]") + for i, op in enumerate(ops_raw) + ) + + inputs_raw = raw.get("inputs") + if ( + not isinstance(inputs_raw, Sequence) + or isinstance(inputs_raw, (str, bytes)) + or not inputs_raw + ): + raise _fail(source, f"{location}.inputs", "must be a non-empty list") + inputs = tuple( + TensorMeta.from_dict( + item, source=source, location=f"{location}.inputs[{i}]" + ) + for i, item in enumerate(inputs_raw) + ) + outputs_raw = raw.get("outputs", []) + if not isinstance(outputs_raw, Sequence) or isinstance( + outputs_raw, (str, bytes) + ): + raise _fail(source, f"{location}.outputs", "must be a list") + outputs = tuple( + TensorMeta.from_dict( + item, source=source, location=f"{location}.outputs[{i}]" + ) + for i, item in enumerate(outputs_raw) + ) + + deps_raw = raw.get("dependencies", []) + if not isinstance(deps_raw, Sequence) or isinstance( + deps_raw, (str, bytes) + ): + raise _fail(source, f"{location}.dependencies", "must be a list") + dependencies = tuple( + _text(dep, source, f"{location}.dependencies[{i}]") + for i, dep in enumerate(deps_raw) + ) + + rejection_raw = raw.get("rejection_reasons", []) + if not isinstance(rejection_raw, Sequence) or isinstance( + rejection_raw, (str, bytes) + ): + raise _fail( + source, f"{location}.rejection_reasons", "must be a list" + ) + rejection_reasons = tuple( + _text(reason, source, f"{location}.rejection_reasons[{i}]") + for i, reason in enumerate(rejection_raw) + ) + + fingerprint = _text( + raw.get("fingerprint"), source, f"{location}.fingerprint" + ) + parent_module = _optional_text( + raw.get("parent_module"), source, f"{location}.parent_module" + ) + expected = graph_fingerprint( + operations=operations, + input_signatures=[item.signature_dict() for item in inputs], + output_signatures=[item.signature_dict() for item in outputs], + safe_constants=raw.get("safe_constants") or {}, + parent_module=parent_module, + ) + # Recomputed fingerprint is the source of truth for equivalent regions. + if fingerprint != expected: + raise _fail( + source, + f"{location}.fingerprint", + "does not match canonical graph fingerprint", + ) + + shape_frequency = raw.get("shape_frequency", {}) + if shape_frequency is None: + shape_frequency = {} + shape_frequency = _mapping( + shape_frequency, source, f"{location}.shape_frequency" + ) + freq: dict[str, int] = {} + for key, count in shape_frequency.items(): + freq[key] = _positive_int( + count, source, f"{location}.shape_frequency.{key}" + ) + + attributes = _metadata_value( + raw.get("attributes", {}), + source, + f"{location}.attributes", + ) + return cls( + name=name, + fingerprint=fingerprint, + operations=operations, + inputs=inputs, + outputs=outputs, + dependencies=dependencies, + safe_constants=_metadata_value( + raw.get("safe_constants", {}), + source, + f"{location}.safe_constants", + ) + or None, + shape_frequency=freq or None, + cuda_time_us=_finite_non_negative( + raw.get("cuda_time_us", 0), + source, + f"{location}.cuda_time_us", + ), + self_cuda_time_us=_finite_non_negative( + raw.get("self_cuda_time_us", 0), + source, + f"{location}.self_cuda_time_us", + ), + calls=_positive_int( + raw.get("calls", 1), source, f"{location}.calls" + ), + parent_module=parent_module, + pattern_family=_optional_text( + raw.get("pattern_family"), source, f"{location}.pattern_family" + ), + rejection_reasons=rejection_reasons, + attributes=attributes if attributes else None, + ) + + @classmethod + def build( + cls, + *, + name: str, + operations: Sequence[str], + inputs: Sequence[TensorMeta], + outputs: Sequence[TensorMeta] = (), + parent_module: str | None = None, + safe_constants: Mapping[str, Any] | None = None, + **kwargs: Any, + ) -> "GraphRegion": + """Construct a region with a recomputed stable fingerprint.""" + if not _NAME_PATTERN.fullmatch(name): + raise ValueError(f"invalid graph region name: {name!r}") + fingerprint = graph_fingerprint( + operations=operations, + input_signatures=[item.signature_dict() for item in inputs], + output_signatures=[item.signature_dict() for item in outputs], + safe_constants=safe_constants, + parent_module=parent_module, + ) + return cls( + name=name, + fingerprint=fingerprint, + operations=tuple(operations), + inputs=tuple(inputs), + outputs=tuple(outputs), + parent_module=parent_module, + safe_constants=dict(safe_constants) if safe_constants else None, + **kwargs, + ) + + def impact_pct(self, total_cuda_time_us: float) -> float: + if total_cuda_time_us <= 0: + return 0.0 + return 100.0 * self.cuda_time_us / total_cuda_time_us + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "name": self.name, + "fingerprint": self.fingerprint, + "operations": list(self.operations), + "dependencies": list(self.dependencies), + "inputs": [item.as_dict() for item in self.inputs], + "outputs": [item.as_dict() for item in self.outputs], + "cuda_time_us": self.cuda_time_us, + "self_cuda_time_us": self.self_cuda_time_us, + "calls": self.calls, + "rejection_reasons": list(self.rejection_reasons), + } + if self.safe_constants: + payload["safe_constants"] = dict(self.safe_constants) + if self.shape_frequency: + payload["shape_frequency"] = dict(self.shape_frequency) + if self.parent_module is not None: + payload["parent_module"] = self.parent_module + if self.pattern_family is not None: + payload["pattern_family"] = self.pattern_family + if self.attributes: + payload["attributes"] = dict(self.attributes) + return payload + + +@dataclass(frozen=True) +class DiscoveryReport: + """Combined profiler + graph-capture report for one workload run.""" + + producer: Mapping[str, Any] + workload: Mapping[str, Any] + environment: Mapping[str, Any] + total_cuda_time_us: float + operators: tuple[OperatorHotspot, ...] + regions: tuple[GraphRegion, ...] = () + graph_breaks: tuple[GraphBreakRecord, ...] = () + unsupported: tuple[UnsupportedOpRecord, ...] = () + source: str = "" + schema_version: int = DISCOVERY_SCHEMA_VERSION + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object = "" + ) -> "DiscoveryReport": + raw = _mapping(raw_value, source, "top level", non_empty=True) + _unknown_fields(raw, _TOP_LEVEL_FIELDS, source, "top level") + version = raw.get("schema_version") + if isinstance(version, bool) or not isinstance(version, int): + raise _fail(source, "schema_version", "must be an integer") + if version != DISCOVERY_SCHEMA_VERSION: + raise _fail( + source, + "schema_version", + f"unsupported version {version}; expected {DISCOVERY_SCHEMA_VERSION}", + ) + producer = dict( + _metadata_value( + _mapping(raw.get("producer"), source, "producer", non_empty=True), + source, + "producer", + ) + ) + workload = dict( + _metadata_value( + _mapping(raw.get("workload"), source, "workload", non_empty=True), + source, + "workload", + ) + ) + environment = dict( + _metadata_value( + _mapping( + raw.get("environment"), + source, + "environment", + non_empty=True, + ), + source, + "environment", + ) + ) + for field in ("name", "version"): + _text(producer.get(field), source, f"producer.{field}") + for field in ("workload_id", "model_id"): + _text(workload.get(field), source, f"workload.{field}") + + operators_raw = raw.get("operators", []) + if not isinstance(operators_raw, Sequence) or isinstance( + operators_raw, (str, bytes) + ): + raise _fail(source, "operators", "must be a list") + operators = tuple( + OperatorHotspot.from_dict( + item, source=source, location=f"operators[{i}]" + ) + for i, item in enumerate(operators_raw) + ) + regions_raw = raw.get("regions", []) + if not isinstance(regions_raw, Sequence) or isinstance( + regions_raw, (str, bytes) + ): + raise _fail(source, "regions", "must be a list") + regions = tuple( + GraphRegion.from_dict( + item, source=source, location=f"regions[{i}]" + ) + for i, item in enumerate(regions_raw) + ) + breaks_raw = raw.get("graph_breaks", []) + if not isinstance(breaks_raw, Sequence) or isinstance( + breaks_raw, (str, bytes) + ): + raise _fail(source, "graph_breaks", "must be a list") + graph_breaks = tuple( + GraphBreakRecord.from_dict( + item, source=source, location=f"graph_breaks[{i}]" + ) + for i, item in enumerate(breaks_raw) + ) + unsupported_raw = raw.get("unsupported", []) + if not isinstance(unsupported_raw, Sequence) or isinstance( + unsupported_raw, (str, bytes) + ): + raise _fail(source, "unsupported", "must be a list") + unsupported = tuple( + UnsupportedOpRecord.from_dict( + item, source=source, location=f"unsupported[{i}]" + ) + for i, item in enumerate(unsupported_raw) + ) + return cls( + producer=producer, + workload=workload, + environment=environment, + total_cuda_time_us=_finite_non_negative( + raw.get("total_cuda_time_us"), + source, + "total_cuda_time_us", + ), + operators=operators, + regions=regions, + graph_breaks=graph_breaks, + unsupported=unsupported, + source=str(source), + schema_version=version, + ) + + def ranked_operators(self) -> tuple[OperatorHotspot, ...]: + return tuple( + sorted( + self.operators, + key=lambda op: (-op.self_cuda_time_us, -op.calls, op.name), + ) + ) + + def ranked_regions(self) -> tuple[GraphRegion, ...]: + return tuple( + sorted( + self.regions, + key=lambda region: ( + -region.cuda_time_us, + -region.calls, + region.name, + ), + ) + ) + + def as_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "producer": dict(self.producer), + "workload": dict(self.workload), + "environment": dict(self.environment), + "total_cuda_time_us": self.total_cuda_time_us, + "operators": [item.as_dict() for item in self.operators], + "regions": [item.as_dict() for item in self.regions], + "graph_breaks": [item.as_dict() for item in self.graph_breaks], + "unsupported": [item.as_dict() for item in self.unsupported], + } + + +def load_discovery_report(path: str | Path) -> DiscoveryReport: + """Load and validate a discovery report without importing torch.""" + file_path = Path(path) + if not file_path.is_file(): + raise DiscoveryError(f"discovery report {file_path!s}: file: not found") + try: + raw = json.loads(file_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise DiscoveryError( + f"discovery report {file_path!s}: JSON: invalid JSON: {exc}" + ) from exc + return DiscoveryReport.from_dict(raw, source=str(file_path)) + + +def write_discovery_report(report: DiscoveryReport, path: str | Path) -> None: + """Atomically write a discovery report.""" + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(output.suffix + ".tmp") + temporary.write_text( + json.dumps(report.as_dict(), indent=2) + "\n", encoding="utf-8" + ) + temporary.replace(output) diff --git a/autokernel/specs/__init__.py b/autokernel/specs/__init__.py new file mode 100644 index 00000000..29eaa8dc --- /dev/null +++ b/autokernel/specs/__init__.py @@ -0,0 +1,90 @@ +"""Kernel specifications: the public description of a benchmarkable operation. + +Typical use:: + + from autokernel.specs import create_builtin_registry, load_spec + + registry = create_builtin_registry() + spec = registry.get("matmul") + + external = load_spec("examples/custom_ops/add.py:SPEC", registry=registry) + +Importing this package never imports ``torch`` and never initializes a GPU. +""" + +from __future__ import annotations + +from .accounting import DT_BYTES, Expression, const, serialize_accounting, size +from .dtypes import ( + CANONICAL_DTYPES, + DTYPE_BYTES, + canonical_dtype_name, + dtype_bytes, + is_canonical_dtype, + resolve_torch_dtype, +) +from .lazy import LazyCallable, lazy_callable +from .loader import ( + SpecCollisionError, + SpecLoadError, + load_spec, + parse_locator, + resolve_spec, +) +from .registry import ( + DuplicateSpecError, + KernelRegistry, + SpecNotFoundError, + builtin_spec_names, + create_builtin_registry, +) +from .types import ( + STANDARD_SIZE_LABELS, + BackwardSpec, + CompileSpec, + EdgeCase, + InputMap, + KernelSpec, + OutputSpec, + SizeMap, + SpecValidationError, + Tolerance, + validate_spec, +) + +__all__ = [ + "CANONICAL_DTYPES", + "DTYPE_BYTES", + "DT_BYTES", + "STANDARD_SIZE_LABELS", + "BackwardSpec", + "CompileSpec", + "DuplicateSpecError", + "EdgeCase", + "Expression", + "InputMap", + "KernelRegistry", + "KernelSpec", + "LazyCallable", + "OutputSpec", + "SizeMap", + "SpecCollisionError", + "SpecLoadError", + "SpecNotFoundError", + "SpecValidationError", + "Tolerance", + "builtin_spec_names", + "canonical_dtype_name", + "const", + "create_builtin_registry", + "dtype_bytes", + "is_canonical_dtype", + "lazy_callable", + "load_spec", + "parse_locator", + "resolve_spec", + "resolve_torch_dtype", + "serialize_accounting", + "size", + "validate_spec", +] diff --git a/autokernel/specs/accounting.py b/autokernel/specs/accounting.py new file mode 100644 index 00000000..626d0e7c --- /dev/null +++ b/autokernel/specs/accounting.py @@ -0,0 +1,250 @@ +"""Serializable arithmetic expressions for FLOP and byte accounting. + +Performance accounting used to live in two places: real lambdas in ``bench.py`` +and Python source strings in ``extract.py``. Both are replaced by a small +expression tree that + +* evaluates like a normal callable (``expr(size)`` / ``expr(size, dt_bytes)``), +* serializes back to a safe numeric source expression for generated kernel + files, and +* reports which size keys it references so specifications can be validated. + +Nothing here evaluates untrusted text: :meth:`Expression.to_source` only emits +numbers, ``s["key"]`` lookups, ``dt_bytes`` and arithmetic operators, and the +loader never calls ``eval`` on specification content. + +Example: + >>> flops = 2 * size("M") * size("N") * size("K") + >>> flops({"M": 2, "N": 3, "K": 4}) + 48 + >>> flops.to_source() + '2 * s[\'M\'] * s[\'N\'] * s[\'K\']' +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping, Union + +__all__ = [ + "DT_BYTES", + "BinaryOp", + "Constant", + "DTypeBytes", + "Expression", + "SizeKey", + "const", + "serialize_accounting", + "size", +] + +Number = Union[int, float] + +# Operator precedence, used to decide where parentheses are required. +_PRECEDENCE: Mapping[str, int] = {"+": 1, "-": 1, "*": 2, "/": 2, "**": 3} + + +class Expression: + """Base class for accounting expressions. + + Instances are callable: ``expr(size)`` for FLOP accounting and + ``expr(size, dt_bytes)`` for byte accounting. + """ + + # -- evaluation ---------------------------------------------------- + def evaluate(self, size: Mapping[str, int], dt_bytes: int | None = None) -> Number: + raise NotImplementedError + + def __call__(self, size: Mapping[str, int], dt_bytes: int | None = None) -> Number: + return self.evaluate(size, dt_bytes) + + # -- introspection ------------------------------------------------- + def size_keys(self) -> frozenset[str]: + """Return every size key referenced by this expression.""" + return frozenset() + + def uses_dtype_bytes(self) -> bool: + """Return True when evaluation requires a ``dt_bytes`` argument.""" + return False + + # -- serialization ------------------------------------------------- + def to_source(self) -> str: + """Return an equivalent Python expression over ``s`` and ``dt_bytes``.""" + raise NotImplementedError + + @property + def precedence(self) -> int: + return 100 + + def _wrapped_source(self, parent_precedence: int, *, is_right: bool = False) -> str: + text = self.to_source() + if self.precedence < parent_precedence or ( + is_right and self.precedence == parent_precedence + ): + return f"({text})" + return text + + # -- operators ----------------------------------------------------- + def __add__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("+", self, _coerce(other)) + + def __radd__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("+", _coerce(other), self) + + def __sub__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("-", self, _coerce(other)) + + def __rsub__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("-", _coerce(other), self) + + def __mul__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("*", self, _coerce(other)) + + def __rmul__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("*", _coerce(other), self) + + def __truediv__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("/", self, _coerce(other)) + + def __rtruediv__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("/", _coerce(other), self) + + def __pow__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("**", self, _coerce(other)) + + def __rpow__(self, other: "Expression | Number") -> "Expression": + return BinaryOp("**", _coerce(other), self) + + +@dataclass(frozen=True) +class Constant(Expression): + """A literal number.""" + + value: Number + + def evaluate(self, size: Mapping[str, int], dt_bytes: int | None = None) -> Number: + return self.value + + def to_source(self) -> str: + return repr(self.value) + + +@dataclass(frozen=True) +class SizeKey(Expression): + """A lookup of one size key, serialized as ``s['key']``.""" + + key: str + + def evaluate(self, size: Mapping[str, int], dt_bytes: int | None = None) -> Number: + try: + return size[self.key] + except KeyError as exc: + raise KeyError( + f"size key {self.key!r} is missing from size mapping {sorted(size)!r}" + ) from exc + + def size_keys(self) -> frozenset[str]: + return frozenset({self.key}) + + def to_source(self) -> str: + return f"s[{self.key!r}]" + + +@dataclass(frozen=True) +class DTypeBytes(Expression): + """The element size, in bytes, of the dtype being benchmarked.""" + + def evaluate(self, size: Mapping[str, int], dt_bytes: int | None = None) -> Number: + if dt_bytes is None: + raise ValueError( + "this expression needs a dtype byte width; call it as expr(size, dt_bytes)" + ) + return dt_bytes + + def uses_dtype_bytes(self) -> bool: + return True + + def to_source(self) -> str: + return "dt_bytes" + + +@dataclass(frozen=True) +class BinaryOp(Expression): + """An arithmetic combination of two expressions.""" + + op: str + left: Expression + right: Expression + + def __post_init__(self) -> None: + if self.op not in _PRECEDENCE: + raise ValueError(f"unsupported operator {self.op!r}") + + def evaluate(self, size: Mapping[str, int], dt_bytes: int | None = None) -> Number: + left = self.left.evaluate(size, dt_bytes) + right = self.right.evaluate(size, dt_bytes) + if self.op == "+": + return left + right + if self.op == "-": + return left - right + if self.op == "*": + return left * right + if self.op == "/": + return left / right + return left ** right + + def size_keys(self) -> frozenset[str]: + return self.left.size_keys() | self.right.size_keys() + + def uses_dtype_bytes(self) -> bool: + return self.left.uses_dtype_bytes() or self.right.uses_dtype_bytes() + + @property + def precedence(self) -> int: + return _PRECEDENCE[self.op] + + def to_source(self) -> str: + precedence = self.precedence + # ``**`` is right-associative; every other supported operator is left. + left = self.left._wrapped_source(precedence, is_right=self.op == "**") + right = self.right._wrapped_source(precedence, is_right=self.op != "**") + return f"{left} {self.op} {right}" + + +#: Element size in bytes of the dtype currently being benchmarked. +DT_BYTES = DTypeBytes() + + +def size(key: str) -> SizeKey: + """Reference one size key, e.g. ``size("M")``.""" + if not isinstance(key, str) or not key: + raise ValueError(f"size key must be a non-empty string, got {key!r}") + return SizeKey(key) + + +def const(value: Number) -> Constant: + """Wrap a literal number as an expression.""" + return Constant(value) + + +def _coerce(value: "Expression | Number") -> Expression: + if isinstance(value, Expression): + return value + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError( + f"expected an accounting Expression or a number, got {type(value).__name__}" + ) + return Constant(value) + + +def serialize_accounting(fn: object) -> str | None: + """Return the source form of an accounting callable, or None if opaque. + + Generated kernel files embed the returned text. Callables that are not + :class:`Expression` instances (arbitrary Python functions supplied by an + external specification) are not serializable, so ``None`` is returned and + callers fall back to referencing the specification itself. + """ + if isinstance(fn, Expression): + return fn.to_source() + return None diff --git a/autokernel/specs/builtins.py b/autokernel/specs/builtins.py new file mode 100644 index 00000000..985dd645 --- /dev/null +++ b/autokernel/specs/builtins.py @@ -0,0 +1,420 @@ +"""Built-in kernel specifications. + +This module is the single source of truth for the nine operations that used to +be described three times: in ``bench.py::KERNEL_CONFIGS``, in ``extract.py``'s +metadata maps, and in the extraction templates. + +Sizes, dtypes, tolerances, edge cases and accounting formulas are carried over +unchanged from the pre-refactor harness; ``tests/test_builtin_specs.py`` freezes +them so a future edit cannot silently shrink benchmark coverage. + +Importing this module does not import ``torch``: reference implementations are +resolved lazily on first call. +""" + +from __future__ import annotations + +from pathlib import Path + +from .accounting import DT_BYTES, size +from .inputs import ( + gen_cross_entropy_inputs, + gen_flash_attention_inputs, + gen_fused_mlp_inputs, + gen_layernorm_inputs, + gen_matmul_inputs, + gen_reduce_inputs, + gen_rmsnorm_inputs, + gen_rotary_embedding_inputs, + gen_softmax_inputs, +) +from .lazy import lazy_callable +from .types import EdgeCase, KernelSpec, Tolerance + +__all__ = ["REPO_ROOT", "builtin_specs", "starter_kernels_for"] + +#: Repository root: ``/autokernel/specs/builtins.py`` -> ````. +REPO_ROOT = Path(__file__).resolve().parents[2] + +_KERNELS_DIR = REPO_ROOT / "kernels" + +# Tolerance sets shared by several operations, spelled out per operation below +# so a change to one never silently changes another. +_TOL_LOOSE = { + "float16": Tolerance(atol=1e-2, rtol=1e-2), + "bfloat16": Tolerance(atol=2e-2, rtol=2e-2), + "float32": Tolerance(atol=1e-4, rtol=1e-4), +} +_TOL_TIGHT = { + "float16": Tolerance(atol=1e-3, rtol=1e-3), + "bfloat16": Tolerance(atol=2e-3, rtol=2e-3), + "float32": Tolerance(atol=1e-5, rtol=1e-5), +} +_TOL_REDUCTION = { + "float16": Tolerance(atol=1e-2, rtol=1e-2), + "bfloat16": Tolerance(atol=1e-1, rtol=5e-2), +} + +_FP16_BF16_FP32 = ("float16", "bfloat16", "float32") +_FP16_BF16 = ("float16", "bfloat16") + + +def starter_kernels_for(name: str) -> dict[str, Path]: + """Return the Triton and CUDA starter kernel paths for a built-in operation.""" + starters: dict[str, Path] = {} + triton_path = _KERNELS_DIR / f"{name}.py" + if triton_path.is_file(): + starters["triton"] = triton_path + cuda_path = _KERNELS_DIR / "cuda" / f"{name}.py" + if cuda_path.is_file(): + starters["cuda"] = cuda_path + return starters + + +def _spec_matmul() -> KernelSpec: + return KernelSpec( + name="matmul", + reference_fn=lazy_callable("reference", "matmul_ref"), + input_generator=gen_matmul_inputs, + sizes=[ + ("tiny", {"M": 128, "N": 128, "K": 128}), + ("small", {"M": 512, "N": 512, "K": 512}), + ("medium", {"M": 1024, "N": 1024, "K": 1024}), + ("large", {"M": 2048, "N": 2048, "K": 2048}), + ("xlarge", {"M": 4096, "N": 4096, "K": 4096}), + ("tall", {"M": 8192, "N": 1024, "K": 1024}), + ("wide", {"M": 1024, "N": 8192, "K": 1024}), + ("deep_k", {"M": 1024, "N": 1024, "K": 8192}), + ("llm_qkv", {"M": 4096, "N": 4096, "K": 512}), + ("llm_mlp", {"M": 4096, "N": 11008, "K": 4096}), + ], + dtypes=_FP16_BF16_FP32, + tolerances=_TOL_LOOSE, + flops_fn=2 * size("M") * size("N") * size("K"), + bytes_fn=( + size("M") * size("K") + size("K") * size("N") + size("M") * size("N") + ) * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_1023", size={"M": 1023, "N": 1023, "K": 1023}), + EdgeCase(name="edge_4097", size={"M": 4097, "N": 4097, "K": 512}), + EdgeCase(name="edge_1537", size={"M": 1537, "N": 1537, "K": 1537}), + ), + shape_keys=("M", "N", "K"), + shape_aliases={"M": "M", "N": "N", "K": "K"}, + starter_kernels=starter_kernels_for("matmul"), + speedup_estimate="2-3x", + ) + + +def _spec_softmax() -> KernelSpec: + return KernelSpec( + name="softmax", + reference_fn=lazy_callable("reference", "softmax_ref"), + input_generator=gen_softmax_inputs, + sizes=[ + ("tiny", {"rows": 32, "cols": 128}), + ("small", {"rows": 256, "cols": 512}), + ("medium", {"rows": 1024, "cols": 1024}), + ("large", {"rows": 4096, "cols": 4096}), + ("xlarge", {"rows": 8192, "cols": 8192}), + ("wide", {"rows": 1024, "cols": 32768}), + ("narrow", {"rows": 32768, "cols": 128}), + ("vocab", {"rows": 4096, "cols": 50257}), + ], + dtypes=_FP16_BF16_FP32, + tolerances=_TOL_TIGHT, + # exp + sub + sum + div + max + flops_fn=5 * size("rows") * size("cols"), + # read + write + bytes_fn=2 * size("rows") * size("cols") * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_1023", size={"rows": 1023, "cols": 1023}), + EdgeCase(name="edge_4097", size={"rows": 4097, "cols": 4097}), + EdgeCase(name="edge_50257", size={"rows": 1024, "cols": 50257}), + ), + shape_keys=("rows", "cols"), + shape_aliases={"M": "rows", "N": "cols", "rows": "rows", "cols": "cols"}, + starter_kernels=starter_kernels_for("softmax"), + speedup_estimate="1.5-3x", + ) + + +def _spec_layernorm() -> KernelSpec: + return KernelSpec( + name="layernorm", + reference_fn=lazy_callable("reference", "layernorm_ref"), + input_generator=gen_layernorm_inputs, + sizes=[ + ("tiny", {"batch": 32, "dim": 128}), + ("small", {"batch": 256, "dim": 512}), + ("medium", {"batch": 1024, "dim": 1024}), + ("large", {"batch": 4096, "dim": 2048}), + ("xlarge", {"batch": 8192, "dim": 4096}), + ("wide", {"batch": 1024, "dim": 8192}), + ("llm_7b", {"batch": 4096, "dim": 4096}), + ("llm_13b", {"batch": 4096, "dim": 5120}), + ], + dtypes=_FP16_BF16_FP32, + tolerances=_TOL_TIGHT, + # mean, var, norm, scale, shift + flops_fn=8 * size("batch") * size("dim"), + bytes_fn=(2 * size("batch") * size("dim") + 2 * size("dim")) * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_1023", size={"batch": 1023, "dim": 1023}), + EdgeCase(name="edge_4097", size={"batch": 4097, "dim": 4097}), + ), + shape_keys=("batch", "dim"), + shape_aliases={ + "M": "batch", + "N": "dim", + "rows": "batch", + "cols": "dim", + "batch": "batch", + "dim": "dim", + }, + starter_kernels=starter_kernels_for("layernorm"), + speedup_estimate="1.5-3x", + ) + + +def _spec_flash_attention() -> KernelSpec: + bhsd = size("batch") * size("heads") * size("seq_len") * size("head_dim") + return KernelSpec( + name="flash_attention", + reference_fn=lazy_callable("reference", "flash_attention_ref"), + input_generator=gen_flash_attention_inputs, + sizes=[ + ("tiny", {"batch": 1, "heads": 4, "seq_len": 64, "head_dim": 64}), + ("small", {"batch": 2, "heads": 8, "seq_len": 256, "head_dim": 64}), + ("medium", {"batch": 2, "heads": 16, "seq_len": 512, "head_dim": 64}), + ("large", {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 64}), + ("xlarge", {"batch": 2, "heads": 32, "seq_len": 2048, "head_dim": 64}), + ("long", {"batch": 1, "heads": 32, "seq_len": 4096, "head_dim": 64}), + ("gqa", {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 128}), + ("llm_7b", {"batch": 1, "heads": 32, "seq_len": 2048, "head_dim": 128}), + ], + dtypes=_FP16_BF16, + tolerances=_TOL_LOOSE, + # 4*B*H*S^2*D FLOPs (Q@K^T + softmax + attn@V) + flops_fn=( + 4 * size("batch") * size("heads") * (size("seq_len") ** 2) * size("head_dim") + ), + # Q, K, V in and one output tensor out + bytes_fn=4 * bhsd * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_127", size={"batch": 1, "heads": 8, "seq_len": 127, "head_dim": 64}), + EdgeCase( + name="edge_1023", size={"batch": 1, "heads": 8, "seq_len": 1023, "head_dim": 64} + ), + ), + shape_keys=("batch", "heads", "seq_len", "head_dim"), + shape_aliases={ + "B": "batch", + "H": "heads", + "N": "seq_len", + "S": "seq_len", + "D": "head_dim", + "batch": "batch", + "heads": "heads", + "seq_len": "seq_len", + "head_dim": "head_dim", + }, + starter_kernels=starter_kernels_for("flash_attention"), + speedup_estimate="2-4x", + ) + + +def _spec_fused_mlp() -> KernelSpec: + return KernelSpec( + name="fused_mlp", + reference_fn=lazy_callable("reference", "fused_mlp_ref"), + input_generator=gen_fused_mlp_inputs, + sizes=[ + ("tiny", {"batch": 32, "dim": 128, "hidden": 256}), + ("small", {"batch": 256, "dim": 512, "hidden": 1024}), + ("medium", {"batch": 1024, "dim": 1024, "hidden": 2048}), + ("large", {"batch": 2048, "dim": 2048, "hidden": 5504}), + ("xlarge", {"batch": 4096, "dim": 4096, "hidden": 11008}), + ("llm_7b", {"batch": 2048, "dim": 4096, "hidden": 11008}), + ("llm_13b", {"batch": 2048, "dim": 5120, "hidden": 13824}), + ], + dtypes=_FP16_BF16_FP32, + tolerances=_TOL_LOOSE, + # gate_proj + up_proj + down_proj + flops_fn=2 * size("batch") * size("dim") * size("hidden") * 3, + bytes_fn=( + size("batch") * size("dim") + + size("hidden") * size("dim") * 3 + + size("batch") * size("dim") + ) * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_1023", size={"batch": 1023, "dim": 1024, "hidden": 2048}), + EdgeCase(name="edge_4097", size={"batch": 4097, "dim": 512, "hidden": 1024}), + ), + shape_keys=("batch", "dim", "hidden"), + shape_aliases={ + "M": "batch", + "N": "hidden", + "K": "dim", + "batch": "batch", + "dim": "dim", + "hidden": "hidden", + }, + starter_kernels=starter_kernels_for("fused_mlp"), + speedup_estimate="2-3x", + ) + + +def _spec_cross_entropy() -> KernelSpec: + return KernelSpec( + name="cross_entropy", + reference_fn=lazy_callable("reference", "cross_entropy_ref"), + input_generator=gen_cross_entropy_inputs, + sizes=[ + ("tiny", {"batch": 32, "vocab": 256}), + ("small", {"batch": 256, "vocab": 1024}), + ("medium", {"batch": 1024, "vocab": 4096}), + ("large", {"batch": 4096, "vocab": 32000}), + ("xlarge", {"batch": 8192, "vocab": 50257}), + ("llama", {"batch": 4096, "vocab": 32000}), + ("gpt2", {"batch": 4096, "vocab": 50257}), + ], + dtypes=_FP16_BF16_FP32, + tolerances={ + "float16": Tolerance(atol=1e-2, rtol=1e-2), + "bfloat16": Tolerance(atol=2e-2, rtol=2e-2), + "float32": Tolerance(atol=1e-5, rtol=1e-5), + }, + # log_softmax + nll + flops_fn=4 * size("batch") * size("vocab"), + bytes_fn=(size("batch") * size("vocab") + size("batch")) * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_1023", size={"batch": 1023, "vocab": 32000}), + EdgeCase(name="edge_50257", size={"batch": 4096, "vocab": 50257}), + ), + shape_keys=("batch", "vocab"), + shape_aliases={"batch": "batch", "vocab": "vocab"}, + starter_kernels=starter_kernels_for("cross_entropy"), + speedup_estimate="1.5-2x", + ) + + +def _spec_rotary_embedding() -> KernelSpec: + return KernelSpec( + name="rotary_embedding", + reference_fn=lazy_callable("reference", "rotary_embedding_ref"), + input_generator=gen_rotary_embedding_inputs, + sizes=[ + ("tiny", {"batch": 1, "heads": 4, "seq_len": 64, "head_dim": 64}), + ("small", {"batch": 2, "heads": 8, "seq_len": 256, "head_dim": 64}), + ("medium", {"batch": 2, "heads": 16, "seq_len": 512, "head_dim": 64}), + ("large", {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 128}), + ("xlarge", {"batch": 2, "heads": 32, "seq_len": 2048, "head_dim": 128}), + ("llm_7b", {"batch": 1, "heads": 32, "seq_len": 2048, "head_dim": 128}), + ("llm_13b", {"batch": 1, "heads": 40, "seq_len": 2048, "head_dim": 128}), + ], + dtypes=_FP16_BF16_FP32, + tolerances=_TOL_TIGHT, + # mul + add per element, x2 (cos and sin parts) + flops_fn=6 * size("batch") * size("heads") * size("seq_len") * size("head_dim"), + bytes_fn=( + size("batch") * size("heads") * size("seq_len") * size("head_dim") * 2 + + size("seq_len") * size("head_dim") + ) * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_127", size={"batch": 1, "heads": 8, "seq_len": 127, "head_dim": 64}), + EdgeCase( + name="edge_1023", size={"batch": 1, "heads": 8, "seq_len": 1023, "head_dim": 128} + ), + ), + shape_keys=("batch", "heads", "seq_len", "head_dim"), + shape_aliases={ + "B": "batch", + "H": "heads", + "N": "seq_len", + "S": "seq_len", + "D": "head_dim", + "batch": "batch", + "heads": "heads", + "seq_len": "seq_len", + "head_dim": "head_dim", + }, + starter_kernels=starter_kernels_for("rotary_embedding"), + speedup_estimate="1.5-2x", + ) + + +def _spec_rmsnorm() -> KernelSpec: + return KernelSpec( + name="rmsnorm", + reference_fn=lazy_callable("reference", "rmsnorm_ref"), + input_generator=gen_rmsnorm_inputs, + sizes=[ + ("small", {"M": 1024, "N": 768}), + ("medium", {"M": 4096, "N": 1024}), + ("large", {"M": 4096, "N": 4096}), + ("llama", {"M": 2048, "N": 4096}), + ], + dtypes=_FP16_BF16, + tolerances=_TOL_REDUCTION, + # square, mean, sqrt, div, mul + flops_fn=6 * size("M") * size("N"), + bytes_fn=(2 * size("M") * size("N") + size("N")) * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_1023", size={"M": 1023, "N": 768}), + EdgeCase(name="edge_4097", size={"M": 4097, "N": 1024}), + ), + shape_keys=("M", "N"), + shape_aliases={"M": "M", "N": "N"}, + starter_kernels=starter_kernels_for("rmsnorm"), + speedup_estimate="1.5-3x", + ) + + +def _spec_reduce() -> KernelSpec: + return KernelSpec( + name="reduce", + reference_fn=lazy_callable("reference", "reduce_sum_ref"), + input_generator=gen_reduce_inputs, + sizes=[ + ("small", {"M": 1024, "N": 1024}), + ("medium", {"M": 4096, "N": 4096}), + ("large", {"M": 8192, "N": 8192}), + ("wide", {"M": 1024, "N": 32768}), + ], + dtypes=_FP16_BF16, + tolerances=_TOL_REDUCTION, + # N additions per row + flops_fn=size("M") * size("N"), + bytes_fn=(size("M") * size("N") + size("M")) * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_1023", size={"M": 1023, "N": 1024}), + EdgeCase(name="edge_4097", size={"M": 4096, "N": 4097}), + ), + shape_keys=("M", "N"), + shape_aliases={"M": "M", "N": "N"}, + starter_kernels=starter_kernels_for("reduce"), + speedup_estimate="1.5-2x", + # Extraction fell back to 4096x4096 (not the 8192x8192 'large' size) + # before the registry existed; keep that fallback byte-for-byte. + default_shape={"M": 4096, "N": 4096}, + ) + + +#: Factories in the order the operations were declared by ``KERNEL_CONFIGS``. +_BUILTIN_FACTORIES = ( + _spec_matmul, + _spec_softmax, + _spec_layernorm, + _spec_flash_attention, + _spec_fused_mlp, + _spec_cross_entropy, + _spec_rotary_embedding, + _spec_rmsnorm, + _spec_reduce, +) + + +def builtin_specs() -> tuple[KernelSpec, ...]: + """Build every built-in specification, in canonical order.""" + return tuple(factory() for factory in _BUILTIN_FACTORIES) diff --git a/autokernel/specs/dtypes.py b/autokernel/specs/dtypes.py new file mode 100644 index 00000000..91ded8ba --- /dev/null +++ b/autokernel/specs/dtypes.py @@ -0,0 +1,77 @@ +"""Canonical dtype names used at the specification boundary. + +Specifications only ever name dtypes with the canonical strings below, which +keeps discovery serializable and importable without ``torch``. Runtime code +translates to ``torch.dtype`` through :func:`resolve_torch_dtype`. +""" + +from __future__ import annotations + +from typing import Any, Mapping + +__all__ = [ + "CANONICAL_DTYPES", + "DTYPE_BYTES", + "canonical_dtype_name", + "dtype_bytes", + "is_canonical_dtype", + "resolve_torch_dtype", +] + +#: Ordered tuple of dtype names a specification may declare. +CANONICAL_DTYPES: tuple[str, ...] = ("float16", "bfloat16", "float32") + +#: Byte width per canonical dtype. Declared statically so byte accounting works +#: without importing torch. +DTYPE_BYTES: Mapping[str, int] = { + "float16": 2, + "bfloat16": 2, + "float32": 4, +} + + +def is_canonical_dtype(name: object) -> bool: + """Return True when ``name`` is one of the canonical dtype strings.""" + return isinstance(name, str) and name in CANONICAL_DTYPES + + +def canonical_dtype_name(dtype: Any) -> str: + """Normalize ``dtype`` (canonical string or ``torch.dtype``) to a canonical name. + + Raises: + ValueError: if the dtype is not supported at the specification boundary. + """ + if is_canonical_dtype(dtype): + return str(dtype) + + text = str(dtype) + if text.startswith("torch."): + text = text[len("torch."):] + if is_canonical_dtype(text): + return text + + raise ValueError( + f"unsupported dtype {dtype!r}; expected one of {', '.join(CANONICAL_DTYPES)}" + ) + + +def dtype_bytes(dtype: Any) -> int: + """Return the byte width of a canonical dtype name or ``torch.dtype``.""" + return DTYPE_BYTES[canonical_dtype_name(dtype)] + + +def resolve_torch_dtype(dtype: Any) -> Any: + """Translate a canonical dtype name to a ``torch.dtype``. + + ``torch`` is imported lazily so specification discovery stays torch-free. + Passing a ``torch.dtype`` through is supported and returns it unchanged + after validation. + """ + import torch # local import: keep module import torch-free + + if isinstance(dtype, torch.dtype): + canonical_dtype_name(dtype) # validate it is supported + return dtype + + name = canonical_dtype_name(dtype) + return getattr(torch, name) diff --git a/autokernel/specs/inputs.py b/autokernel/specs/inputs.py new file mode 100644 index 00000000..be7bdc5d --- /dev/null +++ b/autokernel/specs/inputs.py @@ -0,0 +1,137 @@ +"""Deterministic input generators for the built-in operations. + +Moved verbatim (same tensor creation order, same seeding) out of ``bench.py`` so +one specification owns both the shapes and the inputs for an operation. + +Every generator has the signature:: + + generator(size, dtype, device, seed=42) -> dict[str, Any] + +``dtype`` accepts a canonical dtype name (``"float16"``) or a ``torch.dtype``. +``torch`` is imported inside the generators so importing this module -- and +therefore discovering specifications -- never pulls in torch or a GPU context. +""" + +from __future__ import annotations + +from typing import Any, Mapping + +from .dtypes import resolve_torch_dtype + +__all__ = [ + "gen_cross_entropy_inputs", + "gen_flash_attention_inputs", + "gen_fused_mlp_inputs", + "gen_layernorm_inputs", + "gen_matmul_inputs", + "gen_reduce_inputs", + "gen_rmsnorm_inputs", + "gen_rotary_embedding_inputs", + "gen_softmax_inputs", +] + +SizeMap = Mapping[str, int] + + +def _prepare(dtype: Any, seed: int): + """Resolve the dtype and seed the global RNG (deterministic per seed).""" + import torch + + torch.manual_seed(seed) + return resolve_torch_dtype(dtype) + + +def gen_matmul_inputs(size: SizeMap, dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + dtype = _prepare(dtype, seed) + M, N, K = size["M"], size["N"], size["K"] + A = torch.randn(M, K, device=device, dtype=dtype) + B = torch.randn(K, N, device=device, dtype=dtype) + return {"A": A, "B": B} + + +def gen_softmax_inputs(size: SizeMap, dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + dtype = _prepare(dtype, seed) + rows, cols = size["rows"], size["cols"] + x = torch.randn(rows, cols, device=device, dtype=dtype) + return {"x": x} + + +def gen_layernorm_inputs(size: SizeMap, dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + dtype = _prepare(dtype, seed) + batch, dim = size["batch"], size["dim"] + x = torch.randn(batch, dim, device=device, dtype=dtype) + weight = torch.randn(dim, device=device, dtype=dtype) * 0.1 + bias = torch.randn(dim, device=device, dtype=dtype) * 0.1 + return {"x": x, "weight": weight, "bias": bias} + + +def gen_flash_attention_inputs(size: SizeMap, dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + dtype = _prepare(dtype, seed) + batch, heads = size["batch"], size["heads"] + seq_len, head_dim = size["seq_len"], size["head_dim"] + Q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=dtype) + K = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=dtype) + V = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=dtype) + return {"Q": Q, "K": K, "V": V} + + +def gen_fused_mlp_inputs(size: SizeMap, dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + dtype = _prepare(dtype, seed) + batch, dim, hidden = size["batch"], size["dim"], size["hidden"] + x = torch.randn(batch, dim, device=device, dtype=dtype) + w_gate = torch.randn(hidden, dim, device=device, dtype=dtype) * 0.02 + w_up = torch.randn(hidden, dim, device=device, dtype=dtype) * 0.02 + w_down = torch.randn(dim, hidden, device=device, dtype=dtype) * 0.02 + return {"x": x, "w_gate": w_gate, "w_up": w_up, "w_down": w_down} + + +def gen_cross_entropy_inputs(size: SizeMap, dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + dtype = _prepare(dtype, seed) + batch, vocab = size["batch"], size["vocab"] + logits = torch.randn(batch, vocab, device=device, dtype=dtype) + targets = torch.randint(0, vocab, (batch,), device=device, dtype=torch.long) + return {"logits": logits, "targets": targets} + + +def gen_rotary_embedding_inputs(size: SizeMap, dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + dtype = _prepare(dtype, seed) + batch, heads = size["batch"], size["heads"] + seq_len, head_dim = size["seq_len"], size["head_dim"] + x = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=dtype) + half_dim = head_dim // 2 + cos = torch.randn(seq_len, half_dim, device=device, dtype=dtype) + sin = torch.randn(seq_len, half_dim, device=device, dtype=dtype) + return {"x": x, "cos": cos, "sin": sin} + + +def gen_rmsnorm_inputs(size: SizeMap, dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + dtype = _prepare(dtype, seed) + M, N = size["M"], size["N"] + x = torch.randn(M, N, device=device, dtype=dtype) + weight = torch.randn(N, device=device, dtype=dtype) + return {"x": x, "weight": weight} + + +def gen_reduce_inputs(size: SizeMap, dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + dtype = _prepare(dtype, seed) + M, N = size["M"], size["N"] + x = torch.randn(M, N, device=device, dtype=dtype) + return {"x": x} diff --git a/autokernel/specs/lazy.py b/autokernel/specs/lazy.py new file mode 100644 index 00000000..8c807435 --- /dev/null +++ b/autokernel/specs/lazy.py @@ -0,0 +1,57 @@ +"""Lazily resolved callables. + +Built-in specifications point at reference implementations that live in +``reference.py``, which imports ``torch``. Registry discovery must work without +importing torch (and must never touch a GPU), so specifications hold a +:class:`LazyCallable` that imports the target module on first call instead of at +module import time. +""" + +from __future__ import annotations + +import importlib +from dataclasses import dataclass, field +from typing import Any, Callable + +__all__ = ["LazyCallable", "lazy_callable"] + + +@dataclass(frozen=True) +class LazyCallable: + """A callable that resolves ``module_name:attribute`` on first invocation.""" + + module_name: str + attribute: str + _cache: dict = field(default_factory=dict, repr=False, compare=False) + + def resolve(self) -> Callable[..., Any]: + """Import and return the target callable.""" + cached = self._cache.get("fn") + if cached is not None: + return cached + + module = importlib.import_module(self.module_name) + try: + fn = getattr(module, self.attribute) + except AttributeError as exc: + raise AttributeError( + f"module {self.module_name!r} has no attribute {self.attribute!r}" + ) from exc + if not callable(fn): + raise TypeError( + f"{self.module_name}:{self.attribute} is not callable " + f"(got {type(fn).__name__})" + ) + self._cache["fn"] = fn + return fn + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + return self.resolve()(*args, **kwargs) + + def __str__(self) -> str: + return f"{self.module_name}:{self.attribute}" + + +def lazy_callable(module_name: str, attribute: str) -> LazyCallable: + """Build a :class:`LazyCallable` for ``module_name:attribute``.""" + return LazyCallable(module_name=module_name, attribute=attribute) diff --git a/autokernel/specs/loader.py b/autokernel/specs/loader.py new file mode 100644 index 00000000..e65afac8 --- /dev/null +++ b/autokernel/specs/loader.py @@ -0,0 +1,240 @@ +"""Load kernel specifications supplied from outside the repository. + +Supported locators:: + + package.module:SPEC + /absolute/path/to/spec.py:SPEC + relative/path/to/spec.py:SPEC + +The selected attribute may be a :class:`~autokernel.specs.types.KernelSpec` or a +zero-argument callable returning one. + +Trust boundary: loading a specification imports and executes Python supplied by +the caller, exactly like ``python -c`` would. Only pass locators you trust. +Nothing here calls ``eval``/``exec`` on specification *data*, and ``sys.path`` is +never mutated permanently -- file locators are imported through +``importlib.util.spec_from_file_location`` under a unique module name. +""" + +from __future__ import annotations + +import importlib +import importlib.util +import os +import sys +import uuid +from pathlib import Path +from typing import Any + +from .registry import KernelRegistry, create_builtin_registry +from .types import KernelSpec, SpecValidationError, validate_spec + +__all__ = [ + "SpecCollisionError", + "SpecLoadError", + "load_spec", + "parse_locator", + "resolve_spec", +] + + +class SpecLoadError(ValueError): + """Raised when an external specification cannot be loaded.""" + + +class SpecCollisionError(SpecLoadError): + """Raised when an external specification shadows a registered name.""" + + +def parse_locator(locator: str) -> tuple[str, str]: + """Split ``target:ATTRIBUTE`` into its two halves. + + Windows-style drive letters are handled because the split happens at the + last colon. + """ + if not isinstance(locator, str) or not locator.strip(): + raise SpecLoadError( + "spec locator must be a non-empty string of the form " + "'module:ATTRIBUTE' or 'path/to/spec.py:ATTRIBUTE'" + ) + text = locator.strip() + if ":" not in text: + raise SpecLoadError( + f"invalid spec locator {locator!r}: expected 'module:ATTRIBUTE' or " + f"'path/to/spec.py:ATTRIBUTE' (the attribute name is required)" + ) + target, _, attribute = text.rpartition(":") + target = target.strip() + attribute = attribute.strip() + if not target or not attribute: + raise SpecLoadError( + f"invalid spec locator {locator!r}: both the module/path and the " + f"attribute name are required" + ) + return target, attribute + + +def _looks_like_path(target: str) -> bool: + if target.endswith(".py"): + return True + if os.sep in target: + return True + if os.altsep and os.altsep in target: + return True + return Path(target).is_file() + + +def _import_from_path(target: str, locator: str) -> Any: + path = Path(target).expanduser() + if not path.is_absolute(): + path = (Path.cwd() / path).resolve() + if not path.is_file(): + raise SpecLoadError( + f"cannot load spec {locator!r}: file not found: {path}" + ) + + module_name = f"_autokernel_external_spec_{uuid.uuid4().hex}" + module_spec = importlib.util.spec_from_file_location(module_name, path) + if module_spec is None or module_spec.loader is None: + raise SpecLoadError( + f"cannot load spec {locator!r}: {path} is not an importable Python file" + ) + module = importlib.util.module_from_spec(module_spec) + # Register during exec so dataclasses and relative lookups inside the module + # resolve. The returned objects retain the module globals they need, so the + # temporary UUID module can always be evicted afterward. + sys.modules[module_name] = module + try: + module_spec.loader.exec_module(module) + except Exception as exc: + raise SpecLoadError( + f"cannot load spec {locator!r}: importing {path} raised " + f"{type(exc).__name__}: {exc}" + ) from exc + finally: + sys.modules.pop(module_name, None) + return module + + +def _import_module(target: str, locator: str) -> Any: + try: + return importlib.import_module(target) + except ModuleNotFoundError as exc: + raise SpecLoadError( + f"cannot load spec {locator!r}: module {target!r} not found " + f"({exc}). Use 'path/to/spec.py:ATTRIBUTE' to load from a file." + ) from exc + except Exception as exc: + raise SpecLoadError( + f"cannot load spec {locator!r}: importing module {target!r} raised " + f"{type(exc).__name__}: {exc}" + ) from exc + + +def load_spec( + locator: str, + *, + registry: KernelRegistry | None = None, + override: bool = False, + require_standard_sizes: bool = True, +) -> KernelSpec: + """Load, validate and return the specification named by ``locator``. + + Args: + locator: ``module:ATTRIBUTE`` or ``path/to/spec.py:ATTRIBUTE``. + registry: when given, the loaded name is checked against it for + collisions. + override: allow the loaded specification to shadow a name that already + exists in ``registry``. + require_standard_sizes: require ``small``/``medium``/``large`` sizes. + + Raises: + SpecLoadError: for a missing module or file, a missing attribute, an + attribute of the wrong type, or a factory that misbehaves. + SpecCollisionError: when the name collides and ``override`` is False. + SpecValidationError: when the specification itself is malformed. + """ + target, attribute = parse_locator(locator) + + if _looks_like_path(target): + module = _import_from_path(target, locator) + else: + module = _import_module(target, locator) + + if not hasattr(module, attribute): + available = sorted( + name + for name, value in vars(module).items() + if isinstance(value, KernelSpec) and not name.startswith("_") + ) + hint = f" Available KernelSpec attributes: {', '.join(available)}." if available else "" + raise SpecLoadError( + f"cannot load spec {locator!r}: {getattr(module, '__name__', target)!r} has no " + f"attribute {attribute!r}.{hint}" + ) + + obj = getattr(module, attribute) + if isinstance(obj, KernelSpec): + spec = obj + elif callable(obj): + try: + spec = obj() + except SpecValidationError: + raise + except Exception as exc: + raise SpecLoadError( + f"cannot load spec {locator!r}: calling {attribute!r} raised " + f"{type(exc).__name__}: {exc}" + ) from exc + if not isinstance(spec, KernelSpec): + raise SpecLoadError( + f"cannot load spec {locator!r}: {attribute!r} returned " + f"{type(spec).__name__}, expected a KernelSpec" + ) + else: + raise SpecLoadError( + f"cannot load spec {locator!r}: {attribute!r} is a " + f"{type(obj).__name__}, expected a KernelSpec or a zero-argument " + f"callable returning one" + ) + + validate_spec(spec, require_standard_sizes=require_standard_sizes) + + if registry is not None and registry.contains(spec.name) and not override: + raise SpecCollisionError( + f"cannot load spec {locator!r}: operation name {spec.name!r} is already " + f"registered; rename the spec or pass --spec-override to replace it" + ) + + return spec + + +def resolve_spec( + *, + spec_locator: str | None = None, + name: str | None = None, + registry: KernelRegistry | None = None, + override: bool = False, +) -> tuple[KernelSpec, KernelRegistry]: + """Select the specification for one command invocation. + + Precedence is ``spec_locator`` first, then ``name``. The returned registry is + isolated: an externally loaded specification is registered into it and never + into a process-wide global. + + Returns: + ``(spec, registry)``. + """ + registry = registry if registry is not None else create_builtin_registry() + + if spec_locator: + spec = load_spec(spec_locator, registry=registry, override=override) + registry.register(spec, override=True) + return spec, registry + + if not name: + raise SpecLoadError( + "no operation selected: pass --spec LOCATOR or --kernel NAME" + ) + + return registry.get(name), registry diff --git a/autokernel/specs/registry.py b/autokernel/specs/registry.py new file mode 100644 index 00000000..9db78478 --- /dev/null +++ b/autokernel/specs/registry.py @@ -0,0 +1,120 @@ +"""A small, explicit registry of kernel specifications. + +The registry is an ordinary object, not a module-level global: every command and +every test can build an isolated registry so registering an external +specification never leaks into another run. +""" + +from __future__ import annotations + +from typing import Iterable, Iterator + +from .types import KernelSpec, SpecValidationError, validate_spec + +__all__ = [ + "DuplicateSpecError", + "KernelRegistry", + "SpecNotFoundError", + "builtin_spec_names", + "create_builtin_registry", +] + + +class DuplicateSpecError(ValueError): + """Raised when a name is registered twice without an explicit override.""" + + +class SpecNotFoundError(KeyError): + """Raised when a requested specification name is not registered.""" + + def __str__(self) -> str: # KeyError repr would quote the whole message + return self.args[0] if self.args else super().__str__() + + +class KernelRegistry: + """An ordered collection of :class:`KernelSpec` objects keyed by name. + + Ordering is insertion order, which keeps ``list_names()`` deterministic and + keeps the CLI's "available operations" listing stable. + """ + + def __init__(self, specs: Iterable[KernelSpec] = ()) -> None: + self._specs: dict[str, KernelSpec] = {} + for spec in specs: + self.register(spec) + + # -- mutation ------------------------------------------------------ + def register(self, spec: KernelSpec, *, override: bool = False) -> None: + """Register a specification. + + Args: + spec: the specification to add. + override: replace an existing specification with the same name. + + Raises: + SpecValidationError: if ``spec`` is not a valid specification. + DuplicateSpecError: if the name exists and ``override`` is False. + """ + if not isinstance(spec, KernelSpec): + raise SpecValidationError( + f"can only register KernelSpec objects, got {type(spec).__name__}" + ) + validate_spec(spec) + if spec.name in self._specs and not override: + raise DuplicateSpecError( + f"kernel spec {spec.name!r} is already registered; pass override=True " + f"to replace it" + ) + self._specs[spec.name] = spec + + # -- lookup -------------------------------------------------------- + def get(self, name: str) -> KernelSpec: + """Return the specification registered under ``name``.""" + try: + return self._specs[name] + except KeyError: + available = ", ".join(self.list_names()) or "" + raise SpecNotFoundError( + f"unknown kernel spec {name!r}; available: {available}" + ) from None + + def contains(self, name: str) -> bool: + """Return True when ``name`` is registered.""" + return name in self._specs + + def list_names(self) -> tuple[str, ...]: + """Return registered names in registration order.""" + return tuple(self._specs) + + def specs(self) -> tuple[KernelSpec, ...]: + """Return registered specifications in registration order.""" + return tuple(self._specs.values()) + + # -- dunder -------------------------------------------------------- + def __contains__(self, name: object) -> bool: + return isinstance(name, str) and name in self._specs + + def __iter__(self) -> Iterator[KernelSpec]: + return iter(self._specs.values()) + + def __len__(self) -> int: + return len(self._specs) + + def __repr__(self) -> str: + return f"KernelRegistry({list(self._specs)!r})" + + +def create_builtin_registry() -> KernelRegistry: + """Build a fresh registry containing every built-in operation. + + Importing the built-in specifications does not import ``torch`` and does not + initialize a GPU, so this is safe on a CPU-only machine. + """ + from .builtins import builtin_specs # local import avoids an import cycle + + return KernelRegistry(builtin_specs()) + + +def builtin_spec_names() -> tuple[str, ...]: + """Return the built-in operation names in their canonical order.""" + return create_builtin_registry().list_names() diff --git a/autokernel/specs/types.py b/autokernel/specs/types.py new file mode 100644 index 00000000..5ed88426 --- /dev/null +++ b/autokernel/specs/types.py @@ -0,0 +1,836 @@ +"""Public specification types for kernel operations. + +A :class:`KernelSpec` is the single source of truth for one benchmarkable +operation: its reference implementation, deterministic input generator, sizes, +dtypes, tolerances, edge cases, performance accounting, extraction metadata and +starter kernels. + +Design constraints: + +* dtypes are canonical strings (``float16``, ``bfloat16``, ``float32``) so a + specification can be inspected on a CPU-only machine; translation to + ``torch.dtype`` happens in runtime code only; +* importing this module must not import ``torch`` and must not touch a GPU; +* every validation error names the specification and the offending field. +""" + +from __future__ import annotations + +import math +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Iterable, Mapping, Sequence + +from .accounting import Expression +from .dtypes import CANONICAL_DTYPES, is_canonical_dtype + +__all__ = [ + "STANDARD_SIZE_LABELS", + "BackwardSpec", + "BytesFn", + "CompileSpec", + "EdgeCase", + "FlopsFn", + "InputMap", + "KernelSpec", + "OutputSpec", + "SizeMap", + "SpecValidationError", + "Tolerance", + "validate_spec", +] + +#: Mapping of input name -> tensor (or other value) handed to the candidate. +InputMap = Mapping[str, Any] + +#: Mapping of size key -> dimension, e.g. ``{"M": 1024, "N": 1024}``. +SizeMap = Mapping[str, int] + +FlopsFn = Callable[[SizeMap], "int | float"] +BytesFn = Callable[[SizeMap, int], "int | float"] + +#: Size labels every specification must define so the CLI can select a size +#: without operation-specific knowledge. +STANDARD_SIZE_LABELS: tuple[str, ...] = ("small", "medium", "large") + +_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +class SpecValidationError(ValueError): + """Raised when a kernel specification is malformed. + + The message always identifies the specification and the invalid field. + """ + + +def _fail(spec_name: object, field_name: str, message: str) -> "SpecValidationError": + label = spec_name if isinstance(spec_name, str) and spec_name else "" + return SpecValidationError(f"kernel spec {label!r}: field {field_name!r}: {message}") + + +@dataclass(frozen=True) +class Tolerance: + """Absolute and relative tolerance for one dtype. + + Both values must be finite and non-negative. A NaN or infinite tolerance + would make every comparison pass and silently disable the correctness + gate, so it is rejected at construction time. + """ + + atol: float + rtol: float + + def __post_init__(self) -> None: + for field_name in ("atol", "rtol"): + value = getattr(self, field_name) + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise SpecValidationError( + f"Tolerance.{field_name} must be a number, got {value!r}" + ) + if value != value: # NaN + raise SpecValidationError(f"Tolerance.{field_name} must not be NaN") + if math.isinf(value): + raise SpecValidationError( + f"Tolerance.{field_name} must be finite, got {value!r}" + ) + if value < 0: + raise SpecValidationError( + f"Tolerance.{field_name} must be non-negative, got {value!r}" + ) + object.__setattr__(self, "atol", float(self.atol)) + object.__setattr__(self, "rtol", float(self.rtol)) + + def as_dict(self) -> dict[str, float]: + """Return ``{"atol": ..., "rtol": ...}`` for legacy call sites.""" + return {"atol": self.atol, "rtol": self.rtol} + + +@dataclass(frozen=True) +class EdgeCase: + """One adversarial or awkward shape that must stay covered. + + ``dtype`` of ``None`` means "the specification's first dtype". + ``input_transform`` may post-process the generated inputs, e.g. to force a + degenerate value distribution. + """ + + name: str + size: SizeMap + dtype: str | None = None + seed: int = 42 + input_transform: Callable[[InputMap], InputMap] | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "size", dict(self.size)) + + +def _normalize_paths(value: Iterable[str], field: str) -> tuple[str, ...]: + """Normalize a sequence of output-tree paths, rejecting duplicates.""" + if isinstance(value, (str, bytes)) or not isinstance(value, Iterable): + raise SpecValidationError(f"{field} must be a sequence of path strings") + out: list[str] = [] + for path in value: + if not isinstance(path, str) or not path: + raise SpecValidationError( + f"{field} entries must be non-empty strings, got {path!r}" + ) + if path in out: + raise SpecValidationError(f"{field} contains duplicate path {path!r}") + out.append(path) + if not out: + raise SpecValidationError(f"{field} must contain at least one path") + return tuple(out) + + +@dataclass(frozen=True) +class OutputSpec: + """How a structured output tree participates in correctness checking. + + ``included_paths`` of ``None`` compares every leaf; otherwise only the + listed leaf paths (see :mod:`autokernel.verification.outputs` for the path + syntax) participate, and a configured path that does not exist is an + error. ``compare_non_tensors`` controls whether non-tensor (metadata) + leaves must match exactly. + """ + + included_paths: tuple[str, ...] | None = None + compare_non_tensors: bool = True + + def __post_init__(self) -> None: + if not isinstance(self.compare_non_tensors, bool): + raise SpecValidationError( + "OutputSpec.compare_non_tensors must be a bool, got " + f"{self.compare_non_tensors!r}" + ) + if self.included_paths is not None: + object.__setattr__( + self, + "included_paths", + _normalize_paths(self.included_paths, "OutputSpec.included_paths"), + ) + + +@dataclass(frozen=True) +class BackwardSpec: + """Opt-in gradient verification for an operation. + + ``differentiable_inputs`` names the generated inputs that must receive + gradients. ``output_paths`` selects the tensor output leaves that receive + upstream gradients (``None`` selects every floating tensor leaf). + ``tolerances`` overrides the forward tolerances for gradient comparison, + keyed by canonical dtype. ``enabled_by_default`` lets a specification + request the check without a CLI flag. + """ + + differentiable_inputs: tuple[str, ...] + output_paths: tuple[str, ...] | None = None + tolerances: Mapping[str, Any] | None = None + enabled_by_default: bool = False + + def __post_init__(self) -> None: + if not self.differentiable_inputs: + raise SpecValidationError( + "BackwardSpec.differentiable_inputs must name at least one input" + ) + inputs = _normalize_paths( + self.differentiable_inputs, "BackwardSpec.differentiable_inputs" + ) + object.__setattr__(self, "differentiable_inputs", inputs) + if self.output_paths is not None: + object.__setattr__( + self, + "output_paths", + _normalize_paths(self.output_paths, "BackwardSpec.output_paths"), + ) + if not isinstance(self.enabled_by_default, bool): + raise SpecValidationError( + "BackwardSpec.enabled_by_default must be a bool, got " + f"{self.enabled_by_default!r}" + ) + if self.tolerances is not None: + object.__setattr__( + self, + "tolerances", + _normalize_tolerances("backward_spec", self.tolerances), + ) + + +@dataclass(frozen=True) +class CompileSpec: + """Opt-in ``torch.compile`` verification settings. + + ``fullgraph`` (the default) forbids graph breaks. ``dynamic`` declares + dynamic-shape support, in which case the check runs at least two + compatible shapes through the same compiled callable. + """ + + enabled: bool = False + fullgraph: bool = True + dynamic: bool = False + + def __post_init__(self) -> None: + for field_name in ("enabled", "fullgraph", "dynamic"): + value = getattr(self, field_name) + if not isinstance(value, bool): + raise SpecValidationError( + f"CompileSpec.{field_name} must be a bool, got {value!r}" + ) + + +@dataclass(frozen=True, kw_only=True) +class KernelSpec: + """Everything the harness needs to benchmark and extract one operation. + + Args: + name: identifier-like operation name, unique within a registry. + reference_fn: PyTorch ground truth, called as ``reference_fn(**inputs)``. + input_generator: ``(size, dtype, device, seed) -> InputMap``, deterministic + for a fixed seed. + sizes: ordered label -> size mapping. Must contain ``small``, ``medium`` + and ``large``. Accepts a mapping or a sequence of ``(label, size)`` + pairs (which is also how duplicate labels are detected). + dtypes: canonical dtype names to sweep, in benchmark order. The first + entry is the primary dtype. + tolerances: dtype name -> :class:`Tolerance`. Must cover every declared + dtype; extra entries are allowed. + edge_cases: non-power-of-two and otherwise awkward cases. + flops_fn: ``(size) -> FLOPs``. Prefer an accounting + :class:`~autokernel.specs.accounting.Expression` so extraction can + serialize it. + bytes_fn: ``(size, dt_bytes) -> bytes moved``. + shape_keys: canonical size keys, in display order. Derived from ``sizes`` + when omitted. + shape_aliases: external key name -> canonical key, used to parse shape + strings coming from the profiler. + starter_kernels: backend name -> starter kernel file. + speedup_estimate: human-readable extraction hint, e.g. ``"2-3x"``. + default_shape: extraction fallback when a profiled shape cannot be + parsed. Defaults to the ``large`` size. + output_spec: optional structured-output policy. ``None`` compares every + output leaf and requires non-tensor leaves to match exactly, which + preserves the historical single-tensor behavior. + backward_spec: optional gradient-verification policy. ``None`` means + the operation is forward-only; ``--check-backward`` then fails with + an actionable unsupported message instead of silently skipping. + compile_spec: optional ``torch.compile`` verification settings. + ``None`` behaves like ``CompileSpec()``: the check only runs when + requested, with ``fullgraph=True`` and static shapes. + """ + + name: str + reference_fn: Callable[..., Any] + input_generator: Callable[..., InputMap] + sizes: Mapping[str, SizeMap] | Sequence[tuple[str, SizeMap]] + dtypes: Iterable[str] + tolerances: Mapping[str, Any] + flops_fn: FlopsFn + bytes_fn: BytesFn + edge_cases: Iterable[EdgeCase] = () + shape_keys: Iterable[str] = () + shape_aliases: Mapping[str, str] | Sequence[tuple[str, str]] = () + starter_kernels: Mapping[str, Any] | Sequence[tuple[str, Any]] = () + speedup_estimate: str | None = None + default_shape: SizeMap | None = None + output_spec: OutputSpec | Mapping[str, Any] | None = None + backward_spec: BackwardSpec | Mapping[str, Any] | None = None + compile_spec: CompileSpec | Mapping[str, Any] | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "sizes", _normalize_sizes(self.name, self.sizes)) + object.__setattr__(self, "dtypes", _normalize_dtypes(self.name, self.dtypes)) + object.__setattr__( + self, "tolerances", _normalize_tolerances(self.name, self.tolerances) + ) + object.__setattr__( + self, "edge_cases", _normalize_edge_cases(self.name, self.edge_cases) + ) + object.__setattr__( + self, "shape_keys", _normalize_shape_keys(self.name, self.shape_keys, self.sizes) + ) + object.__setattr__( + self, "shape_aliases", _normalize_aliases(self.name, self.shape_aliases) + ) + object.__setattr__( + self, "starter_kernels", _normalize_starters(self.name, self.starter_kernels) + ) + if self.default_shape is not None: + object.__setattr__( + self, + "default_shape", + _normalize_size_map(self.name, "default_shape", self.default_shape), + ) + object.__setattr__( + self, "output_spec", _coerce_output_spec(self.name, self.output_spec) + ) + object.__setattr__( + self, "backward_spec", _coerce_backward_spec(self.name, self.backward_spec) + ) + object.__setattr__( + self, "compile_spec", _coerce_compile_spec(self.name, self.compile_spec) + ) + # Structural validation happens eagerly. The small/medium/large + # requirement is a *registration* rule (see KernelRegistry.register) so + # tools can still build narrower specifications for inspection. + validate_spec( + self, + require_standard_sizes=False, + check_starter_files=False, + ) + + # -- convenience accessors ----------------------------------------- + @property + def primary_dtype(self) -> str: + """The first declared dtype, used for benchmarking and smoke tests.""" + return self.dtypes[0] + + def size_items(self) -> tuple[tuple[str, dict[str, int]], ...]: + """Return ``((label, size), ...)`` in declaration order.""" + return tuple((label, dict(size)) for label, size in self.sizes.items()) + + def tolerance_for(self, dtype: str) -> Tolerance: + """Return the tolerance declared for a canonical dtype name.""" + from .dtypes import canonical_dtype_name + + name = canonical_dtype_name(dtype) + try: + return self.tolerances[name] + except KeyError as exc: + raise SpecValidationError( + f"kernel spec {self.name!r}: no tolerance declared for dtype {name!r}" + ) from exc + + def extraction_shape(self) -> dict[str, int]: + """Shape used by extraction when a profiled shape cannot be parsed.""" + if self.default_shape is not None: + return dict(self.default_shape) + if "large" in self.sizes: + return dict(self.sizes["large"]) + last_label = next(reversed(list(self.sizes))) + return dict(self.sizes[last_label]) + + def starter_kernel(self, backend: str = "triton") -> Path | None: + """Return the starter kernel path for a backend, or None if undeclared.""" + path = self.starter_kernels.get(backend) + return Path(path) if path is not None else None + + +# --------------------------------------------------------------------------- +# Normalization helpers +# --------------------------------------------------------------------------- + +def _normalize_sizes( + name: object, sizes: Mapping[str, SizeMap] | Sequence[tuple[str, SizeMap]] +) -> dict[str, dict[str, int]]: + if isinstance(sizes, Mapping): + pairs: list[tuple[str, Any]] = list(sizes.items()) + elif isinstance(sizes, (str, bytes)) or not isinstance(sizes, Iterable): + raise _fail(name, "sizes", f"expected a mapping or pairs, got {type(sizes).__name__}") + else: + pairs = [] + for item in sizes: + if not isinstance(item, Sequence) or len(item) != 2: + raise _fail(name, "sizes", f"expected (label, size) pairs, got {item!r}") + pairs.append((item[0], item[1])) + + out: dict[str, dict[str, int]] = {} + for label, size in pairs: + if not isinstance(label, str) or not label: + raise _fail(name, "sizes", f"size label must be a non-empty string, got {label!r}") + if label in out: + raise _fail(name, "sizes", f"duplicate size label {label!r}") + out[label] = _normalize_size_map(name, f"sizes[{label!r}]", size) + return out + + +def _normalize_size_map( + name: object, field_name: str, size: object +) -> dict[str, int]: + """Normalize one shape mapping and reject unusable dimensions.""" + if not isinstance(size, Mapping) or not size: + raise _fail(name, field_name, "must be a non-empty mapping") + normalized: dict[str, int] = {} + for key, value in size.items(): + if not isinstance(key, str) or not key: + raise _fail(name, field_name, f"has a non-string key {key!r}") + if isinstance(value, bool) or not isinstance(value, int): + raise _fail( + name, field_name, f"key {key!r} must be an int, got {value!r}" + ) + if value <= 0: + raise _fail( + name, field_name, f"key {key!r} must be positive, got {value!r}" + ) + normalized[key] = value + return normalized + + +def _normalize_edge_cases( + name: object, edge_cases: Iterable[EdgeCase] +) -> tuple[EdgeCase, ...]: + if isinstance(edge_cases, (str, bytes)) or not isinstance(edge_cases, Iterable): + raise _fail(name, "edge_cases", "expected an iterable of EdgeCase values") + normalized: list[EdgeCase] = [] + for index, edge in enumerate(edge_cases): + if not isinstance(edge, EdgeCase): + raise _fail( + name, "edge_cases", f"expected EdgeCase, got {type(edge).__name__}" + ) + normalized.append( + EdgeCase( + name=edge.name, + size=_normalize_size_map( + name, f"edge_cases[{index}].size", edge.size + ), + dtype=edge.dtype, + seed=edge.seed, + input_transform=edge.input_transform, + ) + ) + return tuple(normalized) + + +def _normalize_dtypes(name: object, dtypes: Iterable[str]) -> tuple[str, ...]: + if isinstance(dtypes, (str, bytes)) or not isinstance(dtypes, Iterable): + raise _fail(name, "dtypes", f"expected an iterable of dtype names, got {dtypes!r}") + out: list[str] = [] + for dtype in dtypes: + if not is_canonical_dtype(dtype): + raise _fail( + name, + "dtypes", + f"unknown dtype {dtype!r}; expected one of {', '.join(CANONICAL_DTYPES)}", + ) + if dtype in out: + raise _fail(name, "dtypes", f"duplicate dtype {dtype!r}") + out.append(str(dtype)) + return tuple(out) + + +def _normalize_tolerances(name: object, tolerances: Mapping[str, Any]) -> dict[str, Tolerance]: + if not isinstance(tolerances, Mapping): + raise _fail( + name, "tolerances", f"expected a mapping, got {type(tolerances).__name__}" + ) + out: dict[str, Tolerance] = {} + for dtype, tol in tolerances.items(): + if not is_canonical_dtype(dtype): + raise _fail( + name, + "tolerances", + f"unknown dtype key {dtype!r}; expected one of {', '.join(CANONICAL_DTYPES)}", + ) + if isinstance(tol, Tolerance): + out[str(dtype)] = tol + elif isinstance(tol, Mapping): + missing = {"atol", "rtol"} - set(tol) + if missing: + raise _fail( + name, + "tolerances", + f"dtype {dtype!r} is missing {sorted(missing)}", + ) + unexpected = set(tol) - {"atol", "rtol"} + if unexpected: + raise _fail( + name, + "tolerances", + f"dtype {dtype!r} has unexpected keys {sorted(unexpected)}", + ) + try: + out[str(dtype)] = Tolerance(atol=tol["atol"], rtol=tol["rtol"]) + except SpecValidationError as exc: + raise _fail(name, "tolerances", f"dtype {dtype!r}: {exc}") from exc + else: + raise _fail( + name, + "tolerances", + f"dtype {dtype!r} must map to a Tolerance, got {type(tol).__name__}", + ) + return out + + +def _coerce_output_spec( + name: object, value: OutputSpec | Mapping[str, Any] | None +) -> OutputSpec | None: + if value is None or isinstance(value, OutputSpec): + return value + if isinstance(value, Mapping): + try: + return OutputSpec(**value) + except TypeError as exc: + raise _fail(name, "output_spec", f"invalid OutputSpec mapping: {exc}") from exc + raise _fail( + name, + "output_spec", + f"expected an OutputSpec, mapping or None, got {type(value).__name__}", + ) + + +def _coerce_backward_spec( + name: object, value: BackwardSpec | Mapping[str, Any] | None +) -> BackwardSpec | None: + if value is None or isinstance(value, BackwardSpec): + return value + if isinstance(value, Mapping): + try: + return BackwardSpec(**value) + except TypeError as exc: + raise _fail(name, "backward_spec", f"invalid BackwardSpec mapping: {exc}") from exc + raise _fail( + name, + "backward_spec", + f"expected a BackwardSpec, mapping or None, got {type(value).__name__}", + ) + + +def _coerce_compile_spec( + name: object, value: CompileSpec | Mapping[str, Any] | None +) -> CompileSpec | None: + if value is None or isinstance(value, CompileSpec): + return value + if isinstance(value, Mapping): + try: + return CompileSpec(**value) + except TypeError as exc: + raise _fail(name, "compile_spec", f"invalid CompileSpec mapping: {exc}") from exc + raise _fail( + name, + "compile_spec", + f"expected a CompileSpec, mapping or None, got {type(value).__name__}", + ) + + +def _normalize_shape_keys( + name: object, shape_keys: Iterable[str], sizes: Mapping[str, SizeMap] +) -> tuple[str, ...]: + keys = tuple(shape_keys) + if not keys: + # Derive from the first declared size, preserving its key order. + for size in sizes.values(): + return tuple(size) + return () + out: list[str] = [] + for key in keys: + if not isinstance(key, str) or not key: + raise _fail(name, "shape_keys", f"shape key must be a non-empty string, got {key!r}") + if key in out: + raise _fail(name, "shape_keys", f"duplicate shape key {key!r}") + out.append(key) + return tuple(out) + + +def _normalize_aliases( + name: object, aliases: Mapping[str, str] | Sequence[tuple[str, str]] +) -> dict[str, str]: + if isinstance(aliases, Mapping): + pairs: list[tuple[Any, Any]] = list(aliases.items()) + elif isinstance(aliases, (str, bytes)) or not isinstance(aliases, Iterable): + raise _fail( + name, "shape_aliases", f"expected a mapping or pairs, got {type(aliases).__name__}" + ) + else: + pairs = [] + for item in aliases: + if not isinstance(item, Sequence) or len(item) != 2: + raise _fail(name, "shape_aliases", f"expected (alias, key) pairs, got {item!r}") + pairs.append((item[0], item[1])) + + out: dict[str, str] = {} + for alias, canonical in pairs: + if not isinstance(alias, str) or not alias: + raise _fail(name, "shape_aliases", f"alias must be a non-empty string, got {alias!r}") + if not isinstance(canonical, str) or not canonical: + raise _fail( + name, + "shape_aliases", + f"alias {alias!r} must resolve to a non-empty string, got {canonical!r}", + ) + if alias in out and out[alias] != canonical: + raise _fail( + name, + "shape_aliases", + f"alias {alias!r} resolves inconsistently to {out[alias]!r} and {canonical!r}", + ) + out[alias] = canonical + return out + + +def _normalize_starters( + name: object, starters: Mapping[str, Any] | Sequence[tuple[str, Any]] +) -> dict[str, Path]: + if isinstance(starters, Mapping): + pairs: list[tuple[Any, Any]] = list(starters.items()) + elif isinstance(starters, (str, bytes)) or not isinstance(starters, Iterable): + raise _fail( + name, + "starter_kernels", + f"expected a mapping or pairs, got {type(starters).__name__}", + ) + else: + pairs = [] + for item in starters: + if not isinstance(item, Sequence) or len(item) != 2: + raise _fail( + name, "starter_kernels", f"expected (backend, path) pairs, got {item!r}" + ) + pairs.append((item[0], item[1])) + + out: dict[str, Path] = {} + for backend, path in pairs: + if not isinstance(backend, str) or not backend: + raise _fail( + name, "starter_kernels", f"backend must be a non-empty string, got {backend!r}" + ) + if backend in out: + raise _fail(name, "starter_kernels", f"duplicate backend {backend!r}") + if not isinstance(path, (str, Path)): + raise _fail( + name, + "starter_kernels", + f"backend {backend!r} must map to a path, got {type(path).__name__}", + ) + out[backend] = Path(path) + return out + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + +def validate_spec( + spec: KernelSpec, + *, + require_standard_sizes: bool = True, + check_starter_files: bool = True, +) -> None: + """Validate a specification, raising :class:`SpecValidationError` on failure. + + Args: + spec: the specification to check. + require_standard_sizes: require ``small``, ``medium`` and ``large`` size + labels so the CLI can pick a size for any operation. + check_starter_files: require every declared starter kernel file to exist. + """ + if not isinstance(spec, KernelSpec): + raise SpecValidationError( + f"expected a KernelSpec, got {type(spec).__name__}" + ) + + name = spec.name + if not isinstance(name, str) or not name: + raise _fail(name, "name", "operation name must be a non-empty string") + if not _NAME_RE.match(name): + raise _fail( + name, + "name", + "operation name must be identifier-like (letters, digits, underscore; " + "not starting with a digit)", + ) + + if not callable(spec.reference_fn): + raise _fail(name, "reference_fn", "reference function must be callable") + if not callable(spec.input_generator): + raise _fail(name, "input_generator", "input generator must be callable") + if not callable(spec.flops_fn): + raise _fail(name, "flops_fn", "FLOP accounting must be callable") + if not callable(spec.bytes_fn): + raise _fail(name, "bytes_fn", "byte accounting must be callable") + + if not spec.sizes: + raise _fail(name, "sizes", "at least one size must be declared") + if require_standard_sizes: + missing = [label for label in STANDARD_SIZE_LABELS if label not in spec.sizes] + if missing: + raise _fail( + name, + "sizes", + f"missing required size label(s) {missing}; declared {sorted(spec.sizes)}", + ) + + if not spec.dtypes: + raise _fail(name, "dtypes", "at least one dtype must be declared") + + missing_tolerances = [d for d in spec.dtypes if d not in spec.tolerances] + if missing_tolerances: + raise _fail( + name, + "tolerances", + f"missing tolerance for declared dtype(s) {missing_tolerances}", + ) + + shape_keys = set(spec.shape_keys) + if not shape_keys: + raise _fail(name, "shape_keys", "at least one shape key must be declared") + for label, size in spec.sizes.items(): + extra = sorted(set(size) - shape_keys) + missing = sorted(shape_keys - set(size)) + if extra or missing: + raise _fail( + name, + "sizes", + f"size {label!r} keys {sorted(size)} do not match shape_keys " + f"{sorted(shape_keys)} (unexpected={extra}, missing={missing})", + ) + + seen_edges: set[str] = set() + for edge in spec.edge_cases: + if not isinstance(edge, EdgeCase): + raise _fail(name, "edge_cases", f"expected EdgeCase, got {type(edge).__name__}") + if not edge.name: + raise _fail(name, "edge_cases", "edge case name must be non-empty") + if edge.name in seen_edges: + raise _fail(name, "edge_cases", f"duplicate edge case name {edge.name!r}") + seen_edges.add(edge.name) + extra = sorted(set(edge.size) - shape_keys) + missing = sorted(shape_keys - set(edge.size)) + if extra or missing: + raise _fail( + name, + "edge_cases", + f"edge case {edge.name!r} keys {sorted(edge.size)} do not match shape_keys " + f"{sorted(shape_keys)} (unexpected={extra}, missing={missing})", + ) + if edge.dtype is not None: + if not is_canonical_dtype(edge.dtype): + raise _fail( + name, "edge_cases", f"edge case {edge.name!r} has unknown dtype {edge.dtype!r}" + ) + if edge.dtype not in spec.dtypes: + raise _fail( + name, + "edge_cases", + f"edge case {edge.name!r} dtype {edge.dtype!r} is not declared in dtypes " + f"{list(spec.dtypes)}", + ) + if isinstance(edge.seed, bool) or not isinstance(edge.seed, int): + raise _fail( + name, "edge_cases", f"edge case {edge.name!r} seed must be an int, got {edge.seed!r}" + ) + if edge.input_transform is not None and not callable(edge.input_transform): + raise _fail( + name, "edge_cases", f"edge case {edge.name!r} input_transform must be callable" + ) + + for alias, canonical in spec.shape_aliases.items(): + if canonical not in shape_keys: + raise _fail( + name, + "shape_aliases", + f"alias {alias!r} resolves to {canonical!r}, which is not a shape key " + f"{sorted(shape_keys)}", + ) + + if isinstance(spec.flops_fn, Expression): + unknown = sorted(spec.flops_fn.size_keys() - shape_keys) + if unknown: + raise _fail( + name, "flops_fn", f"references unknown size key(s) {unknown}" + ) + if spec.flops_fn.uses_dtype_bytes(): + raise _fail(name, "flops_fn", "FLOP accounting must not depend on dtype bytes") + if isinstance(spec.bytes_fn, Expression): + unknown = sorted(spec.bytes_fn.size_keys() - shape_keys) + if unknown: + raise _fail(name, "bytes_fn", f"references unknown size key(s) {unknown}") + + if spec.default_shape is not None: + extra = sorted(set(spec.default_shape) - shape_keys) + missing = sorted(shape_keys - set(spec.default_shape)) + if extra or missing: + raise _fail( + name, + "default_shape", + f"keys {sorted(spec.default_shape)} do not match shape_keys " + f"{sorted(shape_keys)} (unexpected={extra}, missing={missing})", + ) + + if check_starter_files: + for backend, path in spec.starter_kernels.items(): + if not Path(path).is_file(): + raise _fail( + name, + "starter_kernels", + f"backend {backend!r} starter kernel not found: {path}", + ) + + if spec.speedup_estimate is not None and not isinstance(spec.speedup_estimate, str): + raise _fail( + name, + "speedup_estimate", + f"expected a string or None, got {type(spec.speedup_estimate).__name__}", + ) + + for field_name, expected_type in ( + ("output_spec", OutputSpec), + ("backward_spec", BackwardSpec), + ("compile_spec", CompileSpec), + ): + value = getattr(spec, field_name) + if value is not None and not isinstance(value, expected_type): + raise _fail( + name, + field_name, + f"expected a {expected_type.__name__} or None, " + f"got {type(value).__name__}", + ) diff --git a/autokernel/verification/__init__.py b/autokernel/verification/__init__.py new file mode 100644 index 00000000..cf96f081 --- /dev/null +++ b/autokernel/verification/__init__.py @@ -0,0 +1,79 @@ +"""Correctness verification for structured kernel outputs. + +This package generalizes the benchmark harness beyond single-tensor outputs: + +* :mod:`autokernel.verification.outputs` flattens and compares arbitrary + output trees (tensors, tuples, lists, dictionaries, named tuples and + nested combinations) leaf by leaf, with stable diagnostic paths. +* :mod:`autokernel.verification.backward` compares declared input gradients. +* :mod:`autokernel.verification.compile` enforces optional full-graph checks. +* :mod:`autokernel.verification.corpus` validates production shape corpora. +* :mod:`autokernel.verification.results` writes versioned result artifacts. + +Modules here never initialize a GPU at import time; ``torch`` is imported +lazily inside the functions that need it. +""" + +from __future__ import annotations + +from .backward import ( + BackwardReport, + GradientRecord, + check_backward, +) +from .compile import CompileCaseRecord, CompileReport, check_compile +from .corpus import ( + CORPUS_SCHEMA_VERSION, + CorpusCase, + CorpusError, + ShapeCorpus, + load_shape_corpus, + validate_corpus_against_spec, + weighted_aggregate, +) +from .outputs import ( + DEFAULT_TOLERANCE, + LeafRecord, + OutputTreeError, + TreeComparison, + compare_deterministic, + compare_output_trees, + compare_tensor_leaf, + flatten_output_tree, + tree_has_nan_or_inf, +) +from .results import ( + RESULT_SCHEMA_VERSION, + collect_environment_metadata, + result_envelope, + write_result_atomic, +) + +__all__ = [ + "CORPUS_SCHEMA_VERSION", + "DEFAULT_TOLERANCE", + "RESULT_SCHEMA_VERSION", + "BackwardReport", + "CompileCaseRecord", + "CompileReport", + "CorpusCase", + "CorpusError", + "GradientRecord", + "LeafRecord", + "OutputTreeError", + "ShapeCorpus", + "TreeComparison", + "check_backward", + "check_compile", + "collect_environment_metadata", + "compare_deterministic", + "compare_output_trees", + "compare_tensor_leaf", + "flatten_output_tree", + "load_shape_corpus", + "result_envelope", + "tree_has_nan_or_inf", + "validate_corpus_against_spec", + "weighted_aggregate", + "write_result_atomic", +] diff --git a/autokernel/verification/backward.py b/autokernel/verification/backward.py new file mode 100644 index 00000000..37733ca0 --- /dev/null +++ b/autokernel/verification/backward.py @@ -0,0 +1,425 @@ +"""Optional gradient verification for kernel candidates. + +Backward verification is opt-in: it runs only when requested with +``--check-backward`` (or when the spec's ``BackwardSpec.enabled_by_default`` +is set). A specification without a :class:`~autokernel.specs.BackwardSpec` +is forward-only; requesting the check then *fails* with an actionable +unsupported message instead of silently skipping. + +Protocol (per the Week 2 plan): + +1. generate one canonical input mapping (``small`` size, primary dtype); +2. deep-clone tensor inputs for the reference and candidate paths so + neither side shares autograd state; +3. set ``requires_grad=True`` only for inputs declared in + ``BackwardSpec.differentiable_inputs``; +4. run reference and candidate independently; +5. select the tensor output leaves declared by ``output_paths``, or every + floating tensor leaf when omitted; +6. draw deterministic upstream gradients with a fixed-seed generator -- + never only ``output.sum()``, whose symmetry can hide errors; +7. call ``torch.autograd.grad`` with matching inputs and upstreams; +8. compare every requested input gradient by name; +9. report missing gradients, unexpected gradients, shape differences, NaN, + infinity, maximum error and mean error per input; +10. never accumulate gradients: ``autograd.grad`` returns fresh tensors and + the generated inputs are never mutated, so repeated checks cannot + interfere with each other. + +Backward execution is correctness-only; no performance claims are made. +""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import Any, Callable, Mapping + +from ..specs.dtypes import canonical_dtype_name +from ..specs.types import BackwardSpec, KernelSpec, Tolerance +from .outputs import ( + DEFAULT_TOLERANCE, + _is_tensor, + compare_tensor_leaf, + flatten_output_tree, +) + +__all__ = [ + "BackwardReport", + "GradientRecord", + "check_backward", +] + +#: Seed for the upstream-gradient generator. Fixed so every run of the same +#: case compares against identical upstream gradients. +UPSTREAM_SEED = 0x5EED + + +@dataclass(frozen=True) +class GradientRecord: + """Gradient comparison outcome for one declared differentiable input.""" + + input_name: str + status: str # "match" | "mismatch" | "missing" | "unexpected" + reason: str = "" + max_abs_error: float | None = None + mean_abs_error: float | None = None + has_nan: bool = False + has_inf: bool = False + + def as_dict(self) -> dict[str, Any]: + return { + "input_name": self.input_name, + "status": self.status, + "reason": self.reason, + "max_abs_error": self.max_abs_error, + "mean_abs_error": self.mean_abs_error, + "has_nan": self.has_nan, + "has_inf": self.has_inf, + } + + +@dataclass(frozen=True) +class BackwardReport: + """Aggregated backward-verification outcome.""" + + status: str # "PASS" | "FAIL" + reason: str + gradients: tuple[GradientRecord, ...] = () + output_paths: tuple[str, ...] = () + + def as_dict(self) -> dict[str, Any]: + return { + "status": self.status, + "reason": self.reason, + "output_paths": list(self.output_paths), + "gradients": [record.as_dict() for record in self.gradients], + } + + +def _fail_report(reason: str) -> BackwardReport: + return BackwardReport(status="FAIL", reason=reason) + + +def _validate_declared_inputs( + spec: KernelSpec, backward: BackwardSpec, inputs: Mapping[str, Any] +) -> str | None: + """Return an actionable failure reason, or None when inputs are usable.""" + for name in backward.differentiable_inputs: + if name not in inputs: + return ( + f"backward_spec declares differentiable input {name!r} but the " + f"input generator did not produce it (got {sorted(inputs)})" + ) + value = inputs[name] + if not _is_tensor(value): + return ( + f"backward_spec declares differentiable input {name!r} but the " + f"generated value is not a tensor (got {type(value).__name__})" + ) + if not value.is_floating_point(): + return ( + f"backward_spec declares differentiable input {name!r} but the " + f"generated tensor is not floating-point (dtype {value.dtype})" + ) + return None + + +def _clone_inputs(inputs: Mapping[str, Any], differentiable: tuple[str, ...]) -> dict: + """Deep-clone inputs, enabling grad only on declared differentiable names.""" + cloned: dict[str, Any] = {} + for name, value in inputs.items(): + if _is_tensor(value): + tensor = value.detach().clone() + if name in differentiable: + tensor.requires_grad_(True) + cloned[name] = tensor + else: + cloned[name] = copy.deepcopy(value) + return cloned + + +def _select_output_leaves( + reference_out: Any, + candidate_out: Any, + spec: KernelSpec, + backward: BackwardSpec, +) -> tuple[list[tuple[str, Any]], list[tuple[str, Any]], str | None]: + """Pick the floating tensor leaves that receive upstream gradients. + + Returns ``(reference_leaves, candidate_leaves, failure_reason)``. + """ + ref_leaves = flatten_output_tree(reference_out) + cand_leaves = flatten_output_tree(candidate_out) + cand_by_path = dict(cand_leaves) + + if backward.output_paths is not None: + ref_by_path = dict(ref_leaves) + selected_ref: list[tuple[str, Any]] = [] + selected_cand: list[tuple[str, Any]] = [] + for path in backward.output_paths: + if path not in ref_by_path: + return [], [], ( + f"backward_spec.output_paths entry {path!r} does not exist in " + f"the reference output (available: {sorted(ref_by_path)})" + ) + if path not in cand_by_path: + return [], [], ( + f"backward_spec.output_paths entry {path!r} does not exist in " + f"the candidate output" + ) + ref_leaf = ref_by_path[path] + if not (_is_tensor(ref_leaf) and ref_leaf.is_floating_point()): + return [], [], ( + f"backward_spec.output_paths entry {path!r} is not a floating " + f"tensor leaf and cannot receive an upstream gradient" + ) + selected_ref.append((path, ref_leaf)) + selected_cand.append((path, cand_by_path[path])) + return selected_ref, selected_cand, None + + included = ( + set(spec.output_spec.included_paths) + if spec.output_spec is not None and spec.output_spec.included_paths is not None + else None + ) + selected_ref = [ + (path, leaf) + for path, leaf in ref_leaves + if (included is None or path in included) + and _is_tensor(leaf) + and leaf.is_floating_point() + ] + for path, _ in selected_ref: + if path not in cand_by_path: + return [], [], ( + f"candidate output is missing tensor leaf {path!r} needed " + f"for backward comparison" + ) + selected_cand = [(path, cand_by_path[path]) for path, _ in selected_ref] + if not selected_ref: + return [], [], ( + "no floating tensor output leaves available for backward comparison" + ) + return selected_ref, selected_cand, None + + + +def _upstream_gradients( + leaves: list[tuple[str, Any]], device: str, seed: int +) -> list[Any]: + """Deterministic per-leaf upstream gradients from a fixed-seed generator. + + ``randn`` breaks the symmetry that ``output.sum()``-style upstreams have, + so sign and cancellation errors cannot hide. + """ + import torch + + try: + generator = torch.Generator(device=device) + generator_is_device_local = True + except Exception: + generator = torch.Generator() + generator_is_device_local = False + generator.manual_seed(seed) + upstreams = [] + for _, leaf in leaves: + if generator_is_device_local: + upstream = torch.randn( + tuple(leaf.shape), + dtype=leaf.dtype, + device=leaf.device, + generator=generator, + ) + else: + upstream = torch.randn( + tuple(leaf.shape), + dtype=leaf.dtype, + device="cpu", + generator=generator, + ).to(leaf.device) + upstreams.append(upstream) + return upstreams + + +def _gradient_tolerance( + spec: KernelSpec, backward: BackwardSpec, gradient: Any +) -> Tolerance: + tolerances = backward.tolerances if backward.tolerances is not None else spec.tolerances + try: + name = canonical_dtype_name(gradient.dtype) + except ValueError: + return DEFAULT_TOLERANCE + return tolerances.get(name, DEFAULT_TOLERANCE) + + + +def check_backward( + kernel_fn: Callable[..., Any], + spec: KernelSpec, + *, + device: str, + seed: int = UPSTREAM_SEED, +) -> BackwardReport: + """Verify candidate gradients against the reference for one canonical case. + + Never raises for an unsupported or failing candidate: every outcome is a + structured :class:`BackwardReport`. A spec without ``backward_spec`` is + forward-only and yields a FAIL report with an actionable unsupported + message. + """ + import torch + + backward = spec.backward_spec + if backward is None: + return _fail_report( + f"unsupported: kernel spec {spec.name!r} declares no backward_spec; " + f"the operation is forward-only. Add a BackwardSpec to the spec or " + f"drop --check-backward." + ) + + # 1. One canonical input mapping: small size, primary dtype. + if not spec.sizes: + return _fail_report( + f"kernel spec {spec.name!r} declares no sizes; backward " + f"verification needs at least one size" + ) + size_label = "small" if "small" in spec.sizes else next(iter(spec.sizes)) + size = dict(spec.sizes[size_label]) + dtype_name = spec.primary_dtype + base_inputs = spec.input_generator(size, dtype_name, device, seed=42) + + # Declared differentiable inputs must exist and be floating tensors. + invalid = _validate_declared_inputs(spec, backward, base_inputs) + if invalid is not None: + return _fail_report(invalid) + + # 2-3. Independent deep clones; grad only on declared names. + ref_inputs = _clone_inputs(base_inputs, backward.differentiable_inputs) + cand_inputs = _clone_inputs(base_inputs, backward.differentiable_inputs) + + # 4. Independent forward passes. + try: + reference_out = spec.reference_fn(**ref_inputs) + except Exception as exc: + return _fail_report( + f"reference forward failed: {type(exc).__name__}: {exc}" + ) + try: + candidate_out = kernel_fn(**cand_inputs) + except Exception as exc: + return _fail_report( + f"candidate forward failed: {type(exc).__name__}: {exc}" + ) + + # 5. Select the output leaves that receive upstream gradients. + ref_leaves, cand_leaves, failure = _select_output_leaves( + reference_out, candidate_out, spec, backward + ) + if failure is not None: + return _fail_report(failure) + + # 6. Deterministic upstream gradients (identical for both paths). + try: + upstreams = _upstream_gradients(ref_leaves, device, seed) + except Exception as exc: + return _fail_report( + f"upstream gradient generation failed: {type(exc).__name__}: {exc}" + ) + + # 7. Independent autograd.grad calls; allow_unused surfaces missing grads. + grad_inputs_ref = [ref_inputs[name] for name in backward.differentiable_inputs] + grad_inputs_cand = [cand_inputs[name] for name in backward.differentiable_inputs] + try: + ref_grads = torch.autograd.grad( + [leaf for _, leaf in ref_leaves], + grad_inputs_ref, + upstreams, + allow_unused=True, + ) + except Exception as exc: + return _fail_report( + f"reference backward failed: {type(exc).__name__}: {exc}" + ) + try: + cand_grads = torch.autograd.grad( + [leaf for _, leaf in cand_leaves], + grad_inputs_cand, + upstreams, + allow_unused=True, + ) + except Exception as exc: + return _fail_report( + f"candidate backward failed (is the candidate differentiable?): " + f"{type(exc).__name__}: {exc}" + ) + + + # 8-9. Compare every requested gradient by name. + records: list[GradientRecord] = [] + for name, ref_grad, cand_grad in zip( + backward.differentiable_inputs, ref_grads, cand_grads + ): + if ref_grad is None and cand_grad is None: + records.append( + GradientRecord( + input_name=name, + status="missing", + reason="no gradient in reference or candidate; the declared " + "differentiable input does not influence the selected outputs", + ) + ) + continue + if cand_grad is None: + records.append( + GradientRecord( + input_name=name, + status="missing", + reason="candidate produced no gradient but the reference did", + ) + ) + continue + if ref_grad is None: + records.append( + GradientRecord( + input_name=name, + status="unexpected", + reason="candidate produced a gradient where the reference " + "has none", + ) + ) + continue + leaf = compare_tensor_leaf( + f'grad["{name}"]', + cand_grad, + ref_grad, + _gradient_tolerance(spec, backward, ref_grad), + ) + records.append( + GradientRecord( + input_name=name, + status="match" if leaf.match else "mismatch", + reason=leaf.reason, + max_abs_error=leaf.max_abs_error, + mean_abs_error=leaf.mean_abs_error, + has_nan=leaf.has_nan, + has_inf=leaf.has_inf, + ) + ) + + failures = [record for record in records if record.status != "match"] + output_paths = tuple(path for path, _ in ref_leaves) + if failures: + first = failures[0] + return BackwardReport( + status="FAIL", + reason=f"grad[{first.input_name!r}] {first.status}: {first.reason}", + gradients=tuple(records), + output_paths=output_paths, + ) + return BackwardReport( + status="PASS", + reason="", + gradients=tuple(records), + output_paths=output_paths, + ) diff --git a/autokernel/verification/compile.py b/autokernel/verification/compile.py new file mode 100644 index 00000000..5aad77f9 --- /dev/null +++ b/autokernel/verification/compile.py @@ -0,0 +1,194 @@ +"""Optional full-graph ``torch.compile`` verification for kernel candidates.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable + +from ..specs.types import CompileSpec, KernelSpec +from .outputs import compare_output_trees + +__all__ = ["CompileCaseRecord", "CompileReport", "check_compile"] + + +@dataclass(frozen=True) +class CompileCaseRecord: + """Comparison result for one shape passed through the compiled callable.""" + + label: str + size: dict[str, int] + status: str + reason: str = "" + + def as_dict(self) -> dict[str, Any]: + return { + "label": self.label, + "size": dict(self.size), + "status": self.status, + "reason": self.reason, + } + + +@dataclass(frozen=True) +class CompileReport: + """Structured compile-verification outcome.""" + + status: str # "PASS" | "FAIL" | "UNSUPPORTED" + reason: str + cases: tuple[CompileCaseRecord, ...] + environment: dict[str, Any] + + def as_dict(self) -> dict[str, Any]: + return { + "status": self.status, + "reason": self.reason, + "cases": [case.as_dict() for case in self.cases], + "environment": dict(self.environment), + } + + +def _environment(torch: Any, device: str, settings: CompileSpec) -> dict[str, Any]: + try: + import triton + + triton_version = getattr(triton, "__version__", "unknown") + except Exception: + triton_version = None + + gpu_name = None + if device.startswith("cuda"): + try: + gpu_name = torch.cuda.get_device_name(device) + except Exception: + gpu_name = None + return { + "torch_version": getattr(torch, "__version__", "unknown"), + "triton_version": triton_version, + "cuda_version": getattr(getattr(torch, "version", None), "cuda", None), + "device": device, + "gpu_name": gpu_name, + "fullgraph": settings.fullgraph, + "dynamic": settings.dynamic, + } + + +def _selected_cases(spec: KernelSpec, dynamic: bool) -> list[tuple[str, dict[str, int]]]: + items = list(spec.size_items()) + if not items: + return [] + if not dynamic: + label = "small" if "small" in spec.sizes else items[0][0] + return [(label, dict(spec.sizes[label]))] + + # KernelSpec validation guarantees compatible keys. Prefer small/medium + # because compile verification is a correctness gate, not a stress test. + selected: list[tuple[str, dict[str, int]]] = [] + for label in ("small", "medium"): + if label in spec.sizes: + selected.append((label, dict(spec.sizes[label]))) + for label, size in items: + if len(selected) >= 2: + break + if label not in {existing for existing, _ in selected}: + selected.append((label, dict(size))) + return selected[:2] + + +def check_compile( + kernel_fn: Callable[..., Any], + spec: KernelSpec, + *, + device: str, +) -> CompileReport: + """Compile a candidate outside timed regions and compare it to eager. + + The same compiled callable is invoked at least twice per selected shape. + Dynamic specifications exercise two compatible shapes through that single + callable. All failures are returned as structured results. + """ + import torch + + settings = spec.compile_spec or CompileSpec() + environment = _environment(torch, device, settings) + compile_fn = getattr(torch, "compile", None) + if not callable(compile_fn): + return CompileReport( + status="UNSUPPORTED", + reason="torch.compile is unavailable in this PyTorch installation", + cases=(), + environment=environment, + ) + + cases = _selected_cases(spec, settings.dynamic) + if settings.dynamic and len(cases) < 2: + return CompileReport( + status="FAIL", + reason="dynamic compile verification requires at least two declared sizes", + cases=(), + environment=environment, + ) + + def candidate(**inputs: Any) -> Any: + return kernel_fn(**inputs) + + try: + compiled = compile_fn( + candidate, + fullgraph=settings.fullgraph, + dynamic=settings.dynamic, + ) + except Exception as exc: + return CompileReport( + status="FAIL", + reason=f"compiler setup failed: {type(exc).__name__}: {exc}", + cases=(), + environment=environment, + ) + + records: list[CompileCaseRecord] = [] + for label, size in cases: + try: + inputs = spec.input_generator( + size, spec.primary_dtype, device, seed=42 + ) + expected = spec.reference_fn(**inputs) + # torch.compile is lazy. The first call may compile; the second + # proves the compiled callable remains correct on a repeated run. + compiled(**inputs) + actual = compiled(**inputs) + comparison = compare_output_trees( + actual, + expected, + spec.tolerances, + output_spec=spec.output_spec, + ) + status = "PASS" if comparison.match else "FAIL" + records.append( + CompileCaseRecord( + label=label, + size=size, + status=status, + reason="" if comparison.match else comparison.reason, + ) + ) + except Exception as exc: + records.append( + CompileCaseRecord( + label=label, + size=size, + status="FAIL", + reason=f"{type(exc).__name__}: {exc}", + ) + ) + + passed = bool(records) and all(record.status == "PASS" for record in records) + reason = "" if passed else next( + (record.reason for record in records if record.status != "PASS"), + "no compile cases were selected", + ) + return CompileReport( + status="PASS" if passed else "FAIL", + reason=reason, + cases=tuple(records), + environment=environment, + ) diff --git a/autokernel/verification/corpus.py b/autokernel/verification/corpus.py new file mode 100644 index 00000000..e043ef73 --- /dev/null +++ b/autokernel/verification/corpus.py @@ -0,0 +1,281 @@ +"""Production shape corpora: versioned JSON benchmark cases. + +A shape corpus lets production workloads drive benchmarking without editing +Python source. The corpus contains *metadata only* -- shapes, dtypes, weights +and tags. It must never serialize model activations, weights, prompts or any +user data; the schema below has no field that could carry tensor data, and +unknown fields are rejected. + +Schema version 1:: + + { + "schema_version": 1, + "operation": "custom_affine", + "cases": [ + { + "name": "prod-prefill", + "size": {"rows": 4096, "cols": 1024}, + "dtype": "float16", // optional; default: spec primary dtype + "weight": 37, // optional; default: 1 + "tags": ["production"] // optional + } + ] + } + +Validation happens before any GPU allocation: :func:`load_shape_corpus` +parses and structurally validates the file, and +:func:`validate_corpus_against_spec` checks it against the selected +specification (operation name, shape keys, declared dtypes, resolved +duplicates). All failures raise :class:`CorpusError` with the file path and +the offending case. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +from ..specs.dtypes import is_canonical_dtype +from ..specs.types import KernelSpec + +__all__ = [ + "CORPUS_SCHEMA_VERSION", + "CorpusCase", + "CorpusError", + "ShapeCorpus", + "load_shape_corpus", + "validate_corpus_against_spec", + "weighted_aggregate", +] + +#: Only schema version accepted by this implementation. +CORPUS_SCHEMA_VERSION = 1 + +_TOP_LEVEL_KEYS = {"schema_version", "operation", "cases"} +_CASE_KEYS = {"name", "size", "dtype", "weight", "tags"} + + +class CorpusError(ValueError): + """Raised when a shape corpus is malformed or incompatible with a spec.""" + + +@dataclass(frozen=True) +class CorpusCase: + """One production benchmark case (metadata only).""" + + name: str + size: dict[str, int] + dtype: str | None = None # canonical dtype name; None -> spec primary dtype + weight: int = 1 + tags: tuple[str, ...] = () + + +@dataclass(frozen=True) +class ShapeCorpus: + """A validated corpus file.""" + + operation: str + cases: tuple[CorpusCase, ...] + source: str # path the corpus was loaded from, for diagnostics + schema_version: int = CORPUS_SCHEMA_VERSION + + +def _fail(source: object, message: str) -> CorpusError: + label = source if source else "" + return CorpusError(f"shape corpus {label!r}: {message}") + + +def _parse_case(raw: Any, index: int, source: object) -> CorpusCase: + where = f"case #{index}" + if not isinstance(raw, Mapping): + raise _fail(source, f"{where} must be an object, got {type(raw).__name__}") + unknown = sorted(set(raw) - _CASE_KEYS) + if unknown: + raise _fail(source, f"{where} has unknown field(s) {unknown}; " + f"allowed: {sorted(_CASE_KEYS)}") + + name = raw.get("name") + if not isinstance(name, str) or not name: + raise _fail(source, f"{where} ('name') must be a non-empty string") + where = f"case {name!r}" + + size = raw.get("size") + if not isinstance(size, Mapping) or not size: + raise _fail(source, f"{where} ('size') must be a non-empty object") + normalized_size: dict[str, int] = {} + for key, value in size.items(): + if not isinstance(key, str) or not key: + raise _fail(source, f"{where} size keys must be non-empty strings") + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise _fail( + source, f"{where} size {key!r} must be a positive integer, got {value!r}" + ) + normalized_size[key] = value + + dtype = raw.get("dtype") + if dtype is not None and not is_canonical_dtype(dtype): + raise _fail(source, f"{where} has unknown dtype {dtype!r}") + + weight = raw.get("weight", 1) + if isinstance(weight, bool) or not isinstance(weight, int) or weight <= 0: + raise _fail(source, f"{where} ('weight') must be a positive integer, got {weight!r}") + + tags = raw.get("tags", ()) + if isinstance(tags, (str, bytes)) or not isinstance(tags, Sequence): + raise _fail(source, f"{where} ('tags') must be a list of strings") + normalized_tags: list[str] = [] + for tag in tags: + if not isinstance(tag, str) or not tag: + raise _fail(source, f"{where} tags must be non-empty strings, got {tag!r}") + if tag in normalized_tags: + raise _fail(source, f"{where} has duplicate tag {tag!r}") + normalized_tags.append(tag) + + return CorpusCase( + name=name, + size=normalized_size, + dtype=dtype, + weight=weight, + tags=tuple(normalized_tags), + ) + + +def load_shape_corpus(path: str | Path) -> ShapeCorpus: + """Parse and structurally validate a shape corpus file. + + Raises :class:`CorpusError` for a missing file, invalid JSON, an + unsupported schema version, or any malformed case. The check against a + specific operation happens separately in + :func:`validate_corpus_against_spec`. + """ + source = str(path) + file_path = Path(path) + if not file_path.is_file(): + raise _fail(source, "file not found") + try: + text = file_path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError) as exc: + raise _fail(source, f"cannot read file: {exc}") from exc + try: + raw = json.loads(text) + except json.JSONDecodeError as exc: + raise _fail(source, f"invalid JSON: {exc}") from exc + + if not isinstance(raw, Mapping): + raise _fail(source, f"top level must be an object, got {type(raw).__name__}") + unknown = sorted(set(raw) - _TOP_LEVEL_KEYS) + if unknown: + raise _fail(source, f"unknown top-level field(s) {unknown}; " + f"allowed: {sorted(_TOP_LEVEL_KEYS)}") + + version = raw.get("schema_version") + if isinstance(version, bool) or not isinstance(version, int): + raise _fail(source, "('schema_version') must be an integer") + if version != CORPUS_SCHEMA_VERSION: + raise _fail( + source, + f"unsupported schema_version {version}; this harness accepts " + f"{CORPUS_SCHEMA_VERSION}", + ) + + operation = raw.get("operation") + if not isinstance(operation, str) or not operation: + raise _fail(source, "('operation') must be a non-empty string") + + raw_cases = raw.get("cases") + if not isinstance(raw_cases, list) or not raw_cases: + raise _fail(source, "('cases') must be a non-empty list") + + cases = tuple(_parse_case(item, index, source) for index, item in enumerate(raw_cases)) + seen: set[str] = set() + for case in cases: + if case.name in seen: + raise _fail(source, f"duplicate case name {case.name!r}") + seen.add(case.name) + + return ShapeCorpus(operation=operation, cases=cases, source=source) + + + +def validate_corpus_against_spec(corpus: ShapeCorpus, spec: KernelSpec) -> None: + """Check a loaded corpus against the selected specification. + + Raises :class:`CorpusError` when the operation does not match, a case + uses shape keys or a dtype the spec does not declare, or two cases + resolve to the same ``(size, dtype)`` configuration. Runs before any GPU + allocation, so an invalid corpus never costs device memory. + """ + if corpus.operation != spec.name: + raise _fail( + corpus.source, + f"corpus operation {corpus.operation!r} does not match the selected " + f"spec {spec.name!r}", + ) + + shape_keys = set(spec.shape_keys) + resolved: dict[tuple[tuple[tuple[str, int], ...], str], str] = {} + for case in corpus.cases: + extra = sorted(set(case.size) - shape_keys) + missing = sorted(shape_keys - set(case.size)) + if extra or missing: + raise _fail( + corpus.source, + f"case {case.name!r} size keys {sorted(case.size)} do not match " + f"shape_keys {sorted(shape_keys)} (unexpected={extra}, missing={missing})", + ) + dtype = case.dtype if case.dtype is not None else spec.primary_dtype + if dtype not in spec.dtypes: + raise _fail( + corpus.source, + f"case {case.name!r} dtype {dtype!r} is not declared in dtypes " + f"{list(spec.dtypes)}", + ) + key = (tuple(sorted(case.size.items())), dtype) + if key in resolved: + raise _fail( + corpus.source, + f"cases {resolved[key]!r} and {case.name!r} resolve to the same " + f"(size, dtype) configuration; merge their weights instead", + ) + resolved[key] = case.name + + +def weighted_aggregate( + entries: Sequence[Mapping[str, Any]], +) -> dict[str, dict[str, float | int]]: + """Weighted latency aggregation, grouped per dtype. + + Args: + entries: benchmark results, each with ``dtype`` (string), ``weight`` + (positive int), ``kernel_ms`` and ``ref_ms``. + + Returns: + ``{dtype: {"cases": n, "weight": total, "kernel_ms": weighted, + "ref_ms": weighted, "speedup": weighted_ref / weighted_kernel}}``. + Results from different dtypes are never mixed into one aggregate. + """ + groups: dict[str, list[Mapping[str, Any]]] = {} + for entry in entries: + groups.setdefault(str(entry["dtype"]), []).append(entry) + + out: dict[str, dict[str, float | int]] = {} + for dtype, group in groups.items(): + total_weight = sum(int(entry["weight"]) for entry in group) + kernel_ms = ( + sum(float(entry["kernel_ms"]) * int(entry["weight"]) for entry in group) + / total_weight + ) + ref_ms = ( + sum(float(entry["ref_ms"]) * int(entry["weight"]) for entry in group) + / total_weight + ) + out[dtype] = { + "cases": len(group), + "weight": total_weight, + "kernel_ms": kernel_ms, + "ref_ms": ref_ms, + "speedup": (ref_ms / kernel_ms) if kernel_ms > 0 else 0.0, + } + return out diff --git a/autokernel/verification/outputs.py b/autokernel/verification/outputs.py new file mode 100644 index 00000000..94598892 --- /dev/null +++ b/autokernel/verification/outputs.py @@ -0,0 +1,531 @@ +"""Structured output-tree traversal and comparison. + +A kernel may return a single tensor or an arbitrary tree of tensors, tuples, +lists, dictionaries and named tuples. Every leaf receives a stable diagnostic +path:: + + output + output[0] + output.updated_residual + output["aux"][1] + +Comparison rules: + +* the candidate and reference trees must have identical structure (same + containers, same keys, same leaf kinds); +* tensor leaves must share shape and dtype; floating tensors are compared + with the tolerance declared for their dtype, non-floating tensors must be + bitwise equal; +* NaN and infinity are detected per path; +* non-tensor (metadata) leaves must match exactly unless the operation's + :class:`~autokernel.specs.OutputSpec` disables that comparison; +* nothing is silently dropped: a configured ``included_paths`` entry that + does not exist is an error, and mismatched structures fail loudly. + +This module never initializes a GPU; ``torch`` is imported lazily so the +package stays importable on CPU-only machines. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Iterable, Mapping, Sequence + +from ..specs.dtypes import canonical_dtype_name +from ..specs.types import OutputSpec, Tolerance + +__all__ = [ + "DEFAULT_TOLERANCE", + "LeafRecord", + "OutputTreeError", + "TreeComparison", + "compare_deterministic", + "compare_output_trees", + "compare_tensor_leaf", + "flatten_output_tree", + "tree_has_nan_or_inf", +] + +#: Historical fallback when a spec declares no tolerance for a leaf dtype. +DEFAULT_TOLERANCE = Tolerance(atol=1e-2, rtol=1e-2) + + +class OutputTreeError(ValueError): + """Raised when an output tree cannot be traversed as configured.""" + + +def _torch() -> Any: + import torch # local import: keep module import torch-free + + return torch + + +def _is_tensor(value: Any) -> bool: + try: + import torch + except ImportError: # pragma: no cover - torch is a hard dependency + return False + return isinstance(value, torch.Tensor) + + +def _is_namedtuple(value: Any) -> bool: + return ( + isinstance(value, tuple) + and hasattr(value, "_fields") + and all(isinstance(field, str) for field in value._fields) + ) + + +def _dict_key_path(path: str, key: Any) -> str: + return f'{path}["{key}"]' if isinstance(key, str) else f"{path}[{key!r}]" + + +def flatten_output_tree(output: Any) -> tuple[tuple[str, Any], ...]: + """Flatten ``output`` into ``(path, leaf)`` pairs with stable ordering. + + Dictionary keys are visited in sorted order so flattening never depends on + insertion order. Tensors, empty containers and any other value become + leaves; non-empty tuples, lists, dicts and named tuples are traversed. + """ + leaves: list[tuple[str, Any]] = [] + + def walk(node: Any, path: str) -> None: + if _is_tensor(node): + leaves.append((path, node)) + return + if _is_namedtuple(node): + for field in node._fields: + walk(getattr(node, field), f"{path}.{field}") + return + if isinstance(node, Mapping): + if not node: + leaves.append((path, node)) + return + for key in sorted(node, key=lambda k: str(k)): + walk(node[key], _dict_key_path(path, key)) + return + if isinstance(node, (tuple, list)): + if not node: + leaves.append((path, node)) + return + for index, item in enumerate(node): + walk(item, f"{path}[{index}]") + return + leaves.append((path, node)) + + walk(output, "output") + return tuple(leaves) + + +def _leaf_kind(leaf: Any) -> str: + return "tensor" if _is_tensor(leaf) else "metadata" + + +def _filter_leaves( + leaves: Iterable[tuple[str, Any]], + included_paths: tuple[str, ...] | None, + *, + side: str, +) -> list[tuple[str, Any]]: + pairs = list(leaves) + if included_paths is None: + return pairs + available = {path for path, _ in pairs} + missing = [path for path in included_paths if path not in available] + if missing: + raise OutputTreeError( + f"output_spec.included_paths not present in the {side} output: " + f"{missing}; available paths: {sorted(available)}" + ) + included = set(included_paths) + return [(path, leaf) for path, leaf in pairs if path in included] + + + +@dataclass(frozen=True) +class LeafRecord: + """Comparison outcome for one output leaf.""" + + path: str + kind: str # "tensor" | "metadata" + match: bool + reason: str = "" + max_abs_error: float | None = None + mean_abs_error: float | None = None + pct_within_tol: float | None = None + has_nan: bool = False + has_inf: bool = False + + def as_dict(self) -> dict[str, Any]: + return { + "path": self.path, + "kind": self.kind, + "match": self.match, + "reason": self.reason, + "max_abs_error": self.max_abs_error, + "mean_abs_error": self.mean_abs_error, + "pct_within_tol": self.pct_within_tol, + "has_nan": self.has_nan, + "has_inf": self.has_inf, + } + + +@dataclass(frozen=True) +class TreeComparison: + """Aggregated comparison of two output trees.""" + + match: bool + reason: str + leaves: tuple[LeafRecord, ...] + structure_match: bool = True + + @property + def worst_abs_error(self) -> float: + """Largest per-leaf maximum absolute error (0.0 when none).""" + worst = 0.0 + for leaf in self.leaves: + if leaf.max_abs_error is not None: + worst = max(worst, leaf.max_abs_error) + return worst + + @property + def worst_pct_within_tol(self) -> float: + """Smallest per-leaf within-tolerance percentage (100.0 when none).""" + worst = 100.0 + for leaf in self.leaves: + if leaf.pct_within_tol is not None: + worst = min(worst, leaf.pct_within_tol) + return worst + + def has_nan_or_inf(self) -> bool: + return any(leaf.has_nan or leaf.has_inf for leaf in self.leaves) + + def first_failure(self) -> LeafRecord | None: + for leaf in self.leaves: + if not leaf.match: + return leaf + return None + + def leaf_records(self) -> list[dict[str, Any]]: + return [leaf.as_dict() for leaf in self.leaves] + + +def _structure_mismatch_reason( + candidate: Sequence[tuple[str, Any]], expected: Sequence[tuple[str, Any]] +) -> str | None: + cand_map = {path: leaf for path, leaf in candidate} + exp_map = {path: leaf for path, leaf in expected} + extra = [path for path, _ in candidate if path not in exp_map] + missing = [path for path, _ in expected if path not in cand_map] + parts = [] + if missing: + parts.append(f"missing output path(s) {missing}") + if extra: + parts.append(f"unexpected output path(s) {extra}") + if parts: + return "; ".join(parts) + for path, exp_leaf in expected: + cand_leaf = cand_map[path] + if _leaf_kind(cand_leaf) != _leaf_kind(exp_leaf): + return ( + f"leaf kind mismatch at {path}: " + f"{_leaf_kind(cand_leaf)} vs {_leaf_kind(exp_leaf)}" + ) + return None + + +def _tolerance_for( + leaf: Any, + tolerances: Mapping[str, Tolerance], + default: Tolerance, +) -> Tolerance: + """Pick the tolerance declared for a tensor leaf's dtype.""" + if not leaf.is_floating_point(): + return Tolerance(atol=0.0, rtol=0.0) + try: + name = canonical_dtype_name(leaf.dtype) + except ValueError: + # Non-canonical floating dtype (e.g. float64): the spec cannot declare + # a tolerance for it, so apply the harness fallback. + return default + return tolerances.get(name, default) + + +def _metadata_equal(candidate: Any, expected: Any) -> bool: + try: + equal = candidate == expected + except Exception: + return False + if isinstance(equal, bool): + return equal + try: + return bool(equal) + except Exception: + return False + + +def compare_tensor_leaf( + path: str, + candidate: Any, + expected: Any, + tolerance: Tolerance, +) -> LeafRecord: + """Compare two tensor leaves with statistics, NaN/Inf flags and a reason. + + Public so other verifiers (e.g. gradient comparison) can reuse the exact + forward-comparison semantics. + """ + torch = _torch() + is_float = candidate.is_floating_point() + has_nan = bool(torch.isnan(candidate).any().item()) if is_float else False + has_inf = bool(torch.isinf(candidate).any().item()) if is_float else False + + if candidate.shape != expected.shape: + return LeafRecord( + path=path, + kind="tensor", + match=False, + reason=f"shape mismatch: {tuple(candidate.shape)} vs {tuple(expected.shape)}", + max_abs_error=float("inf"), + mean_abs_error=float("inf"), + pct_within_tol=0.0, + has_nan=has_nan, + has_inf=has_inf, + ) + if candidate.dtype != expected.dtype: + return LeafRecord( + path=path, + kind="tensor", + match=False, + reason=f"dtype mismatch: {candidate.dtype} vs {expected.dtype}", + max_abs_error=float("inf"), + mean_abs_error=float("inf"), + pct_within_tol=0.0, + has_nan=has_nan, + has_inf=has_inf, + ) + + if not is_float: + # Integer and boolean leaves carry exact values; a tolerance would be + # meaningless. + match = bool(torch.equal(candidate, expected)) + return LeafRecord( + path=path, + kind="tensor", + match=match, + reason="" if match else "non-floating tensors are not bitwise equal", + max_abs_error=0.0 if match else float("inf"), + mean_abs_error=0.0 if match else float("inf"), + pct_within_tol=100.0 if match else 0.0, + ) + + out_f = candidate.float() + exp_f = expected.float() + if out_f.numel() == 0: + return LeafRecord( + path=path, + kind="tensor", + match=True, + max_abs_error=0.0, + mean_abs_error=0.0, + pct_within_tol=100.0, + ) + + abs_diff = (out_f - exp_f).abs() + max_abs = abs_diff.max().item() + mean_abs = abs_diff.mean().item() + within = ( + (abs_diff <= tolerance.atol + tolerance.rtol * exp_f.abs()).float().mean().item() + * 100.0 + ) + match = bool(torch.allclose(out_f, exp_f, atol=tolerance.atol, rtol=tolerance.rtol)) + reason = "" + if not match: + reason = ( + f"max_abs_error={max_abs:.6e} exceeds " + f"tol(atol={tolerance.atol}, rtol={tolerance.rtol})" + ) + if has_nan: + reason += "; output contains NaN" + elif has_inf: + reason += "; output contains infinity" + return LeafRecord( + path=path, + kind="tensor", + match=match, + reason=reason, + max_abs_error=max_abs, + mean_abs_error=mean_abs, + pct_within_tol=within, + has_nan=has_nan, + has_inf=has_inf, + ) + + +def compare_output_trees( + candidate: Any, + expected: Any, + tolerances: Mapping[str, Tolerance], + *, + output_spec: OutputSpec | None = None, + default_tolerance: Tolerance = DEFAULT_TOLERANCE, + relax: float = 1.0, +) -> TreeComparison: + """Compare two output trees leaf by leaf. + + Args: + candidate: the kernel's output tree. + expected: the reference output tree. + tolerances: canonical dtype name -> :class:`Tolerance`, typically + ``spec.tolerances``. Each floating tensor leaf is compared with + the tolerance declared for its own dtype. + output_spec: optional policy controlling which paths participate and + whether non-tensor leaves must match exactly. + default_tolerance: fallback for leaf dtypes the mapping does not + cover (matches the historical harness default of 1e-2/1e-2). + relax: multiplies both tolerances (used by the numerical-stability + stage for adversarial inputs). + """ + included = output_spec.included_paths if output_spec is not None else None + compare_meta = output_spec.compare_non_tensors if output_spec is not None else True + + cand_leaves = _filter_leaves( + flatten_output_tree(candidate), included, side="candidate" + ) + exp_leaves = _filter_leaves( + flatten_output_tree(expected), included, side="reference" + ) + + mismatch = _structure_mismatch_reason(cand_leaves, exp_leaves) + if mismatch is not None: + return TreeComparison( + match=False, + reason=f"output structure mismatch: {mismatch}", + leaves=(), + structure_match=False, + ) + + records: list[LeafRecord] = [] + exp_map = dict(exp_leaves) + for path, cand_leaf in cand_leaves: + exp_leaf = exp_map[path] + if _is_tensor(cand_leaf): + tol = _tolerance_for(exp_leaf, tolerances, default_tolerance) + if relax != 1.0: + tol = Tolerance(atol=tol.atol * relax, rtol=tol.rtol * relax) + records.append(compare_tensor_leaf(path, cand_leaf, exp_leaf, tol)) + continue + if not compare_meta: + records.append( + LeafRecord(path=path, kind="metadata", match=True, reason="not compared") + ) + continue + equal = _metadata_equal(cand_leaf, exp_leaf) + records.append( + LeafRecord( + path=path, + kind="metadata", + match=equal, + reason="" if equal else f"metadata mismatch: {cand_leaf!r} != {exp_leaf!r}", + ) + ) + + failed = [record for record in records if not record.match] + if failed: + first = failed[0] + reason = first.reason + if first.path != "output": + reason = f"{first.path}: {reason}" + return TreeComparison(match=False, reason=reason, leaves=tuple(records)) + return TreeComparison(match=True, reason="", leaves=tuple(records)) + + + +def compare_deterministic( + first: Any, + other: Any, + *, + output_spec: OutputSpec | None = None, +) -> TreeComparison: + """Bitwise comparison of two runs of the same kernel. + + Every tensor leaf must be bitwise identical (``torch.equal``) and every + compared metadata leaf must be exactly equal; the tree structures must + match. Statistics for differing tensor leaves are reported so failures + stay diagnosable. + """ + torch = _torch() + included = output_spec.included_paths if output_spec is not None else None + compare_meta = output_spec.compare_non_tensors if output_spec is not None else True + first_leaves = _filter_leaves(flatten_output_tree(first), included, side="first run") + other_leaves = _filter_leaves(flatten_output_tree(other), included, side="later run") + + mismatch = _structure_mismatch_reason(first_leaves, other_leaves) + if mismatch is not None: + return TreeComparison( + match=False, + reason=f"output structure mismatch between runs: {mismatch}", + leaves=(), + structure_match=False, + ) + + records: list[LeafRecord] = [] + for (path, a), (_, b) in zip(first_leaves, other_leaves): + if _is_tensor(a): + max_diff: float | None = None + mean_diff: float | None = None + if a.shape == b.shape and a.dtype == b.dtype and a.is_floating_point(): + diff = (a.float() - b.float()).abs() + max_diff = diff.max().item() if diff.numel() else 0.0 + mean_diff = diff.mean().item() if diff.numel() else 0.0 + equal = bool(torch.equal(a, b)) + if equal: + reason = "" + elif max_diff is not None: + reason = f"runs differ (max_diff={max_diff:.6e})" + else: + reason = "runs differ (shape or dtype changed)" + records.append( + LeafRecord( + path=path, + kind="tensor", + match=equal, + reason=reason, + max_abs_error=max_diff, + mean_abs_error=mean_diff, + ) + ) + continue + if not compare_meta: + records.append( + LeafRecord(path=path, kind="metadata", match=True, reason="not compared") + ) + continue + equal = _metadata_equal(a, b) + records.append( + LeafRecord( + path=path, + kind="metadata", + match=equal, + reason="" if equal else f"metadata changed between runs: {a!r} != {b!r}", + ) + ) + + failed = [record for record in records if not record.match] + if failed: + first_fail = failed[0] + reason = first_fail.reason + if first_fail.path != "output": + reason = f"{first_fail.path}: {reason}" + return TreeComparison(match=False, reason=reason, leaves=tuple(records)) + return TreeComparison(match=True, reason="", leaves=tuple(records)) + + +def tree_has_nan_or_inf(tree: Any) -> bool: + """True when any floating tensor leaf contains NaN or infinity.""" + torch = _torch() + for _, leaf in flatten_output_tree(tree): + if _is_tensor(leaf) and leaf.is_floating_point(): + if bool(torch.isnan(leaf).any().item()) or bool(torch.isinf(leaf).any().item()): + return True + return False diff --git a/autokernel/verification/results.py b/autokernel/verification/results.py new file mode 100644 index 00000000..4cfaff6b --- /dev/null +++ b/autokernel/verification/results.py @@ -0,0 +1,111 @@ +"""Versioned, atomic machine-readable benchmark results.""" + +from __future__ import annotations + +import json +import math +import os +import platform +import tempfile +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Mapping + +__all__ = [ + "RESULT_SCHEMA_VERSION", + "collect_environment_metadata", + "result_envelope", + "write_result_atomic", +] + +RESULT_SCHEMA_VERSION = 1 + + +def collect_environment_metadata(device: str) -> dict[str, Any]: + """Collect runtime versions without requiring Triton or a CUDA device.""" + import torch + + try: + import triton + + triton_version = getattr(triton, "__version__", "unknown") + except Exception: + triton_version = None + + gpu_name = None + if device.startswith("cuda"): + try: + gpu_name = torch.cuda.get_device_name(device) + except Exception: + gpu_name = None + return { + "python_version": platform.python_version(), + "platform": platform.platform(), + "torch_version": getattr(torch, "__version__", "unknown"), + "triton_version": triton_version, + "cuda_version": getattr(getattr(torch, "version", None), "cuda", None), + "device": device, + "gpu_name": gpu_name, + } + + +def result_envelope(operation: str, **sections: Any) -> dict[str, Any]: + """Build the stable top-level result record.""" + return { + "schema_version": RESULT_SCHEMA_VERSION, + "created_at": datetime.now(timezone.utc).isoformat(), + "operation": operation, + **sections, + } + + +def _json_default(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if hasattr(value, "as_dict"): + return _json_safe(value.as_dict()) + return str(value) + + +def _json_safe(value: Any) -> Any: + if isinstance(value, float) and not math.isfinite(value): + return str(value) + if isinstance(value, Mapping): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + return value + + +def write_result_atomic(path: str | Path, payload: Mapping[str, Any]) -> Path: + """Atomically replace ``path`` with a complete JSON document.""" + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=destination.parent, + prefix=f".{destination.name}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary = Path(handle.name) + json.dump( + _json_safe(payload), + handle, + indent=2, + sort_keys=True, + default=_json_default, + allow_nan=False, + ) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, destination) + except Exception: + if temporary is not None: + temporary.unlink(missing_ok=True) + raise + return destination diff --git a/autokernel/workload/__init__.py b/autokernel/workload/__init__.py new file mode 100644 index 00000000..be8e2447 --- /dev/null +++ b/autokernel/workload/__init__.py @@ -0,0 +1,45 @@ +"""Versioned FastVideo generation workload manifests.""" + +from .result import ( + RESULT_SCHEMA_VERSION, + GenerationRunResult, + classify_end_to_end, + compare_frame_outputs, + load_generation_result, + write_generation_result, +) +from .types import ( + WORKLOAD_SCHEMA_VERSION, + MeasurementSpec, + ModeEnvSpec, + ModelRef, + ParitySpec, + PerformanceSpec, + RuntimeSpec, + SamplingSpec, + WorkloadError, + WorkloadManifest, + dump_workload, + load_workload, +) + +__all__ = [ + "RESULT_SCHEMA_VERSION", + "WORKLOAD_SCHEMA_VERSION", + "GenerationRunResult", + "MeasurementSpec", + "ModeEnvSpec", + "ModelRef", + "ParitySpec", + "PerformanceSpec", + "RuntimeSpec", + "SamplingSpec", + "WorkloadError", + "WorkloadManifest", + "classify_end_to_end", + "compare_frame_outputs", + "dump_workload", + "load_generation_result", + "load_workload", + "write_generation_result", +] diff --git a/autokernel/workload/_validate.py b/autokernel/workload/_validate.py new file mode 100644 index 00000000..30a8b635 --- /dev/null +++ b/autokernel/workload/_validate.py @@ -0,0 +1,99 @@ +"""Shared JSON-metadata validators for workload and result schemas.""" + +from __future__ import annotations + +import math +from typing import Any, Mapping, Sequence + + +class SchemaError(ValueError): + """Malformed or unsafe schema payload.""" + + +def fail(kind: str, source: object, location: str, message: str) -> SchemaError: + return SchemaError(f"{kind} {source!r}: {location}: {message}") + + +def mapping( + value: Any, + source: object, + location: str, + *, + kind: str, + non_empty: bool = False, + forbidden_keys: set[str] | None = None, +) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or (non_empty and not value): + qualifier = "non-empty " if non_empty else "" + raise fail(kind, source, location, f"must be a {qualifier}object") + forbidden = {k.lower() for k in (forbidden_keys or set())} + for key in value: + if not isinstance(key, str) or not key: + raise fail(kind, source, location, "keys must be non-empty strings") + if key.lower() in forbidden: + raise fail( + kind, + source, + f"{location}.{key}", + "content or secret fields are forbidden", + ) + return value + + +def text(value: Any, source: object, location: str, *, kind: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise fail(kind, source, location, "must be a non-empty string") + return value.strip() + + +def optional_text( + value: Any, source: object, location: str, *, kind: str +) -> str | None: + if value is None: + return None + return text(value, source, location, kind=kind) + + +def positive_int(value: Any, source: object, location: str, *, kind: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise fail(kind, source, location, "must be a positive integer") + return value + + +def non_negative_int( + value: Any, source: object, location: str, *, kind: str +) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise fail(kind, source, location, "must be a non-negative integer") + return value + + +def finite_number( + value: Any, + source: object, + location: str, + *, + kind: str, + minimum: float | None = None, + maximum: float | None = None, +) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise fail(kind, source, location, "must be a finite number") + number = float(value) + if not math.isfinite(number): + raise fail(kind, source, location, "must be a finite number") + if minimum is not None and number < minimum: + raise fail(kind, source, location, f"must be >= {minimum}") + if maximum is not None and number > maximum: + raise fail(kind, source, location, f"must be <= {maximum}") + return number + + +def finite_non_negative( + value: Any, source: object, location: str, *, kind: str +) -> float: + return finite_number(value, source, location, kind=kind, minimum=0.0) + + +def is_sequence(value: Any) -> bool: + return isinstance(value, Sequence) and not isinstance(value, (str, bytes)) diff --git a/autokernel/workload/launcher.py b/autokernel/workload/launcher.py new file mode 100644 index 00000000..3a9ce4f2 --- /dev/null +++ b/autokernel/workload/launcher.py @@ -0,0 +1,395 @@ +"""MotionKernel-side bridge that drives a FastVideo generation launcher. + +The FastVideo checkout owns the GPU process. This module validates workloads, +invokes the shared launcher script in separate processes per mode, and records +resume-friendly stage state under an output directory. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +from .result import ( + GenerationRunResult, + classify_end_to_end, + compare_frame_outputs, + load_generation_result, +) +from .types import WorkloadError, WorkloadManifest, load_workload + +DEFAULT_LAUNCHER_RELATIVE = Path( + "examples/inference/optimizations/generation_launcher.py" +) +STATE_NAME = "launcher_state.json" +STAGES = ("native", "optimized", "compare") + + +@dataclass(frozen=True) +class LauncherPaths: + output_dir: Path + state_path: Path + native_result: Path + optimized_result: Path + comparison_path: Path + + +def _validate_result_identity( + result: GenerationRunResult, + *, + workload_id: str, + requested_mode: str, +) -> None: + if result.workload_id != workload_id: + raise WorkloadError( + f"generation result workload_id {result.workload_id!r} does not " + f"match manifest {workload_id!r}" + ) + accepted_modes = ( + {"optimized", "fused"} if requested_mode == "optimized" + else {requested_mode} + ) + if result.mode not in accepted_modes: + raise WorkloadError( + f"generation result mode {result.mode!r} does not match requested " + f"mode {requested_mode!r}" + ) + + +def _paths(output_dir: str | Path) -> LauncherPaths: + root = Path(output_dir).expanduser().resolve() + return LauncherPaths( + output_dir=root, + state_path=root / STATE_NAME, + native_result=root / "native_result.json", + optimized_result=root / "optimized_result.json", + comparison_path=root / "comparison.json", + ) + + +def _read_state(path: Path) -> dict[str, Any]: + fresh: dict[str, Any] = { + "schema_version": 1, + "completed_stages": [], + "failed_stages": {}, + } + if not path.is_file(): + return fresh + try: + state = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise WorkloadError( + f"launcher state {path!s}: invalid JSON: {exc}" + ) from exc + if not isinstance(state, dict) or state.get("schema_version") != 1: + raise WorkloadError( + f"launcher state {path!s}: unsupported or malformed state; " + "delete the file or run with resume disabled" + ) + merged = {**fresh, **state} + if not isinstance(merged.get("completed_stages"), list): + raise WorkloadError( + f"launcher state {path!s}: completed_stages must be a list" + ) + if not isinstance(merged.get("failed_stages"), dict): + raise WorkloadError( + f"launcher state {path!s}: failed_stages must be an object" + ) + return merged + + +def _write_state(path: Path, state: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(state, indent=2) + "\n", encoding="utf-8") + temporary.replace(path) + + +def resolve_launcher( + fastvideo_checkout: str | Path, + launcher_script: str | Path | None = None, +) -> Path: + """Locate the FastVideo generation launcher script.""" + root = Path(fastvideo_checkout) + if not root.is_dir(): + raise WorkloadError( + f"FastVideo checkout not found: {root}" + ) + script = ( + Path(launcher_script) + if launcher_script is not None + else root / DEFAULT_LAUNCHER_RELATIVE + ) + if not script.is_file(): + raise WorkloadError( + f"generation launcher not found: {script}. " + "Install/update the FastVideo branch that provides " + "examples/inference/optimizations/generation_launcher.py" + ) + return script + + +def build_launcher_command( + *, + python: str, + launcher: Path, + workload: Path, + mode: str, + output_dir: Path, + model_override: str | None = None, +) -> list[str]: + """Construct an argv list for one launcher process (no shell).""" + command = [ + python, + str(launcher), + "--workload", + str(workload), + "--mode", + mode, + "--output-dir", + str(output_dir), + ] + if model_override: + command.extend(["--model", model_override]) + return command + + +def run_mode( + *, + fastvideo_checkout: str | Path, + workload: str | Path | WorkloadManifest, + mode: str, + output_dir: str | Path, + python: str | None = None, + launcher_script: str | Path | None = None, + model_override: str | None = None, + env: Mapping[str, str] | None = None, + check: bool = True, + timeout: float | None = None, +) -> subprocess.CompletedProcess[str]: + """Run one baseline or optimized generation mode in a subprocess.""" + if isinstance(workload, WorkloadManifest): + raise WorkloadError( + "run_mode requires a workload file path so the child process " + "can load the same manifest" + ) + workload_path = Path(workload) + if not workload_path.is_file(): + raise WorkloadError(f"workload file not found: {workload_path}") + + # Validate early on the parent side. + load_workload(workload_path) + + launcher = resolve_launcher(fastvideo_checkout, launcher_script) + command = build_launcher_command( + python=python or sys.executable, + launcher=launcher, + workload=workload_path, + mode=mode, + output_dir=Path(output_dir), + model_override=model_override, + ) + child_env = os.environ.copy() + if env: + child_env.update(env) + # Ensure the FastVideo checkout is importable even when the user has not + # installed it into the active virtualenv. + checkout = str(Path(fastvideo_checkout).resolve()) + existing = child_env.get("PYTHONPATH", "") + child_env["PYTHONPATH"] = ( + checkout if not existing else f"{checkout}{os.pathsep}{existing}" + ) + + logs = Path(output_dir) + logs.mkdir(parents=True, exist_ok=True) + stdout_path = logs / f"{mode}.stdout.log" + stderr_path = logs / f"{mode}.stderr.log" + try: + with stdout_path.open("w", encoding="utf-8") as stdout_file, ( + stderr_path.open("w", encoding="utf-8") + ) as stderr_file: + raw = subprocess.run( + command, + check=False, + text=True, + stdout=stdout_file, + stderr=stderr_file, + env=child_env, + cwd=checkout, + timeout=timeout, + ) + except subprocess.TimeoutExpired as exc: + raise WorkloadError( + f"launcher mode {mode!r} timed out after {timeout} seconds; " + f"logs: {stdout_path}, {stderr_path}" + ) from exc + completed = subprocess.CompletedProcess( + raw.args, + raw.returncode, + stdout_path.read_text(encoding="utf-8"), + stderr_path.read_text(encoding="utf-8"), + ) + if check and completed.returncode != 0: + raise WorkloadError( + f"launcher mode {mode!r} failed with exit {completed.returncode}\n" + f"stdout:\n{completed.stdout}\n" + f"stderr:\n{completed.stderr}" + ) + return completed + + +def run_ab( + *, + fastvideo_checkout: str | Path, + workload: str | Path, + output_dir: str | Path, + python: str | None = None, + launcher_script: str | Path | None = None, + model_override: str | None = None, + modes: Sequence[str] = ("native", "optimized"), + resume: bool = True, + timeout: float | None = None, +) -> dict[str, Any]: + """Run native and optimized modes with resume-friendly stage tracking.""" + paths = _paths(output_dir) + paths.output_dir.mkdir(parents=True, exist_ok=True) + allowed_modes = {"native", "optimized"} + unknown = [mode for mode in modes if mode not in allowed_modes] + if unknown: + raise WorkloadError( + f"unsupported launcher mode(s) {sorted(unknown)}; " + "expected 'native' and/or 'optimized'" + ) + manifest = load_workload(workload) + state = _read_state(paths.state_path) if resume else { + "schema_version": 1, + "completed_stages": [], + "failed_stages": {}, + } + completed = set(state.get("completed_stages", [])) + + results: dict[str, GenerationRunResult] = {} + for mode in modes: + result_path = ( + paths.native_result if mode == "native" else paths.optimized_result + ) + # Accept legacy fused naming for Wan parity. + if mode == "optimized" and not result_path.is_file(): + legacy = paths.output_dir / "fused_result.json" + if legacy.is_file(): + result_path = legacy + + if resume and mode in completed and result_path.is_file(): + loaded = load_generation_result(result_path) + _validate_result_identity( + loaded, + workload_id=manifest.workload_id, + requested_mode=mode, + ) + results[mode] = loaded + continue + + mode_env = {} + if manifest.mode_env is not None: + mode_env = manifest.mode_env.for_mode(mode) + try: + run_mode( + fastvideo_checkout=fastvideo_checkout, + workload=workload, + mode=mode, + output_dir=paths.output_dir, + python=python, + launcher_script=launcher_script, + model_override=model_override, + env=mode_env or None, + check=True, + timeout=timeout, + ) + # Launcher may write mode-specific names. + written = paths.output_dir / f"{mode}_result.json" + if not written.is_file() and mode == "optimized": + written = paths.output_dir / "fused_result.json" + if not written.is_file(): + raise WorkloadError( + f"launcher did not write expected result: {written}" + ) + loaded = load_generation_result(written) + _validate_result_identity( + loaded, + workload_id=manifest.workload_id, + requested_mode=mode, + ) + results[mode] = loaded + completed.add(mode) + state["completed_stages"] = sorted(completed) + state.get("failed_stages", {}).pop(mode, None) + _write_state(paths.state_path, state) + except Exception as exc: + state.setdefault("failed_stages", {})[mode] = str(exc) + _write_state(paths.state_path, state) + raise + + comparison = None + if "native" in results and "optimized" in results: + if not (resume and "compare" in completed and paths.comparison_path.is_file()): + comparison = classify_end_to_end( + results["native"], + results["optimized"], + min_speedup=manifest.performance.min_end_to_end_speedup + if manifest.performance + else 1.01, + max_peak_memory_regression=( + manifest.performance.max_peak_memory_regression + if manifest.performance + else 0.05 + ), + ) + parity_policy = manifest.parity + parity = compare_frame_outputs( + results["native"].frames_path, + results["optimized"].frames_path, + policy=parity_policy.policy if parity_policy else "byte_equal", + atol=( + parity_policy.atol + if parity_policy and parity_policy.atol is not None + else 0.0 + ), + rtol=( + parity_policy.rtol + if parity_policy and parity_policy.rtol is not None + else 0.0 + ), + ) + comparison["parity"] = parity + if not parity["passed"]: + comparison["classification"] = "failed" + comparison["reason"] = ( + f"output parity failed: {parity['reason']}" + ) + temporary = paths.comparison_path.with_suffix( + paths.comparison_path.suffix + ".tmp" + ) + temporary.write_text( + json.dumps(comparison, indent=2) + "\n", encoding="utf-8" + ) + temporary.replace(paths.comparison_path) + completed.add("compare") + state["completed_stages"] = sorted(completed) + _write_state(paths.state_path, state) + else: + comparison = json.loads( + paths.comparison_path.read_text(encoding="utf-8") + ) + + return { + "workload_id": manifest.workload_id, + "output_dir": str(paths.output_dir), + "results": {mode: result.as_dict() for mode, result in results.items()}, + "comparison": comparison, + "state": state, + } diff --git a/autokernel/workload/result.py b/autokernel/workload/result.py new file mode 100644 index 00000000..a1bb1f51 --- /dev/null +++ b/autokernel/workload/result.py @@ -0,0 +1,496 @@ +"""Structured generation-run result schema written by FastVideo launchers. + +MotionKernel validates these artifacts before ranking native-versus-optimized +end-to-end outcomes. Values are metadata and numeric measurements only; frame +tensors live in separate ``.npy`` files referenced by path. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +from ._validate import ( + finite_number, + mapping as _mapping_base, + non_negative_int as _non_negative_int_base, + optional_text as _optional_text_base, + positive_int as _positive_int_base, + text as _text_base, +) +from .types import FORBIDDEN_METADATA_KEYS, WorkloadError + +RESULT_SCHEMA_VERSION = 1 + +_TOP_LEVEL_FIELDS = { + "schema_version", + "mode", + "status", + "workload_id", + "model_id", + "request", + "warmups", + "runs", + "wall_seconds", + "median_wall_seconds", + "generation_seconds", + "peak_memory_mb", + "environment", + "frames_path", + "profiler_path", + "log_path", + "failure_reason", + "stage", +} + +_STATUSES = {"ok", "failed", "skipped"} +_MODES = {"native", "optimized", "fused", "candidate"} + + +def _kind() -> str: + return "generation result" + + +def _text(value: Any, source: object, location: str) -> str: + try: + return _text_base(value, source, location, kind=_kind()) + except Exception as exc: # SchemaError subclass of ValueError + raise WorkloadError(str(exc)) from exc + + +def _optional_text( + value: Any, source: object, location: str +) -> str | None: + try: + return _optional_text_base(value, source, location, kind=_kind()) + except Exception as exc: + raise WorkloadError(str(exc)) from exc + + +def _non_negative_int(value: Any, source: object, location: str) -> int: + try: + return _non_negative_int_base(value, source, location, kind=_kind()) + except Exception as exc: + raise WorkloadError(str(exc)) from exc + + +def _positive_int(value: Any, source: object, location: str) -> int: + try: + return _positive_int_base(value, source, location, kind=_kind()) + except Exception as exc: + raise WorkloadError(str(exc)) from exc + + +def _mapping(value: Any, source: object, location: str, *, non_empty: bool = False): + try: + return _mapping_base( + value, + source, + location, + kind=_kind(), + non_empty=non_empty, + forbidden_keys=FORBIDDEN_METADATA_KEYS, + ) + except Exception as exc: + raise WorkloadError(str(exc)) from exc + + +def _finite_number( + value: Any, + source: object, + location: str, + *, + minimum: float | None = None, +) -> float: + try: + return finite_number( + value, source, location, kind=_kind(), minimum=minimum + ) + except Exception as exc: + raise WorkloadError(str(exc)) from exc + + +def _number_list( + value: Any, source: object, location: str +) -> list[float]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise WorkloadError( + f"generation result {source!r}: {location}: must be a list" + ) + numbers: list[float] = [] + for index, item in enumerate(value): + if item is None: + raise WorkloadError( + f"generation result {source!r}: {location}[{index}]: " + "must be a finite non-negative number (None not allowed)" + ) + numbers.append( + _finite_number( + item, + source, + f"{location}[{index}]", + minimum=0.0, + ) + ) + return numbers + + +def _optional_number_list( + value: Any, source: object, location: str +) -> list[float | None]: + if value is None: + return [] + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise WorkloadError( + f"generation result {source!r}: {location}: must be a list" + ) + result: list[float | None] = [] + for index, item in enumerate(value): + if item is None: + result.append(None) + continue + result.append( + _finite_number( + item, + source, + f"{location}[{index}]", + minimum=0.0, + ) + ) + return result + + +@dataclass(frozen=True) +class GenerationRunResult: + """One mode's measured generation run (native or optimized).""" + + mode: str + status: str + workload_id: str + model_id: str + request: Mapping[str, Any] + warmups: int + runs: int + wall_seconds: tuple[float, ...] + median_wall_seconds: float | None + generation_seconds: tuple[float | None, ...] + peak_memory_mb: tuple[float | None, ...] + environment: Mapping[str, Any] + frames_path: str | None = None + profiler_path: str | None = None + log_path: str | None = None + failure_reason: str | None = None + stage: str = "generate" + schema_version: int = RESULT_SCHEMA_VERSION + source: str = "" + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object = "" + ) -> "GenerationRunResult": + raw = _mapping(raw_value, source, "top level", non_empty=True) + unknown = sorted(set(raw) - _TOP_LEVEL_FIELDS) + if unknown: + raise WorkloadError( + f"generation result {source!r}: top level: " + f"unknown field(s) {unknown}" + ) + + version = raw.get("schema_version") + if isinstance(version, bool) or not isinstance(version, int): + raise WorkloadError( + f"generation result {source!r}: schema_version: " + "must be an integer" + ) + if version != RESULT_SCHEMA_VERSION: + raise WorkloadError( + f"generation result {source!r}: schema_version: " + f"unsupported version {version}; expected {RESULT_SCHEMA_VERSION}" + ) + + mode = _text(raw.get("mode"), source, "mode") + if mode not in _MODES: + raise WorkloadError( + f"generation result {source!r}: mode: " + f"must be one of {sorted(_MODES)}" + ) + status = _text(raw.get("status", "ok"), source, "status") + if status not in _STATUSES: + raise WorkloadError( + f"generation result {source!r}: status: " + f"must be one of {sorted(_STATUSES)}" + ) + + wall = tuple(_number_list(raw.get("wall_seconds", []), source, "wall_seconds")) + median = raw.get("median_wall_seconds") + median_value = ( + None + if median is None + else _finite_number( + median, source, "median_wall_seconds", minimum=0.0 + ) + ) + if status == "ok" and not wall: + raise WorkloadError( + f"generation result {source!r}: wall_seconds: " + "required for successful runs" + ) + + return cls( + mode=mode, + status=status, + workload_id=_text(raw.get("workload_id"), source, "workload_id"), + model_id=_text(raw.get("model_id"), source, "model_id"), + request=dict( + _mapping(raw.get("request", {}), source, "request") + ), + warmups=_non_negative_int(raw.get("warmups", 0), source, "warmups"), + runs=_positive_int(raw.get("runs", 1), source, "runs") + if status == "ok" + else _non_negative_int(raw.get("runs", 0), source, "runs"), + wall_seconds=wall, + median_wall_seconds=median_value, + generation_seconds=tuple( + _optional_number_list( + raw.get("generation_seconds", []), + source, + "generation_seconds", + ) + ), + peak_memory_mb=tuple( + _optional_number_list( + raw.get("peak_memory_mb", []), + source, + "peak_memory_mb", + ) + ), + environment=dict( + _mapping(raw.get("environment", {}), source, "environment") + ), + frames_path=_optional_text( + raw.get("frames_path"), source, "frames_path" + ), + profiler_path=_optional_text( + raw.get("profiler_path"), source, "profiler_path" + ), + log_path=_optional_text(raw.get("log_path"), source, "log_path"), + failure_reason=_optional_text( + raw.get("failure_reason"), source, "failure_reason" + ), + stage=_text(raw.get("stage", "generate"), source, "stage"), + schema_version=version, + source=str(source), + ) + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "schema_version": self.schema_version, + "mode": self.mode, + "status": self.status, + "workload_id": self.workload_id, + "model_id": self.model_id, + "request": dict(self.request), + "warmups": self.warmups, + "runs": self.runs, + "wall_seconds": list(self.wall_seconds), + "median_wall_seconds": self.median_wall_seconds, + "generation_seconds": list(self.generation_seconds), + "peak_memory_mb": list(self.peak_memory_mb), + "environment": dict(self.environment), + "stage": self.stage, + } + if self.frames_path is not None: + payload["frames_path"] = self.frames_path + if self.profiler_path is not None: + payload["profiler_path"] = self.profiler_path + if self.log_path is not None: + payload["log_path"] = self.log_path + if self.failure_reason is not None: + payload["failure_reason"] = self.failure_reason + return payload + + +def load_generation_result(path: str | Path) -> GenerationRunResult: + """Load and validate a launcher result JSON file.""" + file_path = Path(path) + if not file_path.is_file(): + raise WorkloadError(f"generation result {file_path!s}: file: not found") + try: + raw = json.loads(file_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise WorkloadError( + f"generation result {file_path!s}: JSON: invalid JSON: {exc}" + ) from exc + return GenerationRunResult.from_dict(raw, source=str(file_path)) + + +def write_generation_result( + result: GenerationRunResult, path: str | Path +) -> None: + """Atomically write a generation result JSON file.""" + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(output.suffix + ".tmp") + temporary.write_text( + json.dumps(result.as_dict(), indent=2) + "\n", encoding="utf-8" + ) + temporary.replace(output) + + +def compare_frame_outputs( + native_path: str | Path | None, + optimized_path: str | Path | None, + *, + policy: str = "byte_equal", + atol: float = 0.0, + rtol: float = 0.0, +) -> dict[str, Any]: + """Compare saved full-generation frame arrays under the workload policy.""" + if native_path is None or optimized_path is None: + return { + "passed": False, + "policy": policy, + "reason": "one or both frame artifacts are missing", + } + native_file = Path(native_path) + optimized_file = Path(optimized_path) + if not native_file.is_file() or not optimized_file.is_file(): + return { + "passed": False, + "policy": policy, + "reason": "one or both frame artifact paths do not exist", + } + + import numpy as np + + native = np.load(native_file, allow_pickle=False, mmap_mode="r") + optimized = np.load(optimized_file, allow_pickle=False, mmap_mode="r") + same_shape = native.shape == optimized.shape + result: dict[str, Any] = { + "policy": policy, + "native_shape": list(native.shape), + "optimized_shape": list(optimized.shape), + "native_dtype": str(native.dtype), + "optimized_dtype": str(optimized.dtype), + "same_shape": same_shape, + } + if policy == "frames_only": + result["passed"] = same_shape + result["reason"] = ( + "frame shapes match" if same_shape else "frame shapes differ" + ) + return result + if not same_shape: + result["passed"] = False + result["reason"] = "frame shapes differ" + return result + + if policy == "byte_equal": + passed = bool( + native.dtype == optimized.dtype and np.array_equal(native, optimized) + ) + result["passed"] = passed + result["reason"] = ( + "frame arrays are byte equal" + if passed + else "frame arrays are not byte equal" + ) + return result + if policy == "tolerance": + difference = np.abs(native - optimized) + result["max_abs_diff"] = ( + float(difference.max()) if difference.size else 0.0 + ) + result["mean_abs_diff"] = ( + float(difference.mean()) if difference.size else 0.0 + ) + passed = bool( + np.allclose( + native, + optimized, + atol=atol, + rtol=rtol, + equal_nan=False, + ) + ) + result.update( + { + "passed": passed, + "atol": atol, + "rtol": rtol, + "reason": ( + "frame arrays are within tolerance" + if passed + else "frame arrays exceed tolerance" + ), + } + ) + return result + return { + **result, + "passed": False, + "reason": f"unsupported parity policy {policy!r}", + } + + +def classify_end_to_end( + native: GenerationRunResult, + optimized: GenerationRunResult, + *, + min_speedup: float = 1.01, + max_peak_memory_regression: float = 0.05, +) -> dict[str, Any]: + """Classify a native-versus-optimized pair without claiming microbench wins.""" + if native.status != "ok" or optimized.status != "ok": + return { + "classification": "failed", + "reason": "one or both modes failed", + "native_status": native.status, + "optimized_status": optimized.status, + } + if ( + native.median_wall_seconds is None + or optimized.median_wall_seconds is None + or native.median_wall_seconds <= 0 + or optimized.median_wall_seconds <= 0 + ): + return { + "classification": "failed", + "reason": "missing or non-positive median wall times", + } + + speedup = native.median_wall_seconds / optimized.median_wall_seconds + native_mem = [m for m in native.peak_memory_mb if m is not None] + opt_mem = [m for m in optimized.peak_memory_mb if m is not None] + memory_regression = None + if native_mem and opt_mem: + native_peak = max(native_mem) + opt_peak = max(opt_mem) + if native_peak > 0: + memory_regression = (opt_peak - native_peak) / native_peak + + if memory_regression is not None and memory_regression > max_peak_memory_regression: + classification = "regressed" + reason = "peak memory regression exceeds threshold" + elif speedup >= min_speedup: + classification = "improved" + reason = "repeatable end-to-end wall-time improvement" + elif speedup <= (1.0 / min_speedup): + classification = "regressed" + reason = "end-to-end wall time regressed" + else: + classification = "neutral" + reason = "change within timing noise / below promotion threshold" + + return { + "classification": classification, + "reason": reason, + "native_median_wall_seconds": native.median_wall_seconds, + "optimized_median_wall_seconds": optimized.median_wall_seconds, + "end_to_end_speedup": speedup, + "min_end_to_end_speedup": min_speedup, + "peak_memory_regression": memory_regression, + "max_peak_memory_regression": max_peak_memory_regression, + } diff --git a/autokernel/workload/types.py b/autokernel/workload/types.py new file mode 100644 index 00000000..671abbb0 --- /dev/null +++ b/autokernel/workload/types.py @@ -0,0 +1,883 @@ +"""Versioned FastVideo workload manifest contract. + +A workload describes a representative generation job that MotionKernel and a +FastVideo launcher can execute without model-specific Python callables. It is +the shared input for baseline measurement, profiling, candidate validation, +and end-to-end native-versus-optimized comparisons. + +Manifests must never embed tensor values, credentials, weights, or generated +user content. Prompt text is allowed only as an explicit generation input; use +``prompt_file`` when the prompt should stay out of the manifest body. +""" + +from __future__ import annotations + +import json +import math +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +WORKLOAD_SCHEMA_VERSION = 1 + +_TOP_LEVEL_FIELDS = { + "schema_version", + "workload_id", + "description", + "model", + "task", + "prompt", + "prompt_file", + "sampling", + "runtime", + "measurement", + "parity", + "performance", + "mode_env", + "tags", +} +_MODEL_FIELDS = { + "model_id", + "revision", + "trust_remote_code", +} +_SAMPLING_FIELDS = { + "height", + "width", + "num_frames", + "num_inference_steps", + "guidance_scale", + "seed", + "fps", + "dtype", + "attention_backend", +} +_RUNTIME_FIELDS = { + "num_gpus", + "use_fsdp_inference", + "dit_cpu_offload", + "vae_cpu_offload", + "text_encoder_cpu_offload", + "image_encoder_cpu_offload", + "pin_cpu_memory", + "distributed_executor_backend", + "tp_size", + "sp_size", +} +_MEASUREMENT_FIELDS = { + "warmups", + "runs", + "save_frames", + "save_video", +} +_PARITY_FIELDS = { + "policy", + "atol", + "rtol", +} +_PERFORMANCE_FIELDS = { + "min_end_to_end_speedup", + "max_peak_memory_regression", +} +_MODE_ENV_FIELDS = { + "native", + "optimized", +} + +_TASKS = {"t2v", "i2v", "t2i", "i2i"} +_PARITY_POLICIES = {"byte_equal", "tolerance", "frames_only"} +_DTYPES = { + "float16", + "fp16", + "bfloat16", + "bf16", + "float32", + "fp32", +} +_WORKLOAD_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_FORBIDDEN_KEYS = { + "credential", + "credentials", + "data", + "password", + "secret", + "secrets", + "tensor_values", + "token", + "values", + "weights", + "activations", +} +FORBIDDEN_METADATA_KEYS = frozenset(_FORBIDDEN_KEYS) + + +class WorkloadError(ValueError): + """Raised when a workload manifest is malformed or unsafe.""" + + +def _fail(source: object, location: str, message: str) -> WorkloadError: + return WorkloadError(f"workload {source!r}: {location}: {message}") + + +def _mapping( + value: Any, + source: object, + location: str, + *, + non_empty: bool = False, +) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or (non_empty and not value): + qualifier = "non-empty " if non_empty else "" + raise _fail(source, location, f"must be a {qualifier}object") + for key in value: + if not isinstance(key, str) or not key: + raise _fail(source, location, "keys must be non-empty strings") + if key.lower() in _FORBIDDEN_KEYS: + raise _fail( + source, + f"{location}.{key}", + "content or secret fields are forbidden", + ) + return value + + +def _unknown_fields( + raw: Mapping[str, Any], + allowed: set[str], + source: object, + location: str, +) -> None: + unknown = sorted(set(raw) - allowed) + if unknown: + raise _fail(source, location, f"unknown field(s) {unknown}") + + +def _text(value: Any, source: object, location: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise _fail(source, location, "must be a non-empty string") + return value.strip() + + +def _optional_text( + value: Any, source: object, location: str +) -> str | None: + if value is None: + return None + return _text(value, source, location) + + +def _bool(value: Any, source: object, location: str, default: bool) -> bool: + if value is None: + return default + if not isinstance(value, bool): + raise _fail(source, location, "must be a bool") + return value + + +def _positive_int(value: Any, source: object, location: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise _fail(source, location, "must be a positive integer") + return value + + +def _non_negative_int(value: Any, source: object, location: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise _fail(source, location, "must be a non-negative integer") + return value + + +def _finite_number( + value: Any, + source: object, + location: str, + *, + minimum: float | None = None, + maximum: float | None = None, +) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise _fail(source, location, "must be a finite number") + number = float(value) + if not math.isfinite(number): + raise _fail(source, location, "must be a finite number") + if minimum is not None and number < minimum: + raise _fail(source, location, f"must be >= {minimum}") + if maximum is not None and number > maximum: + raise _fail(source, location, f"must be <= {maximum}") + return number + + +def _optional_positive_int( + value: Any, source: object, location: str +) -> int | None: + if value is None: + return None + return _positive_int(value, source, location) + + +def _string_map( + value: Any, source: object, location: str +) -> dict[str, str]: + raw = _mapping(value, source, location) + result: dict[str, str] = {} + for key, item in raw.items(): + result[key] = _text(item, source, f"{location}.{key}") + return result + + +@dataclass(frozen=True) +class ModelRef: + """FastVideo-resolvable model identity.""" + + model_id: str + revision: str | None = None + trust_remote_code: bool = False + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "ModelRef": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _MODEL_FIELDS, source, location) + return cls( + model_id=_text(raw.get("model_id"), source, f"{location}.model_id"), + revision=_optional_text( + raw.get("revision"), source, f"{location}.revision" + ), + trust_remote_code=_bool( + raw.get("trust_remote_code"), + source, + f"{location}.trust_remote_code", + False, + ), + ) + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = {"model_id": self.model_id} + if self.revision is not None: + payload["revision"] = self.revision + if self.trust_remote_code: + payload["trust_remote_code"] = True + return payload + + +@dataclass(frozen=True) +class SamplingSpec: + """Declarative generation parameters for a representative workload.""" + + height: int + width: int + num_frames: int + num_inference_steps: int + guidance_scale: float + seed: int + fps: int | None = None + dtype: str | None = None + attention_backend: str | None = None + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "SamplingSpec": + raw = _mapping(raw_value, source, location, non_empty=True) + _unknown_fields(raw, _SAMPLING_FIELDS, source, location) + dtype = _optional_text(raw.get("dtype"), source, f"{location}.dtype") + if dtype is not None and dtype.lower() not in _DTYPES: + raise _fail( + source, + f"{location}.dtype", + f"must be one of {sorted(_DTYPES)}", + ) + return cls( + height=_positive_int(raw.get("height"), source, f"{location}.height"), + width=_positive_int(raw.get("width"), source, f"{location}.width"), + num_frames=_positive_int( + raw.get("num_frames"), source, f"{location}.num_frames" + ), + num_inference_steps=_positive_int( + raw.get("num_inference_steps"), + source, + f"{location}.num_inference_steps", + ), + guidance_scale=_finite_number( + raw.get("guidance_scale"), + source, + f"{location}.guidance_scale", + minimum=0.0, + ), + seed=_non_negative_int(raw.get("seed"), source, f"{location}.seed"), + fps=_optional_positive_int(raw.get("fps"), source, f"{location}.fps"), + dtype=dtype.lower() if dtype is not None else None, + attention_backend=_optional_text( + raw.get("attention_backend"), + source, + f"{location}.attention_backend", + ), + ) + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "height": self.height, + "width": self.width, + "num_frames": self.num_frames, + "num_inference_steps": self.num_inference_steps, + "guidance_scale": self.guidance_scale, + "seed": self.seed, + } + if self.fps is not None: + payload["fps"] = self.fps + if self.dtype is not None: + payload["dtype"] = self.dtype + if self.attention_backend is not None: + payload["attention_backend"] = self.attention_backend + return payload + + +@dataclass(frozen=True) +class RuntimeSpec: + """Device and offload settings for the FastVideo generator process.""" + + num_gpus: int = 1 + use_fsdp_inference: bool = False + dit_cpu_offload: bool = False + vae_cpu_offload: bool = False + text_encoder_cpu_offload: bool = True + image_encoder_cpu_offload: bool = True + pin_cpu_memory: bool = False + distributed_executor_backend: str | None = None + tp_size: int | None = None + sp_size: int | None = None + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "RuntimeSpec": + if raw_value is None: + return cls() + raw = _mapping(raw_value, source, location) + _unknown_fields(raw, _RUNTIME_FIELDS, source, location) + return cls( + num_gpus=_positive_int( + raw.get("num_gpus", 1), source, f"{location}.num_gpus" + ), + use_fsdp_inference=_bool( + raw.get("use_fsdp_inference"), + source, + f"{location}.use_fsdp_inference", + False, + ), + dit_cpu_offload=_bool( + raw.get("dit_cpu_offload"), + source, + f"{location}.dit_cpu_offload", + False, + ), + vae_cpu_offload=_bool( + raw.get("vae_cpu_offload"), + source, + f"{location}.vae_cpu_offload", + False, + ), + text_encoder_cpu_offload=_bool( + raw.get("text_encoder_cpu_offload"), + source, + f"{location}.text_encoder_cpu_offload", + True, + ), + image_encoder_cpu_offload=_bool( + raw.get("image_encoder_cpu_offload"), + source, + f"{location}.image_encoder_cpu_offload", + True, + ), + pin_cpu_memory=_bool( + raw.get("pin_cpu_memory"), + source, + f"{location}.pin_cpu_memory", + False, + ), + distributed_executor_backend=_optional_text( + raw.get("distributed_executor_backend"), + source, + f"{location}.distributed_executor_backend", + ), + tp_size=_optional_positive_int( + raw.get("tp_size"), source, f"{location}.tp_size" + ), + sp_size=_optional_positive_int( + raw.get("sp_size"), source, f"{location}.sp_size" + ), + ) + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "num_gpus": self.num_gpus, + "use_fsdp_inference": self.use_fsdp_inference, + "dit_cpu_offload": self.dit_cpu_offload, + "vae_cpu_offload": self.vae_cpu_offload, + "text_encoder_cpu_offload": self.text_encoder_cpu_offload, + "image_encoder_cpu_offload": self.image_encoder_cpu_offload, + "pin_cpu_memory": self.pin_cpu_memory, + } + if self.distributed_executor_backend is not None: + payload["distributed_executor_backend"] = ( + self.distributed_executor_backend + ) + if self.tp_size is not None: + payload["tp_size"] = self.tp_size + if self.sp_size is not None: + payload["sp_size"] = self.sp_size + return payload + + +@dataclass(frozen=True) +class MeasurementSpec: + """Warmup/run counts and artifact capture preferences.""" + + warmups: int = 1 + runs: int = 2 + save_frames: bool = True + save_video: bool = False + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "MeasurementSpec": + if raw_value is None: + return cls() + raw = _mapping(raw_value, source, location) + _unknown_fields(raw, _MEASUREMENT_FIELDS, source, location) + return cls( + warmups=_non_negative_int( + raw.get("warmups", 1), source, f"{location}.warmups" + ), + runs=_positive_int(raw.get("runs", 2), source, f"{location}.runs"), + save_frames=_bool( + raw.get("save_frames"), + source, + f"{location}.save_frames", + True, + ), + save_video=_bool( + raw.get("save_video"), + source, + f"{location}.save_video", + False, + ), + ) + + def as_dict(self) -> dict[str, Any]: + return { + "warmups": self.warmups, + "runs": self.runs, + "save_frames": self.save_frames, + "save_video": self.save_video, + } + + +@dataclass(frozen=True) +class ParitySpec: + """Full-output parity policy for native-versus-optimized comparisons.""" + + policy: str = "byte_equal" + atol: float | None = None + rtol: float | None = None + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "ParitySpec": + if raw_value is None: + return cls() + raw = _mapping(raw_value, source, location) + _unknown_fields(raw, _PARITY_FIELDS, source, location) + policy = _text( + raw.get("policy", "byte_equal"), source, f"{location}.policy" + ) + if policy not in _PARITY_POLICIES: + raise _fail( + source, + f"{location}.policy", + f"must be one of {sorted(_PARITY_POLICIES)}", + ) + atol = raw.get("atol") + rtol = raw.get("rtol") + return cls( + policy=policy, + atol=( + None + if atol is None + else _finite_number( + atol, source, f"{location}.atol", minimum=0.0 + ) + ), + rtol=( + None + if rtol is None + else _finite_number( + rtol, source, f"{location}.rtol", minimum=0.0 + ) + ), + ) + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = {"policy": self.policy} + if self.atol is not None: + payload["atol"] = self.atol + if self.rtol is not None: + payload["rtol"] = self.rtol + return payload + + +@dataclass(frozen=True) +class PerformanceSpec: + """Promotion thresholds for end-to-end model-level evaluation.""" + + min_end_to_end_speedup: float = 1.01 + max_peak_memory_regression: float = 0.05 + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "PerformanceSpec": + if raw_value is None: + return cls() + raw = _mapping(raw_value, source, location) + _unknown_fields(raw, _PERFORMANCE_FIELDS, source, location) + return cls( + min_end_to_end_speedup=_finite_number( + raw.get("min_end_to_end_speedup", 1.01), + source, + f"{location}.min_end_to_end_speedup", + minimum=1.0, + ), + max_peak_memory_regression=_finite_number( + raw.get("max_peak_memory_regression", 0.05), + source, + f"{location}.max_peak_memory_regression", + minimum=0.0, + ), + ) + + def as_dict(self) -> dict[str, Any]: + return { + "min_end_to_end_speedup": self.min_end_to_end_speedup, + "max_peak_memory_regression": self.max_peak_memory_regression, + } + + +@dataclass(frozen=True) +class ModeEnvSpec: + """Optional environment variables applied per launcher mode. + + Values are environment-variable assignments only. Do not put Python + callables or model-specific code paths here. + """ + + native: Mapping[str, str] | None = None + optimized: Mapping[str, str] | None = None + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object, location: str + ) -> "ModeEnvSpec": + if raw_value is None: + return cls() + raw = _mapping(raw_value, source, location) + _unknown_fields(raw, _MODE_ENV_FIELDS, source, location) + native = raw.get("native") + optimized = raw.get("optimized") + return cls( + native=( + None + if native is None + else _string_map(native, source, f"{location}.native") + ), + optimized=( + None + if optimized is None + else _string_map(optimized, source, f"{location}.optimized") + ), + ) + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = {} + if self.native is not None: + payload["native"] = dict(self.native) + if self.optimized is not None: + payload["optimized"] = dict(self.optimized) + return payload + + def for_mode(self, mode: str) -> dict[str, str]: + if mode == "native": + return dict(self.native or {}) + if mode in {"optimized", "fused", "candidate"}: + return dict(self.optimized or {}) + raise WorkloadError(f"unknown launcher mode {mode!r}") + + +@dataclass(frozen=True) +class WorkloadManifest: + """Validated generation workload shared by MotionKernel and FastVideo.""" + + workload_id: str + model: ModelRef + task: str + sampling: SamplingSpec + prompt: str | None = None + prompt_file: str | None = None + description: str | None = None + runtime: RuntimeSpec | None = None + measurement: MeasurementSpec | None = None + parity: ParitySpec | None = None + performance: PerformanceSpec | None = None + mode_env: ModeEnvSpec | None = None + tags: tuple[str, ...] = () + source: str = "" + schema_version: int = WORKLOAD_SCHEMA_VERSION + + @classmethod + def from_dict( + cls, raw_value: Any, *, source: object = "" + ) -> "WorkloadManifest": + raw = _mapping(raw_value, source, "top level", non_empty=True) + _unknown_fields(raw, _TOP_LEVEL_FIELDS, source, "top level") + + version = raw.get("schema_version") + if isinstance(version, bool) or not isinstance(version, int): + raise _fail(source, "schema_version", "must be an integer") + if version != WORKLOAD_SCHEMA_VERSION: + raise _fail( + source, + "schema_version", + f"unsupported version {version}; expected {WORKLOAD_SCHEMA_VERSION}", + ) + + workload_id = _text(raw.get("workload_id"), source, "workload_id") + if not _WORKLOAD_ID_PATTERN.fullmatch(workload_id): + raise _fail( + source, + "workload_id", + "must be a short identifier of letters, digits, '.', '_', or '-'", + ) + + task = _text(raw.get("task"), source, "task").lower() + if task not in _TASKS: + raise _fail(source, "task", f"must be one of {sorted(_TASKS)}") + + prompt = _optional_text(raw.get("prompt"), source, "prompt") + prompt_file = _optional_text( + raw.get("prompt_file"), source, "prompt_file" + ) + if prompt is None and prompt_file is None: + raise _fail( + source, + "prompt", + "exactly one of prompt or prompt_file is required", + ) + if prompt is not None and prompt_file is not None: + raise _fail( + source, + "prompt", + "provide only one of prompt or prompt_file", + ) + + tags_raw = raw.get("tags", []) + if not isinstance(tags_raw, Sequence) or isinstance( + tags_raw, (str, bytes) + ): + raise _fail(source, "tags", "must be a list of strings") + tags = tuple( + _text(tag, source, f"tags[{index}]") + for index, tag in enumerate(tags_raw) + ) + if len(tags) != len(set(tags)): + raise _fail(source, "tags", "contains duplicates") + + return cls( + workload_id=workload_id, + model=ModelRef.from_dict( + raw.get("model"), source=source, location="model" + ), + task=task, + sampling=SamplingSpec.from_dict( + raw.get("sampling"), source=source, location="sampling" + ), + prompt=prompt, + prompt_file=prompt_file, + description=_optional_text( + raw.get("description"), source, "description" + ), + runtime=RuntimeSpec.from_dict( + raw.get("runtime"), source=source, location="runtime" + ), + measurement=MeasurementSpec.from_dict( + raw.get("measurement"), source=source, location="measurement" + ), + parity=ParitySpec.from_dict( + raw.get("parity"), source=source, location="parity" + ), + performance=PerformanceSpec.from_dict( + raw.get("performance"), + source=source, + location="performance", + ), + mode_env=ModeEnvSpec.from_dict( + raw.get("mode_env"), source=source, location="mode_env" + ), + tags=tags, + source=str(source), + schema_version=version, + ) + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "schema_version": self.schema_version, + "workload_id": self.workload_id, + "model": self.model.as_dict(), + "task": self.task, + "sampling": self.sampling.as_dict(), + "runtime": (self.runtime or RuntimeSpec()).as_dict(), + "measurement": (self.measurement or MeasurementSpec()).as_dict(), + "parity": (self.parity or ParitySpec()).as_dict(), + "performance": (self.performance or PerformanceSpec()).as_dict(), + } + if self.description is not None: + payload["description"] = self.description + if self.prompt is not None: + payload["prompt"] = self.prompt + if self.prompt_file is not None: + payload["prompt_file"] = self.prompt_file + mode_env = (self.mode_env or ModeEnvSpec()).as_dict() + if mode_env: + payload["mode_env"] = mode_env + if self.tags: + payload["tags"] = list(self.tags) + return payload + + def resolve_prompt(self, *, base_dir: str | Path | None = None) -> str: + """Return the prompt text, loading ``prompt_file`` when needed.""" + if self.prompt is not None: + return self.prompt + assert self.prompt_file is not None + path = Path(self.prompt_file) + if not path.is_absolute(): + root = Path(base_dir) if base_dir is not None else Path(self.source).parent + path = root / path + if not path.is_file(): + raise _fail(self.source, "prompt_file", f"not found: {path}") + text = path.read_text(encoding="utf-8").strip() + if not text: + raise _fail(self.source, "prompt_file", "must be non-empty") + return text + + def generation_request(self, *, base_dir: str | Path | None = None) -> dict[str, Any]: + """Build a FastVideo ``generate`` request dict from this workload.""" + sampling = self.sampling.as_dict() + # FastVideo SamplingConfig does not take dtype/attention_backend. + sampling.pop("dtype", None) + sampling.pop("attention_backend", None) + return { + "prompt": self.resolve_prompt(base_dir=base_dir), + "sampling": sampling, + "output": { + "save_video": (self.measurement or MeasurementSpec()).save_video, + "return_frames": (self.measurement or MeasurementSpec()).save_frames, + }, + } + + def generator_kwargs(self) -> dict[str, Any]: + """Keyword arguments for ``VideoGenerator.from_pretrained``.""" + runtime = self.runtime or RuntimeSpec() + kwargs: dict[str, Any] = { + "num_gpus": runtime.num_gpus, + "use_fsdp_inference": runtime.use_fsdp_inference, + "dit_cpu_offload": runtime.dit_cpu_offload, + "vae_cpu_offload": runtime.vae_cpu_offload, + "text_encoder_cpu_offload": runtime.text_encoder_cpu_offload, + "image_encoder_cpu_offload": runtime.image_encoder_cpu_offload, + "pin_cpu_memory": runtime.pin_cpu_memory, + } + if self.model.revision is not None: + kwargs["revision"] = self.model.revision + if self.model.trust_remote_code: + kwargs["trust_remote_code"] = True + if runtime.distributed_executor_backend is not None: + kwargs["distributed_executor_backend"] = ( + runtime.distributed_executor_backend + ) + if runtime.tp_size is not None: + kwargs["tp_size"] = runtime.tp_size + if runtime.sp_size is not None: + kwargs["sp_size"] = runtime.sp_size + return kwargs + + +def _load_raw_mapping(path: Path) -> Mapping[str, Any]: + text = path.read_text(encoding="utf-8") + suffix = path.suffix.lower() + if suffix in {".yaml", ".yml"}: + try: + import yaml # type: ignore + except ImportError as exc: # pragma: no cover - exercised in envs without PyYAML + raise WorkloadError( + f"workload {path!s}: YAML support requires PyYAML; " + "install motionkernel with PyYAML or use a .json manifest" + ) from exc + raw = yaml.safe_load(text) + elif suffix == ".json": + try: + raw = json.loads(text) + except json.JSONDecodeError as exc: + raise WorkloadError( + f"workload {path!s}: JSON: invalid JSON: {exc}" + ) from exc + else: + raise WorkloadError( + f"workload {path!s}: file: unsupported extension {suffix!r}; " + "use .yaml, .yml, or .json" + ) + if not isinstance(raw, Mapping): + raise WorkloadError(f"workload {path!s}: top level: must be an object") + return raw + + +def load_workload(path: str | Path) -> WorkloadManifest: + """Load and validate a workload manifest without importing torch.""" + file_path = Path(path) + if not file_path.is_file(): + raise WorkloadError(f"workload {file_path!s}: file: not found") + raw = _load_raw_mapping(file_path) + return WorkloadManifest.from_dict(raw, source=str(file_path)) + + +def dump_workload( + workload: WorkloadManifest, + path: str | Path, + *, + fmt: str | None = None, +) -> None: + """Write a workload manifest as YAML or JSON.""" + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + format_name = (fmt or output.suffix.lstrip(".")).lower() + payload = workload.as_dict() + if format_name in {"yaml", "yml"}: + try: + import yaml # type: ignore + except ImportError as exc: # pragma: no cover + raise WorkloadError( + "YAML support requires PyYAML" + ) from exc + text = yaml.safe_dump( + payload, + sort_keys=False, + default_flow_style=False, + ) + elif format_name == "json": + text = json.dumps(payload, indent=2) + "\n" + else: + raise WorkloadError(f"unsupported dump format {format_name!r}") + temporary = output.with_suffix(output.suffix + ".tmp") + temporary.write_text(text, encoding="utf-8") + temporary.replace(output) diff --git a/bench.py b/bench.py index f249e374..ccf570dd 100644 --- a/bench.py +++ b/bench.py @@ -11,9 +11,13 @@ Usage: uv run bench.py # benchmark kernel.py using its KERNEL_TYPE uv run bench.py --kernel matmul # force kernel type + uv run bench.py --spec path/spec.py:SPEC # benchmark an external KernelSpec uv run bench.py --quick # skip stages 3-5, bench only large size uv run bench.py --profile # emit torch profiler trace uv run bench.py --sizes large # benchmark only 'large' size + +Operation metadata (sizes, dtypes, tolerances, edge cases, FLOP/byte accounting) +lives in autokernel/specs/, not in this file. """ from __future__ import annotations @@ -26,12 +30,47 @@ import sys import time import traceback -from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Tuple +from dataclasses import asdict, dataclass +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple import torch import torch.nn.functional as F +# The package lives next to this script; make sure it is importable when bench.py +# is invoked from another working directory. +_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +if _SCRIPT_DIR not in sys.path: + sys.path.insert(0, _SCRIPT_DIR) + +from autokernel.specs import ( # noqa: E402 (path bootstrap must run first) + KernelRegistry, + KernelSpec, + SpecLoadError, + SpecNotFoundError, + SpecValidationError, + create_builtin_registry, + dtype_bytes, + resolve_spec, + resolve_torch_dtype, +) +from autokernel.verification import ( # noqa: E402 + TreeComparison, + check_backward, + check_compile, + collect_environment_metadata, + compare_deterministic, + compare_output_trees, + result_envelope, + tree_has_nan_or_inf, + write_result_atomic, +) +from autokernel.verification.corpus import ( # noqa: E402 + CorpusError, + load_shape_corpus, + validate_corpus_against_spec, + weighted_aggregate, +) + # --------------------------------------------------------------------------- # Timeout helper (cross-platform) # --------------------------------------------------------------------------- @@ -197,450 +236,180 @@ def detect_gpu() -> GPUSpec: # ========================================================================= -# 2. INPUT GENERATORS (deterministic via manual_seed) +# 2. OPERATION METADATA (from the KernelSpec registry) # ========================================================================= - -def gen_matmul_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict: - torch.manual_seed(seed) - M, N, K = size["M"], size["N"], size["K"] - A = torch.randn(M, K, device=device, dtype=dtype) - B = torch.randn(K, N, device=device, dtype=dtype) - return {"A": A, "B": B} +# Input generators, reference wiring, sizes, dtypes, tolerances, edge cases and +# FLOP/byte accounting are owned by autokernel/specs/. This file only translates +# canonical dtype names into torch dtypes at the runtime boundary. -def gen_softmax_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict: - torch.manual_seed(seed) - rows, cols = size["rows"], size["cols"] - x = torch.randn(rows, cols, device=device, dtype=dtype) - return {"x": x} +def _dtype_bytes(dtype: torch.dtype) -> int: + """Return byte-width for a dtype.""" + return torch.tensor([], dtype=dtype).element_size() -def gen_layernorm_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict: - torch.manual_seed(seed) - batch, dim = size["batch"], size["dim"] - x = torch.randn(batch, dim, device=device, dtype=dtype) - weight = torch.ones(dim, device=device, dtype=dtype) - bias = torch.zeros(dim, device=device, dtype=dtype) - return {"x": x, "weight": weight, "bias": bias} +def _spec_sizes(spec: KernelSpec) -> List[Tuple[str, Dict[str, int]]]: + """Ordered ``(label, size)`` pairs, as the harness has always consumed them.""" + return list(spec.size_items()) -def gen_flash_attention_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict: - torch.manual_seed(seed) - batch, heads, seq_len, head_dim = size["batch"], size["heads"], size["seq_len"], size["head_dim"] - Q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=dtype) - K = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=dtype) - V = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=dtype) - return {"Q": Q, "K": K, "V": V} +def _spec_dtypes(spec: KernelSpec) -> List[torch.dtype]: + """Declared dtypes as torch dtypes, in benchmark order.""" + return [resolve_torch_dtype(name) for name in spec.dtypes] -def gen_fused_mlp_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict: - torch.manual_seed(seed) - batch, dim, hidden = size["batch"], size["dim"], size["hidden"] - x = torch.randn(batch, dim, device=device, dtype=dtype) - w_gate = torch.randn(hidden, dim, device=device, dtype=dtype) * 0.02 - w_up = torch.randn(hidden, dim, device=device, dtype=dtype) * 0.02 - w_down = torch.randn(dim, hidden, device=device, dtype=dtype) * 0.02 - return {"x": x, "w_gate": w_gate, "w_up": w_up, "w_down": w_down} +def _spec_tolerances(spec: KernelSpec) -> Dict[torch.dtype, Dict[str, float]]: + """Tolerances keyed by torch dtype.""" + return { + resolve_torch_dtype(name): tol.as_dict() + for name, tol in spec.tolerances.items() + } -def gen_cross_entropy_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict: - torch.manual_seed(seed) - batch, vocab = size["batch"], size["vocab"] - logits = torch.randn(batch, vocab, device=device, dtype=dtype) - targets = torch.randint(0, vocab, (batch,), device=device, dtype=torch.long) - return {"logits": logits, "targets": targets} +def _spec_reference(spec: KernelSpec) -> Callable[[Dict[str, Any]], Any]: + """Adapt ``reference_fn(**inputs)`` to the harness' ``ref_fn(inputs)`` shape.""" + def ref_fn(inputs: Dict[str, Any]) -> Any: + return spec.reference_fn(**inputs) -def gen_rotary_embedding_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict: - torch.manual_seed(seed) - batch, heads, seq_len, head_dim = size["batch"], size["heads"], size["seq_len"], size["head_dim"] - x = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=dtype) - half_dim = head_dim // 2 - cos = torch.randn(seq_len, half_dim, device=device, dtype=dtype) - sin = torch.randn(seq_len, half_dim, device=device, dtype=dtype) - return {"x": x, "cos": cos, "sin": sin} + return ref_fn -def gen_rmsnorm_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict: - torch.manual_seed(seed) - M, N = size["M"], size["N"] - x = torch.randn(M, N, device=device, dtype=dtype) - weight = torch.randn(N, device=device, dtype=dtype) - return {"x": x, "weight": weight} +def _spec_bytes_fn(spec: KernelSpec) -> Callable[[Dict[str, int], torch.dtype], Any]: + """Adapt ``bytes_fn(size, dt_bytes)`` to a torch-dtype call site.""" + def bytes_fn(size: Dict[str, int], dtype: torch.dtype) -> Any: + return spec.bytes_fn(size, _dtype_bytes(dtype)) -def gen_reduce_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict: - torch.manual_seed(seed) - M, N = size["M"], size["N"] - x = torch.randn(M, N, device=device, dtype=dtype) - return {"x": x} + return bytes_fn -# ========================================================================= -# 3. REFERENCE WRAPPERS -# ========================================================================= -# Thin wrappers that call reference.py functions with the right dict keys. +def _spec_edge_cases(spec: KernelSpec) -> List[Tuple[str, Dict[str, int]]]: + """Edge-case ``(label, size)`` pairs for the shape-robustness stage.""" + return [(edge.name, dict(edge.size)) for edge in spec.edge_cases] -def _ref_matmul(inputs: dict) -> torch.Tensor: - import reference - return reference.matmul_ref(inputs["A"], inputs["B"]) -def _ref_softmax(inputs: dict) -> torch.Tensor: - import reference - return reference.softmax_ref(inputs["x"]) +def _legacy_config(spec: KernelSpec) -> Dict[str, Any]: + """Build the pre-registry ``KERNEL_CONFIGS`` entry for one specification.""" + return { + "test_sizes": _spec_sizes(spec), + "test_dtypes": _spec_dtypes(spec), + "tolerances": _spec_tolerances(spec), + "flops_fn": spec.flops_fn, + "bytes_fn": _spec_bytes_fn(spec), + "input_generator": spec.input_generator, + "reference_fn": _spec_reference(spec), + "edge_sizes": _spec_edge_cases(spec), + "spec": spec, + } -def _ref_layernorm(inputs: dict) -> torch.Tensor: - import reference - return reference.layernorm_ref(inputs["x"], inputs["weight"], inputs["bias"]) -def _ref_flash_attention(inputs: dict) -> torch.Tensor: - import reference - return reference.flash_attention_ref(inputs["Q"], inputs["K"], inputs["V"]) +def _build_legacy_configs() -> Dict[str, Dict[str, Any]]: + return {spec.name: _legacy_config(spec) for spec in create_builtin_registry()} -def _ref_fused_mlp(inputs: dict) -> torch.Tensor: - import reference - return reference.fused_mlp_ref(inputs["x"], inputs["w_gate"], inputs["w_up"], inputs["w_down"]) -def _ref_cross_entropy(inputs: dict) -> torch.Tensor: - import reference - return reference.cross_entropy_ref(inputs["logits"], inputs["targets"]) +#: DEPRECATED compatibility view of the old hard-coded configuration table. +#: Derived from the built-in registry; edit autokernel/specs/builtins.py instead. +#: New code should use ``autokernel.specs.create_builtin_registry()``. +KERNEL_CONFIGS: Dict[str, Dict[str, Any]] = _build_legacy_configs() -def _ref_rotary_embedding(inputs: dict) -> torch.Tensor: - import reference - return reference.rotary_embedding_ref(inputs["x"], inputs["cos"], inputs["sin"]) -def _ref_rmsnorm(inputs: dict) -> torch.Tensor: - import reference - return reference.rmsnorm_ref(inputs["x"], inputs["weight"]) +#: Device the harness allocates on. Overridden only by CPU tests; the +#: benchmark itself always measures on the GPU. +BENCH_DEVICE = "cuda" -def _ref_reduce(inputs: dict) -> torch.Tensor: - import reference - return reference.reduce_sum_ref(inputs["x"], dim=-1) +def resolve_operation_name( + spec_name: Optional[str], + kernel_arg: Optional[str], + declared_type: Optional[str], +) -> Optional[str]: + """Apply the operation-selection precedence. -# ========================================================================= -# 4. KERNEL CONFIGS -# ========================================================================= + 1. the name declared by an explicit ``--spec``; + 2. an explicit ``--kernel``; + 3. ``kernel.py::KERNEL_TYPE`` (the historical default). -def _dtype_bytes(dtype: torch.dtype) -> int: - """Return byte-width for a dtype.""" - return torch.tensor([], dtype=dtype).element_size() + Returns None when nothing selects an operation. + """ + return spec_name or kernel_arg or declared_type -KERNEL_CONFIGS: Dict[str, Dict[str, Any]] = { - # ----------------------------------------------------------------- - # MATMUL - # ----------------------------------------------------------------- - "matmul": { - "test_sizes": [ - ("tiny", {"M": 128, "N": 128, "K": 128}), - ("small", {"M": 512, "N": 512, "K": 512}), - ("medium", {"M": 1024, "N": 1024, "K": 1024}), - ("large", {"M": 2048, "N": 2048, "K": 2048}), - ("xlarge", {"M": 4096, "N": 4096, "K": 4096}), - ("tall", {"M": 8192, "N": 1024, "K": 1024}), - ("wide", {"M": 1024, "N": 8192, "K": 1024}), - ("deep_k", {"M": 1024, "N": 1024, "K": 8192}), - ("llm_qkv", {"M": 4096, "N": 4096, "K": 512}), - ("llm_mlp", {"M": 4096, "N": 11008, "K": 4096}), - ], - "test_dtypes": [torch.float16, torch.bfloat16, torch.float32], - "tolerances": { - torch.float16: {"atol": 1e-2, "rtol": 1e-2}, - torch.bfloat16: {"atol": 2e-2, "rtol": 2e-2}, - torch.float32: {"atol": 1e-4, "rtol": 1e-4}, - }, - "flops_fn": lambda s: 2 * s["M"] * s["N"] * s["K"], - "bytes_fn": lambda s, dt: (s["M"] * s["K"] + s["K"] * s["N"] + s["M"] * s["N"]) * _dtype_bytes(dt), - "input_generator": gen_matmul_inputs, - "reference_fn": _ref_matmul, - "edge_sizes": [ - ("edge_1023", {"M": 1023, "N": 1023, "K": 1023}), - ("edge_4097", {"M": 4097, "N": 4097, "K": 512}), - ("edge_1537", {"M": 1537, "N": 1537, "K": 1537}), - ], - }, - - # ----------------------------------------------------------------- - # SOFTMAX - # ----------------------------------------------------------------- - "softmax": { - "test_sizes": [ - ("tiny", {"rows": 32, "cols": 128}), - ("small", {"rows": 256, "cols": 512}), - ("medium", {"rows": 1024, "cols": 1024}), - ("large", {"rows": 4096, "cols": 4096}), - ("xlarge", {"rows": 8192, "cols": 8192}), - ("wide", {"rows": 1024, "cols": 32768}), - ("narrow", {"rows": 32768, "cols": 128}), - ("vocab", {"rows": 4096, "cols": 50257}), - ], - "test_dtypes": [torch.float16, torch.bfloat16, torch.float32], - "tolerances": { - torch.float16: {"atol": 1e-3, "rtol": 1e-3}, - torch.bfloat16: {"atol": 2e-3, "rtol": 2e-3}, - torch.float32: {"atol": 1e-5, "rtol": 1e-5}, - }, - "flops_fn": lambda s: 5 * s["rows"] * s["cols"], # exp + sub + sum + div + max - "bytes_fn": lambda s, dt: 2 * s["rows"] * s["cols"] * _dtype_bytes(dt), # read + write - "input_generator": gen_softmax_inputs, - "reference_fn": _ref_softmax, - "edge_sizes": [ - ("edge_1023", {"rows": 1023, "cols": 1023}), - ("edge_4097", {"rows": 4097, "cols": 4097}), - ("edge_50257", {"rows": 1024, "cols": 50257}), - ], - }, - - # ----------------------------------------------------------------- - # LAYERNORM - # ----------------------------------------------------------------- - "layernorm": { - "test_sizes": [ - ("tiny", {"batch": 32, "dim": 128}), - ("small", {"batch": 256, "dim": 512}), - ("medium", {"batch": 1024, "dim": 1024}), - ("large", {"batch": 4096, "dim": 2048}), - ("xlarge", {"batch": 8192, "dim": 4096}), - ("wide", {"batch": 1024, "dim": 8192}), - ("llm_7b", {"batch": 4096, "dim": 4096}), - ("llm_13b", {"batch": 4096, "dim": 5120}), - ], - "test_dtypes": [torch.float16, torch.bfloat16, torch.float32], - "tolerances": { - torch.float16: {"atol": 1e-3, "rtol": 1e-3}, - torch.bfloat16: {"atol": 2e-3, "rtol": 2e-3}, - torch.float32: {"atol": 1e-5, "rtol": 1e-5}, - }, - "flops_fn": lambda s: 8 * s["batch"] * s["dim"], # mean, var, norm, scale, shift - "bytes_fn": lambda s, dt: (2 * s["batch"] * s["dim"] + 2 * s["dim"]) * _dtype_bytes(dt), - "input_generator": gen_layernorm_inputs, - "reference_fn": _ref_layernorm, - "edge_sizes": [ - ("edge_1023", {"batch": 1023, "dim": 1023}), - ("edge_4097", {"batch": 4097, "dim": 4097}), - ], - }, - - # ----------------------------------------------------------------- - # FLASH ATTENTION - # ----------------------------------------------------------------- - "flash_attention": { - "test_sizes": [ - ("tiny", {"batch": 1, "heads": 4, "seq_len": 64, "head_dim": 64}), - ("small", {"batch": 2, "heads": 8, "seq_len": 256, "head_dim": 64}), - ("medium", {"batch": 2, "heads": 16, "seq_len": 512, "head_dim": 64}), - ("large", {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 64}), - ("xlarge", {"batch": 2, "heads": 32, "seq_len": 2048, "head_dim": 64}), - ("long", {"batch": 1, "heads": 32, "seq_len": 4096, "head_dim": 64}), - ("gqa", {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 128}), - ("llm_7b", {"batch": 1, "heads": 32, "seq_len": 2048, "head_dim": 128}), - ], - "test_dtypes": [torch.float16, torch.bfloat16], - "tolerances": { - torch.float16: {"atol": 1e-2, "rtol": 1e-2}, - torch.bfloat16: {"atol": 2e-2, "rtol": 2e-2}, - torch.float32: {"atol": 1e-4, "rtol": 1e-4}, - }, - # 4*N*S^2*D FLOPs (Q@K^T + softmax + attn@V) - "flops_fn": lambda s: 4 * s["batch"] * s["heads"] * (s["seq_len"] ** 2) * s["head_dim"], - "bytes_fn": lambda s, dt: 3 * s["batch"] * s["heads"] * s["seq_len"] * s["head_dim"] * _dtype_bytes(dt) + \ - s["batch"] * s["heads"] * s["seq_len"] * s["head_dim"] * _dtype_bytes(dt), - "input_generator": gen_flash_attention_inputs, - "reference_fn": _ref_flash_attention, - "edge_sizes": [ - ("edge_127", {"batch": 1, "heads": 8, "seq_len": 127, "head_dim": 64}), - ("edge_1023", {"batch": 1, "heads": 8, "seq_len": 1023, "head_dim": 64}), - ], - }, - - # ----------------------------------------------------------------- - # FUSED MLP (SwiGLU) - # ----------------------------------------------------------------- - "fused_mlp": { - "test_sizes": [ - ("tiny", {"batch": 32, "dim": 128, "hidden": 256}), - ("small", {"batch": 256, "dim": 512, "hidden": 1024}), - ("medium", {"batch": 1024, "dim": 1024, "hidden": 2048}), - ("large", {"batch": 2048, "dim": 2048, "hidden": 5504}), - ("xlarge", {"batch": 4096, "dim": 4096, "hidden": 11008}), - ("llm_7b", {"batch": 2048, "dim": 4096, "hidden": 11008}), - ("llm_13b", {"batch": 2048, "dim": 5120, "hidden": 13824}), - ], - "test_dtypes": [torch.float16, torch.bfloat16, torch.float32], - "tolerances": { - torch.float16: {"atol": 1e-2, "rtol": 1e-2}, - torch.bfloat16: {"atol": 2e-2, "rtol": 2e-2}, - torch.float32: {"atol": 1e-4, "rtol": 1e-4}, - }, - # gate_proj + up_proj + silu + mul + down_proj - "flops_fn": lambda s: 2 * s["batch"] * s["dim"] * s["hidden"] * 3, - "bytes_fn": lambda s, dt: (s["batch"] * s["dim"] + s["hidden"] * s["dim"] * 3 + s["batch"] * s["dim"]) * _dtype_bytes(dt), - "input_generator": gen_fused_mlp_inputs, - "reference_fn": _ref_fused_mlp, - "edge_sizes": [ - ("edge_1023", {"batch": 1023, "dim": 1024, "hidden": 2048}), - ("edge_4097", {"batch": 4097, "dim": 512, "hidden": 1024}), - ], - }, - - # ----------------------------------------------------------------- - # CROSS ENTROPY - # ----------------------------------------------------------------- - "cross_entropy": { - "test_sizes": [ - ("tiny", {"batch": 32, "vocab": 256}), - ("small", {"batch": 256, "vocab": 1024}), - ("medium", {"batch": 1024, "vocab": 4096}), - ("large", {"batch": 4096, "vocab": 32000}), - ("xlarge", {"batch": 8192, "vocab": 50257}), - ("llama", {"batch": 4096, "vocab": 32000}), - ("gpt2", {"batch": 4096, "vocab": 50257}), - ], - "test_dtypes": [torch.float16, torch.bfloat16, torch.float32], - "tolerances": { - torch.float16: {"atol": 1e-2, "rtol": 1e-2}, - torch.bfloat16: {"atol": 2e-2, "rtol": 2e-2}, - torch.float32: {"atol": 1e-5, "rtol": 1e-5}, - }, - # log_softmax + nll - "flops_fn": lambda s: 4 * s["batch"] * s["vocab"], - "bytes_fn": lambda s, dt: (s["batch"] * s["vocab"] + s["batch"]) * _dtype_bytes(dt), - "input_generator": gen_cross_entropy_inputs, - "reference_fn": _ref_cross_entropy, - "edge_sizes": [ - ("edge_1023", {"batch": 1023, "vocab": 32000}), - ("edge_50257", {"batch": 4096, "vocab": 50257}), - ], - }, - - # ----------------------------------------------------------------- - # ROTARY EMBEDDING - # ----------------------------------------------------------------- - "rotary_embedding": { - "test_sizes": [ - ("tiny", {"batch": 1, "heads": 4, "seq_len": 64, "head_dim": 64}), - ("small", {"batch": 2, "heads": 8, "seq_len": 256, "head_dim": 64}), - ("medium", {"batch": 2, "heads": 16, "seq_len": 512, "head_dim": 64}), - ("large", {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 128}), - ("xlarge", {"batch": 2, "heads": 32, "seq_len": 2048, "head_dim": 128}), - ("llm_7b", {"batch": 1, "heads": 32, "seq_len": 2048, "head_dim": 128}), - ("llm_13b", {"batch": 1, "heads": 40, "seq_len": 2048, "head_dim": 128}), - ], - "test_dtypes": [torch.float16, torch.bfloat16, torch.float32], - "tolerances": { - torch.float16: {"atol": 1e-3, "rtol": 1e-3}, - torch.bfloat16: {"atol": 2e-3, "rtol": 2e-3}, - torch.float32: {"atol": 1e-5, "rtol": 1e-5}, - }, - # mul + add per element, x2 (cos and sin parts) - "flops_fn": lambda s: 6 * s["batch"] * s["heads"] * s["seq_len"] * s["head_dim"], - "bytes_fn": lambda s, dt: (s["batch"] * s["heads"] * s["seq_len"] * s["head_dim"] * 2 + - s["seq_len"] * s["head_dim"]) * _dtype_bytes(dt), - "input_generator": gen_rotary_embedding_inputs, - "reference_fn": _ref_rotary_embedding, - "edge_sizes": [ - ("edge_127", {"batch": 1, "heads": 8, "seq_len": 127, "head_dim": 64}), - ("edge_1023", {"batch": 1, "heads": 8, "seq_len": 1023, "head_dim": 128}), - ], - }, - - # ----------------------------------------------------------------- - # RMSNORM - # ----------------------------------------------------------------- - "rmsnorm": { - "test_sizes": [ - ("small", {"M": 1024, "N": 768}), - ("medium", {"M": 4096, "N": 1024}), - ("large", {"M": 4096, "N": 4096}), - ("llama", {"M": 2048, "N": 4096}), - ], - "test_dtypes": [torch.float16, torch.bfloat16], - "tolerances": { - torch.float16: {"atol": 1e-2, "rtol": 1e-2}, - torch.bfloat16: {"atol": 1e-1, "rtol": 5e-2}, - }, - "flops_fn": lambda s: 6 * s["M"] * s["N"], # square, mean, sqrt, div, mul - "bytes_fn": lambda s, dt: (2 * s["M"] * s["N"] + s["N"]) * torch.tensor([], dtype=dt).element_size(), - "input_generator": gen_rmsnorm_inputs, - "reference_fn": _ref_rmsnorm, - "edge_sizes": [ - ("edge_1023", {"M": 1023, "N": 768}), - ("edge_4097", {"M": 4097, "N": 1024}), - ], - }, - - # ----------------------------------------------------------------- - # REDUCE (sum along last dim) - # ----------------------------------------------------------------- - "reduce": { - "test_sizes": [ - ("small", {"M": 1024, "N": 1024}), - ("medium", {"M": 4096, "N": 4096}), - ("large", {"M": 8192, "N": 8192}), - ("wide", {"M": 1024, "N": 32768}), - ], - "test_dtypes": [torch.float16, torch.bfloat16], - "tolerances": { - torch.float16: {"atol": 1e-2, "rtol": 1e-2}, - torch.bfloat16: {"atol": 1e-1, "rtol": 5e-2}, - }, - "flops_fn": lambda s: s["M"] * s["N"], # N additions per row - "bytes_fn": lambda s, dt: (s["M"] * s["N"] + s["M"]) * torch.tensor([], dtype=dt).element_size(), - "input_generator": gen_reduce_inputs, - "reference_fn": _ref_reduce, - "edge_sizes": [ - ("edge_1023", {"M": 1023, "N": 1024}), - ("edge_4097", {"M": 4096, "N": 4097}), - ], - }, -} +def _get_spec_or_exit(registry: KernelRegistry, kernel_type: str) -> KernelSpec: + """Fetch a spec from the registry, preserving the CLI failure contract.""" + try: + return registry.get(kernel_type) + except SpecNotFoundError: + print(f"\nERROR: Unknown kernel type '{kernel_type}'") + print(f" Available: {', '.join(registry.list_names())}") + print("\ncorrectness: FAIL") + print("throughput_tflops: 0.000") + sys.exit(1) -# ========================================================================= -# 5. CORRECTNESS TESTING (5 stages) -# ========================================================================= +def _load_validated_corpus(args: argparse.Namespace, spec: KernelSpec): + """Load and validate ``--shape-corpus`` against the selected spec. -def _compare(output: torch.Tensor, expected: torch.Tensor, atol: float, rtol: float) -> dict: - """Compare two tensors and return statistics.""" - if output.shape != expected.shape: - return { - "match": False, - "reason": f"shape mismatch: {output.shape} vs {expected.shape}", - "max_abs_error": float("inf"), - "mean_abs_error": float("inf"), - "pct_within_tol": 0.0, - } + Returns the validated cases, or None when no corpus was requested. Exits + with the greppable failure contract on any corpus error. Runs before the + candidate module is imported and before any GPU allocation, so a malformed + corpus never executes candidate code or touches the device. + """ + if not args.shape_corpus: + return None + try: + corpus = load_shape_corpus(args.shape_corpus) + validate_corpus_against_spec(corpus, spec) + except CorpusError as e: + print(f"\nERROR: {e}") + print("\ncorrectness: FAIL") + print("throughput_tflops: 0.000") + sys.exit(1) + mode = "corpus-only" if args.shape_corpus_only else "append" + print(f"shape_corpus: {args.shape_corpus} ({len(corpus.cases)} cases, mode={mode})") + return corpus.cases - # Cast both to float32 for comparison - out_f = output.float() - exp_f = expected.float() - abs_diff = (out_f - exp_f).abs() - max_abs = abs_diff.max().item() - mean_abs = abs_diff.mean().item() +# ========================================================================= +# 3. CORRECTNESS TESTING (5 stages) +# ========================================================================= - # Percentage of elements within tolerance - within = (abs_diff <= atol + rtol * exp_f.abs()).float().mean().item() * 100.0 +def _compare_outputs(output: Any, expected: Any, spec: KernelSpec, *, relax: float = 1.0) -> TreeComparison: + """Compare candidate and reference output trees using the spec's policy. + + Single-tensor outputs behave exactly as the historical ``_compare`` did: + the tolerance declared for the benchmark dtype is applied, and the + failure reason is unchanged. Structured outputs are compared leaf by leaf + with stable diagnostic paths. + """ + return compare_output_trees( + output, + expected, + spec.tolerances, + output_spec=spec.output_spec, + relax=relax, + ) - match = torch.allclose(out_f, exp_f, atol=atol, rtol=rtol) - return { - "match": match, - "reason": "" if match else f"max_abs_error={max_abs:.6e} exceeds tol(atol={atol}, rtol={rtol})", - "max_abs_error": max_abs, - "mean_abs_error": mean_abs, - "pct_within_tol": within, - } +def _record_leaves(records: List[Dict[str, Any]], stage: str, case: str, cmp: TreeComparison) -> None: + """Collect per-leaf comparison details for the structured result artifact.""" + for leaf in cmp.leaf_records(): + records.append({"stage": stage, "case": case, **leaf}) -def _has_nan_inf(t: torch.Tensor) -> bool: - """Check for NaN or Inf.""" - return bool(torch.isnan(t).any().item() or torch.isinf(t).any().item()) +def run_correctness( + kernel_fn: Callable, + spec: KernelSpec, + quick: bool = False, + corpus_cases: Optional[Sequence] = None, + corpus_only: bool = False, +) -> dict: + """Run all correctness stages. Returns dict with results. -def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> dict: - """Run all correctness stages. Returns dict with results.""" - device = "cuda" + ``corpus_cases`` appends validated production shapes to the stage-2 shape + sweep; ``corpus_only`` replaces the built-in size/dtype sweep with them. + """ + device = BENCH_DEVICE results = { "smoke_test": "SKIP", "shape_sweep": "SKIP", @@ -650,13 +419,13 @@ def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> d "correctness": "FAIL", } details = [] + leaf_records: List[Dict[str, Any]] = [] all_pass = True - gen_fn = config["input_generator"] - ref_fn = config["reference_fn"] - sizes = config["test_sizes"] - dtypes = config["test_dtypes"] - tols = config["tolerances"] + gen_fn = spec.input_generator + ref_fn = _spec_reference(spec) + sizes = _spec_sizes(spec) + dtypes = _spec_dtypes(spec) # ------------------------------------------------------------------ # Stage 1: SMOKE TEST -- tiny input, tight tolerance @@ -671,22 +440,22 @@ def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> d with _Timeout(30): output = kernel_fn(**inputs) - if _has_nan_inf(output): + if tree_has_nan_or_inf(output): results["smoke_test"] = "FAIL" - details.append(f" smoke: NaN/Inf in output") + details.append(" smoke: NaN/Inf in output") all_pass = False - print(f" FAIL: NaN/Inf in output") + print(" FAIL: NaN/Inf in output") else: - tol = tols.get(dtype0, {"atol": 1e-2, "rtol": 1e-2}) - cmp = _compare(output, expected, **tol) - if cmp["match"]: + cmp = _compare_outputs(output, expected, spec) + _record_leaves(leaf_records, "smoke", tiny_label, cmp) + if cmp.match: results["smoke_test"] = "PASS" - print(f" PASS (max_abs_error={cmp['max_abs_error']:.6e})") + print(f" PASS (max_abs_error={cmp.worst_abs_error:.6e})") else: results["smoke_test"] = "FAIL" - details.append(f" smoke: {cmp['reason']}") + details.append(f" smoke: {cmp.reason}") all_pass = False - print(f" FAIL: {cmp['reason']}") + print(f" FAIL: {cmp.reason}") except BenchTimeoutError: results["smoke_test"] = "FAIL" details.append(" smoke: TIMEOUT") @@ -707,7 +476,8 @@ def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> d if results["smoke_test"] == "FAIL": results["correctness"] = "FAIL" results["details"] = details - print(f"\ncorrectness: FAIL (smoke test failed, aborting remaining stages)") + results["leaf_details"] = leaf_records + print("\ncorrectness: FAIL (smoke test failed, aborting remaining stages)") return results # ------------------------------------------------------------------ @@ -720,54 +490,69 @@ def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> d worst_error = 0.0 worst_case = "" - for label, sz in sizes: - for dtype in dtypes: - sweep_count += 1 - try: - inputs = gen_fn(sz, dtype, device, seed=42) - expected = ref_fn(inputs) - with _Timeout(30): - output = kernel_fn(**inputs) - - if _has_nan_inf(output): - sweep_pass = False - sweep_fail_count += 1 - details.append(f" sweep {label}/{dtype}: NaN/Inf") - print(f" FAIL: {label} {dtype} -> NaN/Inf") - continue - - tol = tols.get(dtype, {"atol": 1e-2, "rtol": 1e-2}) - cmp = _compare(output, expected, **tol) - - if cmp["max_abs_error"] > worst_error: - worst_error = cmp["max_abs_error"] - worst_case = f"{label}/{dtype}" - - if not cmp["match"]: - sweep_pass = False - sweep_fail_count += 1 - details.append(f" sweep {label}/{dtype}: {cmp['reason']}") - print(f" FAIL: {label} {dtype} -> {cmp['reason']}") - else: - print(f" PASS: {label} {dtype} (max_err={cmp['max_abs_error']:.2e}, within_tol={cmp['pct_within_tol']:.1f}%)") + sweep_configs: List[Tuple[str, Dict[str, int], torch.dtype]] = [] + if not corpus_only: + for label, sz in sizes: + for dtype in dtypes: + sweep_configs.append((label, sz, dtype)) + if corpus_cases: + for case in corpus_cases: + case_dtype = resolve_torch_dtype(case.dtype) if case.dtype else dtypes[0] + sweep_configs.append((case.name, dict(case.size), case_dtype)) + + for label, sz, dtype in sweep_configs: + sweep_count += 1 + try: + inputs = gen_fn(sz, dtype, device, seed=42) + expected = ref_fn(inputs) + with _Timeout(30): + output = kernel_fn(**inputs) - except torch.cuda.OutOfMemoryError: - # OOM on larger sizes is acceptable -- just skip - print(f" SKIP: {label} {dtype} -> OOM") - torch.cuda.empty_cache() - continue - except BenchTimeoutError: + if tree_has_nan_or_inf(output): sweep_pass = False sweep_fail_count += 1 - details.append(f" sweep {label}/{dtype}: TIMEOUT") - print(f" FAIL: {label} {dtype} -> TIMEOUT") - except Exception as e: + details.append(f" sweep {label}/{dtype}: NaN/Inf") + print(f" FAIL: {label} {dtype} -> NaN/Inf") + continue + + cmp = _compare_outputs(output, expected, spec) + _record_leaves(leaf_records, "sweep", f"{label}/{dtype}", cmp) + + if cmp.worst_abs_error > worst_error: + worst_error = cmp.worst_abs_error + worst_case = f"{label}/{dtype}" + + if not cmp.match: sweep_pass = False sweep_fail_count += 1 - details.append(f" sweep {label}/{dtype}: {type(e).__name__}: {e}") - print(f" FAIL: {label} {dtype} -> {type(e).__name__}: {e}") - finally: - torch.cuda.empty_cache() + details.append(f" sweep {label}/{dtype}: {cmp.reason}") + print(f" FAIL: {label} {dtype} -> {cmp.reason}") + else: + print( + f" PASS: {label} {dtype} " + f"(max_err={cmp.worst_abs_error:.2e}, " + f"within_tol={cmp.worst_pct_within_tol:.1f}%)" + ) + + except torch.cuda.OutOfMemoryError: + # OOM on larger sizes is acceptable -- just skip + print(f" SKIP: {label} {dtype} -> OOM") + torch.cuda.empty_cache() + continue + except BenchTimeoutError: + sweep_pass = False + sweep_fail_count += 1 + details.append(f" sweep {label}/{dtype}: TIMEOUT") + print(f" FAIL: {label} {dtype} -> TIMEOUT") + except Exception as e: + sweep_pass = False + sweep_fail_count += 1 + details.append( + f" sweep {label}/{dtype}: {type(e).__name__}: {e}" + ) + print(f" FAIL: {label} {dtype} -> {type(e).__name__}: {e}") + finally: + torch.cuda.empty_cache() if sweep_pass: results["shape_sweep"] = f"PASS ({sweep_count} configs, worst_err={worst_error:.2e} at {worst_case})" @@ -786,6 +571,7 @@ def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> d results["edge_cases"] = "SKIP (quick mode)" results["correctness"] = "PASS" if all_pass else "FAIL" results["details"] = details + results["leaf_details"] = leaf_records print(f"\ncorrectness: {results['correctness']} (quick mode: stages 3-5 skipped)") return results @@ -830,25 +616,23 @@ def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> d with _Timeout(30): output = kernel_fn(**transformed) - if _has_nan_inf(output) and not _has_nan_inf(expected): + if tree_has_nan_or_inf(output) and not tree_has_nan_or_inf(expected): stability_pass = False details.append(f" stability {case_name}: NaN/Inf (reference is clean)") print(f" FAIL: {case_name} -> NaN/Inf (reference is clean)") - elif _has_nan_inf(output) and _has_nan_inf(expected): + elif tree_has_nan_or_inf(output) and tree_has_nan_or_inf(expected): # Both have NaN/Inf -- acceptable (e.g. overflow in near_max) print(f" PASS: {case_name} -> both have NaN/Inf (expected overflow)") else: - tol = tols.get(stab_dtype, {"atol": 1e-2, "rtol": 1e-2}) # Relax tolerances for adversarial inputs - relaxed_atol = tol["atol"] * 10 - relaxed_rtol = tol["rtol"] * 10 - cmp = _compare(output, expected, atol=relaxed_atol, rtol=relaxed_rtol) - if cmp["match"]: - print(f" PASS: {case_name} (max_err={cmp['max_abs_error']:.2e})") + cmp = _compare_outputs(output, expected, spec, relax=10.0) + _record_leaves(leaf_records, "stability", case_name, cmp) + if cmp.match: + print(f" PASS: {case_name} (max_err={cmp.worst_abs_error:.2e})") else: stability_pass = False - details.append(f" stability {case_name}: {cmp['reason']}") - print(f" FAIL: {case_name} -> {cmp['reason']}") + details.append(f" stability {case_name}: {cmp.reason}") + print(f" FAIL: {case_name} -> {cmp.reason}") except torch.cuda.OutOfMemoryError: print(f" SKIP: {case_name} -> OOM") @@ -888,11 +672,18 @@ def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> d outputs.append(out_i) for i in range(1, 3): - if not torch.equal(outputs[0], outputs[i]): + cmp = compare_deterministic(outputs[0], outputs[i], output_spec=spec.output_spec) + _record_leaves(leaf_records, "determinism", f"run_0_vs_{i}", cmp) + if not cmp.match: determinism_pass = False - diff = (outputs[0].float() - outputs[i].float()).abs() - details.append(f" determinism: run 0 vs run {i} differ (max_diff={diff.max().item():.6e})") - print(f" FAIL: run 0 vs run {i} differ (max_diff={diff.max().item():.6e})") + failure = cmp.first_failure() + if failure is not None and failure.max_abs_error is not None: + where = "" if failure.path == "output" else f" at {failure.path}" + message = f"run 0 vs run {i} differ{where} (max_diff={failure.max_abs_error:.6e})" + else: + message = f"run 0 vs run {i} differ ({cmp.reason})" + details.append(f" determinism: {message}") + print(f" FAIL: {message}") if determinism_pass: print(" PASS: 3 runs are bitwise identical") @@ -913,46 +704,52 @@ def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> d # ------------------------------------------------------------------ print("\n--- Stage 5: Edge Cases ---") edge_pass = True - edge_sizes = config.get("edge_sizes", []) - if not edge_sizes: + edge_cases = spec.edge_cases + if not edge_cases: results["edge_cases"] = "SKIP (no edge sizes defined)" print(" SKIP: no edge sizes defined") else: - for label, sz in edge_sizes: - for dtype in dtypes[:1]: # test with first dtype only for speed - try: - inputs = gen_fn(sz, dtype, device, seed=42) - expected = ref_fn(inputs) - with _Timeout(30): - output = kernel_fn(**inputs) - - if _has_nan_inf(output) and not _has_nan_inf(expected): - edge_pass = False - details.append(f" edge {label}: NaN/Inf") - print(f" FAIL: {label} -> NaN/Inf") - else: - tol = tols.get(dtype, {"atol": 1e-2, "rtol": 1e-2}) - cmp = _compare(output, expected, **tol) - if cmp["match"]: - print(f" PASS: {label} (max_err={cmp['max_abs_error']:.2e})") - else: - edge_pass = False - details.append(f" edge {label}: {cmp['reason']}") - print(f" FAIL: {label} -> {cmp['reason']}") - - except torch.cuda.OutOfMemoryError: - print(f" SKIP: {label} -> OOM") - torch.cuda.empty_cache() - except BenchTimeoutError: - edge_pass = False - details.append(f" edge {label}: TIMEOUT") - print(f" FAIL: {label} -> TIMEOUT") - except Exception as e: + for edge in edge_cases: + label = edge.name + sz = dict(edge.size) + # An edge case may pin its own dtype; otherwise use the primary + # dtype only, for speed. + dtype = resolve_torch_dtype(edge.dtype) if edge.dtype else dtypes[0] + try: + inputs = gen_fn(sz, dtype, device, seed=edge.seed) + if edge.input_transform is not None: + inputs = edge.input_transform(inputs) + expected = ref_fn(inputs) + with _Timeout(30): + output = kernel_fn(**inputs) + + if tree_has_nan_or_inf(output) and not tree_has_nan_or_inf(expected): edge_pass = False - details.append(f" edge {label}: {type(e).__name__}: {e}") - print(f" FAIL: {label} -> {type(e).__name__}: {e}") - finally: - torch.cuda.empty_cache() + details.append(f" edge {label}: NaN/Inf") + print(f" FAIL: {label} -> NaN/Inf") + else: + cmp = _compare_outputs(output, expected, spec) + _record_leaves(leaf_records, "edge", label, cmp) + if cmp.match: + print(f" PASS: {label} (max_err={cmp.worst_abs_error:.2e})") + else: + edge_pass = False + details.append(f" edge {label}: {cmp.reason}") + print(f" FAIL: {label} -> {cmp.reason}") + + except torch.cuda.OutOfMemoryError: + print(f" SKIP: {label} -> OOM") + torch.cuda.empty_cache() + except BenchTimeoutError: + edge_pass = False + details.append(f" edge {label}: TIMEOUT") + print(f" FAIL: {label} -> TIMEOUT") + except Exception as e: + edge_pass = False + details.append(f" edge {label}: {type(e).__name__}: {e}") + print(f" FAIL: {label} -> {type(e).__name__}: {e}") + finally: + torch.cuda.empty_cache() results["edge_cases"] = "PASS" if edge_pass else "FAIL" if not edge_pass: @@ -962,12 +759,45 @@ def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> d # Final verdict results["correctness"] = "PASS" if all_pass else "FAIL" results["details"] = details + results["leaf_details"] = leaf_records print(f"\ncorrectness: {results['correctness']}") return results +def run_backward_check(kernel_fn: Callable, spec: KernelSpec) -> dict: + """Opt-in gradient verification. Prints the greppable verdict line and + returns the structured report.""" + print("\n=== BACKWARD CORRECTNESS ===") + report = check_backward(kernel_fn, spec, device=BENCH_DEVICE) + if report.status == "PASS": + print(f" upstream outputs: {', '.join(report.output_paths)}") + for record in report.gradients: + print(f" grad[{record.input_name}]: match " + f"(max_err={record.max_abs_error:.2e}, mean_err={record.mean_abs_error:.2e})") + else: + print(f" FAIL: {report.reason}") + for record in report.gradients: + if record.status != "match": + print(f" grad[{record.input_name}]: {record.status}: {record.reason}") + print(f"BACKWARD_CORRECTNESS: {report.status}") + return report.as_dict() + + +def run_compile_check(kernel_fn: Callable, spec: KernelSpec) -> dict: + """Run compile verification outside all timed performance regions.""" + print("\n=== COMPILE CORRECTNESS ===") + report = check_compile(kernel_fn, spec, device=BENCH_DEVICE) + for case in report.cases: + detail = f": {case.reason}" if case.reason else "" + print(f" {case.label}: {case.status}{detail}") + if report.reason and report.status != "PASS": + print(f" {report.reason}") + print(f"COMPILE_CORRECTNESS: {report.status}") + return report.as_dict() + + # ========================================================================= -# 6. PERFORMANCE BENCHMARKING +# 4. PERFORMANCE BENCHMARKING # ========================================================================= def _do_bench(fn: Callable, warmup: int = 25, rep: int = 100) -> float: @@ -1000,20 +830,29 @@ def _do_bench(fn: Callable, warmup: int = 25, rep: int = 100) -> float: return times[len(times) // 2] # median -def run_performance(kernel_fn: Callable, config: dict, gpu: GPUSpec, - sizes_filter: str = "all") -> dict: - """Run performance benchmarks. Returns dict with metrics.""" - device = "cuda" - gen_fn = config["input_generator"] - ref_fn = config["reference_fn"] - flops_fn = config["flops_fn"] - bytes_fn = config["bytes_fn"] - dtypes = config["test_dtypes"] +def run_performance(kernel_fn: Callable, spec: KernelSpec, gpu: GPUSpec, + sizes_filter: str = "all", + corpus_cases: Optional[Sequence] = None, + corpus_only: bool = False) -> dict: + """Run performance benchmarks. Returns dict with metrics. + + ``corpus_cases`` appends production shapes (each benchmarked once; their + weights feed aggregate reporting, not repetition). ``corpus_only`` + benchmarks only those shapes. + """ + device = BENCH_DEVICE + gen_fn = spec.input_generator + ref_fn = _spec_reference(spec) + flops_fn = spec.flops_fn + bytes_fn = _spec_bytes_fn(spec) + dtypes = _spec_dtypes(spec) # Select benchmark size - sizes = config["test_sizes"] + sizes = _spec_sizes(spec) bench_sizes = [] - if sizes_filter == "all": + if corpus_only: + bench_sizes = [] + elif sizes_filter == "all": bench_sizes = sizes else: for label, sz in sizes: @@ -1042,10 +881,21 @@ def run_performance(kernel_fn: Callable, config: dict, gpu: GPUSpec, dtype = dtypes[0] # primary dtype for benchmarking + # (label, size, dtype, weight, source) benchmark configurations. Built-in + # sizes carry weight 1; corpus cases carry their declared weight, which is + # used only for aggregate reporting, never to repeat benchmark loops. + bench_configs: List[Tuple[str, Dict[str, int], torch.dtype, int, str]] = [ + (label, sz, dtype, 1, "builtin") for label, sz in bench_sizes + ] + if corpus_cases: + for case in corpus_cases: + case_dtype = resolve_torch_dtype(case.dtype) if case.dtype else dtype + bench_configs.append((case.name, dict(case.size), case_dtype, case.weight, "corpus")) + all_results = [] primary_result = None - for label, sz in bench_sizes: + for label, sz, dtype, weight, source in bench_configs: print(f"\n Benchmarking: {label} ...") try: flops = flops_fn(sz) @@ -1087,6 +937,8 @@ def run_performance(kernel_fn: Callable, config: dict, gpu: GPUSpec, "label": label, "size": sz, "dtype": str(dtype), + "weight": weight, + "source": source, "flops": flops, "bytes": nbytes, "kernel_latency_us": kernel_us, @@ -1125,21 +977,46 @@ def run_performance(kernel_fn: Callable, config: dict, gpu: GPUSpec, if primary_result is None and all_results: primary_result = all_results[-1] + # Weighted aggregate reporting for corpus cases, grouped per dtype so + # results from different dtypes are never mixed into one aggregate. + corpus_summary = None + corpus_entries = [entry for entry in all_results if entry["source"] == "corpus"] + if corpus_entries: + weighted = weighted_aggregate( + [ + { + "dtype": entry["dtype"], + "weight": entry["weight"], + "kernel_ms": entry["kernel_latency_us"] / 1000.0, + "ref_ms": entry["pytorch_latency_us"] / 1000.0, + } + for entry in corpus_entries + ] + ) + corpus_summary = {"cases": corpus_entries, "weighted": weighted} + print("\n === SHAPE CORPUS: weighted aggregates ===") + for dtype_name, agg in weighted.items(): + print(f" dtype={dtype_name}: cases={agg['cases']}, total_weight={agg['weight']}") + print(f" weighted_kernel_latency_us: {agg['kernel_ms'] * 1000.0:.2f}") + print(f" weighted_pytorch_latency_us: {agg['ref_ms'] * 1000.0:.2f}") + print(f" weighted_speedup_vs_pytorch: {agg['speedup']:.3f}x") + return { "primary": primary_result, "all": all_results, + "corpus": corpus_summary, } # ========================================================================= -# 7. PROFILER (optional) +# 5. PROFILER (optional) # ========================================================================= -def run_profile(kernel_fn: Callable, config: dict): +def run_profile(kernel_fn: Callable, spec: KernelSpec): """Run torch profiler and save a trace.""" - device = "cuda" - gen_fn = config["input_generator"] - sizes = config["test_sizes"] + device = BENCH_DEVICE + gen_fn = spec.input_generator + sizes = _spec_sizes(spec) # Use 'medium' or first size prof_size = None @@ -1150,7 +1027,7 @@ def run_profile(kernel_fn: Callable, config: dict): if prof_size is None: prof_size = sizes[0][1] - dtype = config["test_dtypes"][0] + dtype = _spec_dtypes(spec)[0] inputs = gen_fn(prof_size, dtype, device, seed=42) trace_dir = "./traces" @@ -1188,7 +1065,7 @@ def run_profile(kernel_fn: Callable, config: dict): # ========================================================================= -# 8. MAIN -- orchestrate everything and produce structured output +# 6. MAIN -- orchestrate everything and produce structured output # ========================================================================= def main(): @@ -1197,13 +1074,39 @@ def main(): parser = argparse.ArgumentParser(description="AutoKernel benchmark harness") parser.add_argument("--kernel", type=str, default=None, help="Kernel type to benchmark (default: read from kernel.py)") + parser.add_argument("--spec", type=str, default=None, + help="External KernelSpec locator, e.g. " + "'path/to/spec.py:SPEC' or 'package.module:SPEC'. " + "Takes precedence over --kernel.") + parser.add_argument("--spec-override", action="store_true", + help="Allow --spec to replace a built-in operation of the same name") parser.add_argument("--sizes", type=str, default="all", help="Which sizes to benchmark: small|medium|large|all (default: all)") parser.add_argument("--quick", action="store_true", help="Quick mode: skip correctness stages 3-5, bench only large size") parser.add_argument("--profile", action="store_true", help="Enable torch profiler trace") + parser.add_argument("--shape-corpus", type=str, default=None, metavar="PATH", + help="Versioned JSON shape corpus; validated cases are " + "appended to the built-in sweep and benchmarked once each") + parser.add_argument("--shape-corpus-only", action="store_true", + help="Benchmark only the --shape-corpus cases, skipping the " + "built-in size sweep (requires --shape-corpus)") + parser.add_argument("--check-backward", action="store_true", + help="Also verify gradients against the reference " + "(requires the spec to declare a backward_spec; " + "correctness-only, no performance claims)") + parser.add_argument("--check-compile", action="store_true", + help="Verify torch.compile parity using the spec's compile settings " + "(correctness-only; compilation is never timed)") + parser.add_argument("--result-json", type=str, + default=os.path.join(_SCRIPT_DIR, "workspace", "bench_result.json"), + metavar="PATH", + help="Atomic machine-readable result path " + "(default: workspace/bench_result.json)") args = parser.parse_args() + if args.shape_corpus_only and not args.shape_corpus: + parser.error("--shape-corpus-only requires --shape-corpus PATH") # ------------------------------------------------------------------ # Import the kernel module @@ -1216,6 +1119,41 @@ def main(): kernel_fn = None kernel_type = args.kernel + # ------------------------------------------------------------------ + # Resolve the operation specification. + # Precedence: --spec, then --kernel, then kernel.py::KERNEL_TYPE. + # Loading happens after argument parsing, so `bench.py --help` never + # imports an external specification. + # ------------------------------------------------------------------ + registry: KernelRegistry = create_builtin_registry() + spec: Optional[KernelSpec] = None + if args.spec: + try: + spec, registry = resolve_spec( + spec_locator=args.spec, + registry=registry, + override=args.spec_override, + ) + except (SpecLoadError, SpecValidationError) as e: + print(f"\nERROR: {e}") + print("\ncorrectness: FAIL") + print("throughput_tflops: 0.000") + sys.exit(1) + kernel_type = spec.name + print(f"kernel_spec: {args.spec}") + + # When the operation is already determined (--spec or --kernel), fetch its + # spec and validate the shape corpus *before* importing the candidate + # module: malformed metadata must fail without executing candidate code + # and before any GPU allocation. + kernel_type = spec.name if spec is not None else args.kernel + spec_locked = kernel_type is not None + corpus_cases = None + if spec_locked: + if spec is None: + spec = _get_spec_or_exit(registry, kernel_type) + corpus_cases = _load_validated_corpus(args, spec) + try: # Add cwd to path so 'import kernel' works if os.getcwd() not in sys.path: @@ -1228,46 +1166,50 @@ def main(): kernel_module = importlib.import_module("kernel") kernel_fn = kernel_module.kernel_fn - if kernel_type is None: - kernel_type = getattr(kernel_module, "KERNEL_TYPE", None) - if kernel_type is None: - print("ERROR: kernel.py has no KERNEL_TYPE attribute and --kernel not specified") - sys.exit(1) + declared_type = getattr(kernel_module, "KERNEL_TYPE", None) + if spec_locked: + resolved = kernel_type + else: + resolved = resolve_operation_name(None, args.kernel, declared_type) + if resolved is None: + print("ERROR: kernel.py has no KERNEL_TYPE attribute and --kernel not specified") + sys.exit(1) + if declared_type is not None and declared_type != resolved: + print(f"WARNING: kernel.py declares KERNEL_TYPE '{declared_type}' but " + f"'{resolved}' was requested; benchmarking kernel_fn against " + f"'{resolved}'") + kernel_type = resolved print(f"kernel_type: {kernel_type}") - print(f"kernel_module: kernel.py loaded successfully") + print("kernel_module: kernel.py loaded successfully") except SyntaxError as e: - print(f"\nERROR: kernel.py has a syntax error:") + print("\nERROR: kernel.py has a syntax error:") print(f" {e}") traceback.print_exc() - print(f"\ncorrectness: FAIL") - print(f"throughput_tflops: 0.000") + print("\ncorrectness: FAIL") + print("throughput_tflops: 0.000") sys.exit(1) except Exception as e: - print(f"\nERROR: Failed to import kernel.py:") + print("\nERROR: Failed to import kernel.py:") print(f" {type(e).__name__}: {e}") traceback.print_exc() - print(f"\ncorrectness: FAIL") - print(f"throughput_tflops: 0.000") - sys.exit(1) - - # Validate kernel type - if kernel_type not in KERNEL_CONFIGS: - print(f"\nERROR: Unknown kernel type '{kernel_type}'") - print(f" Available: {', '.join(KERNEL_CONFIGS.keys())}") - print(f"\ncorrectness: FAIL") - print(f"throughput_tflops: 0.000") + print("\ncorrectness: FAIL") + print("throughput_tflops: 0.000") sys.exit(1) - config = KERNEL_CONFIGS[kernel_type] + # The default selection path (kernel.py::KERNEL_TYPE) resolves the spec + # and validates the corpus only after the candidate import. + if not spec_locked: + spec = _get_spec_or_exit(registry, kernel_type) + corpus_cases = _load_validated_corpus(args, spec) # ------------------------------------------------------------------ # GPU Detection # ------------------------------------------------------------------ gpu = detect_gpu() - print(f"\n=== GPU INFO ===") + print("\n=== GPU INFO ===") print(f"gpu_name: {gpu.name}") print(f"gpu_sm_count: {gpu.sm_count}") print(f"gpu_memory_gb: {gpu.memory_gb}") @@ -1281,28 +1223,67 @@ def main(): # ------------------------------------------------------------------ # Correctness # ------------------------------------------------------------------ - print(f"\n=== CORRECTNESS ===") + print("\n=== CORRECTNESS ===") try: - correctness_results = run_correctness(kernel_fn, config, quick=args.quick) + correctness_results = run_correctness( + kernel_fn, spec, quick=args.quick, + corpus_cases=corpus_cases, corpus_only=args.shape_corpus_only, + ) except Exception as e: print(f"\nFATAL: Correctness testing crashed: {type(e).__name__}: {e}") traceback.print_exc() correctness_results = {"correctness": "FAIL", "smoke_test": "CRASH", "shape_sweep": "CRASH", "numerical_stability": "CRASH", "determinism": "CRASH", "edge_cases": "CRASH"} - print(f"\n--- Correctness Summary ---") + print("\n--- Correctness Summary ---") print(f"smoke_test: {correctness_results.get('smoke_test', 'N/A')}") print(f"shape_sweep: {correctness_results.get('shape_sweep', 'N/A')}") print(f"numerical_stability: {correctness_results.get('numerical_stability', 'N/A')}") print(f"determinism: {correctness_results.get('determinism', 'N/A')}") print(f"edge_cases: {correctness_results.get('edge_cases', 'N/A')}") print(f"correctness: {correctness_results['correctness']}") + print(f"FORWARD_CORRECTNESS: {correctness_results['correctness']}") + + # ------------------------------------------------------------------ + # Backward verification (opt-in; correctness-only, never timed) + # ------------------------------------------------------------------ + backward_result = None + backward_requested = args.check_backward or ( + spec.backward_spec is not None and spec.backward_spec.enabled_by_default + ) + if backward_requested: + try: + backward_result = run_backward_check(kernel_fn, spec) + except Exception as e: + print(f"\nFATAL: Backward verification crashed: {type(e).__name__}: {e}") + traceback.print_exc() + backward_result = {"status": "FAIL", "reason": f"crash: {type(e).__name__}: {e}"} + print("BACKWARD_CORRECTNESS: FAIL") + + # ------------------------------------------------------------------ + # Compile verification (opt-in; correctness-only, never timed) + # ------------------------------------------------------------------ + compile_result = None + compile_requested = args.check_compile or ( + spec.compile_spec is not None and spec.compile_spec.enabled + ) + if compile_requested: + try: + compile_result = run_compile_check(kernel_fn, spec) + except Exception as e: + print(f"\nFATAL: Compile verification crashed: {type(e).__name__}: {e}") + traceback.print_exc() + compile_result = { + "status": "FAIL", + "reason": f"crash: {type(e).__name__}: {e}", + } + print("COMPILE_CORRECTNESS: FAIL") # ------------------------------------------------------------------ # Performance # ------------------------------------------------------------------ # Determine primary size info for the header - _perf_sizes = config["test_sizes"] + _perf_sizes = _spec_sizes(spec) _perf_primary_label = None _perf_primary_size = None for _pl, _ps in _perf_sizes: @@ -1312,7 +1293,7 @@ def main(): break if _perf_primary_size is None: _perf_primary_label, _perf_primary_size = _perf_sizes[-1] - _perf_dtype = config["test_dtypes"][0] + _perf_dtype = _spec_dtypes(spec)[0] _size_params = ", ".join(f"{k}={v}" for k, v in _perf_primary_size.items()) print(f"\n=== PERFORMANCE ({_perf_primary_label}: {_size_params}, dtype={_perf_dtype}) ===") @@ -1323,7 +1304,10 @@ def main(): if args.quick: sizes_filter = "large" torch.cuda.reset_peak_memory_stats() - perf_results = run_performance(kernel_fn, config, gpu, sizes_filter=sizes_filter) + perf_results = run_performance( + kernel_fn, spec, gpu, sizes_filter=sizes_filter, + corpus_cases=corpus_cases, corpus_only=args.shape_corpus_only, + ) peak_vram_mb = torch.cuda.max_memory_allocated() / 1024 / 1024 except Exception as e: print(f"\nFATAL: Performance benchmarking crashed: {type(e).__name__}: {e}") @@ -1345,7 +1329,7 @@ def main(): print(f"bytes: {primary['bytes']}") print(f"peak_vram_mb: {peak_vram_mb:.1f}") - print(f"\n=== COMPARISON VS PYTORCH ===") + print("\n=== COMPARISON VS PYTORCH ===") print(f"pytorch_latency_us: {primary['pytorch_latency_us']:.2f}") print(f"pytorch_latency_ms: {primary['pytorch_latency_us'] / 1000.0:.4f}") print(f"kernel_latency_us: {primary['kernel_latency_us']:.2f}") @@ -1354,26 +1338,26 @@ def main(): print(f"pytorch_tflops: {primary['ref_throughput_tflops']:.3f}") print(f"kernel_tflops: {primary['throughput_tflops']:.3f}") else: - print(f"\nlatency_us: 0.00") - print(f"latency_ms: 0.0000") - print(f"throughput_tflops: 0.000") - print(f"bandwidth_gb_s: 0.0") - print(f"pct_peak_compute: 0.0%") - print(f"pct_peak_bandwidth: 0.0%") + print("\nlatency_us: 0.00") + print("latency_ms: 0.0000") + print("throughput_tflops: 0.000") + print("bandwidth_gb_s: 0.0") + print("pct_peak_compute: 0.0%") + print("pct_peak_bandwidth: 0.0%") print(f"peak_vram_mb: {peak_vram_mb:.1f}") - print(f"\n=== COMPARISON VS PYTORCH ===") - print(f"pytorch_latency_us: 0.00") - print(f"pytorch_latency_ms: 0.0000") - print(f"kernel_latency_us: 0.00") - print(f"kernel_latency_ms: 0.0000") - print(f"speedup_vs_pytorch: 0.000x") + print("\n=== COMPARISON VS PYTORCH ===") + print("pytorch_latency_us: 0.00") + print("pytorch_latency_ms: 0.0000") + print("kernel_latency_us: 0.00") + print("kernel_latency_ms: 0.0000") + print("speedup_vs_pytorch: 0.000x") # ------------------------------------------------------------------ # All sizes summary table # ------------------------------------------------------------------ all_perf = perf_results.get("all", []) if len(all_perf) > 1: - print(f"\n=== SIZE SWEEP ===") + print("\n=== SIZE SWEEP ===") print(f"{'size':<12} {'kernel_us':>12} {'pytorch_us':>12} {'speedup':>10} {'tflops':>10} {'%peak':>8}") print("-" * 66) for entry in all_perf: @@ -1386,7 +1370,7 @@ def main(): # ------------------------------------------------------------------ if args.profile: try: - run_profile(kernel_fn, config) + run_profile(kernel_fn, spec) except Exception as e: print(f"\nWARNING: Profiling failed: {type(e).__name__}: {e}") @@ -1396,7 +1380,51 @@ def main(): t_elapsed = time.time() - t_start throughput = primary["throughput_tflops"] if primary else 0.0 - print(f"\n=== FINAL ===") + corpus_identity = None + if args.shape_corpus: + corpus_identity = { + "source": os.path.abspath(args.shape_corpus), + "mode": "only" if args.shape_corpus_only else "append", + "cases": [ + { + "name": case.name, + "size": dict(case.size), + "dtype": case.dtype, + "weight": case.weight, + "tags": list(case.tags), + } + for case in (corpus_cases or ()) + ], + } + result_payload = result_envelope( + kernel_type, + environment=collect_environment_metadata(BENCH_DEVICE), + request={ + "spec": args.spec, + "sizes": args.sizes, + "quick": args.quick, + "profile": args.profile, + "check_backward": backward_requested, + "check_compile": compile_requested, + }, + gpu=asdict(gpu), + shape_corpus=corpus_identity, + forward=correctness_results, + backward=backward_result, + compile=compile_result, + performance={ + **perf_results, + "peak_vram_mb": peak_vram_mb, + }, + bench_time_seconds=t_elapsed, + ) + try: + result_path = write_result_atomic(args.result_json, result_payload) + print(f"result_json: {result_path}") + except Exception as e: + print(f"WARNING: Failed to write result JSON: {type(e).__name__}: {e}") + + print("\n=== FINAL ===") print(f"kernel_type: {kernel_type}") print(f"correctness: {correctness_results['correctness']}") print(f"throughput_tflops: {throughput:.3f}") @@ -1404,8 +1432,8 @@ def main(): print(f"speedup_vs_pytorch: {primary['speedup_vs_pytorch']:.3f}x") print(f"pct_peak_compute: {primary['pct_peak_compute']:.1f}%") else: - print(f"speedup_vs_pytorch: 0.000x") - print(f"pct_peak_compute: 0.0%") + print("speedup_vs_pytorch: 0.000x") + print("pct_peak_compute: 0.0%") print(f"bench_time_seconds: {t_elapsed:.1f}") if t_elapsed > 90: diff --git a/campaign.py b/campaign.py new file mode 100644 index 00000000..bc2b3254 --- /dev/null +++ b/campaign.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +"""Validate and prepare model optimization campaigns. + +Usage: + python campaign.py validate path/to/campaign.json + python campaign.py rank path/to/campaign.json + python campaign.py plan path/to/campaign.json --output workspace/optimization_plan.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +from autokernel.campaign import ( + CampaignError, + load_campaign, + parse_agent_command, + prepare_campaign, + rank_targets, + run_campaign, + write_optimization_plan, +) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Validate and prepare an AutoKernel optimization campaign" + ) + subparsers = parser.add_subparsers(dest="command", required=True) + for name in ("validate", "rank"): + command = subparsers.add_parser(name) + command.add_argument("campaign", type=Path) + plan = subparsers.add_parser("plan") + plan.add_argument("campaign", type=Path) + plan.add_argument( + "--output", + type=Path, + default=Path("workspace/optimization_plan.json"), + ) + prepare = subparsers.add_parser("prepare") + prepare.add_argument("campaign", type=Path) + prepare.add_argument( + "--output-dir", + type=Path, + default=Path("workspace"), + ) + prepare.add_argument( + "--trust-specs", + action="store_true", + help="Allow loading Python spec locators from this campaign", + ) + run = subparsers.add_parser("run") + run.add_argument("campaign", type=Path) + run.add_argument("--budget-hours", type=float, default=10.0) + run.add_argument("--resume", action="store_true") + run.add_argument("--dry-run", action="store_true") + run.add_argument( + "--trust-specs", + action="store_true", + help="Allow loading Python spec locators from this campaign", + ) + run.add_argument( + "--agent-command", + help=( + "Alternative agent command; supports {repo} and {prompt_file} " + "placeholders and is executed without a shell" + ), + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + try: + campaign = load_campaign(args.campaign) + except CampaignError as exc: + print(f"CAMPAIGN_VALIDATION: FAIL\n{exc}", file=sys.stderr) + return 2 + + if args.command == "validate": + print("CAMPAIGN_VALIDATION: PASS") + print(f"workload_id: {campaign.workload['workload_id']}") + print(f"targets: {len(campaign.targets)}") + return 0 + + if args.command == "rank": + rows = [] + for rank, target in enumerate(rank_targets(campaign), start=1): + rows.append( + { + "rank": rank, + "name": target.name, + "operation": target.operation, + "impact_pct": target.impact_pct( + campaign.total_profiled_device_time_us + ), + "calls": target.calls, + "spec_locator": target.spec_locator, + } + ) + print(json.dumps(rows, indent=2)) + return 0 + + if args.command == "plan": + plan = write_optimization_plan(campaign, args.output) + print("CAMPAIGN_PLAN: PASS") + print(f"output: {args.output}") + print(f"targets: {len(plan['kernels_to_optimize'])}") + return 0 + + if args.command == "run": + repo_root = Path(__file__).resolve().parent + prompt_path = repo_root / "workspace" / "overnight_prompt.md" + command = ( + parse_agent_command( + args.agent_command, + repo_root=repo_root, + prompt_path=prompt_path, + ) + if args.agent_command + else None + ) + try: + receipt = run_campaign( + campaign, + repo_root=repo_root, + budget_hours=args.budget_hours, + resume=args.resume, + dry_run=args.dry_run, + trust_specs=args.trust_specs, + agent_command=command, + ) + except CampaignError as exc: + print(f"CAMPAIGN_RUN: FAIL\n{exc}", file=sys.stderr) + return 2 + verdict = ( + "PASS" if receipt["status"] in ("prepared", "completed") else "FAIL" + ) + print(f"CAMPAIGN_RUN: {verdict}") + print(f"status: {receipt['status']}") + print(f"receipt: {repo_root / 'workspace' / 'overnight_receipt.json'}") + return 0 if verdict == "PASS" else 1 + + try: + receipt = prepare_campaign( + campaign, + args.output_dir, + trust_specs=args.trust_specs, + ) + except CampaignError as exc: + print(f"CAMPAIGN_PREPARE: FAIL\n{exc}", file=sys.stderr) + return 2 + print("CAMPAIGN_PREPARE: PASS") + print(f"output: {args.output_dir}") + print(f"targets: {len(receipt['targets'])}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/conftest.py b/conftest.py new file mode 100644 index 00000000..b81783df --- /dev/null +++ b/conftest.py @@ -0,0 +1,16 @@ +"""Pytest bootstrap for the repository root. + +The command-line entry points (``bench.py``, ``extract.py``, ``reference.py``) +live at the repository root and are imported by tests, so the root must be on +``sys.path`` regardless of where pytest is invoked from. +""" + +from __future__ import annotations + +import os +import sys + +REPO_ROOT = os.path.dirname(os.path.abspath(__file__)) + +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) diff --git a/discovery.py b/discovery.py new file mode 100644 index 00000000..d5a622e8 --- /dev/null +++ b/discovery.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +"""Validate discovery reports and rank graph regions by e2e impact. + +Usage: + python discovery.py validate path/to/discovery.json + python discovery.py rank path/to/discovery.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +from autokernel.discovery import ( + DiscoveryError, + load_discovery_report, + load_profiler_export, + rank_operators, + rank_regions, + write_discovery_report, +) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Validate and rank MotionKernel discovery reports" + ) + sub = parser.add_subparsers(dest="command", required=True) + for name in ("validate", "rank"): + cmd = sub.add_parser(name) + cmd.add_argument("report", type=Path) + if name == "rank": + cmd.add_argument( + "--impact-floor", + type=float, + default=0.005, + help="Minimum optimistic e2e improvement to search (default 0.5%%)", + ) + ingest = sub.add_parser( + "ingest-profiler", + help="Convert a FastVideo profiler export into a discovery report", + ) + ingest.add_argument("profile", type=Path) + ingest.add_argument("--output", type=Path, required=True) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + try: + if args.command == "ingest-profiler": + report = load_profiler_export(args.profile) + write_discovery_report(report, args.output) + print("PROFILER_INGEST: PASS") + print(f"operators: {len(report.operators)}") + print(f"output: {args.output}") + return 0 + report = load_discovery_report(args.report) + except DiscoveryError as exc: + print(f"DISCOVERY: FAIL\n{exc}", file=sys.stderr) + return 2 + + if args.command == "validate": + print("DISCOVERY_VALIDATION: PASS") + print(f"workload_id: {report.workload.get('workload_id')}") + print(f"operators: {len(report.operators)}") + print(f"regions: {len(report.regions)}") + print(f"graph_breaks: {len(report.graph_breaks)}") + return 0 + + if args.command == "rank": + ranked = rank_regions( + report.regions, + total_cuda_time_us=report.total_cuda_time_us, + impact_floor=args.impact_floor, + ) + ops = rank_operators( + report.operators, + total_cuda_time_us=report.total_cuda_time_us, + ) + payload = { + "workload_id": report.workload.get("workload_id"), + "total_cuda_time_us": report.total_cuda_time_us, + "operators": [ + { + "name": op.name, + "op_key": op.op_key, + "share_of_e2e": share, + "cuda_time_us": op.cuda_time_us, + "self_cuda_time_us": op.self_cuda_time_us, + "calls": op.calls, + } + for op, share in ops + ], + "regions": [item.as_dict() for item in ranked], + "graph_breaks": [item.as_dict() for item in report.graph_breaks], + } + print(json.dumps(payload, indent=2)) + return 0 + + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/FASTVIDEO_UNIVERSAL_OPTIMIZATION_AGENT_PLAN.md b/docs/FASTVIDEO_UNIVERSAL_OPTIMIZATION_AGENT_PLAN.md new file mode 100644 index 00000000..44f4481c --- /dev/null +++ b/docs/FASTVIDEO_UNIVERSAL_OPTIMIZATION_AGENT_PLAN.md @@ -0,0 +1,463 @@ +# Universal FastVideo optimization agent plan + +## Mission + +Turn MotionKernel into a model-independent optimization system for models that +already run in FastVideo. A user should be able to supply a FastVideo model, +representative workload, GPU budget, and output directory. MotionKernel should +profile the workload, discover worthwhile kernel regions, search for optimized +implementations, validate full-generation correctness and performance, and +produce compatible artifacts plus an honest morning report. + +LTX is the first proof model. Completing this plan means LTX can be optimized +without adding LTX-specific fusion calls to its FastVideo implementation. + +## Repositories + +- MotionKernel: `` (this repository) +- FastVideo: `` +- Existing FastVideo guide: + `/docs/contributing/kernel_optimization.md` +- Existing Wan measurement script: + `/examples/inference/optimizations/wan_fusions_ab.py` +- Existing Wan results: + `docs/WAN_KERNEL_RESULTS.md` + +Before editing either repository, read its `AGENTS.md`, inspect the current +branches and open PRs, and synchronize with the repository's main branch +without discarding unrelated work. + +## Current foundation + +Do not rebuild these pieces: + +- FastVideo metadata-only campaign capture and workload timing. +- MotionKernel campaign validation, ranking, preparation, resumable execution, + terminal receipts, and morning reports. +- `KernelSpec` validation and production shape corpora. +- Opt-in loading of promoted MotionKernel artifacts in FastVideo. +- Three Wan fusion targets and their isolated GB200 measurements. +- A reproducible full-generation native-versus-fused Wan benchmark. + +The Wan result is an important constraint: isolated kernels improved by +approximately 8.6-10.6x, but a 50-step generation remained approximately +36.67 seconds in both modes. The universal system must rank candidates by +expected end-to-end value and avoid spending an overnight budget on regions +whose theoretical model-level impact is negligible. + +## Scope + +Implement the first complete version for: + +- inference only; +- forward kernels only; +- one GPU; +- CUDA tensor graphs; +- models already supported by FastVideo; +- LTX as the end-to-end acceptance model. + +Keep the contracts extensible for training, backward kernels, sequence +parallelism, and multi-GPU execution, but do not block the first usable system +on those capabilities. + +## Required user experience + +The final interface should be equivalent to: + +```bash +motionkernel optimize \ + --fastvideo-checkout /path/to/FastVideo \ + --model Lightricks/LTX-Video \ + --workload workloads/ltx_480p.yaml \ + --budget-hours 10 \ + --output workspace/ltx +``` + +One invocation must perform or resume: + +1. native baseline generation; +2. profiling and graph capture; +3. candidate discovery and impact ranking; +4. safe reference-spec generation; +5. kernel search; +6. isolated correctness and performance validation; +7. full-generation native-versus-optimized validation; +8. artifact packaging and the morning report. + +It is acceptable to return `no_worthwhile_candidate`. It is not acceptable to +claim success based only on isolated operator speedup. + +## Workstream 1: Workload contract and launcher + +Add a versioned workload manifest shared by the FastVideo adapter and +MotionKernel. It must describe: + +- model identifier and optional immutable revision; +- task and prompt or prompt-file reference; +- width, height, frame count, inference steps, guidance, and seed; +- dtype and attention backend; +- device count and distributed settings; +- warmup and measured repetitions; +- output-parity policy and performance threshold. + +Add initial manifests for the existing Wan benchmark and one canonical LTX +text-to-video workload. Do not encode model-specific Python callables in the +manifest. + +Build a FastVideo launcher that: + +- resolves the model through FastVideo's existing registries; +- runs baseline and candidate modes in separate processes; +- captures wall time, generation time, peak allocated memory, environment + identity, and output frames; +- writes structured JSON with stable schemas; +- preserves logs and failure reasons; +- can be resumed without repeating completed stages. + +Exit criteria: + +- The launcher reproduces the existing Wan A/B measurement. +- The same launcher runs an unmodified LTX model from a workload manifest. + +## Workstream 2: Universal profiling and graph capture + +The current `optimization_target` API records known regions. Add automatic +discovery data without requiring model-specific annotations. + +Use two complementary sources: + +1. `torch.profiler` for end-to-end CUDA time, call frequency, launch behavior, + and hotspot attribution. +2. Dynamo/FX graph capture for executable tensor subgraphs and dependency + information. + +Do not require capture of the entire generation pipeline as one graph. Begin +with repeated DiT/module calls and fall back to smaller graphable scopes when +graph breaks occur. Record graph breaks rather than hiding them. + +The capture must contain metadata only by default. Never serialize prompts, +weights, activations, tensor values, credentials, or model outputs into a +campaign. + +For each observed operator or region, capture: + +- stable graph fingerprint; +- ordered operations and dependencies; +- input/output tensor signatures; +- constants that are safe and necessary for semantics; +- shape frequency; +- inclusive and self CUDA time; +- call count; +- parent module scope; +- hardware and software identity. + +Exit criteria: + +- Profiling unmodified Wan and LTX produces ranked operator data. +- Repeated runs produce stable fingerprints for equivalent regions. +- Graph breaks and unsupported operations are visible in the report. + +## Workstream 3: Candidate discovery and value ranking + +Implement a region discovery pass over captured graphs. Start with an +allowlist of pure tensor operations and reject regions with mutation, data +dependent Python control flow, collectives, unsupported custom operators, or +unknown aliasing. + +Initial pattern families: + +- elementwise chains; +- normalization plus modulation; +- residual, gate, and normalization chains; +- activation and projection epilogues; +- RoPE application and Q/K preparation; +- attention pre-processing and post-processing; +- layout, cast, and contiguous-copy chains; +- VAE elementwise and normalization chains. + +Generate overlapping candidate regions, deduplicate them by fingerprint, and +rank them using measured production frequency rather than single-call latency. + +Every candidate must report: + +- observed total GPU time; +- percentage of end-to-end GPU time; +- estimated reducible fraction; +- estimated maximum end-to-end improvement; +- confidence and rejection reasons. + +Enforce a configurable impact floor. By default, do not search a candidate +when its optimistic end-to-end improvement is below 0.5%. Prefer candidates +that can plausibly exceed the 1% promotion threshold. + +Exit criteria: + +- The Wan elementwise targets are discovered but correctly classified as + low-value for the measured 50-step workload. +- LTX produces a short ranked list with evidence for each candidate. + +## Workstream 4: Graph-derived KernelSpec generation + +Create a `KernelSpec` adapter that turns a supported captured FX region into a +search problem. The generated spec must include: + +- an executable PyTorch reference derived from the captured graph; +- ordered input and output contracts; +- a weighted corpus from observed production shapes; +- dtype, layout, device, and alignment constraints; +- numerical tolerances appropriate to output dtype and operation family; +- determinism and edge-case inputs; +- a stable operation and graph fingerprint. + +Do not use arbitrary source from the campaign. Spec generation must operate on +a validated, allowlisted intermediate representation. Fail closed when +semantics cannot be represented safely. + +Initially support the smallest useful ATen subset needed by the highest-value +LTX candidates. Add operations incrementally from actual profile evidence +instead of attempting to support all PyTorch operators upfront. + +Exit criteria: + +- At least one high-value LTX candidate becomes a valid generated + `KernelSpec`. +- The generated eager reference agrees with the original captured region. +- Existing handwritten Wan specs remain supported. + +## Workstream 5: Generic artifact and runtime dispatch + +Replace model-specific enablement with a generic artifact directory or +registry: + +```bash +FASTVIDEO_OPTIMIZED_KERNELS=/path/to/artifacts +``` + +Define a versioned artifact bundle containing: + +- manifest; +- graph fingerprint and operation identity; +- optimized kernel; +- compatibility constraints; +- isolated benchmark results; +- full-generation validation results; +- source campaign identity; +- morning report. + +Artifact selection must check: + +- model and optional model revision; +- graph fingerprint; +- tensor signature; +- GPU architecture; +- PyTorch, CUDA, and Triton compatibility; +- inference/training mode; +- distributed configuration. + +On a mismatch, load failure, or runtime failure, FastVideo must use the native +path and record the fallback reason. Preserve explicit trust boundaries: +FastVideo must never silently import code from an untrusted directory. + +Provide a generic graph-region dispatch mechanism. Do not add +`FASTVIDEO_LTX_FUSIONS` or handwritten LTX conditionals. The LTX source should +remain unchanged except for reusable framework hooks that apply equally to +other models. + +Exit criteria: + +- A promoted artifact is selected by compatibility rather than model-specific + code. +- An incompatible artifact reliably falls back to native execution. +- Disabling the artifact directory has zero behavioral effect. + +## Workstream 6: Orchestration and validation gates + +Extend the unattended runner into the top-level `optimize` workflow. Each +candidate moves through explicit states: + +```text +discovered -> specified -> searching -> operator_validated + -> end_to_end_validated -> promoted +``` + +Terminal alternatives include `unsupported`, `incorrect`, `plateaued`, +`regressed`, `below_impact_floor`, and `budget_exhausted`. + +Before promotion, require: + +- isolated correctness over the weighted shape corpus; +- numerical stability, determinism, and edge checks; +- isolated speedup; +- byte equality or declared full-output tolerance; +- no unacceptable peak-memory regression; +- repeated end-to-end measurements in separate processes; +- a default minimum 1% repeatable end-to-end improvement. + +Use medians and retain individual samples. Treat results within expected timing +noise as neutral, not as speedups. The morning report must distinguish isolated +operator speedup from model-level improvement. + +Exit criteria: + +- Interrupting and resuming does not lose completed work. +- A neutral result such as the current Wan pack is not promoted by default. +- Every terminal run leaves enough structured evidence to reproduce its + conclusion. + +## Workstream 7: LTX proof + +Run the complete workflow on a canonical LTX workload using the provided GPU +cluster. + +The proof must: + +1. use an LTX model already supported by FastVideo; +2. avoid LTX-specific optimization annotations and switches; +3. capture and rank real end-to-end hotspots; +4. generate at least one spec from the captured graph; +5. run the unattended kernel search; +6. perform full native-versus-optimized generation validation; +7. produce a complete artifact bundle and morning report. + +Success is either: + +- at least one promoted artifact with a repeatable end-to-end improvement of + 1% or more and acceptable parity; or +- an evidence-backed `no_worthwhile_candidate` result after the system + correctly evaluates the available regions. + +The framework behavior is the acceptance target; do not fabricate a promotion +to make LTX appear faster. + +## Workstream 8: Generalization audit + +After the LTX proof, run discovery-only campaigns on at least two structurally +different FastVideo models, preferably Cosmos and Kandinsky. Do not optimize +every candidate yet. Confirm: + +- no model-specific source changes are required; +- workload manifests are sufficient; +- graph capture failures are reported clearly; +- candidate fingerprints and compatibility checks remain stable; +- unsupported operation families generate actionable backlog items. + +Use the findings to create a prioritized ATen/graph-pattern support matrix. + +## Required tests + +Keep tests focused on contracts and failure modes: + +- workload schema validation; +- metadata privacy; +- fingerprint stability; +- graph-region safety and rejection; +- impact ranking and Amdahl ceiling calculations; +- generated-reference parity; +- artifact compatibility and native fallback; +- resume state transitions; +- end-to-end result classification. + +GPU validation should use representative production shapes and one canonical +full-generation workload. Do not multiply expensive tests without a distinct +risk they cover. + +## Commit and review strategy + +Keep MotionKernel and FastVideo changes in separate branches and PRs. Prefer +small commits by contract: + +1. workload schema and launcher; +2. profiler and graph capture; +3. discovery and ranking; +4. generated specs; +5. artifact compatibility and generic dispatch; +6. orchestration and validation; +7. LTX evidence and documentation. + +At the end of each workstream: + +- run the relevant focused CPU tests; +- run GPU validation when the workstream first touches CUDA behavior; +- update the plan with deviations and newly discovered gaps; +- commit the coherent result; +- report exact commands, results, commit hashes, and blockers. + +Do not wait for a full code review before starting the next independent +workstream, but do not build later contracts on unresolved correctness issues. + +## Agent operating rules + +- Profile before selecting targets. +- Optimize measured production shapes, not toy inputs. +- Calculate the theoretical end-to-end ceiling before starting a search. +- Preserve native fallbacks. +- Never claim a model improvement from an isolated microbenchmark. +- Never expose or commit SSH keys, kubeconfigs, model credentials, prompts, + activations, weights, or generated user content. +- Do not run GPU workloads on the login node; use SLURM compute allocations. +- Use resumable batch jobs for overnight work. +- Evaluate gaps in this plan while implementing it. Record material gaps, + decisions, and revised acceptance criteria in this document or an adjacent + design document rather than silently working around them. + +## Definition of done + +This project is ready for any FastVideo-supported inference model when a new +model normally requires only: + +1. selecting its existing FastVideo model identifier; +2. adding a declarative workload manifest; +3. running `motionkernel optimize`. + +No model-specific fusion annotations, handwritten `KernelSpec`, environment +switch, or runtime branch should be required. The run must conclude with +reproducible evidence for a promoted speedup or a clear explanation that no +worthwhile compatible kernel was found. + +## Implementation progress + +### Workstream 1 + +MotionKernel PR: https://github.com/aryan5v/motionkernel/pull/9 + +- `autokernel/workload/` — versioned workload schema (`schema_version: 1`), + generation-result schema, end-to-end classification, and FastVideo launcher + bridge with resume state. +- `workloads/wan_t2v_1.3b_480p.yaml` — reproduces the existing Wan A/B request. +- `workloads/ltx_480p.yaml` — canonical LTX-2 distilled T2V proof workload. +- `workload.py` CLI: `validate`, `show`, `validate-result`, `run-ab`. +- CPU tests in `tests/test_workload.py`. + +FastVideo PR: https://github.com/aryan5v/FastVideo/pull/18 + +- `examples/inference/optimizations/generation_launcher.py` — model-agnostic + launcher that loads a workload manifest, runs one mode per process, and + writes structured result JSON. + +### Workstream 2 (in progress) + +MotionKernel: + +- `autokernel/discovery/` — metadata-only discovery report schema, stable graph + fingerprints, pure-tensor allowlist, collective/data-dependent rejection. +- CPU tests in `tests/test_discovery.py`. +- `ranking.py`, `fx_capture.py`, `profiler_parse.py`, and + `profiler_export.py` — impact-floor ranking, CPU FX region capture, + metadata-only profiler ingestion, and exclusive CUDA-time accounting. +- FastVideo collects a dedicated post-warmup `torch.profiler` pass inside the + GPU worker without contaminating clean A/B timing samples. +- Wan 2.1 T2V 1.3B GPU profiling and MotionKernel ingestion completed on a + GB200. The clean 480x832, 49-frame, four-step generation median was 4.4069s; + attention and copy/cast traffic dominated the captured operator data. +- LTX-2 distilled T2V GPU profiling and MotionKernel ingestion also completed + on the same model-agnostic path: 480x768, 97 frames, eight steps, 4.6913s + clean median wall time, 67,802.22 MiB peak allocated CUDA memory, and 3,071 + CPU-side operator rows representing 2.771s of exclusive CUDA time. The + producer excluded all duplicate raw CUDA activity rows. +- `fx_capture.py` now supports model-independent repeated-module FX capture + (`RegionCaptureSession` / `capture_model_regions`): operations, dependencies, + tensor signatures, safe scalars only, graph breaks, unsupported ops, stable + fingerprints, fail-closed mutation/collective/aliasing rejection. Metadata + only — no weights, prompts, tensor values, or source. +- Next framework wall: correlate captured FX regions with measured profiler + hotspots and generate searchable `KernelSpec` objects. diff --git a/docs/WAN_KERNEL_RESULTS.md b/docs/WAN_KERNEL_RESULTS.md new file mode 100644 index 00000000..a6fad94e --- /dev/null +++ b/docs/WAN_KERNEL_RESULTS.md @@ -0,0 +1,69 @@ +# Wan kernel results + +## Gated residual plus affine LayerNorm + +The first Wan target fuses the post-self-attention gated residual update, +FP32 affine LayerNorm, two model-dtype casts, and both output writes into one +Triton launch. + +Validation environment: + +- GPU: NVIDIA GB200 +- PyTorch: 2.11.0+cu128 +- Triton: 3.6.0 +- corpus: `models/wan_gated_residual_norm_corpus.json` +- commit: `2481c35` + +All five forward gates passed: smoke, production-shape sweep, numerical +stability, determinism, and edge cases. The maximum absolute BF16 difference +was `0.03125`, with 100% of values inside the declared tolerance. + +| Wan shape | Fused latency | Eager latency | Speedup | +|---|---:|---:|---:| +| 1.3B, 480p, 49 frames | 51.91 µs | 439.60 µs | 8.468× | +| 1.3B, 480p, 81 frames | 75.39 µs | 684.80 µs | 9.083× | +| 14B, 480p, 49 frames | 153.02 µs | 1344.22 µs | 8.785× | +| 14B, 480p, 49 frames, SP4 | 46.93 µs | 370.87 µs | 7.903× | +| 14B, 480p, 81 frames, SP4 | 67.82 µs | 572.97 µs | 8.449× | + +The equally weighted corpus aggregate was 79.01 µs fused versus 682.49 µs +eager, an 8.638× operator speedup. These are isolated operator results; an +end-to-end Wan benchmark is still required to quantify generation-level +impact. + +## Modulated pre-attention LayerNorm + +The second target fuses FP32 non-affine LayerNorm with Wan's channel-wise +scale and shift. On the same GB200 class, commit `58e09c9` passed smoke, +the full 11-configuration shape sweep, numerical stability, determinism, and +edge cases. All five production corpus cases had 100% of values inside the +declared tolerance. + +| Wan shape | Fused latency | Eager latency | Speedup | +|---|---:|---:|---:| +| 1.3B, 480p, 49 frames | 37.52 µs | 370.64 µs | 9.878× | +| 1.3B, 480p, 81 frames | 54.34 µs | 580.42 µs | 10.682× | +| 14B, 480p, 49 frames | 101.04 µs | 1117.76 µs | 11.062× | +| 14B, 480p, 49 frames, SP4 | 32.94 µs | 305.97 µs | 9.290× | +| 14B, 480p, 81 frames, SP4 | 46.87 µs | 474.74 µs | 10.129× | + +The weighted corpus aggregate was 54.54 µs fused versus 569.91 µs eager, +a 10.449× isolated operator speedup. + +## Post-MLP gated residual + +The third target fuses Wan's FP32 gate multiplication and residual update into +one model-dtype output. Commit `58e09c9` passed the same five correctness +stages across all production and edge shapes on GB200. + +| Wan shape | Fused latency | Eager latency | Speedup | +|---|---:|---:|---:| +| 1.3B, 480p, 49 frames | 36.90 µs | 373.02 µs | 10.108× | +| 1.3B, 480p, 81 frames | 54.06 µs | 587.57 µs | 10.868× | +| 14B, 480p, 49 frames | 103.57 µs | 1182.09 µs | 11.413× | +| 14B, 480p, 49 frames, SP4 | 33.87 µs | 316.75 µs | 9.351× | +| 14B, 480p, 81 frames, SP4 | 49.61 µs | 495.33 µs | 9.984× | + +The weighted corpus aggregate was 55.60 µs fused versus 590.95 µs eager, +a 10.628× isolated operator speedup. All three results remain operator-level; +an end-to-end generation benchmark is the next measurement. diff --git a/docs/WEEK_1_2_AGENT_BRIEF.md b/docs/WEEK_1_2_AGENT_BRIEF.md new file mode 100644 index 00000000..20b19fe7 --- /dev/null +++ b/docs/WEEK_1_2_AGENT_BRIEF.md @@ -0,0 +1,926 @@ +# Agent brief: Week 1 and Week 2 + +## Mission + +Build the first two weeks of the downstream kernel platform: + +1. Week 1: replace hard-coded operation configuration with a stable, + externally loadable `KernelSpec` registry while preserving all existing + behavior. +2. Week 2: generalize correctness verification for structured outputs, custom + production shape corpora, optional backward checks, and + `torch.compile(fullgraph=True)` compatibility. + +This is an implementation assignment, not a design-only exercise. Finish each +week with tested code, documentation, and a reviewable pull request. + +Do not wait for changes or approval from the upstream repository. This is an +independently maintained downstream fork. Preserve the upstream MIT license and +attribution in all derived source. + +## Required operating rules + +- Begin from this fork's current `main`, after the downstream-foundation PR has + merged. +- Keep the Week 1 and Week 2 work in separate pull requests. +- Use branches named: + - `agent/kernel-spec-registry` + - `agent/generalized-verification` +- Never develop directly on `main`. +- Never push to the `upstream` remote. +- Do not use `git reset --hard` in the primary checkout. +- Run autonomous kernel experiments only in a disposable clone or Git + worktree. +- Do not overwrite unrelated local changes. +- Do not silently weaken an existing correctness tolerance or remove an edge + case to make a test pass. +- Preserve the existing no-argument CLI behavior. +- Prefer typed dataclasses and small pure functions over unstructured + dictionaries and new global maps. +- Keep GPU-specific imports and allocation out of registry discovery so CPU + tests can load and inspect every built-in specification. +- Do not claim GPU correctness or performance without recording the actual GPU + run. +- If a requirement cannot be completed, leave the code in a safe state, + document the exact blocker, and do not substitute a fake implementation. + +## Baseline architecture + +Before editing, understand the current ownership: + +- `bench.py` + - contains deterministic input generators; + - contains the hard-coded `KERNEL_CONFIGS` dictionary; + - performs five correctness stages; + - assumes each candidate returns one tensor; + - benchmarks `kernel_fn(**inputs)`. +- `reference.py` + - contains the PyTorch reference functions. +- `extract.py` + - duplicates shape keys, aliases, tolerances, FLOP formulas, byte formulas, + speedup estimates, and starter-kernel lookup; + - generates standalone candidate kernel files. +- `kernel.py` + - declares `KERNEL_TYPE`; + - exports `kernel_fn` with a signature matching the reference function. +- `profile.py` + - classifies GPU kernel names with a hard-coded rule list; + - discovers supported types by scanning `kernels/*.py`. +- `verify.py` + - has model-level replacement strategies for only a subset of operations; + - assumes a simple output comparison path. +- `orchestrate.py` + - consumes extracted kernel metadata and benchmark results. + +The first two weeks must reduce these hard-coded seams without breaking the +existing nine built-in kernel types. + +## Non-goals for Weeks 1 and 2 + +Do not implement the following yet: + +- modulated normalization kernels; +- gated-residual kernels; +- model-specific integration for Wan, Kandinsky, Cosmos, or LTX; +- a distributed GPU scheduler; +- an experiment database; +- package or repository renaming; +- a new web service or user interface; +- automatic graph rewriting for arbitrary model code; +- replacement of the existing profiler classification system; +- broad formatting or unrelated cleanup. + +Week 3 will use the interfaces created here to implement the first fused +operation. Do not pull Week 3 into these changes. + +--- + +# Week 1: custom-operation registry + +## Week 1 outcome + +At the end of Week 1, an external operation must be able to supply its +reference, inputs, sizes, tolerances, edge cases, performance accounting, and +starter kernel through one `KernelSpec`. `bench.py` and `extract.py` must +consume that specification without adding operation-specific branches. + +All nine existing built-in operations must behave exactly as they did before +the refactor. + +## Step 1: capture the compatibility baseline + +Before refactoring: + +1. Record the current built-in names from `bench.py::KERNEL_CONFIGS`. +2. Record the output of: + + ```bash + uv run bench.py --help + uv run extract.py --help + uv run profile.py --help + uv run verify.py --help + ``` + +3. Add CPU tests that freeze: + - the ordered set of built-in operation names; + - each built-in's size labels; + - supported dtype names; + - tolerance values; + - shape-key and alias metadata used by extraction; + - the existing `--kernel` CLI option. +4. If a GPU is available, run a quick baseline for at least: + - `matmul`; + - `layernorm`; + - `rmsnorm`. +5. Save the commands and environment metadata in the PR description. Do not + commit generated benchmark artifacts. + +The compatibility tests are required before moving configuration. They prevent +the refactor from accidentally changing benchmark coverage. + +## Step 2: add a real Python package + +Create the following structure: + +```text +autokernel/ + __init__.py + specs/ + __init__.py + types.py + registry.py + loader.py + builtins.py + inputs.py +tests/ + test_spec_registry.py + test_spec_loader.py + test_builtin_specs.py + fixtures/ + custom_add.py +``` + +Do not move the command-line scripts into the package during this milestone. +They should import the new package. This keeps the migration focused and +preserves current entry points. + +Add a `dev` optional dependency group in `pyproject.toml` containing `pytest`. +Update CI to install the development dependencies and run CPU tests. + +## Step 3: define `KernelSpec` + +Define the public types in `autokernel/specs/types.py`. Use type aliases and +dataclasses rather than an untyped dictionary. + +The initial interface should cover: + +```python +@dataclass(frozen=True) +class Tolerance: + atol: float + rtol: float + + +@dataclass(frozen=True) +class EdgeCase: + name: str + size: Mapping[str, int] + dtype: str | None = None + seed: int = 42 + input_transform: Callable[[InputMap], InputMap] | None = None + + +@dataclass(frozen=True) +class KernelSpec: + name: str + reference_fn: Callable[..., Any] + input_generator: Callable[..., InputMap] + sizes: Mapping[str, Mapping[str, int]] + dtypes: tuple[str, ...] + tolerances: Mapping[str, Tolerance] + edge_cases: tuple[EdgeCase, ...] + flops_fn: Callable[[Mapping[str, int]], int | float] + bytes_fn: Callable[[Mapping[str, int], int], int | float] + shape_keys: tuple[str, ...] + shape_aliases: Mapping[str, str] + starter_kernels: Mapping[str, Path] + speedup_estimate: str | None = None +``` + +The exact spelling can change if tests demonstrate a cleaner API, but retain +these responsibilities. + +### Type and validation requirements + +`KernelSpec` construction or registration must reject: + +- an empty or non-identifier-like operation name; +- duplicate size labels; +- a missing `small`, `medium`, or `large` size for built-ins; +- unknown dtype strings; +- missing tolerances for a declared dtype; +- negative, NaN, or infinite tolerances (a non-finite value would make every + comparison pass and silently disable the correctness gate); +- missing starter-kernel files; +- duplicate shape aliases that resolve inconsistently; +- a reference or input generator that is not callable. + +Validation errors must identify the specification and invalid field. + +Use canonical dtype strings at the specification boundary: + +- `float16` +- `bfloat16` +- `float32` + +Translate to `torch.dtype` only inside runtime code. This keeps discovery +serializable and CPU-friendly. + +Do not put backward or structured-output fields into `KernelSpec` during the +first refactor unless they are optional and unused. Week 2 will add them with +tests. + +## Step 4: implement the registry + +In `autokernel/specs/registry.py`, implement a small registry with: + +- `register(spec: KernelSpec) -> None` +- `get(name: str) -> KernelSpec` +- `list_names() -> tuple[str, ...]` +- `contains(name: str) -> bool` +- a duplicate-registration error; +- deterministic ordering; +- no GPU initialization at import time. + +Provide a function that creates the default registry rather than relying only +on mutable module globals: + +```python +def create_builtin_registry() -> KernelRegistry: + ... +``` + +A fresh registry must be usable in tests without affecting another test. + +## Step 5: migrate built-in metadata + +Create one `KernelSpec` for every current built-in: + +- `matmul` +- `softmax` +- `layernorm` +- `rmsnorm` +- `flash_attention` +- `fused_mlp` +- `cross_entropy` +- `rotary_embedding` +- `reduce` + +Move deterministic input generators out of `bench.py` into +`autokernel/specs/inputs.py`. + +The reference functions may remain in `reference.py` for compatibility. +`builtins.py` may import them. Do not duplicate reference math in multiple +locations. + +Move the following duplicated `extract.py` metadata into the specs: + +- `SHAPE_KEYS` +- `SHAPE_ALIAS_MAP` +- `TOLERANCES_MAP` +- `FLOPS_FN_SRC` responsibilities; +- `BYTES_FN_SRC` responsibilities; +- `SPEEDUP_ESTIMATES`; +- starter-kernel paths. + +Use real callables for FLOP and byte accounting. Do not continue storing Python +function bodies as source strings in the central model. + +If generated candidate files must remain standalone, update extraction to +serialize a safe, generated function representation from known metadata. Do +not call `eval` on external specification content. A preferable solution is to +emit simple numeric expressions through a constrained serializer or have the +generated file retain a reference to its spec. + +Keep compatibility aliases temporarily if other modules still import an old +constant, but mark them as deprecated and derive them from the registry. There +must be only one source of truth. + +## Step 6: add external spec loading + +Implement `autokernel/specs/loader.py`. + +Support: + +```text +package.module:SPEC +/absolute/path/to/spec.py:SPEC +relative/path/to/spec.py:SPEC +``` + +The selected object may be: + +- a `KernelSpec`; or +- a zero-argument callable returning a `KernelSpec`. + +Reject: + +- a missing file or module; +- a missing attribute; +- an attribute returning the wrong type; +- a loaded spec whose declared starter kernel does not exist; +- a spec name that collides with a built-in unless an explicit override flag + is supplied. + +Return actionable errors including the original locator. + +Do not mutate `sys.path` permanently. If loading from a file, use +`importlib.util.spec_from_file_location` with a unique module name. + +## Step 7: connect the CLI + +Add `--spec` to `bench.py` and `extract.py`. + +Required precedence: + +1. explicit `--spec`; +2. explicit `--kernel`; +3. `kernel.py::KERNEL_TYPE`, preserving current behavior. + +Examples: + +```bash +uv run bench.py --spec tests/fixtures/custom_add.py:SPEC +uv run extract.py --spec tests/fixtures/custom_add.py:SPEC --top 1 +``` + +When `--spec` is supplied: + +- load and validate it; +- register it in an isolated registry for that command; +- select it by its declared name; +- do not require editing `KERNEL_CONFIGS`, `reference.py`, or the hard-coded + extraction maps; +- preserve existing `--quick`, `--profile`, `--sizes`, and backend behavior. + +An external spec must not be imported when the CLI only asks for `--help`. + +## Step 8: provide a minimal external example + +Add `examples/custom_ops/add.py` and document it. + +The example should: + +- define a simple PyTorch reference; +- generate deterministic inputs; +- define small, medium, and large sizes; +- declare tolerances and accounting functions; +- point to a starter Triton kernel; +- export `SPEC`; +- run through the same benchmark path as a built-in. + +The example exists to prove extensibility, not performance. Keep it small. + +Do not use modulated normalization as the Week 1 example; that would combine +the registry refactor with the Week 3 feature. + +## Week 1 tests + +At minimum, add tests for: + +- registering and retrieving a valid spec; +- deterministic name ordering; +- duplicate registration; +- every validation error described above; +- loading a module locator; +- loading absolute and relative file locators; +- callable spec factories; +- bad module, path, attribute, and object errors; +- collision behavior; +- exact built-in metadata compatibility; +- deterministic inputs for a fixed seed; +- CLI precedence; +- external example discovery; +- `extract.py` consuming an external spec without an operation-specific map. + +CPU CI must run: + +```bash +uv sync --extra dev +uv run pytest -m "not gpu" +python -m compileall -q . +``` + +Mark GPU tests with `@pytest.mark.gpu` and register the marker in +`pyproject.toml`. + +If a supported GPU is available, also run: + +```bash +uv run bench.py --kernel matmul --quick +uv run bench.py --kernel layernorm --quick +uv run bench.py --kernel rmsnorm --quick +uv run bench.py --spec examples/custom_ops/add.py:SPEC --quick +``` + +## Week 1 acceptance criteria + +Week 1 is complete only when: + +- all nine built-in specs are registered; +- old built-in CLI commands behave unchanged; +- a new external operation runs without editing central operation maps; +- `bench.py` and `extract.py` use `KernelSpec` as their source of operation + metadata; +- no duplicate authoritative metadata remains in `extract.py`; +- registry discovery works on a CPU-only machine; +- CPU CI passes; +- available GPU smoke tests pass; +- README documentation contains a custom-operation quick start; +- the PR documents compatibility evidence and any GPU coverage not run. + +## Week 1 commit sequence + +Prefer small commits in this order: + +1. `Add kernel specification types and registry` +2. `Migrate built-in operation metadata` +3. `Load external kernel specifications` +4. `Use kernel specifications in benchmark and extraction` +5. `Document and test custom operations` + +Do not mix formatting-only changes into these commits. + +--- + +# Week 2: generalized verification + +## Week 2 outcome + +At the end of Week 2, the benchmark harness must correctly validate operations +that return multiple or nested outputs, accept production shape corpora without +editing Python source, optionally compare gradients, and optionally verify +`torch.compile(fullgraph=True)` compatibility. + +Week 2 begins only after the Week 1 registry PR is merged. + +## Step 1: extend the public specification + +Add explicit verification types in `autokernel/specs/types.py`. + +Recommended responsibilities: + +```python +@dataclass(frozen=True) +class OutputSpec: + # Optional paths to tensor leaves that participate in correctness. + included_paths: tuple[str, ...] | None = None + # Whether non-tensor leaves must match exactly. + compare_non_tensors: bool = True + + +@dataclass(frozen=True) +class BackwardSpec: + differentiable_inputs: tuple[str, ...] + output_paths: tuple[str, ...] | None = None + tolerances: Mapping[str, Tolerance] | None = None + enabled_by_default: bool = False + + +@dataclass(frozen=True) +class CompileSpec: + enabled: bool = False + fullgraph: bool = True + dynamic: bool = False +``` + +Add optional fields to `KernelSpec`: + +- `output_spec`; +- `backward_spec`; +- `compile_spec`. + +Existing built-ins must receive defaults that preserve their current behavior. + +Keep integration/replacement hooks out of this PR unless required for an +isolated compile test. Model replacement is a later milestone. + +## Step 2: implement output-tree handling + +Create: + +```text +autokernel/ + verification/ + __init__.py + outputs.py + backward.py + compile.py +``` + +`outputs.py` must handle: + +- a single tensor; +- tuples; +- lists; +- dictionaries with deterministic key order; +- named tuples if they appear in current model outputs; +- nested combinations of the above; +- exact comparison of supported non-tensor leaves when enabled. + +Represent every leaf with a stable path such as: + +```text +output +output[0] +output.updated_residual +output["aux"][1] +``` + +For tensor leaves: + +- require the same tree structure; +- require the same shape; +- require compatible dtype expectations; +- detect NaN and infinity per path; +- compute maximum and mean absolute error per path; +- apply dtype-specific tolerances; +- report the failing path and a concise error summary. + +Do not silently compare only the first tensor. Do not flatten dictionaries in +insertion-dependent ways. Do not coerce mismatched shapes. + +Return structured comparison results suitable for JSON output, not only +printed text. + +## Step 3: refactor benchmark correctness + +Replace the single-tensor assumptions in `bench.py` with the output-tree +comparator. + +Preserve the existing five forward correctness stages: + +1. smoke; +2. shape sweep; +3. numerical stability; +4. determinism; +5. edge cases. + +Structured outputs must pass all applicable stages. + +Determinism means: + +- identical output tree structure across runs; +- every tensor leaf compared; +- exact or configured deterministic tolerance; +- non-tensor leaves unchanged when comparison is enabled. + +Keep the public benchmark summary compatible, but add leaf-level details to a +structured result artifact under `workspace/`. + +Performance timing must call the full operation and must not introduce output +tree traversal inside the timed region. + +## Step 4: add production shape corpora + +Define a versioned JSON schema. The first version should resemble: + +```json +{ + "schema_version": 1, + "operation": "gated_residual_norm", + "cases": [ + { + "name": "example-production-shape", + "size": {"batch": 2, "tokens": 4096, "dim": 3072}, + "dtype": "bfloat16", + "weight": 37, + "tags": ["production", "forward"] + } + ] +} +``` + +Add `--shape-corpus PATH` to `bench.py`. + +Rules: + +- the corpus operation must match the selected spec; +- schema versions must be validated; +- size keys must be accepted by the spec; +- dtypes must be declared by the spec; +- names must be unique; +- weights must be positive integers; +- unknown fields should produce a clear error unless the schema explicitly + allows them; +- duplicate cases should be deterministically deduplicated or rejected; +- committed corpora must not include confidential model inputs or tensor data. + +The corpus contains metadata only. It must never serialize model activations, +weights, prompts, or user data. + +Add merge behavior: + +- default: use the spec's built-in cases; +- `--shape-corpus`: append validated corpus cases; +- `--shape-corpus-only`: use only corpus cases. + +Implementation decisions (recorded by the Week 2 agent): + +- `--shape-corpus-only` is a flag companion to `--shape-corpus PATH` + (argparse rejects it without a path); corpus cases join the stage-2 + correctness sweep and are each benchmarked once in the performance loop, + with weighted aggregates reported per dtype. +- Corpus loading and validation run before the candidate module is imported + whenever the operation is explicitly selected (`--spec`/`--kernel`), and + always before GPU detection: malformed metadata must fail without + executing candidate code. +- Two cases resolving to the same `(size, dtype)` configuration are + rejected with an actionable "merge their weights" message (the plan's + "deduplicated or rejected" choice, made explicit). + +Use `weight` for aggregate reporting, not to repeat allocations or benchmark +loops unnecessarily. + +Report both: + +- unweighted latency by case; +- weighted aggregate latency and speedup. + +Do not mix results from different dtypes into a single unexplained aggregate. + +## Step 5: add optional backward verification + +Backward verification is opt-in during Week 2: + +```bash +uv run bench.py --spec path/to/spec.py:SPEC --check-backward +``` + +Implementation requirements: + +1. Generate one canonical input mapping. +2. Deep-clone tensor inputs for reference and candidate paths. +3. Set `requires_grad=True` only for names declared in + `BackwardSpec.differentiable_inputs`. +4. Run reference and candidate independently. +5. Select tensor output leaves declared by `output_paths`, or all floating + tensor leaves when paths are omitted. +6. Produce deterministic upstream gradients with a fixed seed. Do not use + only `output.sum()` for every test because cancellation and symmetry can + hide errors. +7. Call `torch.autograd.grad` with matching inputs and upstream gradients. +8. Compare every requested input gradient by name. +9. Report missing gradients, unexpected gradients, shape differences, NaN, + infinity, maximum error, and mean error. +10. Do not accumulate gradients between cases. + +If a candidate is intentionally forward-only and no `BackwardSpec` exists, +`--check-backward` must fail with an actionable unsupported message rather +than silently skipping. + +Backward execution is correctness-only in Week 2. Do not add backward +performance claims yet. + +## Step 6: add optional compile verification + +Add: + +```bash +uv run bench.py --spec path/to/spec.py:SPEC --check-compile +``` + +Compile verification must: + +- wrap the candidate call in a stable Python function; +- call `torch.compile(..., fullgraph=True)` by default; +- compile outside all timed performance regions; +- run at least twice after compilation; +- compare compiled output against the eager reference through the same output + tree comparator; +- report graph breaks and compilation exceptions without swallowing them; +- record PyTorch, Triton, CUDA, GPU, and compile-mode metadata. + +When dynamic shape support is declared, test at least two compatible shapes +through the same compiled callable. Otherwise use static shapes. + +Do not report the first-call compilation latency as kernel execution latency. + +Direct Triton kernels may need to be exposed through a compile-aware wrapper. +Prefer PyTorch's supported Triton custom-operation integration rather than +adding graph-break exceptions. Keep the wrapper outside the core +`KernelSpec`; the spec should describe behavior, not own global registration +side effects. + +If `torch.compile` is not available in the installed PyTorch version, emit a +clear unsupported result. Do not mark the check as passed. + +Implementation decisions (recorded by the Week 2 agent): + +- Compile checks run before the performance section and use one stable wrapper + around the candidate. Each selected shape is called twice through the + compiled callable. +- Static checks use the small declared shape. Dynamic checks use small and + medium (or the first two available compatible shapes) through the same + callable. +- `PASS`, `FAIL`, and `UNSUPPORTED` are distinct structured statuses, and + PyTorch, Triton, CUDA, device, GPU and compile-mode metadata are recorded. + +## Step 7: add a structured-output fixture + +Add an example operation under `examples/custom_ops/` that returns: + +```python +{ + "output": tensor, + "aux": (updated_tensor, metadata_value), +} +``` + +It should have a simple differentiable reference and candidate so CPU unit +tests can verify: + +- nested output traversal; +- path generation; +- tensor comparison; +- exact metadata comparison; +- gradient comparison. + +GPU integration can use a small Triton implementation, but the core output and +gradient tests must also run on CPU. + +This fixture proves the framework. It is not the production modulated-norm +kernel. + +## Step 8: preserve machine-readable results + +Define versioned result records for: + +- forward correctness; +- per-output errors; +- backward correctness; +- per-input gradient errors; +- compile correctness; +- shape-corpus identity; +- environment metadata. + +Write JSON atomically under `workspace/`. Use a temporary file followed by +rename so an interrupted run does not leave valid-looking partial JSON. + +Do not break the existing greppable console output that the autonomous agent +loop consumes. Add new stable lines such as: + +```text +FORWARD_CORRECTNESS: PASS +BACKWARD_CORRECTNESS: PASS +COMPILE_CORRECTNESS: PASS +``` + +Only print `PASS` after the complete requested stage succeeds. + +Implementation decisions (recorded by the Week 2 agent): + +- Normal runs write schema version 1 to + `workspace/bench_result.json`; `--result-json PATH` overrides the location. +- The record contains the request, environment and GPU metadata, corpus + identity, forward leaf details, optional backward and compile reports, + performance results and elapsed time. +- The writer creates and fsyncs a temporary file in the destination directory, + then replaces the destination atomically. Non-finite diagnostic values are + encoded as strings so the artifact remains strict JSON. + +## Week 2 tests + +Add CPU tests for: + +- every supported output container; +- nested path stability; +- mismatched structures; +- mismatched shapes; +- per-leaf tolerance selection; +- NaN and infinity reporting; +- non-tensor metadata comparison; +- deterministic output verification; +- shape-corpus schema validation; +- operation mismatch; +- invalid dtype, keys, weights, and duplicates; +- corpus merging and corpus-only behavior; +- weighted aggregate calculations; +- gradient parity; +- missing and unexpected gradients; +- deterministic upstream gradients; +- unsupported backward behavior; +- compile option validation through mocks when a compiler/GPU is unavailable; +- atomic JSON result writing. + +GPU tests must cover: + +- a built-in single-output kernel; +- an external structured-output kernel; +- forward comparison; +- requested backward comparison when implemented; +- full-graph compile comparison. + +Run: + +```bash +uv sync --extra dev +uv run pytest -m "not gpu" +python -m compileall -q . +``` + +On the target GPU, also run the documented GPU suite and save the command, +environment, and result summary in the PR. + +## Week 2 acceptance criteria + +Week 2 is complete only when: + +- existing built-in single-tensor behavior remains unchanged; +- tensor, tuple, list, dictionary, named-tuple, and nested outputs are compared + correctly; +- all output leaves receive stable diagnostic paths; +- a valid external shape corpus is accepted without editing source; +- invalid corpora fail before GPU allocation; +- weighted and unweighted results are reported separately; +- optional backward checks compare every requested gradient; +- optional compile checks use full-graph mode and reference parity; +- compile time is excluded from runtime latency; +- requested checks cannot be silently skipped; +- machine-readable results are versioned and atomically written; +- CPU CI passes; +- available GPU correctness checks pass; +- documentation contains copy-paste examples for all new CLI options. + +## Week 2 commit sequence + +Prefer small commits in this order: + +1. `Compare structured kernel outputs` +2. `Load production shape corpora` +3. `Verify optional kernel gradients` +4. `Check full-graph kernel compilation` +5. `Record generalized verification results` +6. `Document and test verification workflows` + +--- + +# Pull request requirements + +Each pull request must include: + +## Summary + +- what changed; +- why the chosen interface is stable enough for Week 3; +- compatibility impact; +- explicit non-goals. + +## Evidence + +- CPU test command and result; +- compile-all command and result; +- GPU model and software versions, if run; +- exact GPU commands and results; +- before/after CLI compatibility notes; +- sample external-spec invocation; +- sample machine-readable result. + +## Risk assessment + +Address: + +- import-time GPU initialization; +- circular imports between scripts and the new package; +- external module-loading errors; +- arbitrary-code trust boundary for external specs; +- structured-output comparison gaps; +- accidental compile time in latency measurements; +- backward state or gradient accumulation; +- changes to autonomous-loop console parsing. + +## Review boundaries + +The Week 1 PR must not contain Week 2 verification features beyond harmless +optional type placeholders. The Week 2 PR must not contain production fused +kernels or model integrations. + +If either PR becomes too large to review confidently, split it along the commit +sequence above while keeping the acceptance criteria intact. + +# Final handoff + +At the end of Week 2, produce: + +- links to both merged pull requests; +- the final public `KernelSpec` example; +- a list of supported output structures; +- the shape-corpus schema and an example; +- forward, backward, and compile validation examples; +- CPU and GPU validation results; +- known limitations; +- a recommendation for the Week 3 modulated-norm specification. + +Do not begin Week 3 until the handoff demonstrates that a new external, +multi-output operation can be introduced without editing central operation +maps. diff --git a/examples/custom_ops/add.py b/examples/custom_ops/add.py new file mode 100644 index 00000000..939452e0 --- /dev/null +++ b/examples/custom_ops/add.py @@ -0,0 +1,90 @@ +"""Minimal external custom operation: elementwise add. + +This example exists to prove extensibility, not performance. It shows the whole +contract an out-of-tree operation must satisfy: + +* a PyTorch reference, +* a deterministic input generator, +* ``small`` / ``medium`` / ``large`` sizes, +* tolerances per dtype, +* FLOP and byte accounting, +* a starter kernel, +* an exported ``SPEC``. + +Run it through the normal harness:: + + cp examples/custom_ops/add_kernel.py kernel.py + uv run bench.py --spec examples/custom_ops/add.py:SPEC --quick + uv run extract.py --spec examples/custom_ops/add.py:SPEC --top 1 + +No change to ``bench.py``, ``extract.py``, ``reference.py`` or any central +operation map is required. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Mapping + +from autokernel.specs import ( + DT_BYTES, + EdgeCase, + KernelSpec, + Tolerance, + resolve_torch_dtype, + size, +) + +_HERE = Path(__file__).resolve().parent + +#: Starter kernel the agent begins optimizing from. +STARTER_KERNEL = _HERE / "add_kernel.py" + + +def add_ref(x: Any, y: Any) -> Any: + """Reference implementation: elementwise sum of two tensors.""" + return x + y + + +def gen_add_inputs( + size_map: Mapping[str, int], dtype: Any, device: str, seed: int = 42 +) -> dict: + """Deterministic inputs for a fixed seed.""" + import torch + + torch.manual_seed(seed) + torch_dtype = resolve_torch_dtype(dtype) + rows, cols = size_map["rows"], size_map["cols"] + x = torch.randn(rows, cols, device=device, dtype=torch_dtype) + y = torch.randn(rows, cols, device=device, dtype=torch_dtype) + return {"x": x, "y": y} + + +SPEC = KernelSpec( + name="custom_add", + reference_fn=add_ref, + input_generator=gen_add_inputs, + sizes={ + "small": {"rows": 256, "cols": 512}, + "medium": {"rows": 1024, "cols": 1024}, + "large": {"rows": 4096, "cols": 4096}, + }, + dtypes=("float16", "bfloat16", "float32"), + tolerances={ + "float16": Tolerance(atol=1e-3, rtol=1e-3), + "bfloat16": Tolerance(atol=2e-3, rtol=2e-3), + "float32": Tolerance(atol=1e-5, rtol=1e-5), + }, + # one add per element + flops_fn=size("rows") * size("cols"), + # two reads and one write per element + bytes_fn=3 * size("rows") * size("cols") * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_1023", size={"rows": 1023, "cols": 1023}), + EdgeCase(name="edge_4097", size={"rows": 4097, "cols": 129}), + ), + shape_keys=("rows", "cols"), + shape_aliases={"M": "rows", "N": "cols", "rows": "rows", "cols": "cols"}, + starter_kernels={"triton": STARTER_KERNEL}, + speedup_estimate="1.0-1.2x", +) diff --git a/examples/custom_ops/add_kernel.py b/examples/custom_ops/add_kernel.py new file mode 100644 index 00000000..28350606 --- /dev/null +++ b/examples/custom_ops/add_kernel.py @@ -0,0 +1,45 @@ +"""Starter kernel for the external ``custom_add`` example operation. + +Copy this file to ``kernel.py`` and benchmark it against the external spec:: + + cp examples/custom_ops/add_kernel.py kernel.py + uv run bench.py --spec examples/custom_ops/add.py:SPEC --quick + +It is intentionally simple: the example proves that an out-of-tree operation +flows through the same harness, not that elementwise add can be made faster. +""" + +KERNEL_TYPE = "custom_add" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def add_kernel( + x_ptr, y_ptr, out_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = tl.load(y_ptr + offsets, mask=mask, other=0.0) + tl.store(out_ptr + offsets, x + y, mask=mask) + + +def kernel_fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """Entry point called by bench.py. Must match the spec's reference signature.""" + assert x.shape == y.shape, f"shape mismatch: {x.shape} vs {y.shape}" + x_contig = x.contiguous() + y_contig = y.contiguous() + out = torch.empty_like(x_contig) + + n_elements = out.numel() + BLOCK_SIZE = 1024 + grid = (triton.cdiv(n_elements, BLOCK_SIZE),) + + add_kernel[grid](x_contig, y_contig, out, n_elements, BLOCK_SIZE=BLOCK_SIZE) + return out diff --git a/examples/custom_ops/affine.py b/examples/custom_ops/affine.py new file mode 100644 index 00000000..aaa0c6b2 --- /dev/null +++ b/examples/custom_ops/affine.py @@ -0,0 +1,110 @@ +"""Structured-output external operation: affine transform with an aux branch. + +This example proves the Week 2 verification framework: an operation whose +output is a *tree* -- + +.. code-block:: python + + { + "output": y, # tensor + "aux": (residual, n_terms), # (tensor, metadata value) + } + +-- flows through the same five correctness stages as a single-tensor +operation. The reference is differentiable, so the example also exercises +optional backward verification, and it is pure PyTorch so every check runs +on CPU: + +.. code-block:: bash + + cp examples/custom_ops/affine_kernel.py kernel.py + uv run bench.py --spec examples/custom_ops/affine.py:SPEC --quick + uv run bench.py --spec examples/custom_ops/affine.py:SPEC --check-backward + uv run bench.py --spec examples/custom_ops/affine.py:SPEC --check-compile + +The fixture exists to prove the framework. It is not a production kernel. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Mapping + +from autokernel.specs import ( + DT_BYTES, + BackwardSpec, + EdgeCase, + KernelSpec, + Tolerance, + resolve_torch_dtype, + size, +) + +_HERE = Path(__file__).resolve().parent + +#: Starter candidate the agent begins optimizing from. +STARTER_KERNEL = _HERE / "affine_kernel.py" + + +def affine_ref(x: Any, scale: Any, bias: Any) -> dict: + """Reference: ``y = x * scale + bias`` with an aux residual branch. + + Returns a nested output tree: the main tensor, an aux tuple holding the + residual tensor, and a non-tensor metadata value (the number of terms in + the affine expression). + """ + y = x * scale + bias + # Compute the auxiliary branch in float32 and cast once at its boundary. + # This gives eager and Inductor-fused FP16 execution the same numerical + # contract without loosening the output tolerance. + residual = ( + x.float() * scale.float() + bias.float() - x.float() + ).to(x.dtype) + return {"output": y, "aux": (residual, 3)} + + +def gen_affine_inputs( + size_map: Mapping[str, int], dtype: Any, device: str, seed: int = 42 +) -> dict: + """Deterministic inputs for a fixed seed.""" + import torch + + torch.manual_seed(seed) + torch_dtype = resolve_torch_dtype(dtype) + rows, cols = size_map["rows"], size_map["cols"] + x = torch.randn(rows, cols, device=device, dtype=torch_dtype) + scale = torch.randn(cols, device=device, dtype=torch_dtype) + bias = torch.randn(cols, device=device, dtype=torch_dtype) + return {"x": x, "scale": scale, "bias": bias} + + +SPEC = KernelSpec( + name="custom_affine", + reference_fn=affine_ref, + input_generator=gen_affine_inputs, + sizes={ + "small": {"rows": 128, "cols": 256}, + "medium": {"rows": 512, "cols": 512}, + "large": {"rows": 2048, "cols": 2048}, + }, + dtypes=("float16", "float32"), + tolerances={ + "float16": Tolerance(atol=1e-3, rtol=1e-3), + "float32": Tolerance(atol=1e-5, rtol=1e-5), + }, + # mul + add + residual sub per element + flops_fn=3 * size("rows") * size("cols"), + # read x, scale, bias; write output and residual (metadata is tiny) + bytes_fn=(3 * size("rows") + 2) * size("cols") * DT_BYTES, + edge_cases=( + EdgeCase(name="edge_1023", size={"rows": 1023, "cols": 1023}), + EdgeCase(name="edge_single_row", size={"rows": 1, "cols": 1025}), + ), + shape_keys=("rows", "cols"), + shape_aliases={"M": "rows", "N": "cols", "rows": "rows", "cols": "cols"}, + starter_kernels={"pytorch": STARTER_KERNEL}, + speedup_estimate="1.0-1.3x", + backward_spec=BackwardSpec( + differentiable_inputs=("x", "scale", "bias"), + ), +) diff --git a/examples/custom_ops/affine_corpus.json b/examples/custom_ops/affine_corpus.json new file mode 100644 index 00000000..0a76d3d8 --- /dev/null +++ b/examples/custom_ops/affine_corpus.json @@ -0,0 +1,27 @@ +{ + "schema_version": 1, + "operation": "custom_affine", + "cases": [ + { + "name": "prod-prefill", + "size": {"rows": 4096, "cols": 1024}, + "dtype": "float16", + "weight": 7, + "tags": ["production", "forward"] + }, + { + "name": "prod-decode", + "size": {"rows": 64, "cols": 1024}, + "dtype": "float16", + "weight": 93, + "tags": ["production", "forward"] + }, + { + "name": "prod-fp32-fallback", + "size": {"rows": 512, "cols": 512}, + "dtype": "float32", + "weight": 1, + "tags": ["fallback"] + } + ] +} diff --git a/examples/custom_ops/affine_kernel.py b/examples/custom_ops/affine_kernel.py new file mode 100644 index 00000000..7aa8ab80 --- /dev/null +++ b/examples/custom_ops/affine_kernel.py @@ -0,0 +1,25 @@ +"""Starter kernel for the external ``custom_affine`` example operation. + +Copy this file to ``kernel.py`` and benchmark it against the external spec:: + + cp examples/custom_ops/affine_kernel.py kernel.py + uv run bench.py --spec examples/custom_ops/affine.py:SPEC --quick + +The candidate is intentionally plain PyTorch: the fixture proves that a +structured, multi-output operation flows through the harness (including CPU +verification), not that an affine transform can be made faster. Its signature +and output tree must match the spec's reference exactly. +""" + +KERNEL_TYPE = "custom_affine" + +import torch + + +def kernel_fn(x: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor) -> dict: + """Entry point called by bench.py. Must match the reference signature.""" + y = x * scale + bias + residual = ( + x.float() * scale.float() + bias.float() - x.float() + ).to(x.dtype) + return {"output": y, "aux": (residual, 3)} diff --git a/extract.py b/extract.py index 57a59e7e..c11fa3f7 100644 --- a/extract.py +++ b/extract.py @@ -8,6 +8,10 @@ uv run extract.py --kernel-type matmul # extract only matmul kernels uv run extract.py --report path/to/report.json uv run extract.py --backend cuda # use CUDA C++ starter kernels instead of Triton + uv run extract.py --spec path/spec.py:SPEC # extract an external KernelSpec + +Operation metadata (shape aliases, tolerances, FLOP/byte accounting, starter +kernels, speedup estimates) comes from autokernel/specs/, not from this file. """ from __future__ import annotations @@ -29,151 +33,27 @@ DEFAULT_REPORT_PATH = os.path.join(WORKSPACE_DIR, "profile_report.json") OPTIMIZATION_PLAN_PATH = os.path.join(WORKSPACE_DIR, "optimization_plan.json") +# The package lives next to this script; make sure it is importable when +# extract.py is invoked from another working directory. +if SCRIPT_DIR not in sys.path: + sys.path.insert(0, SCRIPT_DIR) -# --------------------------------------------------------------------------- -# Shape key mappings per kernel type -# --------------------------------------------------------------------------- -# Each entry maps op_type -> list of (shape_key_aliases...) so we can parse -# various shape_info string formats from profile_report.json. - -SHAPE_KEYS: Dict[str, List[str]] = { - "matmul": ["M", "N", "K"], - "flash_attention": ["B", "H", "N", "D"], - "layernorm": ["M", "N"], - "softmax": ["M", "N"], - "cross_entropy": ["batch", "vocab"], - "fused_mlp": ["M", "N", "K"], - "rmsnorm": ["M", "N"], - "reduce": ["M", "N"], - "rotary_embedding": ["B", "H", "N", "D"], -} - -# Aliases: profile_report.json may use different key names than bench.py -# Map from alias -> canonical bench.py key, per op_type. -SHAPE_ALIAS_MAP: Dict[str, Dict[str, str]] = { - "matmul": {}, - "flash_attention": { - "B": "batch", "H": "heads", "N": "seq_len", "S": "seq_len", "D": "head_dim", - "batch": "batch", "heads": "heads", "seq_len": "seq_len", "head_dim": "head_dim", - }, - "layernorm": { - "M": "batch", "N": "dim", "rows": "batch", "cols": "dim", - "batch": "batch", "dim": "dim", - }, - "softmax": { - "M": "rows", "N": "cols", "rows": "rows", "cols": "cols", - }, - "cross_entropy": { - "batch": "batch", "vocab": "vocab", - }, - "fused_mlp": { - "M": "batch", "N": "hidden", "K": "dim", - "batch": "batch", "dim": "dim", "hidden": "hidden", - }, - "rmsnorm": { - "M": "M", "N": "N", - }, - "reduce": { - "M": "M", "N": "N", - }, - "rotary_embedding": { - "B": "batch", "H": "heads", "N": "seq_len", "S": "seq_len", "D": "head_dim", - "batch": "batch", "heads": "heads", "seq_len": "seq_len", "head_dim": "head_dim", - }, -} - -# Default tolerances per op_type (matching bench.py structure, serialized for template) -TOLERANCES_MAP: Dict[str, Dict[str, Dict[str, float]]] = { - "matmul": { - "float16": {"atol": 1e-2, "rtol": 1e-2}, - "bfloat16": {"atol": 2e-2, "rtol": 2e-2}, - "float32": {"atol": 1e-4, "rtol": 1e-4}, - }, - "flash_attention": { - "float16": {"atol": 1e-2, "rtol": 1e-2}, - "bfloat16": {"atol": 2e-2, "rtol": 2e-2}, - "float32": {"atol": 1e-4, "rtol": 1e-4}, - }, - "layernorm": { - "float16": {"atol": 1e-3, "rtol": 1e-3}, - "bfloat16": {"atol": 2e-3, "rtol": 2e-3}, - "float32": {"atol": 1e-5, "rtol": 1e-5}, - }, - "softmax": { - "float16": {"atol": 1e-3, "rtol": 1e-3}, - "bfloat16": {"atol": 2e-3, "rtol": 2e-3}, - "float32": {"atol": 1e-5, "rtol": 1e-5}, - }, - "cross_entropy": { - "float16": {"atol": 1e-2, "rtol": 1e-2}, - "bfloat16": {"atol": 2e-2, "rtol": 2e-2}, - "float32": {"atol": 1e-5, "rtol": 1e-5}, - }, - "fused_mlp": { - "float16": {"atol": 1e-2, "rtol": 1e-2}, - "bfloat16": {"atol": 2e-2, "rtol": 2e-2}, - "float32": {"atol": 1e-4, "rtol": 1e-4}, - }, - "rmsnorm": { - "float16": {"atol": 1e-2, "rtol": 1e-2}, - "bfloat16": {"atol": 1e-1, "rtol": 5e-2}, - }, - "reduce": { - "float16": {"atol": 1e-2, "rtol": 1e-2}, - "bfloat16": {"atol": 1e-1, "rtol": 5e-2}, - }, - "rotary_embedding": { - "float16": {"atol": 1e-3, "rtol": 1e-3}, - "bfloat16": {"atol": 2e-3, "rtol": 2e-3}, - "float32": {"atol": 1e-5, "rtol": 1e-5}, - }, -} - -# FLOPS formulas as source strings, per op_type -FLOPS_FN_SRC: Dict[str, str] = { - "matmul": 'return 2 * s["M"] * s["N"] * s["K"]', - "flash_attention": 'return 4 * s["batch"] * s["heads"] * (s["seq_len"] ** 2) * s["head_dim"]', - "layernorm": 'return 8 * s["batch"] * s["dim"]', - "softmax": 'return 5 * s["rows"] * s["cols"]', - "cross_entropy": 'return 4 * s["batch"] * s["vocab"]', - "fused_mlp": 'return 2 * s["batch"] * s["dim"] * s["hidden"] * 3', - "rmsnorm": 'return 6 * s["M"] * s["N"]', - "reduce": 'return s["M"] * s["N"]', - "rotary_embedding": 'return 6 * s["batch"] * s["heads"] * s["seq_len"] * s["head_dim"]', -} - -# BYTES formulas as source strings, per op_type (dt_bytes is passed in) -BYTES_FN_SRC: Dict[str, str] = { - "matmul": 'return (s["M"] * s["K"] + s["K"] * s["N"] + s["M"] * s["N"]) * dt_bytes', - "flash_attention": 'return 4 * s["batch"] * s["heads"] * s["seq_len"] * s["head_dim"] * dt_bytes', - "layernorm": 'return (2 * s["batch"] * s["dim"] + 2 * s["dim"]) * dt_bytes', - "softmax": 'return 2 * s["rows"] * s["cols"] * dt_bytes', - "cross_entropy": 'return (s["batch"] * s["vocab"] + s["batch"]) * dt_bytes', - "fused_mlp": 'return (s["batch"] * s["dim"] + s["hidden"] * s["dim"] * 3 + s["batch"] * s["dim"]) * dt_bytes', - "rmsnorm": 'return (2 * s["M"] * s["N"] + s["N"]) * dt_bytes', - "reduce": 'return (s["M"] * s["N"] + s["M"]) * dt_bytes', - "rotary_embedding": 'return (s["batch"] * s["heads"] * s["seq_len"] * s["head_dim"] * 2 + s["seq_len"] * s["head_dim"]) * dt_bytes', -} - -# Speedup potential heuristic per op_type -SPEEDUP_ESTIMATES: Dict[str, str] = { - "matmul": "2-3x", - "flash_attention": "2-4x", - "layernorm": "1.5-3x", - "softmax": "1.5-3x", - "cross_entropy": "1.5-2x", - "fused_mlp": "2-3x", - "rmsnorm": "1.5-3x", - "reduce": "1.5-2x", - "rotary_embedding": "1.5-2x", -} +from autokernel.specs import ( # noqa: E402 (path bootstrap must run first) + KernelRegistry, + KernelSpec, + SpecLoadError, + SpecValidationError, + create_builtin_registry, + resolve_spec, + serialize_accounting, +) # --------------------------------------------------------------------------- # Shape parsing # --------------------------------------------------------------------------- -def parse_shape_info(shape_info_str: str, op_type: str) -> Optional[Dict[str, int]]: +def parse_shape_info(shape_info_str: str, spec: KernelSpec) -> Optional[Dict[str, int]]: """ Parse a shape_info string like "M=4096, N=4096, K=4096" into a dict. @@ -183,6 +63,9 @@ def parse_shape_info(shape_info_str: str, op_type: str) -> Optional[Dict[str, in - "batch=4096, vocab=32000" - "rows=4096, cols=4096" + Profiler key spellings are mapped to the specification's canonical size keys + through ``spec.shape_aliases``. + Returns None if parsing fails. """ if not shape_info_str or not isinstance(shape_info_str, str): @@ -195,16 +78,15 @@ def parse_shape_info(shape_info_str: str, op_type: str) -> Optional[Dict[str, in raw = {k: int(v) for k, v in pairs} - # Map to canonical bench.py keys using alias map - alias_map = SHAPE_ALIAS_MAP.get(op_type, {}) + # Map to the spec's canonical size keys using its alias map + alias_map = spec.shape_aliases if alias_map: canonical = {} for k, v in raw.items(): mapped_key = alias_map.get(k, k) canonical[mapped_key] = v return canonical - else: - return raw + return raw def shape_to_display(shape: Dict[str, int]) -> str: @@ -220,43 +102,32 @@ def scale_shape(shape: Dict[str, int], factor: float) -> Dict[str, int]: return {k: max(1, int(round(v * factor))) for k, v in shape.items()} -def get_default_shape(op_type: str) -> Dict[str, int]: - """ - Return a reasonable default shape for a given op_type when parsing fails. - Based on the 'large' size from bench.py KERNEL_CONFIGS. - """ - defaults: Dict[str, Dict[str, int]] = { - "matmul": {"M": 2048, "N": 2048, "K": 2048}, - "flash_attention": {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 64}, - "layernorm": {"batch": 4096, "dim": 2048}, - "softmax": {"rows": 4096, "cols": 4096}, - "cross_entropy": {"batch": 4096, "vocab": 32000}, - "fused_mlp": {"batch": 2048, "dim": 2048, "hidden": 5504}, - "rmsnorm": {"M": 4096, "N": 4096}, - "reduce": {"M": 4096, "N": 4096}, - "rotary_embedding": {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 128}, - } - return defaults.get(op_type, {"M": 2048, "N": 2048}) +def get_default_shape(spec: KernelSpec) -> Dict[str, int]: + """Fallback shape for a specification when a profiled shape cannot be parsed.""" + return spec.extraction_shape() # --------------------------------------------------------------------------- # Kernel file generation # --------------------------------------------------------------------------- -def read_starter_kernel(op_type: str, backend: str = "triton") -> Optional[str]: - """Read the starter kernel file. Returns None if not found. - - For backend='triton': reads from kernels/{op_type}.py - For backend='cuda': reads from kernels/cuda/{op_type}.py - """ - if backend == "cuda": - path = os.path.join(KERNELS_DIR, "cuda", f"{op_type}.py") - else: - path = os.path.join(KERNELS_DIR, f"{op_type}.py") - if not os.path.exists(path): +def read_starter_kernel(spec: KernelSpec, backend: str = "triton") -> Optional[str]: + """Read the starter kernel declared by a specification. None if absent.""" + path = spec.starter_kernel(backend) + if path is None or not path.is_file(): return None - with open(path, "r", encoding="utf-8") as f: - return f.read() + return path.read_text(encoding="utf-8") + + +def starter_kernel_display(spec: KernelSpec, backend: str = "triton") -> str: + """Repository-relative starter kernel path, for logs and generated headers.""" + path = spec.starter_kernel(backend) + if path is None: + return f"" + try: + return str(path.relative_to(SCRIPT_DIR)) + except ValueError: + return str(path) def extract_kernel_body(starter_code: str) -> str: @@ -287,8 +158,35 @@ def extract_kernel_body(starter_code: str) -> str: return starter_code +def _accounting_body(fn: object, spec: KernelSpec, spec_locator: Optional[str]) -> str: + """Return the body of a generated accounting function. + + Accounting :class:`~autokernel.specs.accounting.Expression` objects serialize + to a plain numeric expression. Opaque callables (an external spec supplying + an ordinary Python function) cannot be serialized safely, so the generated + file resolves them from the specification at runtime instead of this script + guessing a formula or ``eval``-ing spec content. + """ + source = serialize_accounting(fn) + if source is not None: + return f"return {source}" + if spec_locator: + return ( + f"raise NotImplementedError(\n" + f" 'accounting for {spec.name} is not serializable; load it from '\n" + f" \"the spec: autokernel.specs.load_spec({spec_locator!r})\"\n" + f" )" + ) + return ( + f"raise NotImplementedError(\n" + f" 'accounting for {spec.name} is not serializable; read it from the '\n" + f" 'registered KernelSpec instead'\n" + f" )" + ) + + def generate_kernel_file( - op_type: str, + spec: KernelSpec, rank: int, pct_total: float, model_shape: Dict[str, int], @@ -296,9 +194,11 @@ def generate_kernel_file( gpu_time_ms: float, starter_code: str, backend: str = "triton", + spec_locator: Optional[str] = None, ) -> str: """Generate the complete kernel file content for extraction.""" + op_type = spec.name half_shape = scale_shape(model_shape, 0.5) double_shape = scale_shape(model_shape, 2.0) @@ -306,14 +206,12 @@ def generate_kernel_file( half_display = shape_to_display(half_shape) double_display = shape_to_display(double_shape) - tolerances = TOLERANCES_MAP.get(op_type, { - "float16": {"atol": 1e-2, "rtol": 1e-2}, - "bfloat16": {"atol": 2e-2, "rtol": 2e-2}, - "float32": {"atol": 1e-4, "rtol": 1e-4}, - }) + tolerances = { + dtype: tol.as_dict() for dtype, tol in spec.tolerances.items() + } - flops_fn_body = FLOPS_FN_SRC.get(op_type, 'return 0') - bytes_fn_body = BYTES_FN_SRC.get(op_type, 'return 0') + flops_fn_body = _accounting_body(spec.flops_fn, spec, spec_locator) + bytes_fn_body = _accounting_body(spec.bytes_fn, spec, spec_locator) # Extract the kernel code body (imports + jit functions + kernel_fn) kernel_body = extract_kernel_body(starter_code) @@ -337,6 +235,8 @@ def generate_kernel_file( lines.append(f'KERNEL_TYPE = "{op_type}"') if backend == "cuda": lines.append(f'BACKEND = "cuda"') + if spec_locator: + lines.append(f'KERNEL_SPEC = "{spec_locator}"') lines.append("") # Model-specific shapes @@ -374,7 +274,7 @@ def generate_kernel_file( lines.append("") lines.append(f"# {'=' * 70}") backend_label = "CUDA C++" if backend == "cuda" else "Triton" - backend_dir = f"kernels/cuda/{op_type}.py" if backend == "cuda" else f"kernels/{op_type}.py" + backend_dir = starter_kernel_display(spec, backend) lines.append(f"# {backend_label} kernel code (from {backend_dir})") lines.append(f"# {'=' * 70}") lines.append("") @@ -443,8 +343,8 @@ def generate_optimization_plan( "model_shape": entry["model_shape"], "gpu_time_ms": entry["gpu_time_ms"], "pct_total": entry["pct_total"], - "estimated_speedup_potential": SPEEDUP_ESTIMATES.get( - entry["op_type"], "1.5-2x" + "estimated_speedup_potential": entry.get( + "estimated_speedup_potential", "1.5-2x" ), }) @@ -459,11 +359,30 @@ def generate_optimization_plan( # Main extraction logic # --------------------------------------------------------------------------- +def _synthetic_report_entry(spec: KernelSpec) -> Dict[str, Any]: + """A single extraction target derived from a specification alone. + + Used when ``--spec`` is supplied without a matching profile report entry, so + an external operation can be turned into a starter kernel file without + profiling a model first. + """ + return { + "rank": 1, + "op_type": spec.name, + "pct_total": 0.0, + "gpu_time_ms": 0.0, + "shapes": spec.extraction_shape(), + "autokernel_supported": True, + } + + def extract_kernels( report_path: str, top_n: Optional[int] = None, kernel_type_filter: Optional[str] = None, backend: str = "triton", + spec_locator: Optional[str] = None, + spec_override: bool = False, ) -> None: """Main extraction pipeline.""" @@ -471,30 +390,57 @@ def extract_kernels( print(f"=== AutoKernel Kernel Extractor ({backend_label}) ===") print() + # -- Resolve operation specifications --------------------------------- + # Precedence: --spec, then --kernel-type, then whatever the report names. + registry: KernelRegistry = create_builtin_registry() + external_spec: Optional[KernelSpec] = None + if spec_locator: + try: + external_spec, registry = resolve_spec( + spec_locator=spec_locator, registry=registry, override=spec_override + ) + except (SpecLoadError, SpecValidationError) as e: + print(f"ERROR: {e}") + sys.exit(1) + kernel_type_filter = external_spec.name + print(f"Using external kernel spec: {spec_locator} (operation " + f"'{external_spec.name}')") + print() + # -- Load profile report -- print(f"Reading profile from {report_path}...") report = load_profile_report(report_path) if report is None: - print(f"ERROR: Profile report not found at {report_path}") - print(f" Run the profiler first: uv run profile.py") - sys.exit(1) + if external_spec is None: + print(f"ERROR: Profile report not found at {report_path}") + print(f" Run the profiler first: uv run profile.py") + sys.exit(1) + print(f" No profile report at {report_path}; extracting " + f"'{external_spec.name}' from its specification instead.") + report = {"model_name": f"{external_spec.name} spec", "top_kernels": []} # -- Get model name -- model_name = report.get("model_name", report.get("model", "unknown model")) # -- Get supported kernels -- supported = get_supported_kernels(report) - if not supported: - print("ERROR: No supported kernels found in profile report.") - print(" Ensure the profiler marks kernels with autokernel_supported=True.") - sys.exit(1) # -- Apply filters -- if kernel_type_filter: supported = [k for k in supported if k.get("op_type") == kernel_type_filter] - if not supported: + + if not supported: + if external_spec is not None: + supported = [_synthetic_report_entry(external_spec)] + print(f" Profile report has no '{external_spec.name}' entries; using the " + f"specification's default shape.") + elif kernel_type_filter: print(f"WARNING: No kernels of type '{kernel_type_filter}' found in profile report.") sys.exit(1) + else: + print("ERROR: No supported kernels found in profile report.") + print(" Ensure the profiler marks kernels with autokernel_supported=True.") + sys.exit(1) if top_n is not None: supported = supported[:top_n] @@ -517,8 +463,17 @@ def extract_kernels( gpu_time_ms = kernel_info.get("gpu_time_ms", kernel_info.get("total_gpu_time_ms", 0.0)) shape_info_str = kernel_info.get("shape_info", kernel_info.get("shape", "")) + # Look up the specification that owns this operation + if not registry.contains(op_type): + print(f" WARNING: No kernel specification registered for '{op_type}' " + f"-- skipping. Registered: {', '.join(registry.list_names())}") + skipped += 1 + continue + spec = registry.get(op_type) + entry_locator = spec_locator if external_spec is not None and spec is external_spec else None + # Parse model shape - model_shape = parse_shape_info(shape_info_str, op_type) + model_shape = parse_shape_info(shape_info_str, spec) if model_shape is None: # Try to use a "shapes" dict directly if provided if isinstance(kernel_info.get("shapes"), dict): @@ -526,13 +481,13 @@ def extract_kernels( else: print(f" WARNING: Could not parse shape for {op_type} (rank {rank}), " f"using default shapes.") - model_shape = get_default_shape(op_type) + model_shape = get_default_shape(spec) # Read starter kernel - starter_code = read_starter_kernel(op_type, backend=backend) + starter_code = read_starter_kernel(spec, backend=backend) if starter_code is None: - starter_dir = "kernels/cuda" if backend == "cuda" else "kernels" - print(f" WARNING: No starter kernel found at {starter_dir}/{op_type}.py -- skipping.") + print(f" WARNING: No {backend} starter kernel declared by the " + f"'{op_type}' spec -- skipping.") skipped += 1 continue @@ -544,7 +499,7 @@ def extract_kernels( # Generate the customized kernel file kernel_content = generate_kernel_file( - op_type=op_type, + spec=spec, rank=rank, pct_total=pct_total, model_shape=model_shape, @@ -552,6 +507,7 @@ def extract_kernels( gpu_time_ms=gpu_time_ms, starter_code=starter_code, backend=backend, + spec_locator=entry_locator, ) # Write to workspace @@ -565,8 +521,7 @@ def extract_kernels( print(f" [{position}/{total}] {op_type} (rank {rank}, {pct_total}%) " f"-> {output_relpath}") print(f" Model shape: {shape_display}") - starter_dir = "kernels/cuda" if backend == "cuda" else "kernels" - print(f" Based on: {starter_dir}/{op_type}.py") + print(f" Based on: {starter_kernel_display(spec, backend)}") print() extracted.append({ @@ -576,12 +531,16 @@ def extract_kernels( "gpu_time_ms": gpu_time_ms, "model_shape": model_shape, "output_file": output_relpath, + "estimated_speedup_potential": spec.speedup_estimate or "1.5-2x", }) if not extracted: print("ERROR: No kernels were successfully extracted.") if skipped > 0: - print(f" {skipped} kernel(s) skipped due to missing starter files.") + print( + f" {skipped} kernel(s) skipped " + "(missing spec or starter kernel)." + ) sys.exit(1) # -- Generate optimization plan -- @@ -633,6 +592,18 @@ def main() -> None: default="triton", help="Backend for starter kernels: 'triton' (default) or 'cuda' (native CUDA C++)", ) + parser.add_argument( + "--spec", + type=str, + default=None, + help="External KernelSpec locator, e.g. 'path/to/spec.py:SPEC' or " + "'package.module:SPEC'. Takes precedence over --kernel-type.", + ) + parser.add_argument( + "--spec-override", + action="store_true", + help="Allow --spec to replace a built-in operation of the same name", + ) args = parser.parse_args() @@ -641,6 +612,8 @@ def main() -> None: top_n=args.top, kernel_type_filter=args.kernel_type, backend=args.backend, + spec_locator=args.spec, + spec_override=args.spec_override, ) diff --git a/kernels/layernorm.py b/kernels/layernorm.py index d0d237c3..9835bb53 100644 --- a/kernels/layernorm.py +++ b/kernels/layernorm.py @@ -12,6 +12,7 @@ KERNEL_TYPE = "layernorm" import torch +import torch.nn.functional as F import triton import triton.language as tl @@ -71,6 +72,13 @@ def kernel_fn( """Entry point called by bench.py. Must match reference.layernorm_ref signature.""" assert x.is_cuda + # The strict BF16 gate requires bit-compatible rounding with PyTorch. + # This starter's simple two-pass reduction can differ by one BF16 ULP on + # wide rows, so keep the reference-quality fallback until a Welford Triton + # implementation replaces it. Other dtypes use the custom kernel below. + if x.dtype == torch.bfloat16: + return F.layer_norm(x, (x.shape[-1],), weight, bias, eps) + # Flatten to 2D for row-parallel processing orig_shape = x.shape if x.ndim == 1: diff --git a/kernels/matmul.py b/kernels/matmul.py index 246fe198..79fe170e 100644 --- a/kernels/matmul.py +++ b/kernels/matmul.py @@ -48,7 +48,11 @@ def matmul_kernel( for k in range(0, K, BLOCK_SIZE_K): a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0) b = tl.load(b_ptrs, mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0) - acc += tl.dot(a, b) + # Triton defaults float32 dot products to TF32 on NVIDIA GPUs. The + # specification declares true float32 semantics with a 1e-4 + # correctness gate, so request IEEE inputs explicitly. FP16/BF16 + # tensor-core paths are unaffected. + acc += tl.dot(a, b, input_precision="ieee") a_ptrs += BLOCK_SIZE_K * stride_ak b_ptrs += BLOCK_SIZE_K * stride_bk offs_k += BLOCK_SIZE_K diff --git a/kernels/wan_gated_residual.py b/kernels/wan_gated_residual.py new file mode 100644 index 00000000..38d5c173 --- /dev/null +++ b/kernels/wan_gated_residual.py @@ -0,0 +1,69 @@ +"""Triton baseline for Wan's post-MLP gated residual update.""" + +KERNEL_TYPE = "wan_gated_residual" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _wan_gated_residual_kernel( + residual_ptr, + x_ptr, + gate_ptr, + output_ptr, + tokens, + hidden: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + batch = row // tokens + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden + row_offsets = row * hidden + offsets + gate_offsets = batch * hidden + offsets + residual = tl.load( + residual_ptr + row_offsets, mask=mask, other=0.0 + ).to(tl.float32) + x = tl.load(x_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(gate_ptr + gate_offsets, mask=mask, other=0.0).to( + tl.float32 + ) + tl.store(output_ptr + row_offsets, residual + x * gate, mask=mask) + + +def kernel_fn( + residual: torch.Tensor, + x: torch.Tensor, + gate: torch.Tensor, +) -> torch.Tensor: + """Fuse Wan's MLP gate and residual update into one launch.""" + if not residual.is_cuda: + raise ValueError("wan_gated_residual requires CUDA tensors") + if residual.ndim != 3 or x.shape != residual.shape: + raise ValueError("residual and x must have matching [B, S, D] shapes") + if any(tensor.device != residual.device for tensor in (x, gate)): + raise ValueError("all inputs must be on the residual device") + if not all(tensor.is_contiguous() for tensor in (residual, x, gate)): + raise ValueError("all inputs must be contiguous") + batch, tokens, hidden = residual.shape + if gate.shape != (batch, hidden): + raise ValueError(f"gate must have shape {(batch, hidden)}") + if hidden > 65536: + raise ValueError("hidden dimension exceeds the Triton baseline limit") + + output = torch.empty_like(residual) + block_size = triton.next_power_of_2(hidden) + num_warps = 4 if block_size <= 2048 else 8 + _wan_gated_residual_kernel[(batch * tokens,)]( + residual, + x, + gate, + output, + tokens=tokens, + hidden=hidden, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return output diff --git a/kernels/wan_gated_residual_norm.py b/kernels/wan_gated_residual_norm.py new file mode 100644 index 00000000..c44418f5 --- /dev/null +++ b/kernels/wan_gated_residual_norm.py @@ -0,0 +1,105 @@ +"""Triton baseline for Wan's gated residual + affine LayerNorm transition.""" + +KERNEL_TYPE = "wan_gated_residual_norm" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _wan_gated_residual_norm_kernel( + residual_ptr, + x_ptr, + gate_ptr, + weight_ptr, + bias_ptr, + normalized_ptr, + updated_ptr, + tokens, + hidden: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + batch = row // tokens + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden + row_offsets = row * hidden + offsets + channel_offsets = batch * hidden + offsets + + residual = tl.load( + residual_ptr + row_offsets, mask=mask, other=0.0 + ).to(tl.float32) + x = tl.load(x_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) + gate = tl.load( + gate_ptr + channel_offsets, mask=mask, other=0.0 + ).to(tl.float32) + updated = residual + x * gate + + mean = tl.sum(updated, axis=0) / hidden + centered = tl.where(mask, updated - mean, 0.0) + variance = tl.sum(centered * centered, axis=0) / hidden + normalized = centered * tl.rsqrt(variance + 1e-6) + + weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + bias = tl.load(bias_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + normalized = normalized * weight + bias + + tl.store(normalized_ptr + row_offsets, normalized, mask=mask) + tl.store(updated_ptr + row_offsets, updated, mask=mask) + + +def kernel_fn( + residual: torch.Tensor, + x: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse Wan's post-self-attention transition into one GPU launch.""" + if not residual.is_cuda: + raise ValueError("wan_gated_residual_norm requires CUDA tensors") + if residual.ndim != 3 or x.shape != residual.shape: + raise ValueError("residual and x must have matching [B, S, D] shapes") + if not residual.is_contiguous() or not x.is_contiguous(): + raise ValueError("residual and x must be contiguous") + if any( + tensor.device != residual.device + for tensor in (x, gate, weight, bias) + ): + raise ValueError("all inputs must be on the residual device") + if ( + not gate.is_contiguous() + or not weight.is_contiguous() + or not bias.is_contiguous() + ): + raise ValueError("gate, weight, and bias must be contiguous") + + batch, tokens, hidden = residual.shape + if gate.shape != (batch, hidden): + raise ValueError(f"gate must have shape {(batch, hidden)}") + if weight.shape != (hidden,) or bias.shape != (hidden,): + raise ValueError(f"weight and bias must have shape {(hidden,)}") + if hidden > 65536: + raise ValueError("hidden dimension exceeds the Triton baseline limit") + + normalized = torch.empty_like(residual) + updated = torch.empty_like(residual) + block_size = triton.next_power_of_2(hidden) + num_warps = 4 if block_size <= 2048 else 8 + _wan_gated_residual_norm_kernel[(batch * tokens,)]( + residual, + x, + gate, + weight, + bias, + normalized, + updated, + tokens=tokens, + hidden=hidden, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return normalized, updated diff --git a/kernels/wan_modulated_layer_norm.py b/kernels/wan_modulated_layer_norm.py new file mode 100644 index 00000000..30813882 --- /dev/null +++ b/kernels/wan_modulated_layer_norm.py @@ -0,0 +1,75 @@ +"""Triton baseline for Wan's modulated pre-attention LayerNorm.""" + +KERNEL_TYPE = "wan_modulated_layer_norm" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _wan_modulated_layer_norm_kernel( + x_ptr, + scale_ptr, + shift_ptr, + output_ptr, + tokens, + hidden: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + batch = row // tokens + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden + row_offsets = row * hidden + offsets + modulation_offsets = batch * hidden + offsets + + x = tl.load(x_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) + mean = tl.sum(x, axis=0) / hidden + centered = tl.where(mask, x - mean, 0.0) + variance = tl.sum(centered * centered, axis=0) / hidden + normalized = centered * tl.rsqrt(variance + 1e-6) + scale = tl.load( + scale_ptr + modulation_offsets, mask=mask, other=0.0 + ).to(tl.float32) + shift = tl.load( + shift_ptr + modulation_offsets, mask=mask, other=0.0 + ).to(tl.float32) + output = normalized * (1.0 + scale) + shift + tl.store(output_ptr + row_offsets, output, mask=mask) + + +def kernel_fn( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + """Fuse Wan's non-affine LayerNorm and FP32 modulation.""" + if not x.is_cuda: + raise ValueError("wan_modulated_layer_norm requires CUDA tensors") + if x.ndim != 3: + raise ValueError("x must have shape [B, S, D]") + if any(tensor.device != x.device for tensor in (scale, shift)): + raise ValueError("all inputs must be on the x device") + if not all(tensor.is_contiguous() for tensor in (x, scale, shift)): + raise ValueError("all inputs must be contiguous") + batch, tokens, hidden = x.shape + if scale.shape != (batch, hidden) or shift.shape != (batch, hidden): + raise ValueError(f"scale and shift must have shape {(batch, hidden)}") + if hidden > 65536: + raise ValueError("hidden dimension exceeds the Triton baseline limit") + + output = torch.empty_like(x) + block_size = triton.next_power_of_2(hidden) + num_warps = 4 if block_size <= 2048 else 8 + _wan_modulated_layer_norm_kernel[(batch * tokens,)]( + x, + scale, + shift, + output, + tokens=tokens, + hidden=hidden, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return output diff --git a/models/wan_gated_residual.py b/models/wan_gated_residual.py new file mode 100644 index 00000000..899d7fad --- /dev/null +++ b/models/wan_gated_residual.py @@ -0,0 +1,111 @@ +"""Wan post-MLP gated residual fusion specification.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from autokernel.specs import ( + DT_BYTES, + EdgeCase, + KernelSpec, + Tolerance, + resolve_torch_dtype, + size, +) + +_HERE = Path(__file__).resolve().parent +STARTER_KERNEL = _HERE.parent / "kernels" / "wan_gated_residual.py" + + +def wan_gated_residual_ref(residual: Any, x: Any, gate: Any) -> Any: + """Apply Wan's FP32 gated residual update and model-dtype boundary.""" + return ( + residual.float() + x.float() * gate[:, None, :].float() + ).to(residual.dtype) + + +def gen_wan_gated_residual_inputs( + size_map: Mapping[str, int], + dtype: Any, + device: str, + seed: int = 42, +) -> dict[str, Any]: + """Generate deterministic residual inputs and an FP32 gate.""" + import torch + + generator = torch.Generator(device=device) + generator.manual_seed(seed) + torch_dtype = resolve_torch_dtype(dtype) + batch = size_map["batch"] + tokens = size_map["tokens"] + hidden = size_map["hidden"] + return { + "residual": torch.randn( + batch, + tokens, + hidden, + device=device, + dtype=torch_dtype, + generator=generator, + ), + "x": torch.randn( + batch, + tokens, + hidden, + device=device, + dtype=torch_dtype, + generator=generator, + ), + "gate": torch.randn( + batch, + hidden, + device=device, + dtype=torch.float32, + generator=generator, + ), + } + + +SPEC = KernelSpec( + name="wan_gated_residual", + reference_fn=wan_gated_residual_ref, + input_generator=gen_wan_gated_residual_inputs, + sizes={ + "small": {"batch": 1, "tokens": 20280, "hidden": 1536}, + "medium": {"batch": 1, "tokens": 32760, "hidden": 1536}, + "large": {"batch": 1, "tokens": 20280, "hidden": 5120}, + }, + dtypes=("bfloat16", "float16"), + tolerances={ + "bfloat16": Tolerance(atol=2e-2, rtol=2e-2), + "float16": Tolerance(atol=3e-3, rtol=3e-3), + }, + flops_fn=2 * size("batch") * size("tokens") * size("hidden"), + bytes_fn=( + 3 * size("batch") * size("tokens") * size("hidden") * DT_BYTES + + 4 * size("batch") * size("hidden") + ), + edge_cases=( + EdgeCase( + name="non_power_of_two", + size={"batch": 1, "tokens": 257, "hidden": 1537}, + ), + EdgeCase( + name="batched", + size={"batch": 2, "tokens": 511, "hidden": 1536}, + ), + ), + shape_keys=("batch", "tokens", "hidden"), + shape_aliases={ + "B": "batch", + "S": "tokens", + "D": "hidden", + "batch": "batch", + "tokens": "tokens", + "hidden": "hidden", + }, + starter_kernels={"triton": STARTER_KERNEL}, + speedup_estimate="1.2-2x versus eager PyTorch", +) diff --git a/models/wan_gated_residual_corpus.json b/models/wan_gated_residual_corpus.json new file mode 100644 index 00000000..30458cee --- /dev/null +++ b/models/wan_gated_residual_corpus.json @@ -0,0 +1,41 @@ +{ + "schema_version": 1, + "operation": "wan_gated_residual", + "cases": [ + { + "name": "wan2.1-1.3b-480p-49f", + "size": {"batch": 1, "tokens": 20280, "hidden": 1536}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "1.3b", "480p", "single-gpu"] + }, + { + "name": "wan2.1-1.3b-480p-81f", + "size": {"batch": 1, "tokens": 32760, "hidden": 1536}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "1.3b", "480p", "single-gpu"] + }, + { + "name": "wan2.1-14b-480p-49f", + "size": {"batch": 1, "tokens": 20280, "hidden": 5120}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "14b", "480p", "single-gpu"] + }, + { + "name": "wan2.1-14b-480p-49f-sp4", + "size": {"batch": 1, "tokens": 5070, "hidden": 5120}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "14b", "480p", "sequence-parallel-4"] + }, + { + "name": "wan2.1-14b-480p-81f-sp4", + "size": {"batch": 1, "tokens": 8190, "hidden": 5120}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "14b", "480p", "sequence-parallel-4"] + } + ] +} diff --git a/models/wan_gated_residual_norm.py b/models/wan_gated_residual_norm.py new file mode 100644 index 00000000..d3832504 --- /dev/null +++ b/models/wan_gated_residual_norm.py @@ -0,0 +1,158 @@ +"""Wan self-attention residual/normalization fusion specification. + +This models the transition used after self-attention in every Wan transformer +block: + +1. apply the per-channel attention gate and update the residual stream; +2. compute affine LayerNorm in FP32; and +3. cast both returned streams back to the model dtype. + +The production shapes cover Wan 2.1 1.3B and 14B at common 480p token counts. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Mapping + +from autokernel.specs import ( + DT_BYTES, + EdgeCase, + KernelSpec, + Tolerance, + resolve_torch_dtype, + size, +) + +_HERE = Path(__file__).resolve().parent +STARTER_KERNEL = _HERE.parent / "kernels" / "wan_gated_residual_norm.py" + + +def wan_gated_residual_norm_ref( + residual: Any, + x: Any, + gate: Any, + weight: Any, + bias: Any, +) -> tuple[Any, Any]: + """Return ``(normalized, updated_residual)`` with Wan's dtype boundaries.""" + import torch.nn.functional as F + + output_dtype = residual.dtype + updated_fp32 = residual.float() + x.float() * gate[:, None, :].float() + normalized_fp32 = F.layer_norm( + updated_fp32, + (updated_fp32.shape[-1],), + weight.float(), + bias.float(), + 1e-6, + ) + return normalized_fp32.to(output_dtype), updated_fp32.to(output_dtype) + + +def gen_wan_gated_residual_norm_inputs( + size_map: Mapping[str, int], + dtype: Any, + device: str, + seed: int = 42, +) -> dict[str, Any]: + """Generate deterministic model-dtype activations and FP32 modulation.""" + import torch + + generator = torch.Generator(device=device) + generator.manual_seed(seed) + torch_dtype = resolve_torch_dtype(dtype) + batch = size_map["batch"] + tokens = size_map["tokens"] + hidden = size_map["hidden"] + residual = torch.randn( + batch, + tokens, + hidden, + device=device, + dtype=torch_dtype, + generator=generator, + ) + x = torch.randn( + batch, + tokens, + hidden, + device=device, + dtype=torch_dtype, + generator=generator, + ) + gate = torch.randn( + batch, + hidden, + device=device, + dtype=torch.float32, + generator=generator, + ) + weight = torch.randn( + hidden, + device=device, + dtype=torch.float32, + generator=generator, + ) + bias = torch.randn( + hidden, + device=device, + dtype=torch.float32, + generator=generator, + ) + return { + "residual": residual, + "x": x, + "gate": gate, + "weight": weight, + "bias": bias, + } + + +SPEC = KernelSpec( + name="wan_gated_residual_norm", + reference_fn=wan_gated_residual_norm_ref, + input_generator=gen_wan_gated_residual_norm_inputs, + sizes={ + # Wan 2.1 T2V 1.3B, 480p, common 49-frame token count. + "small": {"batch": 1, "tokens": 20280, "hidden": 1536}, + # Wan 2.1 T2V 1.3B, 480p, common 81-frame token count. + "medium": {"batch": 1, "tokens": 32760, "hidden": 1536}, + # Wan 2.1 14B, 480p. Sequence parallelism reduces tokens per rank. + "large": {"batch": 1, "tokens": 20280, "hidden": 5120}, + }, + dtypes=("bfloat16", "float16"), + tolerances={ + "bfloat16": Tolerance(atol=2e-2, rtol=2e-2), + "float16": Tolerance(atol=3e-3, rtol=3e-3), + }, + # Residual update, mean/variance, normalization, and affine transform. + flops_fn=10 * size("batch") * size("tokens") * size("hidden"), + # Activations: residual + x reads and two output writes. Gate, weight, and + # bias are FP32 vectors and are amortized across the token dimension. + bytes_fn=( + 4 * size("batch") * size("tokens") * size("hidden") * DT_BYTES + + 4 * (size("batch") + 2) * size("hidden") + ), + edge_cases=( + EdgeCase( + name="non_power_of_two", + size={"batch": 1, "tokens": 257, "hidden": 1537}, + ), + EdgeCase( + name="batched", + size={"batch": 2, "tokens": 511, "hidden": 1536}, + ), + ), + shape_keys=("batch", "tokens", "hidden"), + shape_aliases={ + "B": "batch", + "S": "tokens", + "D": "hidden", + "batch": "batch", + "tokens": "tokens", + "hidden": "hidden", + }, + starter_kernels={"triton": STARTER_KERNEL}, + speedup_estimate="7.9-9.1x on GB200 versus eager PyTorch", +) diff --git a/models/wan_gated_residual_norm_corpus.json b/models/wan_gated_residual_norm_corpus.json new file mode 100644 index 00000000..b95612d1 --- /dev/null +++ b/models/wan_gated_residual_norm_corpus.json @@ -0,0 +1,41 @@ +{ + "schema_version": 1, + "operation": "wan_gated_residual_norm", + "cases": [ + { + "name": "wan2.1-1.3b-480p-49f", + "size": {"batch": 1, "tokens": 20280, "hidden": 1536}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "1.3b", "480p", "single-gpu"] + }, + { + "name": "wan2.1-1.3b-480p-81f", + "size": {"batch": 1, "tokens": 32760, "hidden": 1536}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "1.3b", "480p", "single-gpu"] + }, + { + "name": "wan2.1-14b-480p-49f", + "size": {"batch": 1, "tokens": 20280, "hidden": 5120}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "14b", "480p", "single-gpu"] + }, + { + "name": "wan2.1-14b-480p-49f-sp4", + "size": {"batch": 1, "tokens": 5070, "hidden": 5120}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "14b", "480p", "sequence-parallel-4"] + }, + { + "name": "wan2.1-14b-480p-81f-sp4", + "size": {"batch": 1, "tokens": 8190, "hidden": 5120}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "14b", "480p", "sequence-parallel-4"] + } + ] +} diff --git a/models/wan_modulated_layer_norm.py b/models/wan_modulated_layer_norm.py new file mode 100644 index 00000000..6ee173e5 --- /dev/null +++ b/models/wan_modulated_layer_norm.py @@ -0,0 +1,117 @@ +"""Wan pre-attention LayerNorm and scale/shift fusion specification.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from autokernel.specs import ( + DT_BYTES, + EdgeCase, + KernelSpec, + Tolerance, + resolve_torch_dtype, + size, +) + +_HERE = Path(__file__).resolve().parent +STARTER_KERNEL = _HERE.parent / "kernels" / "wan_modulated_layer_norm.py" + + +def wan_modulated_layer_norm_ref( + x: Any, + scale: Any, + shift: Any, +) -> Any: + """Apply Wan's FP32 non-affine LayerNorm and modulation boundary.""" + import torch.nn.functional as F + + normalized = F.layer_norm(x.float(), (x.shape[-1],), None, None, 1e-6) + return (normalized * (1.0 + scale[:, None, :]) + shift[:, None, :]).to( + x.dtype + ) + + +def gen_wan_modulated_layer_norm_inputs( + size_map: Mapping[str, int], + dtype: Any, + device: str, + seed: int = 42, +) -> dict[str, Any]: + """Generate deterministic Wan activations with FP32 modulation.""" + import torch + + generator = torch.Generator(device=device) + generator.manual_seed(seed) + torch_dtype = resolve_torch_dtype(dtype) + batch = size_map["batch"] + tokens = size_map["tokens"] + hidden = size_map["hidden"] + return { + "x": torch.randn( + batch, + tokens, + hidden, + device=device, + dtype=torch_dtype, + generator=generator, + ), + "scale": torch.randn( + batch, + hidden, + device=device, + dtype=torch.float32, + generator=generator, + ), + "shift": torch.randn( + batch, + hidden, + device=device, + dtype=torch.float32, + generator=generator, + ), + } + + +SPEC = KernelSpec( + name="wan_modulated_layer_norm", + reference_fn=wan_modulated_layer_norm_ref, + input_generator=gen_wan_modulated_layer_norm_inputs, + sizes={ + "small": {"batch": 1, "tokens": 20280, "hidden": 1536}, + "medium": {"batch": 1, "tokens": 32760, "hidden": 1536}, + "large": {"batch": 1, "tokens": 20280, "hidden": 5120}, + }, + dtypes=("bfloat16", "float16"), + tolerances={ + "bfloat16": Tolerance(atol=2e-2, rtol=2e-2), + "float16": Tolerance(atol=3e-3, rtol=3e-3), + }, + flops_fn=9 * size("batch") * size("tokens") * size("hidden"), + bytes_fn=( + 2 * size("batch") * size("tokens") * size("hidden") * DT_BYTES + + 8 * size("batch") * size("hidden") + ), + edge_cases=( + EdgeCase( + name="non_power_of_two", + size={"batch": 1, "tokens": 257, "hidden": 1537}, + ), + EdgeCase( + name="batched", + size={"batch": 2, "tokens": 511, "hidden": 1536}, + ), + ), + shape_keys=("batch", "tokens", "hidden"), + shape_aliases={ + "B": "batch", + "S": "tokens", + "D": "hidden", + "batch": "batch", + "tokens": "tokens", + "hidden": "hidden", + }, + starter_kernels={"triton": STARTER_KERNEL}, + speedup_estimate="1.5-3x versus eager PyTorch", +) diff --git a/models/wan_modulated_layer_norm_corpus.json b/models/wan_modulated_layer_norm_corpus.json new file mode 100644 index 00000000..81351d4b --- /dev/null +++ b/models/wan_modulated_layer_norm_corpus.json @@ -0,0 +1,41 @@ +{ + "schema_version": 1, + "operation": "wan_modulated_layer_norm", + "cases": [ + { + "name": "wan2.1-1.3b-480p-49f", + "size": {"batch": 1, "tokens": 20280, "hidden": 1536}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "1.3b", "480p", "single-gpu"] + }, + { + "name": "wan2.1-1.3b-480p-81f", + "size": {"batch": 1, "tokens": 32760, "hidden": 1536}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "1.3b", "480p", "single-gpu"] + }, + { + "name": "wan2.1-14b-480p-49f", + "size": {"batch": 1, "tokens": 20280, "hidden": 5120}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "14b", "480p", "single-gpu"] + }, + { + "name": "wan2.1-14b-480p-49f-sp4", + "size": {"batch": 1, "tokens": 5070, "hidden": 5120}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "14b", "480p", "sequence-parallel-4"] + }, + { + "name": "wan2.1-14b-480p-81f-sp4", + "size": {"batch": 1, "tokens": 8190, "hidden": 5120}, + "dtype": "bfloat16", + "weight": 40, + "tags": ["wan2.1", "14b", "480p", "sequence-parallel-4"] + } + ] +} diff --git a/profile.py b/profile.py index ba7276f0..3933e623 100644 --- a/profile.py +++ b/profile.py @@ -30,6 +30,59 @@ import torch import torch.nn as nn +# This command is intentionally named ``profile.py`` for CLI compatibility, +# but that name also shadows Python's standard-library ``profile`` module when +# the repository root is on ``sys.path``. ``cProfile`` (and therefore +# ``torch.compile``) imports the symbols below from ``profile``. Providing the +# small standard compatibility surface keeps both entry points working. +class _Utils: + """Utility adapter shared with :mod:`cProfile`.""" + + def __init__(self, profiler): + self.profiler = profiler + + def run(self, statement, filename, sort): + profiler = self.profiler() + try: + profiler.run(statement) + except SystemExit: + pass + finally: + self._show(profiler, filename, sort) + + def runctx(self, statement, globals, locals, filename, sort): + profiler = self.profiler() + try: + profiler.runctx(statement, globals, locals) + except SystemExit: + pass + finally: + self._show(profiler, filename, sort) + + @staticmethod + def _show(profiler, filename, sort): + if filename is not None: + profiler.dump_stats(filename) + else: + profiler.print_stats(sort) + + +def run(statement, filename=None, sort=-1): + """Run a statement under :mod:`cProfile`.""" + import cProfile + + return _Utils(cProfile.Profile).run(statement, filename, sort) + + +def runctx(statement, globals, locals, filename=None, sort=-1): + """Run a statement under :mod:`cProfile` with explicit namespaces.""" + import cProfile + + return _Utils(cProfile.Profile).runctx( + statement, globals, locals, filename, sort + ) + + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- diff --git a/pyproject.toml b/pyproject.toml index 3a38026f..410eac84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,14 @@ [project] -name = "autokernel" -version = "1.0.0" -description = "Autonomous AI agents optimizing GPU kernels overnight -- for any PyTorch model" +name = "motionkernel" +version = "0.1.0" +description = "Verified GPU kernel optimization for video generation models" readme = "README.md" requires-python = ">=3.10" dependencies = [ "matplotlib>=3.10.0", "numpy>=2.2.0", "pandas>=2.2.0", + "pyyaml>=6.0", "torch>=2.4.0", "triton>=3.3.0", ] @@ -36,6 +37,16 @@ hf-kernels = [ "kernels>=0.4.0", "huggingface-hub>=0.20.0", ] +# Development: CPU test suite +dev = [ + "pytest>=8.0.0", +] + +[tool.pytest.ini_options] +testpaths = ["tests"] +markers = [ + "gpu: test requires a CUDA GPU (deselect with -m 'not gpu')", +] [tool.uv.sources] torch = [ diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..8b5301b6 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,99 @@ +"""Shared helpers for the CPU test suite.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest + +from autokernel.specs import DT_BYTES, KernelSpec, Tolerance, size + +REPO_ROOT = Path(__file__).resolve().parents[1] +FIXTURES_DIR = Path(__file__).resolve().parent / "fixtures" + + +def _ref(x: Any = None, y: Any = None) -> Any: + return x + + +def _gen(size_map: Any, dtype: Any, device: str, seed: int = 42) -> dict: + return {"x": None} + + +def spec_kwargs(**overrides: Any) -> dict: + """Keyword arguments for a minimal valid :class:`KernelSpec`.""" + base: dict[str, Any] = { + "name": "unit_op", + "reference_fn": _ref, + "input_generator": _gen, + "sizes": { + "small": {"rows": 4, "cols": 4}, + "medium": {"rows": 8, "cols": 8}, + "large": {"rows": 16, "cols": 16}, + }, + "dtypes": ("float16", "float32"), + "tolerances": { + "float16": Tolerance(atol=1e-2, rtol=1e-2), + "float32": Tolerance(atol=1e-5, rtol=1e-5), + }, + "flops_fn": size("rows") * size("cols"), + "bytes_fn": 2 * size("rows") * size("cols") * DT_BYTES, + "shape_keys": ("rows", "cols"), + } + base.update(overrides) + return base + + +def make_spec(**overrides: Any) -> KernelSpec: + """Build a minimal valid specification, overriding any field.""" + return KernelSpec(**spec_kwargs(**overrides)) + + +@pytest.fixture +def repo_root() -> Path: + return REPO_ROOT + + +@pytest.fixture +def fixtures_dir() -> Path: + return FIXTURES_DIR + + +@pytest.fixture +def in_repo_root(monkeypatch: pytest.MonkeyPatch) -> Path: + """Run a test with the repository root as the working directory.""" + monkeypatch.chdir(REPO_ROOT) + return REPO_ROOT + + +@pytest.fixture +def torch_mod(): + """The torch module, skipping the test when torch is unavailable.""" + return pytest.importorskip("torch") + + +def cuda_available() -> bool: + """True when a CUDA device is usable (never raises when torch is absent).""" + try: + import torch + except Exception: + return False + try: + return bool(torch.cuda.is_available()) + except Exception: + return False + + +requires_gpu = pytest.mark.skipif( + not cuda_available(), reason="requires a CUDA GPU" +) + +__all__ = [ + "FIXTURES_DIR", + "REPO_ROOT", + "cuda_available", + "make_spec", + "requires_gpu", + "spec_kwargs", +] diff --git a/tests/fixtures/custom_add.py b/tests/fixtures/custom_add.py new file mode 100644 index 00000000..dabe5fd8 --- /dev/null +++ b/tests/fixtures/custom_add.py @@ -0,0 +1,81 @@ +"""External specification fixture used by the loader and CLI tests. + +Kept deliberately small and CPU-friendly: the tests here care about discovery, +validation and CLI plumbing, not about kernel performance. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Mapping + +from autokernel.specs import DT_BYTES, EdgeCase, KernelSpec, Tolerance, resolve_torch_dtype, size + +REPO_ROOT = Path(__file__).resolve().parents[2] + +#: Reuse the example starter kernel so the fixture declares a real file. +STARTER_KERNEL = REPO_ROOT / "examples" / "custom_ops" / "add_kernel.py" + + +def add_ref(x: Any, y: Any) -> Any: + return x + y + + +def gen_inputs(size_map: Mapping[str, int], dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + torch.manual_seed(seed) + torch_dtype = resolve_torch_dtype(dtype) + rows, cols = size_map["rows"], size_map["cols"] + return { + "x": torch.randn(rows, cols, device=device, dtype=torch_dtype), + "y": torch.randn(rows, cols, device=device, dtype=torch_dtype), + } + + +def _build(name: str = "fixture_add") -> KernelSpec: + return KernelSpec( + name=name, + reference_fn=add_ref, + input_generator=gen_inputs, + sizes={ + "small": {"rows": 8, "cols": 16}, + "medium": {"rows": 32, "cols": 32}, + "large": {"rows": 64, "cols": 64}, + }, + dtypes=("float32",), + tolerances={"float32": Tolerance(atol=1e-5, rtol=1e-5)}, + flops_fn=size("rows") * size("cols"), + bytes_fn=3 * size("rows") * size("cols") * DT_BYTES, + edge_cases=(EdgeCase(name="edge_7", size={"rows": 7, "cols": 7}),), + shape_keys=("rows", "cols"), + shape_aliases={"M": "rows", "N": "cols"}, + starter_kernels={"triton": STARTER_KERNEL}, + speedup_estimate="1.0x", + ) + + +#: A ready-made specification. +SPEC = _build() + + +def SPEC_FACTORY() -> KernelSpec: + """Zero-argument factory returning a specification.""" + return _build() + + +#: A specification whose name collides with a built-in operation. +COLLIDING_SPEC = _build(name="matmul") + +#: Not a specification at all. +NOT_A_SPEC = 42 + + +def BAD_FACTORY() -> int: + """A callable that returns the wrong type.""" + return 42 + + +def RAISING_FACTORY() -> KernelSpec: + """A callable that fails.""" + raise RuntimeError("factory exploded") diff --git a/tests/fixtures/wan_campaign.json b/tests/fixtures/wan_campaign.json new file mode 100644 index 00000000..a0b627d3 --- /dev/null +++ b/tests/fixtures/wan_campaign.json @@ -0,0 +1,88 @@ +{ + "schema_version": 1, + "producer": { + "name": "fastvideo", + "version": "0.2.0" + }, + "workload": { + "workload_id": "wan-t2v-1.3b", + "model_id": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "task": "inference", + "variant_id": "canonical" + }, + "environment": { + "hardware_profile_id": "gb200-1", + "software_profile_id": "torch2.11-cuda12.8-triton3.6" + }, + "total_profiled_device_time_us": 10000.0, + "targets": [ + { + "name": "wan.self_attn_residual_norm", + "operation": "wan_gated_residual_norm", + "kind": "fusion", + "spec_locator": "models/wan_gated_residual_norm.py:SPEC", + "total_device_time_us": 3200.0, + "self_device_time_us": 2400.0, + "calls": 40, + "requires_backward": false, + "observations": [ + { + "name": "b1-s20280-d1536", + "count": 40, + "total_device_time_us": 3200.0, + "inputs": [ + { + "name": "residual", + "shape": [1, 20280, 1536], + "stride": [31150080, 1536, 1], + "dtype": "bfloat16", + "device_type": "cuda", + "requires_grad": false + }, + { + "name": "gate", + "shape": [1, 1, 1536], + "stride": [9216, 1536, 1], + "dtype": "float32", + "device_type": "cuda", + "requires_grad": false + } + ], + "tags": ["wan2.1", "480p"] + } + ], + "attributes": { + "model_family": "wan" + } + }, + { + "name": "wan.rope", + "operation": "rotary_embedding", + "kind": "operator", + "spec_locator": null, + "total_device_time_us": 800.0, + "self_device_time_us": 700.0, + "calls": 80, + "requires_backward": false, + "observations": [ + { + "name": "wan-rope", + "count": 80, + "total_device_time_us": 800.0, + "inputs": [ + { + "name": "query", + "shape": [1, 20280, 12, 128], + "stride": [31150080, 1536, 128, 1], + "dtype": "bfloat16", + "device_type": "cuda", + "requires_grad": false + } + ], + "tags": ["wan2.1"] + } + ], + "attributes": {} + } + ] +} diff --git a/tests/test_backward.py b/tests/test_backward.py new file mode 100644 index 00000000..f0477f7d --- /dev/null +++ b/tests/test_backward.py @@ -0,0 +1,331 @@ +"""Optional backward (gradient) verification on CPU.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, Mapping + +import pytest + +import autokernel.verification.backward as backward_module +from autokernel.specs import ( + DT_BYTES, + BackwardSpec, + KernelSpec, + Tolerance, + load_spec, + resolve_torch_dtype, + size, +) +from autokernel.verification import check_backward + +torch = pytest.importorskip("torch") +bench = pytest.importorskip("bench") + + +def _gen(size_map: Mapping[str, int], dtype: Any, device: str, seed: int = 42) -> dict: + torch.manual_seed(seed) + torch_dtype = resolve_torch_dtype(dtype) + rows, cols = size_map["rows"], size_map["cols"] + return { + "x": torch.randn(rows, cols, device=device, dtype=torch_dtype), + "scale": torch.randn(cols, device=device, dtype=torch_dtype), + "bias": torch.randn(cols, device=device, dtype=torch_dtype), + } + + +def _affine_ref(x: Any, scale: Any, bias: Any) -> dict: + y = x * scale + bias + return {"output": y, "aux": (y - x, 3)} + + +def _spec(**overrides: Any) -> KernelSpec: + kwargs: dict[str, Any] = { + "name": "bwd_affine", + "reference_fn": _affine_ref, + "input_generator": _gen, + "sizes": { + "small": {"rows": 4, "cols": 8}, + "medium": {"rows": 8, "cols": 8}, + "large": {"rows": 16, "cols": 16}, + }, + "dtypes": ("float32",), + "tolerances": {"float32": Tolerance(atol=1e-5, rtol=1e-5)}, + "flops_fn": 3 * size("rows") * size("cols"), + "bytes_fn": 5 * size("rows") * size("cols") * DT_BYTES, + "shape_keys": ("rows", "cols"), + "backward_spec": BackwardSpec(differentiable_inputs=("x", "scale", "bias")), + } + kwargs.update(overrides) + return KernelSpec(**kwargs) + + +# --------------------------------------------------------------------------- +# Parity and mismatch +# --------------------------------------------------------------------------- + +def test_gradient_parity_for_reference_candidate(): + report = check_backward(_affine_ref, _spec(), device="cpu") + assert report.status == "PASS" + assert {r.input_name for r in report.gradients} == {"x", "scale", "bias"} + for record in report.gradients: + assert record.status == "match" + assert record.max_abs_error == 0.0 + # every floating tensor leaf received an upstream gradient + assert set(report.output_paths) == {'output["aux"][0]', 'output["output"]'} + + +def test_empty_sizes_returns_structured_failure(): + spec = _spec() + object.__setattr__(spec, "sizes", {}) + report = check_backward(_affine_ref, spec, device="cpu") + assert report.status == "FAIL" + assert "declares no sizes" in report.reason + + +def test_upstream_generator_fallback_generates_on_cpu_then_moves(monkeypatch): + original_generator = torch.Generator + moves = [] + randn_devices = [] + + def generator(*args, **kwargs): + if kwargs.get("device") == "mps": + raise RuntimeError("device-local generators unsupported") + return original_generator() + + class Generated: + def to(self, device): + moves.append(device) + return self + + def randn(*shape, **kwargs): + randn_devices.append(kwargs["device"]) + return Generated() + + monkeypatch.setattr(torch, "Generator", generator) + monkeypatch.setattr(torch, "randn", randn) + leaf = SimpleNamespace( + shape=(2, 3), + dtype=torch.float32, + device=torch.device("mps"), + ) + + upstreams = backward_module._upstream_gradients( + [("output", leaf)], "mps", seed=123 + ) + + assert len(upstreams) == 1 + assert randn_devices == ["cpu"] + assert moves == [torch.device("mps")] + + +def test_upstream_generation_failure_returns_structured_report(monkeypatch): + monkeypatch.setattr( + backward_module, + "_upstream_gradients", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("rng failed")), + ) + report = check_backward(_affine_ref, _spec(), device="cpu") + assert report.status == "FAIL" + assert "upstream gradient generation failed: RuntimeError: rng failed" in report.reason + + +def test_affine_example_fixture_passes_backward(repo_root): + spec = load_spec(str(repo_root / "examples" / "custom_ops" / "affine.py") + ":SPEC") + report = check_backward(spec.reference_fn, spec, device="cpu") + assert report.status == "PASS" + + +def test_affine_example_byte_accounting(repo_root): + spec = load_spec(str(repo_root / "examples" / "custom_ops" / "affine.py") + ":SPEC") + size_map = {"rows": 4, "cols": 8} + assert spec.bytes_fn(size_map, 2) == (3 * 4 + 2) * 8 * 2 + + +def test_perturbed_candidate_reports_mismatch_with_stats(): + def perturbed(x, scale, bias): + y = x * scale * 2 + bias # wrong gradient wrt x and scale + return {"output": y, "aux": (y - x, 3)} + + report = check_backward(perturbed, _spec(), device="cpu") + assert report.status == "FAIL" + by_name = {r.input_name: r for r in report.gradients} + assert by_name["x"].status == "mismatch" + assert by_name["scale"].status == "mismatch" + assert by_name["bias"].status == "match" # d(y)/d(bias) is unchanged + assert by_name["x"].max_abs_error > 0 + assert by_name["x"].mean_abs_error is not None + + +def test_nan_gradient_is_reported(): + def nan_candidate(x, scale, bias): + y = (x * scale + bias) * torch.where( + torch.arange(x.shape[1], device=x.device) == 0, + torch.tensor(float("nan"), device=x.device), + torch.tensor(1.0, device=x.device), + ) + return {"output": y, "aux": (y - x, 3)} + + report = check_backward(nan_candidate, _spec(), device="cpu") + assert report.status == "FAIL" + assert any(r.has_nan for r in report.gradients) + + +# --------------------------------------------------------------------------- +# Missing and unexpected gradients +# --------------------------------------------------------------------------- + +def test_missing_gradient_is_reported(): + def drops_bias(x, scale, bias): + y = x * scale # bias unused + return {"output": y, "aux": (y - x, 3)} + + report = check_backward(drops_bias, _spec(), device="cpu") + assert report.status == "FAIL" + by_name = {r.input_name: r for r in report.gradients} + assert by_name["bias"].status == "missing" + assert "no gradient" in by_name["bias"].reason + assert by_name["x"].status == "match" + + +def test_unexpected_gradient_is_reported(): + def ref_without_bias(x, scale, bias): + y = x * scale # reference ignores bias... + return {"output": y, "aux": (y - x, 3)} + + spec = _spec(reference_fn=ref_without_bias) + # ...but the candidate uses it + report = check_backward(_affine_ref, spec, device="cpu") + assert report.status == "FAIL" + by_name = {r.input_name: r for r in report.gradients} + assert by_name["bias"].status == "unexpected" + + +# --------------------------------------------------------------------------- +# Unsupported and invalid declarations +# --------------------------------------------------------------------------- + +def test_missing_backward_spec_fails_as_unsupported(): + spec = _spec(backward_spec=None) + report = check_backward(_affine_ref, spec, device="cpu") + assert report.status == "FAIL" + assert "unsupported" in report.reason + assert "forward-only" in report.reason + assert "backward_spec" in report.reason + + +def test_declared_input_missing_from_generator_fails(): + spec = _spec(backward_spec=BackwardSpec(differentiable_inputs=("x", "nope"))) + report = check_backward(_affine_ref, spec, device="cpu") + assert report.status == "FAIL" + assert "'nope'" in report.reason + assert "did not produce" in report.reason + + +def test_non_floating_declared_input_fails(): + def gen_with_int(size_map, dtype, device, seed=42): + inputs = _gen(size_map, dtype, device, seed) + inputs["bias"] = torch.ones(8, dtype=torch.int64) + return inputs + + spec = _spec(input_generator=gen_with_int) + report = check_backward(_affine_ref, spec, device="cpu") + assert report.status == "FAIL" + assert "not floating-point" in report.reason + + +def test_non_differentiable_candidate_fails_actionably(): + def forward_only(x, scale, bias): + with torch.no_grad(): + y = x * scale + bias + return {"output": y, "aux": (y - x, 3)} + + report = check_backward(forward_only, _spec(), device="cpu") + assert report.status == "FAIL" + assert "candidate backward failed" in report.reason + + +# --------------------------------------------------------------------------- +# Output path selection, determinism, accumulation +# --------------------------------------------------------------------------- + +def test_output_paths_restrict_upstream_leaves(): + spec = _spec( + backward_spec=BackwardSpec( + differentiable_inputs=("x", "scale", "bias"), + output_paths=('output["output"]',), + ) + ) + report = check_backward(_affine_ref, spec, device="cpu") + assert report.status == "PASS" + assert report.output_paths == ('output["output"]',) + + +def test_output_paths_must_exist(): + spec = _spec( + backward_spec=BackwardSpec( + differentiable_inputs=("x",), output_paths=("output[9]",) + ) + ) + report = check_backward(_affine_ref, spec, device="cpu") + assert report.status == "FAIL" + assert "does not exist" in report.reason + + +def test_upstream_gradients_are_deterministic(): + def perturbed(x, scale, bias): + y = x * scale * 1.01 + bias + return {"output": y, "aux": (y - x, 3)} + + first = check_backward(perturbed, _spec(), device="cpu") + second = check_backward(perturbed, _spec(), device="cpu") + assert first.status == second.status == "FAIL" + assert [r.as_dict() for r in first.gradients] == [r.as_dict() for r in second.gradients] + + +def test_repeated_checks_do_not_accumulate_state(): + spec = _spec() + for _ in range(3): + report = check_backward(_affine_ref, spec, device="cpu") + assert report.status == "PASS" + assert all(r.max_abs_error == 0.0 for r in report.gradients) + + +def test_backward_tolerances_override_forward_tolerances(): + def slightly_off(x, scale, bias): + y = (x * scale + bias) * 1.001 # 1e-3 relative error in grads + return {"output": y, "aux": (y - x, 3)} + + strict = check_backward(slightly_off, _spec(), device="cpu") + assert strict.status == "FAIL" + + loose = _spec( + backward_spec=BackwardSpec( + differentiable_inputs=("x", "scale", "bias"), + tolerances={"float32": Tolerance(atol=1e-2, rtol=1e-2)}, + ) + ) + relaxed = check_backward(slightly_off, loose, device="cpu") + assert relaxed.status == "PASS" + + +# --------------------------------------------------------------------------- +# bench.py wiring +# --------------------------------------------------------------------------- + +def test_bench_run_backward_check_prints_verdict(monkeypatch, capsys): + monkeypatch.setattr(bench, "BENCH_DEVICE", "cpu") + result = bench.run_backward_check(_affine_ref, _spec()) + captured = capsys.readouterr().out + assert result["status"] == "PASS" + assert "BACKWARD_CORRECTNESS: PASS" in captured + assert "grad[x]" in captured + + +def test_bench_run_backward_check_unsupported(monkeypatch, capsys): + monkeypatch.setattr(bench, "BENCH_DEVICE", "cpu") + result = bench.run_backward_check(_affine_ref, _spec(backward_spec=None)) + captured = capsys.readouterr().out + assert result["status"] == "FAIL" + assert "BACKWARD_CORRECTNESS: FAIL" in captured + assert "unsupported" in captured diff --git a/tests/test_bench_harness.py b/tests/test_bench_harness.py new file mode 100644 index 00000000..33595b2d --- /dev/null +++ b/tests/test_bench_harness.py @@ -0,0 +1,337 @@ +"""The refactored bench harness, exercised on CPU. + +No GPU is required: ``bench.BENCH_DEVICE`` is redirected to ``cpu`` so the five +correctness stages, the spec-driven size/dtype/tolerance plumbing and the +performance loop can be verified without CUDA. +""" + +from __future__ import annotations + +from typing import Any, Mapping + +import pytest + +from autokernel.specs import DT_BYTES, EdgeCase, KernelSpec, Tolerance, resolve_torch_dtype, size + +pytest.importorskip("torch") +bench = pytest.importorskip("bench") + + +def _ref(x: Any, y: Any) -> Any: + return x + y + + +def _gen(size_map: Mapping[str, int], dtype: Any, device: str, seed: int = 42) -> dict: + import torch + + torch.manual_seed(seed) + torch_dtype = resolve_torch_dtype(dtype) + rows, cols = size_map["rows"], size_map["cols"] + return { + "x": torch.randn(rows, cols, device=device, dtype=torch_dtype), + "y": torch.randn(rows, cols, device=device, dtype=torch_dtype), + } + + +def _spec(**overrides: Any) -> KernelSpec: + kwargs: dict[str, Any] = { + "name": "cpu_add", + "reference_fn": _ref, + "input_generator": _gen, + "sizes": { + "small": {"rows": 8, "cols": 16}, + "medium": {"rows": 16, "cols": 16}, + "large": {"rows": 32, "cols": 32}, + }, + "dtypes": ("float32",), + "tolerances": {"float32": Tolerance(atol=1e-5, rtol=1e-5)}, + "flops_fn": size("rows") * size("cols"), + "bytes_fn": 3 * size("rows") * size("cols") * DT_BYTES, + "edge_cases": ( + EdgeCase(name="edge_7", size={"rows": 7, "cols": 7}), + EdgeCase(name="edge_zeros", size={"rows": 5, "cols": 5}, + input_transform=lambda inputs: {k: v * 0 for k, v in inputs.items()}), + ), + "shape_keys": ("rows", "cols"), + } + kwargs.update(overrides) + return KernelSpec(**kwargs) + + +@pytest.fixture +def cpu_device(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(bench, "BENCH_DEVICE", "cpu") + return "cpu" + + +def _good_kernel(x, y): + return x + y + + +def _wrong_kernel(x, y): + return x - y + + +def test_all_five_stages_pass_for_a_correct_candidate(cpu_device, capsys): + results = bench.run_correctness(_good_kernel, _spec(), quick=False) + captured = capsys.readouterr().out + + assert results["correctness"] == "PASS" + assert results["smoke_test"] == "PASS" + assert results["shape_sweep"].startswith("PASS") + assert results["numerical_stability"] == "PASS" + assert results["determinism"] == "PASS" + assert results["edge_cases"] == "PASS" + # The greppable stage banners the agent loop reads must stay put. + for stage in ("Stage 1", "Stage 2", "Stage 3", "Stage 4", "Stage 5"): + assert stage in captured + + +def test_edge_cases_run_every_declared_case(cpu_device, capsys): + bench.run_correctness(_good_kernel, _spec(), quick=False) + captured = capsys.readouterr().out + assert "PASS: edge_7" in captured + assert "PASS: edge_zeros" in captured + + +def test_quick_mode_skips_stages_three_to_five(cpu_device): + results = bench.run_correctness(_good_kernel, _spec(), quick=True) + assert results["correctness"] == "PASS" + assert results["numerical_stability"] == "SKIP (quick mode)" + assert results["determinism"] == "SKIP (quick mode)" + assert results["edge_cases"] == "SKIP (quick mode)" + + +def test_incorrect_candidate_fails_correctness(cpu_device): + results = bench.run_correctness(_wrong_kernel, _spec(), quick=True) + assert results["correctness"] == "FAIL" + assert results["smoke_test"] == "FAIL" + + +def test_missing_edge_cases_report_skip(cpu_device): + results = bench.run_correctness(_good_kernel, _spec(edge_cases=()), quick=False) + assert results["edge_cases"] == "SKIP (no edge sizes defined)" + assert results["correctness"] == "PASS" + + +def test_edge_case_may_pin_its_own_dtype(cpu_device, capsys): + spec = _spec( + dtypes=("float32", "float16"), + tolerances={ + "float32": Tolerance(atol=1e-5, rtol=1e-5), + "float16": Tolerance(atol=1e-2, rtol=1e-2), + }, + edge_cases=(EdgeCase(name="edge_fp16", size={"rows": 6, "cols": 6}, dtype="float16"),), + ) + results = bench.run_correctness(_good_kernel, spec, quick=False) + assert results["edge_cases"] == "PASS" + assert "PASS: edge_fp16" in capsys.readouterr().out + + +@pytest.fixture +def stub_timer(monkeypatch: pytest.MonkeyPatch): + """Replace the GPU timer so the spec-driven plumbing can be checked on CPU. + + Only the timing primitive is stubbed: size selection, dtype resolution and + FLOP/byte accounting all run for real. + """ + calls: list[str] = [] + + def fake_do_bench(fn, warmup: int = 25, rep: int = 100) -> float: + fn() + calls.append("bench") + return 0.5 + + monkeypatch.setattr(bench, "_do_bench", fake_do_bench) + return calls + + +def test_performance_loop_uses_spec_accounting(cpu_device, stub_timer): + spec = _spec() + gpu = bench.GPUSpec(name="cpu-test", peak_tflops_fp16=100.0, peak_bandwidth_gb_s=1000.0) + perf = bench.run_performance(_good_kernel, spec, gpu, sizes_filter="large") + + assert perf["primary"] is not None + entry = perf["primary"] + assert entry["label"] == "large" + assert entry["flops"] == 32 * 32 + assert entry["bytes"] == 3 * 32 * 32 * 4 # float32 + assert entry["dtype"] == "torch.float32" + assert entry["kernel_latency_us"] == pytest.approx(500.0) + assert entry["speedup_vs_pytorch"] == pytest.approx(1.0) + # candidate and reference are both timed + assert len(stub_timer) == 2 + + +def test_performance_reports_every_requested_size(cpu_device, stub_timer): + spec = _spec() + gpu = bench.GPUSpec(name="cpu-test", peak_tflops_fp16=100.0, peak_bandwidth_gb_s=1000.0) + perf = bench.run_performance(_good_kernel, spec, gpu, sizes_filter="all") + labels = [entry["label"] for entry in perf["all"]] + assert labels == ["small", "medium", "large"] + assert perf["primary"]["label"] == "large" + + +# --------------------------------------------------------------------------- +# Structured outputs +# --------------------------------------------------------------------------- + +def _structured_spec(**overrides: Any) -> KernelSpec: + """Same shape contract as ``_spec`` but with a nested output tree.""" + kwargs: dict[str, Any] = { + "name": "cpu_structured", + "reference_fn": _structured_ref, + "input_generator": _gen, + "sizes": { + "small": {"rows": 8, "cols": 16}, + "medium": {"rows": 16, "cols": 16}, + "large": {"rows": 32, "cols": 32}, + }, + "dtypes": ("float32",), + "tolerances": {"float32": Tolerance(atol=1e-5, rtol=1e-5)}, + "flops_fn": size("rows") * size("cols"), + "bytes_fn": 4 * size("rows") * size("cols") * DT_BYTES, + "shape_keys": ("rows", "cols"), + } + kwargs.update(overrides) + return KernelSpec(**kwargs) + + +def _structured_ref(x: Any, y: Any) -> Any: + return {"output": x + y, "aux": (x - y, 2)} + + +def _good_structured_kernel(x: Any, y: Any) -> Any: + return {"output": x + y, "aux": (x - y, 2)} + + +def _wrong_aux_kernel(x: Any, y: Any) -> Any: + return {"output": x + y, "aux": (x - y + 1.0, 2)} + + +def _wrong_metadata_kernel(x: Any, y: Any) -> Any: + return {"output": x + y, "aux": (x - y, 999)} + + +def _dropping_aux_kernel(x: Any, y: Any) -> Any: + return {"output": x + y} + + +def test_structured_candidate_passes_all_stages(cpu_device, capsys): + results = bench.run_correctness(_good_structured_kernel, _structured_spec(), quick=False) + captured = capsys.readouterr().out + + assert results["correctness"] == "PASS" + assert results["smoke_test"] == "PASS" + assert results["determinism"] == "PASS" + # Every tensor leaf is compared, and the paths are stable. + paths = {record["path"] for record in results["leaf_details"]} + assert {'output["output"]', 'output["aux"][0]', 'output["aux"][1]'} <= paths + assert "Stage 1" in captured and "Stage 5" in captured + + +def test_wrong_aux_leaf_fails_with_diagnostic_path(cpu_device): + results = bench.run_correctness(_wrong_aux_kernel, _structured_spec(), quick=True) + assert results["correctness"] == "FAIL" + assert any('output["aux"][0]' in record["path"] and not record["match"] + for record in results["leaf_details"]) + assert any('output["aux"][0]' in detail for detail in results["details"]) + + +def test_wrong_metadata_leaf_fails(cpu_device): + results = bench.run_correctness(_wrong_metadata_kernel, _structured_spec(), quick=True) + assert results["correctness"] == "FAIL" + assert any("metadata mismatch" in record["reason"] for record in results["leaf_details"]) + + +def test_dropped_output_branch_fails_structure_check(cpu_device): + results = bench.run_correctness(_dropping_aux_kernel, _structured_spec(), quick=True) + assert results["correctness"] == "FAIL" + assert any("missing output path" in detail for detail in results["details"]) + + + +# --------------------------------------------------------------------------- +# Shape corpora +# --------------------------------------------------------------------------- + +def _corpus_cases() -> tuple: + from autokernel.verification import CorpusCase + + return ( + CorpusCase(name="prod-a", size={"rows": 8, "cols": 16}, weight=3), + CorpusCase(name="prod-b", size={"rows": 7, "cols": 13}, weight=1), + ) + + +def test_corpus_cases_join_the_shape_sweep(cpu_device, capsys): + results = bench.run_correctness( + _good_kernel, _spec(), quick=True, corpus_cases=_corpus_cases() + ) + captured = capsys.readouterr().out + assert results["correctness"] == "PASS" + assert "PASS: prod-a" in captured + assert "PASS: prod-b" in captured + assert "PASS: small" in captured # built-in sweep still runs + + +def test_corpus_only_replaces_the_builtin_sweep(cpu_device, capsys): + results = bench.run_correctness( + _good_kernel, _spec(), quick=True, + corpus_cases=_corpus_cases(), corpus_only=True, + ) + captured = capsys.readouterr().out + assert results["correctness"] == "PASS" + assert "PASS: prod-a" in captured + assert "PASS: small" not in captured + + +def test_wrong_kernel_fails_a_corpus_case(cpu_device, capsys): + """A candidate that is correct on built-in sizes but wrong on a corpus + shape must fail with the corpus case named.""" + + def wrong_on_odd_rows(x, y): + out = x + y + if x.shape[0] == 7: # prod-b is 7x13 + out = out + 1.0 + return out + + results = bench.run_correctness( + wrong_on_odd_rows, _spec(), quick=True, corpus_cases=_corpus_cases() + ) + captured = capsys.readouterr().out + assert results["correctness"] == "FAIL" + assert "FAIL: prod-b" in captured + assert "PASS: prod-a" in captured + + +def test_performance_benches_corpus_cases_with_weights(cpu_device, stub_timer): + spec = _spec() + gpu = bench.GPUSpec(name="cpu-test", peak_tflops_fp16=100.0, peak_bandwidth_gb_s=1000.0) + perf = bench.run_performance( + _good_kernel, spec, gpu, sizes_filter="all", corpus_cases=_corpus_cases() + ) + + labels = [entry["label"] for entry in perf["all"]] + assert labels == ["small", "medium", "large", "prod-a", "prod-b"] + corpus = perf["corpus"] + assert corpus is not None + assert [entry["weight"] for entry in corpus["cases"]] == [3, 1] + weighted = corpus["weighted"]["torch.float32"] + assert weighted["cases"] == 2 + assert weighted["weight"] == 4 + # the stub timer returns 0.5 ms for every call; weight must not repeat loops + assert weighted["kernel_ms"] == pytest.approx(0.5) + assert len(stub_timer) == 2 * 5 # candidate + reference per configuration + + +def test_performance_corpus_only_skips_builtin_sizes(cpu_device, stub_timer): + spec = _spec() + gpu = bench.GPUSpec(name="cpu-test", peak_tflops_fp16=100.0, peak_bandwidth_gb_s=1000.0) + perf = bench.run_performance( + _good_kernel, spec, gpu, sizes_filter="all", + corpus_cases=_corpus_cases(), corpus_only=True, + ) + labels = [entry["label"] for entry in perf["all"]] + assert labels == ["prod-a", "prod-b"] + diff --git a/tests/test_builtin_specs.py b/tests/test_builtin_specs.py new file mode 100644 index 00000000..4217b6cb --- /dev/null +++ b/tests/test_builtin_specs.py @@ -0,0 +1,336 @@ +"""Compatibility freeze for the nine built-in operations. + +Every value below was captured from ``bench.py::KERNEL_CONFIGS`` and +``extract.py``'s metadata maps *before* the registry refactor. A failure here +means benchmark coverage, tolerances or accounting changed -- which must be a +deliberate, reviewed decision, not a side effect. +""" + +from __future__ import annotations + +import subprocess +import sys + +import pytest +from conftest import REPO_ROOT + +from autokernel.specs import create_builtin_registry, dtype_bytes + +BUILTIN_NAMES = ( + "matmul", + "softmax", + "layernorm", + "flash_attention", + "fused_mlp", + "cross_entropy", + "rotary_embedding", + "rmsnorm", + "reduce", +) + +SIZE_LABELS = { + "matmul": ("tiny", "small", "medium", "large", "xlarge", "tall", "wide", "deep_k", + "llm_qkv", "llm_mlp"), + "softmax": ("tiny", "small", "medium", "large", "xlarge", "wide", "narrow", "vocab"), + "layernorm": ("tiny", "small", "medium", "large", "xlarge", "wide", "llm_7b", "llm_13b"), + "flash_attention": ("tiny", "small", "medium", "large", "xlarge", "long", "gqa", "llm_7b"), + "fused_mlp": ("tiny", "small", "medium", "large", "xlarge", "llm_7b", "llm_13b"), + "cross_entropy": ("tiny", "small", "medium", "large", "xlarge", "llama", "gpt2"), + "rotary_embedding": ("tiny", "small", "medium", "large", "xlarge", "llm_7b", "llm_13b"), + "rmsnorm": ("small", "medium", "large", "llama"), + "reduce": ("small", "medium", "large", "wide"), +} + +DTYPES = { + "matmul": ("float16", "bfloat16", "float32"), + "softmax": ("float16", "bfloat16", "float32"), + "layernorm": ("float16", "bfloat16", "float32"), + "flash_attention": ("float16", "bfloat16"), + "fused_mlp": ("float16", "bfloat16", "float32"), + "cross_entropy": ("float16", "bfloat16", "float32"), + "rotary_embedding": ("float16", "bfloat16", "float32"), + "rmsnorm": ("float16", "bfloat16"), + "reduce": ("float16", "bfloat16"), +} + +TOLERANCES = { + "matmul": {"float16": (1e-2, 1e-2), "bfloat16": (2e-2, 2e-2), "float32": (1e-4, 1e-4)}, + "softmax": {"float16": (1e-3, 1e-3), "bfloat16": (2e-3, 2e-3), "float32": (1e-5, 1e-5)}, + "layernorm": {"float16": (1e-3, 1e-3), "bfloat16": (2e-3, 2e-3), "float32": (1e-5, 1e-5)}, + "flash_attention": {"float16": (1e-2, 1e-2), "bfloat16": (2e-2, 2e-2), "float32": (1e-4, 1e-4)}, + "fused_mlp": {"float16": (1e-2, 1e-2), "bfloat16": (2e-2, 2e-2), "float32": (1e-4, 1e-4)}, + "cross_entropy": {"float16": (1e-2, 1e-2), "bfloat16": (2e-2, 2e-2), "float32": (1e-5, 1e-5)}, + "rotary_embedding": {"float16": (1e-3, 1e-3), "bfloat16": (2e-3, 2e-3), "float32": (1e-5, 1e-5)}, + "rmsnorm": {"float16": (1e-2, 1e-2), "bfloat16": (1e-1, 5e-2)}, + "reduce": {"float16": (1e-2, 1e-2), "bfloat16": (1e-1, 5e-2)}, +} + +LARGE_SIZES = { + "matmul": {"M": 2048, "N": 2048, "K": 2048}, + "softmax": {"rows": 4096, "cols": 4096}, + "layernorm": {"batch": 4096, "dim": 2048}, + "flash_attention": {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 64}, + "fused_mlp": {"batch": 2048, "dim": 2048, "hidden": 5504}, + "cross_entropy": {"batch": 4096, "vocab": 32000}, + "rotary_embedding": {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 128}, + "rmsnorm": {"M": 4096, "N": 4096}, + "reduce": {"M": 8192, "N": 8192}, +} + +# FLOPs and fp16 bytes at the 'large' size, captured pre-refactor. +LARGE_FLOPS = { + "matmul": 17179869184, + "softmax": 83886080, + "layernorm": 67108864, + "flash_attention": 17179869184, + "fused_mlp": 138512695296, + "cross_entropy": 524288000, + "rotary_embedding": 50331648, + "rmsnorm": 100663296, + "reduce": 67108864, +} +LARGE_BYTES_FP16 = { + "matmul": 25165824, + "softmax": 67108864, + "layernorm": 33562624, + "flash_attention": 33554432, + "fused_mlp": 84410368, + "cross_entropy": 262152192, + "rotary_embedding": 33816576, + "rmsnorm": 67117056, + "reduce": 134234112, +} + +EDGE_CASES = { + "matmul": ( + ("edge_1023", {"M": 1023, "N": 1023, "K": 1023}), + ("edge_4097", {"M": 4097, "N": 4097, "K": 512}), + ("edge_1537", {"M": 1537, "N": 1537, "K": 1537}), + ), + "softmax": ( + ("edge_1023", {"rows": 1023, "cols": 1023}), + ("edge_4097", {"rows": 4097, "cols": 4097}), + ("edge_50257", {"rows": 1024, "cols": 50257}), + ), + "layernorm": ( + ("edge_1023", {"batch": 1023, "dim": 1023}), + ("edge_4097", {"batch": 4097, "dim": 4097}), + ), + "flash_attention": ( + ("edge_127", {"batch": 1, "heads": 8, "seq_len": 127, "head_dim": 64}), + ("edge_1023", {"batch": 1, "heads": 8, "seq_len": 1023, "head_dim": 64}), + ), + "fused_mlp": ( + ("edge_1023", {"batch": 1023, "dim": 1024, "hidden": 2048}), + ("edge_4097", {"batch": 4097, "dim": 512, "hidden": 1024}), + ), + "cross_entropy": ( + ("edge_1023", {"batch": 1023, "vocab": 32000}), + ("edge_50257", {"batch": 4096, "vocab": 50257}), + ), + "rotary_embedding": ( + ("edge_127", {"batch": 1, "heads": 8, "seq_len": 127, "head_dim": 64}), + ("edge_1023", {"batch": 1, "heads": 8, "seq_len": 1023, "head_dim": 128}), + ), + "rmsnorm": ( + ("edge_1023", {"M": 1023, "N": 768}), + ("edge_4097", {"M": 4097, "N": 1024}), + ), + "reduce": ( + ("edge_1023", {"M": 1023, "N": 1024}), + ("edge_4097", {"M": 4096, "N": 4097}), + ), +} + +# Alias maps carried over from extract.py::SHAPE_ALIAS_MAP. The identity entries +# for matmul/rmsnorm/reduce make the mapping explicit; behavior is unchanged +# because unmapped keys always passed through. +SHAPE_ALIASES = { + "matmul": {"M": "M", "N": "N", "K": "K"}, + "softmax": {"M": "rows", "N": "cols", "rows": "rows", "cols": "cols"}, + "layernorm": { + "M": "batch", "N": "dim", "rows": "batch", "cols": "dim", + "batch": "batch", "dim": "dim", + }, + "flash_attention": { + "B": "batch", "H": "heads", "N": "seq_len", "S": "seq_len", "D": "head_dim", + "batch": "batch", "heads": "heads", "seq_len": "seq_len", "head_dim": "head_dim", + }, + "fused_mlp": { + "M": "batch", "N": "hidden", "K": "dim", + "batch": "batch", "dim": "dim", "hidden": "hidden", + }, + "cross_entropy": {"batch": "batch", "vocab": "vocab"}, + "rotary_embedding": { + "B": "batch", "H": "heads", "N": "seq_len", "S": "seq_len", "D": "head_dim", + "batch": "batch", "heads": "heads", "seq_len": "seq_len", "head_dim": "head_dim", + }, + "rmsnorm": {"M": "M", "N": "N"}, + "reduce": {"M": "M", "N": "N"}, +} + +SPEEDUP_ESTIMATES = { + "matmul": "2-3x", + "flash_attention": "2-4x", + "layernorm": "1.5-3x", + "softmax": "1.5-3x", + "cross_entropy": "1.5-2x", + "fused_mlp": "2-3x", + "rmsnorm": "1.5-3x", + "reduce": "1.5-2x", + "rotary_embedding": "1.5-2x", +} + +# Extraction fallback shapes, captured from extract.py::get_default_shape. +EXTRACTION_SHAPES = { + "matmul": {"M": 2048, "N": 2048, "K": 2048}, + "flash_attention": {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 64}, + "layernorm": {"batch": 4096, "dim": 2048}, + "softmax": {"rows": 4096, "cols": 4096}, + "cross_entropy": {"batch": 4096, "vocab": 32000}, + "fused_mlp": {"batch": 2048, "dim": 2048, "hidden": 5504}, + "rmsnorm": {"M": 4096, "N": 4096}, + "reduce": {"M": 4096, "N": 4096}, + "rotary_embedding": {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 128}, +} + +INPUT_KEYS = { + "matmul": ("A", "B"), + "softmax": ("x",), + "layernorm": ("x", "weight", "bias"), + "flash_attention": ("Q", "K", "V"), + "fused_mlp": ("x", "w_gate", "w_up", "w_down"), + "cross_entropy": ("logits", "targets"), + "rotary_embedding": ("x", "cos", "sin"), + "rmsnorm": ("x", "weight"), + "reduce": ("x",), +} + + +@pytest.fixture(scope="module") +def registry(): + return create_builtin_registry() + + +def test_builtin_names_and_order(registry): + assert registry.list_names() == BUILTIN_NAMES + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_size_labels(registry, name): + assert tuple(registry.get(name).sizes) == SIZE_LABELS[name] + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_large_size_values(registry, name): + assert registry.get(name).sizes["large"] == LARGE_SIZES[name] + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_dtypes(registry, name): + assert registry.get(name).dtypes == DTYPES[name] + assert registry.get(name).primary_dtype == DTYPES[name][0] + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_tolerances(registry, name): + spec = registry.get(name) + actual = {d: (t.atol, t.rtol) for d, t in spec.tolerances.items()} + assert actual == TOLERANCES[name] + for dtype in spec.dtypes: + assert dtype in spec.tolerances + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_accounting_matches_pre_refactor_values(registry, name): + spec = registry.get(name) + large = spec.sizes["large"] + assert spec.flops_fn(large) == LARGE_FLOPS[name] + assert spec.bytes_fn(large, dtype_bytes("float16")) == LARGE_BYTES_FP16[name] + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_edge_cases(registry, name): + spec = registry.get(name) + actual = tuple((e.name, dict(e.size)) for e in spec.edge_cases) + assert actual == EDGE_CASES[name] + for edge in spec.edge_cases: + assert edge.seed == 42 + assert edge.dtype is None + assert edge.input_transform is None + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_shape_metadata_used_by_extraction(registry, name): + spec = registry.get(name) + assert spec.shape_aliases == SHAPE_ALIASES[name] + assert spec.shape_keys == tuple(LARGE_SIZES[name]) + assert spec.speedup_estimate == SPEEDUP_ESTIMATES[name] + assert spec.extraction_shape() == EXTRACTION_SHAPES[name] + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_starter_kernels_exist(registry, name): + spec = registry.get(name) + assert set(spec.starter_kernels) == {"triton", "cuda"} + for backend in ("triton", "cuda"): + assert spec.starter_kernel(backend).is_file() + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_input_generators_are_deterministic(registry, name, torch_mod): + spec = registry.get(name) + small = spec.sizes["small"] + dtype = "float32" if "float32" in spec.dtypes else spec.primary_dtype + first = spec.input_generator(small, dtype, "cpu", 42) + second = spec.input_generator(small, dtype, "cpu", 42) + third = spec.input_generator(small, dtype, "cpu", 1234) + + assert tuple(first) == INPUT_KEYS[name] + for key in first: + assert torch_mod.equal(first[key], second[key]), key + + changed = any(not torch_mod.equal(first[k], third[k]) for k in first) + assert changed, "a different seed must produce different inputs" + + +@pytest.mark.parametrize("name", BUILTIN_NAMES) +def test_reference_functions_run_on_cpu(registry, name, torch_mod): + spec = registry.get(name) + dtype = "float32" if "float32" in spec.dtypes else spec.primary_dtype + inputs = spec.input_generator(spec.sizes["small"], dtype, "cpu", 42) + output = spec.reference_fn(**inputs) + assert isinstance(output, torch_mod.Tensor) + + +def test_reference_functions_come_from_reference_module(registry): + import reference + + for name in BUILTIN_NAMES: + lazy = registry.get(name).reference_fn + assert lazy.module_name == "reference" + assert getattr(reference, lazy.attribute) is lazy.resolve() + + +def test_discovery_does_not_import_torch(): + """Registry discovery must work on a CPU-only machine without torch.""" + code = ( + "import sys\n" + "from autokernel.specs import create_builtin_registry\n" + "registry = create_builtin_registry()\n" + "assert len(registry) == 9, registry.list_names()\n" + "for spec in registry:\n" + " spec.flops_fn(spec.sizes['large'])\n" + " spec.bytes_fn(spec.sizes['large'], 2)\n" + "assert 'torch' not in sys.modules, 'discovery imported torch'\n" + "print('OK')\n" + ) + result = subprocess.run( + [sys.executable, "-c", code], + cwd=REPO_ROOT, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + assert "OK" in result.stdout diff --git a/tests/test_campaign.py b/tests/test_campaign.py new file mode 100644 index 00000000..361060f1 --- /dev/null +++ b/tests/test_campaign.py @@ -0,0 +1,319 @@ +"""Optimization campaign schema, ranking, and orchestrator bridge.""" + +from __future__ import annotations + +import json +import subprocess +import sys + +import pytest + +from autokernel.campaign import ( + CampaignError, + OptimizationCampaign, + load_campaign, + parse_agent_command, + prepare_campaign, + rank_targets, + run_campaign, + write_optimization_plan, +) + + +def test_load_and_rank_wan_campaign(fixtures_dir): + campaign = load_campaign(fixtures_dir / "wan_campaign.json") + assert campaign.workload["workload_id"] == "wan-t2v-1.3b" + assert [target.name for target in rank_targets(campaign)] == [ + "wan.self_attn_residual_norm", + "wan.rope", + ] + assert rank_targets(campaign)[0].impact_pct( + campaign.total_profiled_device_time_us + ) == pytest.approx(32.0) + + +def test_campaign_rejects_tensor_values(fixtures_dir): + payload = json.loads( + (fixtures_dir / "wan_campaign.json").read_text(encoding="utf-8") + ) + payload["targets"][0]["observations"][0]["inputs"][0]["values"] = [1.0] + with pytest.raises(CampaignError, match="unknown field.*values"): + OptimizationCampaign.from_dict(payload) + + +def test_campaign_rejects_nonfinite_timing(fixtures_dir): + payload = json.loads( + (fixtures_dir / "wan_campaign.json").read_text(encoding="utf-8") + ) + payload["targets"][0]["total_device_time_us"] = float("inf") + with pytest.raises(CampaignError, match="finite non-negative"): + OptimizationCampaign.from_dict(payload) + + +def test_campaign_observation_counts_must_match_calls(fixtures_dir): + payload = json.loads( + (fixtures_dir / "wan_campaign.json").read_text(encoding="utf-8") + ) + payload["targets"][0]["calls"] = 41 + with pytest.raises(CampaignError, match="counts sum to 40"): + OptimizationCampaign.from_dict(payload) + + +def test_campaign_allows_overlapping_target_timings(fixtures_dir): + payload = json.loads( + (fixtures_dir / "wan_campaign.json").read_text(encoding="utf-8") + ) + payload["targets"][1]["total_device_time_us"] = 8_000 + campaign = OptimizationCampaign.from_dict(payload) + assert len(campaign.targets) == 2 + + +@pytest.mark.parametrize("operation", ["../../escape", "bad-name", "bad.name"]) +def test_campaign_rejects_unsafe_operation_identifiers(fixtures_dir, operation): + payload = json.loads( + (fixtures_dir / "wan_campaign.json").read_text(encoding="utf-8") + ) + payload["targets"][0]["operation"] = operation + with pytest.raises(CampaignError, match="safe identifier"): + OptimizationCampaign.from_dict(payload) + + +def test_campaign_rejects_content_hidden_in_attributes(fixtures_dir): + payload = json.loads( + (fixtures_dir / "wan_campaign.json").read_text(encoding="utf-8") + ) + payload["targets"][0]["attributes"]["prompt"] = "do not export me" + with pytest.raises(CampaignError, match="content or secret fields"): + OptimizationCampaign.from_dict(payload) + + +def test_write_plan_bridges_to_existing_orchestrator(fixtures_dir, tmp_path): + campaign = load_campaign(fixtures_dir / "wan_campaign.json") + output = tmp_path / "optimization_plan.json" + plan = write_optimization_plan(campaign, output) + assert output.is_file() + assert plan["kernels_to_optimize"][0] == { + "rank": 1, + "file": "workspace/kernel_wan_gated_residual_norm_1.py", + "op_type": "wan_gated_residual_norm", + "target_name": "wan.self_attn_residual_norm", + "spec_locator": "models/wan_gated_residual_norm.py:SPEC", + "pct_total": 32.0, + "calls": 40, + "requires_backward": False, + } + + +def test_campaign_cli_validate_rank_and_plan(repo_root, fixtures_dir, tmp_path): + campaign = fixtures_dir / "wan_campaign.json" + validate = subprocess.run( + [sys.executable, "campaign.py", "validate", str(campaign)], + cwd=repo_root, + check=False, + capture_output=True, + text=True, + ) + assert validate.returncode == 0 + assert "CAMPAIGN_VALIDATION: PASS" in validate.stdout + + rank = subprocess.run( + [sys.executable, "campaign.py", "rank", str(campaign)], + cwd=repo_root, + check=False, + capture_output=True, + text=True, + ) + assert rank.returncode == 0 + assert json.loads(rank.stdout)[0]["impact_pct"] == 32.0 + + output = tmp_path / "plan.json" + plan = subprocess.run( + [ + sys.executable, + "campaign.py", + "plan", + str(campaign), + "--output", + str(output), + ], + cwd=repo_root, + check=False, + capture_output=True, + text=True, + ) + assert plan.returncode == 0 + assert "CAMPAIGN_PLAN: PASS" in plan.stdout + assert output.is_file() + + +def test_prepare_requires_explicit_trust(fixtures_dir, tmp_path): + campaign = load_campaign(fixtures_dir / "wan_campaign.json") + with pytest.raises(CampaignError, match="trust_specs=True"): + prepare_campaign(campaign, tmp_path) + + +def test_prepare_materializes_starter_and_receipt( + fixtures_dir, tmp_path, monkeypatch +): + payload = json.loads( + (fixtures_dir / "wan_campaign.json").read_text(encoding="utf-8") + ) + payload["targets"] = payload["targets"][:1] + payload["targets"][0][ + "spec_locator" + ] = "models/wan_gated_residual_norm.py:SPEC" + campaign = OptimizationCampaign.from_dict(payload, source="test-campaign") + monkeypatch.chdir(fixtures_dir.parent.parent) + output = tmp_path / "workspace" + receipt = prepare_campaign(campaign, output, trust_specs=True) + candidate = output / "kernel_wan_gated_residual_norm_1.py" + assert candidate.is_file() + assert "KERNEL_TYPE = \"wan_gated_residual_norm\"" in candidate.read_text( + encoding="utf-8" + ) + assert receipt["status"] == "prepared" + assert (output / "optimization_plan.json").is_file() + assert (output / "campaign_receipt.json").is_file() + + +def _single_wan_target(fixtures_dir) -> OptimizationCampaign: + payload = json.loads( + (fixtures_dir / "wan_campaign.json").read_text(encoding="utf-8") + ) + payload["targets"] = payload["targets"][:1] + return OptimizationCampaign.from_dict(payload, source="test-campaign") + + +def test_overnight_dry_run_is_cwd_independent( + repo_root, fixtures_dir, tmp_path, monkeypatch +): + campaign = _single_wan_target(fixtures_dir) + monkeypatch.chdir(tmp_path) + receipt = run_campaign( + campaign, + repo_root=repo_root, + budget_hours=0.25, + dry_run=True, + trust_specs=True, + ) + assert receipt["status"] == "prepared" + prompt = (repo_root / "workspace" / "overnight_prompt.md").read_text( + encoding="utf-8" + ) + assert "--spec models/wan_gated_residual_norm.py:SPEC" in prompt + assert ( + "--shape-corpus models/wan_gated_residual_norm_corpus.json" + in prompt + ) + + +def test_parse_agent_command_substitutes_embedded_placeholders(tmp_path): + command = parse_agent_command( + "agent --repo={repo} --prompt={prompt_file}", + repo_root=tmp_path, + prompt_path=tmp_path / "prompt.md", + ) + assert command == [ + "agent", + f"--repo={tmp_path}", + f"--prompt={tmp_path / 'prompt.md'}", + ] + + +def test_overnight_runner_writes_terminal_morning_report( + repo_root, fixtures_dir +): + campaign = _single_wan_target(fixtures_dir) + command = [ + sys.executable, + "-c", + ( + "import subprocess,sys;" + "subprocess.run([sys.executable,'orchestrate.py','next']," + "check=True)" + ), + ] + receipt = run_campaign( + campaign, + repo_root=repo_root, + budget_hours=0.25, + trust_specs=True, + agent_command=command, + timeout_seconds=10, + ) + assert receipt["status"] == "incomplete" + assert receipt["progress"]["completed_targets"] == 0 + report = repo_root / "workspace" / "morning_report.md" + assert report.is_file() + assert "wan_gated_residual_norm" in report.read_text(encoding="utf-8") + + +def test_overnight_runner_requires_explicit_spec_trust( + repo_root, fixtures_dir +): + campaign = _single_wan_target(fixtures_dir) + with pytest.raises(CampaignError, match="trust"): + run_campaign( + campaign, + repo_root=repo_root, + budget_hours=0.25, + dry_run=True, + ) + + +def test_overnight_runner_records_agent_launch_failure( + repo_root, fixtures_dir, tmp_path +): + campaign = _single_wan_target(fixtures_dir) + missing = tmp_path / "no_such_agent_binary" + receipt = run_campaign( + campaign, + repo_root=repo_root, + budget_hours=0.25, + trust_specs=True, + agent_command=[str(missing)], + timeout_seconds=10, + ) + assert receipt["status"] == "agent_launch_failed" + assert receipt["error"] + assert receipt["finished_at"] + persisted = json.loads( + (repo_root / "workspace" / "overnight_receipt.json").read_text( + encoding="utf-8" + ) + ) + assert persisted["status"] == "agent_launch_failed" + assert (repo_root / "workspace" / "morning_report.md").is_file() + + +def test_run_cli_maps_failure_statuses_to_nonzero_exit( + repo_root, fixtures_dir, tmp_path, capsys +): + import campaign as campaign_cli + + source = fixtures_dir / "wan_campaign.json" + payload = json.loads(source.read_text(encoding="utf-8")) + payload["targets"] = payload["targets"][:1] + campaign_file = tmp_path / "campaign.json" + campaign_file.write_text(json.dumps(payload), encoding="utf-8") + + exit_code = campaign_cli.main( + [ + "run", + str(campaign_file), + "--trust-specs", + "--agent-command", + str(tmp_path / "no_such_agent_binary"), + ] + ) + captured = capsys.readouterr() + assert exit_code == 1 + assert "CAMPAIGN_RUN: FAIL" in captured.out + assert "status: agent_launch_failed" in captured.out + + exit_code = campaign_cli.main( + ["run", str(campaign_file), "--trust-specs", "--dry-run"] + ) + captured = capsys.readouterr() + assert exit_code == 0 + assert "CAMPAIGN_RUN: PASS" in captured.out diff --git a/tests/test_cli_compat.py b/tests/test_cli_compat.py new file mode 100644 index 00000000..bafb7ee8 --- /dev/null +++ b/tests/test_cli_compat.py @@ -0,0 +1,340 @@ +"""CLI compatibility and extraction wiring. + +These tests protect two things the autonomous loop depends on: + +* the pre-existing command-line surface (``--kernel``, ``--sizes``, ``--quick``, + ``--profile``, ``--report``, ``--top``, ``--kernel-type``, ``--backend``); +* extraction consuming a ``KernelSpec`` instead of operation-specific maps. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest +from conftest import FIXTURES_DIR, REPO_ROOT + +from autokernel.specs import create_builtin_registry, load_spec + +FIXTURE_LOCATOR = f"{FIXTURES_DIR / 'custom_add.py'}:SPEC" + + +def run_script(script: str, *args: str, cwd: Path | None = None) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, str(REPO_ROOT / script), *args], + cwd=str(cwd or REPO_ROOT), + capture_output=True, + text=True, + timeout=300, + ) + + +# --------------------------------------------------------------------------- +# bench.py command line +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize( + "flag", ["--kernel", "--sizes", "--quick", "--profile", "--spec", "--spec-override"] +) +def test_bench_help_lists_expected_flags(flag): + result = run_script("bench.py", "--help") + assert result.returncode == 0, result.stderr + assert flag in result.stdout + + +def test_bench_help_does_not_import_the_external_spec(tmp_path: Path): + """`--help` must never execute external specification code.""" + exploding = tmp_path / "exploding_spec.py" + exploding.write_text( + "raise SystemExit('the external spec was imported during --help')\n" + ) + result = run_script("bench.py", "--spec", f"{exploding}:SPEC", "--help") + assert result.returncode == 0, result.stdout + result.stderr + assert "usage: bench.py" in result.stdout + assert "imported during --help" not in result.stdout + result.stderr + + +def test_bench_reports_actionable_error_for_bad_spec(): + result = run_script("bench.py", "--spec", "/nope/missing_spec.py:SPEC") + combined = result.stdout + result.stderr + assert result.returncode == 1 + assert "/nope/missing_spec.py:SPEC" in combined + assert "file not found" in combined + # The greppable failure contract the agent loop parses must survive. + assert "correctness: FAIL" in combined + assert "throughput_tflops: 0.000" in combined + + +def test_bench_loads_external_spec_before_touching_the_candidate(): + """The spec is resolved (and echoed) before kernel.py is imported.""" + result = run_script("bench.py", "--spec", FIXTURE_LOCATOR, "--quick") + combined = result.stdout + result.stderr + assert f"kernel_spec: {FIXTURE_LOCATOR}" in combined + assert "cannot load spec" not in combined + + +def test_bench_spec_collision_is_rejected_without_override(): + colliding = f"{FIXTURES_DIR / 'custom_add.py'}:COLLIDING_SPEC" + result = run_script("bench.py", "--spec", colliding) + combined = result.stdout + result.stderr + assert result.returncode == 1 + assert "already registered" in combined + assert "--spec-override" in combined + + +# --------------------------------------------------------------------------- +# Precedence +# --------------------------------------------------------------------------- + +def test_operation_precedence_spec_then_kernel_then_declared(): + bench = pytest.importorskip("bench") + # --spec wins + assert bench.resolve_operation_name("from_spec", "from_kernel", "declared") == "from_spec" + # --kernel wins over kernel.py + assert bench.resolve_operation_name(None, "from_kernel", "declared") == "from_kernel" + # kernel.py::KERNEL_TYPE is the fallback + assert bench.resolve_operation_name(None, None, "declared") == "declared" + # nothing selected + assert bench.resolve_operation_name(None, None, None) is None + + +def test_legacy_kernel_configs_view_is_derived_from_the_registry(): + bench = pytest.importorskip("bench") + registry = create_builtin_registry() + assert tuple(bench.KERNEL_CONFIGS) == registry.list_names() + entry = bench.KERNEL_CONFIGS["matmul"] + assert entry["spec"] is not None + assert entry["test_sizes"][0] == ("tiny", {"M": 128, "N": 128, "K": 128}) + assert entry["edge_sizes"][0][0] == "edge_1023" + + +def test_bench_reads_metadata_only_from_specs(): + """No operation-specific configuration literals may return to bench.py.""" + source = (REPO_ROOT / "bench.py").read_text() + assert "gen_matmul_inputs" not in source + assert '"test_sizes": [' not in source + assert "_ref_layernorm" not in source + + +# --------------------------------------------------------------------------- +# extract.py command line and generation +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize( + "flag", ["--report", "--top", "--kernel-type", "--backend", "--spec", "--spec-override"] +) +def test_extract_help_lists_expected_flags(flag): + result = run_script("extract.py", "--help") + assert result.returncode == 0, result.stderr + assert flag in result.stdout + + +def test_extract_has_no_duplicated_operation_metadata(): + source = (REPO_ROOT / "extract.py").read_text() + for removed in ( + "SHAPE_KEYS", + "SHAPE_ALIAS_MAP", + "TOLERANCES_MAP", + "FLOPS_FN_SRC", + "BYTES_FN_SRC", + "SPEEDUP_ESTIMATES", + ): + assert removed not in source, f"{removed} must live in the spec registry" + + +def test_extract_parses_profiler_shapes_through_spec_aliases(): + extract = pytest.importorskip("extract") + registry = create_builtin_registry() + + layernorm = registry.get("layernorm") + assert extract.parse_shape_info("M=4096, N=2048", layernorm) == { + "batch": 4096, + "dim": 2048, + } + + attention = registry.get("flash_attention") + assert extract.parse_shape_info("B=1, H=32, N=4096, D=128", attention) == { + "batch": 1, + "heads": 32, + "seq_len": 4096, + "head_dim": 128, + } + + matmul = registry.get("matmul") + assert extract.parse_shape_info("M=8, N=9, K=10", matmul) == {"M": 8, "N": 9, "K": 10} + assert extract.parse_shape_info("", matmul) is None + assert extract.parse_shape_info("no numbers here", matmul) is None + + +def test_extract_default_shapes_match_pre_refactor_values(): + extract = pytest.importorskip("extract") + registry = create_builtin_registry() + assert extract.get_default_shape(registry.get("matmul")) == { + "M": 2048, "N": 2048, "K": 2048 + } + # 'reduce' deliberately keeps its pre-refactor 4096x4096 fallback even though + # its 'large' benchmark size is 8192x8192. + assert extract.get_default_shape(registry.get("reduce")) == {"M": 4096, "N": 4096} + + +def test_extract_generates_a_kernel_file_for_a_builtin(): + extract = pytest.importorskip("extract") + spec = create_builtin_registry().get("matmul") + starter = extract.read_starter_kernel(spec, backend="triton") + assert starter is not None + + content = extract.generate_kernel_file( + spec=spec, + rank=1, + pct_total=42.0, + model_shape={"M": 4096, "N": 4096, "K": 4096}, + model_name="unit-test-model", + gpu_time_ms=1.5, + starter_code=starter, + backend="triton", + ) + + assert 'KERNEL_TYPE = "matmul"' in content + assert "MODEL_SHAPES = {'M': 4096, 'N': 4096, 'K': 4096}" in content + assert "'float16': {'atol': 0.01, 'rtol': 0.01}" in content + assert "return 2 * s['M'] * s['N'] * s['K']" in content + assert "dt_bytes" in content + # The generated file must be valid Python. + compile(content, "generated_matmul.py", "exec") + + +def test_extract_generates_a_kernel_file_for_an_external_spec(): + extract = pytest.importorskip("extract") + spec = load_spec(FIXTURE_LOCATOR) + starter = extract.read_starter_kernel(spec, backend="triton") + assert starter is not None + + content = extract.generate_kernel_file( + spec=spec, + rank=1, + pct_total=0.0, + model_shape=spec.extraction_shape(), + model_name="external-spec", + gpu_time_ms=0.0, + starter_code=starter, + backend="triton", + spec_locator=FIXTURE_LOCATOR, + ) + + assert 'KERNEL_TYPE = "fixture_add"' in content + assert f'KERNEL_SPEC = "{FIXTURE_LOCATOR}"' in content + assert "return s['rows'] * s['cols']" in content + compile(content, "generated_fixture_add.py", "exec") + + +def test_extract_falls_back_to_the_spec_when_accounting_is_opaque(): + extract = pytest.importorskip("extract") + spec = load_spec(FIXTURE_LOCATOR) + body = extract._accounting_body(lambda s: 1, spec, FIXTURE_LOCATOR) + assert "NotImplementedError" in body + assert FIXTURE_LOCATOR in body + compile(f"def FLOPS_FN(s):\n {body}\n", "opaque.py", "exec") + + +def test_extract_plan_uses_spec_speedup_estimates(): + extract = pytest.importorskip("extract") + plan = extract.generate_optimization_plan( + [ + { + "rank": 1, + "op_type": "matmul", + "pct_total": 30.0, + "gpu_time_ms": 2.0, + "model_shape": {"M": 1, "N": 1, "K": 1}, + "output_file": "workspace/kernel_matmul_1.py", + "estimated_speedup_potential": "2-3x", + } + ] + ) + assert plan["total_optimization_targets"] == 1 + assert plan["covered_gpu_time_pct"] == 30.0 + assert plan["kernels_to_optimize"][0]["estimated_speedup_potential"] == "2-3x" + + +def test_extract_synthesizes_a_target_from_a_spec_alone(): + extract = pytest.importorskip("extract") + spec = load_spec(FIXTURE_LOCATOR) + entry = extract._synthetic_report_entry(spec) + assert entry["op_type"] == "fixture_add" + assert entry["autokernel_supported"] is True + assert entry["shapes"] == spec.extraction_shape() + + +# --------------------------------------------------------------------------- +# bench.py shape-corpus command line +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize( + "flag", + [ + "--shape-corpus", + "--shape-corpus-only", + "--check-backward", + "--check-compile", + "--result-json", + ], +) +def test_bench_help_lists_verification_flags(flag): + result = run_script("bench.py", "--help") + assert result.returncode == 0, result.stderr + assert flag in result.stdout + + +def test_bench_corpus_only_requires_corpus_path(): + result = run_script("bench.py", "--shape-corpus-only") + assert result.returncode == 2 + assert "requires --shape-corpus" in result.stderr + + +def test_bench_invalid_corpus_fails_before_gpu_detection(tmp_path: Path): + """An invalid corpus must fail before any GPU probing or allocation.""" + bad = tmp_path / "bad_corpus.json" + bad.write_text(json.dumps({"schema_version": 1, "operation": "matmul", "cases": []})) + result = run_script("bench.py", "--kernel", "matmul", "--shape-corpus", str(bad)) + combined = result.stdout + result.stderr + assert result.returncode == 1 + assert "correctness: FAIL" in combined + assert "throughput_tflops: 0.000" in combined + assert str(bad) in combined + # The GPU info block is printed only after detection; it must not appear. + assert "gpu_name:" not in combined + + +def test_bench_corpus_operation_mismatch_is_actionable(tmp_path: Path): + corpus = tmp_path / "corpus.json" + corpus.write_text(json.dumps({ + "schema_version": 1, + "operation": "some_other_op", + "cases": [{"name": "a", "size": {"M": 4, "N": 4, "K": 4}}], + })) + result = run_script("bench.py", "--kernel", "matmul", "--shape-corpus", str(corpus)) + combined = result.stdout + result.stderr + assert result.returncode == 1 + assert "does not match the selected spec" in combined + + +def test_profile_cli_does_not_break_standard_cprofile_import(): + """The top-level profile.py must preserve cProfile's expected API.""" + result = subprocess.run( + [ + sys.executable, + "-c", + "import cProfile, profile; " + "assert callable(cProfile.run); " + "assert callable(profile.run); " + "assert hasattr(profile, '_Utils')", + ], + cwd=str(REPO_ROOT), + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/test_compile_verification.py b/tests/test_compile_verification.py new file mode 100644 index 00000000..2420e50d --- /dev/null +++ b/tests/test_compile_verification.py @@ -0,0 +1,99 @@ +"""CPU tests for optional torch.compile verification.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from autokernel.specs import CompileSpec, Tolerance +from autokernel.verification.compile import check_compile +from conftest import make_spec + + +def _generator(size: dict[str, int], dtype: str, device: str, seed: int) -> dict: + torch.manual_seed(seed) + return {"x": torch.randn(size["rows"], size["cols"], device=device)} + + +def _spec(**overrides: Any): + return make_spec( + reference_fn=lambda x: {"output": x * 2, "metadata": 7}, + input_generator=_generator, + dtypes=("float32",), + tolerances={"float32": Tolerance(1e-6, 1e-6)}, + **overrides, + ) + + +def test_compile_uses_declared_options_and_runs_twice(monkeypatch): + calls: list[dict[str, Any]] = [] + executions = 0 + + def fake_compile(fn, **options): + calls.append(options) + + def compiled(**inputs): + nonlocal executions + executions += 1 + return fn(**inputs) + + return compiled + + monkeypatch.setattr(torch, "compile", fake_compile) + report = check_compile( + lambda x: {"output": x * 2, "metadata": 7}, + _spec(compile_spec=CompileSpec(fullgraph=True, dynamic=False)), + device="cpu", + ) + + assert report.status == "PASS" + assert calls == [{"fullgraph": True, "dynamic": False}] + assert executions == 2 + assert [case.label for case in report.cases] == ["small"] + + +def test_dynamic_compile_reuses_callable_for_two_shapes(monkeypatch): + compile_calls = 0 + observed_shapes: list[tuple[int, ...]] = [] + + def fake_compile(fn, **options): + nonlocal compile_calls + compile_calls += 1 + + def compiled(**inputs): + observed_shapes.append(tuple(inputs["x"].shape)) + return fn(**inputs) + + return compiled + + monkeypatch.setattr(torch, "compile", fake_compile) + report = check_compile( + lambda x: {"output": x * 2, "metadata": 7}, + _spec(compile_spec=CompileSpec(dynamic=True)), + device="cpu", + ) + + assert report.status == "PASS" + assert compile_calls == 1 + assert len(report.cases) == 2 + assert observed_shapes == [(4, 4), (4, 4), (8, 8), (8, 8)] + + +def test_compile_mismatch_is_not_passed(monkeypatch): + monkeypatch.setattr(torch, "compile", lambda fn, **options: fn) + report = check_compile( + lambda x: {"output": x * 3, "metadata": 7}, + _spec(), + device="cpu", + ) + assert report.status == "FAIL" + assert report.cases[0].status == "FAIL" + assert "max_abs_error" in report.reason + + +def test_compile_unavailable_is_unsupported(monkeypatch): + monkeypatch.delattr(torch, "compile", raising=False) + report = check_compile(lambda x: x, _spec(), device="cpu") + assert report.status == "UNSUPPORTED" + assert "unavailable" in report.reason diff --git a/tests/test_correlation.py b/tests/test_correlation.py new file mode 100644 index 00000000..967cbffb --- /dev/null +++ b/tests/test_correlation.py @@ -0,0 +1,688 @@ +"""Tests for offline correlation between profiler rows and FX regions.""" + +from __future__ import annotations + +import pytest + +from autokernel.discovery import ( + DiscoveryReport, + GraphRegion, + OperatorHotspot, + TensorMeta, + correlate_discovery_report, + correlate_profiler_to_regions, +) + + +def _tensor(name: str = "x", shape: tuple[int, ...] = (1, 128, 64)) -> TensorMeta: + """Helper to create a simple tensor metadata for testing.""" + stride: list[int] = [] + running = 1 + for dim in reversed(shape): + stride.append(running) + running *= max(dim, 1) + return TensorMeta( + name=name, + shape=shape, + stride=tuple(reversed(stride)), + dtype="bfloat16", + device_type="cuda", + ) + + +def test_correlate_profiler_to_regions_basic_match(): + """Test basic profiler-to-region correlation with exact parent module match.""" + # Create profiler rows + profiler_rows = [ + OperatorHotspot( + name="aten::mm", + op_key="aten::mm", + calls=100, + cuda_time_us=6000.0, + self_cuda_time_us=5900.0, + parent_module="blocks.0.attn", + ), + OperatorHotspot( + name="aten::mul", + op_key="aten::mul", + calls=40, + cuda_time_us=400.0, + self_cuda_time_us=400.0, + parent_module="blocks.0.norm", + ), + ] + + # Create FX regions + fx_regions = [ + GraphRegion.build( + name="attention", + operations=["aten::mm", "aten::add"], + inputs=[_tensor("q"), _tensor("k")], + parent_module="blocks.0.attn", + calls=100, + ), + GraphRegion.build( + name="normalization", + operations=["aten::mul", "aten::layer_norm"], + inputs=[_tensor("x")], + parent_module="blocks.0.norm", + calls=40, + ), + ] + + # Correlate + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=10000.0, + ) + + # Should have 2 regions and no unmatched rows + assert len(correlated) == 2 + assert unmatched == () + + # Check that timing data was populated + attn_region = next(r for r in correlated if r.name == "attention") + assert attn_region.self_cuda_time_us == 5900.0 + assert attn_region.calls == 100 + assert attn_region.attributes["matched_profiler_rows"] == 1 + + norm_region = next(r for r in correlated if r.name == "normalization") + assert norm_region.self_cuda_time_us == 400.0 + assert norm_region.calls == 40 + assert norm_region.attributes["matched_profiler_rows"] == 1 + + +def test_correlate_unmatched_profiler_rows(): + """Test that unmatched profiler rows are returned separately.""" + profiler_rows = [ + OperatorHotspot( + name="aten::mm", + op_key="aten::mm", + calls=100, + cuda_time_us=6000.0, + self_cuda_time_us=5900.0, + parent_module="blocks.0.attn", + ), + OperatorHotspot( + name="custom::unknown_op", + op_key="custom::unknown_op", + calls=10, + cuda_time_us=100.0, + self_cuda_time_us=100.0, + parent_module="unknown.module", + ), + ] + + fx_regions = [ + GraphRegion.build( + name="attention", + operations=["aten::mm", "aten::add"], + inputs=[_tensor("q"), _tensor("k")], + parent_module="blocks.0.attn", + calls=100, + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=10000.0, + ) + + # Only the captured region is returned; the custom op row is unmatched + assert len(correlated) == 1 + assert correlated[0].name == "attention" + assert len(unmatched) == 1 + assert unmatched[0].op_key == "custom::unknown_op" + assert unmatched[0].self_cuda_time_us == 100.0 + assert unmatched[0].calls == 10 + + +def test_deduplicate_equivalent_regions(): + """Test that equivalent regions are deduplicated by fingerprint.""" + profiler_rows = [ + OperatorHotspot( + name="aten::mul", + op_key="aten::mul", + calls=80, + cuda_time_us=800.0, + self_cuda_time_us=800.0, + parent_module="blocks.0.norm", + ), + ] + + # Create two equivalent regions (same operations, inputs, parent_module, etc.) + # Note: fingerprint includes parent_module, so they must be identical to deduplicate + fx_regions = [ + GraphRegion.build( + name="norm_0", + operations=["aten::mul", "aten::layer_norm"], + inputs=[_tensor("x")], + parent_module="blocks.0.norm", + calls=40, + shape_frequency={"shape1": 40}, + ), + GraphRegion.build( + name="norm_1", + operations=["aten::mul", "aten::layer_norm"], + inputs=[_tensor("x")], + parent_module="blocks.0.norm", # Same parent_module for same fingerprint + calls=40, + shape_frequency={"shape1": 40}, + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=10000.0, + ) + + # Should deduplicate to one region + assert len(correlated) == 1 + + # Should aggregate calls and shape frequencies + region = correlated[0] + assert region.calls == 80 # Aggregated from both regions + assert region.shape_frequency == {"shape1": 80} # Aggregated + + +def test_exclusive_time_without_double_counting(): + """Test that nested scopes are not double-counted.""" + # Simulate nested profiler rows + profiler_rows = [ + OperatorHotspot( + name="blocks.0.forward", + op_key="blocks.0.forward", + calls=10, + cuda_time_us=5000.0, + self_cuda_time_us=1000.0, # Exclusive time (excluding children) + parent_module="blocks.0", + ), + OperatorHotspot( + name="aten::mm", + op_key="aten::mm", + calls=10, + cuda_time_us=4000.0, + self_cuda_time_us=4000.0, + parent_module="blocks.0.attn", + ), + ] + + fx_regions = [ + GraphRegion.build( + name="block_forward", + operations=["aten::mm", "aten::add"], + inputs=[_tensor("x")], + parent_module="blocks.0", + calls=10, + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=10000.0, + ) + + region = correlated[0] + # Should use exclusive time to avoid double-counting + # The region should get the appropriate time based on matched rows + assert region.self_cuda_time_us > 0 + # Total should not exceed the sum of exclusive times + assert region.self_cuda_time_us <= 5000.0 + + +def test_synthetic_cpu_fixtures_produce_timed_regions(): + """Test that synthetic CPU fixtures produce non-zero timed regions.""" + # Create profiler rows with synthetic CPU timing + profiler_rows = [ + OperatorHotspot( + name="aten::mul", + op_key="aten::mul", + calls=50, + cuda_time_us=0.0, # CPU-only + self_cuda_time_us=0.0, + cpu_time_us=500.0, # CPU time + parent_module="cpu.module", + ), + ] + + fx_regions = [ + GraphRegion.build( + name="cpu_region", + operations=["aten::mul", "aten::add"], + inputs=[_tensor("x")], + parent_module="cpu.module", + calls=50, + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=0.0, # CPU-only + ) + + assert len(correlated) == 1 + region = correlated[0] + # Should still have calls even if CUDA time is zero + assert region.calls == 50 + # CPU regions should be tracked even with zero CUDA time + assert region.attributes["matched_profiler_rows"] == 1 + + +def test_correlate_discovery_report_integration(): + """Test the high-level correlate_discovery_report function.""" + profiler_export_rows = [ + { + "name": "aten::mm", + "op_key": "aten::mm", + "calls": 100, + "cuda_time_us": 6000.0, + "self_cuda_time_us": 5900.0, + "parent_module": "blocks.0.attn", + }, + ] + + fx_report = DiscoveryReport.from_dict( + { + "schema_version": 1, + "producer": {"name": "fastvideo", "version": "test"}, + "workload": {"workload_id": "test", "model_id": "test"}, + "environment": {"hardware_profile_id": "cpu", "software_profile_id": "test"}, + "total_cuda_time_us": 10000.0, + "operators": [], + "regions": [ + GraphRegion.build( + name="attention", + operations=["aten::mm", "aten::add"], + inputs=[_tensor("q"), _tensor("k")], + parent_module="blocks.0.attn", + calls=100, + ).as_dict() + ], + "graph_breaks": [], + "unsupported": [], + } + ) + + correlated_report = correlate_discovery_report(profiler_export_rows, fx_report) + + # Should have operators from profiler + assert len(correlated_report.operators) == 1 + assert correlated_report.operators[0].name == "aten::mm" + + # Should have populated regions + assert len(correlated_report.regions) == 1 + assert correlated_report.regions[0].self_cuda_time_us == 5900.0 + + # Other fields should be preserved + assert correlated_report.workload == fx_report.workload + assert correlated_report.environment == fx_report.environment + + +def test_correlate_discovery_report_keeps_unmatched_as_diagnostic(): + profiler_export_rows = [ + { + "name": "aten::unmatched", + "calls": 3, + "cuda_time_us": 30.0, + "self_cuda_time_us": 30.0, + }, + ] + fx_report = DiscoveryReport.from_dict( + { + "schema_version": 1, + "producer": {"name": "fastvideo", "version": "test"}, + "workload": {"workload_id": "test", "model_id": "test"}, + "environment": {"hardware_profile_id": "cpu"}, + "total_cuda_time_us": 30.0, + "operators": [], + "regions": [ + GraphRegion.build( + name="captured", + operations=["aten::add"], + inputs=[_tensor("x")], + ).as_dict(), + ], + } + ) + + correlated = correlate_discovery_report( + profiler_export_rows, + fx_report, + ) + + assert all( + region.name != "unmatched_profiler_rows" + for region in correlated.regions + ) + assert any( + item.op_name == "profiler::unmatched" + for item in correlated.unsupported + ) + + +def test_overload_qualified_op_key_matches_normalized_operations(): + """Profiler rows keep overload suffixes; region ops are normalized.""" + profiler_rows = [ + OperatorHotspot( + name="aten::add.Tensor", + op_key="aten::add.Tensor", + calls=20, + cuda_time_us=200.0, + self_cuda_time_us=200.0, + ), + ] + fx_regions = [ + GraphRegion.build( + name="residual", + operations=["aten::add"], + inputs=[_tensor("x")], + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=1000.0, + ) + + assert unmatched == () + assert correlated[0].attributes["matched_profiler_rows"] == 1 + assert correlated[0].self_cuda_time_us == 200.0 + + +def test_correlated_report_with_unmatched_rows_round_trips(): + """A report produced from unmatched rows must survive serialization.""" + profiler_export_rows = [ + { + "name": "custom::unknown_op", + "calls": 3, + "cuda_time_us": 30.0, + "self_cuda_time_us": 30.0, + }, + ] + fx_report = DiscoveryReport.from_dict( + { + "schema_version": 1, + "producer": {"name": "fastvideo", "version": "test"}, + "workload": {"workload_id": "test", "model_id": "test"}, + "environment": {"hardware_profile_id": "cpu"}, + "total_cuda_time_us": 30.0, + "operators": [], + "regions": [ + GraphRegion.build( + name="captured", + operations=["aten::add"], + inputs=[_tensor("x")], + ).as_dict(), + ], + } + ) + + correlated = correlate_discovery_report(profiler_export_rows, fx_report) + reloaded = DiscoveryReport.from_dict(correlated.as_dict()) + assert len(reloaded.regions) == 1 + assert any( + item.op_name == "profiler::unmatched" + for item in reloaded.unsupported + ) + + +def test_confidence_calculation(): + """Test that confidence is calculated correctly.""" + profiler_rows = [ + OperatorHotspot( + name="aten::mul", + op_key="aten::mul", + calls=100, + cuda_time_us=1000.0, + self_cuda_time_us=1000.0, + parent_module="blocks.0.norm", + ), + ] + + fx_regions = [ + GraphRegion.build( + name="safe_region", + operations=["aten::mul", "aten::add"], # Safe operations + inputs=[_tensor("x")], + parent_module="blocks.0.norm", + calls=100, + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=10000.0, + ) + + region = correlated[0] + # Should have high confidence: safe + matched + high call count + assert region.attributes["confidence"] >= 0.8 + + +def test_e2e_improvement_estimation(): + """Test that end-to-end improvement is estimated correctly.""" + profiler_rows = [ + OperatorHotspot( + name="aten::mm", + op_key="aten::mm", + calls=100, + cuda_time_us=2000.0, # 20% of total + self_cuda_time_us=2000.0, + parent_module="blocks.0.attn", + ), + ] + + fx_regions = [ + GraphRegion.build( + name="attention", + operations=["aten::mm", "aten::add"], + inputs=[_tensor("x")], + parent_module="blocks.0.attn", + calls=100, + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=10000.0, + ) + + region = correlated[0] + # 20% e2e share * 0.9 reducible fraction = 18% max improvement + assert region.attributes["e2e_share_pct"] == 20.0 + assert region.attributes["estimated_max_e2e_improvement"] == pytest.approx(0.18, rel=0.01) + + +def test_shape_frequency_aggregation(): + """Test that shape frequencies are aggregated correctly.""" + profiler_rows = [ + OperatorHotspot( + name="aten::mul", + op_key="aten::mul", + calls=50, + cuda_time_us=500.0, + self_cuda_time_us=500.0, + input_shapes=[(2, 128)], # Shape metadata + parent_module="blocks.0.norm", + ), + ] + + fx_regions = [ + GraphRegion.build( + name="norm", + operations=["aten::mul", "aten::layer_norm"], + inputs=[_tensor("x", (2, 128))], + parent_module="blocks.0.norm", + calls=50, + shape_frequency={"2x128": 30}, # Existing shape frequency + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=10000.0, + ) + + region = correlated[0] + # Should merge shape frequencies from both sources + assert region.shape_frequency is not None + # Should have at least the original frequency + assert sum(region.shape_frequency.values()) >= 30 + + +def test_rejection_reasons_aggregation(): + """Test that rejection reasons are aggregated from multiple sources.""" + profiler_rows = [ + OperatorHotspot( + name="aten::mul_", + op_key="aten::mul_", + calls=10, + cuda_time_us=100.0, + self_cuda_time_us=100.0, + parent_module="blocks.0.norm", + ), + ] + + fx_regions = [ + GraphRegion.build( + name="unsafe_region", + operations=["aten::mul_"], # In-place operation + inputs=[_tensor("x")], + parent_module="blocks.0.norm", + calls=10, + rejection_reasons=("custom_reason",), + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=10000.0, + ) + + region = correlated[0] + # Should have both safety rejection and custom reason + assert len(region.rejection_reasons) >= 1 + # Should include in-place mutation rejection + assert any("in-place" in reason.lower() for reason in region.rejection_reasons) + + +def test_hierarchical_parent_module_matching(): + """Test that hierarchical parent module relationships are matched.""" + profiler_rows = [ + OperatorHotspot( + name="aten::mm", + op_key="aten::mm", + calls=100, + cuda_time_us=5000.0, + self_cuda_time_us=5000.0, + parent_module="blocks.0.attn.q_proj", # More specific + ), + ] + + fx_regions = [ + GraphRegion.build( + name="attention", + operations=["aten::mm", "aten::add"], + inputs=[_tensor("x")], + parent_module="blocks.0.attn", # Less specific + calls=100, + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=10000.0, + ) + + # Should still match due to hierarchical relationship + assert len(correlated) == 1 + assert correlated[0].attributes["matched_profiler_rows"] == 1 + + +def test_exact_region_range_disambiguates_shared_operation(): + profiler_rows = [ + OperatorHotspot( + name="blocks.89abcdef", + op_key="blocks.89abcdef", + calls=12, + cuda_time_us=800.0, + self_cuda_time_us=50.0, + parent_module="blocks", + ), + ] + fx_regions = [ + GraphRegion.build( + name="blocks.01234567", + operations=["aten::mul"], + inputs=[_tensor("x", (2, 128))], + parent_module="blocks", + calls=12, + ), + GraphRegion.build( + name="blocks.89abcdef", + operations=["aten::mul"], + inputs=[_tensor("x", (8, 128))], + parent_module="blocks", + calls=12, + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=1000.0, + ) + + target = next(region for region in correlated if region.name == "blocks.89abcdef") + other = next(region for region in correlated if region.name == "blocks.01234567") + assert target.cuda_time_us == 800.0 + assert target.self_cuda_time_us == 800.0 + assert target.attributes["matched_profiler_rows"] == 1 + assert other.cuda_time_us == 0.0 + + +def test_ambiguous_op_only_row_remains_unmatched(): + profiler_rows = [ + OperatorHotspot( + name="aten::mul", + op_key="aten::mul", + calls=100, + cuda_time_us=900.0, + self_cuda_time_us=900.0, + ), + ] + fx_regions = [ + GraphRegion.build( + name="first", + operations=["aten::mul"], + inputs=[_tensor("x", (2, 128))], + parent_module="blocks", + ), + GraphRegion.build( + name="second", + operations=["aten::mul"], + inputs=[_tensor("x", (8, 128))], + parent_module="blocks", + ), + ] + + correlated, unmatched = correlate_profiler_to_regions( + profiler_rows, + fx_regions, + total_cuda_time_us=1000.0, + ) + + assert all(region.cuda_time_us == 0.0 for region in correlated) + assert len(unmatched) == 1 + assert unmatched[0].op_key == "aten::mul" diff --git a/tests/test_discovery.py b/tests/test_discovery.py new file mode 100644 index 00000000..97725a54 --- /dev/null +++ b/tests/test_discovery.py @@ -0,0 +1,217 @@ +"""Discovery report schema, fingerprint stability, and region safety.""" + +from __future__ import annotations + +import json + +import pytest + +from autokernel.discovery import ( + DiscoveryError, + DiscoveryReport, + GraphRegion, + TensorMeta, + graph_fingerprint, + is_region_safe, + load_discovery_report, + reject_region, + write_discovery_report, +) + + +def _tensor(name: str = "x") -> TensorMeta: + return TensorMeta( + name=name, + shape=(1, 128, 64), + stride=(8192, 64, 1), + dtype="bfloat16", + device_type="cuda", + ) + + +def test_fingerprint_stable_across_equivalent_regions(): + ops = ("aten::mul", "aten::add", "aten::layer_norm") + inputs = [_tensor("residual").signature_dict(), _tensor("gate").signature_dict()] + a = graph_fingerprint(operations=ops, input_signatures=inputs) + b = graph_fingerprint(operations=ops, input_signatures=inputs) + assert a == b + assert len(a) == 32 + + different = graph_fingerprint( + operations=("aten::add", "aten::mul"), + input_signatures=inputs, + ) + assert different != a + + +def test_fingerprint_rejects_nonfinite_and_unsupported_constants(): + with pytest.raises(ValueError, match="finite"): + graph_fingerprint( + operations=("aten::add",), + input_signatures=[_tensor().signature_dict()], + safe_constants={"scale": float("inf")}, + ) + with pytest.raises(ValueError, match="unsupported"): + graph_fingerprint( + operations=("aten::add",), + input_signatures=[_tensor().signature_dict()], + safe_constants={"opaque": object()}, + ) + + +def test_rejects_in_place_mutation_and_invalid_region_name(): + assert any("in-place mutation" in reason for reason in reject_region(["aten::add_"])) + with pytest.raises(ValueError, match="invalid graph region name"): + GraphRegion.build( + name="invalid name", + operations=["aten::add"], + inputs=[_tensor()], + ) + + +def test_graph_region_build_and_report_roundtrip(tmp_path): + region = GraphRegion.build( + name="elementwise.mul_add", + operations=["aten::mul", "aten::add"], + inputs=[_tensor("a"), _tensor("b")], + outputs=[_tensor("out")], + parent_module="blocks.0", + cuda_time_us=1200.0, + self_cuda_time_us=1100.0, + calls=40, + pattern_family="elementwise_chain", + shape_frequency={"b1-s128-d64": 40}, + ) + assert is_region_safe(region.operations) + assert not region.rejection_reasons + + report = DiscoveryReport.from_dict( + { + "schema_version": 1, + "producer": {"name": "fastvideo", "version": "test"}, + "workload": { + "workload_id": "ltx-t2v-480p", + "model_id": "FastVideo/LTX2-Distilled-Diffusers", + }, + "environment": { + "hardware_profile_id": "test-gpu", + "software_profile_id": "torch-test", + }, + "total_cuda_time_us": 10000.0, + "operators": [ + { + "name": "aten::mm", + "op_key": "aten::mm", + "calls": 100, + "cuda_time_us": 6000.0, + "self_cuda_time_us": 5900.0, + }, + { + "name": "aten::mul", + "op_key": "aten::mul", + "calls": 40, + "cuda_time_us": 400.0, + "self_cuda_time_us": 400.0, + }, + ], + "regions": [region.as_dict()], + "graph_breaks": [ + { + "scope": "dit.forward", + "reason": "data-dependent branching", + "count": 2, + } + ], + "unsupported": [ + { + "op_name": "custom::flash", + "reason": "unsupported custom operator", + "count": 1, + } + ], + } + ) + ranked = report.ranked_operators() + assert ranked[0].op_key == "aten::mm" + assert ranked[0].impact_pct(report.total_cuda_time_us) == pytest.approx(59.0) + assert report.graph_breaks[0].reason.startswith("data-dependent") + + path = tmp_path / "discovery.json" + write_discovery_report(report, path) + loaded = load_discovery_report(path) + assert loaded.as_dict() == report.as_dict() + + +def test_reject_collectives_and_data_dependent_ops(): + reasons = reject_region( + ["aten::mul", "aten::all_reduce", "aten::item", "custom::foo"] + ) + assert any("collective" in r for r in reasons) + assert any("data-dependent" in r for r in reasons) + assert any("custom" in r for r in reasons) + assert not is_region_safe(["aten::mul", "c10d::all_reduce_"]) + + +def test_allow_elementwise_chain(): + assert is_region_safe( + ["aten::mul", "aten::add", "aten::layer_norm", "aten::silu"] + ) + + +def test_discovery_rejects_prompt_payload(): + payload = { + "schema_version": 1, + "producer": {"name": "fastvideo", "version": "test"}, + "workload": {"workload_id": "x", "model_id": "m"}, + "environment": {"hardware_profile_id": "h", "software_profile_id": "s"}, + "total_cuda_time_us": 1.0, + "operators": [], + "regions": [], + } + payload["operators"] = [ + { + "name": "aten::add", + "op_key": "aten::add", + "calls": 1, + "cuda_time_us": 1.0, + "self_cuda_time_us": 1.0, + "attributes": {"prompt": "secret"}, + } + ] + with pytest.raises(DiscoveryError, match="secret fields"): + DiscoveryReport.from_dict(payload) + + +def test_wan_like_elementwise_region_is_safe_but_low_impact(): + """Wan fusion shapes are discoverable; ranking must still use e2e share.""" + region = GraphRegion.build( + name="wan.gated_residual_norm", + operations=[ + "aten::mul", + "aten::add", + "aten::layer_norm", + ], + inputs=[ + TensorMeta( + "residual", + (1, 20280, 1536), + (31150080, 1536, 1), + "bfloat16", + "cuda", + ), + TensorMeta( + "gate", + (1, 1, 1536), + (1536, 1536, 1), + "float32", + "cuda", + ), + ], + cuda_time_us=50.0, + calls=40, + pattern_family="residual_gate_norm", + ) + assert is_region_safe(region.operations) + # ~0.5% of a 10ms e2e window => below default 0.5% search floor when + # optimistic reducible fraction is considered in WS3. + assert region.impact_pct(10_000.0) == pytest.approx(0.5) diff --git a/tests/test_fx_region_capture.py b/tests/test_fx_region_capture.py new file mode 100644 index 00000000..3fdbd70f --- /dev/null +++ b/tests/test_fx_region_capture.py @@ -0,0 +1,197 @@ +"""CPU tests for model-independent Dynamo/FX region capture.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from autokernel.discovery import ( + DiscoveryReport, + RegionCaptureSession, + capture_model_regions, + capture_module_region, + is_region_safe, + reject_region, +) + + +class _PureBlock(nn.Module): + def forward(self, x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.silu(x * scale + x) + + +class _MutatingBlock(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + x.add_(1.0) + return x + + +class _TinyStack(nn.Module): + def __init__(self) -> None: + super().__init__() + self.block0 = _PureBlock() + self.block1 = _PureBlock() + + def forward(self, x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + x = self.block0(x, scale) + x = self.block1(x, scale) + return x + + +def test_cpu_module_produces_graph_region_with_ops_and_deps(): + module = _PureBlock() + x = torch.randn(2, 8) + scale = torch.randn(2, 8) + result = capture_module_region( + module, + (x, scale), + name="test.pure_block", + parent_module="blocks.0", + ) + assert result.region is not None + region = result.region + assert region.operations + assert region.inputs + assert region.outputs + assert region.parent_module == "blocks.0" + assert region.fingerprint + # Dependencies should encode at least one edge for a multi-op graph. + assert isinstance(region.dependencies, tuple) + # Metadata only — no forbidden keys + payload = region.as_dict() + assert "values" not in str(payload).lower() or "values" not in payload + for key in ("weights", "prompt", "tensor_values", "source_code"): + assert key not in payload + + +def test_equivalent_graphs_stable_fingerprints(): + module = _PureBlock() + a = capture_module_region( + module, + (torch.randn(2, 8), torch.randn(2, 8)), + name="test.pure_block", + parent_module="blocks.0", + ) + b = capture_module_region( + module, + (torch.randn(2, 8), torch.randn(2, 8)), + name="test.pure_block", + parent_module="blocks.0", + ) + assert a.region is not None and b.region is not None + # Same structure + same shapes/dtypes => same fingerprint (values ignored). + assert a.region.operations == b.region.operations + assert a.region.fingerprint == b.region.fingerprint + + +def test_fingerprint_changes_with_shape(): + module = _PureBlock() + small = capture_module_region( + module, + (torch.randn(2, 8), torch.randn(2, 8)), + name="test.pure_block", + parent_module="blocks.0", + ) + large = capture_module_region( + module, + (torch.randn(4, 16), torch.randn(4, 16)), + name="test.pure_block", + parent_module="blocks.0", + ) + assert small.region is not None and large.region is not None + assert small.region.fingerprint != large.region.fingerprint + + +def test_mutation_rejected_fail_closed(): + module = _MutatingBlock() + result = capture_module_region( + module, + (torch.randn(2, 4),), + name="test.mutating", + parent_module="bad", + ) + # Either trace fails or region is rejected for in-place mutation. + if result.region is not None: + assert result.region.rejection_reasons + assert not is_region_safe(result.region.operations) + else: + assert result.graph_breaks + + +def test_collectives_and_custom_ops_rejected(): + reasons = reject_region( + ["aten::mul", "c10d::all_reduce_", "custom::flash_attn"] + ) + assert any("collective" in r for r in reasons) + assert any("custom" in r for r in reasons) + + +def test_repeated_module_capture_session_builds_report(): + model = _TinyStack() + x = torch.randn(2, 8) + scale = torch.randn(2, 8) + + with capture_model_regions( + model, + predicate=lambda name, _m: name in {"block0", "block1"}, + name_prefix="dit", + ) as session: + # Multiple calls → aggregated call counts / shape frequency. + for _ in range(3): + model(x, scale) + + regions = session.regions() + assert regions + # Same pure block structure should collapse to one fingerprint class + # (two scopes may still share ops; at least one region with calls>=3). + assert any(region.calls >= 3 for region in regions) or any( + sum((region.shape_frequency or {}).values()) >= 3 for region in regions + ) + + report = session.to_discovery_report( + workload={"workload_id": "unit-fx", "model_id": "toy"}, + total_cuda_time_us=0.0, + ) + assert isinstance(report, DiscoveryReport) + assert report.regions + # Graph breaks / unsupported must be present as fields (may be empty for pure). + assert isinstance(report.graph_breaks, tuple) + assert isinstance(report.unsupported, tuple) + # Round-trip + again = DiscoveryReport.from_dict(report.as_dict()) + assert [r.fingerprint for r in again.regions] == [ + r.fingerprint for r in report.regions + ] + + +def test_graph_breaks_visible_for_untraceable_control_flow(): + class _DataDependent(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.sum().item() > 0: + return x * 2 + return x + + result = capture_module_region( + _DataDependent(), + (torch.randn(2, 2),), + name="test.data_dependent", + ) + # symbolic_trace typically fails on data-dependent Python control flow + assert result.region is None or result.graph_breaks or result.region.rejection_reasons + if result.region is None: + assert result.graph_breaks + assert any("fx_trace_failed" in b.reason for b in result.graph_breaks) + + +def test_session_register_explicit_module(): + block = _PureBlock() + session = RegionCaptureSession(name_prefix="leaf") + session.register_module(block, scope="blocks.0") + for _ in range(2): + block(torch.randn(1, 4), torch.randn(1, 4)) + session.close() + regions = session.regions() + assert len(regions) == 1 + assert regions[0].calls == 2 + assert regions[0].parent_module == "blocks.0" + assert regions[0].shape_frequency diff --git a/tests/test_gpu_smoke.py b/tests/test_gpu_smoke.py new file mode 100644 index 00000000..9db60f2c --- /dev/null +++ b/tests/test_gpu_smoke.py @@ -0,0 +1,117 @@ +"""GPU smoke tests for the specification path. + +Deselected by default with ``-m "not gpu"``. On a CUDA machine:: + + uv run pytest -m gpu + +These run the real starter kernels through the real harness, so a failure here +means the registry refactor changed kernel behavior. +""" + +from __future__ import annotations + +import pytest + +from autokernel.specs import create_builtin_registry, load_spec +from autokernel.verification import check_backward, check_compile +from conftest import REPO_ROOT, requires_gpu + +pytestmark = [pytest.mark.gpu, requires_gpu] + + +def _load_kernel_fn(path): + """Import ``kernel_fn`` from a starter kernel file without touching sys.path.""" + import importlib.util + import uuid + + module_spec = importlib.util.spec_from_file_location( + f"_starter_{uuid.uuid4().hex}", path + ) + module = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(module) + return module.kernel_fn + + +@pytest.mark.parametrize("name", ["matmul", "layernorm", "rmsnorm"]) +def test_builtin_starter_kernels_pass_correctness(name): + bench = pytest.importorskip("bench") + spec = create_builtin_registry().get(name) + kernel_fn = _load_kernel_fn(spec.starter_kernel("triton")) + + results = bench.run_correctness(kernel_fn, spec, quick=True) + assert results["correctness"] == "PASS", results.get("details") + + +def test_external_spec_runs_through_the_same_harness(): + bench = pytest.importorskip("bench") + example = REPO_ROOT / "examples" / "custom_ops" / "add.py" + spec = load_spec(f"{example}:SPEC", registry=create_builtin_registry()) + kernel_fn = _load_kernel_fn(spec.starter_kernel("triton")) + + results = bench.run_correctness(kernel_fn, spec, quick=False) + assert results["correctness"] == "PASS", results.get("details") + + +@pytest.mark.parametrize( + "operation", + [ + "wan_gated_residual_norm", + "wan_modulated_layer_norm", + "wan_gated_residual", + ], +) +def test_wan_starter_kernel_passes_correctness(operation): + bench = pytest.importorskip("bench") + model = REPO_ROOT / "models" / f"{operation}.py" + spec = load_spec(f"{model}:SPEC") + kernel_fn = _load_kernel_fn(spec.starter_kernel("triton")) + + results = bench.run_correctness(kernel_fn, spec, quick=True) + assert results["correctness"] == "PASS", results.get("details") + + +def test_wan_starter_rejects_noncontiguous_and_cross_device_inputs(): + import torch + + model = REPO_ROOT / "models" / "wan_gated_residual_norm.py" + spec = load_spec(f"{model}:SPEC") + kernel_fn = _load_kernel_fn(spec.starter_kernel("triton")) + residual = torch.randn(1, 2, 8, device="cuda", dtype=torch.bfloat16) + x = torch.randn_like(residual) + gate = torch.randn(1, 16, device="cuda")[:, ::2] + weight = torch.randn(8, device="cuda") + bias = torch.randn(8, device="cuda") + + with pytest.raises(ValueError, match="gate, weight, and bias"): + kernel_fn(residual, x, gate, weight, bias) + with pytest.raises(ValueError, match="residual device"): + kernel_fn(residual, x, gate.contiguous(), weight, bias.cpu()) + + +def test_external_spec_performance_path(): + bench = pytest.importorskip("bench") + example = REPO_ROOT / "examples" / "custom_ops" / "add.py" + spec = load_spec(f"{example}:SPEC") + kernel_fn = _load_kernel_fn(spec.starter_kernel("triton")) + + gpu = bench.detect_gpu() + perf = bench.run_performance(kernel_fn, spec, gpu, sizes_filter="large") + primary = perf["primary"] + assert primary is not None + assert primary["kernel_latency_us"] > 0 + assert primary["bytes"] == 3 * 4096 * 4096 * 2 # float16 primary dtype + + +def test_structured_external_spec_forward_backward_and_compile(): + bench = pytest.importorskip("bench") + example = REPO_ROOT / "examples" / "custom_ops" / "affine.py" + spec = load_spec(f"{example}:SPEC") + kernel_fn = _load_kernel_fn(spec.starter_kernel("pytorch")) + + forward = bench.run_correctness(kernel_fn, spec, quick=True) + backward = check_backward(kernel_fn, spec, device="cuda") + compiled = check_compile(kernel_fn, spec, device="cuda") + + assert forward["correctness"] == "PASS", forward.get("details") + assert backward.status == "PASS", backward.reason + assert compiled.status == "PASS", compiled.reason diff --git a/tests/test_output_trees.py b/tests/test_output_trees.py new file mode 100644 index 00000000..3dc74dbb --- /dev/null +++ b/tests/test_output_trees.py @@ -0,0 +1,283 @@ +"""Structured output-tree comparison on CPU.""" + +from __future__ import annotations + +from collections import namedtuple + +import pytest + +from autokernel.specs import OutputSpec, Tolerance +from autokernel.verification import ( + OutputTreeError, + compare_deterministic, + compare_output_trees, + flatten_output_tree, + tree_has_nan_or_inf, +) + +torch = pytest.importorskip("torch") + +TOLS = { + "float16": Tolerance(atol=1e-2, rtol=1e-2), + "float32": Tolerance(atol=1e-5, rtol=1e-5), +} + +Point = namedtuple("Point", ["x", "y"]) + + +def _t(*values: float, dtype: torch.dtype = torch.float32) -> torch.Tensor: + return torch.tensor(values, dtype=dtype) + + +# --------------------------------------------------------------------------- +# Flattening and paths +# --------------------------------------------------------------------------- + +def test_flatten_single_tensor_has_root_path(): + leaves = flatten_output_tree(_t(1.0)) + assert [path for path, _ in leaves] == ["output"] + + +def test_flatten_tuple_and_list_paths(): + leaves = flatten_output_tree((_t(1.0), [_t(2.0), _t(3.0)])) + assert [path for path, _ in leaves] == ["output[0]", "output[1][0]", "output[1][1]"] + + +def test_flatten_dict_paths_are_sorted_not_insertion_ordered(): + first = flatten_output_tree({"z": _t(1.0), "a": _t(2.0)}) + second = flatten_output_tree({"a": _t(2.0), "z": _t(1.0)}) + assert [path for path, _ in first] == ['output["a"]', 'output["z"]'] + assert [path for path, _ in first] == [path for path, _ in second] + + +def test_flatten_namedtuple_uses_field_names(): + leaves = flatten_output_tree(Point(x=_t(1.0), y=_t(2.0))) + assert [path for path, _ in leaves] == ["output.x", "output.y"] + + +def test_flatten_nested_combination_matches_brief_example(): + tree = {"output": _t(1.0), "aux": (_t(2.0), 3)} + leaves = flatten_output_tree(tree) + assert [path for path, _ in leaves] == ['output["aux"][0]', 'output["aux"][1]', 'output["output"]'] + assert leaves[1][1] == 3 # metadata leaf kept as-is + + +def test_flatten_empty_containers_become_leaves(): + leaves = flatten_output_tree({"empty": {}, "items": []}) + assert [path for path, _ in leaves] == ['output["empty"]', 'output["items"]'] + + +# --------------------------------------------------------------------------- +# Comparison: matching trees +# --------------------------------------------------------------------------- + +def test_single_tensor_match(): + cmp = compare_output_trees(_t(1.0, 2.0), _t(1.0, 2.0), TOLS) + assert cmp.match + assert cmp.structure_match + assert cmp.worst_abs_error == 0.0 + assert cmp.leaf_records()[0]["path"] == "output" + + +def test_nested_tree_match_with_metadata(): + ref = {"output": _t(1.0), "aux": (_t(2.0), 3)} + cmp = compare_output_trees(ref, ref, TOLS) + assert cmp.match + assert len(cmp.leaves) == 3 + kinds = {leaf.path: leaf.kind for leaf in cmp.leaves} + assert kinds['output["aux"][1]'] == "metadata" + + +def test_close_values_within_dtype_tolerance(): + candidate = _t(1.0, dtype=torch.float16) + 1e-3 + expected = _t(1.0, dtype=torch.float16) + cmp = compare_output_trees(candidate, expected, TOLS) + assert cmp.match + + +# --------------------------------------------------------------------------- +# Comparison: failures +# --------------------------------------------------------------------------- + +def test_shape_mismatch_fails_with_path(): + cmp = compare_output_trees( + {"output": torch.zeros(2, 3)}, {"output": torch.zeros(3, 2)}, TOLS + ) + assert not cmp.match + assert "shape mismatch" in cmp.reason + + +def test_single_tensor_reason_matches_legacy_format(): + cmp = compare_output_trees(_t(1.0), _t(2.0), TOLS) + assert not cmp.match + assert cmp.reason.startswith("max_abs_error=") + assert "exceeds tol(atol=" in cmp.reason + + +def test_nested_failure_reason_is_prefixed_with_path(): + candidate = {"output": _t(1.0), "aux": (_t(2.0), 3)} + expected = {"output": _t(1.0), "aux": (_t(9.0), 3)} + cmp = compare_output_trees(candidate, expected, TOLS) + assert not cmp.match + assert cmp.reason.startswith('output["aux"][0]: ') + + +def test_missing_output_path_fails(): + cmp = compare_output_trees(_t(1.0), (_t(1.0), _t(2.0)), TOLS) + assert not cmp.match + assert not cmp.structure_match + assert "missing output path(s)" in cmp.reason + + +def test_unexpected_output_path_fails(): + cmp = compare_output_trees((_t(1.0), _t(2.0)), _t(1.0), TOLS) + assert not cmp.match + assert "unexpected output path(s)" in cmp.reason + + +def test_leaf_kind_mismatch_fails(): + cmp = compare_output_trees({"a": _t(1.0)}, {"a": 1.0}, TOLS) + assert not cmp.match + assert "leaf kind mismatch" in cmp.reason + + +def test_dtype_mismatch_fails(): + cmp = compare_output_trees(_t(1.0, dtype=torch.float16), _t(1.0), TOLS) + assert not cmp.match + assert "dtype mismatch" in cmp.reason + + +def test_per_leaf_tolerance_selection_by_leaf_dtype(): + """A float16 leaf gets the float16 tolerance, a float32 leaf float32's.""" + candidate = ( + _t(1.0, dtype=torch.float16) + 5e-3, # within float16 tol (1e-2) + _t(1.0, dtype=torch.float32) + 5e-3, # outside float32 tol (1e-5) + ) + expected = (_t(1.0, dtype=torch.float16), _t(1.0, dtype=torch.float32)) + cmp = compare_output_trees(candidate, expected, TOLS) + assert not cmp.match + by_path = {leaf.path: leaf for leaf in cmp.leaves} + assert by_path["output[0]"].match + assert not by_path["output[1]"].match + + +def test_integer_tensors_compare_exactly(): + cmp = compare_output_trees(_t(1, 2, dtype=torch.int64), _t(1, 2, dtype=torch.int64), TOLS) + assert cmp.match + cmp = compare_output_trees(_t(1, 2, dtype=torch.int64), _t(1, 3, dtype=torch.int64), TOLS) + assert not cmp.match + assert "bitwise" in cmp.leaves[0].reason + + +# --------------------------------------------------------------------------- +# NaN / infinity +# --------------------------------------------------------------------------- + +def test_nan_detected_per_path(): + candidate = {"ok": _t(1.0), "bad": _t(float("nan"))} + expected = {"ok": _t(1.0), "bad": _t(0.0)} + cmp = compare_output_trees(candidate, expected, TOLS) + assert not cmp.match + by_path = {leaf.path: leaf for leaf in cmp.leaves} + assert by_path['output["bad"]'].has_nan + assert not by_path['output["ok"]'].has_nan + assert "NaN" in by_path['output["bad"]'].reason + + +def test_infinity_detected_per_path(): + cmp = compare_output_trees(_t(float("inf")), _t(1.0), TOLS) + assert not cmp.match + assert cmp.leaves[0].has_inf + assert "infinity" in cmp.leaves[0].reason + + +def test_tree_has_nan_or_inf_traverses_containers(): + assert not tree_has_nan_or_inf({"a": (_t(1.0), 3)}) + assert tree_has_nan_or_inf({"a": (_t(float("nan")), 3)}) + assert tree_has_nan_or_inf((_t(1.0), [_t(float("inf"))])) + + +# --------------------------------------------------------------------------- +# Metadata comparison policy +# --------------------------------------------------------------------------- + +def test_metadata_mismatch_fails_by_default(): + cmp = compare_output_trees((_t(1.0), 3), (_t(1.0), 4), TOLS) + assert not cmp.match + assert "metadata mismatch" in cmp.leaves[1].reason + + +def test_metadata_mismatch_ignored_when_disabled(): + policy = OutputSpec(compare_non_tensors=False) + cmp = compare_output_trees((_t(1.0), 3), (_t(1.0), 4), TOLS, output_spec=policy) + assert cmp.match + assert cmp.leaves[1].reason == "not compared" + + +def test_included_paths_restricts_comparison(): + policy = OutputSpec(included_paths=('output["output"]',)) + candidate = {"output": _t(1.0), "aux": (_t(9.0), 4)} + expected = {"output": _t(1.0), "aux": (_t(2.0), 3)} + cmp = compare_output_trees(candidate, expected, TOLS, output_spec=policy) + assert cmp.match + assert [leaf.path for leaf in cmp.leaves] == ['output["output"]'] + + +def test_included_paths_must_exist(): + policy = OutputSpec(included_paths=("output[9]",)) + with pytest.raises(OutputTreeError, match="not present"): + compare_output_trees((_t(1.0),), (_t(1.0),), TOLS, output_spec=policy) + + +# --------------------------------------------------------------------------- +# Determinism +# --------------------------------------------------------------------------- + +def test_determinism_bitwise_identical_trees(): + tree = {"output": _t(1.5, -2.5), "aux": (_t(3.0), 3)} + cmp = compare_deterministic(tree, tree) + assert cmp.match + + +def test_determinism_detects_any_leaf_difference(): + first = {"output": _t(1.0), "aux": (_t(2.0), 3)} + second = {"output": _t(1.0), "aux": (_t(2.0 + 1e-3), 3)} + cmp = compare_deterministic(first, second) + assert not cmp.match + failure = cmp.first_failure() + assert failure is not None + assert failure.path == 'output["aux"][0]' + assert failure.max_abs_error is not None + + +def test_determinism_reports_distinct_mean_and_max_errors(): + cmp = compare_deterministic(_t(0.0, 0.0), _t(0.0, 2.0)) + failure = cmp.first_failure() + assert failure is not None + assert failure.max_abs_error == pytest.approx(2.0) + assert failure.mean_abs_error == pytest.approx(1.0) + + +def test_determinism_detects_metadata_change(): + cmp = compare_deterministic((_t(1.0), 3), (_t(1.0), 4)) + assert not cmp.match + assert "metadata changed" in cmp.leaves[1].reason + + +def test_determinism_detects_structure_change(): + cmp = compare_deterministic(_t(1.0), (_t(1.0),)) + assert not cmp.match + assert not cmp.structure_match + + +# --------------------------------------------------------------------------- +# Relax factor (numerical-stability stage) +# --------------------------------------------------------------------------- + +def test_relax_multiplies_tolerances(): + candidate = _t(1.0) + 5e-5 + expected = _t(1.0) + strict = compare_output_trees(candidate, expected, TOLS) + assert not strict.match + relaxed = compare_output_trees(candidate, expected, TOLS, relax=10.0) + assert relaxed.match diff --git a/tests/test_ranking_fx.py b/tests/test_ranking_fx.py new file mode 100644 index 00000000..a9f27a73 --- /dev/null +++ b/tests/test_ranking_fx.py @@ -0,0 +1,413 @@ +"""Ranking, profiler parse, and CPU FX capture — real shipped entry points.""" + +from __future__ import annotations + +import json + +import pytest +import torch +import torch.nn as nn + +from autokernel.discovery.fingerprint import graph_fingerprint +from autokernel.discovery import ( + DEFAULT_IMPACT_FLOOR, + GraphRegion, + OperatorHotspot, + TensorMeta, + capture_callable_region, + capture_module_region, + load_discovery_report, + optimistic_e2e_improvement, + parse_key_averages_rows, + profiler_export_to_report, + rank_operators, + rank_regions, +) + + +class _GatedResidual(nn.Module): + """Tiny pure-tensor stand-in for residual * gate + residual patterns.""" + + def forward(self, residual: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + return residual + residual * gate + + +def test_optimistic_e2e_and_impact_floor_on_wan_like_share(): + # 0.5% of e2e with 90% reducible => 0.45% optimistic < 0.5% floor + share = 0.005 + improvement = optimistic_e2e_improvement(share, reducible_fraction=0.9) + assert improvement == pytest.approx(0.0045) + assert improvement < DEFAULT_IMPACT_FLOOR + + +def test_rank_regions_marks_low_value_and_high_value(): + low = GraphRegion.build( + name="wan.elementwise", + operations=["aten::mul", "aten::add", "aten::layer_norm"], + inputs=[ + TensorMeta("r", (1, 128, 64), (8192, 64, 1), "bfloat16", "cpu"), + TensorMeta("g", (1, 1, 64), (64, 64, 1), "float32", "cpu"), + ], + cuda_time_us=50.0, + self_cuda_time_us=50.0, + calls=40, + ) + high = GraphRegion.build( + name="hot.epilogue", + operations=["aten::mul", "aten::add", "aten::silu"], + inputs=[ + TensorMeta("x", (1, 128, 64), (8192, 64, 1), "float16", "cpu"), + ], + cuda_time_us=2500.0, + self_cuda_time_us=2500.0, + calls=40, + ) + ranked = rank_regions( + [low, high], + total_cuda_time_us=10_000.0, + impact_floor=0.005, + reducible_fraction=0.9, + ) + assert ranked[0].region.name == "hot.epilogue" + assert ranked[0].search_worthy is True + assert ranked[0].estimated_max_e2e_improvement == 0.225 # 25% * 0.9 + assert ranked[1].region.name == "wan.elementwise" + assert ranked[1].search_worthy is False + assert any("below_impact_floor" in r for r in ranked[1].rejection_reasons) + + +def test_parse_key_averages_rows_real_aliases(): + rows = [ + { + "name": "aten::mm", + "cuda_time_total": 6000.0, + "self_cuda_time_total": 5900.0, + "count": 100, + }, + { + "key": "aten::mul", + "cuda_time_ms": 0.4, + "self_cuda_time_ms": 0.4, + "calls": 40, + }, + ] + ops = parse_key_averages_rows(rows) + assert ops[0].op_key == "aten::mm" + assert ops[0].cuda_time_us == 6000.0 + assert ops[0].calls == 100 + assert ops[1].op_key == "aten::mul" + assert ops[1].cuda_time_us == 400.0 + + +def test_capture_module_region_cpu_fx(): + module = _GatedResidual() + residual = torch.randn(2, 8, 16) + gate = torch.randn(2, 1, 16) + result = capture_module_region( + module, + (residual, gate), + name="test.gated_residual", + parent_module="blocks.0", + ) + assert result.region is not None + assert result.region.name == "test.gated_residual" + assert len(result.region.operations) >= 1 + assert result.region.fingerprint + # Fingerprint stable across re-capture with same parent_module + again = capture_module_region( + module, + (residual, gate), + name="test.gated_residual", + parent_module="blocks.0", + ) + assert again.region is not None + assert again.region.operations == result.region.operations + assert again.region.fingerprint == result.region.fingerprint + # Ranking can consume the region with injected timing + timed = GraphRegion.build( + name=result.region.name, + operations=result.region.operations, + inputs=result.region.inputs, + outputs=result.region.outputs, + parent_module=result.region.parent_module, + pattern_family=result.region.pattern_family, + rejection_reasons=result.region.rejection_reasons, + cuda_time_us=2000.0, + self_cuda_time_us=2000.0, + calls=32, + ) + ranked = rank_regions([timed], total_cuda_time_us=10_000.0) + assert ranked[0].search_worthy is True + + +def test_capture_callable_region(): + def pure(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x * y + x + + x = torch.randn(4, 8) + y = torch.randn(4, 8) + result = capture_callable_region(pure, (x, y), name="test.mul_add") + assert result.region is not None + joined = " ".join(result.region.operations).lower() + assert "mul" in joined or "add" in joined + + +def test_discovery_cli_rank_entry_point(tmp_path): + from discovery import main as discovery_main + + region = GraphRegion.build( + name="cli.region", + operations=["aten::mul", "aten::add"], + inputs=[TensorMeta("x", (1, 8, 8), (64, 8, 1), "float16", "cpu")], + cuda_time_us=3000.0, + calls=10, + ) + payload = { + "schema_version": 1, + "producer": {"name": "test", "version": "0"}, + "workload": {"workload_id": "unit", "model_id": "m"}, + "environment": { + "hardware_profile_id": "h", + "software_profile_id": "s", + }, + "total_cuda_time_us": 10000.0, + "operators": [ + { + "name": "aten::mm", + "op_key": "aten::mm", + "calls": 1, + "cuda_time_us": 5000.0, + "self_cuda_time_us": 5000.0, + } + ], + "regions": [region.as_dict()], + "graph_breaks": [], + "unsupported": [], + } + path = tmp_path / "d.json" + path.write_text(json.dumps(payload), encoding="utf-8") + assert discovery_main(["validate", str(path)]) == 0 + assert discovery_main(["rank", str(path)]) == 0 + + +def test_profiler_export_ingestion_and_cli(tmp_path): + from discovery import main as discovery_main + + export = { + "schema_version": 1, + "producer": {"name": "fastvideo", "version": "1"}, + "workload": {"workload_id": "ltx-unit", "model_id": "ltx"}, + "environment": {"torch": "2.x", "gpu_name": "unit"}, + "total_cuda_time_us": 100.0, + "rows": [ + { + "name": "ProfilerStep*", + "calls": 1, + "cuda_time_us": 100.0, + "self_cuda_time_us": 10.0, + "cpu_time_us": 20.0, + }, + { + "name": "aten::mul", + "calls": 4, + "cuda_time_us": 90.0, + "self_cuda_time_us": 90.0, + "cpu_time_us": 5.0, + "input_shapes": [[1, 8], [1, 8]], + }, + ], + } + report = profiler_export_to_report(export) + assert report.total_cuda_time_us == 100.0 + assert [op.op_key for op in report.operators] == [ + "ProfilerStep", + "aten::mul", + ] + + source = tmp_path / "profile.json" + output = tmp_path / "discovery.json" + source.write_text(json.dumps(export), encoding="utf-8") + assert ( + discovery_main( + ["ingest-profiler", str(source), "--output", str(output)] + ) + == 0 + ) + loaded = load_discovery_report(output) + assert len(loaded.operators) == 2 + + +def test_operator_ranking_uses_self_cuda_time(): + nested_scope = OperatorHotspot( + name="attention_scope", + op_key="attention_scope", + calls=1, + cuda_time_us=100.0, + self_cuda_time_us=1.0, + cpu_time_us=0.0, + input_shapes=(), + parent_module=None, + source="torch_profiler", + ) + kernel = OperatorHotspot( + name="attention_kernel", + op_key="attention_kernel", + calls=1, + cuda_time_us=80.0, + self_cuda_time_us=80.0, + cpu_time_us=0.0, + input_shapes=(), + parent_module=None, + source="torch_profiler", + ) + + ranked = rank_operators( + [nested_scope, kernel], + total_cuda_time_us=100.0, + ) + + assert ranked[0] == (kernel, pytest.approx(0.8)) + assert ranked[1] == (nested_scope, pytest.approx(0.01)) + + +def test_profiler_export_rejects_nested_secret_metadata(): + export = { + "schema_version": 1, + "producer": {"name": "fastvideo", "version": "1"}, + "workload": {"workload_id": "unit", "model_id": "model"}, + "environment": {"runtime": {"token": "do-not-store"}}, + "total_cuda_time_us": 1.0, + "rows": [ + { + "name": "aten::add", + "calls": 1, + "cuda_time_us": 1.0, + "self_cuda_time_us": 1.0, + } + ], + } + with pytest.raises(ValueError, match="forbidden"): + profiler_export_to_report(export) + + +def _capture_export(**overrides): + """A FastVideo-shaped export that also carries the optional FX capture block.""" + inputs = [ + { + "name": "input_0", + "shape": [2, 4], + "stride": [4, 1], + "dtype": "float32", + "device_type": "cpu", + "requires_grad": False, + } + ] + operations = ["aten::mul", "aten::silu", "aten::add"] + fingerprint = graph_fingerprint( + operations=operations, + input_signatures=[ + {k: v for k, v in inputs[0].items() if k != "name"} + ], + output_signatures=[], + safe_constants={}, + parent_module="transformer.blocks", + ) + export = { + "schema_version": 1, + "producer": {"name": "fastvideo", "version": "1"}, + "workload": {"workload_id": "unit", "model_id": "any-dit"}, + "environment": {"torch": "2.x", "gpu_name": "unit"}, + "total_cuda_time_us": 90.0, + "rows": [ + { + "name": "aten::mul", + "calls": 4, + "cuda_time_us": 90.0, + "self_cuda_time_us": 90.0, + "cpu_time_us": 5.0, + } + ], + "capture": { + "capture_schema_version": 1, + "tracer": "symbolic", + "scopes": ["transformer.blocks"], + "errors": [], + }, + "regions": [ + { + "name": "transformer.blocks.7dff7ade", + "fingerprint": fingerprint, + "operations": operations, + "dependencies": ["0->1", "1->2"], + "inputs": inputs, + "outputs": [], + "cuda_time_us": 0.0, + "self_cuda_time_us": 0.0, + "calls": 120, + "rejection_reasons": [], + "shape_frequency": {"input_0:2x4:float32": 120}, + "parent_module": "transformer.blocks", + } + ], + "graph_breaks": [ + {"scope": "transformer.blocks", "reason": "empty_graph", "count": 2} + ], + "unsupported": [ + { + "op_name": "module::attn", + "reason": "module::attn: nested module not expanded", + "count": 1, + "scope": "transformer.blocks", + } + ], + } + export.update(overrides) + return export + + +def test_profiler_export_ingests_optional_fx_capture_block(): + report = profiler_export_to_report(_capture_export()) + + assert len(report.operators) == 1 + assert len(report.regions) == 1 + region = report.regions[0] + assert region.parent_module == "transformer.blocks" + assert region.operations == ("aten::mul", "aten::silu", "aten::add") + assert region.calls == 120 + assert dict(region.shape_frequency) == {"input_0:2x4:float32": 120} + assert [item.reason for item in report.graph_breaks] == ["empty_graph"] + assert report.unsupported[0].op_name == "module::attn" + + +def test_profiler_export_without_capture_still_loads(): + export = _capture_export() + for key in ("capture", "regions", "graph_breaks", "unsupported"): + export.pop(key) + + report = profiler_export_to_report(export) + assert report.regions == () + assert report.graph_breaks == () + + +def test_profiler_export_rejects_capture_payload_without_version_block(): + export = _capture_export() + export.pop("capture") + + with pytest.raises(ValueError, match="capture"): + profiler_export_to_report(export) + + +def test_profiler_export_rejects_unknown_capture_schema_version(): + export = _capture_export() + export["capture"] = {"capture_schema_version": 2} + + with pytest.raises(ValueError, match="capture_schema_version"): + profiler_export_to_report(export) + + +def test_profiler_export_rejects_tampered_region_fingerprint(): + export = _capture_export() + export["regions"][0]["fingerprint"] = "0" * 32 + + with pytest.raises(ValueError, match="fingerprint"): + profiler_export_to_report(export) diff --git a/tests/test_result_writer.py b/tests/test_result_writer.py new file mode 100644 index 00000000..066cd5a3 --- /dev/null +++ b/tests/test_result_writer.py @@ -0,0 +1,70 @@ +"""Tests for versioned atomic result artifacts.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import pytest + +import autokernel.verification.results as results + + +def test_result_envelope_is_versioned(): + payload = results.result_envelope("affine", forward={"status": "PASS"}) + assert payload["schema_version"] == results.RESULT_SCHEMA_VERSION + assert payload["operation"] == "affine" + assert payload["created_at"].endswith("+00:00") + + +def test_write_result_atomic_replaces_destination(tmp_path: Path): + destination = tmp_path / "nested" / "bench_result.json" + destination.parent.mkdir() + destination.write_text('{"old": true}\n') + + written = results.write_result_atomic(destination, {"new": True}) + + assert written == destination + assert json.loads(destination.read_text()) == {"new": True} + assert list(destination.parent.glob("*.tmp")) == [] + assert list(destination.parent.glob(".*.tmp")) == [] + + +def test_write_result_atomic_serializes_non_finite_values_safely(tmp_path: Path): + destination = tmp_path / "bench_result.json" + results.write_result_atomic(destination, {"bad": float("nan")}) + assert json.loads(destination.read_text()) == {"bad": "nan"} + assert list(tmp_path.glob(".*.tmp")) == [] + + +def test_write_result_atomic_sanitizes_values_from_as_dict(tmp_path: Path): + @dataclass + class Record: + error: float + + def as_dict(self): + return {"nested": {"error": self.error}} + + destination = tmp_path / "bench_result.json" + results.write_result_atomic(destination, {"record": Record(float("inf"))}) + assert json.loads(destination.read_text()) == { + "record": {"nested": {"error": "inf"}} + } + + +def test_write_result_atomic_preserves_old_file_on_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + destination = tmp_path / "bench_result.json" + destination.write_text('{"old": true}\n') + + def fail_replace(source, target): + raise OSError("interrupted") + + monkeypatch.setattr(results.os, "replace", fail_replace) + with pytest.raises(OSError, match="interrupted"): + results.write_result_atomic(destination, {"new": True}) + + assert json.loads(destination.read_text()) == {"old": True} + assert list(tmp_path.glob(".*.tmp")) == [] diff --git a/tests/test_shape_corpus.py b/tests/test_shape_corpus.py new file mode 100644 index 00000000..129a7c77 --- /dev/null +++ b/tests/test_shape_corpus.py @@ -0,0 +1,278 @@ +"""Shape-corpus loading, validation and weighted aggregation.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from autokernel.verification import ( + CorpusError, + load_shape_corpus, + validate_corpus_against_spec, + weighted_aggregate, +) +from conftest import make_spec + + +def _write(tmp_path: Path, payload: object, name: str = "corpus.json") -> Path: + path = tmp_path / name + path.write_text(json.dumps(payload)) + return path + + +def _valid_payload(**overrides): + payload = { + "schema_version": 1, + "operation": "unit_op", + "cases": [ + {"name": "prod-a", "size": {"rows": 100, "cols": 200}, "dtype": "float16", + "weight": 3, "tags": ["production"]}, + {"name": "prod-b", "size": {"rows": 7, "cols": 9}}, + ], + } + payload.update(overrides) + return payload + + +def _spec(**overrides): + return make_spec(**overrides) + + +# --------------------------------------------------------------------------- +# Happy path +# --------------------------------------------------------------------------- + +def test_load_valid_corpus(tmp_path): + corpus = load_shape_corpus(_write(tmp_path, _valid_payload())) + assert corpus.operation == "unit_op" + assert corpus.schema_version == 1 + assert len(corpus.cases) == 2 + first, second = corpus.cases + assert first.name == "prod-a" + assert first.size == {"rows": 100, "cols": 200} + assert first.dtype == "float16" + assert first.weight == 3 + assert first.tags == ("production",) + # defaults + assert second.dtype is None + assert second.weight == 1 + assert second.tags == () + + +def test_valid_corpus_passes_spec_validation(tmp_path): + corpus = load_shape_corpus(_write(tmp_path, _valid_payload())) + validate_corpus_against_spec(corpus, _spec()) # must not raise + + +def test_example_corpus_matches_example_spec(repo_root): + from autokernel.specs import load_spec + + corpus = load_shape_corpus(repo_root / "examples" / "custom_ops" / "affine_corpus.json") + spec = load_spec(str(repo_root / "examples" / "custom_ops" / "affine.py") + ":SPEC") + validate_corpus_against_spec(corpus, spec) + + +# --------------------------------------------------------------------------- +# File and schema errors +# --------------------------------------------------------------------------- + +def test_missing_file_fails(tmp_path): + with pytest.raises(CorpusError, match="file not found"): + load_shape_corpus(tmp_path / "nope.json") + + +def test_invalid_json_fails(tmp_path): + path = tmp_path / "bad.json" + path.write_text("{not json") + with pytest.raises(CorpusError, match="invalid JSON"): + load_shape_corpus(path) + + +def test_non_utf8_corpus_fails_with_corpus_error(tmp_path): + path = tmp_path / "bad-encoding.json" + path.write_bytes(b"\xff\xfe") + with pytest.raises(CorpusError, match="cannot read file"): + load_shape_corpus(path) + + +def test_top_level_must_be_object(tmp_path): + with pytest.raises(CorpusError, match="top level must be an object"): + load_shape_corpus(_write(tmp_path, [1, 2, 3])) + + +def test_unknown_top_level_field_fails(tmp_path): + with pytest.raises(CorpusError, match="unknown top-level field"): + load_shape_corpus(_write(tmp_path, _valid_payload(tensors={"a": [1, 2]}))) + + +@pytest.mark.parametrize("version", [None, "1", 1.0, True]) +def test_schema_version_must_be_integer(tmp_path, version): + payload = _valid_payload() + if version is None: + del payload["schema_version"] + else: + payload["schema_version"] = version + with pytest.raises(CorpusError, match="schema_version.*must be an integer"): + load_shape_corpus(_write(tmp_path, payload)) + + +def test_unsupported_schema_version_fails(tmp_path): + with pytest.raises(CorpusError, match="unsupported schema_version 2"): + load_shape_corpus(_write(tmp_path, _valid_payload(schema_version=2))) + + +def test_operation_must_be_non_empty(tmp_path): + with pytest.raises(CorpusError, match="operation"): + load_shape_corpus(_write(tmp_path, _valid_payload(operation=""))) + + +def test_cases_must_be_non_empty(tmp_path): + with pytest.raises(CorpusError, match="cases.*non-empty"): + load_shape_corpus(_write(tmp_path, _valid_payload(cases=[]))) + + +# --------------------------------------------------------------------------- +# Case-level errors +# --------------------------------------------------------------------------- + +def _one_case(tmp_path: Path, case: dict) -> Path: + return _write(tmp_path, _valid_payload(cases=[case])) + + +def test_case_must_be_object(tmp_path): + with pytest.raises(CorpusError, match="case #0 must be an object"): + load_shape_corpus(_write(tmp_path, _valid_payload(cases=[42]))) + + +def test_unknown_case_field_fails(tmp_path): + case = {"name": "a", "size": {"rows": 1, "cols": 1}, "activations": [0.1]} + with pytest.raises(CorpusError, match="unknown field"): + load_shape_corpus(_one_case(tmp_path, case)) + + +@pytest.mark.parametrize("name", ["", 7, None]) +def test_bad_case_name_fails(tmp_path, name): + case = {"name": name, "size": {"rows": 1, "cols": 1}} + with pytest.raises(CorpusError, match="name"): + load_shape_corpus(_one_case(tmp_path, case)) + + +def test_duplicate_case_names_fail(tmp_path): + payload = _valid_payload() + payload["cases"][1] = dict(payload["cases"][0]) + with pytest.raises(CorpusError, match="duplicate case name 'prod-a'"): + load_shape_corpus(_write(tmp_path, payload)) + + +def test_size_must_be_non_empty_mapping(tmp_path): + with pytest.raises(CorpusError, match="size"): + load_shape_corpus(_one_case(tmp_path, {"name": "a", "size": {}})) + + +@pytest.mark.parametrize("value", [0, -3, 1.5, "64", True]) +def test_size_values_must_be_positive_ints(tmp_path, value): + case = {"name": "a", "size": {"rows": value, "cols": 4}} + with pytest.raises(CorpusError, match="positive integer"): + load_shape_corpus(_one_case(tmp_path, case)) + + +def test_unknown_case_dtype_fails(tmp_path): + case = {"name": "a", "size": {"rows": 1, "cols": 1}, "dtype": "float64"} + with pytest.raises(CorpusError, match="unknown dtype"): + load_shape_corpus(_one_case(tmp_path, case)) + + +@pytest.mark.parametrize("weight", [0, -2, 2.5, "3", False]) +def test_weight_must_be_positive_int(tmp_path, weight): + case = {"name": "a", "size": {"rows": 1, "cols": 1}, "weight": weight} + with pytest.raises(CorpusError, match="weight.*positive integer"): + load_shape_corpus(_one_case(tmp_path, case)) + + +def test_tags_must_be_strings(tmp_path): + case = {"name": "a", "size": {"rows": 1, "cols": 1}, "tags": ["ok", 3]} + with pytest.raises(CorpusError, match="tags"): + load_shape_corpus(_one_case(tmp_path, case)) + + +def test_duplicate_tags_fail(tmp_path): + case = {"name": "a", "size": {"rows": 1, "cols": 1}, "tags": ["x", "x"]} + with pytest.raises(CorpusError, match="duplicate tag"): + load_shape_corpus(_one_case(tmp_path, case)) + + +# --------------------------------------------------------------------------- +# Spec compatibility +# --------------------------------------------------------------------------- + +def test_operation_mismatch_fails(tmp_path): + corpus = load_shape_corpus(_write(tmp_path, _valid_payload(operation="other_op"))) + with pytest.raises(CorpusError, match="does not match the selected spec"): + validate_corpus_against_spec(corpus, _spec()) + + +def test_size_keys_must_match_shape_keys(tmp_path): + corpus = load_shape_corpus( + _one_case(tmp_path, {"name": "a", "size": {"rows": 1, "wrong": 2}}) + ) + with pytest.raises(CorpusError, match="do not match shape_keys"): + validate_corpus_against_spec(corpus, _spec()) + + +def test_case_dtype_must_be_declared_by_spec(tmp_path): + case = {"name": "a", "size": {"rows": 1, "cols": 1}, "dtype": "bfloat16"} + corpus = load_shape_corpus(_one_case(tmp_path, case)) + with pytest.raises(CorpusError, match="not declared in dtypes"): + validate_corpus_against_spec(corpus, _spec()) + + +def test_resolved_duplicate_cases_fail(tmp_path): + """Two cases resolving to the same (size, dtype) config are rejected.""" + payload = _valid_payload( + cases=[ + {"name": "a", "size": {"rows": 4, "cols": 4}}, # dtype None -> primary float16 + {"name": "b", "size": {"rows": 4, "cols": 4}, "dtype": "float16"}, + ] + ) + corpus = load_shape_corpus(_write(tmp_path, payload)) + with pytest.raises(CorpusError, match="resolve to the same"): + validate_corpus_against_spec(corpus, _spec()) + + +# --------------------------------------------------------------------------- +# Weighted aggregation +# --------------------------------------------------------------------------- + +def test_weighted_aggregate_math(): + entries = [ + {"dtype": "torch.float16", "weight": 3, "kernel_ms": 1.0, "ref_ms": 3.0}, + {"dtype": "torch.float16", "weight": 1, "kernel_ms": 3.0, "ref_ms": 9.0}, + ] + agg = weighted_aggregate(entries) + assert set(agg) == {"torch.float16"} + group = agg["torch.float16"] + assert group["cases"] == 2 + assert group["weight"] == 4 + assert group["kernel_ms"] == pytest.approx((3 * 1.0 + 1 * 3.0) / 4) + assert group["ref_ms"] == pytest.approx((3 * 3.0 + 1 * 9.0) / 4) + assert group["speedup"] == pytest.approx(group["ref_ms"] / group["kernel_ms"]) + + +def test_weighted_aggregate_never_mixes_dtypes(): + entries = [ + {"dtype": "torch.float16", "weight": 1, "kernel_ms": 1.0, "ref_ms": 2.0}, + {"dtype": "torch.float32", "weight": 1, "kernel_ms": 5.0, "ref_ms": 10.0}, + ] + agg = weighted_aggregate(entries) + assert set(agg) == {"torch.float16", "torch.float32"} + assert agg["torch.float16"]["kernel_ms"] == pytest.approx(1.0) + assert agg["torch.float32"]["kernel_ms"] == pytest.approx(5.0) + + +def test_weighted_aggregate_zero_kernel_latency_guards_division(): + agg = weighted_aggregate( + [{"dtype": "torch.float16", "weight": 1, "kernel_ms": 0.0, "ref_ms": 1.0}] + ) + assert agg["torch.float16"]["speedup"] == 0.0 diff --git a/tests/test_spec_loader.py b/tests/test_spec_loader.py new file mode 100644 index 00000000..fa1ba174 --- /dev/null +++ b/tests/test_spec_loader.py @@ -0,0 +1,300 @@ +"""External specification loading.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from autokernel.specs import ( + KernelRegistry, + KernelSpec, + SpecCollisionError, + SpecLoadError, + SpecValidationError, + create_builtin_registry, + load_spec, + parse_locator, + resolve_spec, +) +from conftest import FIXTURES_DIR, make_spec + +FIXTURE_FILE = FIXTURES_DIR / "custom_add.py" + + +# --------------------------------------------------------------------------- +# Locator parsing +# --------------------------------------------------------------------------- + +def test_parse_locator_splits_module_and_attribute(): + assert parse_locator("package.module:SPEC") == ("package.module", "SPEC") + + +def test_parse_locator_splits_at_last_colon(): + assert parse_locator(r"C:\ops\spec.py:SPEC") == (r"C:\ops\spec.py", "SPEC") + + +@pytest.mark.parametrize("bad", ["", " ", "no_attribute", "module:", ":SPEC", None, 5]) +def test_parse_locator_rejects_malformed_input(bad): + with pytest.raises(SpecLoadError): + parse_locator(bad) + + +# --------------------------------------------------------------------------- +# File locators +# --------------------------------------------------------------------------- + +def test_load_absolute_file_locator(): + spec = load_spec(f"{FIXTURE_FILE}:SPEC") + assert isinstance(spec, KernelSpec) + assert spec.name == "fixture_add" + + +def test_load_relative_file_locator(in_repo_root): + spec = load_spec("tests/fixtures/custom_add.py:SPEC") + assert spec.name == "fixture_add" + + +def test_load_callable_factory_from_file(): + spec = load_spec(f"{FIXTURE_FILE}:SPEC_FACTORY") + assert spec.name == "fixture_add" + + +def test_missing_file_reports_locator(): + locator = "/definitely/not/here/spec.py:SPEC" + with pytest.raises(SpecLoadError) as exc: + load_spec(locator) + message = str(exc.value) + assert locator in message + assert "file not found" in message + + +def test_file_that_fails_to_import_reports_locator(tmp_path: Path): + bad = tmp_path / "explodes.py" + bad.write_text("raise RuntimeError('boom')\n") + with pytest.raises(SpecLoadError) as exc: + load_spec(f"{bad}:SPEC") + assert "boom" in str(exc.value) + assert str(bad) in str(exc.value) + + +def test_file_loading_does_not_mutate_sys_path(tmp_path: Path): + before = list(sys.path) + load_spec(f"{FIXTURE_FILE}:SPEC") + assert sys.path == before + + +def test_file_loading_does_not_leak_temporary_modules(): + before = set(sys.modules) + load_spec(f"{FIXTURE_FILE}:SPEC") + load_spec(f"{FIXTURE_FILE}:SPEC") + added = [name for name in set(sys.modules) - before if "external_spec" in name] + assert added == [] + + +# --------------------------------------------------------------------------- +# Module locators +# --------------------------------------------------------------------------- + +def test_load_module_locator(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + module_dir = tmp_path / "pkgdir" + module_dir.mkdir() + (module_dir / "external_op.py").write_text( + "from autokernel.specs import DT_BYTES, KernelSpec, Tolerance, size\n" + "\n" + "def ref(x):\n" + " return x\n" + "\n" + "def gen(size_map, dtype, device, seed=42):\n" + " return {'x': None}\n" + "\n" + "SPEC = KernelSpec(\n" + " name='module_op',\n" + " reference_fn=ref,\n" + " input_generator=gen,\n" + " sizes={'small': {'n': 1}, 'medium': {'n': 2}, 'large': {'n': 3}},\n" + " dtypes=('float32',),\n" + " tolerances={'float32': Tolerance(atol=1e-5, rtol=1e-5)},\n" + " flops_fn=size('n'),\n" + " bytes_fn=size('n') * DT_BYTES,\n" + " shape_keys=('n',),\n" + ")\n" + ) + monkeypatch.syspath_prepend(str(module_dir)) + spec = load_spec("external_op:SPEC") + assert spec.name == "module_op" + + +def test_load_module_locator_callable_factory(): + # A built-in factory doubles as a module-locator + zero-argument factory case. + spec = load_spec("autokernel.specs.builtins:_spec_matmul") + assert spec.name == "matmul" + + +def test_missing_module_reports_locator(): + with pytest.raises(SpecLoadError) as exc: + load_spec("definitely_not_a_module:SPEC") + assert "definitely_not_a_module:SPEC" in str(exc.value) + assert "not found" in str(exc.value) + + +# --------------------------------------------------------------------------- +# Attribute and type errors +# --------------------------------------------------------------------------- + +def test_missing_attribute_lists_available_specs(): + with pytest.raises(SpecLoadError) as exc: + load_spec(f"{FIXTURE_FILE}:NOPE") + message = str(exc.value) + assert "has no attribute 'NOPE'" in message + assert "SPEC" in message + + +def test_attribute_of_wrong_type_is_rejected(): + with pytest.raises(SpecLoadError, match="expected a KernelSpec"): + load_spec(f"{FIXTURE_FILE}:NOT_A_SPEC") + + +def test_factory_returning_wrong_type_is_rejected(): + with pytest.raises(SpecLoadError, match="returned int"): + load_spec(f"{FIXTURE_FILE}:BAD_FACTORY") + + +def test_factory_that_raises_is_reported(): + with pytest.raises(SpecLoadError, match="factory exploded"): + load_spec(f"{FIXTURE_FILE}:RAISING_FACTORY") + + +def test_invalid_spec_in_file_is_rejected(tmp_path: Path): + bad = tmp_path / "bad_spec.py" + bad.write_text( + "from autokernel.specs import DT_BYTES, KernelSpec, Tolerance, size\n" + "SPEC = KernelSpec(\n" + " name='bad op',\n" + " reference_fn=lambda x: x,\n" + " input_generator=lambda *a, **k: {},\n" + " sizes={'small': {'n': 1}, 'medium': {'n': 2}, 'large': {'n': 3}},\n" + " dtypes=('float32',),\n" + " tolerances={'float32': Tolerance(atol=1e-5, rtol=1e-5)},\n" + " flops_fn=size('n'),\n" + " bytes_fn=size('n') * DT_BYTES,\n" + " shape_keys=('n',),\n" + ")\n" + ) + with pytest.raises(SpecLoadError) as exc: + load_spec(f"{bad}:SPEC") + assert "identifier-like" in str(exc.value) + + +def test_starter_kernel_must_exist_for_loaded_spec(tmp_path: Path): + bad = tmp_path / "missing_starter.py" + bad.write_text( + "from autokernel.specs import DT_BYTES, KernelSpec, Tolerance, size\n" + "SPEC = KernelSpec(\n" + " name='no_starter',\n" + " reference_fn=lambda x: x,\n" + " input_generator=lambda *a, **k: {},\n" + " sizes={'small': {'n': 1}, 'medium': {'n': 2}, 'large': {'n': 3}},\n" + " dtypes=('float32',),\n" + " tolerances={'float32': Tolerance(atol=1e-5, rtol=1e-5)},\n" + " flops_fn=size('n'),\n" + " bytes_fn=size('n') * DT_BYTES,\n" + " shape_keys=('n',),\n" + " starter_kernels={'triton': '/nope/does_not_exist.py'},\n" + ")\n" + ) + with pytest.raises(SpecValidationError) as exc: + load_spec(f"{bad}:SPEC") + assert "starter kernel not found" in str(exc.value) + + +# --------------------------------------------------------------------------- +# Collisions and selection +# --------------------------------------------------------------------------- + +def test_collision_with_builtin_is_rejected_by_default(): + registry = create_builtin_registry() + with pytest.raises(SpecCollisionError, match="already registered"): + load_spec(f"{FIXTURE_FILE}:COLLIDING_SPEC", registry=registry) + + +def test_collision_is_allowed_with_override(): + registry = create_builtin_registry() + spec = load_spec(f"{FIXTURE_FILE}:COLLIDING_SPEC", registry=registry, override=True) + assert spec.name == "matmul" + + +def test_no_collision_when_registry_not_supplied(): + spec = load_spec(f"{FIXTURE_FILE}:COLLIDING_SPEC") + assert spec.name == "matmul" + + +def test_resolve_spec_prefers_locator_over_name(): + spec, registry = resolve_spec( + spec_locator=f"{FIXTURE_FILE}:SPEC", name="matmul" + ) + assert spec.name == "fixture_add" + assert registry.contains("fixture_add") + assert registry.contains("matmul") + + +def test_resolve_spec_falls_back_to_name(): + spec, _registry = resolve_spec(name="rmsnorm") + assert spec.name == "rmsnorm" + + +def test_resolve_spec_requires_a_selection(): + with pytest.raises(SpecLoadError, match="no operation selected"): + resolve_spec() + + +def test_resolve_spec_registers_into_the_supplied_registry_only(): + isolated = KernelRegistry([make_spec(name="only_here")]) + _spec, registry = resolve_spec( + spec_locator=f"{FIXTURE_FILE}:SPEC", registry=isolated + ) + assert registry is isolated + assert isolated.list_names() == ("only_here", "fixture_add") + assert not create_builtin_registry().contains("fixture_add") + + +def test_resolve_spec_collision_respects_override(): + registry = create_builtin_registry() + with pytest.raises(SpecCollisionError): + resolve_spec(spec_locator=f"{FIXTURE_FILE}:COLLIDING_SPEC", registry=registry) + spec, registry2 = resolve_spec( + spec_locator=f"{FIXTURE_FILE}:COLLIDING_SPEC", + registry=create_builtin_registry(), + override=True, + ) + assert spec.name == "matmul" + assert registry2.get("matmul") is spec + + +# --------------------------------------------------------------------------- +# The shipped example +# --------------------------------------------------------------------------- + +def test_example_custom_op_is_discoverable(repo_root: Path): + example = repo_root / "examples" / "custom_ops" / "add.py" + spec = load_spec(f"{example}:SPEC", registry=create_builtin_registry()) + assert spec.name == "custom_add" + assert set(spec.sizes) >= {"small", "medium", "large"} + assert spec.starter_kernel("triton").is_file() + assert spec.flops_fn(spec.sizes["large"]) == 4096 * 4096 + assert spec.bytes_fn(spec.sizes["large"], 2) == 3 * 4096 * 4096 * 2 + + +def test_example_custom_op_inputs_are_deterministic(torch_mod, repo_root: Path): + example = repo_root / "examples" / "custom_ops" / "add.py" + spec = load_spec(f"{example}:SPEC") + first = spec.input_generator(spec.sizes["small"], "float32", "cpu", 42) + second = spec.input_generator(spec.sizes["small"], "float32", "cpu", 42) + other = spec.input_generator(spec.sizes["small"], "float32", "cpu", 7) + assert set(first) == {"x", "y"} + for key in first: + assert torch_mod.equal(first[key], second[key]) + assert not torch_mod.equal(first["x"], other["x"]) + expected = spec.reference_fn(**first) + assert torch_mod.equal(expected, first["x"] + first["y"]) diff --git a/tests/test_spec_registry.py b/tests/test_spec_registry.py new file mode 100644 index 00000000..8ab9dc05 --- /dev/null +++ b/tests/test_spec_registry.py @@ -0,0 +1,483 @@ +"""Registry, validation and accounting behavior.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from autokernel.specs import ( + DT_BYTES, + BackwardSpec, + CompileSpec, + DuplicateSpecError, + EdgeCase, + KernelRegistry, + KernelSpec, + OutputSpec, + SpecNotFoundError, + SpecValidationError, + Tolerance, + canonical_dtype_name, + create_builtin_registry, + dtype_bytes, + serialize_accounting, + size, + validate_spec, +) +from conftest import make_spec, spec_kwargs + +# --------------------------------------------------------------------------- +# Registry behavior +# --------------------------------------------------------------------------- + +def test_register_and_get_round_trip(): + registry = KernelRegistry() + spec = make_spec(name="alpha") + registry.register(spec) + + assert registry.get("alpha") is spec + assert registry.contains("alpha") + assert "alpha" in registry + assert len(registry) == 1 + assert list(registry) == [spec] + + +def test_list_names_is_registration_ordered(): + registry = KernelRegistry( + [make_spec(name="zeta"), make_spec(name="alpha"), make_spec(name="mid")] + ) + assert registry.list_names() == ("zeta", "alpha", "mid") + # Repeated calls are stable. + assert registry.list_names() == registry.list_names() + + +def test_duplicate_registration_is_rejected(): + registry = KernelRegistry([make_spec(name="alpha")]) + with pytest.raises(DuplicateSpecError, match="already registered"): + registry.register(make_spec(name="alpha")) + + +def test_duplicate_registration_with_override_replaces(): + first = make_spec(name="alpha") + second = make_spec(name="alpha", speedup_estimate="9x") + registry = KernelRegistry([first]) + registry.register(second, override=True) + assert registry.get("alpha") is second + assert registry.list_names() == ("alpha",) + + +def test_unknown_name_lists_available_specs(): + registry = KernelRegistry([make_spec(name="alpha")]) + with pytest.raises(SpecNotFoundError) as exc: + registry.get("nope") + assert "alpha" in str(exc.value) + + +def test_registry_rejects_non_specs(): + registry = KernelRegistry() + with pytest.raises(SpecValidationError): + registry.register("not a spec") # type: ignore[arg-type] + + +def test_fresh_registries_are_isolated(): + a = create_builtin_registry() + b = create_builtin_registry() + a.register(make_spec(name="only_in_a")) + assert a.contains("only_in_a") + assert not b.contains("only_in_a") + assert b.list_names() == create_builtin_registry().list_names() + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("bad_name", ["", " ", "2fast", "has space", "has-dash", None, 7]) +def test_reject_non_identifier_names(bad_name): + with pytest.raises(SpecValidationError, match="name"): + make_spec(name=bad_name) + + +def test_reject_duplicate_size_labels(): + sizes = [ + ("small", {"rows": 1, "cols": 1}), + ("medium", {"rows": 2, "cols": 2}), + ("large", {"rows": 3, "cols": 3}), + ("small", {"rows": 4, "cols": 4}), + ] + with pytest.raises(SpecValidationError, match="duplicate size label 'small'"): + make_spec(sizes=sizes) + + +@pytest.mark.parametrize("missing", ["small", "medium", "large"]) +def test_reject_missing_standard_size(missing): + sizes = { + label: {"rows": 4, "cols": 4} + for label in ("small", "medium", "large") + if label != missing + } + spec = make_spec(sizes=sizes) + with pytest.raises(SpecValidationError, match="missing required size label"): + KernelRegistry().register(spec) + + +def test_reject_unknown_dtype(): + with pytest.raises(SpecValidationError, match="unknown dtype"): + make_spec(dtypes=("float8",)) + + +def test_reject_duplicate_dtype(): + with pytest.raises(SpecValidationError, match="duplicate dtype"): + make_spec(dtypes=("float16", "float16")) + + +def test_reject_missing_tolerance_for_declared_dtype(): + with pytest.raises(SpecValidationError, match="missing tolerance for declared dtype"): + make_spec( + dtypes=("float16", "bfloat16"), + tolerances={"float16": Tolerance(atol=1e-2, rtol=1e-2)}, + ) + + +@pytest.mark.parametrize("atol,rtol", [(-1e-3, 1e-3), (1e-3, -1e-3)]) +def test_reject_negative_tolerances(atol, rtol): + with pytest.raises(SpecValidationError, match="non-negative"): + Tolerance(atol=atol, rtol=rtol) + + +def test_reject_negative_tolerance_through_mapping(): + with pytest.raises(SpecValidationError, match="non-negative"): + make_spec(dtypes=("float16",), tolerances={"float16": {"atol": -1.0, "rtol": 0.0}}) + + +@pytest.mark.parametrize("atol,rtol", [(float("nan"), 1e-3), (1e-3, float("nan"))]) +def test_reject_nan_tolerances(atol, rtol): + with pytest.raises(SpecValidationError, match="must not be NaN"): + Tolerance(atol=atol, rtol=rtol) + + +@pytest.mark.parametrize( + "atol,rtol", + [ + (float("inf"), 1e-3), + (1e-3, float("inf")), + (float("-inf"), 1e-3), + (1e-3, float("-inf")), + ], +) +def test_reject_infinite_tolerances(atol, rtol): + with pytest.raises(SpecValidationError, match="must be finite"): + Tolerance(atol=atol, rtol=rtol) + + +def test_reject_infinite_tolerance_through_mapping(): + with pytest.raises(SpecValidationError, match="must be finite"): + make_spec( + dtypes=("float16",), + tolerances={"float16": {"atol": float("inf"), "rtol": 0.0}}, + ) + + +def test_extra_tolerances_are_allowed(): + spec = make_spec( + dtypes=("float16",), + tolerances={ + "float16": Tolerance(atol=1e-2, rtol=1e-2), + "float32": Tolerance(atol=1e-5, rtol=1e-5), + }, + ) + assert set(spec.tolerances) == {"float16", "float32"} + + +def test_reject_missing_starter_kernel_file(tmp_path: Path): + missing = tmp_path / "nope.py" + spec = make_spec(starter_kernels={"triton": missing}) + with pytest.raises(SpecValidationError, match="starter kernel not found"): + KernelRegistry().register(spec) + + +def test_accept_existing_starter_kernel_file(tmp_path: Path): + present = tmp_path / "starter.py" + present.write_text("KERNEL_TYPE = 'unit_op'\n") + spec = make_spec(starter_kernels={"triton": present}) + assert spec.starter_kernel("triton") == present + assert spec.starter_kernel("cuda") is None + + +def test_reject_inconsistent_shape_aliases(): + aliases = [("M", "rows"), ("M", "cols")] + with pytest.raises(SpecValidationError, match="resolves inconsistently"): + make_spec(shape_aliases=aliases) + + +def test_reject_alias_to_unknown_shape_key(): + with pytest.raises(SpecValidationError, match="not a shape key"): + make_spec(shape_aliases={"M": "not_a_key"}) + + +@pytest.mark.parametrize("field", ["reference_fn", "input_generator"]) +def test_reject_non_callable_reference_and_generator(field): + with pytest.raises(SpecValidationError, match="must be callable"): + make_spec(**{field: "not callable"}) + + +def test_reject_size_keys_that_do_not_match_shape_keys(): + with pytest.raises(SpecValidationError, match="do not match shape_keys"): + make_spec( + sizes={ + "small": {"rows": 1, "cols": 1}, + "medium": {"rows": 2, "cols": 2}, + "large": {"rows": 3, "oops": 3}, + } + ) + + +def test_reject_non_positive_size(): + with pytest.raises(SpecValidationError, match="must be positive"): + make_spec( + sizes={ + "small": {"rows": 0, "cols": 1}, + "medium": {"rows": 2, "cols": 2}, + "large": {"rows": 3, "cols": 3}, + } + ) + + +@pytest.mark.parametrize("bad_dimension", [0, -1, 1.5, True]) +def test_reject_invalid_edge_case_dimension(bad_dimension): + edge = EdgeCase( + name="bad", size={"rows": bad_dimension, "cols": 1} + ) + with pytest.raises(SpecValidationError, match=r"edge_cases\[0\]\.size"): + make_spec(edge_cases=(edge,)) + + +@pytest.mark.parametrize("bad_dimension", [0, -1, 1.5, True]) +def test_reject_invalid_default_shape_dimension(bad_dimension): + with pytest.raises(SpecValidationError, match="default_shape"): + make_spec(default_shape={"rows": bad_dimension, "cols": 1}) + + +def test_empty_starter_kernels_are_valid_for_benchmark_only_specs(): + spec = make_spec(starter_kernels={}) + assert spec.starter_kernel("triton") is None + + +def test_reject_duplicate_edge_case_names(): + edges = ( + EdgeCase(name="dup", size={"rows": 1, "cols": 1}), + EdgeCase(name="dup", size={"rows": 2, "cols": 2}), + ) + with pytest.raises(SpecValidationError, match="duplicate edge case name"): + make_spec(edge_cases=edges) + + +def test_reject_edge_case_dtype_not_declared(): + edges = (EdgeCase(name="e", size={"rows": 1, "cols": 1}, dtype="bfloat16"),) + with pytest.raises(SpecValidationError, match="is not declared in dtypes"): + make_spec(dtypes=("float32",), tolerances={"float32": Tolerance(1e-5, 1e-5)}, edge_cases=edges) + + +def test_reject_edge_case_shape_mismatch(): + edges = (EdgeCase(name="e", size={"rows": 1}),) + with pytest.raises(SpecValidationError, match="do not match shape_keys"): + make_spec(edge_cases=edges) + + +def test_reject_flops_expression_with_unknown_key(): + with pytest.raises(SpecValidationError, match="unknown size key"): + make_spec(flops_fn=size("nope") * 2) + + +def test_reject_flops_expression_that_uses_dtype_bytes(): + with pytest.raises(SpecValidationError, match="must not depend on dtype bytes"): + make_spec(flops_fn=size("rows") * DT_BYTES) + + +def test_validation_error_names_spec_and_field(): + with pytest.raises(SpecValidationError) as exc: + make_spec(name="named_op", dtypes=("float8",)) + message = str(exc.value) + assert "'named_op'" in message + assert "'dtypes'" in message + + +def test_validate_spec_can_skip_standard_size_requirement(): + narrow = make_spec(sizes={"small": {"rows": 1, "cols": 1}}) + # Construction and explicit opt-out are fine... + validate_spec(narrow, require_standard_sizes=False) + # ...but registration enforces the standard labels. + with pytest.raises(SpecValidationError, match="missing required size label"): + validate_spec(narrow) + + +def test_spec_is_immutable(): + spec = make_spec() + with pytest.raises(Exception): + spec.name = "other" # type: ignore[misc] + + +def test_defensive_copy_of_sizes(): + sizes = { + "small": {"rows": 1, "cols": 1}, + "medium": {"rows": 2, "cols": 2}, + "large": {"rows": 3, "cols": 3}, + } + spec = make_spec(sizes=sizes) + sizes["small"]["rows"] = 999 + assert spec.sizes["small"]["rows"] == 1 + + +# --------------------------------------------------------------------------- +# Accounting expressions +# --------------------------------------------------------------------------- + +def test_accounting_expression_evaluation_and_source(): + flops = 2 * size("M") * size("N") * size("K") + assert flops({"M": 2, "N": 3, "K": 4}) == 48 + assert flops.to_source() == "2 * s['M'] * s['N'] * s['K']" + assert flops.size_keys() == {"M", "N", "K"} + assert not flops.uses_dtype_bytes() + + +def test_accounting_expression_parenthesizes_by_precedence(): + expr = (size("a") + size("b")) * DT_BYTES + assert expr.to_source() == "(s['a'] + s['b']) * dt_bytes" + assert expr({"a": 1, "b": 2}, 4) == 12 + assert expr.uses_dtype_bytes() + + +def test_accounting_power_is_right_associative_in_source(): + expr = 4 * (size("s") ** 2) + assert expr.to_source() == "4 * s['s'] ** 2" + assert expr({"s": 3}) == 36 + + +def test_accounting_supports_reflected_power(): + expr = 2 ** size("s") + assert expr.to_source() == "2 ** s['s']" + assert expr({"s": 4}) == 16 + + +def test_accounting_expression_requires_dtype_bytes_when_used(): + expr = size("rows") * DT_BYTES + with pytest.raises(ValueError, match="dtype byte width"): + expr({"rows": 4}) + + +def test_accounting_expression_reports_missing_size_key(): + expr = size("rows") + with pytest.raises(KeyError, match="rows"): + expr({"cols": 4}) + + +def test_serialize_accounting_returns_none_for_opaque_callable(): + assert serialize_accounting(lambda s: 1) is None + assert serialize_accounting(size("rows")) == "s['rows']" + + +def test_serialized_accounting_is_valid_python(): + spec = create_builtin_registry().get("matmul") + source = serialize_accounting(spec.bytes_fn) + compiled = compile(source, "", "eval") + assert eval(compiled, {"s": {"M": 2, "N": 3, "K": 4}, "dt_bytes": 2}) == ( + spec.bytes_fn({"M": 2, "N": 3, "K": 4}, 2) + ) + + +# --------------------------------------------------------------------------- +# dtype helpers +# --------------------------------------------------------------------------- + +def test_dtype_helpers_are_torch_free(): + assert dtype_bytes("float16") == 2 + assert dtype_bytes("bfloat16") == 2 + assert dtype_bytes("float32") == 4 + assert canonical_dtype_name("float32") == "float32" + with pytest.raises(ValueError, match="unsupported dtype"): + canonical_dtype_name("float8") + + +def test_resolve_torch_dtype_translates_only_in_runtime(torch_mod): + from autokernel.specs import resolve_torch_dtype + + assert resolve_torch_dtype("bfloat16") is torch_mod.bfloat16 + assert resolve_torch_dtype(torch_mod.float32) is torch_mod.float32 + assert canonical_dtype_name(torch_mod.float16) == "float16" + + +def test_validate_spec_rejects_non_spec_objects(): + with pytest.raises(SpecValidationError, match="expected a KernelSpec"): + validate_spec(object()) # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# Verification policy types (Week 2) +# --------------------------------------------------------------------------- + +def test_verification_policies_default_to_none(): + spec = make_spec() + assert spec.output_spec is None + assert spec.backward_spec is None + assert spec.compile_spec is None + + +def test_output_spec_accepts_mapping_coercion(): + spec = make_spec(output_spec={"included_paths": ("output[0]",), "compare_non_tensors": False}) + assert isinstance(spec.output_spec, OutputSpec) + assert spec.output_spec.included_paths == ("output[0]",) + assert spec.output_spec.compare_non_tensors is False + + +def test_output_spec_rejects_duplicate_paths(): + with pytest.raises(SpecValidationError, match="duplicate path"): + OutputSpec(included_paths=("output", "output")) + + +def test_output_and_backward_specs_reject_empty_output_paths(): + with pytest.raises(SpecValidationError, match="at least one path"): + OutputSpec(included_paths=()) + with pytest.raises(SpecValidationError, match="at least one path"): + BackwardSpec(differentiable_inputs=("x",), output_paths=()) + + +def test_output_spec_rejects_unknown_mapping_keys(): + with pytest.raises(SpecValidationError, match="output_spec"): + make_spec(output_spec={"nonsense": True}) + + +def test_backward_spec_requires_differentiable_inputs(): + with pytest.raises(SpecValidationError, match="at least one input"): + BackwardSpec(differentiable_inputs=()) + + +def test_backward_spec_normalizes_tolerances(): + policy = BackwardSpec( + differentiable_inputs=("x",), + tolerances={"float32": {"atol": 1e-4, "rtol": 1e-4}}, + ) + assert policy.tolerances["float32"] == Tolerance(atol=1e-4, rtol=1e-4) + + +def test_backward_spec_rejects_unknown_tolerance_dtype(): + with pytest.raises(SpecValidationError, match="unknown dtype key"): + BackwardSpec(differentiable_inputs=("x",), tolerances={"float64": Tolerance(0.0, 0.0)}) + + +def test_backward_spec_rejects_infinite_tolerances(): + with pytest.raises(SpecValidationError, match="must be finite"): + BackwardSpec( + differentiable_inputs=("x",), + tolerances={"float32": {"atol": float("inf"), "rtol": 0.0}}, + ) + + +def test_compile_spec_rejects_non_bool_flags(): + with pytest.raises(SpecValidationError, match="CompileSpec.fullgraph must be a bool"): + CompileSpec(fullgraph="yes") + + +def test_kernel_spec_rejects_wrongly_typed_policies(): + with pytest.raises(SpecValidationError, match="backward_spec"): + make_spec(backward_spec="not-a-spec") diff --git a/tests/test_wan_additional_targets.py b/tests/test_wan_additional_targets.py new file mode 100644 index 00000000..fb8cabe2 --- /dev/null +++ b/tests/test_wan_additional_targets.py @@ -0,0 +1,68 @@ +"""CPU contracts for the remaining Wan overnight-campaign targets.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from autokernel.verification import ( + load_shape_corpus, + validate_corpus_against_spec, +) +from models.wan_gated_residual import ( + SPEC as GATED_RESIDUAL_SPEC, +) +from models.wan_gated_residual import ( + gen_wan_gated_residual_inputs, + wan_gated_residual_ref, +) +from models.wan_modulated_layer_norm import ( + SPEC as MODULATED_NORM_SPEC, +) +from models.wan_modulated_layer_norm import ( + gen_wan_modulated_layer_norm_inputs, + wan_modulated_layer_norm_ref, +) + + +def test_additional_wan_specs_cover_production_corpora(repo_root): + for spec, filename in ( + (GATED_RESIDUAL_SPEC, "wan_gated_residual_corpus.json"), + (MODULATED_NORM_SPEC, "wan_modulated_layer_norm_corpus.json"), + ): + corpus = load_shape_corpus(repo_root / "models" / filename) + validate_corpus_against_spec(corpus, spec) + assert len(corpus.cases) == 5 + assert {case.size["hidden"] for case in corpus.cases} == { + 1536, + 5120, + } + + +def test_wan_gated_residual_reference_matches_dtype_boundary(): + size = {"batch": 2, "tokens": 3, "hidden": 5} + inputs = gen_wan_gated_residual_inputs( + size, "bfloat16", "cpu", seed=13 + ) + actual = wan_gated_residual_ref(**inputs) + expected = ( + inputs["residual"].float() + + inputs["x"].float() * inputs["gate"][:, None, :] + ).to(torch.bfloat16) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_wan_modulated_norm_reference_matches_dtype_boundary(): + size = {"batch": 2, "tokens": 3, "hidden": 5} + inputs = gen_wan_modulated_layer_norm_inputs( + size, "bfloat16", "cpu", seed=17 + ) + actual = wan_modulated_layer_norm_ref(**inputs) + normalized = F.layer_norm( + inputs["x"].float(), (size["hidden"],), None, None, 1e-6 + ) + expected = ( + normalized * (1 + inputs["scale"][:, None, :]) + + inputs["shift"][:, None, :] + ).to(torch.bfloat16) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/test_wan_gated_residual_norm.py b/tests/test_wan_gated_residual_norm.py new file mode 100644 index 00000000..7aa1e7f1 --- /dev/null +++ b/tests/test_wan_gated_residual_norm.py @@ -0,0 +1,69 @@ +"""CPU contract tests for the first Wan production kernel specification.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from autokernel.verification.corpus import ( + load_shape_corpus, + validate_corpus_against_spec, +) +from models.wan_gated_residual_norm import ( + SPEC, + gen_wan_gated_residual_norm_inputs, + wan_gated_residual_norm_ref, +) + + +def test_wan_spec_covers_production_widths_and_dtypes(): + assert SPEC.sizes["small"]["hidden"] == 1536 + assert SPEC.sizes["large"]["hidden"] == 5120 + assert SPEC.dtypes == ("bfloat16", "float16") + assert SPEC.starter_kernels["triton"].name == "wan_gated_residual_norm.py" + + +def test_wan_production_corpus_matches_spec(repo_root): + corpus = load_shape_corpus( + repo_root / "models" / "wan_gated_residual_norm_corpus.json" + ) + validate_corpus_against_spec(corpus, SPEC) + assert len(corpus.cases) == 5 + assert {case.size["hidden"] for case in corpus.cases} == {1536, 5120} + assert any("sequence-parallel-4" in case.tags for case in corpus.cases) + + +def test_wan_input_generator_is_deterministic_on_cpu(): + size = {"batch": 2, "tokens": 3, "hidden": 5} + first = gen_wan_gated_residual_norm_inputs(size, "bfloat16", "cpu", seed=7) + second = gen_wan_gated_residual_norm_inputs(size, "bfloat16", "cpu", seed=7) + assert first.keys() == second.keys() + for name in first: + torch.testing.assert_close(first[name], second[name], rtol=0, atol=0) + + +def test_wan_reference_matches_explicit_fastvideo_dtype_boundary(): + size = {"batch": 2, "tokens": 3, "hidden": 5} + inputs = gen_wan_gated_residual_norm_inputs( + size, "bfloat16", "cpu", seed=11 + ) + normalized, updated = wan_gated_residual_norm_ref(**inputs) + + expected_updated_fp32 = ( + inputs["residual"].float() + + inputs["x"].float() * inputs["gate"][:, None, :] + ) + expected_normalized = F.layer_norm( + expected_updated_fp32, + (size["hidden"],), + inputs["weight"].float(), + inputs["bias"].float(), + 1e-6, + ).to(torch.bfloat16) + + assert normalized.dtype == torch.bfloat16 + assert updated.dtype == torch.bfloat16 + torch.testing.assert_close(normalized, expected_normalized, rtol=0, atol=0) + torch.testing.assert_close( + updated, expected_updated_fp32.to(torch.bfloat16), rtol=0, atol=0 + ) diff --git a/tests/test_workload.py b/tests/test_workload.py new file mode 100644 index 00000000..dd1b0a83 --- /dev/null +++ b/tests/test_workload.py @@ -0,0 +1,365 @@ +"""Workload manifest schema, result classification, and launcher bridge.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from autokernel.workload import ( + WORKLOAD_SCHEMA_VERSION, + WorkloadError, + WorkloadManifest, + load_workload, +) +from autokernel.workload.launcher import ( + build_launcher_command, + resolve_launcher, + run_ab, +) +from autokernel.workload.result import ( + GenerationRunResult, + classify_end_to_end, + compare_frame_outputs, + load_generation_result, + write_generation_result, +) + + +REPO_ROOT = Path(__file__).resolve().parents[1] +WORKLOADS = REPO_ROOT / "workloads" + + +def test_load_wan_and_ltx_manifests(): + wan = load_workload(WORKLOADS / "wan_t2v_1.3b_480p.yaml") + ltx = load_workload(WORKLOADS / "ltx_480p.yaml") + assert wan.schema_version == WORKLOAD_SCHEMA_VERSION + assert wan.workload_id == "wan-t2v-1.3b-480p" + assert wan.model.model_id == "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + assert wan.sampling.height == 480 + assert wan.sampling.width == 832 + assert wan.mode_env is not None + assert wan.mode_env.for_mode("native")["FASTVIDEO_WAN_FUSIONS"] == "0" + assert wan.mode_env.for_mode("optimized")["FASTVIDEO_WAN_FUSIONS"] == "1" + + assert ltx.workload_id == "ltx-t2v-480p" + assert "LTX" in ltx.model.model_id or "ltx" in ltx.model.model_id.lower() + assert ltx.task == "t2v" + assert ltx.performance is not None + assert ltx.performance.min_end_to_end_speedup == pytest.approx(1.01) + + +def test_workload_roundtrip_json(tmp_path): + wan = load_workload(WORKLOADS / "wan_t2v_1.3b_480p.yaml") + path = tmp_path / "wan.json" + path.write_text(json.dumps(wan.as_dict(), indent=2), encoding="utf-8") + again = load_workload(path) + assert again.as_dict() == wan.as_dict() + + +def test_generation_request_matches_wan_ab_shape(): + wan = load_workload(WORKLOADS / "wan_t2v_1.3b_480p.yaml") + request = wan.generation_request() + assert "prompt" in request + assert request["sampling"]["height"] == 480 + assert request["sampling"]["width"] == 832 + assert request["sampling"]["num_frames"] == 49 + assert request["sampling"]["num_inference_steps"] == 4 + assert request["sampling"]["guidance_scale"] == 5.0 + assert request["sampling"]["seed"] == 1024 + assert request["output"]["return_frames"] is True + kwargs = wan.generator_kwargs() + assert kwargs["num_gpus"] == 1 + assert kwargs["text_encoder_cpu_offload"] is True + + +def test_rejects_secret_fields(): + payload = load_workload(WORKLOADS / "ltx_480p.yaml").as_dict() + payload["runtime"]["password"] = "nope" # noqa: S105 - rejection fixture + with pytest.raises(WorkloadError, match="secret fields"): + WorkloadManifest.from_dict(payload) + + +def test_requires_prompt_or_prompt_file(): + payload = load_workload(WORKLOADS / "ltx_480p.yaml").as_dict() + del payload["prompt"] + with pytest.raises(WorkloadError, match="prompt or prompt_file"): + WorkloadManifest.from_dict(payload) + + +def test_prompt_file_resolution(tmp_path): + prompt_path = tmp_path / "prompt.txt" + prompt_path.write_text("hello from file\n", encoding="utf-8") + payload = load_workload(WORKLOADS / "ltx_480p.yaml").as_dict() + del payload["prompt"] + payload["prompt_file"] = "prompt.txt" + manifest = WorkloadManifest.from_dict(payload, source=str(tmp_path / "w.yaml")) + assert manifest.resolve_prompt(base_dir=tmp_path) == "hello from file" + + +def test_rejects_unknown_top_level_field(): + payload = load_workload(WORKLOADS / "ltx_480p.yaml").as_dict() + payload["callable"] = "models.foo:bar" + with pytest.raises(WorkloadError, match="unknown field"): + WorkloadManifest.from_dict(payload) + + +def test_rejects_bad_schema_version(): + payload = load_workload(WORKLOADS / "ltx_480p.yaml").as_dict() + payload["schema_version"] = 99 + with pytest.raises(WorkloadError, match="unsupported version"): + WorkloadManifest.from_dict(payload) + + +def test_result_schema_and_classification(tmp_path): + native = GenerationRunResult.from_dict( + { + "schema_version": 1, + "mode": "native", + "status": "ok", + "workload_id": "wan-t2v-1.3b-480p", + "model_id": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "request": {"prompt": "x"}, + "warmups": 1, + "runs": 2, + "wall_seconds": [36.7, 36.6], + "median_wall_seconds": 36.65, + "generation_seconds": [30.0, 30.1], + "peak_memory_mb": [20000.0, 20010.0], + "environment": {"cuda": "12.8"}, + } + ) + optimized = GenerationRunResult.from_dict( + { + "schema_version": 1, + "mode": "optimized", + "status": "ok", + "workload_id": "wan-t2v-1.3b-480p", + "model_id": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "request": {"prompt": "x"}, + "warmups": 1, + "runs": 2, + "wall_seconds": [36.8, 36.5], + "median_wall_seconds": 36.65, + "generation_seconds": [30.0, 30.0], + "peak_memory_mb": [20000.0, 20000.0], + "environment": {"cuda": "12.8"}, + } + ) + verdict = classify_end_to_end(native, optimized) + assert verdict["classification"] == "neutral" + + improved = GenerationRunResult.from_dict( + { + **optimized.as_dict(), + "wall_seconds": [30.0, 30.2], + "median_wall_seconds": 30.1, + } + ) + assert classify_end_to_end(native, improved)["classification"] == "improved" + + path = tmp_path / "native_result.json" + write_generation_result(native, path) + loaded = load_generation_result(path) + assert loaded.median_wall_seconds == pytest.approx(36.65) + + +def test_build_launcher_command_and_resolve(tmp_path): + checkout = tmp_path / "FastVideo" + script = ( + checkout + / "examples" + / "inference" + / "optimizations" + / "generation_launcher.py" + ) + script.parent.mkdir(parents=True) + script.write_text("# stub\n", encoding="utf-8") + resolved = resolve_launcher(checkout) + assert resolved == script + command = build_launcher_command( + python="python3", + launcher=script, + workload=WORKLOADS / "wan_t2v_1.3b_480p.yaml", + mode="native", + output_dir=tmp_path / "out", + ) + assert command[:2] == ["python3", str(script)] + assert "--mode" in command and "native" in command + + +def test_run_ab_resume(tmp_path, monkeypatch): + checkout = tmp_path / "FastVideo" + script = ( + checkout + / "examples" + / "inference" + / "optimizations" + / "generation_launcher.py" + ) + script.parent.mkdir(parents=True) + script.write_text("# stub\n", encoding="utf-8") + out = tmp_path / "ab" + out.mkdir() + + calls: list[str] = [] + + def fake_run_mode(**kwargs): + import numpy as np + + mode = kwargs["mode"] + calls.append(mode) + frames_path = out / f"{mode}_frames.npy" + np.save(frames_path, np.zeros((2, 4, 4, 3), dtype=np.uint8)) + result = GenerationRunResult.from_dict( + { + "schema_version": 1, + "mode": mode if mode != "fused" else "optimized", + "status": "ok", + "workload_id": "wan-t2v-1.3b-480p", + "model_id": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "request": {}, + "warmups": 1, + "runs": 1, + "wall_seconds": [10.0 if mode == "native" else 9.0], + "median_wall_seconds": 10.0 if mode == "native" else 9.0, + "generation_seconds": [8.0], + "peak_memory_mb": [1000.0], + "environment": {}, + "frames_path": str(frames_path), + } + ) + write_generation_result(result, out / f"{mode}_result.json") + + class _Done: + returncode = 0 + stdout = "" + stderr = "" + + return _Done() + + monkeypatch.setattr( + "autokernel.workload.launcher.run_mode", fake_run_mode + ) + first = run_ab( + fastvideo_checkout=checkout, + workload=WORKLOADS / "wan_t2v_1.3b_480p.yaml", + output_dir=out, + resume=True, + ) + assert first["comparison"]["classification"] in {"improved", "neutral"} + assert calls == ["native", "optimized"] + + second = run_ab( + fastvideo_checkout=checkout, + workload=WORKLOADS / "wan_t2v_1.3b_480p.yaml", + output_dir=out, + resume=True, + ) + assert calls == ["native", "optimized"] # no re-run + assert second["comparison"] is not None + assert second["comparison"]["parity"]["passed"] is True + + +def test_compare_frame_outputs_policies(tmp_path): + import numpy as np + + native = tmp_path / "native.npy" + same = tmp_path / "same.npy" + close = tmp_path / "close.npy" + np.save(native, np.array([[[[10, 20, 30]]]], dtype=np.uint8)) + np.save(same, np.array([[[[10, 20, 30]]]], dtype=np.uint8)) + np.save(close, np.array([[[[11, 20, 30]]]], dtype=np.uint8)) + + exact = compare_frame_outputs(native, same, policy="byte_equal") + assert exact["passed"] is True + mismatch = compare_frame_outputs(native, close, policy="byte_equal") + assert mismatch["passed"] is False + frames_only = compare_frame_outputs(native, close, policy="frames_only") + assert frames_only["passed"] is True + tolerant = compare_frame_outputs( + native, + close, + policy="tolerance", + atol=1.0, + ) + assert tolerant["passed"] is True + + +def test_cli_validate(monkeypatch): + from workload import main + + assert main(["validate", str(WORKLOADS / "ltx_480p.yaml")]) == 0 + assert main(["show", str(WORKLOADS / "wan_t2v_1.3b_480p.yaml")]) == 0 + + +def test_read_state_rejects_corrupt_json(tmp_path): + from autokernel.workload.launcher import _read_state + + path = tmp_path / "launcher_state.json" + path.write_text("{not-json", encoding="utf-8") + with pytest.raises(WorkloadError, match="invalid JSON"): + _read_state(path) + + +def test_run_ab_rejects_unknown_modes(tmp_path, monkeypatch): + from autokernel.workload.launcher import run_ab + + checkout = tmp_path / "FastVideo" + script = ( + checkout + / "examples" + / "inference" + / "optimizations" + / "generation_launcher.py" + ) + script.parent.mkdir(parents=True) + script.write_text("# stub\n", encoding="utf-8") + with pytest.raises(WorkloadError, match="unsupported launcher mode"): + run_ab( + fastvideo_checkout=checkout, + workload=WORKLOADS / "wan_t2v_1.3b_480p.yaml", + output_dir=tmp_path / "out", + modes=("native", "candidate"), + ) + + +def test_classify_end_to_end_zero_optimized_median(): + from autokernel.workload.result import GenerationRunResult, classify_end_to_end + + native = GenerationRunResult.from_dict( + { + "schema_version": 1, + "mode": "native", + "status": "ok", + "workload_id": "w", + "model_id": "m", + "request": {}, + "warmups": 0, + "runs": 1, + "wall_seconds": [1.0], + "median_wall_seconds": 1.0, + "generation_seconds": [1.0], + "peak_memory_mb": [1.0], + "environment": {}, + } + ) + optimized = GenerationRunResult.from_dict( + { + "schema_version": 1, + "mode": "optimized", + "status": "ok", + "workload_id": "w", + "model_id": "m", + "request": {}, + "warmups": 0, + "runs": 1, + "wall_seconds": [0.0], + "median_wall_seconds": 0.0, + "generation_seconds": [0.0], + "peak_memory_mb": [1.0], + "environment": {}, + } + ) + assert classify_end_to_end(native, optimized)["classification"] == "failed" diff --git a/workload.py b/workload.py new file mode 100644 index 00000000..143fb83e --- /dev/null +++ b/workload.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +"""Validate and drive versioned FastVideo workload manifests. + +Usage: + python workload.py validate workloads/ltx_480p.yaml + python workload.py show workloads/wan_t2v_1.3b_480p.yaml + python workload.py run-ab \\ + --fastvideo-checkout /path/to/FastVideo \\ + --workload workloads/wan_t2v_1.3b_480p.yaml \\ + --output workspace/wan_ab +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +from autokernel.workload import WorkloadError, load_workload +from autokernel.workload.launcher import run_ab +from autokernel.workload.result import load_generation_result + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Validate and execute FastVideo workload manifests" + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + validate = subparsers.add_parser( + "validate", help="Validate a workload manifest" + ) + validate.add_argument("workload", type=Path) + + show = subparsers.add_parser( + "show", help="Print the normalized workload JSON" + ) + show.add_argument("workload", type=Path) + + run = subparsers.add_parser( + "run-ab", + help="Run native and optimized modes via a FastVideo launcher", + ) + run.add_argument("--fastvideo-checkout", type=Path, required=True) + run.add_argument("--workload", type=Path, required=True) + run.add_argument("--output", type=Path, required=True) + run.add_argument("--model", help="Optional model_id override") + run.add_argument( + "--launcher-script", + type=Path, + help="Override path to generation_launcher.py", + ) + run.add_argument( + "--no-resume", + action="store_true", + help="Ignore completed stages and re-run everything", + ) + run.add_argument( + "--modes", + default="native,optimized", + help="Comma-separated modes (default: native,optimized)", + ) + + result = subparsers.add_parser( + "validate-result", + help="Validate a generation_launcher result JSON", + ) + result.add_argument("result", type=Path) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + try: + if args.command == "validate": + workload = load_workload(args.workload) + print("WORKLOAD_VALIDATION: PASS") + print(f"workload_id: {workload.workload_id}") + print(f"model_id: {workload.model.model_id}") + print(f"task: {workload.task}") + return 0 + + if args.command == "show": + workload = load_workload(args.workload) + print(json.dumps(workload.as_dict(), indent=2)) + return 0 + + if args.command == "validate-result": + result = load_generation_result(args.result) + print("RESULT_VALIDATION: PASS") + print(f"mode: {result.mode}") + print(f"status: {result.status}") + print(f"workload_id: {result.workload_id}") + return 0 + + if args.command == "run-ab": + payload = run_ab( + fastvideo_checkout=args.fastvideo_checkout, + workload=args.workload, + output_dir=args.output, + launcher_script=args.launcher_script, + model_override=args.model, + modes=tuple( + mode.strip() + for mode in args.modes.split(",") + if mode.strip() + ), + resume=not args.no_resume, + ) + print(json.dumps(payload, indent=2)) + return 0 + except WorkloadError as exc: + print(f"WORKLOAD: FAIL\n{exc}", file=sys.stderr) + return 2 + + print(f"unknown command: {args.command}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/workloads/ltx_480p.yaml b/workloads/ltx_480p.yaml new file mode 100644 index 00000000..888a2e51 --- /dev/null +++ b/workloads/ltx_480p.yaml @@ -0,0 +1,46 @@ +# Canonical LTX text-to-video workload for the universal optimization agent. +# Uses a FastVideo-supported LTX-2 distilled checkpoint so discovery and +# end-to-end validation do not require LTX-specific fusion annotations. +schema_version: 1 +workload_id: ltx-t2v-480p +description: > + Canonical LTX proof workload for MotionKernel's model-independent optimizer. + Distilled checkpoint keeps overnight GPU budgets practical while still + exercising a full FastVideo generation path. +model: + model_id: FastVideo/LTX2-Distilled-Diffusers +task: t2v +prompt: > + A red fox trots through fresh powder snow at golden hour, soft rim light, + gentle camera pan following the animal, cinematic nature documentary style. +sampling: + height: 480 + width: 768 + num_frames: 97 + num_inference_steps: 8 + guidance_scale: 1.0 + seed: 1024 + fps: 24 + dtype: bfloat16 +runtime: + num_gpus: 1 + use_fsdp_inference: false + dit_cpu_offload: false + vae_cpu_offload: false + text_encoder_cpu_offload: true + pin_cpu_memory: false +measurement: + warmups: 1 + runs: 2 + save_frames: true + save_video: false +parity: + policy: byte_equal +performance: + min_end_to_end_speedup: 1.01 + max_peak_memory_regression: 0.05 +tags: + - ltx + - t2v + - 480p + - proof diff --git a/workloads/wan_t2v_1.3b_480p.yaml b/workloads/wan_t2v_1.3b_480p.yaml new file mode 100644 index 00000000..54fd28c7 --- /dev/null +++ b/workloads/wan_t2v_1.3b_480p.yaml @@ -0,0 +1,49 @@ +# Canonical Wan text-to-video workload matching the existing A/B benchmark in +# FastVideo examples/inference/optimizations/wan_fusions_ab.py. +schema_version: 1 +workload_id: wan-t2v-1.3b-480p +description: > + Representative Wan 2.1 1.3B T2V workload used for end-to-end native-versus-fused + measurement. Default measurement uses a short step count suitable for CI-scale + dry runs; raise sampling.num_inference_steps to 50 for production ranking. +model: + model_id: Wan-AI/Wan2.1-T2V-1.3B-Diffusers +task: t2v +prompt: > + A curious raccoon peers through a vibrant field of yellow sunflowers, soft + natural light, steady cinematic camera. +sampling: + height: 480 + width: 832 + num_frames: 49 + num_inference_steps: 4 + guidance_scale: 5.0 + seed: 1024 + dtype: bfloat16 +runtime: + num_gpus: 1 + use_fsdp_inference: false + dit_cpu_offload: false + vae_cpu_offload: false + text_encoder_cpu_offload: true + pin_cpu_memory: false +measurement: + warmups: 1 + runs: 2 + save_frames: true + save_video: false +parity: + policy: byte_equal +performance: + min_end_to_end_speedup: 1.01 + max_peak_memory_regression: 0.05 +mode_env: + native: + FASTVIDEO_WAN_FUSIONS: "0" + optimized: + FASTVIDEO_WAN_FUSIONS: "1" +tags: + - wan + - t2v + - 480p + - baseline