diff --git a/.github/workflows/pypi_publish.yml b/.github/workflows/pypi_publish.yml index a6c7573..6c9a5c2 100644 --- a/.github/workflows/pypi_publish.yml +++ b/.github/workflows/pypi_publish.yml @@ -59,9 +59,46 @@ jobs: run: | pytest tests ${{ matrix.pytest-args }} + # The ahead-of-time C compiler (`fnnx.extras.compilers.c`). It needs the `compiler` + # extra and a system C compiler, which the ubuntu image ships, so it gets its own leg + # rather than a matrix entry: the lean legs exclude it through the existing + # `--ignore-glob='*test_extra_*'`, and the full legs skip it cleanly when the extra or + # a C compiler is absent. `onnx` is pinned to one exact release instead of taken from + # the extra's supported range, because that release defines the schema set and node + # corpus the conformance ledger and pass-list ratchet are keyed to. `scikit-learn` and + # `skl2onnx` are what the ONNX-ML tests convert their models with, where the node corpus + # is too thin to cover an op; those tests skip cleanly without them. + compiler-test: + if: | + github.event_name == 'pull_request' || + github.event_name == 'workflow_dispatch' || + (github.event_name == 'release' && + startsWith(github.ref, 'refs/tags/') && + contains(github.ref, 'py_v')) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: '**/pyproject.toml' + + - name: Install dependencies + run: | + pip install -e "src/python[core,test-essential,compiler]" \ + "onnx==1.22.*" pytest-xdist scikit-learn skl2onnx + + - name: Run compiler tests + working-directory: src/python + run: | + pytest tests/test_extra_compiler_*.py -n auto + deploy: runs-on: ubuntu-latest - needs: [test] + needs: [test, compiler-test] environment: pypi permissions: id-token: write diff --git a/src/python/examples/torch_to_c.py b/src/python/examples/torch_to_c.py new file mode 100644 index 0000000..abe7b6e --- /dev/null +++ b/src/python/examples/torch_to_c.py @@ -0,0 +1,389 @@ +"""Hands-on demo: torch model -> ONNX -> FNNX bundle -> self-contained C99 header. + +Builds a two-stage torch model (feature standardization, then a small MLP), exports each +stage to ONNX, wires the two into an FNNX ``pipeline`` bundle, and compiles that bundle +with ``fnnx.extras.compilers.c`` into a single C header with no dependencies beyond libm. +The artifact is then built into a shared library and its outputs are checked against the +FNNX runtime (onnxruntime) and against torch itself. + +Usage (from src/python, with torch, onnx, onnxruntime, numpy and a C compiler installed): + + python examples/torch_to_c.py + +Everything is written under ./_fnnx_c_demo/ so it can be inspected afterwards: + + _fnnx_c_demo/onnx/{scale,mlp}.onnx the exported stages + _fnnx_c_demo/scoring.fnnx/ the FNNX pipeline bundle + _fnnx_c_demo/c/scoring.h the compiled pipeline + its compile report + _fnnx_c_demo/c/main.c a C program that includes it, built with cc + _fnnx_c_demo/c_mlp/mlp.h the `mlp` stage compiled straight from ONNX + +The compiler has two entrypoints, both shown below: + * compile_bundle(...) -- an FNNX pipeline bundle, one C entrypoint per node plus the + pipeline glue that runs them in order + * compile_onnx(...) -- a bare `.onnx` model, no FNNX packaging involved +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import onnx +import torch +from torch import nn + +from fnnx.extras.compilers.c import compile_bundle, compile_onnx +from fnnx.runtime import Runtime + +OUT_DIR = Path.cwd() / "_fnnx_c_demo" +FEATURES = 3 +OPSET = 17 +MAX_BATCH = 8 +SEED = 20260727 + + +# -------------------------------------------------------------------------------------- +# 1. The torch model, as two stages +# -------------------------------------------------------------------------------------- + + +class Standardize(nn.Module): + """`(x - mean) / std`, with the statistics baked in as buffers.""" + + mean: torch.Tensor + std: torch.Tensor + + def __init__(self, mean: torch.Tensor, std: torch.Tensor) -> None: + super().__init__() + self.register_buffer("mean", mean) + self.register_buffer("std", std) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.mean) / self.std + + +class MLP(nn.Module): + def __init__(self, features: int, hidden: int) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Linear(features, hidden), nn.ReLU(), nn.Linear(hidden, 1) + ) + + def forward(self, z: torch.Tensor) -> torch.Tensor: + return self.net(z) + + +def build_stages() -> tuple[Standardize, MLP]: + torch.manual_seed(SEED) + scale = Standardize( + mean=torch.tensor([0.5, -1.0, 2.0]), std=torch.tensor([1.5, 0.25, 3.0]) + ) + return scale.eval(), MLP(FEATURES, hidden=8).eval() + + +# -------------------------------------------------------------------------------------- +# 2. torch -> ONNX +# -------------------------------------------------------------------------------------- + + +def export_onnx( + module: nn.Module, path: Path, *, input_name: str, output_name: str +) -> Path: + """Export one stage with a symbolic batch axis named `batch`. + + The dimension name matters downstream: the compiler either fixes `batch` at compile + time or turns it into a per-call argument, and it is addressed by this name. + """ + path.parent.mkdir(parents=True, exist_ok=True) + torch.onnx.export( + module, + (torch.zeros(1, FEATURES),), + str(path), + input_names=[input_name], + output_names=[output_name], + dynamic_axes={input_name: {0: "batch"}, output_name: {0: "batch"}}, + opset_version=OPSET, + ) + return path + + +# -------------------------------------------------------------------------------------- +# 3. ONNX -> FNNX pipeline bundle +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Stage: + """One pipeline node: an ONNX model, the edges it reads and writes, and its shapes.""" + + id: str + model_path: Path + inputs: tuple[str, ...] + outputs: tuple[str, ...] + input_shapes: tuple[tuple[int | str, ...], ...] + output_shapes: tuple[tuple[int | str, ...], ...] + + +def _tensor_spec(shape: tuple[int | str, ...]) -> dict[str, Any]: + return {"dtype": "Array[float32]", "shape": list(shape)} + + +def _manifest_tensor(name: str, shape: tuple[int | str, ...]) -> dict[str, Any]: + return {"name": name, "content_type": "NDJSON", **_tensor_spec(shape)} + + +def _op_attributes(model_path: Path) -> dict[str, Any]: + """Read back what the exporter actually produced, rather than restating it.""" + model = onnx.load(str(model_path), load_external_data=False) + return { + "opsets": [ + {"domain": opset.domain or "ai.onnx", "version": opset.version} + for opset in model.opset_import + ], + "requires_ort_extensions": False, + "has_external_data": False, + "onnx_ir_version": model.ir_version, + } + + +def write_bundle( + directory: Path, + stages: list[Stage], + *, + name: str, + inputs: list[dict[str, Any]], + outputs: list[dict[str, Any]], +) -> Path: + """Write an unpacked FNNX `pipeline` bundle: the manifest, the ops, and the wiring.""" + shutil.rmtree(directory, ignore_errors=True) + directory.mkdir(parents=True) + + manifest = { + "variant": "pipeline", + "name": name, + "version": "1.0.0", + "description": "Standardize features, then score them with a small MLP.", + "producer_name": "fnnx-examples", + "producer_version": "1.0.0", + "producer_tags": ["torch", "demo"], + "inputs": inputs, + "outputs": outputs, + "dynamic_attributes": [], + "env_vars": [], + } + ops = [ + { + "id": stage.id, + "op": "ONNX_v1", + "inputs": [_tensor_spec(shape) for shape in stage.input_shapes], + "outputs": [_tensor_spec(shape) for shape in stage.output_shapes], + "attributes": _op_attributes(stage.model_path), + "dynamic_attributes": {}, + } + for stage in stages + ] + variant_config = { + "nodes": [ + { + "op_instance_id": stage.id, + "inputs": list(stage.inputs), + "outputs": list(stage.outputs), + "extra_dynattrs": {}, + } + for stage in stages + ] + } + + for filename, document in ( + ("manifest.json", manifest), + ("ops.json", ops), + ("variant_config.json", variant_config), + ("dtypes.json", {}), + ("env.json", {}), + ("meta.json", []), + ): + (directory / filename).write_text(json.dumps(document, indent=2) + "\n") + + for stage in stages: + artifacts = directory / "ops_artifacts" / stage.id + artifacts.mkdir(parents=True) + shutil.copyfile(stage.model_path, artifacts / "model.onnx") + return directory + + +# -------------------------------------------------------------------------------------- +# 4. Compile, build, run +# -------------------------------------------------------------------------------------- + + +def print_report(result: Any) -> None: + report = result.report + memory = report["memory"] + runtime_dims = ( + ", ".join(f"{dim['name']}<={dim['max']}" for dim in report["runtime_dims"]) + or "none" + ) + nodes = ", ".join(f"{node['symbol']}()" for node in report["nodes"]) or "none" + print( + f" header: {result.header_path} ({result.header_path.stat().st_size} B)" + ) + print(f" report: {result.report_path}") + print(f" entrypoint: {report['entrypoint']['symbol']}()") + print(f" node entries: {nodes}") + print(f" opsets: {report['opsets']}") + print(f" fixed dims: {report['dim_bindings'] or 'none'}") + print(f" runtime dims: {runtime_dims}") + print(f" kernels: {len(report['kernels'])} ({', '.join(report['kernels'])})") + print( + f" static mem: {memory['static_bytes']} B " + f"(weights {memory['weights_bytes']}, arena {memory['arena_bytes']})" + ) + + +def compare(label: str, actual: np.ndarray, expected: np.ndarray) -> None: + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-6) + print(f" {label:<28} max |diff| = {np.abs(actual - expected).max():.3e} OK") + + +# -------------------------------------------------------------------------------------- +# 5. The artifact as a C program: what the header is actually for +# -------------------------------------------------------------------------------------- + +C_MAIN = """\ +#define SCORING_IMPLEMENTATION +#include "scoring.h" +#include + +int main(void) +{{ + /* Buffers are sized by the macros the header states, i.e. for the maximum + batch; this call passes {batch} rows and only those are read and written. */ + float x[SCORING_INPUT_X_COUNT] = {{{values}}}; + float score[SCORING_OUTPUT_SCORE_COUNT]; + + if (scoring_run({batch}, x, score) != SCORING_OK) {{ + return 1; + }} + for (int row = 0; row < {batch}; ++row) {{ + printf("%.7f\\n", score[row]); + }} + return 0; +}} +""" + + +def run_as_c_program(directory: Path, x: np.ndarray) -> np.ndarray: + """Compile a C program against the emitted header and return what it prints. + + Nothing of FNNX is involved past this point: the header, a C compiler, and libm. + """ + values = ", ".join(f"{value:.9g}f" for value in x.ravel()) + source = directory / "main.c" + binary = directory / "main" + source.write_text(C_MAIN.format(batch=x.shape[0], values=values)) + subprocess.run( + [ + "cc", + "-std=c99", + "-Wall", + "-Wextra", + "-Werror", + "-O2", + str(source), + "-o", + str(binary), + "-lm", + ], + check=True, + ) + output = subprocess.run([str(binary)], check=True, capture_output=True, text=True) + return np.array([[float(line)] for line in output.stdout.split()], dtype=np.float32) + + +def main() -> None: + scale, mlp = build_stages() + onnx_dir = OUT_DIR / "onnx" + scale_path = export_onnx( + scale, onnx_dir / "scale.onnx", input_name="x", output_name="z" + ) + mlp_path = export_onnx(mlp, onnx_dir / "mlp.onnx", input_name="z", output_name="s") + print(f"exported {scale_path} and {mlp_path}") + + row = ("batch", FEATURES) + score = ("batch", 1) + bundle = write_bundle( + OUT_DIR / "scoring.fnnx", + [ + Stage("scale", scale_path, ("x",), ("z",), (row,), (row,)), + Stage("mlp", mlp_path, ("z",), ("score",), (row,), (score,)), + ], + name="scoring", + inputs=[_manifest_tensor("x", row)], + outputs=[_manifest_tensor("score", score)], + ) + print(f"wrote FNNX pipeline bundle {bundle}") + + x = np.random.default_rng(SEED).normal(size=(4, FEATURES)).astype(np.float32) + with torch.no_grad(): + torch_score = mlp(scale(torch.from_numpy(x))).numpy() + runtime_score = Runtime(str(bundle)).compute({"x": x}, {})["score"] + + # `batch` stays a per-call argument: buffers are sized for MAX_BATCH and every + # entrypoint takes the actual size. Use `dim_bindings={"batch": 4}` instead to bake a + # single size in; unbound symbolic dimensions default to 1. + print("\n=== compile_bundle: the whole pipeline as one header ===") + result = compile_bundle(bundle, OUT_DIR / "c", runtime_dims={"batch": MAX_BATCH}) + print_report(result) + + # Builds the header under `-std=c99 -Wall -Wextra -Werror` and binds it via ctypes. + compiled = result.load() + print("\n running the compiled artifact:") + compare("pipeline vs torch", compiled.run({"x": x})["score"], torch_score) + compare("pipeline vs fnnx runtime", compiled.run({"x": x})["score"], runtime_score) + # Every node is callable on its own, by the id `ops.json` gives it. + compare( + "node `scale` vs torch", + compiled.run_node("scale", {"x": x})["z"], + scale(torch.from_numpy(x)).detach().numpy(), + ) + other = ( + np.random.default_rng(SEED + 1).normal(size=(7, FEATURES)).astype(np.float32) + ) + with torch.no_grad(): + expected = mlp(scale(torch.from_numpy(other))).numpy() + compare("same artifact, batch of 7", compiled.run({"x": other})["score"], expected) + + # Without a bundle there is no manifest name to take the symbol prefix from, so the + # graph's own name is used unless `prefix` says otherwise. + print("\n=== compile_onnx: one ONNX model, no bundle ===") + onnx_result = compile_onnx( + mlp_path, OUT_DIR / "c_mlp", dim_bindings={"batch": 4}, prefix="mlp" + ) + print_report(onnx_result) + with torch.no_grad(): + expected_mlp = mlp(torch.from_numpy(x)).numpy() + print() + compare("mlp vs torch", onnx_result.load().run({"z": x})["s"], expected_mlp) + + print("\n=== the header on its own: a C program, cc, and libm ===") + compare( + "compiled C binary vs torch", run_as_c_program(OUT_DIR / "c", x), torch_score + ) + + print( + "\nThe same two compilations from the command line:\n" + f" python -m fnnx.extras.compilers.c {bundle} " + f"-o {OUT_DIR / 'c'} --runtime-dim batch={MAX_BATCH}\n" + f" python -m fnnx.extras.compilers.c {mlp_path} " + f"-o {OUT_DIR / 'c_mlp'} --dim batch=4 --prefix mlp" + ) + + +if __name__ == "__main__": + main() diff --git a/src/python/examples/xgboost_to_c.py b/src/python/examples/xgboost_to_c.py new file mode 100644 index 0000000..d9ad844 --- /dev/null +++ b/src/python/examples/xgboost_to_c.py @@ -0,0 +1,459 @@ +"""Hands-on demo: XGBoost model -> ONNX-ML -> self-contained C99 header. + +Trains a small XGBoost regressor and classifier, converts both to ONNX with onnxmltools, +and compiles them with ``fnnx.extras.compilers.c`` into single C headers -- gradient +boosted trees as straight-line C, with no runtime, no allocation and no ONNX at inference +time. Each artifact is then built into a shared library and checked against XGBoost's own +predictions, and the classifier is additionally driven from a plain C program. + +Usage (from src/python; needs xgboost, onnxmltools, onnx, numpy and a C compiler): + + pip install onnxmltools + python examples/xgboost_to_c.py + +Everything is written under ./_fnnx_xgb_demo/: + + _fnnx_xgb_demo/reg/xgb_reg.h the regressor, `xgb_reg_run()` + _fnnx_xgb_demo/reg5/xgb_reg5.h the same regressor, re-encoded for ai.onnx.ml 5 + _fnnx_xgb_demo/clf/xgb_clf.h the classifier, `xgb_clf_run()` + _fnnx_xgb_demo/clf/main.c a C program that includes it, built with cc + +Both models compile to `ai.onnx.ml` kernels: `TreeEnsembleRegressor` and +`TreeEnsembleClassifier` become a table-driven tree walk, plus the post-transform the +classifier needs. The converters still emit the opset-1 encoding, so section 3 rewrites the +regressor into the opset-5 `TreeEnsemble` that replaced the pair, and checks that both +compile to the same thing. To package a model as an FNNX bundle instead of a bare `.onnx` +file -- and to get a per-node entrypoint per stage -- see `torch_to_c.py`. +""" + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path +from typing import Any, cast + +import numpy as np +import xgboost +from onnx import ModelProto, NodeProto, TensorProto, helper +from onnx.reference import ReferenceEvaluator +from onnxmltools.convert import convert_xgboost +from onnxmltools.convert.common.data_types import FloatTensorType + +from fnnx.extras.compilers.c import compile_onnx + +ML_DOMAIN = "ai.onnx.ml" + +OUT_DIR = Path.cwd() / "_fnnx_xgb_demo" +FEATURES = 4 +SAMPLES = 400 +TREES = 12 +MAX_BATCH = 32 +SEED = 20260727 + + +# -------------------------------------------------------------------------------------- +# 1. The models +# -------------------------------------------------------------------------------------- + + +def training_data() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + rng = np.random.default_rng(SEED) + x = rng.normal(size=(SAMPLES, FEATURES)).astype(np.float32) + target = 2.0 * x[:, 0] - x[:, 1] + 0.5 * x[:, 2] * x[:, 3] + y = (target + rng.normal(scale=0.1, size=SAMPLES)).astype(np.float32) + return x, y, (y > 0).astype(np.int64) + + +def train() -> tuple[xgboost.XGBRegressor, xgboost.XGBClassifier, np.ndarray]: + x, y, labels = training_data() + settings = {"n_estimators": TREES, "max_depth": 3, "random_state": SEED} + regressor = xgboost.XGBRegressor(**settings).fit(x, y) + classifier = xgboost.XGBClassifier(**settings).fit(x, labels) + return regressor, classifier, x + + +# -------------------------------------------------------------------------------------- +# 2. XGBoost -> ONNX +# -------------------------------------------------------------------------------------- + + +def to_onnx(model: Any) -> ModelProto: + """Convert one fitted XGBoost model, with the row axis exported under a name. + + `FloatTensorType([None, FEATURES])` -- the usual spelling -- leaves that axis + *anonymous*, and the compiler has no name to bind: it pins the axis to 1 and the + artifact serves single rows only. A string dimension is what makes `batch` addressable + by `runtime_dims` / `dim_bindings` further down. + """ + return convert_xgboost( + model, initial_types=[("x", FloatTensorType(["batch", FEATURES]))] + ) + + +def rename_tensor(model: ModelProto, old: str, new: str) -> ModelProto: + """Rename a graph tensor in place; the converter calls every output `variable`.""" + for node in model.graph.node: + node.input[:] = [new if name == old else name for name in node.input] + node.output[:] = [new if name == old else name for name in node.output] + for value in model.graph.output: + if value.name == old: + value.name = new + return model + + +# -------------------------------------------------------------------------------------- +# 3. ai.onnx.ml 1 -> ai.onnx.ml 5, by hand +# -------------------------------------------------------------------------------------- + +# Opset 5 replaced the `TreeEnsembleRegressor`/`TreeEnsembleClassifier` pair with a single +# `TreeEnsemble`, re-encoded around the walk rather than around the trees: +# +# * interior nodes and leaves are two index spaces, not one. The legacy families +# interleave `LEAF` entries among the branch entries and key everything by +# (`nodes_treeids`, `nodes_nodeids`); v5 drops both id families, keeps only interior +# nodes in `nodes_*`, and marks each child with `nodes_trueleafs`/`nodes_falseleafs` +# to say which space its index addresses. +# * `tree_roots` replaces the tree ids: one entry per tree, an index into `nodes_*`. +# * leaf scores move from the (`target_treeids`, `target_nodeids`, `target_ids`, +# `target_weights`) quadruple to `leaf_targetids` + `leaf_weights`, one entry per leaf. +# * `nodes_modes` becomes a uint8 tensor of enumerators instead of strings, and +# `nodes_values` becomes the `nodes_splits` tensor -- typed attributes, so a float64 +# ensemble no longer needs the `*_as_tensor` twins opset 3 had added. +# * `base_values` is gone, as are `classlabels_*`; a classifier is now just a +# `TreeEnsemble` with several targets and whatever post-transform the caller wants. +# +# The compiler reads both, and its own tables are close to the v5 form -- which is why +# `ops/tree.py` converts the legacy encoding into it rather than the other way round. + +_MODES = { + "BRANCH_LEQ": 0, + "BRANCH_LT": 1, + "BRANCH_GTE": 2, + "BRANCH_GT": 3, + "BRANCH_EQ": 4, + "BRANCH_NEQ": 5, +} +_AGGREGATES = {"AVERAGE": 0, "SUM": 1, "MIN": 2, "MAX": 3} +_SUM = 1 +_LEAF = "LEAF" + + +def node_attributes(node: NodeProto) -> dict[str, Any]: + return {entry.name: helper.get_attribute_value(entry) for entry in node.attribute} + + +def to_opset5(model: ModelProto) -> ModelProto: + """Re-encode a legacy `TreeEnsembleRegressor` graph as an opset-5 `TreeEnsemble`. + + Only what this example produces is handled: one regressor node, `NONE` post-transform, + and no set-membership tests (the legacy encoding has none). Anything else raises rather + than emitting an ensemble that scores differently. + """ + legacy = model.graph.node[0] + attributes = node_attributes(legacy) + aggregate = _AGGREGATES[attributes.get("aggregate_function", b"SUM").decode()] + modes = [mode.decode() for mode in attributes["nodes_modes"]] + tree_ids = list(attributes["nodes_treeids"]) + node_ids = list(attributes["nodes_nodeids"]) + missing = list(attributes.get("nodes_missing_value_tracks_true", [0] * len(modes))) + + # The two index spaces, assigned in the order the legacy families list their entries. + branches: dict[tuple[int, int], int] = {} + leaves: dict[tuple[int, int], int] = {} + for index, mode in enumerate(modes): + space = leaves if mode == _LEAF else branches + space[tree_ids[index], node_ids[index]] = len(space) + + def child(tree: int, node_id: int) -> tuple[int, int]: + """The child's index, and whether that index addresses the leaf space.""" + key = (tree, node_id) + return (leaves[key], 1) if key in leaves else (branches[key], 0) + + features: list[int] = [] + node_modes: list[int] = [] + splits: list[float] = [] + true_ids: list[int] = [] + false_ids: list[int] = [] + true_leafs: list[int] = [] + false_leafs: list[int] = [] + tracks: list[int] = [] + for index, mode in enumerate(modes): + if mode == _LEAF: + continue + tree = tree_ids[index] + true_id, true_leaf = child(tree, attributes["nodes_truenodeids"][index]) + false_id, false_leaf = child(tree, attributes["nodes_falsenodeids"][index]) + features.append(attributes["nodes_featureids"][index]) + node_modes.append(_MODES[mode]) + splits.append(attributes["nodes_values"][index]) + true_ids.append(true_id) + false_ids.append(false_id) + true_leafs.append(true_leaf) + false_leafs.append(false_leaf) + tracks.append(missing[index]) + + weights, targets = _leaf_scores(attributes, leaves) + _fold_base_values(attributes, aggregate, leaves, weights, targets) + + roots = [] + for tree in sorted(set(tree_ids)): + # The root is the tree's first entry, which is how the compiler reads it too. + first = tree_ids.index(tree) + if modes[first] == _LEAF: + raise ValueError( + f"Tree {tree} is a single leaf; v5 spells that as a root whose two " + "children are the same leaf, which this example does not emit." + ) + roots.append(branches[tree, node_ids[first]]) + + ensemble = helper.make_node( + "TreeEnsemble", + [legacy.input[0]], + [legacy.output[0]], + domain=ML_DOMAIN, + n_targets=int(attributes["n_targets"]), + aggregate_function=aggregate, + post_transform=0, + tree_roots=roots, + leaf_targetids=targets, + leaf_weights=_float_tensor("leaf_weights", weights), + nodes_featureids=features, + nodes_truenodeids=true_ids, + nodes_falsenodeids=false_ids, + nodes_trueleafs=true_leafs, + nodes_falseleafs=false_leafs, + nodes_missing_value_tracks_true=tracks, + nodes_splits=_float_tensor("nodes_splits", splits), + nodes_modes=helper.make_tensor( + "nodes_modes", TensorProto.UINT8, [len(node_modes)], node_modes + ), + ) + return helper.make_model( + helper.make_graph( + [ensemble], + model.graph.name, + list(model.graph.input), + list(model.graph.output), + ), + opset_imports=[helper.make_opsetid(ML_DOMAIN, 5)], + ) + + +def _leaf_scores( + attributes: dict[str, Any], leaves: dict[tuple[int, int], int] +) -> tuple[list[float], list[int]]: + """`leaf_weights` and `leaf_targetids`, from the legacy `target_*` quadruple.""" + weights = [0.0] * len(leaves) + targets = [0] * len(leaves) + for tree, node_id, target, weight in zip( + attributes["target_treeids"], + attributes["target_nodeids"], + attributes["target_ids"], + attributes["target_weights"], + ): + index = leaves[tree, node_id] + weights[index] = weight + targets[index] = target + return weights, targets + + +def _fold_base_values( + attributes: dict[str, Any], + aggregate: int, + leaves: dict[tuple[int, int], int], + weights: list[float], + targets: list[int], +) -> None: + """Push the dropped `base_values` into the first tree's leaves. + + v5 has no such attribute. Under SUM the fold is exact -- and marginally *more* accurate + than the legacy artifact, which adds the base once at the end in float32 while this + carries it through the accumulation, exactly as XGBoost does. + """ + base = list(attributes.get("base_values", [])) + if not any(base): + return + if aggregate != _SUM: + raise ValueError( + "`base_values` can only be folded into the leaves under SUM aggregation." + ) + first = min(tree for tree, _ in leaves) + for (tree, _), index in leaves.items(): + if tree == first: + weights[index] += base[targets[index]] + + +def _float_tensor(name: str, values: list[float]) -> TensorProto: + return helper.make_tensor(name, TensorProto.FLOAT, [len(values)], values) + + +# -------------------------------------------------------------------------------------- +# 3. Compile, build, run +# -------------------------------------------------------------------------------------- + + +def print_report(result: Any) -> None: + report = result.report + memory = report["memory"] + runtime_dims = ( + ", ".join(f"{dim['name']}<={dim['max']}" for dim in report["runtime_dims"]) + or "none" + ) + signature = ", ".join( + f"{tensor['dtype']}{list(tensor['shape'])} {tensor['name']}" + for tensor in report["entrypoint"]["inputs"] + report["entrypoint"]["outputs"] + ) + print( + f" header: {result.header_path} ({result.header_path.stat().st_size} B)" + ) + print(f" entrypoint: {report['entrypoint']['symbol']}({signature})") + print(f" opsets: {report['opsets']}") + print(f" runtime dims: {runtime_dims}") + print(f" kernels: {', '.join(report['kernels'])}") + print( + f" static mem: {memory['static_bytes']} B " + f"(weights {memory['weights_bytes']}, arena {memory['arena_bytes']})" + ) + + +def compare(label: str, actual: np.ndarray, expected: np.ndarray) -> None: + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-6) + print(f" {label:<30} max |diff| = {np.abs(actual - expected).max():.3e} OK") + + +def compare_exact(label: str, actual: np.ndarray, expected: np.ndarray) -> None: + np.testing.assert_array_equal(actual, expected) + print(f" {label:<30} {len(actual)} rows identical OK") + + +# -------------------------------------------------------------------------------------- +# 4. The classifier as a C program +# -------------------------------------------------------------------------------------- + +C_MAIN = """\ +#define XGB_CLF_IMPLEMENTATION +#include "xgb_clf.h" +#include + +int main(void) +{{ + float x[XGB_CLF_INPUT_X_COUNT] = {{{values}}}; + int64_t label[XGB_CLF_OUTPUT_LABEL_COUNT]; + float probabilities[XGB_CLF_OUTPUT_PROBABILITIES_COUNT]; + + if (xgb_clf_run({rows}, x, label, probabilities) != XGB_CLF_OK) {{ + return 1; + }} + for (int row = 0; row < {rows}; ++row) {{ + printf("%lld %.7f\\n", (long long) label[row], probabilities[row * 2 + 1]); + }} + return 0; +}} +""" + + +def run_as_c_program(directory: Path, x: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Build a C program against the emitted header; returns what it prints per row. + + Nothing of FNNX, ONNX or XGBoost is involved past this point: the header, a C + compiler, and libm for the logistic post-transform. + """ + source = directory / "main.c" + source.write_text( + C_MAIN.format( + rows=x.shape[0], values=", ".join(f"{value:.9g}f" for value in x.ravel()) + ) + ) + binary = directory / "main" + flags = ["-std=c99", "-Wall", "-Wextra", "-Werror", "-O2"] + subprocess.run(["cc", *flags, str(source), "-o", str(binary), "-lm"], check=True) + printed = subprocess.run( + [str(binary)], check=True, capture_output=True, text=True + ).stdout.split("\n") + rows = [line.split() for line in printed if line] + return ( + np.array([int(row[0]) for row in rows], dtype=np.int64), + np.array([float(row[1]) for row in rows], dtype=np.float32), + ) + + +def main() -> None: + shutil.rmtree(OUT_DIR, ignore_errors=True) + regressor, classifier, x = train() + batch = x[:MAX_BATCH] + print(f"trained {TREES} trees per model on {SAMPLES}x{FEATURES} rows") + + print("\n=== XGBRegressor -> TreeEnsembleRegressor -> C ===") + model = rename_tensor(to_onnx(regressor), "variable", "score") + result = compile_onnx( + model, OUT_DIR / "reg", runtime_dims={"batch": MAX_BATCH}, prefix="xgb_reg" + ) + print_report(result) + compiled = result.load() + legacy_scores = compiled.run({"x": batch})["score"].ravel() + print() + compare("compiled vs xgboost", legacy_scores, regressor.predict(batch)) + # The same artifact serves any batch up to the maximum it was compiled for. + compare( + "same artifact, 5 rows", + compiled.run({"x": x[:5]})["score"].ravel(), + regressor.predict(x[:5]), + ) + + print("\n=== the same regressor re-encoded as ai.onnx.ml 5 `TreeEnsemble` ===") + modern = to_opset5(model) + modern_result = compile_onnx( + modern, OUT_DIR / "reg5", runtime_dims={"batch": MAX_BATCH}, prefix="xgb_reg5" + ) + print_report(modern_result) + modern_scores = modern_result.load().run({"x": batch})["score"].ravel() + print() + # The reference evaluator checks the re-encoding itself, not just what the compiler + # makes of it: a v5 graph this compiler mis-read would still have to score the same. + evaluated = cast( + list[np.ndarray], ReferenceEvaluator(modern).run(None, {"x": batch}) + ) + reference = evaluated[0].ravel() + compare("opset 5 vs onnx reference", modern_scores, reference) + compare("opset 5 vs xgboost", modern_scores, regressor.predict(batch)) + compare("opset 5 vs opset 1 artifact", modern_scores, legacy_scores) + print( + f" {'same tables either way':<30} " + f"{modern_result.report['memory']['static_bytes']} B vs " + f"{result.report['memory']['static_bytes']} B, kernel " + f"{modern_result.report['kernels'][0].split('_', 3)[-1]}" + ) + + print("\n=== XGBClassifier -> TreeEnsembleClassifier -> C ===") + result = compile_onnx( + to_onnx(classifier), + OUT_DIR / "clf", + runtime_dims={"batch": MAX_BATCH}, + prefix="xgb_clf", + ) + print_report(result) + compiled = result.load() + outputs = compiled.run({"x": batch}) + print() + compare_exact("labels vs xgboost", outputs["label"], classifier.predict(batch)) + compare( + "probabilities vs xgboost", + outputs["probabilities"], + classifier.predict_proba(batch), + ) + + print("\n=== the classifier header on its own: a C program, cc, and libm ===") + labels, positive = run_as_c_program(OUT_DIR / "clf", batch) + compare_exact("C binary labels vs xgboost", labels, classifier.predict(batch)) + compare("C binary p(class 1) vs xgboost", positive, outputs["probabilities"][:, 1]) + + print( + "\nOn the command line, once the `.onnx` files are on disk:\n" + f" python -m fnnx.extras.compilers.c model.onnx -o {OUT_DIR / 'reg'} " + f"--runtime-dim batch={MAX_BATCH} --prefix xgb_reg" + ) + + +if __name__ == "__main__": + main() diff --git a/src/python/fnnx/extras/compilers/__init__.py b/src/python/fnnx/extras/compilers/__init__.py new file mode 100644 index 0000000..ac3b275 --- /dev/null +++ b/src/python/fnnx/extras/compilers/__init__.py @@ -0,0 +1 @@ +"""Ahead-of-time compilers that turn FNNX bundles into standalone artifacts.""" diff --git a/src/python/fnnx/extras/compilers/c/__init__.py b/src/python/fnnx/extras/compilers/c/__init__.py new file mode 100644 index 0000000..3dad015 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/__init__.py @@ -0,0 +1,39 @@ +"""Compiler from FNNX bundles and ONNX models to a self-contained C99 header.""" + +from importlib import import_module +from typing import TYPE_CHECKING, Any + +from fnnx.extras.compilers.c.errors import CompileError, HarnessError +from fnnx.extras.compilers.c.result import CompileResult + +if TYPE_CHECKING: + from fnnx.extras.compilers.c.bundle import compile_bundle + from fnnx.extras.compilers.c.harness import CompiledModel, load_compiled + from fnnx.extras.compilers.c.onnx.api import compile_onnx + +__all__ = [ + "CompileError", + "CompileResult", + "CompiledModel", + "HarnessError", + "compile_bundle", + "compile_onnx", + "load_compiled", +] + +# Compiling needs the optional `onnx` package and the load-and-run harness needs numpy and +# a system C compiler; importing them lazily keeps the error types importable without +# either, and leaves the missing-dependency messages intact. +_LAZY = { + "compile_bundle": "fnnx.extras.compilers.c.bundle", + "compile_onnx": "fnnx.extras.compilers.c.onnx.api", + "CompiledModel": "fnnx.extras.compilers.c.harness", + "load_compiled": "fnnx.extras.compilers.c.harness", +} + + +def __getattr__(name: str) -> Any: + module = _LAZY.get(name) + if module is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(import_module(module), name) diff --git a/src/python/fnnx/extras/compilers/c/__main__.py b/src/python/fnnx/extras/compilers/c/__main__.py new file mode 100644 index 0000000..79b5fc7 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/__main__.py @@ -0,0 +1,150 @@ +"""Command-line entry: `python -m fnnx.extras.compilers.c -o `.""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from typing import TYPE_CHECKING + +from fnnx.extras.compilers.c.errors import CompileError + +if TYPE_CHECKING: + from fnnx.extras.compilers.c.result import CompileResult + +ONNX_SUFFIX = ".onnx" + + +def main(argv: Sequence[str] | None = None) -> int: + arguments = _parser().parse_args(argv) + try: + result = _compile(arguments) + except (CompileError, OSError) as error: + # OSError covers an unusable `-o`: a path that is a file, or one nothing may + # write to. Nothing the compiler itself raises reaches here as an OSError. + print(f"error: {error}", file=sys.stderr) + return 1 + print("\n".join(_summary(arguments.source, result))) + return 0 + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m fnnx.extras.compilers.c", + description=( + "Compile an FNNX pipeline bundle or an ONNX model into a single " + "self-contained C99 header." + ), + ) + parser.add_argument( + "source", + type=Path, + help="FNNX bundle (directory or tar) or `.onnx` model file to compile", + ) + parser.add_argument( + "-o", + "--output-dir", + type=Path, + required=True, + help="directory the header and the compile report are written to", + ) + parser.add_argument( + "--dim", + action="append", + default=[], + dest="dims", + type=_dim_binding, + metavar="NAME=VALUE", + help=( + "bind a symbolic dimension to a size; repeatable. " + "Dimensions left unbound default to 1" + ), + ) + parser.add_argument( + "--runtime-dim", + action="append", + default=[], + dest="runtime_dims", + type=_dim_binding, + metavar="NAME=MAX", + help=( + "leave a symbolic dimension to be sized per call, up to MAX; repeatable. " + "Buffers are sized for MAX and every entrypoint takes the actual value" + ), + ) + parser.add_argument( + "--prefix", + help=( + "prefix for the emitted files and every public symbol " + "(default: the model's own name)" + ), + ) + return parser + + +def _dim_binding(text: str) -> tuple[str, int]: + name, separator, size = text.partition("=") + if not separator or not name: + raise argparse.ArgumentTypeError(f"expected NAME=VALUE, got `{text}`") + try: + return name, int(size) + except ValueError: + raise argparse.ArgumentTypeError( + f"dimension `{name}` needs an integer size, got `{size}`" + ) from None + + +def _compile(arguments: argparse.Namespace) -> CompileResult: + return _entrypoint(arguments.source)( + arguments.source, + arguments.output_dir, + dim_bindings=dict(arguments.dims), + runtime_dims=dict(arguments.runtime_dims), + prefix=arguments.prefix, + ) + + +def _entrypoint(source: Path) -> Callable[..., CompileResult]: + """The compiler the source asks for: `.onnx` is a model, anything else a bundle.""" + try: + from fnnx.extras.compilers.c import compile_bundle, compile_onnx + except ModuleNotFoundError as error: + # The optional `onnx` dependency, missing; its own message names what to install. + raise CompileError(str(error)) from error + return compile_onnx if source.suffix == ONNX_SUFFIX else compile_bundle + + +def _summary(source: Path, result: CompileResult) -> list[str]: + report = result.report + memory = report["memory"] + fields = [ + ("header", str(result.header_path)), + ("report", str(result.report_path)), + ("entrypoint", f"{report['entrypoint']['symbol']}()"), + ("opsets", _pairs(report["opsets"])), + ("dimensions", _pairs(report["dim_bindings"]) or "none"), + ( + "runtime dimensions", + ", ".join(f"{dim['name']}<={dim['max']}" for dim in report["runtime_dims"]) + or "none", + ), + ("kernels", str(len(report["kernels"]))), + ( + "static memory", + f"{memory['static_bytes']} bytes (weights {memory['weights_bytes']}, " + f"arena {memory['arena_bytes']})", + ), + ] + width = max(len(label) for label, _ in fields) + 2 + return [f"Compiled `{source}`:"] + [ + f" {label + ':':<{width}}{value}" for label, value in fields + ] + + +def _pairs(mapping: Mapping[str, int]) -> str: + return ", ".join(f"{name}={value}" for name, value in mapping.items()) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/python/fnnx/extras/compilers/c/bundle.py b/src/python/fnnx/extras/compilers/c/bundle.py new file mode 100644 index 0000000..9ba48ba --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/bundle.py @@ -0,0 +1,926 @@ +"""The FNNX bundle layer: reading a pipeline bundle and compiling it to one header. + +The ONNX core knows nothing about FNNX. This module holds everything that does: reading and +validating the bundle, handing each node to a node compiler, and emitting the pipeline glue +that calls the compiled nodes in order. +""" + +from __future__ import annotations + +import json +import os +import re +import shutil +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from onnx import ModelProto, TypeProto, ValueInfoProto, helper + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.api import write_artifact +from fnnx.extras.compilers.c.onnx.codegen import ( + DEFAULT_PREFIX, + ArtifactScope, + IOTensor, + NodeEntry, + Program, + StaticBuffer, + build_program, + reserve_dim_parameters, +) +from fnnx.extras.compilers.c.onnx.dtypes import C_TYPES, numpy_dtype_name +from fnnx.extras.compilers.c.onnx.emit import UniqueNames, sanitize_identifier +from fnnx.extras.compilers.c.onnx.frontend import prepare_model +from fnnx.extras.compilers.c.onnx.kernels import CFunction +from fnnx.extras.compilers.c.onnx.loader import load_model +from fnnx.extras.compilers.c.onnx.runtime_dims import RuntimeDim, resolve_runtime_dims +from fnnx.extras.compilers.c.onnx.shapes import ( + UNBOUND_DIM_DEFAULT, + drop_shadowed_inputs, +) +from fnnx.extras.compilers.c.onnx.specialize import specialize +from fnnx.extras.compilers.c.result import CompileResult +from fnnx.handlers._common import unpack_model +from fnnx.validators.model_schema import ( + validate_manifest, + validate_op_instances, + validate_variant, +) + +PIPELINE_VARIANT = "pipeline" +ONNX_OP = "ONNX_v1" + +MANIFEST_FILE = "manifest.json" +OPS_FILE = "ops.json" +VARIANT_FILE = "variant_config.json" +ARTIFACTS_DIR = "ops_artifacts" +ONNX_MODEL_FILE = "model.onnx" + +NDJSON_CONTENT_TYPE = "NDJSON" + +_ARRAY_DTYPE = re.compile(r"^Array\[(.+)\]$") + +# The FNNX element names the artifact has a C type for. An `Array[...]` of anything else — +# a float16, a runtime string — is a compile error naming the entry that asks for it. +ELEMENT_TYPES = {numpy_dtype_name(elem_type): elem_type for elem_type in C_TYPES} + + +def compile_bundle( + bundle_path: str | os.PathLike[str], + output_dir: str | os.PathLike[str], + *, + dim_bindings: Mapping[str, int] | None = None, + runtime_dims: Mapping[str, int] | None = None, + prefix: str | None = None, +) -> CompileResult: + """Compile an FNNX pipeline bundle into a single self-contained C99 header. + + `bundle_path` is a bundle directory or a tar-packaged bundle, the two forms the runtime + accepts. `prefix` defaults to the manifest's name, sanitized to a C identifier. + `runtime_dims` maps a symbolic dimension to the largest size the artifact must serve, + leaving the actual size to each call. Compilation is all-or-nothing: nothing is written + unless the whole bundle compiles. + """ + dims = resolve_runtime_dims(runtime_dims, dim_bindings) + source = Path(bundle_path) + if not source.exists(): + raise CompileError(f"FNNX bundle not found: `{source}`.") + directory, temporary = _unpack(source) + try: + bundle = read_bundle(directory, source_name=source.name) + + def build(bindings: Mapping[str, int]) -> Program: + return _PipelineBuilder(bundle, dict(bindings), prefix, dims).build() + + program = ( + specialize(build, dims, dim_bindings=dim_bindings or {}) + if dims + else build(dim_bindings or {}) + ) + finally: + if temporary: + shutil.rmtree(directory, ignore_errors=True) + return write_artifact( + program, + output_dir, + options={ + "prefix": prefix, + "dim_bindings": dict(sorted((dim_bindings or {}).items())), + "runtime_dims": dict((runtime_dims or {}).items()), + }, + ) + + +# -------------------------------------------------------------------------------------- +# The bundle, as the compiler needs it +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class IOSpec: + """One tensor of an FNNX op spec: its element type and its (partly symbolic) shape.""" + + elem_type: int + shape: tuple[int | str, ...] + + def bind(self, dim_bindings: Mapping[str, int]) -> tuple[int, ...]: + return tuple( + size + if isinstance(size, int) + else dim_bindings.get(size, UNBOUND_DIM_DEFAULT) + for size in self.shape + ) + + +@dataclass(frozen=True) +class ManifestTensor: + """A pipeline input or output as the manifest declares it. + + An empty `shape` leaves the extent unconstrained: it then follows from the nodes the + tensor is wired to, and a non-empty one has to agree with what those nodes say. + """ + + name: str + elem_type: int + shape: tuple[int | str, ...] + + +@dataclass(frozen=True) +class OpInstance: + """One entry of `ops.json`, with the directory holding its artifacts.""" + + id: str + op: str + inputs: tuple[IOSpec, ...] + outputs: tuple[IOSpec, ...] + attributes: Mapping[str, Any] + artifact_dir: Path + + +@dataclass(frozen=True) +class PipelineNode: + """One `variant_config` node: the op instance it runs and the edges it is wired to.""" + + instance: OpInstance + inputs: tuple[str, ...] + outputs: tuple[str, ...] + + +@dataclass(frozen=True) +class Bundle: + """A pipeline bundle, read. + + `name` is the manifest's own name, which the artifact's prefix defaults to and which is + empty for the many bundles that carry none; `label` is what the header calls the bundle + when it says where it came from, and falls back to the file it was read from. + """ + + name: str + label: str + inputs: tuple[ManifestTensor, ...] + outputs: tuple[ManifestTensor, ...] + nodes: tuple[PipelineNode, ...] + + +def read_bundle(directory: Path, *, source_name: str) -> Bundle: + """Read an unpacked bundle directory, validating it against the C compiler's contract.""" + manifest = _read_json(directory / MANIFEST_FILE) + ops = _read_json(directory / OPS_FILE) + variant_config = _read_json(directory / VARIANT_FILE) + + _validate(validate_manifest, manifest, MANIFEST_FILE) + variant = manifest.get("variant") + if variant != PIPELINE_VARIANT: + raise CompileError( + f"The C compiler compiles `{PIPELINE_VARIANT}` bundles; this one is " + f"`{variant}`." + ) + _validate( + lambda config: validate_variant(variant, config), variant_config, VARIANT_FILE + ) + _reject_unknown_ops(ops) + _validate(validate_op_instances, ops, OPS_FILE) + _reject_dynamic_attributes(manifest, ops, variant_config) + + instances = { + instance["id"]: _read_op_instance(instance, directory) for instance in ops + } + name = manifest.get("name") or "" + return Bundle( + name=name, + label=name or source_name, + inputs=_read_manifest_tensors(manifest["inputs"], "input"), + outputs=_read_manifest_tensors(manifest["outputs"], "output"), + nodes=tuple( + _read_pipeline_node(node, instances) for node in variant_config["nodes"] + ), + ) + + +def _read_manifest_tensors( + entries: Sequence[Mapping[str, Any]], role: str +) -> tuple[ManifestTensor, ...]: + tensors = [] + seen: set[str] = set() + for entry in entries: + label = f"Manifest {role} `{entry['name']}`" + content_type = entry.get("content_type") + if content_type != NDJSON_CONTENT_TYPE: + raise CompileError( + f"{label} has content type `{content_type}`; the C compiler compiles only " + f"`{NDJSON_CONTENT_TYPE}` tensors." + ) + if entry["name"] in seen: + # The entrypoint takes one parameter per manifest tensor, and the edge of that + # name can only be wired to one of them; the other would be a parameter nothing + # reads, which the artifact's own `-Werror` build contract refuses. + raise CompileError( + f"{label} is declared twice; every {role} needs its own name." + ) + seen.add(entry["name"]) + tensors.append( + ManifestTensor( + name=entry["name"], + elem_type=_element_type(entry["dtype"], label), + shape=_read_shape(entry["shape"], label), + ) + ) + return tuple(tensors) + + +def _read_op_instance(entry: Mapping[str, Any], directory: Path) -> OpInstance: + label = f"Op instance `{entry['id']}`" + return OpInstance( + id=entry["id"], + op=entry["op"], + inputs=_read_specs(entry["inputs"], f"{label} input"), + outputs=_read_specs(entry["outputs"], f"{label} output"), + attributes=entry["attributes"], + artifact_dir=directory / ARTIFACTS_DIR / entry["id"], + ) + + +def _read_specs(entries: Sequence[Mapping[str, Any]], role: str) -> tuple[IOSpec, ...]: + return tuple( + IOSpec( + elem_type=_element_type(entry["dtype"], f"{role} {index}"), + shape=_read_shape(entry["shape"], f"{role} {index}"), + ) + for index, entry in enumerate(entries) + ) + + +def _read_pipeline_node( + node: Mapping[str, Any], instances: Mapping[str, OpInstance] +) -> PipelineNode: + instance = instances.get(node["op_instance_id"]) + if instance is None: + raise CompileError( + f"Pipeline node references op instance `{node['op_instance_id']}`, which " + f"`{OPS_FILE}` does not define." + ) + inputs = tuple(node["inputs"]) + outputs = tuple(node["outputs"]) + for role, wired, specs in ( + ("input", inputs, instance.inputs), + ("output", outputs, instance.outputs), + ): + if len(wired) != len(specs): + raise CompileError( + f"Pipeline node `{instance.id}` is wired to {len(wired)} {role}(s), but " + f"its op spec declares {len(specs)}." + ) + return PipelineNode(instance=instance, inputs=inputs, outputs=outputs) + + +def _read_shape(shape: Sequence[Any], label: str) -> tuple[int | str, ...]: + dims: list[int | str] = [] + for size in shape: + if isinstance(size, int) and not isinstance(size, bool) and size >= 0: + dims.append(size) + elif isinstance(size, str): + dims.append(size) + else: + raise CompileError( + f"{label} has shape entry {size!r}, which is neither a non-negative " + "dimension size nor a dimension name." + ) + return tuple(dims) + + +def _element_type(dtype: str, label: str) -> int: + match = _ARRAY_DTYPE.match(dtype) + if match is None: + raise CompileError( + f"{label} has dtype `{dtype}`; the C compiler compiles only `Array[...]` " + "tensors." + ) + elem_type = ELEMENT_TYPES.get(match.group(1)) + if elem_type is None: + raise CompileError( + f"{label} has element type `{match.group(1)}`, which the C compiler does not " + f"support; supported types are {', '.join(sorted(ELEMENT_TYPES))}." + ) + return elem_type + + +def _reject_unknown_ops(ops: Any) -> None: + """Dispatch every op instance through the node-compiler registry, before anything else + reads `ops.json`: an op the C compiler has no compiler for is the more useful error.""" + if not isinstance(ops, list) or any(not isinstance(entry, dict) for entry in ops): + raise CompileError(f"Bundle `{OPS_FILE}` must hold a list of op instances.") + for instance in ops: + if instance.get("op") not in NODE_COMPILERS: + raise CompileError( + f"Op instance `{instance.get('id')}` runs op `{instance.get('op')}`, " + "which the C compiler has no node compiler for; it compiles " + f"{', '.join(f'`{name}`' for name in NODE_COMPILERS)}." + ) + + +def _reject_dynamic_attributes( + manifest: Mapping[str, Any], + ops: Sequence[Mapping[str, Any]], + variant_config: Mapping[str, Any], +) -> None: + """Refuse a bundle that takes attribute values per call. + + Everything the artifact does is fixed at compile time, so a dynamic attribute cannot be + honoured; compiling as if it were absent would silently ignore what the caller passes. + """ + declared: list[tuple[str, Iterable[Any]]] = [ + ("the manifest", [entry["name"] for entry in manifest["dynamic_attributes"]]) + ] + declared += [ + (f"op instance `{instance['id']}`", instance["dynamic_attributes"]) + for instance in ops + ] + declared += [ + (f"pipeline node `{node['op_instance_id']}`", node["extra_dynattrs"]) + for node in variant_config["nodes"] + ] + for owner, names in declared: + listed = ", ".join(f"`{name}`" for name in sorted(names)) + if listed: + raise CompileError( + f"The C compiler does not support dynamic attributes, and {owner} " + f"declares {listed}." + ) + + +def _read_json(path: Path) -> Any: + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError) as error: + raise CompileError( + f"Could not read `{path.name}` from the bundle: {error}" + ) from error + + +def _validate(validator: Callable[[Any], None], document: Any, filename: str) -> None: + try: + validator(document) + except Exception as error: + raise CompileError(f"Bundle `{filename}` is not valid: {error}") from error + + +def _unpack(source: Path) -> tuple[Path, bool]: + try: + directory, temporary = unpack_model(os.fspath(source)) + except Exception as error: + raise CompileError(f"Could not open FNNX bundle `{source}`: {error}") from error + return Path(directory), temporary + + +# -------------------------------------------------------------------------------------- +# Node compilers +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class NodeCompilation: + """What a node compiler is handed, and the scope it emits its definitions into. + + A compiler returns a `Program` whose inputs, outputs and body are the node's entry + function — one pointer parameter per op-spec tensor, in spec order — and whose weights, + scratch and functions are the definitions that entry needs. Pipeline codegen depends on + nothing else, so a future FNNX op type plugs in here. + """ + + instance: OpInstance + prefix: str + scope: ArtifactScope + dim_bindings: Mapping[str, int] + runtime_dims: tuple[RuntimeDim, ...] = () + + +NodeCompiler = Callable[[NodeCompilation], Program] + + +def compile_onnx_node(compilation: NodeCompilation) -> Program: + """Compile an `ONNX_v1` node: its `model.onnx`, specialized to its op spec's shapes. + + Everything that can fail runs under one handler, so a failure anywhere — in the file + layout, in the ONNX core, in the agreement between spec and graph — names the op + instance the pipeline knows the node by. + """ + instance = compilation.instance + try: + program = _compile_onnx_model(compilation) + except CompileError as error: + raise CompileError(f"Op instance `{instance.id}`: {error}") from error + return program + + +def _compile_onnx_model(compilation: NodeCompilation) -> Program: + instance = compilation.instance + if instance.attributes.get("requires_ort_extensions"): + raise CompileError( + "the op requires the onnxruntime extensions, which the C compiler does not " + "implement." + ) + path = instance.artifact_dir / ONNX_MODEL_FILE + if not path.is_file(): + raise CompileError( + f"`{ONNX_MODEL_FILE}` is missing from `{instance.artifact_dir}`." + ) + loaded = load_model(path) + _declare_spec_types(loaded.model, instance) + prepared = prepare_model(loaded, dim_bindings=compilation.dim_bindings) + program = build_program( + prepared, + prefix=compilation.prefix, + scope=compilation.scope, + runtime_dims=compilation.runtime_dims, + ) + _verify_signature(program, instance, compilation.dim_bindings) + return program + + +NODE_COMPILERS: dict[str, NodeCompiler] = {ONNX_OP: compile_onnx_node} + + +def _declare_spec_types(model: ModelProto, instance: OpInstance) -> None: + """Give the ONNX graph the I/O types the FNNX op spec declares. + + The spec is the contract the pipeline is wired on, and it is where the symbolic + dimension names live — a converter usually leaves the graph's own batch dimension + nameless — so it is what the graph gets specialized to. Whatever the graph states + concretely has to agree with it. + """ + drop_shadowed_inputs(model) + _declare_side(model.graph.input, instance.inputs, "input") + _declare_side(model.graph.output, instance.outputs, "output") + + +def _declare_side( + entries: Sequence[ValueInfoProto], specs: Sequence[IOSpec], role: str +) -> None: + if len(entries) != len(specs): + raise CompileError( + f"its op spec declares {len(specs)} {role}(s), but its ONNX graph has " + f"{len(entries)}." + ) + for index, (entry, spec) in enumerate(zip(entries, specs)): + label = f"{role} {index} (`{entry.name}`)" + _check_declared_type(entry.type, spec, label) + entry.type.CopyFrom( + helper.make_tensor_type_proto(spec.elem_type, list(spec.shape)) + ) + + +def _check_declared_type(declared: TypeProto, spec: IOSpec, label: str) -> None: + kind = declared.WhichOneof("value") + if kind is not None and kind != "tensor_type": + raise CompileError( + f"{label} is a `{kind}` in the ONNX graph, which the C compiler does not " + "support; only tensors can be compiled." + ) + tensor_type = declared.tensor_type + if tensor_type.elem_type and tensor_type.elem_type != spec.elem_type: + raise CompileError( + f"{label} is `{numpy_dtype_name(tensor_type.elem_type)}` in the ONNX graph, " + f"but `{numpy_dtype_name(spec.elem_type)}` in the op spec." + ) + if not tensor_type.HasField("shape"): + return + dims = tensor_type.shape.dim + if len(dims) != len(spec.shape): + raise CompileError( + f"{label} has rank {len(dims)} in the ONNX graph, but rank " + f"{len(spec.shape)} in the op spec." + ) + for axis, (dim, size) in enumerate(zip(dims, spec.shape)): + if dim.WhichOneof("value") != "dim_value" or not isinstance(size, int): + continue + if dim.dim_value != size: + raise CompileError( + f"{label} has size {dim.dim_value} on axis {axis} in the ONNX graph, but " + f"{size} in the op spec." + ) + + +def _verify_signature( + program: Program, instance: OpInstance, dim_bindings: Mapping[str, int] +) -> None: + """Check the compiled entry against the op spec the pipeline wires the node by.""" + for role, tensors, specs in ( + ("input", program.inputs, instance.inputs), + ("output", program.outputs, instance.outputs), + ): + if len(tensors) != len(specs): + raise CompileError( + f"its op spec declares {len(specs)} {role}(s), but the compiled graph " + f"takes {len(tensors)}." + ) + for index, (tensor, spec) in enumerate(zip(tensors, specs)): + expected = spec.bind(dim_bindings) + if (tensor.elem_type, tensor.shape) != (spec.elem_type, expected): + raise CompileError( + f"{role} {index} (`{tensor.name}`) compiles to " + f"{numpy_dtype_name(tensor.elem_type)}{list(tensor.shape)}, but its " + f"op spec declares {numpy_dtype_name(spec.elem_type)}" + f"{list(expected)}." + ) + + +# -------------------------------------------------------------------------------------- +# Pipeline codegen +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _Edge: + """Where a pipeline edge's data lives, as a C expression, plus the type it holds.""" + + expr: str + elem_type: int + shape: tuple[int, ...] + + +@dataclass(frozen=True) +class _Port: + """One end of an edge: a node's input or output, and the tensor it compiled to.""" + + label: str + tensor: IOTensor + + +@dataclass +class _PipelineBuilder: + bundle: Bundle + dim_bindings: Mapping[str, int] + requested_prefix: str | None + runtime_dims: tuple[RuntimeDim, ...] = () + + prefix: str = field(init=False) + scope: ArtifactScope = field(init=False) + names: UniqueNames = field(default_factory=UniqueNames, init=False) + produced: dict[str, str] = field(default_factory=dict, init=False) + + def __post_init__(self) -> None: + source = ( + self.requested_prefix + if self.requested_prefix is not None + else self.bundle.name + ) + self.prefix = sanitize_identifier(source, fallback=DEFAULT_PREFIX) + self.scope = ArtifactScope(self.prefix, self.names) + reserve_dim_parameters(self.names, self.runtime_dims) + + def build(self) -> Program: + # Handed out before any node is compiled, so that the pipeline entrypoint's + # parameters read as the manifest names rather than as whatever a node claimed. + parameters = [ + self.names.assign(tensor.name, fallback=role) + for role, tensors in ( + ("input", self.bundle.inputs), + ("output", self.bundle.outputs), + ) + for tensor in tensors + ] + ordered = _order_nodes(self.bundle) + programs = self._compile_nodes(ordered) + entries = tuple( + NodeEntry( + id=instance_id, + symbol=f"{program.prefix}_run", + inputs=program.inputs, + outputs=program.outputs, + body=program.body, + body_owners=program.body_owners, + ) + for instance_id, program in programs.items() + ) + + ports = _collect_ports(ordered, programs) + self.produced = { + name: node.instance.id for node in ordered for name in node.outputs + } + split = len(self.bundle.inputs) + inputs = self._pipeline_tensors( + self.bundle.inputs, parameters[:split], "INPUT", ports + ) + outputs = self._pipeline_tensors( + self.bundle.outputs, parameters[split:], "OUTPUT", ports + ) + edges, buffers = self._plan_edges(ordered, ports, (*inputs, *outputs)) + body, body_owners = self._emit_body(ordered, entries, edges, inputs) + return Program( + prefix=self.prefix, + graph_name=self.bundle.label, + source=f"FNNX bundle `{self.bundle.label}`", + opsets=_merged_opsets(programs.values()), + dim_bindings=_merged_bindings(programs.values()), + inputs=inputs, + outputs=outputs, + weights=_merged_buffers(program.weights for program in programs.values()), + scratch=_merged_buffers( + [buffers, *(program.scratch for program in programs.values())] + ), + functions=_merged_functions(programs.values()), + body=body, + labels=tuple( + table for program in programs.values() for table in program.labels + ), + nodes=entries, + runtime_dims=self.runtime_dims, + body_owners=body_owners, + ) + + def _compile_nodes(self, ordered: Sequence[PipelineNode]) -> dict[str, Program]: + """Compile each op instance once, in the order the pipeline runs them. + + Two pipeline nodes may run the same op instance; they then share one entrypoint, + called twice on different buffers. + """ + programs: dict[str, Program] = {} + for node in ordered: + instance = node.instance + if instance.id in programs: + continue + prefix = self.names.assign( + f"{self.prefix}_node_{instance.id}", fallback=f"{self.prefix}_node" + ) + programs[instance.id] = NODE_COMPILERS[instance.op]( + NodeCompilation( + instance=instance, + prefix=prefix, + scope=self.scope, + dim_bindings=self.dim_bindings, + runtime_dims=self.runtime_dims, + ) + ) + return programs + + def _pipeline_tensors( + self, + declared: Sequence[ManifestTensor], + parameters: Sequence[str], + role: str, + ports: Mapping[str, list[_Port]], + ) -> tuple[IOTensor, ...]: + """The manifest tensors as the artifact exposes them, sized from the nodes they + are wired to and checked against whatever the manifest itself declares.""" + tensors = [] + for tensor, c_name in zip(declared, parameters): + elem_type, shape = self._derived_type(tensor, role, ports) + if elem_type != tensor.elem_type: + raise CompileError( + f"Manifest {role.lower()} `{tensor.name}` is declared " + f"`Array[{numpy_dtype_name(tensor.elem_type)}]`, but the node it is " + f"wired to takes `Array[{numpy_dtype_name(elem_type)}]`." + ) + if tensor.shape: + expected = IOSpec(elem_type, tensor.shape).bind(self.dim_bindings) + if expected != shape: + raise CompileError( + f"Manifest {role.lower()} `{tensor.name}` declares shape " + f"{list(expected)}, but the node it is wired to takes " + f"{list(shape)}." + ) + tensors.append( + IOTensor( + name=tensor.name, + c_name=c_name, + macro=f"{self.prefix.upper()}_{role}_{c_name.upper()}", + elem_type=elem_type, + shape=shape, + owner=self.produced.get(tensor.name, ""), + ) + ) + return tuple(tensors) + + def _derived_type( + self, declared: ManifestTensor, role: str, ports: Mapping[str, list[_Port]] + ) -> tuple[int, tuple[int, ...]]: + # An output only counts as wired where a node *writes* it: one that merely appears + # as some node's input is a name the caller's output buffer would take over, and + # the node would then read the buffer it was supposed to be read into. + if role == "OUTPUT" and declared.name not in self.produced: + raise CompileError( + f"Manifest output `{declared.name}` is produced by no pipeline node." + ) + used = ports.get(declared.name) + if used: + tensor = _agreed_tensor(declared.name, used) + return tensor.elem_type, tensor.shape + if not declared.shape: + raise CompileError( + f"Manifest input `{declared.name}` is read by no pipeline node and " + "declares no shape, so the buffer it needs cannot be sized." + ) + return declared.elem_type, IOSpec(declared.elem_type, declared.shape).bind( + self.dim_bindings + ) + + def _plan_edges( + self, + ordered: Sequence[PipelineNode], + ports: Mapping[str, list[_Port]], + exposed: Sequence[IOTensor], + ) -> tuple[dict[str, _Edge], tuple[StaticBuffer, ...]]: + """Give every edge a buffer: the caller's, or a static one of the pipeline's own. + + An edge the manifest exposes is the caller's buffer, written there by its producer + and read from there by every downstream node, so a fan-out that is also a pipeline + output needs no copy. + """ + edges = { + tensor.name: _Edge(tensor.c_name, tensor.elem_type, tensor.shape) + for tensor in exposed + } + buffers = [] + for node in ordered: + for name in node.outputs: + if name in edges: + continue + tensor = _agreed_tensor(name, ports[name]) + symbol = self.names.assign( + f"{self.prefix}_e_{name}", fallback=f"{self.prefix}_e" + ) + edges[name] = _Edge(symbol, tensor.elem_type, tensor.shape) + buffers.append( + StaticBuffer(name, symbol, tensor.elem_type, tensor.shape, None) + ) + return edges, tuple(buffers) + + def _emit_body( + self, + ordered: Sequence[PipelineNode], + entries: Sequence[NodeEntry], + edges: Mapping[str, _Edge], + inputs: Sequence[IOTensor], + ) -> tuple[tuple[str, ...], tuple[str, ...]]: + symbols = {entry.id: entry.symbol for entry in entries} + read = {name for node in ordered for name in node.inputs} + dims = [dim.c_name for dim in self.runtime_dims] + body = [ + f"(void){tensor.c_name};" for tensor in inputs if tensor.name not in read + ] + owners = [""] * len(body) + if ordered: + body.append("int status;") + owners.append("") + for node in ordered: + arguments = ", ".join( + dims + [edges[name].expr for name in (*node.inputs, *node.outputs)] + ) + body.append(f"status = {symbols[node.instance.id]}({arguments});") + body.append( + f"if (status != {self.prefix.upper()}_OK) {{\n return status;\n}}" + ) + owners += [node.instance.id] * 2 + return tuple(body), tuple(owners) + + +def _order_nodes(bundle: Bundle) -> tuple[PipelineNode, ...]: + """Topological order of the pipeline nodes, stable in the order they are declared.""" + available = {tensor.name for tensor in bundle.inputs} + producers: dict[str, PipelineNode] = {} + for node in bundle.nodes: + for name in node.outputs: + if name in available: + raise CompileError( + f"Pipeline node `{node.instance.id}` writes `{name}`, which is also a " + "manifest input." + ) + if name in producers: + raise CompileError( + f"Pipeline nodes `{producers[name].instance.id}` and " + f"`{node.instance.id}` both write `{name}`." + ) + producers[name] = node + for node in bundle.nodes: + for name in node.inputs: + if name not in available and name not in producers: + raise CompileError( + f"Pipeline node `{node.instance.id}` reads `{name}`, which no " + "manifest input and no node produces." + ) + + remaining = list(bundle.nodes) + ordered: list[PipelineNode] = [] + while remaining: + for index, node in enumerate(remaining): + if all(name in available for name in node.inputs): + del remaining[index] + ordered.append(node) + available.update(node.outputs) + break + else: + blocked = ", ".join(f"`{node.instance.id}`" for node in remaining) + raise CompileError( + f"The pipeline has a cycle: nodes {blocked} each wait on another's output." + ) + return tuple(ordered) + + +def _collect_ports( + ordered: Sequence[PipelineNode], programs: Mapping[str, Program] +) -> dict[str, list[_Port]]: + """Every use of every edge, so that the ends of an edge can be checked against each + other: one buffer cannot hold two shapes.""" + ports: dict[str, list[_Port]] = {} + for node in ordered: + program = programs[node.instance.id] + for role, names, tensors in ( + ("input", node.inputs, program.inputs), + ("output", node.outputs, program.outputs), + ): + for index, (name, tensor) in enumerate(zip(names, tensors)): + label = f"{role} {index} of node `{node.instance.id}`" + ports.setdefault(name, []).append(_Port(label, tensor)) + return ports + + +def _agreed_tensor(name: str, ports: Sequence[_Port]) -> IOTensor: + first = ports[0] + for other in ports[1:]: + if (other.tensor.elem_type, other.tensor.shape) != ( + first.tensor.elem_type, + first.tensor.shape, + ): + raise CompileError( + f"Pipeline edge `{name}` is {numpy_dtype_name(first.tensor.elem_type)}" + f"{list(first.tensor.shape)} as {first.label}, but " + f"{numpy_dtype_name(other.tensor.elem_type)}" + f"{list(other.tensor.shape)} as {other.label}." + ) + return first.tensor + + +def _merged_opsets(programs: Iterable[Program]) -> dict[str, int]: + """The opset each domain is compiled at, the highest winning where nodes differ. + + Nodes are compiled independently and may import different versions of a domain; the + merged map is reported metadata, not something dispatch reads back. + """ + merged: dict[str, int] = {} + for program in programs: + for domain, version in program.opsets.items(): + merged[domain] = max(merged.get(domain, version), version) + return merged + + +def _merged_bindings(programs: Iterable[Program]) -> dict[str, int]: + merged: dict[str, int] = {} + for program in programs: + merged.update(program.dim_bindings) + return merged + + +def _merged_functions(programs: Iterable[Program]) -> tuple[CFunction, ...]: + """Every kernel the nodes emitted, each once: they are named from the artifact-wide + prefix, so nodes using the same kernel at the same types share one definition.""" + merged: dict[str, CFunction] = {} + for program in programs: + for function in program.functions: + existing = merged.setdefault(function.name, function) + if existing.definition != function.definition: + raise CompileError( + f"Kernel `{function.name}` was emitted twice with different " + "definitions; a kernel name must encode everything its code " + "depends on." + ) + return tuple(merged.values()) + + +def _merged_buffers( + groups: Iterable[Sequence[StaticBuffer]], +) -> tuple[StaticBuffer, ...]: + """Every static buffer the nodes reserved, each once, sized for the largest claim. + + Only kernel scratch is ever claimed twice — its symbol is shared between the nodes that + call the kernel — and sharing it stays safe under the artifact's one-call-at-a-time + contract, exactly as it is between the nodes of a single graph. + """ + merged: dict[str, StaticBuffer] = {} + for group in groups: + for buffer in group: + reserved = merged.get(buffer.symbol) + if reserved is None or reserved.elem_count < buffer.elem_count: + merged[buffer.symbol] = buffer + return tuple(merged.values()) diff --git a/src/python/fnnx/extras/compilers/c/errors.py b/src/python/fnnx/extras/compilers/c/errors.py new file mode 100644 index 0000000..ed204fe --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/errors.py @@ -0,0 +1,6 @@ +class CompileError(Exception): + """Raised when a model cannot be compiled to C.""" + + +class HarnessError(Exception): + """Raised when a compiled artifact cannot be built, loaded, or driven from Python.""" diff --git a/src/python/fnnx/extras/compilers/c/harness.py b/src/python/fnnx/extras/compilers/c/harness.py new file mode 100644 index 0000000..8a1fc35 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/harness.py @@ -0,0 +1,444 @@ +"""Building a compiled artifact into a shared library and driving it from Python. + +This is tooling around the artifact, not part of it: nothing here influences the generated +C. It needs numpy and a system C compiler, and nothing else — in particular not `onnx`, +so an artifact can be exercised wherever it was copied to. +""" + +from __future__ import annotations + +import ctypes +import json +import os +import shutil +import subprocess +import tempfile +import weakref +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from importlib.util import find_spec +from pathlib import Path +from typing import Any + +from fnnx.extras.compilers.c.errors import HarnessError + +if find_spec("numpy") is None: + # ModuleNotFoundError rather than a bare ImportError so that callers (and + # `pytest.importorskip`) can tell a missing optional dependency from a broken one. + raise ModuleNotFoundError( + "The FNNX C load-and-run harness requires numpy. " + 'Install it with `pip install "fnnx[core]"`.', + name="numpy", + ) + +import numpy # noqa: E402 + +# The artifact's build contract: what the generated header must compile cleanly under. +STRICT_FLAGS = ("-std=c99", "-Wall", "-Wextra", "-Werror", "-Werror=vla") +SHARED_FLAGS = ("-fPIC", "-shared") + +COMPILER_CANDIDATES = ("cc", "gcc", "clang") + +_REQUIRED_REPORT_FIELDS = ("prefix", "header", "entrypoint") + + +def load_compiled( + path: str | os.PathLike[str], *, compiler: str | None = None +) -> CompiledModel: + """Build a compiled artifact into a shared library and bind its entrypoints. + + `path` is the emitted header or its compile report, which sits beside it; the report is + what drives the binding. `compiler` overrides the detected system C compiler. + """ + report_path = _resolve_report_path(Path(path)) + try: + report = json.loads(report_path.read_text(encoding="utf-8")) + except (OSError, ValueError) as error: + raise HarnessError( + f"Could not read the compile report `{report_path}`: {error}" + ) from error + missing = [field for field in _REQUIRED_REPORT_FIELDS if field not in report] + if missing: + raise HarnessError( + f"The compile report `{report_path}` is missing the " + f"{', '.join(f'`{field}`' for field in missing)} field(s); it does not " + "describe an artifact of this compiler." + ) + + build_dir = Path(tempfile.mkdtemp(prefix="fnnx_c_harness_")) + try: + library_path = _build_shared_library( + report_path.parent, report, build_dir, compiler=compiler + ) + model = CompiledModel(report, library_path) + except BaseException: + _discard(build_dir) + raise + # The library stays mapped once loaded, so the build directory only has to outlive the + # load itself; tying it to the model keeps the artifact directory free of build output. + weakref.finalize(model, _discard, build_dir) + return model + + +@dataclass(frozen=True) +class TensorSpec: + """One tensor of an entrypoint's signature, as the compile report describes it. + + `shape` is the buffer's capacity. Where the artifact has runtime dimensions, `axes` + says how each extent follows from them — `(None, size)` for a fixed axis, and + `(dimension, factor)` for one that scales — so the shape a particular call works at is + `shape_at`. + """ + + name: str + dtype: numpy.dtype + shape: tuple[int, ...] + axes: tuple[tuple[str | None, int], ...] = () + + def shape_at(self, dims: Mapping[str, int]) -> tuple[int, ...]: + if not self.axes: + return self.shape + return tuple( + size if dim is None else size * dims[dim] for dim, size in self.axes + ) + + +@dataclass(frozen=True) +class RuntimeDimSpec: + """A dimension the caller sizes per call, and the maximum it was compiled for.""" + + name: str + maximum: int + + +class CompiledModel: + """A compiled artifact, built into a shared library and bound through ctypes. + + Not reentrant, following the artifact's own contract: one in-flight call per model. + """ + + def __init__(self, report: Mapping[str, Any], library_path: Path) -> None: + self.report = dict(report) + self.library_path = library_path + try: + self._library = ctypes.CDLL(str(library_path)) + except OSError as error: + raise HarnessError( + f"Could not load the built shared library `{library_path}`: {error}" + ) from error + self._dims = tuple( + RuntimeDimSpec(str(dim["name"]), int(dim["max"])) + for dim in report.get("runtime_dims", ()) + ) + self._entry = self._bind(report["entrypoint"], "The compiled model") + self._nodes = { + str(node["id"]): self._bind(node, f"Node `{node['id']}`") + for node in report.get("nodes", ()) + } + + @property + def inputs(self) -> tuple[TensorSpec, ...]: + return self._entry.inputs + + @property + def outputs(self) -> tuple[TensorSpec, ...]: + return self._entry.outputs + + @property + def node_ids(self) -> tuple[str, ...]: + return tuple(self._nodes) + + @property + def runtime_dims(self) -> tuple[RuntimeDimSpec, ...]: + return self._dims + + def run( + self, + inputs: Mapping[str, Any] | None = None, + *, + dims: Mapping[str, int] | None = None, + **named: Any, + ) -> dict[str, numpy.ndarray]: + """Run the whole graph on named numpy arrays, returning its named outputs. + + Inputs may be passed as a mapping — tensor names need not be Python identifiers — + or as keyword arguments, and are validated against the compiled shapes and dtypes + before the C entrypoint is called. `dims` gives the size of each runtime dimension + for this call; one left out is read off the inputs that scale with it. + """ + return self._entry.run(_merge_inputs(inputs, named), dims or {}) + + def run_node( + self, + node_id: str, + inputs: Mapping[str, Any] | None = None, + *, + dims: Mapping[str, int] | None = None, + **named: Any, + ) -> dict[str, numpy.ndarray]: + """Run a single node's entrypoint, by the node id the compile report lists.""" + entry = self._nodes.get(str(node_id)) + if entry is None: + available = ", ".join(f"`{name}`" for name in self._nodes) or "none" + raise HarnessError( + f"The compiled artifact exposes no entrypoint for node `{node_id}`; " + f"it exposes: {available}." + ) + return entry.run(_merge_inputs(inputs, named), dims or {}) + + def _bind(self, description: Mapping[str, Any], label: str) -> _Entrypoint: + symbol = description["symbol"] + try: + function = getattr(self._library, symbol) + except AttributeError: + raise HarnessError( + f"The shared library built from `{self.report['header']}` exports no " + f"symbol `{symbol}`, which the compile report names as an entrypoint." + ) from None + inputs = _tensor_specs(description["inputs"]) + outputs = _tensor_specs(description["outputs"]) + function.restype = ctypes.c_int + function.argtypes = [ctypes.c_int32] * len(self._dims) + [ctypes.c_void_p] * ( + len(inputs) + len(outputs) + ) + return _Entrypoint(label, symbol, inputs, outputs, function, self._dims) + + +@dataclass +class _Entrypoint: + """A bound C entrypoint: the tensors it takes and the callable behind its symbol.""" + + label: str + symbol: str + inputs: tuple[TensorSpec, ...] + outputs: tuple[TensorSpec, ...] + call: Callable[..., int] + dims: tuple[RuntimeDimSpec, ...] = () + + def run( + self, values: Mapping[str, Any], dims: Mapping[str, int] + ) -> dict[str, numpy.ndarray]: + self._check_names(values) + sizes = self._resolve_dims(values, dims) + arguments = tuple( + self._checked(spec, values[spec.name], sizes) for spec in self.inputs + ) + results = { + spec.name: numpy.empty(spec.shape_at(sizes), dtype=spec.dtype) + for spec in self.outputs + } + buffers = (*arguments, *results.values()) + status = self.call( + *[sizes[dim.name] for dim in self.dims], + *[buffer.ctypes.data for buffer in buffers], + ) + if status != 0: + raise HarnessError( + f"{self.label}: `{self.symbol}` returned status {status}." + ) + return results + + def _check_names(self, values: Mapping[str, Any]) -> None: + expected = {spec.name for spec in self.inputs} + missing = sorted(expected - set(values)) + unexpected = sorted(set(values) - expected) + if missing or unexpected: + details = [] + if missing: + details.append(f"missing {_quoted(missing)}") + if unexpected: + details.append(f"unexpected {_quoted(unexpected)}") + raise HarnessError( + f"{self.label}: {' and '.join(details)}; it takes " + f"{_quoted(spec.name for spec in self.inputs) or 'no inputs'}." + ) + + def _resolve_dims( + self, values: Mapping[str, Any], given: Mapping[str, int] + ) -> dict[str, int]: + """The size of every runtime dimension for this call, stated or read off the inputs.""" + unknown = sorted(set(given) - {dim.name for dim in self.dims}) + if unknown: + named = _quoted(dim.name for dim in self.dims) or "none" + raise HarnessError( + f"{self.label}: {_quoted(unknown)} is not a runtime dimension of this " + f"artifact; it has {named}." + ) + sizes = {} + for dim in self.dims: + size = given.get(dim.name) + if size is None: + size = self._infer_dim(dim, values) + if not isinstance(size, (int, numpy.integer)) or isinstance(size, bool): + raise HarnessError( + f"{self.label}: runtime dimension `{dim.name}` needs an integer " + f"size, got {size!r}." + ) + if not 1 <= size <= dim.maximum: + raise HarnessError( + f"{self.label}: runtime dimension `{dim.name}` is {size}, outside " + f"the [1, {dim.maximum}] the artifact was compiled for." + ) + sizes[dim.name] = int(size) + return sizes + + def _infer_dim(self, dim: RuntimeDimSpec, values: Mapping[str, Any]) -> int: + for spec in self.inputs: + for axis, (name, factor) in enumerate(spec.axes): + if name != dim.name: + continue + extent = numpy.shape(values[spec.name]) + if axis >= len(extent): + continue + size, remainder = divmod(extent[axis], factor) + if remainder: + raise HarnessError( + f"{self.label}: input `{spec.name}` is {extent[axis]} long on " + f"axis {axis}, which is not {factor} times a size of runtime " + f"dimension `{dim.name}`." + ) + return size + raise HarnessError( + f"{self.label}: no input's shape depends on runtime dimension " + f"`{dim.name}`, so its size has to be passed as `dims={{'{dim.name}': ...}}`." + ) + + def _checked( + self, spec: TensorSpec, value: Any, dims: Mapping[str, int] + ) -> numpy.ndarray: + array = numpy.asarray(value) + if array.dtype != spec.dtype: + raise HarnessError( + f"{self.label}: input `{spec.name}` has dtype `{array.dtype}`, but the " + f"artifact was compiled for `{spec.dtype}`." + ) + expected = spec.shape_at(dims) + if array.shape != expected: + raise HarnessError( + f"{self.label}: input `{spec.name}` has shape {array.shape}, but the " + f"artifact was compiled for {expected}." + ) + return numpy.ascontiguousarray(array) + + +def _tensor_specs( + descriptions: Sequence[Mapping[str, Any]], +) -> tuple[TensorSpec, ...]: + return tuple( + TensorSpec( + name=description["name"], + dtype=numpy.dtype(description["dtype"]), + shape=tuple(description["shape"]), + axes=_axes(description.get("runtime_shape")), + ) + for description in descriptions + ) + + +def _axes( + runtime_shape: Sequence[Any] | None, +) -> tuple[tuple[str | None, int], ...]: + if not runtime_shape: + return () + return tuple( + (None, int(axis)) + if isinstance(axis, int) + else (str(axis["dim"]), int(axis["coefficient"])) + for axis in runtime_shape + ) + + +def _merge_inputs( + mapping: Mapping[str, Any] | None, named: dict[str, Any] +) -> dict[str, Any]: + values = dict(mapping) if mapping is not None else {} + duplicates = sorted(set(values) & set(named)) + if duplicates: + raise HarnessError( + f"Input(s) {_quoted(duplicates)} were given both in the mapping and as " + "keyword arguments." + ) + values.update(named) + return values + + +def _resolve_report_path(path: Path) -> Path: + if path.suffix == ".json": + report = path + elif path.suffix == ".h": + report = path.with_name(f"{path.stem}_report.json") + else: + raise HarnessError( + f"`{path}` is neither a generated header (`.h`) nor a compile report " + "(`.json`); pass one of the two files a compilation emitted." + ) + if not report.is_file(): + raise HarnessError(f"Compile report not found: `{report}`.") + return report + + +def _build_shared_library( + artifact_dir: Path, + report: Mapping[str, Any], + build_dir: Path, + *, + compiler: str | None, +) -> Path: + header = artifact_dir / report["header"] + if not header.is_file(): + raise HarnessError( + f"The header `{header}` the compile report names is missing." + ) + unit = build_dir / "implementation.c" + unit.write_text( + f"#define {report['prefix'].upper()}_IMPLEMENTATION\n" + f'#include "{header.name}"\n', + encoding="utf-8", + ) + library = build_dir / f"{report['prefix']}.so" + command = [ + _find_compiler(compiler), + *STRICT_FLAGS, + *SHARED_FLAGS, + f"-I{artifact_dir}", + str(unit), + "-o", + str(library), + "-lm", + ] + try: + process = subprocess.run(command, capture_output=True, text=True) + except OSError as error: + raise HarnessError( + f"Could not run the C compiler `{command[0]}`: {error}" + ) from error + if process.returncode != 0: + raise HarnessError( + f"Building `{header.name}` as a shared library failed " + f"(`{' '.join(command)}`):\n{process.stderr.strip()}" + ) + return library + + +def _find_compiler(requested: str | None) -> str: + if requested is not None: + if shutil.which(requested) is None: + raise HarnessError(f"The requested C compiler `{requested}` was not found.") + return requested + candidates = [os.environ.get("CC"), *COMPILER_CANDIDATES] + for candidate in candidates: + if candidate and shutil.which(candidate): + return candidate + raise HarnessError( + "No system C compiler was found; the load-and-run harness needs one of " + f"{', '.join(COMPILER_CANDIDATES)} on PATH, the `CC` environment variable set, " + "or the `compiler` argument." + ) + + +def _discard(directory: Path) -> None: + shutil.rmtree(directory, ignore_errors=True) + + +def _quoted(names: Iterable[str]) -> str: + return ", ".join(f"`{name}`" for name in names) diff --git a/src/python/fnnx/extras/compilers/c/onnx/__init__.py b/src/python/fnnx/extras/compilers/c/onnx/__init__.py new file mode 100644 index 0000000..4f370d6 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/__init__.py @@ -0,0 +1,12 @@ +"""FNNX-agnostic core of the C compiler: ONNX model in, C out.""" + +from importlib.util import find_spec + +if find_spec("onnx") is None: + # ModuleNotFoundError rather than a bare ImportError so that callers (and + # `pytest.importorskip`) can tell a missing optional dependency from a broken one. + raise ModuleNotFoundError( + "The FNNX C compiler requires the `onnx` package. " + 'Install it with `pip install "fnnx[compiler]"`.', + name="onnx", + ) diff --git a/src/python/fnnx/extras/compilers/c/onnx/api.py b/src/python/fnnx/extras/compilers/c/onnx/api.py new file mode 100644 index 0000000..e1505a7 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/api.py @@ -0,0 +1,174 @@ +"""The standalone ONNX-to-C entrypoint and the compile report it writes.""" + +from __future__ import annotations + +import json +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from onnx import ModelProto, TensorProto + +from fnnx import __version__ +from fnnx.extras.compilers.c.onnx.codegen import IOTensor, Program, build_program +from fnnx.extras.compilers.c.onnx.dtypes import c_type, element_size, numpy_dtype_name +from fnnx.extras.compilers.c.onnx.frontend import prepare_model +from fnnx.extras.compilers.c.onnx.header import buffer_bytes, render_header +from fnnx.extras.compilers.c.onnx.loader import display_domain, load_model +from fnnx.extras.compilers.c.onnx.runtime_dims import ShapeTerm, resolve_runtime_dims +from fnnx.extras.compilers.c.onnx.specialize import specialize +from fnnx.extras.compilers.c.result import CompileResult + +COMPILER = "fnnx.extras.compilers.c" + + +def compile_onnx( + source: str | os.PathLike[str] | ModelProto, + output_dir: str | os.PathLike[str], + *, + dim_bindings: Mapping[str, int] | None = None, + runtime_dims: Mapping[str, int] | None = None, + prefix: str | None = None, +) -> CompileResult: + """Compile an ONNX model into a single self-contained C99 header plus a report. + + `source` is a path to a `.onnx` file or an in-memory `ModelProto`. `prefix` defaults to + the graph name, sanitized to a C identifier. `runtime_dims` maps a symbolic dimension to + the largest size the artifact must serve, leaving the actual size to each call; every + other symbolic dimension is fixed at compile time. Compilation is all-or-nothing: + nothing is written unless the whole model compiles. + """ + dims = resolve_runtime_dims(runtime_dims, dim_bindings) + loaded = load_model(source) + + def build(bindings: Mapping[str, int]) -> Program: + return build_program( + prepare_model(loaded, dim_bindings=bindings), + prefix=prefix, + runtime_dims=dims, + ) + + program = ( + specialize(build, dims, dim_bindings=dim_bindings or {}) + if dims + else build(dim_bindings or {}) + ) + return write_artifact( + program, + output_dir, + options={ + "prefix": prefix, + "dim_bindings": dict(sorted((dim_bindings or {}).items())), + "runtime_dims": dict((runtime_dims or {}).items()), + }, + ) + + +def write_artifact( + program: Program, output_dir: str | os.PathLike[str], *, options: dict[str, Any] +) -> CompileResult: + """Render `program` into its output directory, creating the directory if needed. + + The last step of every compilation, and the only one that writes: a model that fails to + compile leaves no partial artifact behind. + """ + report = build_report(program, options=options) + header = render_header(program) + directory = Path(output_dir) + directory.mkdir(parents=True, exist_ok=True) + header_path = directory / report["header"] + report_path = directory / f"{program.prefix}_report.json" + header_path.write_text(header, encoding="utf-8") + report_path.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return CompileResult( + header_path=header_path, report_path=report_path, report=report + ) + + +def build_report(program: Program, *, options: dict[str, Any]) -> dict[str, Any]: + """Machine-readable description of the artifact, for tooling and for the harness.""" + weights_bytes = buffer_bytes(program.weights) + arena_bytes = buffer_bytes(program.scratch) + return { + "fnnx_version": __version__, + "compiler": COMPILER, + "prefix": program.prefix, + "header": f"{program.prefix}.h", + "graph": program.graph_name, + "options": options, + "dim_bindings": dict(sorted(program.dim_bindings.items())), + "runtime_dims": [ + { + "name": dim.name, + "max": dim.maximum, + "parameter": dim.c_name, + "macro": dim.macro(program.prefix), + } + for dim in program.runtime_dims + ], + "opsets": { + display_domain(domain): version + for domain, version in sorted(program.opsets.items()) + }, + "kernels": [function.name for function in program.functions], + "class_labels": [ + { + "tensor": table.tensor, + "symbol": table.symbol, + "macro": table.macro, + "dtype": "str" if table.elem_type == TensorProto.STRING else "int64", + "values": list(table.values), + } + for table in program.labels + ], + "memory": { + "weights_bytes": weights_bytes, + "arena_bytes": arena_bytes, + "static_bytes": weights_bytes + arena_bytes, + }, + "entrypoint": { + "symbol": f"{program.prefix}_run", + "inputs": [_tensor_report(tensor) for tensor in program.inputs], + "outputs": [_tensor_report(tensor) for tensor in program.outputs], + }, + "nodes": [ + { + "id": node.id, + "symbol": node.symbol, + "inputs": [_tensor_report(tensor) for tensor in node.inputs], + "outputs": [_tensor_report(tensor) for tensor in node.outputs], + } + for node in program.nodes + ], + } + + +def _tensor_report(tensor: IOTensor) -> dict[str, Any]: + """One tensor of an entrypoint's signature; `shape` is the buffer's capacity. + + With runtime dimensions in play the buffer is sized for their maxima while a call works + at whatever sizes it passes, so `runtime_shape` says how each axis follows from them — + which is what the load-and-run harness sizes its arrays by. + """ + report = { + "name": tensor.name, + "c_name": tensor.c_name, + "macro": tensor.macro, + "dtype": numpy_dtype_name(tensor.elem_type), + "c_type": c_type(tensor.elem_type), + "shape": list(tensor.shape), + "elem_count": tensor.elem_count, + "bytes": tensor.elem_count * element_size(tensor.elem_type), + } + if tensor.runtime_shape: + report["runtime_shape"] = [_axis_report(term) for term in tensor.runtime_shape] + return report + + +def _axis_report(term: ShapeTerm) -> Any: + if term.dim is None: + return term.size + return {"dim": term.dim, "coefficient": term.coefficient} diff --git a/src/python/fnnx/extras/compilers/c/onnx/codegen.py b/src/python/fnnx/extras/compilers/c/onnx/codegen.py new file mode 100644 index 0000000..15cb713 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/codegen.py @@ -0,0 +1,717 @@ +"""Planning a prepared graph into C: buffers, kernel dispatch, and the entrypoint body.""" + +from __future__ import annotations + +import math +import re +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field + +from onnx import GraphProto, NodeProto, TensorProto, TypeProto, helper + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import element_type_name +from fnnx.extras.compilers.c.onnx.emit import UniqueNames, sanitize_identifier +from fnnx.extras.compilers.c.onnx.frontend import PreparedModel +from fnnx.extras.compilers.c.onnx.functions import ( + MAX_EXPANSION_DEPTH, + Expansion, + expand_function, +) +from fnnx.extras.compilers.c.onnx.kernels import ( + KERNELS, + CFunction, + ConstantData, + KernelGenerator, + NodeContext, + ScratchBuffer, + TensorRef, +) +from fnnx.extras.compilers.c.onnx.loader import display_domain, normalize_domain +from fnnx.extras.compilers.c.onnx.registry import KernelSpec +from fnnx.extras.compilers.c.onnx.runtime_dims import RuntimeDim, ShapeTerm +from fnnx.extras.compilers.c.onnx.shapes import graph_label, static_shape, tensor_types + +DEFAULT_PREFIX = "fnnx_model" + + +@dataclass(frozen=True) +class IOTensor: + """A tensor the caller provides a buffer for: a graph input or output. + + `shape` is the buffer's capacity — the extents at the runtime dimensions' maxima — and + `runtime_shape`, present only where the artifact has runtime dimensions, says which of + those extents scale with which dimension. + """ + + name: str + c_name: str + macro: str + elem_type: int + shape: tuple[int, ...] + runtime_shape: tuple[ShapeTerm, ...] = () + owner: str = "" + + @property + def elem_count(self) -> int: + return math.prod(self.shape) + + +@dataclass(frozen=True) +class StaticBuffer: + """A `static` array the implementation owns: an embedded weight, or scratch space.""" + + name: str + symbol: str + elem_type: int + shape: tuple[int, ...] + tensor: TensorProto | None + + @property + def elem_count(self) -> int: + return math.prod(self.shape) + + @property + def declared_count(self) -> int: + """C99 has no zero-length arrays, yet a zero-element tensor still needs an address.""" + return max(1, self.elem_count) + + +@dataclass(frozen=True) +class LabelTable: + """Class names the header publishes alongside the output tensor they describe.""" + + tensor: str + symbol: str + macro: str + elem_type: int + values: tuple[str, ...] | tuple[int, ...] + + +@dataclass(frozen=True) +class NodeEntry: + """A per-node entrypoint the artifact publishes beside the whole-model one. + + `id` is the FNNX op-instance id the caller knows the node by; `symbol` is what that id + was sanitized to, which is what the header actually declares. + """ + + id: str + symbol: str + inputs: tuple[IOTensor, ...] + outputs: tuple[IOTensor, ...] + body: tuple[str, ...] + body_owners: tuple[str, ...] = () + + +@dataclass(frozen=True) +class ArtifactScope: + """What a graph shares with the other graphs compiled into the same header. + + `prefix` is the artifact-wide prefix kernels, their scratch and their constant tables + are named from, so that graphs using the same kernel share one definition; `symbols` + hands out every other identifier, keeping the graphs' own definitions distinct. + """ + + prefix: str + symbols: UniqueNames + + +@dataclass(frozen=True) +class Program: + """Everything the header renderer needs, with every ordering already fixed. + + `body_owners` names, per statement, the node the statement was emitted for; it is what + lets a failure to compile a whole family of shapes point back at a node of the model. + """ + + prefix: str + graph_name: str + source: str + opsets: dict[str, int] + dim_bindings: dict[str, int] + inputs: tuple[IOTensor, ...] + outputs: tuple[IOTensor, ...] + weights: tuple[StaticBuffer, ...] + scratch: tuple[StaticBuffer, ...] + functions: tuple[CFunction, ...] + body: tuple[str, ...] + labels: tuple[LabelTable, ...] = () + nodes: tuple[NodeEntry, ...] = () + runtime_dims: tuple[RuntimeDim, ...] = () + body_owners: tuple[str, ...] = () + + +def build_program( + prepared: PreparedModel, + *, + prefix: str | None = None, + scope: ArtifactScope | None = None, + runtime_dims: tuple[RuntimeDim, ...] = (), +) -> Program: + """Plan the C code for a prepared graph. + + `prefix` defaults to the graph name, sanitized to a C identifier, and to `fnnx_model` + when that leaves nothing. `scope` compiles this graph as one of several sharing a + header — kernels deduplicated between them, identifiers kept apart — and defaults to a + scope of the graph's own; a caller passing one owns reserving `runtime_dims`' parameter + names in it, so that no tensor is emitted under a name a parameter already carries. + """ + return _ProgramBuilder(prepared, prefix, scope, runtime_dims).build() + + +def reserve_dim_parameters( + names: UniqueNames, runtime_dims: Sequence[RuntimeDim] +) -> None: + """Claim the entrypoint parameters the runtime dimensions take, before any tensor does.""" + for dim in runtime_dims: + taken = names.assign(dim.c_name, fallback=dim.c_name) + assert taken == dim.c_name + + +@dataclass +class _Slot: + """Where a tensor's data lives, as a C expression, plus its static type.""" + + expr: str + elem_type: int + shape: tuple[int, ...] + + +def _slot_for(ref: TensorRef) -> _Slot: + return _Slot(ref.expr, ref.elem_type, ref.shape) + + +def _initializers(graph: GraphProto) -> dict[str, TensorProto]: + return {initializer.name: initializer for initializer in graph.initializer} + + +def _bound(ref: TensorRef | None) -> TensorRef: + """An operand an expansion binds; it only ever binds ones the node actually passes.""" + if ref is None: + raise CompileError("the function body binds an operand the node does not pass.") + return ref + + +@dataclass +class _ProgramBuilder: + prepared: PreparedModel + requested_prefix: str | None + requested_scope: ArtifactScope | None = None + runtime_dims: tuple[RuntimeDim, ...] = () + + prefix: str = field(init=False) + scope: ArtifactScope = field(init=False) + types: dict[str, TypeProto] = field(init=False) + constants: dict[str, TensorProto] = field(init=False) + slots: dict[str, _Slot] = field(default_factory=dict, init=False) + depth: int = field(default=0, init=False) + referenced: set[str] = field(default_factory=set, init=False) + weights: list[StaticBuffer] = field(default_factory=list, init=False) + scratch: list[StaticBuffer] = field(default_factory=list, init=False) + kernel_scratch: dict[str, StaticBuffer] = field(default_factory=dict, init=False) + constants_data: dict[str, TensorProto] = field(default_factory=dict, init=False) + functions: dict[str, CFunction] = field(default_factory=dict, init=False) + statements: list[str] = field(default_factory=list, init=False) + statement_owners: list[str] = field(default_factory=list, init=False) + + def __post_init__(self) -> None: + graph = self.prepared.model.graph + source = ( + self.requested_prefix if self.requested_prefix is not None else graph.name + ) + self.prefix = sanitize_identifier(source, fallback=DEFAULT_PREFIX) + if self.requested_scope is not None: + self.scope = self.requested_scope + else: + self.scope = ArtifactScope(self.prefix, UniqueNames()) + reserve_dim_parameters(self.scope.symbols, self.runtime_dims) + self.types = tensor_types(graph) + self.constants = _initializers(graph) + + @property + def graph(self) -> GraphProto: + return self.prepared.model.graph + + @property + def names(self) -> UniqueNames: + """One namespace for parameters and static buffers alike, across every graph. + + A parameter that happened to match a buffer symbol would shadow it inside the + entrypoint; a buffer two graphs of one artifact both named would be defined twice. + """ + return self.scope.symbols + + def _emit(self, statements: Iterable[str], owner: str) -> None: + for statement in statements: + self.statements.append(statement) + self.statement_owners.append(owner) + + def build(self) -> Program: + inputs = self._plan_inputs() + self._plan_weights() + outputs, copies = self._plan_outputs() + labels = self._plan_labels(outputs) + for node in self.graph.node: + self._emit_node(node) + self._emit(copies, "") + unused = [ + f"(void){tensor.c_name};" + for tensor in (*inputs, *outputs) + if tensor.c_name not in self.referenced + ] + return Program( + prefix=self.prefix, + graph_name=self.graph.name, + source=f"ONNX graph `{self.graph.name}`", + opsets=dict(self.prepared.opsets), + dim_bindings=dict(self.prepared.dim_bindings), + inputs=inputs, + outputs=outputs, + # A buffer no emitted statement names — a weight nothing reads, or the + # intermediate of a zero-element copy that compiles to nothing — would be a + # `static` the C compiler warns about, so only the referenced ones are emitted. + weights=tuple( + weight for weight in self.weights if weight.symbol in self.referenced + ), + scratch=tuple( + buffer for buffer in self.scratch if buffer.symbol in self.referenced + ), + functions=tuple(self.functions.values()), + body=tuple(unused + self.statements), + labels=labels, + runtime_dims=self.runtime_dims, + body_owners=tuple([""] * len(unused) + self.statement_owners), + ) + + def _plan_labels(self, outputs: tuple[IOTensor, ...]) -> tuple[LabelTable, ...]: + """Give every class-label table a symbol and a macro, keyed to its output tensor.""" + by_name = {tensor.name: tensor for tensor in outputs} + tables = [] + for labels in self.prepared.class_labels: + tensor = by_name.get(labels.tensor) + if tensor is None: + raise CompileError( + f"Graph `{graph_label(self.graph)}`: the class labels of " + f"`{labels.tensor}` describe a tensor the graph does not output." + ) + symbol = self.names.assign( + f"{self.prefix}_classlabels_{tensor.c_name}", + fallback=f"{self.prefix}_classlabels", + ) + tables.append( + LabelTable( + tensor=tensor.name, + symbol=symbol, + # Derived from the symbol rather than from the output's own macro family, + # because the symbol is the name `self.names` has already made unique: + # two tables keying one tensor would otherwise share a macro and define + # it twice with the two lengths. + macro=f"{symbol.upper()}_COUNT", + elem_type=labels.elem_type, + values=labels.values, + ) + ) + return tuple(tables) + + def _plan_inputs(self) -> tuple[IOTensor, ...]: + inputs = [] + for entry in self.graph.input: + elem_type, shape = self._tensor_type(entry.name, f"input `{entry.name}`") + c_name = self.names.assign(entry.name, fallback="input") + self.slots[entry.name] = _Slot(c_name, elem_type, shape) + inputs.append( + self._io_tensor(entry.name, c_name, "INPUT", elem_type, shape) + ) + return tuple(inputs) + + def _plan_weights(self) -> None: + """Embed the initializers the emitted code actually reads, in graph order. + + One already bound to a buffer is left alone: a function body takes the operands its + caller fixes as initializers so that folding can read them, and the caller's own + buffer already holds those bytes. + """ + referenced = {name for node in self.graph.node for name in node.input if name} + referenced |= {entry.name for entry in self.graph.output} + for initializer in self.graph.initializer: + if initializer.name not in referenced or initializer.name in self.slots: + continue + symbol = self._storage_symbol("w", initializer.name) + shape = tuple(initializer.dims) + self.slots[initializer.name] = _Slot(symbol, initializer.data_type, shape) + self.weights.append( + StaticBuffer( + initializer.name, symbol, initializer.data_type, shape, initializer + ) + ) + + def _plan_outputs(self) -> tuple[tuple[IOTensor, ...], list[str]]: + """Give every graph output a parameter, copying into it where it aliases a buffer. + + An output a node computes is written straight into the caller's buffer, and every + downstream consumer reads it from there. One that aliases an input, a folded + constant or an earlier output is copied instead, after all nodes have run. + """ + produced = { + name: node.name or f"" + for node in self.graph.node + for name in node.output + if name + } + outputs = [] + copies = [] + for entry in self.graph.output: + source = self.slots.get(entry.name) + c_name = self.names.assign(entry.name, fallback="output") + if source is None: + if entry.name not in produced: + raise CompileError( + f"Graph `{graph_label(self.graph)}`: output `{entry.name}` is not " + "produced by any node, input or initializer." + ) + elem_type, shape = self._tensor_type( + entry.name, f"output `{entry.name}`" + ) + self.slots[entry.name] = _Slot(c_name, elem_type, shape) + else: + elem_type, shape = source.elem_type, source.shape + count = math.prod(shape) + if count: + self.referenced.add(c_name) + self._mark_used(source) + copies.append( + f"memcpy({c_name}, {source.expr}, " + f"{count}u * sizeof(*{c_name}));" + ) + outputs.append( + self._io_tensor( + entry.name, + c_name, + "OUTPUT", + elem_type, + shape, + owner=produced.get(entry.name, ""), + ) + ) + return tuple(outputs), copies + + def _emit_node(self, node: NodeProto) -> None: + """Compile one node: its native kernel, else its ONNX function body, else an error.""" + label = node.name or f"" + domain = normalize_domain(node.domain) + opset_version = self.prepared.opsets.get(domain) + if opset_version is None: + raise CompileError( + f"Graph `{graph_label(self.graph)}`: node `{label}` uses domain " + f"`{display_domain(domain)}`, which the model does not import an opset for." + ) + inputs = tuple(self._read_ref(name, label) for name in node.input) + outputs = tuple(self._write_ref(name) for name in node.output) + spec = KERNELS.select(domain, node.op_type, opset_version) + if spec is not None: + self._emit_kernel(spec, node, domain, opset_version, inputs, outputs, label) + elif not self._expand_node(node, domain, opset_version, label, inputs, outputs): + raise KERNELS.unsupported_op_error( + domain, node.op_type, opset_version, node_name=label + ) + + def _emit_kernel( + self, + spec: KernelSpec[KernelGenerator], + node: NodeProto, + domain: str, + opset_version: int, + inputs: tuple[TensorRef | None, ...], + outputs: tuple[TensorRef | None, ...], + label: str, + ) -> None: + context = NodeContext( + node=node, + domain=domain, + opset_version=opset_version, + since_version=spec.since_version, + prefix=self.scope.prefix, + inputs=inputs, + outputs=outputs, + constants=self.constants, + ) + emission = spec.generator(context) + for function in emission.functions: + self._add_function(function) + for constant in emission.constants: + self._embed_constant(constant) + for buffer in emission.scratch: + self._reserve_scratch(buffer) + self._mark_emitted(context, emission.statements) + self._emit(emission.statements, label) + + def _expand_node( + self, + node: NodeProto, + domain: str, + opset_version: int, + label: str, + inputs: tuple[TensorRef | None, ...], + outputs: tuple[TensorRef | None, ...], + ) -> bool: + """Compile the node through the function body ONNX defines for its op. + + False means ONNX defines no body, leaving the node unsupported. Anything the body + itself cannot compile is an error naming the node it was expanded for, so a failure + several expansions deep still points back at the model's own node. + """ + if self.depth >= MAX_EXPANSION_DEPTH: + raise CompileError( + f"Node `{label}`: ONNX function bodies nested more than " + f"{MAX_EXPANSION_DEPTH} levels deep; `{node.op_type}` (domain " + f"`{display_domain(domain)}`) does not expand into primitive ops." + ) + input_types = tuple( + None + if ref is None + else helper.make_tensor_type_proto(ref.elem_type, list(ref.shape)) + for ref in inputs + ) + # The values the graph fixes go with the types: a body that computes its own result + # shape from an operand needs the operand, not just its extents. + input_values = { + index: self.constants[ref.name] + for index, ref in enumerate(inputs) + if ref is not None and ref.name in self.constants + } + try: + expansion = expand_function( + node, domain, opset_version, input_types, input_values + ) + if expansion is None: + return False + self._emit_expansion(expansion, inputs, outputs, label) + except CompileError as error: + raise CompileError( + f"Node `{label}`: compiling the ONNX function body of `{node.op_type}` " + f"(domain `{display_domain(domain)}`, opset version {opset_version}) " + f"failed: {error}" + ) from error + return True + + def _emit_expansion( + self, + expansion: Expansion, + inputs: tuple[TensorRef | None, ...], + outputs: tuple[TensorRef | None, ...], + label: str, + ) -> None: + """Emit a prepared function body against the expanded node's own buffers. + + The body is a graph of its own — its tensor names, types and opsets are unrelated to + the enclosing graph's — so it is compiled in a scope of its own, with only the + caller's buffers shared: the body writes its outputs straight into them. + """ + outer = (self.prepared, self.types, self.constants, self.slots) + graph = expansion.prepared.model.graph + self.prepared = expansion.prepared + self.types = tensor_types(graph) + self.constants = _initializers(graph) + self.slots = {} + self.depth += 1 + try: + for name, index in expansion.inputs: + self.slots[name] = _slot_for(_bound(inputs[index])) + self._plan_weights() + for name, index in expansion.outputs: + self._bind_body_output(name, _bound(outputs[index])) + for node in graph.node: + self._emit_node(node) + for name, index in expansion.outputs: + self._copy_body_output(name, _bound(outputs[index]), label) + finally: + self.prepared, self.types, self.constants, self.slots = outer + self.depth -= 1 + + def _bind_body_output(self, name: str, ref: TensorRef) -> None: + """Point a body output at the caller's buffer, so the body writes straight into it. + + A name the body already binds — an input it passes through, or a constant folding + resolved it to — keeps that binding and is copied out afterwards instead. + """ + declared = self._tensor_type(name, f"output `{name}`") + if declared != (ref.elem_type, ref.shape): + raise CompileError( + f"the body computes `{name}` as " + f"{element_type_name(declared[0])}{list(declared[1])}, but the node's " + f"output is {element_type_name(ref.elem_type)}{list(ref.shape)}." + ) + self.slots.setdefault(name, _slot_for(ref)) + + def _copy_body_output(self, name: str, ref: TensorRef, label: str) -> None: + slot = self.slots[name] + count = math.prod(slot.shape) + if slot.expr == ref.expr or not count: + return + self.referenced.add(ref.expr) + self._mark_used(slot) + self._emit( + (f"memcpy({ref.expr}, {slot.expr}, {count}u * sizeof(*{ref.expr}));",), + label, + ) + + def _mark_emitted(self, context: NodeContext, statements: tuple[str, ...]) -> None: + """Record the buffers the emitted call sites actually name. + + An operand the kernel drops — Gemm's `C` when beta is zero, say — must not keep a + weight or an entrypoint parameter alive, or the artifact stops building under + `-Wunused-const-variable` and `-Wunused-parameter`. + """ + text = "\n".join(statements) + for ref in (*context.inputs, *context.outputs): + if ref is not None and re.search(rf"\b{re.escape(ref.expr)}\b", text): + self.referenced.add(ref.expr) + + def _read_ref(self, name: str, node_label: str) -> TensorRef | None: + if not name: + return None + slot = self.slots.get(name) + if slot is None: + raise CompileError( + f"Graph `{graph_label(self.graph)}`: node `{node_label}` reads tensor " + f"`{name}`, which no input, initializer or preceding node defines." + ) + return TensorRef(name, slot.elem_type, slot.shape, slot.expr) + + def _write_ref(self, name: str) -> TensorRef | None: + if not name: + return None + slot = self.slots.get(name) + if slot is None: + elem_type, shape = self._tensor_type(name, f"tensor `{name}`") + symbol = self._storage_symbol("t", name) + slot = _Slot(symbol, elem_type, shape) + self.slots[name] = slot + self.scratch.append(StaticBuffer(name, symbol, elem_type, shape, None)) + return TensorRef(name, slot.elem_type, slot.shape, slot.expr) + + def _add_function(self, function: CFunction) -> None: + existing = self.functions.get(function.name) + if existing is not None: + if existing.definition != function.definition: + raise CompileError( + f"Kernel `{function.name}` was emitted twice with different " + "definitions; a kernel name must encode everything its code " + "depends on." + ) + return + if self.names.is_taken(function.name): + # Kernel names are composed from the prefix rather than handed out by + # `self.names`, so that nodes sharing a kernel agree on it; a tensor that + # sanitizes to the same identifier would shadow the function inside the + # entrypoint, and the build would fail on an opaque C diagnostic. + raise CompileError( + f"Kernel `{function.name}` collides with the identifier already emitted " + f"for a tensor of that name in graph `{graph_label(self.graph)}`; " + "compile with an explicit `prefix` to keep the two apart." + ) + self.functions[function.name] = function + + def _embed_constant(self, constant: ConstantData) -> None: + """Embed a table a kernel reads from its attributes, once per distinct table. + + It is a weight in every way that matters — `static const` data the call site names, + counted in the reported footprint — so it is planned as one; what differs is only + that it comes from an attribute rather than from an initializer. + """ + existing = self.constants_data.get(constant.symbol) + if existing is not None: + if existing.SerializeToString() != constant.tensor.SerializeToString(): + raise CompileError( + f"Constant table `{constant.symbol}` was emitted twice with different " + "contents; its symbol must encode everything the data depends on." + ) + else: + if self.names.is_taken(constant.symbol): + # Composed from the prefix rather than handed out by `self.names`, so that + # nodes reading the same table agree on it; a tensor sanitizing to the same + # identifier would shadow the data inside the entrypoint. + raise CompileError( + f"Constant table `{constant.symbol}` collides with the identifier " + f"already emitted for a tensor of that name in graph " + f"`{graph_label(self.graph)}`; compile with an explicit `prefix` to " + "keep the two apart." + ) + self.constants_data[constant.symbol] = constant.tensor + self.weights.append( + StaticBuffer( + constant.symbol, + constant.symbol, + constant.tensor.data_type, + tuple(constant.tensor.dims), + constant.tensor, + ) + ) + self.referenced.add(constant.symbol) + + def _reserve_scratch(self, buffer: ScratchBuffer) -> None: + """Set aside the working storage a kernel asked for, as one static buffer. + + Nodes calling the same kernel at different shapes ask for different amounts, so the + buffer grows to the largest of them rather than being reserved once per node; they + run one after another, which is what makes sharing it safe. + """ + reserved = self.kernel_scratch.get(buffer.symbol) + if reserved is not None and reserved.elem_type != buffer.elem_type: + raise CompileError( + f"Kernel scratch `{buffer.symbol}` was reserved at two element types; a " + "scratch symbol must encode everything its storage depends on." + ) + if reserved is None and self.names.is_taken(buffer.symbol): + # Composed from the prefix rather than handed out by `self.names`, for the + # reason kernel names are; a tensor sanitizing to the same identifier would + # shadow the buffer inside the entrypoint. + raise CompileError( + f"Kernel scratch `{buffer.symbol}` collides with the identifier already " + f"emitted for a tensor of that name in graph `{graph_label(self.graph)}`; " + "compile with an explicit `prefix` to keep the two apart." + ) + if reserved is None or reserved.elem_count < buffer.elem_count: + grown = StaticBuffer( + buffer.symbol, + buffer.symbol, + buffer.elem_type, + (buffer.elem_count,), + None, + ) + if reserved is None: + self.scratch.append(grown) + else: + self.scratch[self.scratch.index(reserved)] = grown + self.kernel_scratch[buffer.symbol] = grown + self.referenced.add(buffer.symbol) + + def _mark_used(self, slot: _Slot) -> None: + self.referenced.add(slot.expr) + + def _io_tensor( + self, + name: str, + c_name: str, + role: str, + elem_type: int, + shape: tuple[int, ...], + owner: str = "", + ) -> IOTensor: + macro = f"{self.prefix.upper()}_{role}_{c_name.upper()}" + return IOTensor(name, c_name, macro, elem_type, shape, owner=owner) + + def _storage_symbol(self, kind: str, name: str) -> str: + return self.names.assign( + f"{self.prefix}_{kind}_{name}", fallback=f"{self.prefix}_{kind}" + ) + + def _tensor_type(self, name: str, role: str) -> tuple[int, tuple[int, ...]]: + type_proto = self.types.get(name) + shape = static_shape(type_proto) + if type_proto is None or shape is None: + raise CompileError( + f"Graph `{graph_label(self.graph)}`: {role} has no static tensor type." + ) + return type_proto.tensor_type.elem_type, shape diff --git a/src/python/fnnx/extras/compilers/c/onnx/dtypes.py b/src/python/fnnx/extras/compilers/c/onnx/dtypes.py new file mode 100644 index 0000000..8a50d83 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/dtypes.py @@ -0,0 +1,70 @@ +"""ONNX element types the C compiler supports, and the C types they are emitted as.""" + +from __future__ import annotations + +from onnx import TensorProto, helper + +from fnnx.extras.compilers.c.errors import CompileError + +# bool becomes a byte holding 0/1 rather than C99's `_Bool`: every tensor buffer then has +# an explicit, fixed-width layout that is identical across compilers and ABIs. +C_TYPES: dict[int, str] = { + TensorProto.FLOAT: "float", + TensorProto.DOUBLE: "double", + TensorProto.INT8: "int8_t", + TensorProto.INT16: "int16_t", + TensorProto.INT32: "int32_t", + TensorProto.INT64: "int64_t", + TensorProto.UINT8: "uint8_t", + TensorProto.UINT16: "uint16_t", + TensorProto.UINT32: "uint32_t", + TensorProto.UINT64: "uint64_t", + TensorProto.BOOL: "uint8_t", +} + +# The floating-point element types: the ones whose kernels have to reckon with NaN and +# signed zero, and whose arithmetic differs from the integer families'. +FLOAT_TYPES = frozenset({TensorProto.FLOAT, TensorProto.DOUBLE}) + +# The unsigned integer element types. A kernel branching on `value < 0` is not merely dead +# code for these — it is a `-Wtype-limits` diagnostic, which the artifact's `-Werror` build +# contract turns into a failure — so their kernels drop the negative branch instead. +UNSIGNED_TYPES = frozenset( + { + TensorProto.UINT8, + TensorProto.UINT16, + TensorProto.UINT32, + TensorProto.UINT64, + } +) + + +def element_type_name(elem_type: int) -> str: + try: + return TensorProto.DataType.Name(elem_type) + except ValueError: + return f"UNKNOWN({elem_type})" + + +def is_supported(elem_type: int) -> bool: + return elem_type in C_TYPES + + +def numpy_dtype_name(elem_type: int) -> str: + """The numpy name (`float32`, `bool`, ...) callers bind this element type to.""" + return helper.tensor_dtype_to_np_dtype(elem_type).name + + +def element_size(elem_type: int) -> int: + return helper.tensor_dtype_to_np_dtype(elem_type).itemsize + + +def c_type(elem_type: int) -> str: + try: + return C_TYPES[elem_type] + except KeyError: + raise CompileError( + f"Element type `{element_type_name(elem_type)}` is not supported by the C " + f"compiler; supported types are " + f"{', '.join(sorted(element_type_name(t) for t in C_TYPES))}." + ) from None diff --git a/src/python/fnnx/extras/compilers/c/onnx/emit.py b/src/python/fnnx/extras/compilers/c/onnx/emit.py new file mode 100644 index 0000000..edbcded --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/emit.py @@ -0,0 +1,189 @@ +"""Primitives for emitting C source: identifiers, literals, and initializer lists.""" + +from __future__ import annotations + +import math +import string +from collections.abc import Iterable, Iterator +from typing import Any + +from onnx import TensorProto + +C_KEYWORDS = frozenset( + { + "auto", + "break", + "case", + "char", + "const", + "continue", + "default", + "do", + "double", + "else", + "enum", + "extern", + "float", + "for", + "goto", + "if", + "inline", + "int", + "long", + "register", + "restrict", + "return", + "short", + "signed", + "sizeof", + "static", + "struct", + "switch", + "typedef", + "union", + "unsigned", + "void", + "volatile", + "while", + "_Bool", + "_Complex", + "_Imaginary", + } +) + +# The nonzero status an entrypoint returns for an argument it cannot serve; the enumerator +# is the artifact's prefix followed by this, and a kernel that validates an operand at run +# time returns it by that name. +INVALID_ARGUMENT_STATUS = "ERROR_INVALID_ARGUMENT" + +_IDENTIFIER_CHARS = frozenset(string.ascii_letters + string.digits + "_") +_STRING_ESCAPES = {'"': '\\"', "\\": "\\\\", "?": "\\?"} +_UNSIGNED_TYPES = frozenset({TensorProto.UINT8, TensorProto.UINT16, TensorProto.UINT32}) + +# Decimal precision that reads back as the exact same value, per IEEE-754 binary32/binary64. +_FLOAT_DIGITS = 9 +_DOUBLE_DIGITS = 17 + +_LINE_WIDTH = 88 + + +def sanitize_identifier(name: str, *, fallback: str) -> str: + """Turn an arbitrary ONNX name into a valid C identifier, deterministically.""" + cleaned = "".join(char if char in _IDENTIFIER_CHARS else "_" for char in name) + if not any(char.isalnum() for char in cleaned): + return fallback + if not cleaned[0].isalpha(): + # C identifiers cannot start with a digit, and a leading underscore is reserved + # for the implementation at file scope. + cleaned = f"v_{cleaned}" + if cleaned in C_KEYWORDS: + cleaned = f"{cleaned}_" + return cleaned + + +class UniqueNames: + """Hands out distinct C identifiers, disambiguating collisions deterministically. + + Uniqueness is case-insensitive so that the uppercase macro names derived from these + identifiers stay distinct too. + """ + + def __init__(self) -> None: + self._taken: set[str] = set() + + def is_taken(self, name: str) -> bool: + return name.upper() in self._taken + + def assign(self, name: str, *, fallback: str) -> str: + base = sanitize_identifier(name, fallback=fallback) + candidate = base + suffix = 2 + while candidate.upper() in self._taken: + candidate = f"{base}_{suffix}" + suffix += 1 + self._taken.add(candidate.upper()) + return candidate + + +def scalar_literal(value: Any, elem_type: int) -> str: + """A C literal that reads back as exactly `value` at element type `elem_type`.""" + if elem_type == TensorProto.FLOAT: + return _float_literal(float(value), digits=_FLOAT_DIGITS, suffix="f") + if elem_type == TensorProto.DOUBLE: + return _float_literal(float(value), digits=_DOUBLE_DIGITS, suffix="") + if elem_type == TensorProto.BOOL: + return "1" if value else "0" + return _integer_literal(int(value), elem_type) + + +def string_literal(value: str) -> str: + """A C string literal holding exactly `value`'s UTF-8 bytes. + + Everything outside printable ASCII goes in as an octal escape rather than a hex one: + `\\x` escapes are greedy, so a `\\xff` followed by a literal `a` would read as a single + out-of-range character, while an octal escape stops after three digits. `?` is escaped + too, since a run of them forms a trigraph the C99 preprocessor still rewrites. + """ + pieces = [] + for byte in value.encode("utf-8"): + char = chr(byte) + if char in _STRING_ESCAPES: + pieces.append(_STRING_ESCAPES[char]) + elif 0x20 <= byte < 0x7F: + pieces.append(char) + else: + pieces.append(f"\\{byte:03o}") + return '"' + "".join(pieces) + '"' + + +def initializer_lines( + literals: Iterable[str], *, indent: str = " " +) -> Iterator[str]: + """Wrap literals into `{ ... }`-body lines, one comma-separated run per line.""" + line = indent + for literal in literals: + piece = f"{literal}," + if line != indent and len(line) + len(piece) > _LINE_WIDTH: + yield line.rstrip() + line = indent + line += f"{piece} " + if line.strip(): + yield line.rstrip() + + +def comment_safe(text: str) -> str: + """Neutralize block-comment delimiters so arbitrary names can go in comments. + + A nested `/*` is not a syntax error but is a `-Wall` diagnostic, which the artifact's + `-Werror` build contract turns into a failure, so it is broken up like `*/` is. + """ + return text.replace("*/", "* /").replace("/*", "/ *") + + +def _float_literal(value: float, *, digits: int, suffix: str) -> str: + # NaN and the infinities have no decimal form; 's macros are constant + # expressions usable in the static initializers weights are emitted as. + if math.isnan(value): + return "NAN" + if math.isinf(value): + return "INFINITY" if value > 0 else "-INFINITY" + text = f"{value:.{digits}g}" + if "." not in text and "e" not in text: + text += ".0" + return text + suffix + + +def _integer_literal(value: int, elem_type: int) -> str: + if elem_type == TensorProto.INT64: + if value == -(2**63): + # `-9223372036854775808` is negation applied to a constant too large to be + # signed; the minimum has to be built from the maximum. + return "(-INT64_C(9223372036854775807) - INT64_C(1))" + return f"INT64_C({value})" + if elem_type == TensorProto.UINT64: + return f"UINT64_C({value})" + if elem_type == TensorProto.INT32 and value == -(2**31): + return "(-2147483647 - 1)" + if elem_type in _UNSIGNED_TYPES: + return f"{value}u" + return str(value) diff --git a/src/python/fnnx/extras/compilers/c/onnx/folding.py b/src/python/fnnx/extras/compilers/c/onnx/folding.py new file mode 100644 index 0000000..48598a2 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/folding.py @@ -0,0 +1,289 @@ +"""Compile-time constant folding through the official ONNX reference evaluator.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from functools import lru_cache + +import numpy as np +import onnx.defs +from onnx import ( + AttributeProto, + GraphProto, + ModelProto, + NodeProto, + TensorProto, + TypeProto, + helper, +) +from onnx.numpy_helper import from_array +from onnx.reference import ReferenceEvaluator + +# The evaluator's versioned implementation classes are the mechanical record of where it +# distinguishes an op's historical semantics; onnx exposes them only through this builder. +from onnx.reference.ops._op_list import ( + _build_registered_operators as _standard_implementations, +) +from onnx.reference.ops.aionnxml._op_list import ( + _build_registered_operators as _ml_implementations, +) + +from fnnx.extras.compilers.c.onnx.loader import ( + ML_DOMAIN, + STANDARD_DOMAIN, + normalize_domain, +) +from fnnx.extras.compilers.c.onnx.shapes import static_shape, tensor_types + +# Ops whose output is a function of its input's shape alone, so a statically-shaped input +# makes them constant even though its values are not known until run time. +_SHAPE_ONLY_OPS = frozenset({"Shape", "Size"}) + +# Ops that draw from a random distribution: their output is not a function of their inputs, +# so folding one would bake a single draw into the artifact — silently turning an +# unsupported op into a wrong-but-compiling constant, and making the output non-deterministic +# across compiles. `Dropout` is listed because it samples a mask in training mode; the +# inference-mode identity it also serves is compiled as a kernel instead. +NONDETERMINISTIC_OPS = frozenset( + { + "Bernoulli", + "Dropout", + "Multinomial", + "RandomNormal", + "RandomNormalLike", + "RandomUniform", + "RandomUniformLike", + } +) + + +def evaluator_is_version_faithful( + domain: str, op_type: str, opset_version: int +) -> bool: + """Whether `onnx.reference` implements the op exactly as `opset_version` defines it. + + The evaluator carries a versioned implementation class per revision whose semantics it + distinguishes, and applies the newest one to every later opset; an op it has no class + for at all it runs by expanding the ONNX function body of the schema the imported opset + selects, which is the newest revision's body from that revision on. Either way it is + faithful when the requested opset selects a schema revision a versioned class + implements, or one at or after the newest revision the evaluator distinguishes — the + schema history's own newest, where there are no versioned classes to go by. Anything + older gets the wrong semantics silently, and is refused here rather than folded. + """ + normalized = normalize_domain(domain) + revision = _schema_revision(normalized, op_type, opset_version) + if revision is None: + return False + implemented = _implemented_revisions(normalized, op_type) + if implemented and revision in implemented: + return True + newest = max(implemented) if implemented else _latest_revision(normalized, op_type) + return newest is not None and revision >= newest + + +def fold_constants(model: ModelProto, opsets: Mapping[str, int]) -> int: + """Replace every node with compile-time-constant inputs by its computed value. + + Returns the number of nodes folded. Nodes the evaluator cannot vouch for or cannot + execute are left in place for kernel compilation: folding is what makes shape + computations static, never a correctness requirement. + """ + graph = model.graph + constants = {initializer.name: initializer for initializer in graph.initializer} + types = tensor_types(graph) + kept: list[NodeProto] = [] + folded: list[TensorProto] = [] + + for node in graph.node: + inputs = _constant_inputs(node, constants, types, opsets) + results = None if inputs is None else _evaluate(node, inputs, opsets) + if results is None: + kept.append(node) + continue + for value in results: + constants[value.name] = value + folded.append(value) + + if not folded: + return 0 + node_count = len(graph.node) + del graph.node[:] + graph.node.extend(kept) + graph.initializer.extend(folded) + return node_count - len(kept) + + +def prune_unused_initializers(graph: GraphProto) -> None: + referenced = _referenced_names(graph) + kept = [ + initializer + for initializer in graph.initializer + if initializer.name in referenced + ] + if len(kept) != len(graph.initializer): + del graph.initializer[:] + graph.initializer.extend(kept) + + +def prune_stale_value_info(graph: GraphProto) -> None: + """Drop intermediate shapes for tensors no node produces any more.""" + produced = {name for node in graph.node for name in node.output if name} + kept = [entry for entry in graph.value_info if entry.name in produced] + if len(kept) != len(graph.value_info): + del graph.value_info[:] + graph.value_info.extend(kept) + + +def _constant_inputs( + node: NodeProto, + constants: Mapping[str, TensorProto], + types: Mapping[str, TypeProto], + opsets: Mapping[str, int], +) -> list[TensorProto] | None: + """Constant tensors the node reads, or None when it cannot be folded.""" + domain = normalize_domain(node.domain) + if _draws_at_random(node): + return None + opset_version = opsets.get(domain) + if opset_version is None or not evaluator_is_version_faithful( + domain, node.op_type, opset_version + ): + return None + + inputs: dict[str, TensorProto] = {} + for name in (*node.input, *sorted(_outer_scope_names(node))): + if not name or name in inputs: + continue + constant = constants.get(name) + if constant is None and domain == STANDARD_DOMAIN: + constant = _shape_only_placeholder(node, name, types) + if constant is None: + return None + inputs[name] = constant + return list(inputs.values()) + + +def _draws_at_random(node: NodeProto) -> bool: + """Whether the node — or any node its subgraphs would run — samples a distribution.""" + if ( + normalize_domain(node.domain) == STANDARD_DOMAIN + and node.op_type in NONDETERMINISTIC_OPS + ): + return True + return any( + _draws_at_random(inner) for graph in _subgraphs(node) for inner in graph.node + ) + + +def _shape_only_placeholder( + node: NodeProto, name: str, types: Mapping[str, TypeProto] +) -> TensorProto | None: + """A zero tensor standing in for a `Shape`/`Size` input of known static shape.""" + if node.op_type not in _SHAPE_ONLY_OPS: + return None + shape = static_shape(types.get(name)) + if shape is None: + return None + try: + dtype = helper.tensor_dtype_to_np_dtype(types[name].tensor_type.elem_type) + return from_array(np.zeros(shape, dtype=dtype), name) + except Exception: + return None + + +def _evaluate( + node: NodeProto, inputs: list[TensorProto], opsets: Mapping[str, int] +) -> list[TensorProto] | None: + outputs = [name for name in node.output if name] + graph = helper.make_graph( + [node], + "constant_fold", + [], + [helper.make_empty_tensor_value_info(name) for name in outputs], + initializer=inputs, + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid(domain, version) for domain, version in opsets.items() + ], + ) + try: + results = ReferenceEvaluator(model).run(None, {}) + folded = [ + from_array(np.asarray(value), name) for name, value in zip(outputs, results) + ] + except Exception: + return None + return folded + + +def _outer_scope_names(node: NodeProto) -> set[str]: + """Names a node's subgraphs read from the enclosing graph.""" + names: set[str] = set() + for subgraph in _subgraphs(node): + names |= _free_names(subgraph) + return names + + +def _free_names(graph: GraphProto) -> set[str]: + bound = {entry.name for entry in graph.input} | { + initializer.name for initializer in graph.initializer + } + free: set[str] = set() + for node in graph.node: + free |= {name for name in node.input if name and name not in bound} + free |= {name for name in _outer_scope_names(node) if name not in bound} + bound |= {name for name in node.output if name} + return free + + +def _referenced_names(graph: GraphProto) -> set[str]: + names = {entry.name for entry in graph.output} + for node in graph.node: + names |= {name for name in node.input if name} + names |= _outer_scope_names(node) + return names + + +def _subgraphs(node: NodeProto) -> Iterator[GraphProto]: + for attribute in node.attribute: + if attribute.type == AttributeProto.GRAPH: + yield attribute.g + elif attribute.type == AttributeProto.GRAPHS: + yield from attribute.graphs + + +def _schema_revision(domain: str, op_type: str, opset_version: int) -> int | None: + try: + return onnx.defs.get_schema(op_type, opset_version, domain).since_version + except onnx.defs.SchemaError: + return None + + +@lru_cache(maxsize=None) +def _latest_revision(domain: str, op_type: str) -> int | None: + revisions = [ + schema.since_version + for schema in onnx.defs.get_all_schemas_with_history() + if schema.name == op_type and schema.domain == domain + ] + return max(revisions) if revisions else None + + +@lru_cache(maxsize=None) +def _implemented_revisions(domain: str, op_type: str) -> frozenset[int] | None: + """Opset versions the reference evaluator implements separately, or None if it has no + implementation for the op at all (it may still expand a function body, whose semantics + this check cannot vouch for).""" + if domain == STANDARD_DOMAIN: + registry = _standard_implementations() + elif domain == ML_DOMAIN: + registry = _ml_implementations() + else: + return None + implementations = registry.get(op_type) + if implementations is None: + return None + return frozenset(version for version in implementations if isinstance(version, int)) diff --git a/src/python/fnnx/extras/compilers/c/onnx/frontend.py b/src/python/fnnx/extras/compilers/c/onnx/frontend.py new file mode 100644 index 0000000..e399e24 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/frontend.py @@ -0,0 +1,58 @@ +"""The compiler frontend: from a loaded model to a fully static, verified graph.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +from onnx import ModelProto + +from fnnx.extras.compilers.c.onnx import folding, shapes +from fnnx.extras.compilers.c.onnx.loader import LoadedModel +from fnnx.extras.compilers.c.onnx.verify import verify_static +from fnnx.extras.compilers.c.onnx.zipmap import ClassLabels, remove_zipmap + + +@dataclass(frozen=True) +class PreparedModel: + """A graph whose every tensor has a static shape and a supported element type.""" + + model: ModelProto + opsets: dict[str, int] + dim_bindings: dict[str, int] + class_labels: tuple[ClassLabels, ...] = () + + +def prepare_model( + loaded: LoadedModel, *, dim_bindings: Mapping[str, int] | None = None +) -> PreparedModel: + """Bind symbolic dimensions, fold constants to fixpoint, and verify the result.""" + model = ModelProto() + model.CopyFrom(loaded.model) + bindings = dim_bindings or {} + + shapes.drop_shadowed_inputs(model) + # Before the declared types are snapshotted: the entry a removed `ZipMap` leaves behind + # describes a sequence of maps, which is nothing the tensor taking its place could fall + # back on. + class_labels = remove_zipmap(model) + declared_outputs = shapes.declared_output_types(model.graph) + applied = shapes.bind_dims(model, bindings) + shapes.drop_empty_shape_operands(model) + shapes.state_stft_onesided(model) + model = shapes.infer_shapes(model) + while folding.fold_constants(model, loaded.opsets): + model = shapes.infer_shapes(model) + applied.update( + shapes.apply_declared_output_shapes(model, declared_outputs, bindings) + ) + folding.prune_unused_initializers(model.graph) + folding.prune_stale_value_info(model.graph) + + verify_static(model) + return PreparedModel( + model=model, + opsets=dict(loaded.opsets), + dim_bindings=applied, + class_labels=class_labels, + ) diff --git a/src/python/fnnx/extras/compilers/c/onnx/functions.py b/src/python/fnnx/extras/compilers/c/onnx/functions.py new file mode 100644 index 0000000..d270e32 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/functions.py @@ -0,0 +1,297 @@ +"""ONNX function bodies as compilable sub-models: what dispatch falls back to. + +An op no registered kernel serves is compiled through the body ONNX itself defines for it — +the very body `onnx.reference` expands, so the compiler and the oracle it is tested against +work from one definition of the op. + +The body is prepared as a **standalone model under its own opset imports** rather than +inlined into the caller's graph. ONNX writes a body against whatever opset suits it, which +need not be the one importing the op — `Relu`'s body is written at opset 18 for a schema +introduced at 14 — so inlining would either interpret body nodes at an opset they were not +written for, or produce a graph (opset 14 importing a node added at 15) that ONNX's own +shape inference rejects. Keeping the body separate means each half is compiled at exactly +the opset it was written for; only the emitted C is spliced together. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType + +import onnx.defs +from onnx import ( + AttributeProto, + FunctionProto, + ModelProto, + NodeProto, + TensorProto, + TypeProto, + helper, +) + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.frontend import PreparedModel, prepare_model +from fnnx.extras.compilers.c.onnx.loader import ( + LoadedModel, + display_domain, + resolve_opsets, +) + +# A body is compiled like any other model, so an op that is function-defined inside a +# function body expands in turn. ONNX's own bodies nest a few levels deep; the cap turns a +# body that expanded into itself into a compile error rather than a stack overflow. +MAX_EXPANSION_DEPTH = 8 + + +@dataclass(frozen=True) +class Binding: + """How a function body's graph I/O binds to the expanded node's. + + `inputs` and `outputs` pair a name in the body's graph with the index of the node + operand it stands for; operands the node leaves out have no entry. + """ + + inputs: tuple[tuple[str, int], ...] + outputs: tuple[tuple[str, int], ...] + + +@dataclass(frozen=True) +class FunctionModel(Binding): + """A function body as a standalone model, before the frontend's passes have run.""" + + model: ModelProto + + +@dataclass(frozen=True) +class Expansion(Binding): + """The same body once the frontend has made it static.""" + + prepared: PreparedModel + + +def function_body( + node: NodeProto, + domain: str, + opset_version: int, + input_types: Sequence[TypeProto | None], +) -> FunctionProto | None: + """The body ONNX defines for the node's op at `opset_version`, or None if it defines none. + + `input_types` holds the static type of each of the node's inputs, None where the node + leaves an optional operand out. Context-dependent bodies are built from the node and + those types, which is how an op like `Clip` collapses to just the comparisons its + optional inputs call for. + """ + schema = _schema(node.op_type, domain, opset_version) + return None if schema is None else _body(schema, node, input_types) + + +def function_model( + node: NodeProto, + domain: str, + opset_version: int, + input_types: Sequence[TypeProto | None], + input_values: Mapping[int, TensorProto] = MappingProxyType({}), +) -> FunctionModel | None: + """The op's ONNX function body as a standalone model, or None if ONNX defines none. + + The body becomes a model under its own opset imports, with the node's input types + (`input_types`, one per node input) as its graph inputs and the node's attributes + substituted into it. `input_values` carries the operands the caller's graph already + fixes, by position: they become initializers as well as inputs, which is what lets a body + compute its own result shape from one — `CenterCropPad` reads the extents it crops to as + a tensor, and without the value the body's `Pad` would take a shape no folding can settle. + """ + schema = _schema(node.op_type, domain, opset_version) + if schema is None: + return None + body = _body(schema, node, input_types) + if body is None: + return None + + provided = _provided_types(node, input_types) + inputs = tuple( + (formal, index) for index, formal in enumerate(body.input) if index in provided + ) + outputs = tuple( + (formal, index) + for index, formal in enumerate(body.output) + if index < len(node.output) and node.output[index] + ) + return FunctionModel( + inputs=inputs, + outputs=outputs, + model=_body_model(node, body, schema, provided, inputs, outputs, input_values), + ) + + +def expand_function( + node: NodeProto, + domain: str, + opset_version: int, + input_types: Sequence[TypeProto | None], + input_values: Mapping[int, TensorProto] = MappingProxyType({}), +) -> Expansion | None: + """Prepare the op's ONNX function body as a standalone static model, or None if it has none. + + The body goes through the same frontend as a top-level model — shape inference, + constant folding, static verification — under its own opset imports. + """ + built = function_model(node, domain, opset_version, input_types, input_values) + if built is None: + return None + model = built.model + return Expansion( + prepared=prepare_model(LoadedModel(model=model, opsets=resolve_opsets(model))), + inputs=built.inputs, + outputs=built.outputs, + ) + + +def _schema(op_type: str, domain: str, opset_version: int) -> onnx.defs.OpSchema | None: + try: + return onnx.defs.get_schema(op_type, opset_version, domain) + except onnx.defs.SchemaError: + return None + + +def _body( + schema: onnx.defs.OpSchema, + node: NodeProto, + input_types: Sequence[TypeProto | None], +) -> FunctionProto | None: + # `onnx`'s stubs declare only part of `OpSchema`'s pybind11 surface; these three members + # exist at runtime, and are the same ones `onnx.reference` dispatches a function on. + if schema.has_function: # type: ignore[attr-defined] + return schema.function_body + if not schema.has_context_dependent_function: # type: ignore[attr-defined] + return None + types = [entry if entry is not None else TypeProto() for entry in input_types] + try: + payload = schema.get_context_dependent_function( # type: ignore[attr-defined] + node.SerializeToString(), [entry.SerializeToString() for entry in types] + ) + except Exception as error: + raise CompileError( + f"ONNX could not build the function body of `{node.op_type}` (domain " + f"`{display_domain(schema.domain)}`) for this node: {error}" + ) from error + body = FunctionProto() + body.ParseFromString(payload) + return body + + +def _provided_types( + node: NodeProto, input_types: Sequence[TypeProto | None] +) -> dict[int, TypeProto]: + """Type of every operand the node actually passes, by position.""" + provided = {} + for index, name in enumerate(node.input): + if not name: + continue + type_proto = input_types[index] if index < len(input_types) else None + if type_proto is None: + raise CompileError(f"the type of operand `{name}` is not known.") + provided[index] = type_proto + return provided + + +def _body_model( + node: NodeProto, + body: FunctionProto, + schema: onnx.defs.OpSchema, + provided: Mapping[int, TypeProto], + inputs: tuple[tuple[str, int], ...], + outputs: tuple[tuple[str, int], ...], + values: Mapping[int, TensorProto], +) -> ModelProto: + """The body as a model computing the node's outputs from the node's input types.""" + omitted = { + formal: "" for index, formal in enumerate(body.input) if index not in provided + } + attributes = _attributes(node, schema) + nodes = _reachable( + [_resolved(entry, attributes, omitted) for entry in body.node], + {name for name, _ in outputs}, + ) + initializers = [] + for formal, index in inputs: + value = values.get(index) + if value is None: + continue + renamed = TensorProto() + renamed.CopyFrom(value) + renamed.name = formal + initializers.append(renamed) + graph = helper.make_graph( + nodes, + f"{node.op_type}_function_body", + [helper.make_value_info(formal, provided[index]) for formal, index in inputs], + [helper.make_empty_tensor_value_info(formal) for formal, _ in outputs], + initializer=initializers, + ) + return helper.make_model(graph, opset_imports=list(body.opset_import)) + + +def _attributes( + node: NodeProto, schema: onnx.defs.OpSchema +) -> dict[str, AttributeProto]: + """What the body's attribute references resolve against: the node's, then the defaults. + + A body reads the caller's attributes by name, so an attribute the node leaves at its + default has to be filled in from the schema — otherwise the body node referencing it + (typically a `Constant` holding the value) is left without one. + """ + attributes = {entry.name: entry for entry in node.attribute} + for name, formal in schema.attributes.items(): + if ( + name not in attributes + and formal.default_value.type != AttributeProto.UNDEFINED + ): + attributes[name] = formal.default_value + return attributes + + +def _resolved( + node: NodeProto, + attributes: Mapping[str, AttributeProto], + omitted: Mapping[str, str], +) -> NodeProto: + """A body node with its attribute references substituted and omitted operands blanked.""" + resolved = NodeProto() + resolved.CopyFrom(node) + del resolved.input[:] + resolved.input.extend(omitted.get(name, name) for name in node.input) + del resolved.attribute[:] + for attribute in node.attribute: + if not attribute.ref_attr_name: + resolved.attribute.append(attribute) + continue + supplied = attributes.get(attribute.ref_attr_name) + if supplied is None: + # Neither the node nor the schema gives the attribute a value, so the body + # node keeps none either and falls back on its own op's default. + continue + substituted = resolved.attribute.add() + substituted.CopyFrom(supplied) + substituted.name = attribute.name + return resolved + + +def _reachable(nodes: Sequence[NodeProto], wanted: set[str]) -> list[NodeProto]: + """The body nodes that contribute to the outputs the caller asked for. + + A body computes every output its op declares; a caller that omits an optional one must + not have to compile the ops that would have produced it. Function bodies are in + topological order, so one backwards pass suffices. + """ + live = set(wanted) + kept = [] + for node in reversed(nodes): + if not any(name in live for name in node.output if name): + continue + kept.append(node) + live |= {name for name in node.input if name} + kept.reverse() + return kept diff --git a/src/python/fnnx/extras/compilers/c/onnx/header.py b/src/python/fnnx/extras/compilers/c/onnx/header.py new file mode 100644 index 0000000..ed399ab --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/header.py @@ -0,0 +1,431 @@ +"""Rendering a planned program as one self-contained, STB-style C99 header.""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator + +from onnx import TensorProto +from onnx.numpy_helper import to_array + +from fnnx import __version__ +from fnnx.extras.compilers.c.onnx.codegen import ( + IOTensor, + LabelTable, + Program, + StaticBuffer, +) +from fnnx.extras.compilers.c.onnx.dtypes import c_type, element_size, numpy_dtype_name +from fnnx.extras.compilers.c.onnx.emit import ( + INVALID_ARGUMENT_STATUS, + comment_safe, + initializer_lines, + scalar_literal, + string_literal, +) +from fnnx.extras.compilers.c.onnx.loader import display_domain + + +def render_header(program: Program) -> str: + """The complete `.h` text: declarations, then a guarded implementation.""" + upper = program.prefix.upper() + lines: list[str] = [ + *_preamble(program), + f"#ifndef {upper}_H_INCLUDED", + f"#define {upper}_H_INCLUDED", + "", + "#include ", + "", + "#ifdef __cplusplus", + 'extern "C" {', + "#endif", + "", + *_status_enum(program), + *_footprint_macros(program), + *_runtime_dim_macros(program), + *_io_macros(program), + *_label_declarations(program), + *_entry_declarations(program), + "#ifdef __cplusplus", + "}", + "#endif", + "", + f"#endif /* {upper}_H_INCLUDED */", + "", + f"#ifdef {upper}_IMPLEMENTATION", + f"#ifndef {upper}_IMPLEMENTATION_INCLUDED", + f"#define {upper}_IMPLEMENTATION_INCLUDED", + "", + "#include ", + "#include ", + "#include ", + "", + "#ifdef __cplusplus", + 'extern "C" {', + "#endif", + "", + *_label_definitions(program), + *_weight_definitions(program), + *_scratch_definitions(program), + *_kernel_definitions(program), + *_entry_definitions(program), + "#ifdef __cplusplus", + "}", + "#endif", + "", + f"#endif /* {upper}_IMPLEMENTATION_INCLUDED */", + f"#endif /* {upper}_IMPLEMENTATION */", + ] + return "\n".join(lines) + "\n" + + +def buffer_bytes(buffers: Iterable[StaticBuffer]) -> int: + return sum( + buffer.declared_count * element_size(buffer.elem_type) for buffer in buffers + ) + + +def _preamble(program: Program) -> list[str]: + upper = program.prefix.upper() + opsets = ", ".join( + f"{display_domain(domain)}={version}" + for domain, version in sorted(program.opsets.items()) + ) + bindings = ( + ", ".join( + f"{name}={value}" for name, value in sorted(program.dim_bindings.items()) + ) + or "none" + ) + weights = buffer_bytes(program.weights) + arena = buffer_bytes(program.scratch) + per_node = ( + [f" * `{program.prefix}_node__run` runs one node on the caller's buffers."] + if program.nodes + else [] + ) + return [ + "/*", + f" * {program.prefix}.h - generated by FNNX {__version__} " + "(fnnx.extras.compilers.c)", + f" * from {comment_safe(program.source)}.", + " *", + " * Self-contained C99 inference code: no dynamic allocation, no I/O, " + + ("every buffer" if program.runtime_dims else "every shape"), + " * " + + ("sized" if program.runtime_dims else "fixed") + + " at compile time, and libm the only thing to link against. Builds clean", + " * under -std=c99 -Wall -Wextra -Werror -Werror=vla.", + " *", + " * Usage: include this header wherever the declarations are needed, and in exactly", + " * one translation unit ask for the implementation as well:", + " *", + f" * #define {upper}_IMPLEMENTATION", + f' * #include "{program.prefix}.h"', + " *", + f" * `{program.prefix}_run` runs the whole model, returning {upper}_OK on success.", + *per_node, + " * Every buffer is caller-provided, contiguous and row-major, of exactly the shape", + f" * the {upper}_* macros below describe.", + *_runtime_dim_preamble(program), + " *", + " * Not reentrant: the implementation owns static scratch buffers, so at most one", + " * call per compiled model may be in flight at a time.", + " *", + f" * Static memory: {weights + arena} bytes " + f"({weights} of weights, {arena} of scratch)", + f" * Opset imports: {opsets}", + f" * Dimension bindings: {bindings}", + *_runtime_dim_summary(program), + " */", + ] + + +def _runtime_dim_preamble(program: Program) -> list[str]: + if not program.runtime_dims: + return [] + upper = program.prefix.upper() + listed = ", ".join(f"`{dim.c_name}`" for dim in program.runtime_dims) + return [ + " *", + f" * Runtime dimensions: every entrypoint takes {listed} ahead of its buffers, as", + f" * the sizes this call works at. Each must lie in [1, {upper}_DIM_*_MAX]; outside", + f" * that the call returns {upper}_{INVALID_ARGUMENT_STATUS} and writes nothing.", + " * Buffers are sized for the maxima, but a call reads and writes only the leading", + " * elements of the shape its dimension values give — the shape below is the", + " * capacity, not what any one call passes.", + ] + + +def _runtime_dim_summary(program: Program) -> list[str]: + if not program.runtime_dims: + return [] + listed = ", ".join(f"{dim.name}<={dim.maximum}" for dim in program.runtime_dims) + return [f" * Runtime dimensions: {listed}"] + + +def _status_enum(program: Program) -> list[str]: + upper = program.prefix.upper() + return [ + "/* Status codes returned by the entrypoints. */", + f"typedef enum {program.prefix}_status {{", + f" {upper}_OK = 0,", + f" {upper}_{INVALID_ARGUMENT_STATUS} = 1", + f"}} {program.prefix}_status;", + "", + ] + + +def _footprint_macros(program: Program) -> list[str]: + upper = program.prefix.upper() + weights = buffer_bytes(program.weights) + arena = buffer_bytes(program.scratch) + return [ + "/* Static memory footprint, in bytes: constant data, scratch space, and the sum. */", + f"#define {upper}_WEIGHTS_BYTES {weights}", + f"#define {upper}_ARENA_BYTES {arena}", + f"#define {upper}_STATIC_BYTES {weights + arena}", + "", + ] + + +def _runtime_dim_macros(program: Program) -> list[str]: + lines: list[str] = [] + for dim in program.runtime_dims: + lines.append( + f"/* Runtime dimension `{comment_safe(dim.name)}`: the largest value the " + f"`{dim.c_name}` parameter may take. */" + ) + lines.append(f"#define {dim.macro(program.prefix)} {dim.maximum}") + lines.append("") + return lines + + +def _io_macros(program: Program) -> list[str]: + lines = _tensor_macros(program.inputs, program.outputs, owner="") + for node in program.nodes: + lines.extend( + _tensor_macros(node.inputs, node.outputs, owner=f" of node `{node.id}`") + ) + return lines + + +def _tensor_macros( + inputs: Iterable[IOTensor], outputs: Iterable[IOTensor], *, owner: str +) -> list[str]: + lines: list[str] = [] + for role, tensors in (("Input", inputs), ("Output", outputs)): + for tensor in tensors: + lines.append( + f"/* {role} `{comment_safe(tensor.name)}`{comment_safe(owner)}: " + f"{_tensor_label(tensor)}. */" + ) + lines.append(f"#define {tensor.macro}_RANK {len(tensor.shape)}") + lines.extend( + f"#define {tensor.macro}_DIM_{axis} {size}" + for axis, size in enumerate(tensor.shape) + ) + lines.append(f"#define {tensor.macro}_COUNT {tensor.elem_count}") + lines.append("") + return lines + + +def _tensor_label(tensor: IOTensor) -> str: + """The tensor's type, with any axis that scales named after the dimension it scales with.""" + if not tensor.runtime_shape: + return _type_label(tensor.elem_type, tensor.shape) + axes = [ + str(term.size) + if term.dim is None + else (term.dim if term.coefficient == 1 else f"{term.coefficient}*{term.dim}") + for term in tensor.runtime_shape + ] + return ( + f"{numpy_dtype_name(tensor.elem_type)}[{', '.join(axes)}], " + f"at most {_type_label(tensor.elem_type, tensor.shape)}" + ) + + +def _label_declarations(program: Program) -> list[str]: + """The class-label tables, declared: metadata for the caller, not data any kernel reads. + + They carry external linkage rather than being `static` like the buffers below: they are + part of what the artifact publishes, and a `static` array no emitted statement names + would fail the header's own `-Wall -Werror` build contract. + """ + lines: list[str] = [] + for table in program.labels: + lines.append( + f"/* Class labels of output `{comment_safe(table.tensor)}`, " + f"one per element of its trailing axis. */" + ) + lines.append(f"#define {table.macro} {len(table.values)}") + lines.append( + f"extern const {_label_c_type(table.elem_type)} " + f"{table.symbol}[{table.macro}];" + ) + lines.append("") + return lines + + +def _label_definitions(program: Program) -> list[str]: + lines: list[str] = [] + for table in program.labels: + lines.append( + f"const {_label_c_type(table.elem_type)} {table.symbol}[{table.macro}] = {{" + ) + lines.extend(initializer_lines(_label_literals(table))) + lines.append("};") + lines.append("") + return lines + + +def _label_literals(table: LabelTable) -> list[str]: + if table.elem_type == TensorProto.STRING: + return [string_literal(str(value)) for value in table.values] + return [scalar_literal(value, table.elem_type) for value in table.values] + + +def _label_c_type(elem_type: int) -> str: + return "char* const" if elem_type == TensorProto.STRING else c_type(elem_type) + + +def _entry_declarations(program: Program) -> list[str]: + lines = [ + "/* Runs the whole model. Returns one of the status codes above. */", + f"int {program.prefix}_run(" + f"{_signature(program, program.inputs, program.outputs)});", + "", + ] + for node in program.nodes: + lines.append( + f"/* Runs node `{comment_safe(node.id)}` on buffers of the caller's own. */" + ) + lines.append( + f"int {node.symbol}({_signature(program, node.inputs, node.outputs)});" + ) + lines.append("") + return lines + + +def _entry_definitions(program: Program) -> list[str]: + lines: list[str] = [] + # Node entrypoints first: the whole-model one calls them, and a reader follows the + # header top to bottom. + for node in program.nodes: + lines.extend( + _entry_definition( + program, node.symbol, node.inputs, node.outputs, node.body + ) + ) + lines.extend( + _entry_definition( + program, + f"{program.prefix}_run", + program.inputs, + program.outputs, + program.body, + ) + ) + return lines + + +def _entry_definition( + program: Program, + symbol: str, + inputs: tuple[IOTensor, ...], + outputs: tuple[IOTensor, ...], + body: Iterable[str], +) -> list[str]: + lines = [f"int {symbol}({_signature(program, inputs, outputs)})", "{"] + for statement in (*_dim_guard(program), *body): + lines.extend(f" {line}" if line else "" for line in statement.splitlines()) + lines.append(f" return {program.prefix.upper()}_OK;") + lines.append("}") + lines.append("") + return lines + + +def _dim_guard(program: Program) -> list[str]: + """Refuse a dimension value outside the range the artifact was compiled for. + + First thing in every entrypoint, before a single buffer is written: a rejected call + leaves the caller's outputs exactly as it found them. + """ + if not program.runtime_dims: + return [] + tests = " || ".join( + f"{dim.c_name} < 1 || {dim.c_name} > {dim.macro(program.prefix)}" + for dim in program.runtime_dims + ) + status = f"{program.prefix.upper()}_{INVALID_ARGUMENT_STATUS}" + return [f"if ({tests}) {{\n return {status};\n}}"] + + +def _signature( + program: Program, inputs: Iterable[IOTensor], outputs: Iterable[IOTensor] +) -> str: + parameters = [f"int32_t {dim.c_name}" for dim in program.runtime_dims] + parameters += [ + f"const {c_type(tensor.elem_type)}* {tensor.c_name}" for tensor in inputs + ] + parameters += [f"{c_type(tensor.elem_type)}* {tensor.c_name}" for tensor in outputs] + return ", ".join(parameters) if parameters else "void" + + +def _weight_definitions(program: Program) -> list[str]: + lines: list[str] = [] + for buffer in program.weights: + literals = list(_literals(buffer)) + if not literals: + literals = ["0"] + lines.append(_buffer_comment("Initializer", buffer)) + lines.append( + f"static const {c_type(buffer.elem_type)} " + f"{buffer.symbol}[{buffer.declared_count}] = {{" + ) + lines.extend(initializer_lines(literals)) + lines.append("};") + lines.append("") + return lines + + +def _scratch_definitions(program: Program) -> list[str]: + lines: list[str] = [] + for buffer in program.scratch: + lines.append(_buffer_comment("Intermediate", buffer)) + lines.append( + f"static {c_type(buffer.elem_type)} " + f"{buffer.symbol}[{buffer.declared_count}];" + ) + lines.append("") + return lines + + +def _kernel_definitions(program: Program) -> list[str]: + lines: list[str] = [] + for function in program.functions: + lines.extend(function.definition.splitlines()) + lines.append("") + return lines + + +def _literals(buffer: StaticBuffer) -> Iterator[str]: + assert buffer.tensor is not None + values = to_array(buffer.tensor).reshape(-1).tolist() + return (scalar_literal(value, buffer.elem_type) for value in values) + + +def _buffer_comment(role: str, buffer: StaticBuffer) -> str: + note = ( + " (zero-element; C99 forbids a zero-length array)" + if buffer.elem_count == 0 + else "" + ) + return ( + f"/* {role} `{comment_safe(buffer.name)}`: " + f"{_type_label(buffer.elem_type, buffer.shape)}{note}. */" + ) + + +def _type_label(elem_type: int, shape: tuple[int, ...]) -> str: + return f"{numpy_dtype_name(elem_type)}[{', '.join(str(size) for size in shape)}]" diff --git a/src/python/fnnx/extras/compilers/c/onnx/kernels.py b/src/python/fnnx/extras/compilers/c/onnx/kernels.py new file mode 100644 index 0000000..385910b --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/kernels.py @@ -0,0 +1,269 @@ +"""The kernel-generator convention and the registry ONNX kernels register into.""" + +from __future__ import annotations + +import hashlib +import math +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import onnx.defs +from onnx import NodeProto, TensorProto, helper +from onnx.numpy_helper import from_array, to_array + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.registry import ( + KernelRegistry, + latest_semantic_revision, +) + + +@dataclass(frozen=True) +class TensorRef: + """A tensor a kernel reads or writes, with the C expression naming its buffer.""" + + name: str + elem_type: int + shape: tuple[int, ...] + expr: str + + @property + def elem_count(self) -> int: + return math.prod(self.shape) + + +@dataclass(frozen=True) +class CFunction: + """A shared `static` kernel definition; `name` is what it is deduplicated by.""" + + name: str + definition: str + + +@dataclass(frozen=True) +class NodeContext: + """Everything a kernel generator needs to emit code for one node. + + `inputs`/`outputs` hold None where the node omits an optional operand. `prefix` is the + artifact-wide symbol prefix: kernel names build on it so that every static definition in + the emitted header is unique, while kernels shared between nodes still deduplicate. + `constants` holds the graph's compile-time values, for the operands an op reads as + configuration rather than as data. + """ + + node: NodeProto + domain: str + opset_version: int + since_version: int + prefix: str + inputs: tuple[TensorRef | None, ...] + outputs: tuple[TensorRef | None, ...] + constants: Mapping[str, TensorProto] = field(default_factory=dict) + + @property + def label(self) -> str: + return self.node.name or f"" + + def require_input(self, index: int) -> TensorRef: + return self._require(self.inputs, index, "input") + + def optional_input(self, index: int) -> TensorRef | None: + return self.inputs[index] if index < len(self.inputs) else None + + def require_output(self, index: int) -> TensorRef: + return self._require(self.outputs, index, "output") + + def attribute(self, name: str, default: Any) -> Any: + """The node's `name` attribute as a Python value, or `default` when absent.""" + for entry in self.node.attribute: + if entry.name == name: + return helper.get_attribute_value(entry) + return default + + def float_attribute(self, name: str) -> float: + """The node's `name` attribute, defaulting to the one the op's schema declares. + + Reading the default off the schema rather than restating it keeps a kernel from + drifting from the value ONNX's own tooling — and the reference evaluator — applies. + """ + return float( + self.attribute(name, self._schema().attributes[name].default_value.f) + ) + + def int_attribute(self, name: str) -> int: + """The node's `name` integer attribute, defaulting to the schema's own default.""" + return int( + self.attribute(name, self._schema().attributes[name].default_value.i) + ) + + def string_attribute(self, name: str) -> str: + """The node's `name` string attribute, defaulting to the schema's own default.""" + value = self.attribute(name, self._schema().attributes[name].default_value.s) + return value.decode("utf-8") if isinstance(value, bytes) else str(value) + + def constant_input(self, index: int) -> np.ndarray | None: + """The compile-time value of input `index`, or None when it has none. + + Only an operand every path through the graph fixes has one: an initializer, or a + tensor constant folding resolved. A runtime value reads as None. + """ + operand = self.inputs[index] if index < len(self.inputs) else None + if operand is None: + return None + tensor = self.constants.get(operand.name) + return None if tensor is None else to_array(tensor) + + def _schema(self) -> onnx.defs.OpSchema: + return onnx.defs.get_schema(self.node.op_type, self.since_version, self.domain) + + def _require( + self, operands: tuple[TensorRef | None, ...], index: int, role: str + ) -> TensorRef: + operand = operands[index] if index < len(operands) else None + if operand is None: + raise CompileError( + f"Node `{self.label}`: op `{self.node.op_type}` requires {role} {index}, " + "which this node leaves out." + ) + return operand + + +@dataclass(frozen=True) +class ScratchBuffer: + """Static working storage a kernel needs beyond the tensors it reads and writes. + + A few kernels cannot compute in place — a determinant eliminates on a copy of its matrix + — and the artifact allocates nothing, so the space is reserved at compile time and counted + in the reported footprint like every other buffer. `symbol` is what it is deduplicated by: + the nodes that share a kernel share its buffer, sized for the largest of them, which is + safe under the artifact's one-call-at-a-time contract. + """ + + symbol: str + elem_type: int + elem_count: int + + +@dataclass(frozen=True) +class ConstantData: + """A table a kernel reads from `static const` storage rather than from an operand. + + ONNX-ML carries an op's parameters in attributes rather than in inputs — a scaler's + per-feature offsets, an encoder's categories, an ensemble's nodes — and those tables run + long enough that passing them as compound literals would put them on the stack and leave + them out of the reported footprint. `symbol` is what the data is deduplicated by, so it + has to encode the contents; `constant_data` builds one that does. + """ + + symbol: str + tensor: TensorProto + + +@dataclass(frozen=True) +class NodeEmission: + """What a kernel generator contributes: shared functions plus the call site.""" + + functions: tuple[CFunction, ...] + statements: tuple[str, ...] + scratch: tuple[ScratchBuffer, ...] = () + constants: tuple[ConstantData, ...] = () + + +KernelGenerator = Callable[[NodeContext], NodeEmission] + +KERNELS: KernelRegistry[KernelGenerator] = KernelRegistry() + + +def register_kernel( + domain: str, op_type: str, versions: Iterable[int], generator: KernelGenerator +) -> None: + """Register `generator` at each listed schema revision the installed `onnx` defines. + + ONNX bumps an op's `since_version` on every spec change, including the type-constraint + additions that leave the emitted code identical, so one generator usually covers several + revisions — and each has to be registered, or the semantic-revision guard rejects the + kernel at the newer opset. `versions` is therefore the explicit claim of which revisions + this generator implements: ones the installed `onnx` package does not define are skipped + (keeping kernels installable across the supported `onnx` range), and ones that are not + listed are left to the guard rather than silently served with older semantics. + """ + for version in versions: + if latest_semantic_revision(domain, op_type, version) == version: + KERNELS.register(domain, op_type, version, generator) + + +def copy_tensor(source: TensorRef, result: TensorRef) -> NodeEmission: + """`source`'s elements written into `result`, where no per-element code is needed. + + A straight copy: an identity, or a reinterpretation of the same bytes at another element + type of the same width. A shared kernel would not earn its keep over one `memcpy`. + """ + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + return NodeEmission( + functions=(), + statements=( + f"memcpy({result.expr}, {source.expr}, " + f"{result.elem_count}u * sizeof(*{result.expr}));", + ), + ) + + +def constant_data( + context: NodeContext, role: str, values: np.ndarray +) -> tuple[ConstantData, str]: + """`values` as static constant data, and the C expression naming it. + + The symbol carries a digest of the contents, so two nodes reading the same table share + one definition and two different tables can never collide on a name. + """ + array = np.ascontiguousarray(values) + digest = hashlib.sha256( + f"{array.dtype.str}{array.shape}".encode() + array.tobytes() + ) + symbol = ( + f"{context.prefix}_{context.node.op_type.lower()}_{role}_" + f"{digest.hexdigest()[:12]}" + ) + return ConstantData(symbol, from_array(array, symbol)), symbol + + +def broadcast_strides( + source: TensorRef, shape: tuple[int, ...], *, node_label: str +) -> tuple[int, ...]: + """Row-major strides addressing `source` while iterating a tensor of `shape`. + + A stride is zero on every axis `source` is broadcast along, so the same element is read + for every coordinate on that axis; `source` is aligned to the trailing axes, as ONNX's + broadcasting rules prescribe. + """ + if len(source.shape) > len(shape): + raise _broadcast_error(source, shape, node_label) + padded = (1,) * (len(shape) - len(source.shape)) + source.shape + strides = [] + stride = 1 + for size, target in zip(reversed(padded), reversed(shape)): + if size == target: + strides.append(stride) + elif size == 1: + strides.append(0) + else: + raise _broadcast_error(source, shape, node_label) + stride *= size + return tuple(reversed(strides)) + + +def _broadcast_error( + source: TensorRef, shape: tuple[int, ...], node_label: str +) -> CompileError: + return CompileError( + f"Node `{node_label}`: tensor `{source.name}` of shape {list(source.shape)} does " + f"not broadcast to {list(shape)}." + ) + + +# Imported for the side effect of registering the kernels; the import sits at the bottom +# because every op module builds on the definitions above. +from fnnx.extras.compilers.c.onnx import ops # noqa: E402, F401 diff --git a/src/python/fnnx/extras/compilers/c/onnx/loader.py b/src/python/fnnx/extras/compilers/c/onnx/loader.py new file mode 100644 index 0000000..292887b --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/loader.py @@ -0,0 +1,161 @@ +"""Model loading, external-data resolution, and per-domain opset resolution.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path + +import onnx +import onnx.defs +from onnx import ModelProto + +# onnx exposes no public walker over every tensor of a model (initializers plus the +# ones nested in subgraph attributes), which is what external-data detection needs. +from onnx.external_data_helper import ( + _get_all_tensors, + load_external_data_for_model, + uses_external_data, +) + +from fnnx.extras.compilers.c.errors import CompileError + +STANDARD_DOMAIN = "" +ML_DOMAIN = "ai.onnx.ml" +SUPPORTED_DOMAINS = (STANDARD_DOMAIN, ML_DOMAIN) + +_DOMAIN_ALIASES = {"ai.onnx": STANDARD_DOMAIN} + + +def normalize_domain(domain: str) -> str: + return _DOMAIN_ALIASES.get(domain, domain) + + +def display_domain(domain: str) -> str: + normalized = normalize_domain(domain) + return "ai.onnx" if normalized == STANDARD_DOMAIN else normalized + + +@lru_cache(maxsize=None) +def max_supported_opset(domain: str) -> int: + """Highest opset version the installed `onnx` package defines for `domain`.""" + normalized = normalize_domain(domain) + if normalized == STANDARD_DOMAIN: + return onnx.defs.onnx_opset_version() + versions = [ + schema.since_version + for schema in onnx.defs.get_all_schemas_with_history() + if schema.domain == normalized + ] + if not versions: + raise ValueError( + f"The installed `onnx` package defines no schemas for domain `{normalized}`." + ) + return max(versions) + + +@dataclass(frozen=True) +class LoadedModel: + model: ModelProto + opsets: dict[str, int] + + def opset_for(self, domain: str) -> int: + normalized = normalize_domain(domain) + version = self.opsets.get(normalized) + if version is None: + raise CompileError( + f"Model imports no opset for domain `{display_domain(normalized)}`." + ) + return version + + +def load_model( + source: str | os.PathLike[str] | ModelProto, + *, + base_dir: str | os.PathLike[str] | None = None, +) -> LoadedModel: + """Load an ONNX model, embed its external data, and resolve its opset imports. + + `base_dir` is where external tensor files are looked up; it defaults to the directory + holding `source` when a path is given, and is required for an in-memory proto whose + tensors live in external files. The caller's proto is never mutated. + """ + if isinstance(source, ModelProto): + model = ModelProto() + model.CopyFrom(source) + data_dir = Path(base_dir) if base_dir is not None else None + else: + path = Path(source) + if not path.is_file(): + raise CompileError(f"ONNX model file not found: `{path}`.") + try: + model = onnx.load_model(os.fspath(path), load_external_data=False) + except Exception as exc: + raise CompileError(f"Failed to parse ONNX model `{path}`: {exc}") from exc + data_dir = Path(base_dir) if base_dir is not None else path.parent + + if model.ir_version > onnx.IR_VERSION: + raise CompileError( + f"Model IR version {model.ir_version} is newer than the installed `onnx` package " + f"({onnx.__version__}) supports (at most {onnx.IR_VERSION}); " + "upgrade `onnx` to compile this model." + ) + _embed_external_data(model, data_dir) + return LoadedModel(model=model, opsets=resolve_opsets(model)) + + +def resolve_opsets(model: ModelProto) -> dict[str, int]: + """Map every imported domain, normalized, to the opset version the model requests.""" + opsets: dict[str, int] = {} + for imported in model.opset_import: + domain = normalize_domain(imported.domain) + if domain not in SUPPORTED_DOMAINS: + supported = " and ".join(display_domain(d) for d in SUPPORTED_DOMAINS) + raise CompileError( + f"Model imports unsupported opset domain `{imported.domain}`; " + f"the C compiler supports only {supported}." + ) + if imported.version < 1: + raise CompileError( + f"Model imports invalid opset version {imported.version} " + f"for domain `{display_domain(domain)}`." + ) + maximum = max_supported_opset(domain) + if imported.version > maximum: + raise CompileError( + f"Model imports opset version {imported.version} for domain " + f"`{display_domain(domain)}`, but the installed `onnx` package " + f"({onnx.__version__}) defines at most version {maximum} for that domain; " + "upgrade `onnx` to compile this model." + ) + previous = opsets.get(domain) + if previous is not None and previous != imported.version: + raise CompileError( + f"Model imports conflicting opset versions {previous} and {imported.version} " + f"for domain `{display_domain(domain)}`." + ) + opsets[domain] = imported.version + if not opsets: + raise CompileError("Model imports no opsets.") + return opsets + + +def _embed_external_data(model: ModelProto, data_dir: Path | None) -> None: + external = [ + tensor for tensor in _get_all_tensors(model) if uses_external_data(tensor) + ] + if not external: + return + if data_dir is None: + names = ", ".join(f"`{tensor.name}`" for tensor in external[:3]) + raise CompileError( + f"Model stores tensors ({names}) in external files, but no directory to resolve " + "them against is known; pass `base_dir` when compiling an in-memory model." + ) + try: + load_external_data_for_model(model, os.fspath(data_dir)) + except Exception as exc: + raise CompileError( + f"Failed to load external tensor data from `{data_dir}`: {exc}" + ) from exc diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/__init__.py b/src/python/fnnx/extras/compilers/c/onnx/ops/__init__.py new file mode 100644 index 0000000..a7fab97 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/__init__.py @@ -0,0 +1,33 @@ +"""Kernel generators, imported for the side effect of registering them.""" + +from fnnx.extras.compilers.c.onnx.ops import ( # noqa: F401 + activations, + attention, + casts, + conv, + cumulative, + einsum, + elementwise, + gather, + gemm, + generate, + linear_attention, + logic, + loss, + ml, + normalization, + pad, + pool, + quantize, + recurrent, + reduce, + resize, + sampling, + scatter, + signal, + softmax, + svm, + tfidf, + tree, + views, +) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/activations.py b/src/python/fnnx/extras/compilers/c/onnx/ops/activations.py new file mode 100644 index 0000000..98f3000 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/activations.py @@ -0,0 +1,229 @@ +"""Activation kernels, each emitted as the formula the ONNX spec defines it by.""" + +from __future__ import annotations + +import math +from functools import partial +from string import Template + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import FLOAT_TYPES, c_type +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import ( + Scalar, + combiner, + elementwise, + expand, + pointwise, +) + +# Revision 1 of the older activations carried the legacy `consumed_inputs` attribute; the +# revisions listed after the first are the ones that only widened type constraints. +_RELU_VERSIONS = (6, 13, 14) +_SIGMOID_VERSIONS = (6, 13) +_LEAKY_RELU_VERSIONS = (6, 16) +# PRelu-7 replaced the legacy broadcast attributes with numpy-style broadcasting of `slope`. +_PRELU_VERSIONS = (7, 9, 16) +_ELU_VERSIONS = (6, 22) +_SELU_VERSIONS = (6, 22) +_CELU_VERSIONS = (12,) +_HARD_SIGMOID_VERSIONS = (6, 22) +_SOFTPLUS_VERSIONS = (1, 22) +_SOFTSIGN_VERSIONS = (1, 22) +_SHRINK_VERSIONS = (9,) +_THRESHOLDED_RELU_VERSIONS = (10, 22) +_GELU_VERSIONS = (20,) + +# The constants of Gelu's two formulas, which ONNX writes as `sqrt(2)`, `sqrt(2/pi)` and +# `0.044715`; the square roots are taken in the element type, as the op's function body +# takes them. +_TWO_OVER_PI = 2.0 / math.pi +_GELU_CUBIC = 0.044715 +_GELU_MODES = ("none", "tanh") + +# The reference's numerically stable form: whichever branch is taken, the exponent is of a +# negative number, so it underflows to zero instead of overflowing to infinity. +_SIGMOID_TEMPLATE = Template("""\ +static $element $name($element x) +{ + return x > $zero + ? $one / ($one + exp$f(-x)) + : exp$f(x) / ($one + exp$f(x)); +}""") + +# Activations whose formula is a plain expression over the operand and its attributes. +# `$element`, `$zero`, `$one` and `$f` come from the element type; the attribute names are +# passed to the kernel as parameters, so one kernel serves every node running the op. +_ATTRIBUTE_ACTIVATIONS: dict[str, tuple[tuple[int, ...], tuple[str, ...], str]] = { + "Elu": ( + _ELU_VERSIONS, + ("alpha",), + "(x0 > $zero) ? x0 : alpha * (exp$f(x0) - $one)", + ), + "LeakyRelu": (_LEAKY_RELU_VERSIONS, ("alpha",), "(x0 > $zero) ? x0 : x0 * alpha"), + "Selu": ( + _SELU_VERSIONS, + ("alpha", "gamma"), + "((x0 > $zero) ? x0 : exp$f(x0) * alpha - alpha) * gamma", + ), + "ThresholdedRelu": ( + _THRESHOLDED_RELU_VERSIONS, + ("alpha",), + "(x0 > alpha) ? x0 : $zero", + ), +} + +# The same, without attributes. +_PLAIN_ACTIVATIONS: dict[str, tuple[tuple[int, ...], str]] = { + "Softplus": (_SOFTPLUS_VERSIONS, "log$f(exp$f(x0) + $one)"), + "Softsign": (_SOFTSIGN_VERSIONS, "x0 / (fabs$f(x0) + $one)"), +} + + +def _relu(context: NodeContext) -> NodeEmission: + result = context.require_output(0) + # The spec is `max(0, x)`, evaluated as numpy's `maximum`: NaN propagates and -0 comes + # out as +0, both of which a plain `value > 0 ? value : 0` would get wrong. + guard = ( + "x0 > $zero || isnan(x0)" if result.elem_type in FLOAT_TYPES else "x0 > $zero" + ) + return pointwise(context, f"({guard}) ? x0 : $zero") + + +def _sigmoid(context: NodeContext) -> NodeEmission: + result = context.require_output(0) + name = f"{context.prefix}_sigmoid_{c_type(result.elem_type)}" + helper = CFunction( + name, + expand(_SIGMOID_TEMPLATE.safe_substitute(name=name), result.elem_type), + ) + return pointwise(context, f"{name}(x0)", helpers=(helper,)) + + +def _attribute_activation( + context: NodeContext, *, names: tuple[str, ...], template: str +) -> NodeEmission: + result = context.require_output(0) + return pointwise( + context, + template, + scalars=tuple( + Scalar(name, result.elem_type, context.float_attribute(name)) + for name in names + ), + ) + + +def _prelu(context: NodeContext) -> NodeEmission: + result = context.require_output(0) + return elementwise( + context, + expression=expand("(x0 > $zero) ? x0 : x0 * x1", result.elem_type), + operands=(context.require_input(0), context.require_input(1)), + result=result, + ) + + +def _celu(context: NodeContext) -> NodeEmission: + """Celu: `max(0, x) + min(0, alpha * (exp(x / alpha) - 1))`, as ONNX defines it.""" + result = context.require_output(0) + largest = combiner(context, result.elem_type, largest=True) + smallest = combiner(context, result.elem_type, largest=False) + return pointwise( + context, + f"{largest.name}($zero, x0) + " + f"{smallest.name}($zero, alpha * (exp$f(x0 / alpha) - $one))", + scalars=(Scalar("alpha", result.elem_type, context.float_attribute("alpha")),), + helpers=(largest, smallest), + ) + + +def _hard_sigmoid(context: NodeContext) -> NodeEmission: + result = context.require_output(0) + largest = combiner(context, result.elem_type, largest=True) + smallest = combiner(context, result.elem_type, largest=False) + return pointwise( + context, + f"{largest.name}($zero, {smallest.name}($one, x0 * alpha + beta))", + scalars=tuple( + Scalar(name, result.elem_type, context.float_attribute(name)) + for name in ("alpha", "beta") + ), + helpers=(largest, smallest), + ) + + +def _gelu(context: NodeContext) -> NodeEmission: + """Gelu in whichever of its two forms the `approximate` attribute selects. + + Grouped as `(0.5 * x) * (1 + phi)`, the order of the function body ONNX defines the op + by: at a large negative x the second factor underflows to zero, and only this grouping + leaves the sign on the zero that comes out of it. + """ + result = context.require_output(0) + constant = partial(scalar_literal, elem_type=result.elem_type) + if _approximate(context) == "tanh": + phi = ( + f"tanh$f(sqrt$f({constant(_TWO_OVER_PI)}) * " + f"(x0 + {constant(_GELU_CUBIC)} * pow$f(x0, {constant(3.0)})))" + ) + variant = "_tanh" + else: + phi = f"erf$f(x0 / sqrt$f({constant(2.0)}))" + variant = "_erf" + return pointwise( + context, f"({constant(0.5)} * x0) * ($one + {phi})", variant=variant + ) + + +def _approximate(context: NodeContext) -> str: + value = context.attribute("approximate", b"none") + mode = value.decode() if isinstance(value, bytes) else str(value) + if mode not in _GELU_MODES: + raise CompileError( + f"Node `{context.label}`: Gelu's `approximate` attribute is `{mode}`, but " + f"ONNX defines only {' and '.join(f'`{name}`' for name in _GELU_MODES)}." + ) + return mode + + +def _shrink(context: NodeContext) -> NodeEmission: + """Shrink, which the reference evaluates in double for the integer element types too. + + Adding the bias in double and rounding once on the way back gives the element type's own + arithmetic for the floating-point families, and numpy's promotion for the integer ones. + """ + return pointwise( + context, + "($element)((x0 < -lambd) ? (x0 + bias) : ((x0 > lambd) ? (x0 - bias) : 0))", + scalars=tuple( + Scalar(name, TensorProto.DOUBLE, context.float_attribute(name)) + for name in ("lambd", "bias") + ), + ) + + +for _op_type, (_versions, _names, _template) in _ATTRIBUTE_ACTIVATIONS.items(): + register_kernel( + "", + _op_type, + _versions, + partial(_attribute_activation, names=_names, template=_template), + ) +for _op_type, (_versions, _template) in _PLAIN_ACTIVATIONS.items(): + register_kernel("", _op_type, _versions, partial(pointwise, template=_template)) +register_kernel("", "Relu", _RELU_VERSIONS, _relu) +register_kernel("", "Sigmoid", _SIGMOID_VERSIONS, _sigmoid) +register_kernel("", "PRelu", _PRELU_VERSIONS, _prelu) +register_kernel("", "Celu", _CELU_VERSIONS, _celu) +register_kernel("", "HardSigmoid", _HARD_SIGMOID_VERSIONS, _hard_sigmoid) +register_kernel("", "Gelu", _GELU_VERSIONS, _gelu) +register_kernel("", "Shrink", _SHRINK_VERSIONS, _shrink) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/attention.py b/src/python/fnnx/extras/compilers/c/onnx/ops/attention.py new file mode 100644 index 0000000..b082179 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/attention.py @@ -0,0 +1,1087 @@ +"""Attention and RotaryEmbedding: the transformer primitives, as one loop nest each. + +ONNX defines both as functions, and the compiler's default for a function is to inline its +body. That was tried and does not work here: both bodies compute the shapes they reshape, +slice and expand to through `Shape`/`Size`/`Concat`/`Where` chains written at opset 23, and +constant folding refuses to execute a node whose (op, opset) the reference evaluator is not +version-faithful for — which, at the opset those bodies are written in, is every one of +them. The shape operands stay run-time values, the static-shape verifier rejects the body, +and nothing is compiled. So these two get native kernels, which is exactly the case the +"native kernels only where expansion proves insufficient" rule is for. + +An attention head is one loop nest: for each (batch item, query head, query position) a row +of scores over the key positions, softmaxed, and read back as a weighted sum of the value +rows. Everything that differs between two nodes is geometry rather than code — how the +operands are laid out, how long the sequences are, which optional operands are present — so +one shared kernel per (element type, softmax precision, mask flavour) serves every node, +with the strides, the extents and the attribute values as arguments. + +The two layouts ONNX allows — `(batch, head, sequence, size)` and +`(batch, sequence, head * size)` — are one tensor at two sets of strides, so they do not +fork the kernel: the call site passes the batch, head and sequence strides of whichever it +has, and the head size is contiguous in both. + +**The arithmetic is `double` whatever the tensors hold.** The reference evaluator scales `Q` +and `K` by a numpy float64 scalar, which promotes everything from the QK product through the +softmax to float64, and rounds back to the tensor's own type only at the very end; a kernel +computing a float32 model in float32 would be a different operator. `softmax_precision` is +the one place ONNX narrows that back, and it is what the softmax type is read from. + +`past_key`/`past_value` are never concatenated into working storage: the kernel reads a key +position out of the past cache or out of `K` depending on where it falls. The concatenation +ONNX calls `present_key`/`present_value` is written by a small kernel of its own, and only +when the node asks for it. + +RotaryEmbedding rotates `rotary_embedding_dim / 2` coordinate pairs per (batch item, +position, head) and copies the tail beyond them through unchanged. Which two lanes a pair +is, and where the rotated pair is written back, is the whole of what `interleaved` changes, +so that too is an argument rather than a second kernel. + +Where the reference evaluator and the schema prose disagree, the reference is what is +compiled — both test suites take their expected values from it — and every place that +happens is commented below. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from string import Template + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + FLOAT_TYPES, + c_type, + element_type_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + ScratchBuffer, + TensorRef, + broadcast_strides, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + checked_call, + kernel_name, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import combiner, extents, math_suffix + +# Attention arrived at opset 23 and gained `nonpad_kv_seqlen` at 24; the kernel serves both, +# reading the operand as absent at the older one. RotaryEmbedding has had one revision. +_ATTENTION_VERSIONS = (23, 24) +_ROTARY_VERSIONS = (23,) + +_PAST_KEY_INPUT = 4 +_PAST_VALUE_INPUT = 5 +_NONPAD_INPUT = 6 + +# The four intermediates `qk_matmul_output` may carry. ONNX's text distinguishes 0 ("the +# output of qk matmul") from 1 ("the output after the softcap operation"), but the reference +# evaluator softcaps the tensor it reports before it ever looks at the mode, so the two are +# one and the same value; the reference is what is compiled. +_OUTPUT_MODES = (0, 1, 2, 3) + + +# -------------------------------------------------------------------------------------- +# Attention +# -------------------------------------------------------------------------------------- + +# The bias added to one score. `attn_mask`'s last axis is the key positions, and ONNX pads it +# out to the total key length rather than broadcasting it, so it is addressed here while +# every axis before it arrives as a broadcast stride. The causal offset is the strict upper +# triangle of the columns from `past_seq` on, which is where the reference's `np.triu(..., 1)` +# lands, and `nonpad_kv_seqlen` masks off every column at or past a batch item's own length. +# `causal_stride` is which row of that triangle a query row reads: see `_causal_row_stride`. +_BIAS_TEMPLATE = Template("""\ +static double $name( + const $mask* mask, + const size_t* strides, + size_t mask_len, + const int64_t* nonpad, + size_t item, + size_t head, + size_t row, + size_t column, + size_t past_seq, + int causal, + size_t causal_stride) +{ + double bias = 0.0; + if (mask != NULL) { + const size_t base = + item * strides[0] + head * strides[1] + row * strides[2]; +$read + } + if (causal && column >= past_seq + && column - past_seq > row * causal_stride) { + bias += -INFINITY; + } + if (nonpad != NULL && (int64_t)column >= nonpad[item]) { + bias += -INFINITY; + } + return bias; +}""") + +# A float mask is the bias itself, and the columns past its own length are the negative +# infinity ONNX pads it with. +_ADDITIVE_MASK = """\ + bias = column < mask_len ? (double)mask[base + column] : -INFINITY;""" + +# A boolean mask marks the entries that take part. The reference evaluator writes the two +# cases as different expressions — `(1 - mask) * -inf` under `is_causal`, `(1 - mask)` with +# its ones replaced by -inf otherwise — and the first turns an entry that *does* take part +# into `0 * -inf`, which is a NaN that then poisons its whole softmax row. That asymmetry is +# not in the prose; it is what the oracle computes, so it is what is compiled. +_BOOLEAN_MASK = """\ + const int taken = column < mask_len && mask[base + column] != 0; + bias = taken ? (causal ? (double)NAN : 0.0) : -INFINITY;""" + +# One query row against every key position: the scaled dot products, the softcap, the bias, +# the softmax, and the value rows read back through it. `scale` is already the square root +# the reference takes of it, and it multiplies `Q` and `K` separately rather than the product +# once — an algebraically equal rearrangement would round differently. +_ATTENTION_TEMPLATE = Template("""\ +static void $name( + $element* y, + $element* qk_out, + const $element* q, + const $element* k, + const $element* v, + const $element* past_k, + const $element* past_v, + const $mask* mask, + const int64_t* nonpad, + $soft* scores, + const size_t* q_strides, + const size_t* k_strides, + const size_t* v_strides, + const size_t* y_strides, + const size_t* mask_strides, + size_t batch, + size_t q_heads, + size_t kv_heads, + size_t q_seq, + size_t kv_seq, + size_t past_seq, + size_t head_size, + size_t v_head_size, + size_t repeats, + size_t mask_len, + double scale, + double softcap, + int causal, + size_t causal_stride, + int mode) +{ + const size_t total = past_seq + kv_seq; + size_t item, head, row, column, lane; + for (item = 0; item < batch; ++item) { + for (head = 0; head < q_heads; ++head) { + /* Grouped-query attention repeats each key/value head `repeats` times over + adjacent query heads, which is what `np.repeat` interleaves them as. */ + const size_t kv_head = head / repeats; + const size_t past_base = (item * kv_heads + kv_head) * past_seq; + const $element* q_head = q + item * q_strides[0] + head * q_strides[1]; + const $element* k_head = k + item * k_strides[0] + kv_head * k_strides[1]; + const $element* v_head = v + item * v_strides[0] + kv_head * v_strides[1]; + $element* y_head = y + item * y_strides[0] + head * y_strides[1]; + for (row = 0; row < q_seq; ++row) { + const $element* q_row = q_head + row * q_strides[2]; + const size_t reported = + ((item * q_heads + head) * q_seq + row) * total; + $soft largest = -INFINITY; + $soft weight = $soft_zero; + for (column = 0; column < total; ++column) { + const $element* k_row = column < past_seq + ? past_k + (past_base + column) * head_size + : k_head + (column - past_seq) * k_strides[2]; + double score = 0.0; + for (lane = 0; lane < head_size; ++lane) { + score += ((double)q_row[lane] * scale) + * ((double)k_row[lane] * scale); + } + if (softcap > 0.0) { + score = tanh(score / softcap) * softcap; + } + if (qk_out != NULL && mode <= 1) { + qk_out[reported + column] = ($element)score; + } + score += $bias(mask, mask_strides, mask_len, nonpad, + item, head, row, column, past_seq, causal, causal_stride); + if (qk_out != NULL && mode == 2) { + qk_out[reported + column] = ($element)score; + } + scores[column] = ($soft)score; + } + /* Max-subtracted, as the reference's softmax is: a row that is entirely + -inf therefore leaves `-inf - -inf`, and comes out NaN. */ + for (column = 0; column < total; ++column) { + largest = $maximum(largest, scores[column]); + } + for (column = 0; column < total; ++column) { + scores[column] = exp$soft_suffix(scores[column] - largest); + weight += scores[column]; + } + for (column = 0; column < total; ++column) { + scores[column] /= weight; + if (qk_out != NULL && mode == 3) { + qk_out[reported + column] = ($element)scores[column]; + } + } + for (lane = 0; lane < v_head_size; ++lane) { + $product weighted = $product_zero; + for (column = 0; column < total; ++column) { + const $element* v_row = column < past_seq + ? past_v + (past_base + column) * v_head_size + : v_head + (column - past_seq) * v_strides[2]; + weighted += + ($product)scores[column] * ($product)v_row[lane]; + } + y_head[row * y_strides[2] + lane] = ($element)weighted; + } + } + } + } +}""") + +# `present_key` and `present_value` are the past cache and the incoming keys or values +# concatenated along the sequence axis — always as the 4-D cache layout, whichever layout the +# incoming operand itself has, which is what the strides are for. +_PRESENT_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* past, + const $element* current, + const size_t* strides, + size_t batch, + size_t heads, + size_t past_seq, + size_t kv_seq, + size_t head_size) +{ + const size_t total = past_seq + kv_seq; + size_t item, head, step, lane; + for (item = 0; item < batch; ++item) { + for (head = 0; head < heads; ++head) { + const $element* source = + current + item * strides[0] + head * strides[1]; + $element* written = out + (item * heads + head) * total * head_size; + for (step = 0; step < past_seq; ++step) { + const size_t cached = + ((item * heads + head) * past_seq + step) * head_size; + for (lane = 0; lane < head_size; ++lane) { + written[step * head_size + lane] = past[cached + lane]; + } + } + for (step = 0; step < kv_seq; ++step) { + for (lane = 0; lane < head_size; ++lane) { + written[(past_seq + step) * head_size + lane] = + source[step * strides[2] + lane]; + } + } + } + } +}""") + + +@dataclass(frozen=True) +class _Geometry: + """A node's extents, and where each operand's elements sit under its layout. + + Rank 4 packs `(batch, head, sequence, size)` and rank 3 `(batch, sequence, head * size)`; + the kernel walks both through one batch, head and sequence stride, since the head size is + contiguous either way. The caches and `qk_matmul_output` are 4-D whichever rank the node + reads, which is what `cache_shape` and `qk_shape` state. + """ + + rank: int + batch: int + q_heads: int + kv_heads: int + q_seq: int + kv_seq: int + past_seq: int + head_size: int + v_head_size: int + + @property + def total_seq(self) -> int: + return self.past_seq + self.kv_seq + + @property + def repeats(self) -> int: + """How many query heads share one key/value head, at least one.""" + return max(1, self.q_heads // self.kv_heads) if self.kv_heads else 1 + + @property + def rows(self) -> int: + """How many score rows the node computes; nothing is emitted for none at all.""" + return self.batch * self.q_heads * self.q_seq + + def strides(self, heads: int, seq: int, size: int) -> tuple[int, int, int]: + """`(batch, head, sequence)` strides of an operand at this node's layout.""" + if self.rank == 4: + return (heads * seq * size, seq * size, size) + return (seq * heads * size, size, heads * size) + + @property + def y_shape(self) -> tuple[int, ...]: + if self.rank == 4: + return (self.batch, self.q_heads, self.q_seq, self.v_head_size) + return (self.batch, self.q_seq, self.q_heads * self.v_head_size) + + @property + def qk_shape(self) -> tuple[int, ...]: + return (self.batch, self.q_heads, self.q_seq, self.total_seq) + + def cache_shape(self, size: int) -> tuple[int, ...]: + return (self.batch, self.kv_heads, self.total_seq, size) + + +def _attention(context: NodeContext) -> NodeEmission: + geometry = _geometry(context) + element = _element_type(context) + results = tuple( + context.outputs[index] if index < len(context.outputs) else None + for index in range(4) + ) + verify_shape(context, context.require_output(0), geometry.y_shape) + present = _present_emission(context, geometry, element, results) + scoring = _scoring_emission(context, geometry, element, results) + return NodeEmission( + functions=present.functions + scoring.functions, + statements=present.statements + scoring.statements, + scratch=scoring.scratch, + ) + + +def _present_emission( + context: NodeContext, + geometry: _Geometry, + element: int, + results: tuple[TensorRef | None, ...], +) -> NodeEmission: + """The key and value caches the node asks to have written back.""" + name = kernel_name(context, "present", c_type(element)) + # `present_key` is output 1 and reads `K`, which is input 1; `present_value` is output 2 + # and reads `V`, which is input 2 — the same index on either side. + caches = ( + (1, _PAST_KEY_INPUT, geometry.head_size), + (2, _PAST_VALUE_INPUT, geometry.v_head_size), + ) + statements = [] + for slot, past_index, size in caches: + result = results[slot] + if result is None: + continue + verify_shape(context, result, geometry.cache_shape(size)) + if result.elem_count == 0: + continue + current = context.require_input(slot) + past = context.optional_input(past_index) + statements.append( + call_kernel( + name, + [ + result.expr, + "NULL" if past is None else past.expr, + current.expr, + extents(geometry.strides(geometry.kv_heads, geometry.kv_seq, size)), + f"{geometry.batch}u", + f"{geometry.kv_heads}u", + f"{geometry.past_seq}u", + f"{geometry.kv_seq}u", + f"{size}u", + ], + ) + ) + if not statements: + return NodeEmission(functions=(), statements=()) + definition = _PRESENT_TEMPLATE.substitute(name=name, element=c_type(element)) + return NodeEmission( + functions=(CFunction(name, definition),), statements=tuple(statements) + ) + + +def _scoring_emission( + context: NodeContext, + geometry: _Geometry, + element: int, + results: tuple[TensorRef | None, ...], +) -> NodeEmission: + """The attention itself: the kernel, the row of scores it works in, and the call.""" + mask = context.optional_input(3) + mask_type = _mask_type(context, mask, element) + soft = _softmax_type(context) + product = _product_type(soft, element) + mode = _output_mode(context) + reported = results[3] + if reported is not None: + verify_shape(context, reported, geometry.qk_shape) + if geometry.rows == 0: + return NodeEmission(functions=(), statements=()) + + bias = _bias_function(context, mask_type) + largest = combiner(context, soft, largest=True) + name = kernel_name( + context, c_type(element), f"soft{c_type(soft)}", f"mask{c_type(mask_type)}" + ) + definition = _ATTENTION_TEMPLATE.substitute( + name=name, + element=c_type(element), + mask=c_type(mask_type), + soft=c_type(soft), + product=c_type(product), + soft_zero=scalar_literal(0, soft), + product_zero=scalar_literal(0, product), + soft_suffix=math_suffix(soft), + maximum=largest.name, + bias=bias.name, + ) + scratch = ScratchBuffer( + kernel_name(context, "scores", c_type(soft)), soft, geometry.total_seq + ) + arguments = _scoring_arguments(context, geometry, mask, reported, scratch, mode) + return NodeEmission( + functions=(bias, largest, CFunction(name, definition)), + statements=(call_kernel(name, arguments),), + scratch=(scratch,), + ) + + +def _scoring_arguments( + context: NodeContext, + geometry: _Geometry, + mask: TensorRef | None, + reported: TensorRef | None, + scratch: ScratchBuffer, + mode: int, +) -> list[str]: + optional = { + index: context.optional_input(index) + for index in (_PAST_KEY_INPUT, _PAST_VALUE_INPUT, _NONPAD_INPUT) + } + causal = int(context.int_attribute("is_causal") != 0) + return [ + context.require_output(0).expr, + "NULL" if reported is None else reported.expr, + *(context.require_input(index).expr for index in range(3)), + *( + "NULL" if operand is None else operand.expr + for operand in ( + optional[_PAST_KEY_INPUT], + optional[_PAST_VALUE_INPUT], + mask, + optional[_NONPAD_INPUT], + ) + ), + scratch.symbol, + extents(geometry.strides(geometry.q_heads, geometry.q_seq, geometry.head_size)), + extents( + geometry.strides(geometry.kv_heads, geometry.kv_seq, geometry.head_size) + ), + extents( + geometry.strides(geometry.kv_heads, geometry.kv_seq, geometry.v_head_size) + ), + extents( + geometry.strides(geometry.q_heads, geometry.q_seq, geometry.v_head_size) + ), + extents(_mask_strides(context, mask, geometry)), + f"{geometry.batch}u", + f"{geometry.q_heads}u", + f"{geometry.kv_heads}u", + f"{geometry.q_seq}u", + f"{geometry.kv_seq}u", + f"{geometry.past_seq}u", + f"{geometry.head_size}u", + f"{geometry.v_head_size}u", + f"{geometry.repeats}u", + f"{0 if mask is None else mask.shape[-1]}u", + scalar_literal(_scale(context, geometry.head_size), TensorProto.DOUBLE), + scalar_literal(context.float_attribute("softcap"), TensorProto.DOUBLE), + str(causal), + f"{_causal_row_stride(context, mask, causal)}u", + str(mode), + ] + + +def _bias_function(context: NodeContext, mask_type: int) -> CFunction: + name = kernel_name(context, "bias", c_type(mask_type)) + read = _BOOLEAN_MASK if mask_type == TensorProto.BOOL else _ADDITIVE_MASK + return CFunction( + name, + _BIAS_TEMPLATE.substitute(name=name, mask=c_type(mask_type), read=read), + ) + + +# -------------------------------------------------------------------------------------- +# Reading an Attention node's geometry and types +# -------------------------------------------------------------------------------------- + + +def _geometry(context: NodeContext) -> _Geometry: + query, key, value = (context.require_input(index) for index in range(3)) + rank = len(query.shape) + if rank not in (3, 4) or len(key.shape) != rank or len(value.shape) != rank: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `Q`, `K` and `V` as three tensors " + f"of one rank — 4 for `(batch, head, sequence, size)`, 3 for " + f"`(batch, sequence, head * size)` — but they have shapes " + f"{list(query.shape)}, {list(key.shape)} and {list(value.shape)}." + ) + if rank == 4: + batch, q_heads, q_seq, head_size = query.shape + kv_heads, kv_seq, key_size = key.shape[1:] + value_heads, value_seq, v_head_size = value.shape[1:] + _head_count(context, "q_num_heads", q_heads) + _head_count(context, "kv_num_heads", kv_heads) + else: + q_heads = _head_count(context, "q_num_heads", None) + kv_heads = _head_count(context, "kv_num_heads", None) + batch, q_seq, q_hidden = query.shape + kv_seq, key_hidden = key.shape[1:] + value_seq, value_hidden = value.shape[1:] + head_size = _head_size(context, query, q_hidden, q_heads) + key_size = _head_size(context, key, key_hidden, kv_heads) + v_head_size = _head_size(context, value, value_hidden, kv_heads) + value_heads = kv_heads + _verify_operands( + context, + batch=batch, + heads=(q_heads, kv_heads, value_heads), + sizes=(head_size, key_size), + sequences=(kv_seq, value_seq), + ) + past_seq = _past_length(context, batch, kv_heads, head_size, v_head_size) + _verify_nonpad(context, batch) + return _Geometry( + rank=rank, + batch=batch, + q_heads=q_heads, + kv_heads=kv_heads, + q_seq=q_seq, + kv_seq=kv_seq, + past_seq=past_seq, + head_size=head_size, + v_head_size=v_head_size, + ) + + +def _head_count(context: NodeContext, name: str, inferred: int | None) -> int: + """The head count for one side of the node, from its shapes or from its attributes. + + ONNX defines the two attributes for the 3-D layout, where nothing else says how the + hidden axis splits. The 4-D layout carries the count in the shape and the reference + evaluator reads it there; a node stating a different one describes two tensors, and there + is no telling which of them it meant. + """ + declared = context.attribute(name, None) + if inferred is None: + if declared is None: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `Q`, `K` and `V` as rank 3, " + f"where only `{name}` says how the hidden axis splits into heads, and this " + "node leaves it out." + ) + return int(declared) + if declared is not None and int(declared) != inferred: + raise CompileError( + f"Node `{context.label}`: `Attention` states `{name}` {int(declared)}, but its " + f"rank-4 operands carry {inferred} head(s) on the axis ONNX reads that count " + "off." + ) + return inferred + + +def _head_size( + context: NodeContext, operand: TensorRef, hidden: int, heads: int +) -> int: + if heads <= 0 or hidden % heads: + raise CompileError( + f"Node `{context.label}`: `Attention` splits `{operand.name}`'s hidden axis of " + f"{hidden} into {heads} head(s), which does not divide it." + ) + return hidden // heads + + +def _verify_operands( + context: NodeContext, + *, + batch: int, + heads: tuple[int, int, int], + sizes: tuple[int, int], + sequences: tuple[int, int], +) -> None: + """Refuse to emit a kernel whose addressing disagrees with the operands it is handed.""" + query, key, value = (context.require_input(index) for index in range(3)) + if key.shape[0] != batch or value.shape[0] != batch: + raise CompileError( + f"Node `{context.label}`: `Attention` attends one batch, but `{query.name}`, " + f"`{key.name}` and `{value.name}` carry {batch}, {key.shape[0]} and " + f"{value.shape[0]} item(s)." + ) + if sizes[0] != sizes[1]: + raise CompileError( + f"Node `{context.label}`: `Attention` contracts `{query.name}` against " + f"`{key.name}` over the head size, but theirs are {sizes[0]} and {sizes[1]}." + ) + q_heads, kv_heads, value_heads = heads + if kv_heads != value_heads or sequences[0] != sequences[1]: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `{key.name}` and `{value.name}` at " + f"the same key positions, but they carry {kv_heads} head(s) of " + f"{sequences[0]} against {value_heads} of {sequences[1]}." + ) + if q_heads % kv_heads if kv_heads else q_heads: + raise CompileError( + f"Node `{context.label}`: `Attention` shares each of the {kv_heads} key/value " + f"head(s) between the query heads that follow it, which needs {q_heads} query " + f"head(s) to be a multiple of {kv_heads}." + ) + + +def _past_length( + context: NodeContext, batch: int, kv_heads: int, head_size: int, v_head_size: int +) -> int: + """How many cached key positions precede `K`, with both caches checked against them.""" + past_key = context.optional_input(_PAST_KEY_INPUT) + past_value = context.optional_input(_PAST_VALUE_INPUT) + given = [cache for cache in (past_key, past_value) if cache is not None] + if not given: + return 0 + if past_key is None or past_value is None: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `past_key` and `past_value` as one " + "cache, so ONNX defines them as used together; this node reads only " + f"`{given[0].name}`." + ) + if len(past_key.shape) != 4: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `{past_key.name}` as " + "`(batch, kv_num_heads, past_sequence_length, head_size)`, but it has shape " + f"{list(past_key.shape)}." + ) + past_seq = past_key.shape[2] + # Pairs rather than a mapping: a node may name one tensor for both caches, and two + # `TensorRef`s of one tensor are equal, so a mapping would drop one of the two checks. + expected = ( + (past_key, (batch, kv_heads, past_seq, head_size)), + (past_value, (batch, kv_heads, past_seq, v_head_size)), + ) + for operand, shape in expected: + if operand.shape != shape: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `{operand.name}` as a cache of " + f"shape {list(shape)}, but it has shape {list(operand.shape)}." + ) + return past_seq + + +def _verify_nonpad(context: NodeContext, batch: int) -> None: + nonpad = context.optional_input(_NONPAD_INPUT) + if nonpad is not None and nonpad.shape != (batch,): + raise CompileError( + f"Node `{context.label}`: `Attention` reads `{nonpad.name}` as one key length " + f"per batch item — a tensor of shape [{batch}] — but it has shape " + f"{list(nonpad.shape)}." + ) + + +def _element_type(context: NodeContext) -> int: + """The element type `Q`, `K`, `V` and the caches share. + + ONNX constrains the query/key side and the value side separately, so a model may state + two different floating-point types for them. The reference evaluator then computes in + whatever numpy promotes the pair to, at every step of the chain; rather than reproduce + that promotion lattice on a combination nothing in the corpus or the sweep exercises, + the compiler serves the case where the operands agree. + """ + query = context.require_input(0) + element = query.elem_type + for index in (1, 2, _PAST_KEY_INPUT, _PAST_VALUE_INPUT): + operand = context.optional_input(index) + if operand is not None and operand.elem_type != element: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `{query.name}` as " + f"`{element_type_name(element)}` and `{operand.name}` as " + f"`{element_type_name(operand.elem_type)}`; this compiler attends one " + "element type, and the two have to agree." + ) + return element + + +def _mask_type(context: NodeContext, mask: TensorRef | None, element: int) -> int: + """The element type the mask is read at, which is the flavour of bias the kernel adds. + + ONNX's type constraint admits every numeric type, while its own text defines the operand + as "a boolean mask ... or a float mask of the same type as query, key, value" — and the + reference evaluator cannot evaluate an integer mask under `is_causal` at all, since an + integer array holds no -inf. Only the two forms the text names are compiled. A node + without a mask takes the additive kernel with a null pointer. + """ + if mask is None: + return element + if mask.elem_type in (TensorProto.BOOL, element): + return mask.elem_type + raise CompileError( + f"Node `{context.label}`: `Attention` reads `{mask.name}` as " + f"`{element_type_name(mask.elem_type)}`, but ONNX defines `attn_mask` as a boolean " + f"mask or a float mask of the operands' own type, which is " + f"`{element_type_name(element)}` here." + ) + + +def _softmax_type(context: NodeContext) -> int: + """The element type the softmax runs in. + + `double` by default, and not the tensors' own type: the reference evaluator's scaling of + `Q` and `K` is by a numpy float64 scalar, so everything the softmax reads has already + been promoted to float64 whatever the model holds. `softmax_precision` is what narrows + it back. + """ + declared = context.attribute("softmax_precision", None) + if declared is None: + return TensorProto.DOUBLE + precision = int(declared) + if precision not in FLOAT_TYPES: + raise CompileError( + f"Node `{context.label}`: `Attention` states a `softmax_precision` of " + f"`{element_type_name(precision)}`, but ONNX defines the attribute as the " + "floating-point precision the softmax runs in, and this compiler computes in " + f"{' and '.join(sorted(element_type_name(t) for t in FLOAT_TYPES))}." + ) + return precision + + +def _product_type(soft: int, element: int) -> int: + """The type the softmax weights are read back against the values in. + + numpy promotes the two operands of that product, so it is the wider of them. + """ + if soft == TensorProto.DOUBLE or element == TensorProto.DOUBLE: + return TensorProto.DOUBLE + return TensorProto.FLOAT + + +def _output_mode(context: NodeContext) -> int: + mode = context.int_attribute("qk_matmul_output_mode") + if mode not in _OUTPUT_MODES: + raise CompileError( + f"Node `{context.label}`: `Attention` states a `qk_matmul_output_mode` of " + f"{mode}, but ONNX defines only " + f"{', '.join(str(value) for value in _OUTPUT_MODES)}." + ) + return mode + + +def _scale(context: NodeContext, head_size: int) -> float: + """What `Q` and `K` are each multiplied by before their product. + + ONNX's `scale` scales the product; the reference evaluator takes its square root and + applies that to both operands, so that is what the kernel is handed. Both edges follow + numpy rather than Python: the square root of a negative is a NaN, and the default of a + head of no elements is an infinity that no dot product ever reads. + """ + declared = context.attribute("scale", None) + if declared is None: + value = math.inf if head_size == 0 else 1.0 / math.sqrt(head_size) + else: + value = float(declared) + return math.sqrt(value) if value >= 0.0 else math.nan + + +def _mask_strides( + context: NodeContext, mask: TensorRef | None, geometry: _Geometry +) -> tuple[int, ...]: + """Strides addressing `attn_mask` while walking batch item, query head and query row. + + Its last axis is the key positions, which ONNX pads out to the total key length rather + than broadcasting, so the kernel addresses that one itself; everything before it + broadcasts numpy-style onto the three axes the loops walk. + """ + if mask is None: + return (0, 0, 0) + walked = (geometry.batch, geometry.q_heads, geometry.q_seq) + if not 1 <= len(mask.shape) <= len(walked) + 1: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `{mask.name}` of shape " + f"{list(mask.shape)}, but ONNX defines `attn_mask` as broadcastable to " + f"`(batch, q_num_heads, q_sequence_length, total_sequence_length)`." + ) + if mask.shape[-1] > geometry.total_seq: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `{mask.name}` along " + f"{mask.shape[-1]} key position(s), but the node attends " + f"{geometry.total_seq}; the reference evaluator pads a shorter mask out to " + "the key axis and refuses a longer one outright, so nothing says what the " + "columns past the end would mean." + ) + leading = replace(mask, shape=mask.shape[:-1]) + stretched = broadcast_strides(leading, walked, node_label=context.label) + # `broadcast_strides` reads the strides off the shape it is given, which is the mask + # without its key axis; the real ones carry that axis' extent as a factor. + return tuple(stride * mask.shape[-1] for stride in stretched) + + +def _causal_row_stride( + context: NodeContext, mask: TensorRef | None, causal: int +) -> int: + """Whether the causal triangle advances with the query row, or one row serves them all. + + The reference evaluator adds the triangle *into the mask* — `np.triu` over the extents it + reads off `attn_mask.shape[-2:]` — and only then broadcasts the sum onto the scores. So a + mask carrying one row on the query axis, which is the `(batch, 1, 1, total)` padding mask + a decoder passes alongside `is_causal`, gets the triangle's first row and masks every + query position alike rather than by its own place in the sequence. The schema's prose + describes the other reading; the oracle both suites compare against computes this one. + """ + if not causal or mask is None: + return 1 + if len(mask.shape) < 2: + raise CompileError( + f"Node `{context.label}`: `Attention` reads `{mask.name}` of shape " + f"{list(mask.shape)} under `is_causal`, where the reference evaluator takes the " + "triangle's extent from the mask's own query axis — a mask of rank 1 has none, " + "and the reference refuses such a node outright, so nothing vouches for what " + "this one would compute." + ) + return 0 if mask.shape[-2] == 1 else 1 + + +# -------------------------------------------------------------------------------------- +# RotaryEmbedding +# -------------------------------------------------------------------------------------- + +# One rotation per (batch item, position, head, pair). `interleaved` picks the two lanes a +# pair is made of — adjacent ones rather than the two halves — and writes the rotated pair +# back into those same lanes, so it changes addressing and nothing else. The lanes past +# `rotary_embedding_dim` are copied through untouched, which is what makes a partial rotation +# partial. +_ROTARY_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const $element* cosine, + const $element* sine, + const int64_t* positions, + const size_t* strides, + const size_t* cache_strides, + const size_t* position_strides, + size_t batch, + size_t seq, + size_t heads, + size_t head_size, + size_t rotary_half, + size_t cache_rows, + size_t cache_row, + int interleaved) +{ + size_t item, step, head, pair, lane; + for (item = 0; item < batch; ++item) { + for (step = 0; step < seq; ++step) { + size_t cache_base = + item * cache_strides[0] + step * cache_strides[1]; + if (positions != NULL) { + ptrdiff_t position = (ptrdiff_t)positions[ + item * position_strides[0] + step * position_strides[1]]; + /* numpy's own indexing, which is what the reference gathers with: a + negative index counts from the end, anything else is out of range. */ + if (position < 0) { + position += (ptrdiff_t)cache_rows; + } + if (position < 0 || position >= (ptrdiff_t)cache_rows) { + return 1; + } + cache_base = (size_t)position * cache_row; + } + for (head = 0; head < heads; ++head) { + const size_t base = + item * strides[0] + step * strides[1] + head * strides[2]; + const size_t angles = cache_base + head * cache_strides[2]; + for (pair = 0; pair < rotary_half; ++pair) { + const size_t low = interleaved ? 2 * pair : pair; + const size_t high = + interleaved ? 2 * pair + 1 : pair + rotary_half; + const $element c = cosine[angles + pair]; + const $element s = sine[angles + pair]; + const $element x1 = in[base + low]; + const $element x2 = in[base + high]; + out[base + low] = c * x1 - s * x2; + out[base + high] = s * x1 + c * x2; + } + for (lane = 2 * rotary_half; lane < head_size; ++lane) { + out[base + lane] = in[base + lane]; + } + } + } + } + return 0; +}""") + + +def _rotary_embedding(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + result = context.require_output(0) + verify_shape(context, result, source.shape) + batch, seq, heads, head_size = _rotary_layout(context, source) + rotary = context.int_attribute("rotary_embedding_dim") or head_size + if rotary % 2 or not 0 <= rotary <= head_size: + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` rotates the first {rotary} lane(s) " + f"of a head of {head_size}, which ONNX splits into pairs — so it has to be even " + "and no wider than the head." + ) + positions = context.optional_input(3) + cache_strides = _cache_strides(context, positions, (batch, seq, heads), rotary // 2) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = _rotary_element_type(context) + name = kernel_name(context, c_type(element)) + definition = _ROTARY_TEMPLATE.substitute(name=name, element=c_type(element)) + cosine = context.require_input(1) + arguments = [ + result.expr, + source.expr, + cosine.expr, + context.require_input(2).expr, + "NULL" if positions is None else positions.expr, + extents(_rotary_strides(len(source.shape), seq, heads, head_size)), + extents(cache_strides), + extents( + (0, 0) + if positions is None + else broadcast_strides(positions, (batch, seq), node_label=context.label) + ), + f"{batch}u", + f"{seq}u", + f"{heads}u", + f"{head_size}u", + f"{rotary // 2}u", + f"{cosine.shape[0]}u", + f"{cosine.shape[-1]}u", + str(int(context.int_attribute("interleaved") != 0)), + ] + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(checked_call(context, name, arguments),), + ) + + +def _rotary_layout( + context: NodeContext, source: TensorRef +) -> tuple[int, int, int, int]: + """`(batch, sequence, heads, head size)`, whichever of the two layouts the node reads.""" + if len(source.shape) == 4: + batch, heads, seq, head_size = source.shape + declared = context.attribute("num_heads", None) + if declared is not None and int(declared) != heads: + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` states `num_heads` " + f"{int(declared)}, but its rank-4 `{source.name}` carries {heads} head(s) " + "on the axis ONNX reads that count off." + ) + return batch, seq, heads, head_size + if len(source.shape) != 3: + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` reads `{source.name}` as " + "`(batch, head, sequence, size)` or `(batch, sequence, head * size)`, but it " + f"has shape {list(source.shape)}." + ) + heads = context.attribute("num_heads", None) + if heads is None or int(heads) <= 0: + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` reads `{source.name}` as rank 3, " + "where only `num_heads` says how the hidden axis splits into heads, and this " + f"node states {'none' if heads is None else int(heads)}." + ) + batch, seq, hidden = source.shape + if hidden % int(heads): + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` splits `{source.name}`'s hidden " + f"axis of {hidden} into {int(heads)} head(s), which does not divide it." + ) + return batch, seq, int(heads), hidden // int(heads) + + +def _rotary_strides( + rank: int, seq: int, heads: int, head_size: int +) -> tuple[int, int, int]: + """`(batch, sequence, head)` strides; the result carries the operand's own layout back.""" + if rank == 4: + return (heads * seq * head_size, head_size, seq * head_size) + return (seq * heads * head_size, heads * head_size, head_size) + + +def _cache_strides( + context: NodeContext, + positions: TensorRef | None, + walked: tuple[int, int, int], + half: int, +) -> tuple[int, ...]: + """Strides addressing the sine and cosine caches over batch item, position and head. + + With `position_ids` the caches are gathered by row, which the kernel addresses itself, so + it walks none of these axes. Without it they carry the batch and the position and are + stretched over the heads, exactly as the reference's `expand_dims` at the head axis does. + """ + cosine = context.require_input(1) + sine = context.require_input(2) + if cosine.shape != sine.shape: + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` reads `{cosine.name}` and " + f"`{sine.name}` as one pair of caches, but they have shapes " + f"{list(cosine.shape)} and {list(sine.shape)}." + ) + if not cosine.shape or cosine.shape[-1] != half: + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` rotates {half} pair(s) per head, so " + f"ONNX ends its caches on that axis, but `{cosine.name}` has shape " + f"{list(cosine.shape)}." + ) + expected = 2 if positions is not None else 3 + if len(cosine.shape) != expected: + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` reads `{cosine.name}` as rank " + f"{expected} — `(max_position_id_plus_1, rotary_embedding_dim / 2)` when " + "`position_ids` gathers it, `(batch, sequence, rotary_embedding_dim / 2)` when " + f"it stands as it is — but it has shape {list(cosine.shape)}." + ) + if positions is not None: + # Exactly rank 2, not "at most": the reference gathers the caches by these indices + # and then inserts the head axis with `expand_dims(..., 2)`, which for a lower rank + # lands past the angles instead of before them and rotates by a transposed cache. + if len(positions.shape) != 2: + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` reads `{positions.name}` as " + f"`(batch, sequence)`, but it has shape {list(positions.shape)}." + ) + return (0, 0, 0) + stretched = replace(cosine, shape=(*cosine.shape[:-1], 1, half)) + return broadcast_strides(stretched, (*walked, half), node_label=context.label)[:3] + + +def _rotary_element_type(context: NodeContext) -> int: + element = context.require_input(0).elem_type + for index in (1, 2): + operand = context.require_input(index) + if operand.elem_type != element: + raise CompileError( + f"Node `{context.label}`: `RotaryEmbedding` rotates " + f"`{element_type_name(element)}` values by `{operand.name}`, which is " + f"`{element_type_name(operand.elem_type)}`; ONNX defines the operand and " + "its caches as one type." + ) + return element + + +register_kernel("", "Attention", _ATTENTION_VERSIONS, _attention) +register_kernel("", "RotaryEmbedding", _ROTARY_VERSIONS, _rotary_embedding) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/axes.py b/src/python/fnnx/extras/compilers/c/onnx/ops/axes.py new file mode 100644 index 0000000..1f23b7d --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/axes.py @@ -0,0 +1,222 @@ +"""Grouping a tensor by axes: the addressing every axis-wise kernel is emitted from. + +A reduction, a softmax and a cumulative sum all walk the same two nested loops: one over the +groups the axes an op does *not* name leave behind, and one over the elements of a group. +This module turns a static shape and a set of axes into the compile-time literals those loops +take — an extent and a stride per axis — and emits the one helper that turns a loop counter +back into an offset into the tensor's row-major buffer. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass +from string import Template + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.emit import INVALID_ARGUMENT_STATUS +from fnnx.extras.compilers.c.onnx.kernels import CFunction, NodeContext, TensorRef +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents + +# The parameters every axis-wise kernel takes after its buffers: the two loop bounds, and the +# extents and strides that address each loop's coordinates. +GROUP_PARAMETERS = """\ + size_t group_count, + size_t group_size, + int kept_rank, + const size_t* kept_shape, + const size_t* kept_strides, + int reduced_rank, + const size_t* reduced_shape, + const size_t* reduced_strides""" + +# The offset of a group's first element, and of an element within a group; both decompose a +# linear counter into coordinates, which is what `$offset` is for. +GROUP_BASE = "$offset(group, kept_rank, kept_shape, kept_strides)" +GROUP_ELEMENT = "$offset($index, reduced_rank, reduced_shape, reduced_strides)" + +_OFFSET_TEMPLATE = Template("""\ +static size_t $name( + size_t index, + int rank, + const size_t* shape, + const size_t* strides) +{ + size_t offset = 0; + int axis; + for (axis = rank - 1; axis >= 0; --axis) { + offset += (index % shape[axis]) * strides[axis]; + index /= shape[axis]; + } + return offset; +}""") + + +@dataclass(frozen=True) +class Grouping: + """A tensor's axes split into the ones an op runs across and the ones it runs along. + + Each entry is the axis's extent and its stride into the tensor's row-major buffer, so a + kernel addresses the original tensor without any of it being copied or transposed. + `kept_axes` and `reduced_axes` are the positions those entries came from, which is what + an operand addressed from the same coordinates — a normalization's scale — is split by. + """ + + kept: tuple[tuple[int, int], ...] + reduced: tuple[tuple[int, int], ...] + kept_axes: tuple[int, ...] + reduced_axes: tuple[int, ...] + + @property + def group_count(self) -> int: + return math.prod(extent for extent, _ in self.kept) + + @property + def group_size(self) -> int: + return math.prod(extent for extent, _ in self.reduced) + + @property + def arguments(self) -> list[str]: + """Call-site literals for `GROUP_PARAMETERS`, in order.""" + return [ + f"{self.group_count}u", + f"{self.group_size}u", + *_axis_arguments(self.kept), + *_axis_arguments(self.reduced), + ] + + +def group_axes(shape: Sequence[int], axes: Sequence[int]) -> Grouping: + """Split `shape` into the axes `axes` names and the ones it leaves.""" + strides = row_major_strides(shape) + named = frozenset(axes) + kept_axes = tuple(axis for axis in range(len(shape)) if axis not in named) + return Grouping( + kept=tuple((shape[axis], strides[axis]) for axis in kept_axes), + reduced=tuple((shape[axis], strides[axis]) for axis in axes), + kept_axes=kept_axes, + reduced_axes=tuple(axes), + ) + + +def row_major_strides(shape: Sequence[int]) -> tuple[int, ...]: + strides = [] + stride = 1 + for extent in reversed(shape): + strides.append(stride) + stride *= extent + return tuple(reversed(strides)) + + +def normalize_axis(context: NodeContext, axis: int, rank: int) -> int: + """An ONNX axis, which may count from the end, as an index into a rank-`rank` tensor.""" + resolved = axis + rank if axis < 0 else axis + if not 0 <= resolved < rank: + raise CompileError( + f"Node `{context.label}`: axis {axis} is out of range for the rank-{rank} " + f"tensor `{context.require_input(0).name}`." + ) + return resolved + + +def normalize_axes( + context: NodeContext, axes: Sequence[int], rank: int +) -> tuple[int, ...]: + """The axes an op names, resolved and sorted so a kernel walks them in memory order.""" + resolved = tuple(normalize_axis(context, int(axis), rank) for axis in axes) + if len(set(resolved)) != len(resolved): + raise CompileError( + f"Node `{context.label}`: axes {[int(axis) for axis in axes]} name the same " + "dimension more than once." + ) + return tuple(sorted(resolved)) + + +def verify_group_count( + context: NodeContext, grouping: Grouping, result: TensorRef +) -> None: + """Refuse to emit a kernel that would write past the result buffer. + + The groups are counted from the operand's shape and the axes the node names, while the + buffer is sized from the shape ONNX inferred for the result; a disagreement between the + two is a compiler bug, and this is where it stops rather than where it corrupts memory. + """ + if grouping.group_count != result.elem_count: + raise CompileError( + f"Node `{context.label}`: the axes it names leave {grouping.group_count} " + f"group(s), but its output `{result.name}` holds {result.elem_count} element(s)." + ) + + +def verify_shape( + context: NodeContext, result: TensorRef, expected: Sequence[int] +) -> None: + """Refuse to emit a kernel whose addressing disagrees with the buffer it writes. + + The extents are derived from the operands and the node's own attributes, while the buffer + is sized from the shape ONNX inferred for the result; a disagreement between the two is a + compiler bug, and this is where it stops rather than where it corrupts memory. + """ + if result.shape != tuple(expected): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` addresses a result of shape " + f"{list(expected)}, but its output `{result.name}` holds " + f"{list(result.shape)}." + ) + + +def verify_same_shape( + context: NodeContext, source: TensorRef, result: TensorRef +) -> None: + """Refuse to emit an axis-wise kernel whose two buffers are not laid out alike. + + Softmax and the cumulative folds write a result of the operand's own shape, so one + grouping addresses both; a disagreement is a compiler bug, and this is where it stops. + """ + if source.shape != result.shape: + raise CompileError( + f"Node `{context.label}`: `{source.name}` has shape {list(source.shape)} but " + f"its result `{result.name}` has shape {list(result.shape)}; this op leaves " + "the shape alone." + ) + + +def offset_helper(prefix: str) -> CFunction: + """The shared index-to-offset function every axis-wise kernel calls.""" + name = f"{prefix}_axis_offset" + return CFunction(name, _OFFSET_TEMPLATE.substitute(name=name)) + + +def kernel_name(context: NodeContext, *parts: str) -> str: + """A kernel name encoding the op and everything else its code depends on.""" + return "_".join((context.prefix, context.node.op_type.lower(), *parts)) + + +def call_kernel(name: str, arguments: Sequence[str]) -> str: + return f"{name}(\n " + ",\n ".join(arguments) + ");" + + +def checked_call(context: NodeContext, name: str, arguments: Sequence[str]) -> str: + """A call site for a kernel that validates an operand's values at run time. + + Such a kernel returns nonzero for a value ONNX leaves undefined — an index outside the + axis it addresses — and the entrypoint passes that on as the argument error the status + enum exists for, rather than reading past a buffer. + """ + call = call_kernel(name, arguments).rstrip(";") + return "\n".join( + [ + f"if ({call} != 0) {{", + f" return {context.prefix.upper()}_{INVALID_ARGUMENT_STATUS};", + "}", + ] + ) + + +def _axis_arguments(axes: Sequence[tuple[int, int]]) -> list[str]: + return [ + str(len(axes)), + extents([extent for extent, _ in axes]), + extents([stride for _, stride in axes]), + ] diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/broadcast.py b/src/python/fnnx/extras/compilers/c/onnx/ops/broadcast.py new file mode 100644 index 0000000..3a829c1 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/broadcast.py @@ -0,0 +1,251 @@ +"""The loop every elementwise kernel is emitted from. + +An elementwise op is a C expression over one scalar per operand: this module turns that +expression into a shared `static` kernel and the call site invoking it. Operands broadcast +onto the result's shape numpy-style, addressed through strides that are zero on every axis +they are stretched along; when every operand already has the result's shape — the common +case — the same expression is emitted into a flat loop that pays no index arithmetic. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from string import Template + +from onnx import TensorProto + +from fnnx.extras.compilers.c.onnx.dtypes import FLOAT_TYPES, c_type +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + broadcast_strides, +) + +_POINTWISE_TEMPLATE = Template("""\ +static void $name( +$parameters, + size_t count) +{ + size_t index; + for (index = 0; index < count; ++index) { +$reads + out[index] = $expression; + } +}""") + +_BROADCAST_TEMPLATE = Template("""\ +static void $name( +$parameters, + size_t count, + int rank, + const size_t* shape, +$stride_parameters) +{ + size_t index; + for (index = 0; index < count; ++index) { + size_t remainder = index; +$offsets + int axis; + for (axis = rank - 1; axis >= 0; --axis) { + size_t coordinate = remainder % shape[axis]; + remainder /= shape[axis]; +$accumulate + } +$reads + out[index] = $expression; + } +}""") + +# numpy's `minimum`/`maximum`, which ONNX's Min, Max, Clip, HardSigmoid and Celu are all +# defined through: NaN wins over anything, and everything else that is not strictly better +# than the second operand yields the second operand — which is what makes `minimum(-0, 0)` +# come out `+0` and `minimum(0, -0)` come out `-0`. C's `fmin`/`fmax` do the opposite with +# NaN, so they are never used. +_COMBINE_TEMPLATE = Template("""\ +static $element $name($element left, $element right) +{ + return ($test) ? left : right; +}""") + + +@dataclass(frozen=True) +class Scalar: + """An attribute value a kernel reads as a parameter rather than as an inlined literal. + + Kernels are shared across nodes, so the values that differ between two nodes running the + same op — Elu's alpha, Shrink's bias — have to arrive as arguments, or every distinct + attribute value would emit a kernel of its own. + """ + + name: str + elem_type: int + value: float | int + + +def elementwise( + context: NodeContext, + *, + expression: str, + operands: Sequence[TensorRef], + result: TensorRef, + scalars: Sequence[Scalar] = (), + helpers: Sequence[CFunction] = (), + variant: str = "", +) -> NodeEmission: + """A kernel writing `expression` into every element of `result`. + + `expression` is C over the locals `x0`, `x1`, ... — one per operand, of that operand's + element type — and over the `scalars`' names. `variant` distinguishes kernels whose code + differs for a reason the operand types do not capture, such as an attribute that selects + between two formulas; `helpers` are functions the expression calls, emitted first. + """ + aligned = all(operand.shape == result.shape for operand in operands) + name = _kernel_name(context, operands, result, variant, aligned) + parameters = [f" {c_type(result.elem_type)}* out"] + parameters += [ + f" const {c_type(operand.elem_type)}* in{index}" + for index, operand in enumerate(operands) + ] + parameters += [ + f" {c_type(scalar.elem_type)} {scalar.name}" for scalar in scalars + ] + arguments = [result.expr, *(operand.expr for operand in operands)] + arguments += [scalar_literal(scalar.value, scalar.elem_type) for scalar in scalars] + + if aligned: + definition = _POINTWISE_TEMPLATE.substitute( + name=name, + parameters=",\n".join(parameters), + reads=_reads(operands, lambda index: "index"), + expression=expression, + ) + call = _call(name, [*arguments, f"{result.elem_count}u"]) + else: + strides = [ + broadcast_strides(operand, result.shape, node_label=context.label) + for operand in operands + ] + definition = _BROADCAST_TEMPLATE.substitute( + name=name, + parameters=",\n".join(parameters), + stride_parameters=",\n".join( + f" const size_t* strides{index}" for index in range(len(operands)) + ), + offsets="\n".join( + f" size_t offset{index} = 0;" for index in range(len(operands)) + ), + accumulate="\n".join( + f" offset{index} += coordinate * strides{index}[axis];" + for index in range(len(operands)) + ), + reads=_reads(operands, lambda index: f"offset{index}"), + expression=expression, + ) + call = _call( + name, + [ + *arguments, + f"{result.elem_count}u", + str(len(result.shape)), + extents(result.shape), + *(extents(stride) for stride in strides), + ], + ) + return NodeEmission( + functions=(*helpers, CFunction(name, definition)), statements=(call,) + ) + + +def pointwise( + context: NodeContext, + template: str, + *, + scalars: Sequence[Scalar] = (), + helpers: Sequence[CFunction] = (), + variant: str = "", +) -> NodeEmission: + """A kernel for a one-operand op, whose formula is `template` over the local `x0`.""" + result = context.require_output(0) + return elementwise( + context, + expression=expand(template, result.elem_type), + operands=(context.require_input(0),), + result=result, + scalars=scalars, + helpers=helpers, + variant=variant, + ) + + +def expand(template: str, elem_type: int) -> str: + """Fill in what a kernel expression takes from its element type. + + `$f` is the libm suffix, `$one` and `$zero` are literals of that type, and `$element` + is its C type. + """ + return Template(template).substitute( + f=math_suffix(elem_type), + one=scalar_literal(1, elem_type), + zero=scalar_literal(0, elem_type), + element=c_type(elem_type), + ) + + +def combiner(context: NodeContext, elem_type: int, *, largest: bool) -> CFunction: + """numpy's `minimum` or `maximum` at `elem_type`, as a function to call per element.""" + comparison = ">" if largest else "<" + test = f"left {comparison} right" + if elem_type in FLOAT_TYPES: + test = f"{test} || isnan(left)" + name = f"{context.prefix}_{'maximum' if largest else 'minimum'}_{c_type(elem_type)}" + return CFunction( + name, + _COMBINE_TEMPLATE.substitute(name=name, element=c_type(elem_type), test=test), + ) + + +def extents(values: Sequence[int]) -> str: + """Shapes and strides as a compound literal; rank 0 gets an unread placeholder.""" + literals = ", ".join(f"{value}u" for value in values) or "0u" + return f"(const size_t[]){{{literals}}}" + + +def math_suffix(elem_type: int) -> str: + """The libm suffix selecting the overload for this element type: `sinf` against `sin`.""" + return "" if elem_type == TensorProto.DOUBLE else "f" + + +def _reads(operands: Sequence[TensorRef], offset: Callable[[int], str]) -> str: + return "\n".join( + f" const {c_type(operand.elem_type)} x{index} = " + f"in{index}[{offset(index)}];" + for index, operand in enumerate(operands) + ) + + +def _call(name: str, arguments: Sequence[str]) -> str: + return f"{name}(\n " + ",\n ".join(arguments) + ");" + + +def _kernel_name( + context: NodeContext, + operands: Sequence[TensorRef], + result: TensorRef, + variant: str, + aligned: bool, +) -> str: + """A name encoding everything the emitted code depends on, and nothing else. + + Two nodes running the same op reach the same kernel — and so share one definition — only + when their operand types, arity, broadcasting and formula all agree; anything else would + be two kernels colliding on one name. + """ + types = [c_type(operand.elem_type) for operand in operands] + types.append(c_type(result.elem_type)) + tag = f"{len(operands)}_{types[0]}" if len(set(types)) == 1 else "_".join(types) + form = "" if aligned else "_bcast" + return f"{context.prefix}_{context.node.op_type.lower()}{variant}{form}_{tag}" diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/casts.py b/src/python/fnnx/extras/compilers/c/onnx/ops/casts.py new file mode 100644 index 0000000..19138da --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/casts.py @@ -0,0 +1,72 @@ +"""Casts: converting a tensor's values to another element type, or reading its bits as one.""" + +from __future__ import annotations + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type, element_type_name +from fnnx.extras.compilers.c.onnx.kernels import ( + NodeContext, + NodeEmission, + copy_tensor, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import elementwise + +# Cast-1 took `to` as a type name rather than the type id it has been since 6; every +# revision after that widened the type constraints alone (9 strings, 13 bfloat16, 19 float8 +# and the `saturate` attribute that goes with it, 21 int4, 23 float4, 24 and 25 more of the +# same). None of those types are compilable, so one generator serves every revision — and +# the target type is read off the graph rather than off `to`, which shape inference has +# already applied. +_CAST_VERSIONS = (6, 9, 13, 19, 21, 23, 24, 25) +_BITCAST_VERSIONS = (26,) + + +def _cast(context: NodeContext) -> NodeEmission: + """Cast between the compilable element types, which is C's own conversion. + + That is what ONNX specifies for all but the boolean target: truncation toward zero out + of the floating-point families — undefined out of the target's range, as ONNX leaves it + too — and, between the integer families, the modular reinterpretation of the low bits + that every two's-complement target performs. + """ + source = context.require_input(0) + result = context.require_output(0) + if source.elem_type == result.elem_type: + return copy_tensor(source, result) + element = c_type(result.elem_type) + # A boolean target is the one rule of ONNX's own: every nonzero value is true, NaN + # included. It also needs a variant of its own, since `bool` and `uint8` are one C type: + # without it a model casting in both directions would name two formulas alike. + to_bool = result.elem_type == TensorProto.BOOL + return elementwise( + context, + expression=f"({element})(x0 != 0)" if to_bool else f"({element})x0", + operands=(source,), + result=result, + variant="_to_bool" if to_bool else "", + ) + + +def _bitcast(context: NodeContext) -> NodeEmission: + """BitCast: the same bytes, read at another element type. + + ONNX defines the op only between types of equal width — its own type inference rejects + anything else before a kernel is reached, and a revision that relaxed that would not be + served by this generator — so moving the bytes is the whole operation. + """ + source = context.require_input(0) + result = context.require_output(0) + if result.elem_type == TensorProto.BOOL: + raise CompileError( + f"Node `{context.label}`: BitCast to `BOOL` is not supported by the C " + "compiler; a boolean tensor is emitted as bytes holding 0 or 1, which the " + f"bits of a `{element_type_name(source.elem_type)}` need not be." + ) + return copy_tensor(source, result) + + +register_kernel("", "Cast", _CAST_VERSIONS, _cast) +register_kernel("", "BitCast", _BITCAST_VERSIONS, _bitcast) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/conv.py b/src/python/fnnx/extras/compilers/c/onnx/ops/conv.py new file mode 100644 index 0000000..1a20604 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/conv.py @@ -0,0 +1,978 @@ +"""The convolutions: the sliding window, and where each of its taps lands in the operand. + +Every convolution is the same walk whatever its rank: for each output position, each tap of +the filter reads the operand at `position * stride + tap * dilation - pad` along each spatial +axis, skipping the taps that fall outside it. So there is one kernel per element type, taking +the geometry — the extents, strides, dilations and pads of each axis — as compile-time +literals, and `group` splits the channels into independent stacks addressed by the same walk. + +`ConvTranspose` walks the same geometry backwards. It is the gradient of a convolution, so +each of *its* output positions reads the operand at `(position + pad - tap * dilation) / +stride` — the same relation solved the other way, which drops the taps the stride does not +divide. `DeformConv` keeps the forward walk but shifts each tap by an offset it reads at run +time, so the position it samples falls between elements and is interpolated. `Col2Im` runs the +backward walk on its own, with no filter to weight by: it is the same geometry folding a stack +of blocks back into the image they were cut from. + +The geometry itself is resolved at compile time rather than in the kernel: `auto_pad` and +`output_shape` turn into concrete pads — through the shared reading of the window attributes +in `window.py`, plus the backward walk's own arithmetic here — and the result shape those pads +imply is checked against the one ONNX inferred, so a disagreement between the two stops the +compile instead of writing past a buffer. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass +from string import Template + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type, element_type_name +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + kernel_name, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents, math_suffix +from fnnx.extras.compilers.c.onnx.ops.window import ( + auto_pad_mode, + declared_pads, + offsets, + output_extents, + resolve_pads, + spatial_attribute, + spatial_extents, +) + +# The geometry every walk of a sliding window takes, in the order `Geometry.arguments` fills +# it; the quantized convolutions take the same block after their own operands. +WINDOW_PARAMETERS = """\ + size_t batch_count, + size_t groups, + size_t group_channels, + size_t group_filters, + size_t input_size, + size_t output_size, + size_t window_size, + int spatial_rank, + const size_t* input_shape, + const size_t* output_shape, + const size_t* window_shape, + const size_t* strides, + const size_t* dilations, + const ptrdiff_t* pads""" + +_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + const $element* weights, + const $element* bias, +$parameters) +{ + size_t batch, group, filter, position, tap, channel; + for (batch = 0; batch < batch_count; ++batch) { + for (group = 0; group < groups; ++group) { + const $element* plane = + in + (batch * groups + group) * group_channels * input_size; + for (filter = 0; filter < group_filters; ++filter) { + const size_t channel_index = group * group_filters + filter; + const $element* window = + weights + channel_index * group_channels * window_size; + $element* result = + out + (batch * groups * group_filters + channel_index) * output_size; + for (position = 0; position < output_size; ++position) { + $element sum = (bias != NULL) ? bias[channel_index] : $zero; + for (tap = 0; tap < window_size; ++tap) { + size_t remaining_position = position; + size_t remaining_tap = tap; + size_t offset = 0; + size_t stride = 1; + int inside = 1; + int axis; + for (axis = spatial_rank - 1; axis >= 0; --axis) { + const ptrdiff_t coordinate = + (ptrdiff_t)(remaining_position % output_shape[axis]) + * (ptrdiff_t)strides[axis] + + (ptrdiff_t)(remaining_tap % window_shape[axis]) + * (ptrdiff_t)dilations[axis] + - pads[axis]; + remaining_position /= output_shape[axis]; + remaining_tap /= window_shape[axis]; + if (coordinate < 0 + || coordinate >= (ptrdiff_t)input_shape[axis]) { + inside = 0; + } else { + offset += (size_t)coordinate * stride; + } + stride *= input_shape[axis]; + } + if (inside) { + for (channel = 0; channel < group_channels; ++channel) { + sum += plane[channel * input_size + offset] + * window[channel * window_size + tap]; + } + } + } + result[position] = sum; + } + } + } + } +}""") + +_TRANSPOSE_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + const $element* weights, + const $element* bias, +$parameters) +{ + size_t batch, group, filter, position, tap, channel; + for (batch = 0; batch < batch_count; ++batch) { + for (group = 0; group < groups; ++group) { + const $element* plane = + in + (batch * groups + group) * group_channels * input_size; + for (filter = 0; filter < group_filters; ++filter) { + const size_t channel_index = group * group_filters + filter; + /* The filter's taps for every channel of this group, which `W` holds one + input channel at a time: (C, M / group, ...). */ + const $element* stack = + weights + + (group * group_channels * group_filters + filter) * window_size; + $element* result = + out + (batch * groups * group_filters + channel_index) * output_size; + for (position = 0; position < output_size; ++position) { + $element sum = (bias != NULL) ? bias[channel_index] : $zero; + for (tap = 0; tap < window_size; ++tap) { + size_t remaining_position = position; + size_t remaining_tap = tap; + size_t offset = 0; + size_t stride = 1; + int inside = 1; + int axis; + for (axis = spatial_rank - 1; axis >= 0; --axis) { + const ptrdiff_t step = (ptrdiff_t)strides[axis]; + const ptrdiff_t reach = + (ptrdiff_t)(remaining_position % output_shape[axis]) + + pads[axis] + - (ptrdiff_t)(remaining_tap % window_shape[axis]) + * (ptrdiff_t)dilations[axis]; + const ptrdiff_t coordinate = reach / step; + remaining_position /= output_shape[axis]; + remaining_tap /= window_shape[axis]; + /* A reach the stride does not divide lands between two of the + operand's elements, where this tap contributes nothing. */ + if (reach % step != 0 || coordinate < 0 + || coordinate >= (ptrdiff_t)input_shape[axis]) { + inside = 0; + } else { + offset += (size_t)coordinate * stride; + } + stride *= input_shape[axis]; + } + if (inside) { + for (channel = 0; channel < group_channels; ++channel) { + sum += plane[channel * input_size + offset] + * stack[channel * group_filters * window_size + tap]; + } + } + } + result[position] = sum; + } + } + } + } +}""") + +_SAMPLE_TEMPLATE = Template("""\ +static $element $name( + const $element* plane, + size_t rows, + size_t columns, + $element row, + $element column) +{ + $element row_floor, column_floor, row_fraction, column_fraction, total; + ptrdiff_t top, left; + int down, right; + /* At or past either border every corner falls outside the plane, and a coordinate that + is not a number cannot be floored into an index at all: both sample nothing. */ + if (!(row > -$one) || !(row < ($element)rows) + || !(column > -$one) || !(column < ($element)columns)) { + return $zero; + } + row_floor = floor$f(row); + column_floor = floor$f(column); + top = (ptrdiff_t)row_floor; + left = (ptrdiff_t)column_floor; + row_fraction = row - row_floor; + column_fraction = column - column_floor; + total = $zero; + for (down = 0; down < 2; ++down) { + const ptrdiff_t sampled_row = top + down; + if (sampled_row < 0 || sampled_row >= (ptrdiff_t)rows) { + continue; + } + for (right = 0; right < 2; ++right) { + const ptrdiff_t sampled_column = left + right; + if (sampled_column < 0 || sampled_column >= (ptrdiff_t)columns) { + continue; + } + total += (down ? row_fraction : $one - row_fraction) + * (right ? column_fraction : $one - column_fraction) + * plane[(size_t)sampled_row * columns + (size_t)sampled_column]; + } + } + return total; +}""") + +_DEFORM_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + const $element* weights, + const $element* offsets, + const $element* bias, + const $element* mask, + size_t batch_count, + size_t groups, + size_t group_channels, + size_t group_filters, + size_t offset_channels, + size_t offset_groups, + size_t input_rows, + size_t input_columns, + size_t output_rows, + size_t output_columns, + size_t window_rows, + size_t window_columns, + size_t row_stride, + size_t column_stride, + size_t row_dilation, + size_t column_dilation, + ptrdiff_t row_pad, + ptrdiff_t column_pad) +{ + const size_t input_size = input_rows * input_columns; + const size_t output_size = output_rows * output_columns; + const size_t window_size = window_rows * window_columns; + size_t batch, group, filter, row, column, channel, tap_row, tap_column; + for (batch = 0; batch < batch_count; ++batch) { + for (group = 0; group < groups; ++group) { + for (filter = 0; filter < group_filters; ++filter) { + const size_t channel_index = group * group_filters + filter; + for (row = 0; row < output_rows; ++row) { + for (column = 0; column < output_columns; ++column) { + const size_t position = row * output_columns + column; + $element sum = (bias != NULL) ? bias[channel_index] : $zero; + for (channel = 0; channel < group_channels; ++channel) { + const size_t source = group * group_channels + channel; + const $element* plane = + in + + (batch * groups * group_channels + source) + * input_size; + /* One deformation per (offset group, tap): a plane of `mask` + weights, and two planes of `offsets` -- a coordinate per + spatial axis. */ + const size_t deformation = + (batch * offset_groups + source / offset_channels) + * window_size; + for (tap_row = 0; tap_row < window_rows; ++tap_row) { + for (tap_column = 0; + tap_column < window_columns; + ++tap_column) { + const size_t tap = + tap_row * window_columns + tap_column; + const size_t shift = + 2 * (deformation + tap) * output_size + position; + const $element sampled_row = + ($element)((ptrdiff_t)(row * row_stride + + tap_row * row_dilation) - row_pad) + + offsets[shift]; + const $element sampled_column = + ($element)((ptrdiff_t)(column * column_stride + + tap_column * column_dilation) - column_pad) + + offsets[shift + output_size]; + $element weight = + weights[(channel_index * group_channels + channel) + * window_size + tap]; + if (mask != NULL) { + weight *= + mask[(deformation + tap) * output_size + + position]; + } + sum += $sample( + plane, + input_rows, + input_columns, + sampled_row, + sampled_column) * weight; + } + } + } + out[(batch * groups * group_filters + channel_index) * output_size + + position] = sum; + } + } + } + } + } +}""") + +_COL2IM_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + size_t plane_count, + size_t column_count, + size_t image_size, + size_t window_size, + int spatial_rank, + const size_t* image_shape, + const size_t* block_shape, + const size_t* column_shape, + const size_t* strides, + const size_t* dilations, + const ptrdiff_t* pads) +{ + size_t plane, position, tap; + for (plane = 0; plane < plane_count; ++plane) { + const $element* columns = in + plane * window_size * column_count; + $element* result = out + plane * image_size; + for (position = 0; position < image_size; ++position) { + $element sum = $zero; + for (tap = 0; tap < window_size; ++tap) { + size_t remaining_position = position; + size_t remaining_tap = tap; + size_t offset = 0; + size_t stride = 1; + int inside = 1; + int axis; + for (axis = spatial_rank - 1; axis >= 0; --axis) { + const ptrdiff_t step = (ptrdiff_t)strides[axis]; + const ptrdiff_t reach = + (ptrdiff_t)(remaining_position % image_shape[axis]) + + pads[axis] + - (ptrdiff_t)(remaining_tap % block_shape[axis]) + * (ptrdiff_t)dilations[axis]; + const ptrdiff_t coordinate = reach / step; + remaining_position /= image_shape[axis]; + remaining_tap /= block_shape[axis]; + /* A reach the stride does not divide falls between two block positions, + where no block placed this tap here. */ + if (reach % step != 0 || coordinate < 0 + || coordinate >= (ptrdiff_t)column_shape[axis]) { + inside = 0; + } else { + offset += (size_t)coordinate * stride; + } + stride *= column_shape[axis]; + } + if (inside) { + sum += columns[tap * column_count + offset]; + } + } + result[position] = sum; + } + } +}""") + +# Conv arrived at opset 1, 11 revised `auto_pad` and 22 widened the element types. Only 22 is +# claimed: it is the revision the reference evaluator is version-faithful for and the one +# every Conv test in the backend corpus imports, so it is the only one anything can vouch +# for. A model importing an older one gets the unsupported-version error. ConvTranspose (1, +# 11, 22) and DeformConv (19, 22) are claimed at their newest revision for the same reason. +_VERSIONS = (22,) + +# Col2Im has had one revision only, the one it arrived at. +_COL2IM_VERSIONS = (18,) + +# DeformConv samples between elements, which the compiler only emits for two spatial axes: +# ONNX's reference evaluator implements no other rank, and the backend corpus tests none, so +# nothing could vouch for the code an N-d sampler would emit. +_DEFORM_SPATIAL_RANK = 2 + + +@dataclass(frozen=True) +class Geometry: + """The convolution's shape, resolved to the literals the kernel walks it with.""" + + batch_count: int + groups: int + group_channels: int + group_filters: int + input_shape: tuple[int, ...] + output_shape: tuple[int, ...] + window_shape: tuple[int, ...] + strides: tuple[int, ...] + dilations: tuple[int, ...] + pads: tuple[int, ...] + + @property + def result_shape(self) -> tuple[int, ...]: + return ( + self.batch_count, + self.groups * self.group_filters, + *self.output_shape, + ) + + @property + def arguments(self) -> list[str]: + """Call-site literals for the geometry parameters a window kernel takes.""" + return [ + f"{self.batch_count}u", + f"{self.groups}u", + f"{self.group_channels}u", + f"{self.group_filters}u", + f"{math.prod(self.input_shape)}u", + f"{math.prod(self.output_shape)}u", + f"{math.prod(self.window_shape)}u", + str(len(self.output_shape)), + extents(self.input_shape), + extents(self.output_shape), + extents(self.window_shape), + extents(self.strides), + extents(self.dilations), + offsets(self.pads), + ] + + +def _conv(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + weights = context.require_input(1) + bias = context.optional_input(2) + result = context.require_output(0) + geometry = convolution_geometry(context, source, weights) + verify_bias(context, bias, geometry.groups * geometry.group_filters) + verify_shape(context, result, geometry.result_shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + name = kernel_name(context, element) + definition = _TEMPLATE.substitute( + name=name, + element=element, + zero=scalar_literal(0, result.elem_type), + parameters=WINDOW_PARAMETERS, + ) + call = call_kernel( + name, + [ + result.expr, + source.expr, + weights.expr, + "NULL" if bias is None else bias.expr, + *geometry.arguments, + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def convolution_geometry( + context: NodeContext, source: TensorRef, weights: TensorRef +) -> Geometry: + """The walk a forward convolution takes, shared with the quantized convolutions.""" + rank = _spatial_rank(context, source, weights) + groups = _groups(context) + filters, group_channels = weights.shape[0], weights.shape[1] + if source.shape[1] != group_channels * groups or filters % groups: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` in {groups} group(s) takes " + f"{group_channels * groups} input and a multiple of {groups} output " + f"channel(s), but `{source.name}` carries {source.shape[1]} and " + f"`{weights.name}` produces {filters}." + ) + + window_shape = weights.shape[2:] + _verify_kernel_shape(context, window_shape) + strides = spatial_attribute(context, "strides", rank, 1) + dilations = spatial_attribute(context, "dilations", rank, 1) + begins, ends = resolve_pads( + context, source.shape[2:], window_shape, dilations, strides + ) + return Geometry( + batch_count=source.shape[0], + groups=groups, + group_channels=group_channels, + group_filters=filters // groups, + input_shape=source.shape[2:], + output_shape=output_extents( + source.shape[2:], window_shape, dilations, strides, begins, ends + ), + window_shape=window_shape, + strides=strides, + dilations=dilations, + pads=begins, + ) + + +# -------------------------------------------------------------------------------------- +# The transposed convolution +# -------------------------------------------------------------------------------------- + + +def _conv_transpose(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + weights = context.require_input(1) + bias = context.optional_input(2) + result = context.require_output(0) + geometry = _transpose_geometry(context, source, weights) + verify_bias(context, bias, geometry.groups * geometry.group_filters) + verify_shape(context, result, geometry.result_shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + name = kernel_name(context, element) + definition = _TRANSPOSE_TEMPLATE.substitute( + name=name, + element=element, + zero=scalar_literal(0, result.elem_type), + parameters=WINDOW_PARAMETERS, + ) + call = call_kernel( + name, + [ + result.expr, + source.expr, + weights.expr, + "NULL" if bias is None else bias.expr, + *geometry.arguments, + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _transpose_geometry( + context: NodeContext, source: TensorRef, weights: TensorRef +) -> Geometry: + rank = _spatial_rank(context, source, weights) + groups = _groups(context) + channels, group_filters = weights.shape[0], weights.shape[1] + # `W` is (C, M / group, ...) here rather than Conv's (M, C / group, ...): a transposed + # convolution scatters each input channel over the filters instead of gathering. + if source.shape[1] != channels or channels % groups: + raise CompileError( + f"Node `{context.label}`: `ConvTranspose` in {groups} group(s) takes a filter " + f"holding one stack per input channel, in a multiple of {groups}, but " + f"`{source.name}` carries {source.shape[1]} channel(s) against " + f"`{weights.name}`'s {channels}." + ) + + window_shape = weights.shape[2:] + _verify_kernel_shape(context, window_shape) + strides = spatial_attribute(context, "strides", rank, 1) + dilations = spatial_attribute(context, "dilations", rank, 1) + begins, output_shape = _transpose_pads( + context, source.shape[2:], window_shape, dilations, strides + ) + return Geometry( + batch_count=source.shape[0], + groups=groups, + group_channels=channels // groups, + group_filters=group_filters, + input_shape=source.shape[2:], + output_shape=output_shape, + window_shape=window_shape, + strides=strides, + dilations=dilations, + pads=begins, + ) + + +def _transpose_pads( + context: NodeContext, + input_shape: Sequence[int], + window_shape: Sequence[int], + dilations: Sequence[int], + strides: Sequence[int], +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """The pad before each spatial axis and the result's extents, as ONNX derives them. + + A transposed convolution reaches `stride * (extent - 1) + output_padding + the window's + dilated span` elements, and the pads crop that reach down to the result. Stating an + `output_shape` or an `auto_pad` mode picks the result first and derives the pads from it + instead, which is what makes the stride's ambiguity — a stride of 2 maps two operand + extents onto one result extent — expressible. + """ + rank = len(input_shape) + mode = auto_pad_mode(context) + declared = declared_pads(context, rank, mode) + output_padding = spatial_attribute(context, "output_padding", rank, 0, minimum=0) + reach = tuple( + stride * (extent - 1) + padding + (window - 1) * dilation + 1 + for extent, window, dilation, stride, padding in zip( + input_shape, window_shape, dilations, strides, output_padding + ) + ) + requested = spatial_extents(context, "output_shape", rank, minimum=0) + if declared is not None: + begins, ends = declared + return begins, requested or tuple( + span - begin - end for span, begin, end in zip(reach, begins, ends) + ) + if requested is None and mode in ("NOTSET", "VALID"): + return (0,) * rank, reach + + # ONNX's SAME modes pad so that the result measures `extent * stride`; an explicit + # `output_shape` names it outright. Either way the pads are what is left over, split + # between the two ends — with the odd one going to the end SAME_UPPER names, as the + # spec's own equations put it. A requested result the reach falls short of needs no pad + # at all: the positions past the reach are ones no tap contributes to. A SAME mode is + # left unclamped instead, so that a window whose span is narrower than its stride — where + # ONNX's own shape inference stops agreeing with its equations — is caught by the result + # shape rather than resolved here. + output_shape = requested or tuple( + extent * stride for extent, stride in zip(input_shape, strides) + ) + totals = [span - extent for span, extent in zip(reach, output_shape)] + if requested is not None: + totals = [max(total, 0) for total in totals] + begins = tuple( + total // 2 if mode == "SAME_UPPER" else total - total // 2 for total in totals + ) + return begins, output_shape + + +# -------------------------------------------------------------------------------------- +# The deformable convolution +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _Deformation: + """What a deformable convolution reads beyond a plain one: offsets, and their groups.""" + + geometry: Geometry + offset_groups: int + + @property + def offset_shape(self) -> tuple[int, ...]: + """`offset`: two coordinates per tap per offset group, at every result position.""" + return ( + self.geometry.batch_count, + self.offset_groups * math.prod(self.geometry.window_shape) * 2, + *self.geometry.output_shape, + ) + + @property + def mask_shape(self) -> tuple[int, ...]: + return (self.offset_shape[0], self.offset_shape[1] // 2, *self.offset_shape[2:]) + + +def _deform_conv(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + weights = context.require_input(1) + offsets = context.require_input(2) + bias = context.optional_input(3) + mask = context.optional_input(4) + result = context.require_output(0) + deformation = _deform_geometry(context, source, weights) + geometry = deformation.geometry + _verify_operand(context, offsets, deformation.offset_shape) + _verify_operand(context, mask, deformation.mask_shape) + verify_bias(context, bias, geometry.groups * geometry.group_filters) + verify_shape(context, result, geometry.result_shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + sampler = CFunction( + f"{context.prefix}_bilinear_{element}", + _SAMPLE_TEMPLATE.substitute( + name=f"{context.prefix}_bilinear_{element}", + element=element, + f=math_suffix(result.elem_type), + one=scalar_literal(1, result.elem_type), + zero=scalar_literal(0, result.elem_type), + ), + ) + name = kernel_name(context, element) + definition = _DEFORM_TEMPLATE.substitute( + name=name, + element=element, + zero=scalar_literal(0, result.elem_type), + sample=sampler.name, + ) + call = call_kernel( + name, + [ + result.expr, + source.expr, + weights.expr, + offsets.expr, + "NULL" if bias is None else bias.expr, + "NULL" if mask is None else mask.expr, + f"{geometry.batch_count}u", + f"{geometry.groups}u", + f"{geometry.group_channels}u", + f"{geometry.group_filters}u", + f"{geometry.groups * geometry.group_channels // deformation.offset_groups}u", + f"{deformation.offset_groups}u", + *(f"{extent}u" for extent in geometry.input_shape), + *(f"{extent}u" for extent in geometry.output_shape), + *(f"{extent}u" for extent in geometry.window_shape), + *(f"{stride}u" for stride in geometry.strides), + *(f"{dilation}u" for dilation in geometry.dilations), + *(str(pad) for pad in geometry.pads), + ], + ) + return NodeEmission( + functions=(sampler, CFunction(name, definition)), statements=(call,) + ) + + +def _deform_geometry( + context: NodeContext, source: TensorRef, weights: TensorRef +) -> _Deformation: + rank = _spatial_rank(context, source, weights) + if rank != _DEFORM_SPATIAL_RANK: + raise CompileError( + f"Node `{context.label}`: `DeformConv` is compiled for {_DEFORM_SPATIAL_RANK} " + f"spatial axes — the rank ONNX's reference evaluator and backend tests cover — " + f"but `{source.name}` of shape {list(source.shape)} has {rank}." + ) + groups = _groups(context) + offset_groups = context.int_attribute("offset_group") + filters, group_channels = weights.shape[0], weights.shape[1] + if source.shape[1] != group_channels * groups or filters % groups: + raise CompileError( + f"Node `{context.label}`: `DeformConv` in {groups} group(s) takes " + f"{group_channels * groups} input and a multiple of {groups} output " + f"channel(s), but `{source.name}` carries {source.shape[1]} and " + f"`{weights.name}` produces {filters}." + ) + if offset_groups < 1 or source.shape[1] % offset_groups: + raise CompileError( + f"Node `{context.label}`: `DeformConv` splits its {source.shape[1]} input " + f"channel(s) into {offset_groups} offset group(s); ONNX defines " + "`offset_group` as a positive count that divides them." + ) + + window_shape = weights.shape[2:] + _verify_kernel_shape(context, window_shape) + strides = spatial_attribute(context, "strides", rank, 1) + dilations = spatial_attribute(context, "dilations", rank, 1) + begins, ends = declared_pads(context, rank, "NOTSET") or ((0,) * rank, (0,) * rank) + return _Deformation( + geometry=Geometry( + batch_count=source.shape[0], + groups=groups, + group_channels=group_channels, + group_filters=filters // groups, + input_shape=source.shape[2:], + output_shape=output_extents( + source.shape[2:], window_shape, dilations, strides, begins, ends + ), + window_shape=window_shape, + strides=strides, + dilations=dilations, + pads=begins, + ), + offset_groups=offset_groups, + ) + + +# -------------------------------------------------------------------------------------- +# Reading the geometry off the node +# -------------------------------------------------------------------------------------- + + +def _spatial_rank(context: NodeContext, source: TensorRef, weights: TensorRef) -> int: + """The number of axes the window slides along, once both operands agree on it.""" + op_type = context.node.op_type + if len(source.shape) < 3: + raise CompileError( + f"Node `{context.label}`: `{op_type}` takes a batch of multi-channel signals " + f"— a tensor of rank 3 or more — but `{source.name}` has shape " + f"{list(source.shape)}." + ) + if len(weights.shape) != len(source.shape): + raise CompileError( + f"Node `{context.label}`: `{op_type}` convolves `{source.name}` of shape " + f"{list(source.shape)} with `{weights.name}` of shape " + f"{list(weights.shape)}; ONNX defines both as rank " + f"{len(source.shape)} — two leading axes and one per spatial axis." + ) + return len(source.shape) - 2 + + +def _groups(context: NodeContext) -> int: + groups = context.int_attribute("group") + if groups < 1: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` splits its channels into " + f"{groups} group(s); ONNX defines `group` as a positive count." + ) + return groups + + +def verify_bias(context: NodeContext, bias: TensorRef | None, channels: int) -> None: + if bias is not None and bias.shape != (channels,): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` takes one bias per output " + f"channel, but `{bias.name}` has shape {list(bias.shape)} against {channels} " + "channel(s)." + ) + + +def _verify_operand( + context: NodeContext, operand: TensorRef | None, expected: Sequence[int] +) -> None: + """Refuse an operand the emitted addressing would read outside of. + + `offset` and `mask` are indexed by the geometry the other operands and the attributes + fix, so one shaped for a different geometry is a compile error rather than a read past + the end of a buffer. + """ + if operand is not None and operand.shape != tuple(expected): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` addresses `{operand.name}` " + f"as {list(expected)}, but it has shape {list(operand.shape)}." + ) + + +def _verify_kernel_shape(context: NodeContext, window_shape: tuple[int, ...]) -> None: + """Refuse a `kernel_shape` that disagrees with the filter the node is actually handed. + + ONNX defines the attribute as inferred from `W` when absent, and says nothing about what + a node stating a different one means — the reference pads for one shape and convolves + with the other — so it is a compile error rather than a guess. + """ + declared = context.attribute("kernel_shape", None) + if declared is not None and tuple(int(value) for value in declared) != window_shape: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` declares `kernel_shape` " + f"{[int(value) for value in declared]}, but the filter it is handed measures " + f"{list(window_shape)}." + ) + + +# -------------------------------------------------------------------------------------- +# Folding the columns back into an image +# -------------------------------------------------------------------------------------- + + +def _col2im(context: NodeContext) -> NodeEmission: + """Col2Im: every block written back over the image positions it was read from. + + The columns hold one value per (tap, block position) pair, and a block position is where + a window of `block_shape` sat: the same geometry a convolution slides. So each image + position sums the taps that reached it, found by the transposed convolution's own walk — + `(position + pad - tap * dilation) / stride` — which is the relation solved for the block + the tap came from, and which drops the taps the stride does not divide. + """ + source = context.require_input(0) + result = context.require_output(0) + image_shape = _shape_operand(context, 1, "image_shape") + block_shape = _shape_operand(context, 2, "block_shape") + rank = len(image_shape) + if len(block_shape) != rank: + raise CompileError( + f"Node `{context.label}`: `Col2Im` was given {len(block_shape)} block extent(s) " + f"for a {rank}-dimensional image; ONNX defines one per spatial axis." + ) + if len(source.shape) != 3: + raise CompileError( + f"Node `{context.label}`: `Col2Im` takes a batch of column stacks — a tensor of " + f"rank 3 — but `{source.name}` has shape {list(source.shape)}." + ) + + if result.elem_type == TensorProto.BOOL: + raise CompileError( + f"Node `{context.label}`: `Col2Im` of a " + f"`{element_type_name(result.elem_type)}` tensor is not supported by the C " + "compiler; summing truth values has no defined result." + ) + + window_size = math.prod(block_shape) + if window_size < 1 or source.shape[1] % window_size != 0: + raise CompileError( + f"Node `{context.label}`: `Col2Im` folds blocks of {window_size} value(s), " + f"which does not divide the {source.shape[1]} row(s) of `{source.name}` into " + "whole channels." + ) + channels = source.shape[1] // window_size + verify_shape(context, result, (source.shape[0], channels, *image_shape)) + + dilations = spatial_attribute(context, "dilations", rank, 1) + strides = spatial_attribute(context, "strides", rank, 1) + begins, ends = declared_pads(context, rank, "NOTSET") or ((0,) * rank, (0,) * rank) + column_shape = output_extents( + image_shape, block_shape, dilations, strides, begins, ends + ) + if ( + any(extent < 1 for extent in column_shape) + or math.prod(column_shape) != source.shape[2] + ): + raise CompileError( + f"Node `{context.label}`: `Col2Im` places {list(column_shape)} block(s) over an " + f"image of {list(image_shape)}, but `{source.name}` holds {source.shape[2]} " + "column(s)." + ) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + name = f"{context.prefix}_col2im_{element}" + definition = _COL2IM_TEMPLATE.substitute( + name=name, element=element, zero=scalar_literal(0, result.elem_type) + ) + call = call_kernel( + name, + [ + result.expr, + source.expr, + f"{source.shape[0] * channels}u", + f"{source.shape[2]}u", + f"{math.prod(image_shape)}u", + f"{window_size}u", + str(rank), + extents(image_shape), + extents(block_shape), + extents(column_shape), + extents(strides), + extents(dilations), + offsets(begins), + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _shape_operand(context: NodeContext, index: int, role: str) -> tuple[int, ...]: + values = context.constant_input(index) + if values is None: + raise CompileError( + f"Node `{context.label}`: `Col2Im` takes its `{role}` from " + f"`{context.require_input(index).name}`, which is not known at compile time; " + "the shape of the result then depends on input data, which the C compiler " + "cannot compile." + ) + extents_ = tuple(int(value) for value in values.reshape(-1)) + if any(extent < 1 for extent in extents_): + raise CompileError( + f"Node `{context.label}`: `Col2Im` was given `{role}` {list(extents_)}; ONNX " + "defines them as extents, which are positive." + ) + return extents_ + + +register_kernel("", "Conv", _VERSIONS, _conv) +register_kernel("", "ConvTranspose", _VERSIONS, _conv_transpose) +register_kernel("", "DeformConv", _VERSIONS, _deform_conv) +register_kernel("", "Col2Im", _COL2IM_VERSIONS, _col2im) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/cumulative.py b/src/python/fnnx/extras/compilers/c/onnx/ops/cumulative.py new file mode 100644 index 0000000..d491bec --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/cumulative.py @@ -0,0 +1,188 @@ +"""CumSum and CumProd: the running fold along one axis. + +The axis is an operand rather than an attribute, and models — the ONNX backend corpus among +them — do pass it at run time. It decides which elements the scan visits but nothing about +any shape, so a run-time axis still compiles to fully static code: one call site per axis the +operand's rank allows, chosen by a switch, with an out-of-range value returning the argument +error the status enum exists for. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from functools import partial +from string import Template + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type, numpy_dtype_name +from fnnx.extras.compilers.c.onnx.emit import INVALID_ARGUMENT_STATUS +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + GROUP_PARAMETERS, + call_kernel, + group_axes, + kernel_name, + normalize_axis, + offset_helper, + verify_same_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import expand + +# CumSum-14 only added bfloat16 to CumSum-11's type constraints; CumProd arrived at 26. +_CUM_SUM_VERSIONS = (11, 14) +_CUM_PROD_VERSIONS = (26,) + +# `exclusive` and `reverse` are parameters rather than four kernels: the call site passes +# literals, so a C compiler folds the branches away and the artifact still carries one scan. +_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, +$parameters, + int exclusive, + int reverse) +{ + size_t group, index; + for (group = 0; group < group_count; ++group) { + const size_t base = $offset(group, kept_rank, kept_shape, kept_strides); + $element total = $identity; + for (index = 0; index < group_size; ++index) { + const size_t position = base + $offset( + reverse ? group_size - 1 - index : index, + reduced_rank, reduced_shape, reduced_strides); + const $element x = in[position]; + if (exclusive) { + out[position] = total; + } + total = ($element)($combine); + if (!exclusive) { + out[position] = total; + } + } + } +}""") + + +def _cumulative(context: NodeContext, *, identity: str, combine: str) -> NodeEmission: + source = context.require_input(0) + result = context.require_output(0) + axis_operand = context.require_input(1) + verify_same_shape(context, source, result) + rank = len(source.shape) + if rank == 0: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` scans along an axis, but " + f"`{source.name}` is a scalar and has none." + ) + if axis_operand.elem_count != 1: + raise CompileError( + f"Node `{context.label}`: the axis of `{context.node.op_type}` comes from " + f"`{axis_operand.name}`, which holds {axis_operand.elem_count} values; ONNX " + "defines it as a single one." + ) + + offset = offset_helper(context.prefix) + name = kernel_name(context, numpy_dtype_name(result.elem_type)) + definition = _TEMPLATE.substitute( + name=name, + element=c_type(result.elem_type), + parameters=GROUP_PARAMETERS, + offset=offset.name, + identity=expand(identity, result.elem_type), + combine=expand(combine, result.elem_type), + ) + + def call(axis: int) -> str: + grouping = group_axes(source.shape, (axis,)) + return call_kernel( + name, + [ + result.expr, + source.expr, + *grouping.arguments, + str(context.int_attribute("exclusive")), + str(context.int_attribute("reverse")), + ], + ) + + fixed = context.constant_input(1) + if fixed is not None: + return NodeEmission( + functions=(offset, CFunction(name, definition)), + statements=( + call(normalize_axis(context, int(fixed.reshape(-1)[0]), rank)), + ), + ) + normalize = _normalize_helper(context.prefix) + return NodeEmission( + functions=(offset, normalize, CFunction(name, definition)), + statements=(_dispatch(context, axis_operand.expr, rank, call, normalize),), + ) + + +def _normalize_helper(prefix: str) -> CFunction: + """An axis counted from the end, resolved against a rank, both known only as values.""" + name = f"{prefix}_normalized_axis" + return CFunction( + name, + "\n".join( + [ + f"static int64_t {name}(int64_t axis, int64_t rank)", + "{", + " return (axis < 0) ? (axis + rank) : axis;", + "}", + ] + ), + ) + + +def _dispatch( + context: NodeContext, + operand: str, + rank: int, + call: Callable[[int], str], + normalize: CFunction, +) -> str: + """The scan for whichever axis the operand names at run time, or an argument error. + + The axis is resolved through a function rather than into a local, so that the statement + introduces no identifier of its own — one would shadow the entrypoint parameter a tensor + of the same name is emitted as. + """ + cases = [] + for axis in range(rank): + cases.append(f"case {axis}:") + cases.extend(_indented(call(axis), " ")) + cases.append(" break;") + return "\n".join( + [ + f"switch ({normalize.name}((int64_t){operand}[0], {rank})) {{", + *cases, + "default:", + f" return {context.prefix.upper()}_{INVALID_ARGUMENT_STATUS};", + "}", + ] + ) + + +def _indented(statement: str, indent: str) -> Sequence[str]: + return [f"{indent}{line}" if line else "" for line in statement.splitlines()] + + +register_kernel( + "", + "CumSum", + _CUM_SUM_VERSIONS, + partial(_cumulative, identity="$zero", combine="total + x"), +) +register_kernel( + "", + "CumProd", + _CUM_PROD_VERSIONS, + partial(_cumulative, identity="$one", combine="total * x"), +) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/einsum.py b/src/python/fnnx/extras/compilers/c/onnx/ops/einsum.py new file mode 100644 index 0000000..4869a29 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/einsum.py @@ -0,0 +1,379 @@ +"""Einsum: the equation read as addressing, and the sum it leaves behind. + +ONNX defines Einsum as numpy's, and numpy's is a statement about coordinates: every label in +the equation names an extent, every operand's elements are addressed by the labels its term +carries, and the result is the sum — over the labels the output term leaves out — of the +operands multiplied together. The equation is therefore a compile-time object, one stride per +operand per label, and what is left for the kernel is two loops: one over the result's +elements, one over the labels being summed. Every equation of a given arity and element type +comes out as the same kernel, told apart only by the extents and strides its call site passes. + +A label repeated inside one term is a diagonal, and needs no code of its own: the strides of +the axes it names add up, so stepping that one coordinate steps along the diagonal. A label a +term does not carry contributes a zero stride, which is what makes an outer product, a summed +axis and a stretched operand the same loop. + +numpy stretches a labelled axis of extent 1 against another operand's, and ONNX's own shape +inference does not. Where that disagreement reaches the result — the equation computes a +wider tensor than the buffer ONNX sized — the shape check refuses the node rather than +writing past it; where it stays inside a summed label, the sum is numpy's, which is what the +reference evaluator computes. +""" + +from __future__ import annotations + +import math +import re +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from string import Template + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + kernel_name, + row_major_strides, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents + +# Einsum has had a single revision since it arrived at opset 12. +_VERSIONS = (12,) + +# A term is letters with at most one `...` among them; numpy reads nothing else, and neither +# does this. +_TERM = re.compile(r"([a-zA-Z]*)(\.\.\.)?([a-zA-Z]*)") + +# What each axis an ellipsis covers is labelled with. The dots make it unspellable in an +# equation, so it can never collide with a label the equation names itself. +_BROADCAST_LABEL = "...{}" + +_TEMPLATE = Template("""\ +static void $name( + $element* out, +$operands, + size_t count, + int rank, + const size_t* shape, +$result_strides, + size_t summed_count, + int summed_rank, + const size_t* summed_shape, +$summed_strides) +{ + size_t index, term; + for (index = 0; index < count; ++index) { + size_t remainder = index; + $element total = $zero; +$bases + int axis; + for (axis = rank - 1; axis >= 0; --axis) { + const size_t coordinate = remainder % shape[axis]; + remainder /= shape[axis]; +$walk_result + } + for (term = 0; term < summed_count; ++term) { + size_t rest = term; +$offsets + int summed; + for (summed = summed_rank - 1; summed >= 0; --summed) { + const size_t coordinate = rest % summed_shape[summed]; + rest /= summed_shape[summed]; +$walk_summed + } + total += $product; + } + out[index] = total; + } +}""") + + +@dataclass(frozen=True) +class _Equation: + """The equation as addressing: one label per axis of each operand and of the result.""" + + terms: tuple[tuple[str, ...], ...] + result: tuple[str, ...] + extents: Mapping[str, int] + + @property + def summed(self) -> tuple[str, ...]: + """The labels the result leaves out, in the order the terms first name them.""" + kept = frozenset(self.result) + ordered = dict.fromkeys(label for term in self.terms for label in term) + return tuple(label for label in ordered if label not in kept) + + +def _einsum(context: NodeContext) -> NodeEmission: + operands = tuple( + context.require_input(index) for index in range(len(context.node.input)) + ) + result = context.require_output(0) + equation = _parse(context, operands) + verify_shape( + context, result, [equation.extents[label] for label in equation.result] + ) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + summed = equation.summed + strides = [ + _label_strides(term, operand.shape, equation.extents) + for term, operand in zip(equation.terms, operands) + ] + element = c_type(result.elem_type) + name = kernel_name(context, str(len(operands)), element) + arguments = [ + result.expr, + *(operand.expr for operand in operands), + f"{result.elem_count}u", + str(len(equation.result)), + extents([equation.extents[label] for label in equation.result]), + *(_stride_literal(stride, equation.result) for stride in strides), + f"{math.prod(equation.extents[label] for label in summed)}u", + str(len(summed)), + extents([equation.extents[label] for label in summed]), + *(_stride_literal(stride, summed) for stride in strides), + ] + definition = _definition(name, element, len(operands), result.elem_type) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call_kernel(name, arguments),), + ) + + +def _definition(name: str, element: str, arity: int, elem_type: int) -> str: + """The kernel for `arity` operands of `element`, which every equation shares.""" + operands = range(arity) + return _TEMPLATE.substitute( + name=name, + element=element, + zero=scalar_literal(0, elem_type), + operands=",\n".join(f" const {element}* in{index}" for index in operands), + result_strides=",\n".join( + f" const size_t* result_strides{index}" for index in operands + ), + summed_strides=",\n".join( + f" const size_t* summed_strides{index}" for index in operands + ), + bases="\n".join(f" size_t base{index} = 0;" for index in operands), + walk_result="\n".join( + f" base{index} += coordinate * result_strides{index}[axis];" + for index in operands + ), + offsets="\n".join( + f" size_t offset{index} = base{index};" for index in operands + ), + walk_summed="\n".join( + f" offset{index} += coordinate * summed_strides{index}[summed];" + for index in operands + ), + product=" * ".join(f"in{index}[offset{index}]" for index in operands), + ) + + +def _stride_literal(strides: Mapping[str, int], labels: Sequence[str]) -> str: + return extents([strides.get(label, 0) for label in labels]) + + +def _parse(context: NodeContext, operands: Sequence[TensorRef]) -> _Equation: + """The equation, with every ellipsis expanded against the operands' own ranks. + + A term that does not address the operand it is paired with — one term too few, or a rank + the term does not cover — is refused here as well as by ONNX's own shape inference, which + is what a model reaches first: it would otherwise be a stride addressing a buffer by + another operand's shape, and this is where that stops rather than where it reads past one. + """ + text = _equation_text(context) + term_texts, result_text = _split(context, text) + if len(term_texts) != len(operands): + raise CompileError( + f"Node `{context.label}`: the equation `{text}` states {len(term_texts)} " + f"term(s) for {len(operands)} operand(s); Einsum takes one term per operand." + ) + terms = [_split_term(context, term) for term in term_texts] + covered = [ + _ellipsis_rank(context, text, term, operand) + for text, term, operand in zip(term_texts, terms, operands) + ] + broadcast = tuple( + _BROADCAST_LABEL.format(axis) for axis in range(max(covered, default=0)) + ) + expanded = tuple( + (*head, *broadcast[len(broadcast) - count :], *tail) + for (head, _, tail), count in zip(terms, covered) + ) + return _Equation( + terms=expanded, + result=_result_labels(context, result_text, expanded, broadcast), + extents=_measure(context, expanded, operands), + ) + + +def _equation_text(context: NodeContext) -> str: + """The `equation` attribute, with the spaces numpy ignores taken out. + + An equation of nothing at all is refused rather than read as the scalar term numpy takes + it for: ONNX's own reference implementation rejects it outright, so nothing states what + such a node computes. + """ + value = context.attribute("equation", b"") + text = value.decode() if isinstance(value, bytes) else str(value) + stripped = text.strip().replace(" ", "") + if not stripped: + raise CompileError( + f"Node `{context.label}`: its `equation` is empty, which ONNX's Einsum does " + "not define." + ) + return stripped + + +def _split(context: NodeContext, text: str) -> tuple[list[str], str | None]: + """The equation's input terms, and its output term where it states one.""" + parts = text.split("->") + if len(parts) > 2: + raise CompileError( + f"Node `{context.label}`: the equation `{text}` states its output more than " + "once; ONNX's Einsum writes `->` at most once." + ) + return parts[0].split(","), parts[1] if len(parts) == 2 else None + + +def _split_term( + context: NodeContext, text: str +) -> tuple[tuple[str, ...], bool, tuple[str, ...]]: + """One term as the labels before its ellipsis, whether it has one, and the ones after.""" + match = _TERM.fullmatch(text) + if match is None: + raise CompileError( + f"Node `{context.label}`: `{text}` is not a term ONNX's Einsum defines: a term " + "is written as letters, with at most one `...` among them." + ) + head, ellipsis, tail = match.groups() + return tuple(head), ellipsis is not None, tuple(tail) + + +def _ellipsis_rank( + context: NodeContext, + text: str, + term: tuple[tuple[str, ...], bool, tuple[str, ...]], + operand: TensorRef, +) -> int: + """How many of the operand's axes this term's ellipsis stands for.""" + head, ellipsis, tail = term + named = len(head) + len(tail) + if named > len(operand.shape) or (not ellipsis and named != len(operand.shape)): + raise CompileError( + f"Node `{context.label}`: the term `{text}` names {named} label(s) for " + f"`{operand.name}`, which has shape {list(operand.shape)}; a term names one " + "label per axis, unless it carries `...` for the rest." + ) + return len(operand.shape) - named if ellipsis else 0 + + +def _result_labels( + context: NodeContext, + text: str | None, + terms: Sequence[tuple[str, ...]], + broadcast: tuple[str, ...], +) -> tuple[str, ...]: + """The result's axes: the output term's labels, or the ones implicit mode leaves. + + Implicit mode is numpy's: the axes an ellipsis covers come first, then every label + exactly one axis of the whole equation carries, in alphabetical order. + """ + carried = Counter(label for term in terms for label in term) + if text is None: + return ( + *broadcast, + *sorted( + label + for label, count in carried.items() + if count == 1 and label not in broadcast + ), + ) + head, ellipsis, tail = _split_term(context, text) + labels = (*head, *(broadcast if ellipsis else ()), *tail) + for position, label in enumerate(labels): + if label in labels[:position]: + raise CompileError( + f"Node `{context.label}`: the output term `{text}` names label `{label}` " + "more than once; each axis of the result is one label." + ) + if label not in carried: + raise CompileError( + f"Node `{context.label}`: the output term `{text}` names label `{label}`, " + "which no operand's term carries." + ) + return labels + + +def _measure( + context: NodeContext, + terms: Sequence[tuple[str, ...]], + operands: Sequence[TensorRef], +) -> dict[str, int]: + """Every label's extent, over each axis of each operand that carries it. + + numpy reads the two positions a label can repeat in differently: inside one term it is a + diagonal, which is defined only over axes measuring alike, while across two operands it + broadcasts, an extent of 1 stretching to the other's. + """ + sizes: dict[str, int] = {} + for term, operand in zip(terms, operands): + carried: dict[str, int] = {} + for label, extent in zip(term, operand.shape): + diagonal = carried.setdefault(label, extent) + if diagonal != extent: + raise CompileError( + f"Node `{context.label}`: the term for `{operand.name}` names label " + f"`{label}` on two axes of shape {list(operand.shape)} measuring " + f"{diagonal} and {extent}; a repeated label is a diagonal, which ONNX " + "defines only over axes of equal extent." + ) + for label, extent in carried.items(): + sizes[label] = _stretched(context, sizes.get(label), extent, label) + return sizes + + +def _stretched( + context: NodeContext, measured: int | None, extent: int, label: str +) -> int: + """One label's extent so far against another operand's, which a 1 stretches to.""" + if measured is None or measured in (extent, 1): + return extent + if extent == 1: + return measured + raise CompileError( + f"Node `{context.label}`: label `{label}` measures {measured} on one operand and " + f"{extent} on another; ONNX's Einsum stretches an extent of 1 against another " + "operand's, and defines nothing for two extents that differ otherwise." + ) + + +def _label_strides( + term: Sequence[str], shape: Sequence[int], sizes: Mapping[str, int] +) -> dict[str, int]: + """How far one step along each label moves through this operand's row-major buffer. + + The strides of every axis a label names add up, which is what walks a diagonal; an axis + the operand is stretched along contributes nothing, which is what broadcasts it. + """ + strides: dict[str, int] = {} + for label, stride, extent in zip(term, row_major_strides(shape), shape): + if extent == sizes[label]: + strides[label] = strides.get(label, 0) + stride + return strides + + +register_kernel("", "Einsum", _VERSIONS, _einsum) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/elementwise.py b/src/python/fnnx/extras/compilers/c/onnx/ops/elementwise.py new file mode 100644 index 0000000..428f33f --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/elementwise.py @@ -0,0 +1,400 @@ +"""Elementwise math: broadcasting arithmetic, the variadic families, and pointwise math.""" + +from __future__ import annotations + +from functools import partial +from string import Template + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + FLOAT_TYPES, + UNSIGNED_TYPES, + c_type, + element_type_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + copy_tensor, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import ( + Scalar, + combiner, + elementwise, + expand, + math_suffix, + pointwise, +) + +# Numpy-style multidirectional broadcasting arrived at opset 7; 13 and 14 only widened the +# type constraints. +_ARITHMETIC_VERSIONS = (7, 13, 14) +_ARITHMETIC_OPERATORS = {"Add": "+", "Sub": "-", "Mul": "*", "Div": "/"} + +# The variadic families took equal shapes until opset 8 introduced broadcasting; the later +# revisions widened types. Sum and Mean were never defined for the integer families. +_EXTREMUM_VERSIONS = (8, 12, 13) +_ACCUMULATE_VERSIONS = (8, 13) + +# Pow-7 introduced broadcasting, 12 let the exponent carry its own type, and 13 and 15 only +# widened the type constraints. +_POW_VERSIONS = (7, 12, 13, 15) +_MOD_VERSIONS = (10, 13) +# Identity has never changed for tensors; every revision listed adds element or container +# types, and the ones this compiler does not support are rejected before dispatch. +_IDENTITY_VERSIONS = (1, 13, 14, 16, 19, 21, 23, 24, 25) +# Dropout-7 is inference-only; 10 made the mask output boolean; 12 added the ratio and +# training_mode inputs; 13 and 22 only changed which element types are allowed. The kernel +# takes the mask's element type from the graph, so the same code serves all of them. +_DROPOUT_VERSIONS = (7, 10, 12, 13, 22) +# Clip-6 carries its bounds as attributes and 11 moved them into optional inputs; 12 and 13 +# widened the types. +_CLIP_ATTRIBUTE_VERSIONS = (6,) +_CLIP_INPUT_VERSIONS = (11, 12, 13) + +# Ops with one operand whose result is a libm call or a short expression over it, each with +# the schema revisions the formula covers. The legacy revision 1 of the older ones carried a +# `consumed_inputs` attribute and is left unregistered. +_UNARY_MATH: dict[str, tuple[tuple[int, ...], str]] = { + "Acos": ((7, 22), "acos$f(x0)"), + "Acosh": ((9, 22), "acosh$f(x0)"), + "Asin": ((7, 22), "asin$f(x0)"), + "Asinh": ((9, 22), "asinh$f(x0)"), + "Atan": ((7, 22), "atan$f(x0)"), + "Atanh": ((9, 22), "atanh$f(x0)"), + "Ceil": ((6, 13), "ceil$f(x0)"), + "Cos": ((7, 22), "cos$f(x0)"), + "Cosh": ((9, 22), "cosh$f(x0)"), + "Erf": ((9, 13), "erf$f(x0)"), + "Exp": ((6, 13), "exp$f(x0)"), + "Floor": ((6, 13), "floor$f(x0)"), + "Log": ((6, 13), "log$f(x0)"), + "Reciprocal": ((6, 13), "$one / x0"), + # ONNX rounds halves to even, which `rint` does under C's default rounding mode. + "Round": ((11, 22), "rint$f(x0)"), + "Sin": ((7, 22), "sin$f(x0)"), + "Sinh": ((9, 22), "sinh$f(x0)"), + "Sqrt": ((6, 13), "sqrt$f(x0)"), + "Tan": ((7, 22), "tan$f(x0)"), + "Tanh": ((6, 13), "tanh$f(x0)"), +} + +# Abs, Neg and Sign serve the integer families as well, where the sign has to be read off a +# comparison rather than from libm. +_SIGN_VERSIONS = (9, 13) +_ABS_VERSIONS = _NEG_VERSIONS = (6, 13) + +_FLOORED_MOD_TEMPLATE = Template("""\ +static $element $name($element left, $element right) +{ + $element remainder = left % right; + /* C truncates toward zero; ONNX's fmod=0 takes the divisor's sign, as numpy does. */ + if (remainder != 0 && ((remainder < 0) != (right < 0))) { + remainder += right; + } + return remainder; +}""") + +# Integer exponentiation, which numpy performs exactly rather than through `pow`. The +# squaring runs in the unsigned counterpart of the element type so that an overflow wraps — +# defined behaviour, and the same result numpy's integer power gives — instead of being +# undefined signed overflow. +_INTEGER_POW_TEMPLATE = Template("""\ +static $element $name($element base, $exponent exponent) +{ +$guard { + $unsigned accumulator = 1u; + $unsigned factor = ($unsigned)base; + uint64_t remaining = (uint64_t)exponent; + while (remaining > 0) { + if (remaining & 1) { + accumulator = ($unsigned)(accumulator * factor); + } + factor = ($unsigned)(factor * factor); + remaining >>= 1; + } + return ($element)accumulator; + } +}""") + +# ONNX leaves a negative integer exponent undefined and numpy refuses to evaluate it, so +# there is no behaviour to match here — only one that has to be defined. +_NEGATIVE_EXPONENT_GUARD = """\ + if (exponent < 0) { + return 0; + } +""" + +_UNSIGNED_COUNTERPARTS: dict[int, int] = { + TensorProto.INT8: TensorProto.UINT8, + TensorProto.INT16: TensorProto.UINT16, + TensorProto.INT32: TensorProto.UINT32, + TensorProto.INT64: TensorProto.UINT64, +} + + +def _arithmetic(context: NodeContext, *, operator: str) -> NodeEmission: + return elementwise( + context, + expression=f"x0 {operator} x1", + operands=(context.require_input(0), context.require_input(1)), + result=context.require_output(0), + ) + + +def _extremum(context: NodeContext, *, largest: bool) -> NodeEmission: + """Min or Max over any number of operands, folded left as the reference folds them.""" + operands = tuple( + context.require_input(index) for index in range(len(context.inputs)) + ) + result = context.require_output(0) + expression = "x0" + helpers: tuple[CFunction, ...] = () + if len(operands) > 1: + helper = combiner(context, result.elem_type, largest=largest) + helpers = (helper,) + for index in range(1, len(operands)): + expression = f"{helper.name}({expression}, x{index})" + return elementwise( + context, + expression=expression, + operands=operands, + result=result, + helpers=helpers, + ) + + +def _accumulate(context: NodeContext, *, average: bool) -> NodeEmission: + """Sum or Mean: the operands added in order, and for Mean divided by how many there are.""" + operands = tuple( + context.require_input(index) for index in range(len(context.inputs)) + ) + result = context.require_output(0) + total = " + ".join(f"x{index}" for index in range(len(operands))) + expression = ( + f"({total}) / {scalar_literal(len(operands), result.elem_type)}" + if average + else total + ) + return elementwise(context, expression=expression, operands=operands, result=result) + + +def _mod(context: NodeContext) -> NodeEmission: + left = context.require_input(0) + right = context.require_input(1) + result = context.require_output(0) + truncated = bool(context.attribute("fmod", 0)) + if result.elem_type in FLOAT_TYPES and not truncated: + raise CompileError( + f"Node `{context.label}`: Mod on `{element_type_name(result.elem_type)}` " + "tensors requires the `fmod` attribute to be 1, as the ONNX spec does." + ) + helpers: tuple[CFunction, ...] = () + if result.elem_type in FLOAT_TYPES: + expression = f"fmod{math_suffix(result.elem_type)}(x0, x1)" + variant = "_truncated" + elif truncated or result.elem_type in UNSIGNED_TYPES: + # For the unsigned families both definitions agree: no remainder is ever negative. + expression = "x0 % x1" + variant = "_truncated" + else: + name = f"{context.prefix}_floored_mod_{c_type(result.elem_type)}" + helpers = ( + CFunction( + name, + _FLOORED_MOD_TEMPLATE.substitute( + name=name, element=c_type(result.elem_type) + ), + ), + ) + expression = f"{name}(x0, x1)" + variant = "_floored" + return elementwise( + context, + expression=expression, + operands=(left, right), + result=result, + helpers=helpers, + variant=variant, + ) + + +def _pow(context: NodeContext) -> NodeEmission: + """Pow, whose exponent carries an element type of its own from opset 12 on.""" + base = context.require_input(0) + exponent = context.require_input(1) + result = context.require_output(0) + helpers: tuple[CFunction, ...] = () + if result.elem_type in FLOAT_TYPES and exponent.elem_type == result.elem_type: + expression = f"pow{math_suffix(result.elem_type)}(x0, x1)" + elif result.elem_type in FLOAT_TYPES or exponent.elem_type in FLOAT_TYPES: + # numpy promotes a mixed pair to float64 and casts the result back to the base's + # type, so the double-precision call is what has to be matched here. + expression = f"({c_type(result.elem_type)})pow((double)x0, (double)x1)" + else: + helper = _integer_pow(context, result.elem_type, exponent.elem_type) + helpers = (helper,) + expression = f"{helper.name}(x0, x1)" + return elementwise( + context, + expression=expression, + operands=(base, exponent), + result=result, + helpers=helpers, + ) + + +def _integer_pow(context: NodeContext, elem_type: int, exponent_type: int) -> CFunction: + name = f"{context.prefix}_integer_pow_{c_type(elem_type)}_{c_type(exponent_type)}" + return CFunction( + name, + _INTEGER_POW_TEMPLATE.substitute( + name=name, + element=c_type(elem_type), + exponent=c_type(exponent_type), + unsigned=c_type(_UNSIGNED_COUNTERPARTS[elem_type]), + guard="" if exponent_type in UNSIGNED_TYPES else _NEGATIVE_EXPONENT_GUARD, + ), + ) + + +def _abs(context: NodeContext) -> NodeEmission: + result = context.require_output(0) + if result.elem_type in FLOAT_TYPES: + return pointwise(context, "fabs$f(x0)") + if result.elem_type in UNSIGNED_TYPES: + return pointwise(context, "x0") + return pointwise(context, "(x0 < $zero) ? ($element)-x0 : x0") + + +def _sign(context: NodeContext) -> NodeEmission: + result = context.require_output(0) + if result.elem_type in UNSIGNED_TYPES: + return pointwise(context, "($element)(x0 > $zero)") + # numpy's sign leaves NaN alone; every zero comes out `+0` whatever sign it went in + # with, which is what the difference of the two comparisons already gives. + negative = "($element)((x0 > $zero) - (x0 < $zero))" + if result.elem_type in FLOAT_TYPES: + return pointwise(context, f"isnan(x0) ? x0 : {negative}") + return pointwise(context, negative) + + +def _clip_from_attributes(context: NodeContext) -> NodeEmission: + """Clip up to opset 10, whose bounds are attributes defaulting to the float32 extremes.""" + result = context.require_output(0) + low = float(context.attribute("min", _FLOAT_MIN)) + high = float(context.attribute("max", _FLOAT_MAX)) + largest = combiner(context, result.elem_type, largest=True) + smallest = combiner(context, result.elem_type, largest=False) + return pointwise( + context, + f"{smallest.name}(high, {largest.name}(low, x0))", + scalars=( + Scalar("low", result.elem_type, low), + Scalar("high", result.elem_type, high), + ), + helpers=(largest, smallest), + ) + + +def _clip_from_inputs(context: NodeContext) -> NodeEmission: + """Clip from opset 11 on, whose bounds are optional scalar inputs. + + ONNX defines the result as numpy's `clip`, which applies the lower bound first: a lower + bound above the upper one yields the upper one, and a NaN bound wins outright. Which + operand a *tie* yields — all that separates `+0` from `-0` here — differs between the + forms numpy evaluates: with both bounds it is `minimum(max, maximum(min, x))`, keeping + the data's zero, while a single bound goes through plain `maximum(x, min)`, keeping the + bound's. Both are reproduced here rather than unified. + """ + result = context.require_output(0) + source = context.require_input(0) + largest = combiner(context, result.elem_type, largest=True) + smallest = combiner(context, result.elem_type, largest=False) + bounds = { + tag: (bound, helper) + for tag, index, helper in (("lo", 1, largest), ("hi", 2, smallest)) + if (bound := context.optional_input(index)) is not None + } + if not bounds: + return copy_tensor(source, result) + if len(bounds) == 2: + expression = f"{smallest.name}(x2, {largest.name}(x1, x0))" + else: + ((_, helper),) = bounds.values() + expression = f"{helper.name}(x0, x1)" + return elementwise( + context, + expression=expression, + operands=(source, *(bound for bound, _ in bounds.values())), + result=result, + helpers=tuple(helper for _, helper in bounds.values()), + variant=f"_{''.join(bounds)}", + ) + + +def _dropout(context: NodeContext) -> NodeEmission: + """Dropout in inference mode: the data unchanged, and a mask of ones where asked for. + + Training mode samples a mask, which no compiled artifact can reproduce, so it is a + compile error unless the graph proves the mode off. + """ + training_mode = context.optional_input(2) + mode = context.constant_input(2) + if training_mode is not None and (mode is None or mode.any()): + raise CompileError( + f"Node `{context.label}`: Dropout is supported in inference mode only, but " + f"`{training_mode.name}` is not a compile-time false; the training-mode mask " + "is drawn at random and cannot be compiled." + ) + result = context.require_output(0) + emission = copy_tensor(context.require_input(0), result) + mask = context.outputs[1] if len(context.outputs) > 1 else None + if mask is None: + return emission + ones = elementwise( + context, + expression=expand("$one", mask.elem_type), + operands=(), + result=mask, + variant="_mask", + ) + return NodeEmission( + functions=emission.functions + ones.functions, + statements=emission.statements + ones.statements, + ) + + +def _identity(context: NodeContext) -> NodeEmission: + return copy_tensor(context.require_input(0), context.require_output(0)) + + +# The float32 extremes ONNX gives as Clip-6's default bounds. +_FLOAT_MAX = 3.4028234663852886e38 +_FLOAT_MIN = -_FLOAT_MAX + + +for _op_type, _operator in _ARITHMETIC_OPERATORS.items(): + register_kernel( + "", _op_type, _ARITHMETIC_VERSIONS, partial(_arithmetic, operator=_operator) + ) +for _op_type, (_versions, _template) in _UNARY_MATH.items(): + register_kernel("", _op_type, _versions, partial(pointwise, template=_template)) +register_kernel("", "Min", _EXTREMUM_VERSIONS, partial(_extremum, largest=False)) +register_kernel("", "Max", _EXTREMUM_VERSIONS, partial(_extremum, largest=True)) +register_kernel("", "Sum", _ACCUMULATE_VERSIONS, partial(_accumulate, average=False)) +register_kernel("", "Mean", _ACCUMULATE_VERSIONS, partial(_accumulate, average=True)) +register_kernel("", "Mod", _MOD_VERSIONS, _mod) +register_kernel("", "Pow", _POW_VERSIONS, _pow) +register_kernel("", "Abs", _ABS_VERSIONS, _abs) +register_kernel("", "Neg", _NEG_VERSIONS, partial(pointwise, template="-x0")) +register_kernel("", "Sign", _SIGN_VERSIONS, _sign) +register_kernel("", "Clip", _CLIP_ATTRIBUTE_VERSIONS, _clip_from_attributes) +register_kernel("", "Clip", _CLIP_INPUT_VERSIONS, _clip_from_inputs) +register_kernel("", "Dropout", _DROPOUT_VERSIONS, _dropout) +register_kernel("", "Identity", _IDENTITY_VERSIONS, _identity) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/gather.py b/src/python/fnnx/extras/compilers/c/onnx/ops/gather.py new file mode 100644 index 0000000..b075356 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/gather.py @@ -0,0 +1,525 @@ +"""The ops that read a tensor at positions decided at run time. + +Gather, GatherElements and GatherND all address `data` through an index operand, and TopK +through a ranking it computes; either way the positions are values rather than shapes, so +the addressing is a loop instead of the compile-time strides the views are emitted from. +What stays static is the result's shape, which follows from the operands' shapes alone. + +An index operand comes from the caller, so every one of them is normalized the way ONNX +defines it — a negative index counted back from the end of the axis — and then bounds +checked: a kernel returns nonzero for an index outside its axis and the entrypoint passes +that on as an argument error, rather than reading past a buffer. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from functools import partial +from string import Template + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import FLOAT_TYPES, c_type +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + checked_call, + kernel_name, + normalize_axis, + row_major_strides, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import expand, extents + +# Gather-11 defined the negative index Gather-1 left undefined, and 13 added bfloat16; +# GatherElements arrived at 11 and gained bfloat16 at 13. Accepting a negative index at +# Gather-1 serves more than that revision defines, never something else in its place. +_GATHER_VERSIONS = (1, 11, 13) +_GATHER_ELEMENTS_VERSIONS = (11, 13) + +# GatherND-12 added `batch_dims`, which the generator reads as 0 where the schema has none; +# 13 added bfloat16. +_GATHER_ND_VERSIONS = (11, 12, 13) + +# TopK moved `k` from an attribute to an operand at 10 and gained `largest` and `sorted` at +# 11; 24 added bfloat16. The attribute form is the same selection read from another place, +# so one generator serves it, told where to look. +_TOP_K_ATTRIBUTE_VERSIONS = (1,) +_TOP_K_OPERAND_VERSIONS = (10, 11, 24) + +_GATHER_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const $index* indices, + size_t outer, + size_t extent, + size_t inner, + size_t index_count) +{ + size_t before, chosen; + for (before = 0; before < outer; ++before) { + for (chosen = 0; chosen < index_count; ++chosen) { + ptrdiff_t position = (ptrdiff_t)indices[chosen]; + if (position < 0) { + position += (ptrdiff_t)extent; + } + if (position < 0 || position >= (ptrdiff_t)extent) { + return 1; + } + memcpy( + out + (before * index_count + chosen) * inner, + in + (before * extent + (size_t)position) * inner, + inner * sizeof(*out)); + } + } + return 0; +}""") + +_GATHER_ELEMENTS_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const $index* indices, + size_t count, + int rank, + const size_t* shape, + const size_t* strides, + int axis, + size_t extent) +{ + size_t index; + for (index = 0; index < count; ++index) { + size_t remainder = index; + size_t source = 0; + int walked; + ptrdiff_t position = (ptrdiff_t)indices[index]; + for (walked = rank - 1; walked >= 0; --walked) { + const size_t coordinate = remainder % shape[walked]; + remainder /= shape[walked]; + if (walked != axis) { + source += coordinate * strides[walked]; + } + } + if (position < 0) { + position += (ptrdiff_t)extent; + } + if (position < 0 || position >= (ptrdiff_t)extent) { + return 1; + } + out[index] = in[source + (size_t)position * strides[axis]]; + } + return 0; +}""") + +_GATHER_ND_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const $index* indices, + size_t batch, + size_t outer, + size_t depth, + size_t slice_size, + size_t batch_stride, + const size_t* extents, + const size_t* strides) +{ + size_t batch_index, row, level; + for (batch_index = 0; batch_index < batch; ++batch_index) { + for (row = 0; row < outer; ++row) { + const size_t block = batch_index * outer + row; + size_t source = batch_index * batch_stride; + for (level = 0; level < depth; ++level) { + ptrdiff_t position = (ptrdiff_t)indices[block * depth + level]; + if (position < 0) { + position += (ptrdiff_t)extents[level]; + } + if (position < 0 || position >= (ptrdiff_t)extents[level]) { + return 1; + } + source += (size_t)position * strides[level]; + } + memcpy(out + block * slice_size, in + source, slice_size * sizeof(*out)); + } + } + return 0; +}""") + +# The selection is a partial sort: each pass takes the element that comes first among those +# the passes before it left, which is one scan of the group per result. `$precedes` is a +# strict total order on (value, position) pairs, so "what the passes before left" is +# everything the last selection precedes — no marker array, and no allocation. +_TOP_K_TEMPLATE = Template("""\ +static void $name( + $element* values, + int64_t* indices, + const $element* in, + size_t outer, + size_t extent, + size_t inner, + size_t wanted) +{ + size_t before, after, rank, position; + for (before = 0; before < outer; ++before) { + for (after = 0; after < inner; ++after) { + $element taken = $zero; + ptrdiff_t taken_at = -1; + for (rank = 0; rank < wanted; ++rank) { + $element best = $zero; + ptrdiff_t best_at = -1; + for (position = 0; position < extent; ++position) { + const $element candidate = + in[(before * extent + position) * inner + after]; + if (taken_at >= 0 && + !$precedes(taken, (size_t)taken_at, candidate, position)) { + continue; + } + if (best_at < 0 || + $precedes(candidate, position, best, (size_t)best_at)) { + best = candidate; + best_at = (ptrdiff_t)position; + } + } + values[(before * wanted + rank) * inner + after] = best; + indices[(before * wanted + rank) * inner + after] = (int64_t)best_at; + taken = best; + taken_at = best_at; + } + } + } +}""") + +# Ties are broken by position, which is what makes the order strict and total; the value +# comparison itself is numpy's, since that is what the reference evaluator sorts with — a +# NaN counts as larger than every number, at either end of the ranking. +_PRECEDES_TEMPLATE = Template("""\ +static int $name($element left, size_t left_at, $element right, size_t right_at) +{ + if ($better) { + return 1; + } + if ($worse) { + return 0; + } + return left_at < right_at; +}""") + + +def _gather(context: NodeContext) -> NodeEmission: + """Gather: the operand sliced along one axis at each index the operand names.""" + data = context.require_input(0) + indices = context.require_input(1) + result = context.require_output(0) + rank = len(data.shape) + if rank == 0: + raise CompileError( + f"Node `{context.label}`: `Gather` reads along an axis of `{data.name}`, " + "which is a scalar and has none." + ) + axis = normalize_axis(context, context.int_attribute("axis"), rank) + verify_shape( + context, + result, + (*data.shape[:axis], *indices.shape, *data.shape[axis + 1 :]), + ) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + outer, extent, inner = _split_at(data.shape, axis) + name = _indexed_name(context, data, indices) + return NodeEmission( + functions=( + CFunction( + name, + _GATHER_TEMPLATE.substitute( + name=name, + element=c_type(data.elem_type), + index=c_type(indices.elem_type), + ), + ), + ), + statements=( + checked_call( + context, + name, + [ + result.expr, + data.expr, + indices.expr, + f"{outer}u", + f"{extent}u", + f"{inner}u", + f"{indices.elem_count}u", + ], + ), + ), + ) + + +def _gather_elements(context: NodeContext) -> NodeEmission: + """GatherElements: one element of the operand per index, at the index's own coordinates.""" + data = context.require_input(0) + indices = context.require_input(1) + result = context.require_output(0) + rank = len(data.shape) + if rank == 0 or len(indices.shape) != rank: + raise CompileError( + f"Node `{context.label}`: `GatherElements` reads `{data.name}` of rank {rank} " + f"through `{indices.name}` of rank {len(indices.shape)}; ONNX defines the two " + "as having the same rank, and at least one axis." + ) + axis = normalize_axis(context, context.int_attribute("axis"), rank) + for other in range(rank): + if other != axis and indices.shape[other] != data.shape[other]: + raise CompileError( + f"Node `{context.label}`: `GatherElements` gathers along axis {axis}, so " + f"`{indices.name}` of shape {list(indices.shape)} and `{data.name}` of " + f"shape {list(data.shape)} have to agree on every other axis." + ) + verify_shape(context, result, indices.shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + name = _indexed_name(context, data, indices) + return NodeEmission( + functions=( + CFunction( + name, + _GATHER_ELEMENTS_TEMPLATE.substitute( + name=name, + element=c_type(data.elem_type), + index=c_type(indices.elem_type), + ), + ), + ), + statements=( + checked_call( + context, + name, + [ + result.expr, + data.expr, + indices.expr, + f"{result.elem_count}u", + str(rank), + extents(indices.shape), + extents(row_major_strides(data.shape)), + str(axis), + f"{data.shape[axis]}u", + ], + ), + ), + ) + + +def _gather_nd(context: NodeContext) -> NodeEmission: + """GatherND: a slice of the operand per index tuple, the leading axes shared as batches.""" + data = context.require_input(0) + indices = context.require_input(1) + result = context.require_output(0) + # GatherND-11 has no `batch_dims`, so the default is read here rather than off the + # schema, which carries no entry to read it from at that revision. + batch_dims = int(context.attribute("batch_dims", 0)) + rank = len(data.shape) + if not indices.shape: + raise CompileError( + f"Node `{context.label}`: `GatherND` takes its index tuples from the last axis " + f"of `{indices.name}`, which is a scalar and has none." + ) + depth = indices.shape[-1] + if not 0 <= batch_dims < min(rank, len(indices.shape)): + raise CompileError( + f"Node `{context.label}`: `GatherND` shares {batch_dims} batch dimension(s) " + f"between `{data.name}` of rank {rank} and `{indices.name}` of rank " + f"{len(indices.shape)}; ONNX defines it as fewer than either rank." + ) + if data.shape[:batch_dims] != indices.shape[:batch_dims]: + raise CompileError( + f"Node `{context.label}`: `GatherND` shares {batch_dims} batch dimension(s), " + f"but `{data.name}` of shape {list(data.shape)} and `{indices.name}` of shape " + f"{list(indices.shape)} disagree on them." + ) + if depth > rank - batch_dims: + raise CompileError( + f"Node `{context.label}`: `GatherND` indexes {depth} dimension(s) of " + f"`{data.name}`, which has {rank - batch_dims} left after its batch dimensions." + ) + verify_shape( + context, + result, + (*indices.shape[:-1], *data.shape[batch_dims + depth :]), + ) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + batch = math.prod(data.shape[:batch_dims]) + strides = row_major_strides(data.shape) + slice_size = math.prod(data.shape[batch_dims + depth :]) + name = _indexed_name(context, data, indices) + return NodeEmission( + functions=( + CFunction( + name, + _GATHER_ND_TEMPLATE.substitute( + name=name, + element=c_type(data.elem_type), + index=c_type(indices.elem_type), + ), + ), + ), + statements=( + checked_call( + context, + name, + [ + result.expr, + data.expr, + indices.expr, + f"{batch}u", + f"{math.prod(indices.shape[batch_dims:-1])}u", + f"{depth}u", + f"{slice_size}u", + f"{math.prod(data.shape[batch_dims:])}u", + extents(data.shape[batch_dims : batch_dims + depth]), + extents(strides[batch_dims : batch_dims + depth]), + ], + ), + ), + ) + + +def _top_k(context: NodeContext, *, from_attribute: bool) -> NodeEmission: + """TopK: the `k` first elements along one axis under the ranking `largest` selects. + + `k` has to be fixed at compile time — it is the extent of the result's axis, and a value + the graph only computes at run time would make that shape depend on input data. The + ranking itself is the reference evaluator's: values first, and among equal values the + smaller position, so the selection is defined even where a group holds duplicates. + """ + source = context.require_input(0) + values = context.require_output(0) + indices = context.require_output(1) + rank = len(source.shape) + if rank == 0: + raise CompileError( + f"Node `{context.label}`: `TopK` ranks along an axis of `{source.name}`, which " + "is a scalar and has none." + ) + axis = normalize_axis(context, context.int_attribute("axis"), rank) + wanted = _wanted(context, from_attribute=from_attribute) + if not 0 <= wanted <= source.shape[axis]: + raise CompileError( + f"Node `{context.label}`: `TopK` asks for {wanted} element(s) along axis " + f"{axis} of `{source.name}`, which holds {source.shape[axis]}." + ) + expected = (*source.shape[:axis], wanted, *source.shape[axis + 1 :]) + verify_shape(context, values, expected) + verify_shape(context, indices, expected) + if values.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + # `largest` and `sorted` arrived at 11; the revisions before it always take the largest + # in sorted order, which is what the two attributes default to. Sorted output answers + # `sorted=0` as well, which asks for the same elements in no particular order. + largest = bool(int(context.attribute("largest", 1))) + precedes = _precedes(context, source.elem_type, largest=largest) + outer, extent, inner = _split_at(source.shape, axis) + name = kernel_name( + context, "largest" if largest else "smallest", c_type(source.elem_type) + ) + return NodeEmission( + functions=( + precedes, + CFunction( + name, + _TOP_K_TEMPLATE.substitute( + name=name, + element=c_type(source.elem_type), + zero=expand("$zero", source.elem_type), + precedes=precedes.name, + ), + ), + ), + statements=( + call_kernel( + name, + [ + values.expr, + indices.expr, + source.expr, + f"{outer}u", + f"{extent}u", + f"{inner}u", + f"{wanted}u", + ], + ), + ), + ) + + +def _wanted(context: NodeContext, *, from_attribute: bool) -> int: + """How many elements TopK selects, from wherever this revision states it.""" + if from_attribute: + return context.int_attribute("k") + operand = context.require_input(1) + fixed = context.constant_input(1) + if fixed is None or fixed.size != 1: + raise CompileError( + f"Node `{context.label}`: `TopK` takes `k` from `{operand.name}`, which holds " + f"{'no single value' if fixed is not None else 'no value'} known at compile " + "time; the shape of the result then depends on input data, which the C " + "compiler cannot compile." + ) + return int(fixed.reshape(-1)[0]) + + +def _precedes(context: NodeContext, elem_type: int, *, largest: bool) -> CFunction: + """The order TopK ranks by, as a function of two values and their positions.""" + comparison = ">" if largest else "<" + better = f"left {comparison} right" + worse = f"right {comparison} left" + if elem_type in FLOAT_TYPES: + # numpy sorts a NaN above every number, which puts it first when the largest come + # first and last when the smallest do. + outranks = "isnan(left) && !isnan(right)" + outranked = "isnan(right) && !isnan(left)" + if not largest: + outranks, outranked = outranked, outranks + better = f"{better} || ({outranks})" + worse = f"{worse} || ({outranked})" + name = kernel_name( + context, "before", "largest" if largest else "smallest", c_type(elem_type) + ) + return CFunction( + name, + _PRECEDES_TEMPLATE.substitute( + name=name, element=c_type(elem_type), better=better, worse=worse + ), + ) + + +def _split_at(shape: Sequence[int], axis: int) -> tuple[int, int, int]: + """`shape` as the three factors an axis-wise kernel walks: before it, it, and after.""" + return math.prod(shape[:axis]), shape[axis], math.prod(shape[axis + 1 :]) + + +def _indexed_name(context: NodeContext, data: TensorRef, indices: TensorRef) -> str: + return kernel_name(context, c_type(data.elem_type), c_type(indices.elem_type)) + + +register_kernel("", "Gather", _GATHER_VERSIONS, _gather) +register_kernel("", "GatherElements", _GATHER_ELEMENTS_VERSIONS, _gather_elements) +register_kernel("", "GatherND", _GATHER_ND_VERSIONS, _gather_nd) +register_kernel( + "", "TopK", _TOP_K_ATTRIBUTE_VERSIONS, partial(_top_k, from_attribute=True) +) +register_kernel( + "", "TopK", _TOP_K_OPERAND_VERSIONS, partial(_top_k, from_attribute=False) +) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/gemm.py b/src/python/fnnx/extras/compilers/c/onnx/ops/gemm.py new file mode 100644 index 0000000..767ccec --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/gemm.py @@ -0,0 +1,462 @@ +"""The dense matrix kernels: Gemm, MatMul and Det. + +All three walk a row-major buffer as a stack of matrices, and all three sum a product over an +inner axis in an order the spec leaves open. What separates them is how the matrices are +addressed: Gemm's two transposes and its broadcast bias become strides, MatMul's batch axes +broadcast against each other, and Det walks one matrix at a time down a copy of it. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from string import Template + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import FLOAT_TYPES, c_type +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + ScratchBuffer, + TensorRef, + broadcast_strides, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + kernel_name, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents, math_suffix + +# Both transposes and C's broadcast become strides, so one kernel per element type covers +# every attribute combination. +_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* left, + const $element* right, + const $element* bias, + size_t rows, + size_t columns, + size_t inner, + size_t left_row_stride, + size_t left_inner_stride, + size_t right_inner_stride, + size_t right_column_stride, + size_t bias_row_stride, + size_t bias_column_stride, + $scalar alpha, + $scalar beta) +{ + size_t row, column, index; + for (row = 0; row < rows; ++row) { + for (column = 0; column < columns; ++column) { + $element sum = $zero; + for (index = 0; index < inner; ++index) { + sum += left[row * left_row_stride + index * left_inner_stride] + * right[index * right_inner_stride + column * right_column_stride]; + } +$store + } + } +}""") + +# The floating-point families scale in the element type, as the reference does when it +# multiplies a float array by alpha. +_FLOAT_STORE = Template("""\ + sum *= alpha; + if (bias != NULL) { + sum += beta * bias[row * bias_row_stride + column * bias_column_stride]; + } + out[row * columns + column] = sum;""") + +# The integer families scale in double and truncate on the way back, as the reference does +# when the float alpha promotes an integer dot product to float64. +_INTEGER_STORE = Template("""\ + { + double scaled = (double)sum * alpha; + if (bias != NULL) { + scaled += beta + * (double)bias[row * bias_row_stride + + column * bias_column_stride]; + } + out[row * columns + column] = ($element)scaled; + }""") + +# Numpy-style broadcasting of C arrived at opset 7; 9 and 13 widened the types, and 11 made +# C optional, which the kernel already handles. +_GEMM_VERSIONS = (7, 9, 11, 13) + + +def _gemm(context: NodeContext) -> NodeEmission: + left = context.require_input(0) + right = context.require_input(1) + bias = context.optional_input(2) + result = context.require_output(0) + alpha = float(context.attribute("alpha", 1.0)) + beta = float(context.attribute("beta", 1.0)) + transpose_left = bool(context.attribute("transA", 0)) + transpose_right = bool(context.attribute("transB", 0)) + + rows, inner = _oriented_shape(context, left, transpose_left) + right_inner, columns = _oriented_shape(context, right, transpose_right) + if inner != right_inner: + raise CompileError( + f"Node `{context.label}`: Gemm operands `{left.name}` and `{right.name}` do " + f"not share an inner dimension ({inner} against {right_inner})." + ) + left_strides = _oriented_strides(left.shape, transpose_left) + right_strides = _oriented_strides(right.shape, transpose_right) + bias_strides: tuple[int, ...] + if bias is None or beta == 0.0: + # The reference drops C whenever beta is zero, so `0 * inf` never reaches the sum. + bias_expr, bias_strides = "NULL", (0, 0) + else: + bias_expr = bias.expr + bias_strides = broadcast_strides( + bias, (rows, columns), node_label=context.label + ) + + element = c_type(result.elem_type) + scales_in_element_type = result.elem_type in FLOAT_TYPES + scalar_type = result.elem_type if scales_in_element_type else TensorProto.DOUBLE + name = f"{context.prefix}_gemm_{element}" + store = _FLOAT_STORE if scales_in_element_type else _INTEGER_STORE + definition = _TEMPLATE.substitute( + name=name, + element=element, + scalar=c_type(scalar_type), + zero=scalar_literal(0, result.elem_type), + store=store.substitute(element=element), + ) + call = "\n".join( + [ + f"{name}(", + f" {result.expr}, {left.expr}, {right.expr}, {bias_expr},", + f" {rows}u, {columns}u, {inner}u,", + f" {left_strides[0]}u, {left_strides[1]}u, " + f"{right_strides[0]}u, {right_strides[1]}u,", + f" {bias_strides[0]}u, {bias_strides[1]}u,", + f" {scalar_literal(alpha, scalar_type)}, " + f"{scalar_literal(beta, scalar_type)});", + ] + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _oriented_shape( + context: NodeContext, operand: TensorRef, transposed: bool +) -> tuple[int, int]: + """The operand's shape as the multiplication sees it, after its transpose attribute.""" + if len(operand.shape) != 2: + raise CompileError( + f"Node `{context.label}`: Gemm takes 2-D operands, but `{operand.name}` has " + f"shape {list(operand.shape)}." + ) + rows, columns = operand.shape + return (columns, rows) if transposed else (rows, columns) + + +def _oriented_strides(shape: tuple[int, ...], transposed: bool) -> tuple[int, int]: + """Strides along the two axes of the oriented operand, into its row-major buffer.""" + return (1, shape[1]) if transposed else (shape[1], 1) + + +register_kernel("", "Gemm", _GEMM_VERSIONS, _gemm) + + +# MatMul is numpy's `matmul`: everything before the last two axes is a batch the two operands +# broadcast against each other, so the batch coordinate becomes an offset into each of them +# and the matrix product itself is the same three loops Gemm runs. +# The geometry of a batched product, in the order `MatrixProduct.arguments` fills it; the +# quantized products take the same block after their own operands. +PRODUCT_PARAMETERS = """\ + size_t batch_count, + int batch_rank, + const size_t* batch_shape, + const size_t* left_batch_strides, + const size_t* right_batch_strides, + size_t rows, + size_t columns, + size_t inner""" + +_MATMUL_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* left, + const $element* right, +$parameters) +{ + size_t batch, row, column, index; + for (batch = 0; batch < batch_count; ++batch) { + size_t left_base = 0; + size_t right_base = 0; + size_t remainder = batch; + int axis; + for (axis = batch_rank - 1; axis >= 0; --axis) { + const size_t coordinate = remainder % batch_shape[axis]; + remainder /= batch_shape[axis]; + left_base += coordinate * left_batch_strides[axis]; + right_base += coordinate * right_batch_strides[axis]; + } + for (row = 0; row < rows; ++row) { + for (column = 0; column < columns; ++column) { + $element sum = $zero; + for (index = 0; index < inner; ++index) { + sum += left[left_base + row * inner + index] + * right[right_base + index * columns + column]; + } + out[(batch * rows + row) * columns + column] = sum; + } + } + } +}""") + +# MatMul's semantics have not changed since opset 1: 9 added the integer families and 13 +# bfloat16. Only 13 is claimed all the same, because it is the only revision anything can +# vouch for — the reference evaluator is version-faithful there and the corpus's own MatMul +# tests import it, while nothing checks 1 or 9. A model importing one of those gets the +# unsupported-version error rather than a kernel no oracle has ever seen. +_MATMUL_VERSIONS = (13,) + + +@dataclass(frozen=True) +class MatrixProduct: + """Where every matrix of a batched product sits, and the shape the product comes to. + + The quantized products walk the same operands as `MatMul`, so the addressing is resolved + once here and each kernel differs only in what it accumulates and how it stores it. + """ + + rows: int + columns: int + inner: int + batch_shape: tuple[int, ...] + left_batch_strides: tuple[int, ...] + right_batch_strides: tuple[int, ...] + result_shape: tuple[int, ...] + + @property + def arguments(self) -> list[str]: + """Call-site literals for the geometry parameters a batched product's kernel takes.""" + return [ + f"{math.prod(self.batch_shape)}u", + str(len(self.batch_shape)), + extents(self.batch_shape), + extents(self.left_batch_strides), + extents(self.right_batch_strides), + f"{self.rows}u", + f"{self.columns}u", + f"{self.inner}u", + ] + + +def matrix_product( + context: NodeContext, left: TensorRef, right: TensorRef +) -> MatrixProduct: + (rows, inner), left_batch = _matrix_stack(context, left, column_vector=False) + (right_inner, columns), right_batch = _matrix_stack( + context, right, column_vector=True + ) + if inner != right_inner: + raise CompileError( + f"Node `{context.label}`: {context.node.op_type} operands `{left.name}` and " + f"`{right.name}` do not share an inner dimension ({inner} against " + f"{right_inner})." + ) + batch_shape = _broadcast_shape(context, left_batch, right_batch) + return MatrixProduct( + rows=rows, + columns=columns, + inner=inner, + batch_shape=batch_shape, + left_batch_strides=_batch_strides( + context, left, left_batch, batch_shape, rows * inner + ), + right_batch_strides=_batch_strides( + context, right, right_batch, batch_shape, inner * columns + ), + # A promoted rank-1 operand contributes an axis of extent 1 that ONNX drops from + # the result again, which leaves the row-major layout — and so the addressing — + # unchanged. + result_shape=( + *batch_shape, + *((rows,) if len(left.shape) > 1 else ()), + *((columns,) if len(right.shape) > 1 else ()), + ), + ) + + +def _matmul(context: NodeContext) -> NodeEmission: + left = context.require_input(0) + right = context.require_input(1) + result = context.require_output(0) + product = matrix_product(context, left, right) + verify_shape(context, result, product.result_shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + name = kernel_name(context, element) + definition = _MATMUL_TEMPLATE.substitute( + name=name, + element=element, + zero=scalar_literal(0, result.elem_type), + parameters=PRODUCT_PARAMETERS, + ) + call = call_kernel(name, [result.expr, left.expr, right.expr, *product.arguments]) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _matrix_stack( + context: NodeContext, operand: TensorRef, *, column_vector: bool +) -> tuple[tuple[int, int], tuple[int, ...]]: + """The operand's trailing matrix and the batch axes in front of it. + + numpy — and so ONNX — reads a rank-1 operand as the single row or column that makes the + product defined, and drops that axis from the result again. + """ + if not operand.shape: + raise CompileError( + f"Node `{context.label}`: {context.node.op_type} multiplies matrices, but " + f"`{operand.name}` is a scalar." + ) + if len(operand.shape) == 1: + (extent,) = operand.shape + return ((extent, 1) if column_vector else (1, extent)), () + return (operand.shape[-2], operand.shape[-1]), operand.shape[:-2] + + +def _broadcast_shape( + context: NodeContext, left: tuple[int, ...], right: tuple[int, ...] +) -> tuple[int, ...]: + """The shape two batches broadcast onto, numpy-style, aligned at their trailing axes.""" + rank = max(len(left), len(right)) + padded_left = (1,) * (rank - len(left)) + left + padded_right = (1,) * (rank - len(right)) + right + shape = [] + for one, other in zip(padded_left, padded_right): + if one != other and 1 not in (one, other): + raise CompileError( + f"Node `{context.label}`: the batch shapes {list(left)} and {list(right)} " + "of its operands do not broadcast against each other." + ) + # An axis of 1 stretches to whatever the other side is, and that includes 0: a batch + # of no matrices against a batch of one is still a batch of no matrices, so `max` + # would be wrong exactly where a zero-element operand is involved. + shape.append(other if one == 1 else one) + return tuple(shape) + + +def _batch_strides( + context: NodeContext, + operand: TensorRef, + batch: tuple[int, ...], + batch_shape: tuple[int, ...], + matrix_size: int, +) -> tuple[int, ...]: + """Strides addressing the operand's matrices while iterating the broadcast batch.""" + return tuple( + stride * matrix_size + for stride in broadcast_strides( + replace(operand, shape=batch), batch_shape, node_label=context.label + ) + ) + + +register_kernel("", "MatMul", _MATMUL_VERSIONS, _matmul) + + +# The determinant, by the same LU factorization with partial pivoting LAPACK runs, so that +# the pivots multiplied together here are the ones the reference's `numpy.linalg.det` +# multiplies. Elimination is destructive, hence the copy: the operand is read-only, and the +# artifact has nowhere else to put a working matrix. +_DET_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + $element* work, + size_t batch_count, + size_t order) +{ + size_t batch, step, row, column, pivot; + for (batch = 0; batch < batch_count; ++batch) { + $element determinant = $one; + memcpy(work, in + batch * order * order, order * order * sizeof(*work)); + for (step = 0; step < order; ++step) { + pivot = step; + for (row = step + 1; row < order; ++row) { + if ($absolute(work[row * order + step]) + > $absolute(work[pivot * order + step])) { + pivot = row; + } + } + if (pivot != step) { + for (column = step; column < order; ++column) { + const $element swapped = work[step * order + column]; + work[step * order + column] = work[pivot * order + column]; + work[pivot * order + column] = swapped; + } + determinant = -determinant; + } + determinant *= work[step * order + step]; + if (work[step * order + step] == $zero) { + break; + } + for (row = step + 1; row < order; ++row) { + const $element factor = + work[row * order + step] / work[step * order + step]; + for (column = step + 1; column < order; ++column) { + work[row * order + column] -= + factor * work[step * order + column]; + } + } + } + out[batch] = determinant; + } +}""") + +# Det arrived at 11 and 22 added bfloat16. Only 22 is claimed, for the reason MatMul claims +# only 13: it is the revision the reference evaluator is faithful for and the one both corpus +# tests import, and nothing vouches for 11. +_DET_VERSIONS = (22,) + + +def _det(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + result = context.require_output(0) + if len(source.shape) < 2 or source.shape[-1] != source.shape[-2]: + raise CompileError( + f"Node `{context.label}`: Det takes square matrices, but `{source.name}` has " + f"shape {list(source.shape)}." + ) + verify_shape(context, result, source.shape[:-2]) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + order = source.shape[-1] + element = c_type(result.elem_type) + name = kernel_name(context, element) + definition = _DET_TEMPLATE.substitute( + name=name, + element=element, + absolute=f"fabs{math_suffix(result.elem_type)}", + one=scalar_literal(1, result.elem_type), + zero=scalar_literal(0, result.elem_type), + ) + work = ScratchBuffer(f"{name}_work", result.elem_type, order * order) + call = call_kernel( + name, + [result.expr, source.expr, work.symbol, f"{result.elem_count}u", f"{order}u"], + ) + return NodeEmission( + functions=(CFunction(name, definition),), statements=(call,), scratch=(work,) + ) + + +register_kernel("", "Det", _DET_VERSIONS, _det) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/generate.py b/src/python/fnnx/extras/compilers/c/onnx/ops/generate.py new file mode 100644 index 0000000..df2a3d6 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/generate.py @@ -0,0 +1,405 @@ +"""The ops whose result is a function of its own coordinates. + +OneHot, EyeLike and Trilu each decide an element from where it sits: whether the coordinate +along one axis is the index an operand names, whether two coordinates are a fixed distance +apart, whether one is above the other. ReverseSequence is the same idea one step on — the +coordinate along the time axis is mapped to another coordinate, per row, by a length the +caller supplies. None of them needs an index into a buffer, only the loop that produces the +coordinates, so each is one kernel over the result's own shape. + +`ConstantOfShape` and `Range` belong to the same family and need no kernel at all. Both read +their entire result — its shape included — out of their operands, so the compiler accepts +them only where the graph fixes those operands, and where it does the folding pass resolves +them into an initializer before dispatch is ever reached. +""" + +from __future__ import annotations + +import math +from string import Template + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import FLOAT_TYPES, UNSIGNED_TYPES, c_type +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + checked_call, + kernel_name, + normalize_axis, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import expand, math_suffix + +# OneHot-11 only clarified how a negative axis is counted; EyeLike-22 and Trilu's single +# revision widened nothing this compiler compiles differently. ReverseSequence has one +# revision of its own. +_ONE_HOT_VERSIONS = (9, 11) +_EYE_LIKE_VERSIONS = (9, 22) +_TRILU_VERSIONS = (14,) +_REVERSE_SEQUENCE_VERSIONS = (10,) + +_ONE_HOT_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $index* indices, + const $element* values, + size_t outer, + size_t depth, + size_t inner) +{ + size_t before, position, after; + for (before = 0; before < outer; ++before) { + for (position = 0; position < depth; ++position) { + for (after = 0; after < inner; ++after) { + const $index folded = + $fold(indices[before * inner + after], ($index)depth); + out[(before * depth + position) * inner + after] = + (folded == ($index)position) ? values[1] : values[0]; + } + } + } +}""") + +# ONNX folds an index into the range the depth allows, which is what the reference evaluator +# computes with numpy's `mod`: the result takes the sign of the depth, where C's `%` takes +# the sign of the index. An index of a floating-point type folds the same way and then +# matches no position unless it is a whole number, exactly as comparing the two would. +_FOLD_TEMPLATE = Template("""\ +static $element $name($element value, $element depth) +{ + const $element folded = $modulo; + return $adjust; +}""") + +_EYE_LIKE_TEMPLATE = Template("""\ +static void $name($element* out, size_t rows, size_t columns, int64_t offset) +{ + size_t row, column; + for (row = 0; row < rows; ++row) { + for (column = 0; column < columns; ++column) { + out[row * columns + column] = + ((int64_t)column - (int64_t)row == offset) ? $one : $zero; + } + } +}""") + +_TRILU_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + size_t batch, + size_t rows, + size_t columns, + int64_t offset) +{ + size_t matrix, row, column; + for (matrix = 0; matrix < batch; ++matrix) { + for (row = 0; row < rows; ++row) { + for (column = 0; column < columns; ++column) { + const size_t position = (matrix * rows + row) * columns + column; + out[position] = + ((int64_t)column - (int64_t)row $keep offset) ? in[position] : $zero; + } + } + } +}""") + +_REVERSE_SEQUENCE_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const int64_t* lengths, + size_t batch, + size_t time, + size_t inner, + size_t batch_stride, + size_t time_stride) +{ + size_t row, step; + for (row = 0; row < batch; ++row) { + const int64_t length = lengths[row]; + if (length < 1 || (size_t)length > time) { + return 1; + } + for (step = 0; step < time; ++step) { + const size_t source = + (step < (size_t)length) ? (size_t)length - 1 - step : step; + memcpy( + out + row * batch_stride + step * time_stride, + in + row * batch_stride + source * time_stride, + inner * sizeof(*out)); + } + } + return 0; +}""") + + +def _one_hot(context: NodeContext) -> NodeEmission: + """OneHot: the depth axis set at the position each index names, elsewhere the off value. + + `depth` decides the extent of that axis, so it has to be fixed at compile time; the two + values it selects between decide nothing about any shape and are read at run time. + """ + indices = context.require_input(0) + values = context.require_input(2) + result = context.require_output(0) + rank = len(indices.shape) + axis = normalize_axis(context, context.int_attribute("axis"), rank + 1) + depth = _depth(context) + verify_shape(context, result, (*indices.shape[:axis], depth, *indices.shape[axis:])) + if values.elem_count != 2: + raise CompileError( + f"Node `{context.label}`: `OneHot` takes its off and on values from " + f"`{values.name}`, which holds {values.elem_count}; ONNX defines it as two." + ) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + fold = _fold(context, indices.elem_type) + name = kernel_name(context, c_type(result.elem_type), c_type(indices.elem_type)) + return NodeEmission( + functions=( + fold, + CFunction( + name, + _ONE_HOT_TEMPLATE.substitute( + name=name, + element=c_type(result.elem_type), + index=c_type(indices.elem_type), + fold=fold.name, + ), + ), + ), + statements=( + call_kernel( + name, + [ + result.expr, + indices.expr, + values.expr, + f"{math.prod(indices.shape[:axis])}u", + f"{depth}u", + f"{math.prod(indices.shape[axis:])}u", + ], + ), + ), + ) + + +def _depth(context: NodeContext) -> int: + operand = context.require_input(1) + fixed = context.constant_input(1) + if fixed is None or fixed.size != 1: + raise CompileError( + f"Node `{context.label}`: `OneHot` takes its depth from `{operand.name}`, " + f"which holds {'no single value' if fixed is not None else 'no value'} known " + "at compile time; the shape of the result then depends on input data, which " + "the C compiler cannot compile." + ) + return int(fixed.reshape(-1)[0]) + + +def _fold(context: NodeContext, elem_type: int) -> CFunction: + """An index folded into `[0, depth)`, at whatever type the indices are given in.""" + element = c_type(elem_type) + modulo = ( + f"fmod{math_suffix(elem_type)}(value, depth)" + if elem_type in FLOAT_TYPES + else "value % depth" + ) + # An unsigned index is already inside the range, and comparing one against zero is a + # diagnostic the artifact's `-Werror` build turns into a failure. + adjust = ( + "folded" + if elem_type in UNSIGNED_TYPES + else "(folded < 0) ? folded + depth : folded" + ) + name = f"{context.prefix}_onehot_fold_{element}" + return CFunction( + name, + _FOLD_TEMPLATE.substitute( + name=name, element=element, modulo=modulo, adjust=adjust + ), + ) + + +def _eye_like(context: NodeContext) -> NodeEmission: + """EyeLike: ones on one diagonal of a matrix the operand's shape describes. + + Nothing of the operand but its shape is read, which is why the operand itself does not + appear in the emitted call at all. + """ + source = context.require_input(0) + result = context.require_output(0) + if len(source.shape) != 2: + raise CompileError( + f"Node `{context.label}`: `EyeLike` takes the shape of `{source.name}`, which " + f"is {list(source.shape)}; ONNX defines the op over 2-D tensors." + ) + verify_shape(context, result, source.shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + rows, columns = result.shape + name = kernel_name(context, c_type(result.elem_type)) + return NodeEmission( + functions=( + CFunction( + name, + _EYE_LIKE_TEMPLATE.substitute( + name=name, + element=c_type(result.elem_type), + one=expand("$one", result.elem_type), + zero=expand("$zero", result.elem_type), + ), + ), + ), + statements=( + call_kernel( + name, + [ + result.expr, + f"{rows}u", + f"{columns}u", + str(context.int_attribute("k")), + ], + ), + ), + ) + + +def _trilu(context: NodeContext) -> NodeEmission: + """Trilu: one triangle of each matrix kept, the other zeroed. + + The diagonal to cut along is an operand, and one that decides nothing about any shape, + so a graph computing it at run time still compiles: the kernel reads it as a value. + """ + source = context.require_input(0) + result = context.require_output(0) + if len(source.shape) < 2: + raise CompileError( + f"Node `{context.label}`: `Trilu` cuts each matrix of `{source.name}`, which " + f"has shape {list(source.shape)}; ONNX defines the op from rank 2 up." + ) + verify_shape(context, result, source.shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + upper = bool(context.int_attribute("upper")) + rows, columns = source.shape[-2:] + name = kernel_name(context, "upper" if upper else "lower", c_type(result.elem_type)) + return NodeEmission( + functions=( + CFunction( + name, + _TRILU_TEMPLATE.substitute( + name=name, + element=c_type(result.elem_type), + keep=">=" if upper else "<=", + zero=expand("$zero", result.elem_type), + ), + ), + ), + statements=( + call_kernel( + name, + [ + result.expr, + source.expr, + f"{math.prod(source.shape[:-2])}u", + f"{rows}u", + f"{columns}u", + _offset(context), + ], + ), + ), + ) + + +def _offset(context: NodeContext) -> str: + """Which diagonal Trilu cuts along, as a C expression of type `int64_t`.""" + operand = context.optional_input(1) + if operand is None: + return "0" + if operand.elem_count != 1: + raise CompileError( + f"Node `{context.label}`: `Trilu` takes the diagonal from `{operand.name}`, " + f"which holds {operand.elem_count} values; ONNX defines it as a single one." + ) + return f"{operand.expr}[0]" + + +def _reverse_sequence(context: NodeContext) -> NodeEmission: + """ReverseSequence: the first `sequence_lens[b]` steps of each batch reversed. + + The lengths are values, not shapes, so they are read at run time — and validated there: + ONNX defines them only within the time axis, and a longer one would read past the buffer. + """ + source = context.require_input(0) + lengths = context.require_input(1) + result = context.require_output(0) + rank = len(source.shape) + if rank < 2: + raise CompileError( + f"Node `{context.label}`: `ReverseSequence` reverses along the time axis of " + f"`{source.name}`, which has shape {list(source.shape)}; ONNX defines the op " + "from rank 2 up." + ) + batch_axis = context.int_attribute("batch_axis") + time_axis = context.int_attribute("time_axis") + if {batch_axis, time_axis} != {0, 1}: + raise CompileError( + f"Node `{context.label}`: `ReverseSequence` runs over batch axis {batch_axis} " + f"and time axis {time_axis}; ONNX defines the two as 0 and 1 in either order." + ) + verify_shape(context, result, source.shape) + batch = source.shape[batch_axis] + if lengths.shape != (batch,): + raise CompileError( + f"Node `{context.label}`: `ReverseSequence` takes its lengths from " + f"`{lengths.name}` of shape {list(lengths.shape)}; ONNX defines one per batch, " + f"of which axis {batch_axis} of `{source.name}` has {batch}." + ) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + inner = math.prod(source.shape[2:]) + strides = (source.shape[1] * inner, inner) + name = kernel_name(context, c_type(result.elem_type)) + return NodeEmission( + functions=( + CFunction( + name, + _REVERSE_SEQUENCE_TEMPLATE.substitute( + name=name, element=c_type(result.elem_type) + ), + ), + ), + statements=( + checked_call( + context, + name, + [ + result.expr, + source.expr, + lengths.expr, + f"{batch}u", + f"{source.shape[time_axis]}u", + f"{inner}u", + f"{strides[batch_axis]}u", + f"{strides[time_axis]}u", + ], + ), + ), + ) + + +register_kernel("", "OneHot", _ONE_HOT_VERSIONS, _one_hot) +register_kernel("", "EyeLike", _EYE_LIKE_VERSIONS, _eye_like) +register_kernel("", "Trilu", _TRILU_VERSIONS, _trilu) +register_kernel("", "ReverseSequence", _REVERSE_SEQUENCE_VERSIONS, _reverse_sequence) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/linear_attention.py b/src/python/fnnx/extras/compilers/c/onnx/ops/linear_attention.py new file mode 100644 index 0000000..5e35e08 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/linear_attention.py @@ -0,0 +1,482 @@ +"""LinearAttention: the recurrence ONNX drives with a `Scan`, walked one token at a time. + +The op carries one state matrix of shape `(d_k, d_v)` per batch item and per key/value head, +and walks the sequence forward through it. A step optionally decays the state by `exp(g_t)` +along the key dimension, optionally corrects the token's value by what the state already +answers for its key — `v_t <- beta_t * (v_t - S^T k_t)`, the delta rule — writes the outer +product `k_t (x) v_t` into the state, and then reads the *updated* state back through the +query: `o_t = scale * q_t^T S`. Nothing crosses a batch item or a key/value head, so those +two are the outer loops and everything inside them is sequential. + +**Why a kernel and not the function body.** The rest of this op family compiles through the +body ONNX itself defines; this one cannot. `LinearAttention`'s body drives the recurrence +with a `Scan` over the time axis, and a `Scan` whose trip count is a run-time tensor is not +something constant folding can resolve away — it lands on the v1 unsupported surface +(`verify.py`'s `CONTROL_FLOW_OPS`), which is what the corpus's own `..._expanded` models, +the pre-inlined bodies, are refused for. SPEC's "native kernels only where expansion proves +insufficient" is exactly this case. + +**What varies between call sites is geometry, not code.** Every extent and stride is a +kernel argument, so one shared `static` function per element type serves every node. +`update_rule` reaches the kernel as nothing at all: it decides only which of `decay` and +`beta` a node passes, and the reference refuses every other combination outright, so the two +operands being NULL or not *is* the rule. The two packings ONNX allows each of them — a +per-head scalar or a per-key-dimension vector for the decay, a per-head or a whole-batch +scalar for beta — are likewise two strides rather than two kernels, and grouped-query +attention is one more loop bound: `np.repeat` along the head axis makes the query heads one +key/value head serves consecutive. + +**Where the recurrence accumulates.** The reference converts every operand to float32 and +runs the whole recurrence there whatever the tensors hold, casting the result back to the +query's type and `present_state` to `past_state`'s. The op's type constraints admit float16, +bfloat16 and float, of which this compiler supports only the last, so the accumulator and +the tensors' own type coincide — which is what lets the running state live directly in the +`present_state` output buffer, updated in place across the whole sequence, with no scratch +and no rounding between steps. The kernel still names the accumulator type in its own right, +and the C compiler rejects the call site if the two ever stop coinciding. + +`chunk_size` is read here by nothing: ONNX documents it as a tuning hint for a +chunk-parallel implementation, and its own reference implementation ignores it outright. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from string import Template + +import onnx.defs +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + c_type, + element_type_name, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import call_kernel, kernel_name, verify_shape +from fnnx.extras.compilers.c.onnx.ops.broadcast import math_suffix + +# LinearAttention arrived at opset 27 and has had one revision. +_VERSIONS = (27,) + +# The reference evaluator converts every operand to float32 and runs the recurrence there +# regardless of what the tensors hold; float is the one of the op's three element types this +# compiler supports, so the two coincide, and the kernel says which is which. +_ACCUMULATOR = TensorProto.FLOAT + +_QUERY, _KEY, _VALUE, _PAST_STATE, _DECAY, _BETA = range(6) +_OUTPUT, _PRESENT_STATE = range(2) + +# Which of the two optional gates each rule reads. The reference raises both ways — an +# operand a rule does not read is as much an error as one it needs and does not get — so +# these two flags are the whole of what `update_rule` decides. +_UPDATE_RULES: dict[str, tuple[bool, bool]] = { + "linear": (False, False), + "gated": (True, False), + "delta": (False, True), + "gated_delta": (True, True), +} + +# The recurrence, per batch item and key/value head. Three properties shape the loops: +# the state's columns are independent — the delta correction of column `m` reads and then +# rewrites only column `m` — so no copy of the pre-write state is needed anywhere; the read +# happens after the write, which is what makes `o_t` a function of `S_t` rather than +# `S_{t-1}`; and the state matrix is `present_state` itself, so the sequence ends with the +# answer already in place. +_KERNEL_TEMPLATE = Template("""\ +static void $name( + $element* out, + $accumulate* state, + const $element* query, + const $element* key, + const $element* value, + const $element* past, + const $element* decay, + const $element* beta, + size_t batch, + size_t steps, + size_t q_heads, + size_t kv_heads, + size_t group, + size_t d_k, + size_t d_v, + size_t decay_row, + size_t decay_head_stride, + size_t decay_dim_stride, + size_t beta_row, + size_t beta_head_stride, + double scale) +{ + const size_t query_row = q_heads * d_k; + const size_t key_row = kv_heads * d_k; + const size_t value_row = kv_heads * d_v; + const size_t out_row = q_heads * d_v; + const size_t cells = d_k * d_v; + size_t item, head, step, share, row, column; + for (item = 0; item < batch; ++item) { + for (head = 0; head < kv_heads; ++head) { + const size_t origin = (item * kv_heads + head) * cells; + $accumulate* carried = state + origin; + for (row = 0; row < cells; ++row) { + carried[row] = past == NULL ? $zero : ($accumulate)past[origin + row]; + } + for (step = 0; step < steps; ++step) { + const size_t token = item * steps + step; + const $element* k_t = key + token * key_row + head * d_k; + const $element* v_t = value + token * value_row + head * d_v; + if (decay != NULL) { + /* The gate is in log space and broadcasts along the value dimension; + a per-head scalar reaches every key dimension through a zero + stride. */ + const $element* g_t = + decay + token * decay_row + head * decay_head_stride; + for (row = 0; row < d_k; ++row) { + $accumulate* line = carried + row * d_v; + const $accumulate factor = + exp$f(($accumulate)g_t[row * decay_dim_stride]); + for (column = 0; column < d_v; ++column) { + line[column] *= factor; + } + } + } + for (column = 0; column < d_v; ++column) { + $accumulate written = ($accumulate)v_t[column]; + if (beta != NULL) { + /* The delta rule writes only what this key is not already answered + with, at the rate beta names; the state it reads is the decayed + one, which is why the gate runs first. */ + $accumulate retrieved = $zero; + for (row = 0; row < d_k; ++row) { + retrieved += + carried[row * d_v + column] * ($accumulate)k_t[row]; + } + written = ($accumulate)beta[token * beta_row + + head * beta_head_stride] * (written - retrieved); + } + for (row = 0; row < d_k; ++row) { + carried[row * d_v + column] += ($accumulate)k_t[row] * written; + } + } + for (share = 0; share < group; ++share) { + /* Grouped-query attention repeats the state along the head axis, so + the query heads this one serves are `group` consecutive ones. */ + const size_t reader = head * group + share; + const $element* q_t = query + token * query_row + reader * d_k; + $element* answer = out + token * out_row + reader * d_v; + for (column = 0; column < d_v; ++column) { + $accumulate total = $zero; + for (row = 0; row < d_k; ++row) { + total += ($accumulate)q_t[row] * carried[row * d_v + column]; + } + /* The derived scale is a numpy float64 in the reference, so the + product widens and rounds once on the way into the result; a + scale the node states outright is a Python float, which numpy + weakens to the sum's own type. The two differ by that rounding + alone, and this takes the wider of them. */ + answer[column] = ($element)(scale * (double)total); + } + } + } + } + } +}""") + + +@dataclass(frozen=True) +class _Packing: + """Where one of the two optional gates keeps the value a `(token, head, key dim)` reads. + + ONNX packs both into `(B, T, L)` and lets `L` say which granularity they carry, so the + difference between them is three strides rather than two kernels: `row` is one token's + worth, and a granularity coarser than the axis it feeds addresses it with a zero stride. + """ + + row: int + head_stride: int + dim_stride: int + + @staticmethod + def absent_if(packing: _Packing | None) -> _Packing: + return _Packing(0, 0, 0) if packing is None else packing + + @property + def arguments(self) -> list[str]: + return [f"{self.row}u", f"{self.head_stride}u", f"{self.dim_stride}u"] + + +@dataclass(frozen=True) +class _Geometry: + """A node's shape, read off its operands and the two head counts it declares.""" + + batch: int + steps: int + q_heads: int + kv_heads: int + d_k: int + d_v: int + + @property + def group(self) -> int: + return self.q_heads // self.kv_heads + + @property + def output_shape(self) -> tuple[int, ...]: + return (self.batch, self.steps, self.q_heads * self.d_v) + + @property + def state_shape(self) -> tuple[int, ...]: + return (self.batch, self.kv_heads, self.d_k, self.d_v) + + @property + def arguments(self) -> list[str]: + return [ + f"{self.batch}u", + f"{self.steps}u", + f"{self.q_heads}u", + f"{self.kv_heads}u", + f"{self.group}u", + f"{self.d_k}u", + f"{self.d_v}u", + ] + + +def _linear_attention(context: NodeContext) -> NodeEmission: + gating, delta = _update_rule(context) + geometry = _geometry(context) + decay = _decay_packing(context, geometry) if gating else None + beta = _beta_packing(context, geometry) if delta else None + operands = tuple(context.optional_input(index) for index in range(_BETA + 1)) + results = (context.require_output(_OUTPUT), context.require_output(_PRESENT_STATE)) + _verify_element_types(context, operands, results) + verify_shape(context, results[_OUTPUT], geometry.output_shape) + verify_shape(context, results[_PRESENT_STATE], geometry.state_shape) + + elem_type = context.require_input(_QUERY).elem_type + name = kernel_name(context, numpy_dtype_name(elem_type)) + definition = _KERNEL_TEMPLATE.substitute( + name=name, + element=c_type(elem_type), + accumulate=c_type(_ACCUMULATOR), + f=math_suffix(_ACCUMULATOR), + zero=scalar_literal(0, _ACCUMULATOR), + ) + arguments = [ + results[_OUTPUT].expr, + results[_PRESENT_STATE].expr, + *(_operand(ref) for ref in operands), + *geometry.arguments, + # An absent gate reaches the kernel as the NULL pointer its branch reads, so the + # strides that would have placed it are zeros nothing ever addresses through. + *_Packing.absent_if(decay).arguments, + # Beta is one value per `(token, head)`, so it has no key dimension to stride along. + *_Packing.absent_if(beta).arguments[:2], + scalar_literal(_scale(context, geometry), TensorProto.DOUBLE), + ] + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call_kernel(name, arguments),), + ) + + +def _operand(ref: TensorRef | None) -> str: + return "NULL" if ref is None else ref.expr + + +def _update_rule(context: NodeContext) -> tuple[bool, bool]: + """The rule's `(gating, delta correction)`, checked against the operands the node passes. + + ONNX ties the two optional operands to the rule in both directions: `decay` is required + by `gated` and `gated_delta` and forbidden by the other two, `beta` by `delta` and + `gated_delta`. The reference raises on every other combination rather than ignoring the + stray operand, so there is nothing for a kernel to compute for one. + """ + schema = onnx.defs.get_schema( + context.node.op_type, context.since_version, context.domain + ) + default = schema.attributes["update_rule"].default_value.s + rule = context.attribute("update_rule", default).decode() + if rule not in _UPDATE_RULES: + raise CompileError( + f"Node `{context.label}`: `LinearAttention`'s `update_rule` is `{rule}`, which " + f"is not one of the recurrences ONNX defines: {', '.join(_UPDATE_RULES)}." + ) + reads = _UPDATE_RULES[rule] + for index, name, needed in zip((_DECAY, _BETA), ("decay", "beta"), reads): + given = context.optional_input(index) is not None + if given != needed: + raise CompileError( + f"Node `{context.label}`: `LinearAttention`'s `update_rule` is `{rule}`, " + f"which {'requires' if needed else 'forbids'} the `{name}` input, but this " + f"node {'leaves it out' if needed else 'passes one'}." + ) + return reads + + +def _geometry(context: NodeContext) -> _Geometry: + """The extents the kernel walks, from the packed operands and the declared head counts. + + `d_k` is the query's own head width and `d_v` the value's; the key carries `d_k` too, + which is what makes the outer product it writes the shape of the state. + """ + query, key, value = ( + context.require_input(index) for index in (_QUERY, _KEY, _VALUE) + ) + for ref in (query, key, value): + _verify_rank(context, ref, 3, "packed as (B, T, H * D)") + q_heads = context.int_attribute("q_num_heads") + kv_heads = context.int_attribute("kv_num_heads") + if q_heads <= 0 or kv_heads <= 0 or q_heads % kv_heads != 0: + raise CompileError( + f"Node `{context.label}`: `LinearAttention` shares each of its {kv_heads} " + f"key/value head(s) between an equal number of its {q_heads} query head(s), so " + "`q_num_heads` has to be a positive multiple of `kv_num_heads`." + ) + d_k = _head_width(context, query, q_heads, "query") + d_v = _head_width(context, value, kv_heads, "value") + geometry = _Geometry( + batch=query.shape[0], + steps=query.shape[1], + q_heads=q_heads, + kv_heads=kv_heads, + d_k=d_k, + d_v=d_v, + ) + _verify_packed(context, key, (geometry.batch, geometry.steps, kv_heads * d_k)) + _verify_packed(context, value, (geometry.batch, geometry.steps, kv_heads * d_v)) + past = context.optional_input(_PAST_STATE) + if past is not None and past.shape != geometry.state_shape: + raise CompileError( + f"Node `{context.label}`: `LinearAttention` carries a state of shape " + f"{list(geometry.state_shape)}, but its `past_state` `{past.name}` has shape " + f"{list(past.shape)}." + ) + return geometry + + +def _head_width(context: NodeContext, ref: TensorRef, heads: int, role: str) -> int: + """How wide one head's slice of a packed operand is, which has to divide evenly.""" + packed = ref.shape[2] + if packed % heads != 0: + raise CompileError( + f"Node `{context.label}`: `LinearAttention` packs {heads} head(s) into the " + f"last dimension of its {role} `{ref.name}`, which holds {packed} element(s) " + "and is not divisible by that." + ) + return packed // heads + + +def _decay_packing(context: NodeContext, geometry: _Geometry) -> _Packing: + """Where the decay gate of a `(token, head, key dim)` sits in a `(B, T, L)` operand. + + `L` says the granularity: one value per head, broadcast across the key dimensions, or + one per key dimension. The per-head reading is tried first, as the reference does, so + the two agree where a `d_k` of one makes both apply and mean the same thing. + """ + decay = context.require_input(_DECAY) + _verify_rank(context, decay, 3, "packed as (B, T, H_kv) or (B, T, H_kv * d_k)") + heads, packed = geometry.kv_heads, decay.shape[2] + if packed == heads: + packing = _Packing(row=heads, head_stride=1, dim_stride=0) + elif packed == heads * geometry.d_k: + packing = _Packing(row=packed, head_stride=geometry.d_k, dim_stride=1) + else: + raise CompileError( + f"Node `{context.label}`: `LinearAttention` reads its `decay` `{decay.name}` " + f"either per head — {heads} value(s) — or per key dimension — " + f"{heads * geometry.d_k} — but its last dimension holds {packed}." + ) + _verify_packed(context, decay, (geometry.batch, geometry.steps, packed)) + return packing + + +def _beta_packing(context: NodeContext, geometry: _Geometry) -> _Packing: + """Where the update rate of a `(token, head)` sits in a `(B, T, L)` operand. + + One value per head, or a single one the whole batch item shares — which the reference + reaches by broadcasting the head axis, and this by a zero stride. + """ + beta = context.require_input(_BETA) + _verify_rank(context, beta, 3, "packed as (B, T, H_kv) or (B, T, 1)") + heads, packed = geometry.kv_heads, beta.shape[2] + if packed not in (heads, 1): + raise CompileError( + f"Node `{context.label}`: `LinearAttention` reads its `beta` `{beta.name}` " + f"either per head — {heads} value(s) — or as one the heads share, but its last " + f"dimension holds {packed}." + ) + _verify_packed(context, beta, (geometry.batch, geometry.steps, packed)) + return _Packing(row=packed, head_stride=1 if packed == heads else 0, dim_stride=0) + + +def _scale(context: NodeContext, geometry: _Geometry) -> float: + """What the read is scaled by: the node's own factor, or `1/sqrt(d_k)` for its default. + + ONNX states the default as the attribute value 0, which no model could mean literally — + it would answer every query with zero — and the reference reads it as the request for + the derived factor. A `d_k` of zero leaves that an infinity, which is what numpy's own + division by `sqrt(0)` yields and what the whole result then rests on: with no key + dimension to sum over, every answer is `inf * 0`, a NaN, in the reference and here alike. + """ + scale = context.float_attribute("scale") + if scale != 0.0: + return scale + return math.inf if geometry.d_k == 0 else 1.0 / math.sqrt(geometry.d_k) + + +def _verify_rank(context: NodeContext, ref: TensorRef, rank: int, form: str) -> None: + if len(ref.shape) != rank: + raise CompileError( + f"Node `{context.label}`: `LinearAttention` reads `{ref.name}` as a rank-" + f"{rank} tensor {form}, but it has shape {list(ref.shape)}." + ) + + +def _verify_packed( + context: NodeContext, ref: TensorRef, expected: tuple[int, ...] +) -> None: + """Refuse to emit a kernel whose addressing disagrees with an operand it is handed. + + Every packed operand shares the batch and the sequence with the query, and its last + dimension is fixed by the head counts the node declares; a graph that states another + shape describes a different op, and this is where that stops rather than where it reads + past a buffer. + """ + if ref.shape != expected: + raise CompileError( + f"Node `{context.label}`: `LinearAttention` addresses `{ref.name}` as a tensor " + f"of shape {list(expected)}, but it holds {list(ref.shape)}." + ) + + +def _verify_element_types( + context: NodeContext, + operands: tuple[TensorRef | None, ...], + results: tuple[TensorRef, ...], +) -> None: + """Refuse every element type but the one the recurrence accumulates in. + + ONNX types the state independently of the activations, and the reference accumulates in + float32 whatever either of them holds. The running state lives in the `present_state` + buffer here, which is sound exactly while that buffer holds the accumulator's own type; + the compiler supports no other of the op's three types anyway, so the two never part. + """ + for ref in (*operands, *results): + if ref is not None and ref.elem_type != _ACCUMULATOR: + raise CompileError( + f"Node `{context.label}`: `LinearAttention` accumulates in " + f"`{element_type_name(_ACCUMULATOR)}`, which is the only one of its element " + f"types the C compiler serves, but `{ref.name}` is " + f"`{element_type_name(ref.elem_type)}`." + ) + + +register_kernel("", "LinearAttention", _VERSIONS, _linear_attention) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/logic.py b/src/python/fnnx/extras/compilers/c/onnx/ops/logic.py new file mode 100644 index 0000000..71bf7ec --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/logic.py @@ -0,0 +1,157 @@ +"""The ops whose result is a decision: comparisons, logic, bit manipulation, selection.""" + +from __future__ import annotations + +from functools import partial + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type, element_size +from fnnx.extras.compilers.c.onnx.kernels import ( + NodeContext, + NodeEmission, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import elementwise, expand, pointwise + +# Predicates over two operands, with the schema revisions each expression covers. Revision 1 +# of the older ones carried the legacy broadcast attributes and is left unregistered; the +# revisions listed after the first only widened the type constraints. And, Or and Xor take +# boolean operands, whose emitted bytes are 0 or 1, so `!=` is exclusive-or on them. +_BINARY_PREDICATES: dict[str, tuple[tuple[int, ...], str]] = { + "And": ((7,), "x0 && x1"), + "Equal": ((7, 11, 13, 19), "x0 == x1"), + "Greater": ((7, 9, 13), "x0 > x1"), + "GreaterOrEqual": ((12, 16), "x0 >= x1"), + "Less": ((7, 9, 13), "x0 < x1"), + "LessOrEqual": ((12, 16), "x0 <= x1"), + "Or": ((7,), "x0 || x1"), + "Xor": ((7,), "x0 != x1"), +} + +_NOT_VERSIONS = (1,) +_BITWISE_VERSIONS = (18,) +_BITWISE_OPERATORS = {"BitwiseAnd": "&", "BitwiseOr": "|", "BitwiseXor": "^"} +_BIT_SHIFT_VERSIONS = (11,) +# IsInf-10 and IsNaN-9/13 predate the float8 types the later revisions accept; neither op's +# behaviour on the compilable types has changed. +_IS_INF_VERSIONS = (10, 20) +_IS_NAN_VERSIONS = (9, 13, 20) +# Where-9 already broadcasts all three operands; 16 only widened the type constraints. +_WHERE_VERSIONS = (9, 16) + +_SHIFT_OPERATORS = {"LEFT": "<<", "RIGHT": ">>"} + + +def _binary(context: NodeContext, *, expression: str) -> NodeEmission: + """A two-operand kernel computing `expression` over `x0` and `x1`. + + The result is cast to the op's own output type: boolean for the predicates, the + operands' type for the bitwise family, both of which `$element` resolves to. + """ + result = context.require_output(0) + return elementwise( + context, + expression=expand(f"($element)({expression})", result.elem_type), + operands=(context.require_input(0), context.require_input(1)), + result=result, + ) + + +def _bit_shift(context: NodeContext) -> NodeEmission: + """BitShift, whose operands are unsigned, so both directions shift in zeros. + + A shift by the operand's own width or more is undefined in C and unstated by ONNX; numpy + — the spec's executable form, and this compiler's oracle — yields zero, which the guard + reproduces rather than leaving to whatever the target's shift instruction does. + """ + result = context.require_output(0) + direction = _shift_direction(context) + width = 8 * element_size(result.elem_type) + # The narrow types promote to `int`, where a left shift can overflow into the sign bit; + # shifting in the widest unsigned type instead keeps every intermediate defined. + promoted = c_type(TensorProto.UINT64 if width == 64 else TensorProto.UINT32) + shift = f"(({promoted})x0 {_SHIFT_OPERATORS[direction]} x1)" + return elementwise( + context, + expression=expand( + f"($element)(x1 >= {width}u ? 0u : {shift})", result.elem_type + ), + operands=(context.require_input(0), context.require_input(1)), + result=result, + variant=f"_{direction.lower()}", + ) + + +def _shift_direction(context: NodeContext) -> str: + value = context.attribute("direction", b"") + direction = value.decode() if isinstance(value, bytes) else str(value) + if direction not in _SHIFT_OPERATORS: + raise CompileError( + f"Node `{context.label}`: BitShift's `direction` attribute is " + f"`{direction}`, but ONNX defines only " + f"{' and '.join(f'`{name}`' for name in _SHIFT_OPERATORS)}." + ) + return direction + + +def _is_inf(context: NodeContext) -> NodeEmission: + """IsInf, whose two attributes select which of the infinities count as one.""" + result = context.require_output(0) + positive = bool(context.int_attribute("detect_positive")) + negative = bool(context.int_attribute("detect_negative")) + if not positive and not negative: + # No operand: nothing is detected, so reading one would leave the kernel with an + # unused local, which the artifact's `-Werror` build contract does not allow. + return elementwise( + context, + expression=expand("$zero", result.elem_type), + operands=(), + result=result, + variant="_never", + ) + if positive and negative: + return pointwise(context, "($element)(isinf(x0) != 0)", variant="_any") + sign = ">" if positive else "<" + return pointwise( + context, + f"($element)(isinf(x0) && x0 {sign} 0)", + variant="_positive" if positive else "_negative", + ) + + +def _where(context: NodeContext) -> NodeEmission: + return elementwise( + context, + expression="x0 ? x1 : x2", + operands=tuple(context.require_input(index) for index in range(3)), + result=context.require_output(0), + ) + + +for _op_type, (_versions, _expression) in _BINARY_PREDICATES.items(): + register_kernel("", _op_type, _versions, partial(_binary, expression=_expression)) +for _op_type, _operator in _BITWISE_OPERATORS.items(): + register_kernel( + "", + _op_type, + _BITWISE_VERSIONS, + partial(_binary, expression=f"x0 {_operator} x1"), + ) +register_kernel("", "Not", _NOT_VERSIONS, partial(pointwise, template="($element)!x0")) +register_kernel( + "", + "BitwiseNot", + _BITWISE_VERSIONS, + partial(pointwise, template="($element)~x0"), +) +register_kernel("", "BitShift", _BIT_SHIFT_VERSIONS, _bit_shift) +register_kernel("", "IsInf", _IS_INF_VERSIONS, _is_inf) +register_kernel( + "", + "IsNaN", + _IS_NAN_VERSIONS, + partial(pointwise, template="($element)(isnan(x0) != 0)"), +) +register_kernel("", "Where", _WHERE_VERSIONS, _where) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/loss.py b/src/python/fnnx/extras/compilers/c/onnx/ops/loss.py new file mode 100644 index 0000000..45cf358 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/loss.py @@ -0,0 +1,327 @@ +"""SoftmaxCrossEntropyLoss: a log-softmax along the class axis, read at the labels. + +The op reads its operand as instances by classes by everything else — `[N, C, D...]`, which a +row-major buffer already lays out as `N` blocks of `C` planes of `D` elements — so the class +axis is walked by a stride and no reshape is ever emitted. Each `(instance, position)` pair +takes one pass for the largest logit, one for the exponentials, and one read at the label: +the loss of an entry is the negated log-softmax at its own class, and the optional `log_prob` +output is that same log-softmax written out in full. + +The logarithm is taken of the *normalized probability* — `log(exp(x - max) / total)` rather +than the algebraically equal `x - max - log(total)` — because that is what ONNX's own +LogSoftmax and the reference evaluator compute, and the two differ where the exponential +underflows: the first answers a logit far below the largest with `-inf`, the second with a +large finite number. + +`reduction`, the optional `weights` operand and whether `ignore_index` is set at all decide +how the per-entry losses are folded, and each combination is a kernel of its own rather than a +run-time branch — an entry ignored contributes nothing to either side of the weighted mean, +and a kernel that takes no weights must not carry a parameter it never reads past the +artifact's `-Werror` build. The index's *value* is an ordinary argument, as every other +attribute here is: it changes what a shared kernel skips, not the code that skips it. + +Labels are validated at run time, since a label outside the class axis is an out-of-bounds +read that no static check can rule out, and the kernel returns the artifact's +invalid-argument status for one. ONNX defines the labels as class indices, so that is the +whole of the range; the reference raises on one at or past the class count but reads a +negative one from the end of the axis the way numpy indexes, which is an artifact of how it +gathers rather than anything the op is defined to compute. +""" + +from __future__ import annotations + +import math +from string import Template + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + FLOAT_TYPES, + c_type, + element_type_name, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + checked_call, + kernel_name, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import math_suffix + +# SoftmaxCrossEntropyLoss-12 differs from 13 only in the type constraints it allows; nothing +# vouches for the older revision, whose backend tests the corpus does not carry, so only the +# current one is served and an older import gets the standard unsupported-version error. +_VERSIONS = (13,) + +_REDUCTIONS = ("none", "mean", "sum") + +# The three passes over one entry's class axis, shared by both kernels below: the largest +# logit, the sum of the exponentials around it, and the probability of a single class. +_SOFTMAX_PASSES = """\ + $element largest = -INFINITY; + $element total = $zero; + for (cls = 0; cls < classes; ++cls) { + const $element logit = in[base + cls * inner]; + if (logit > largest || isnan(logit)) { + largest = logit; + } + } + for (cls = 0; cls < classes; ++cls) { + total += exp$f(in[base + cls * inner] - largest); + }""" + +_LOG_PROB_TEMPLATE = f"""\ +static void $name( + $element* out, + const $element* in, + size_t rows, + size_t classes, + size_t inner) +{{ + size_t row, position, cls; + for (row = 0; row < rows; ++row) {{ + for (position = 0; position < inner; ++position) {{ + const size_t base = row * classes * inner + position; +{_SOFTMAX_PASSES} + for (cls = 0; cls < classes; ++cls) {{ + out[base + cls * inner] = + log$f(exp$f(in[base + cls * inner] - largest) / total); + }} + }} + }} +}}""" + +# `$ignored` skips an entry the node's `ignore_index` names, `$weight` is what the entry +# contributes to the weighted mean's denominator, and `$write` folds it: the three parts the +# attribute combination decides. The label check runs before the class axis is read, so a +# label the reference would raise on never addresses the buffer. +_LOSS_TEMPLATE = f"""\ +static int $name( + $element* out, + const $element* in, + const $label* labels, +$weight_parameter\ +$ignore_parameter\ + size_t rows, + size_t classes, + size_t inner) +{{ +$accumulators\ + size_t row, position, cls; + for (row = 0; row < rows; ++row) {{ + for (position = 0; position < inner; ++position) {{ + const size_t entry = row * inner + position; + const size_t base = row * classes * inner + position; + const int64_t label = (int64_t)labels[entry]; +$ignored\ + if (label < 0 || label >= (int64_t)classes) {{ + return 1; + }} +{_SOFTMAX_PASSES} + {{ + const $element chosen = exp$f( + in[base + (size_t)label * inner] - largest) / total; + const $element weight = $weight; + const $element value = weight * -log$f(chosen); +$write\ + }} + }} + }} +$finish\ + return 0; +}}""" + +_IGNORED_NONE = """\ + if (label == ignore_index) { + out[entry] = $zero; + continue; + } +""" + +_IGNORED_REDUCED = """\ + if (label == ignore_index) { + continue; + } +""" + +_ACCUMULATORS = """\ + $element loss_total = $zero; + $element weight_total = $zero; +""" + +_WRITE_ELEMENT = """\ + out[entry] = value; +""" + +_WRITE_ACCUMULATED = """\ + loss_total += value; + weight_total += weight; +""" + +# The weighted mean divides by the weights it actually summed, not by the number of entries: +# an ignored one contributes to neither side, and one whose weight is zero contributes to +# neither either. A denominator of zero divides zero by zero, which is the NaN the reference +# answers with rather than a guess at what it should be. +_FINISH_MEAN = """\ + out[0] = loss_total / weight_total; +""" + +_FINISH_SUM = """\ + out[0] = loss_total; + (void)weight_total; +""" + + +def _softmax_cross_entropy_loss(context: NodeContext) -> NodeEmission: + scores = context.require_input(0) + labels = context.require_input(1) + result = context.require_output(0) + weights = context.optional_input(2) + reduction = _reduction(context) + rows, classes, inner = _partition(context, scores, labels) + + if scores.elem_type not in FLOAT_TYPES: + raise CompileError( + f"Node `{context.label}`: SoftmaxCrossEntropyLoss takes `{scores.name}` as " + f"logits, so it must be a floating-point tensor; this one is " + f"`{element_type_name(scores.elem_type)}`." + ) + if weights is not None and weights.shape != (classes,): + raise CompileError( + f"Node `{context.label}`: `{weights.name}` gives one weight per class, so ONNX " + f"gives it shape [{classes}]; this model gives it {list(weights.shape)}." + ) + verify_shape(context, result, labels.shape if reduction == "none" else ()) + + functions = [] + statements = [] + log_prob = context.outputs[1] if len(context.outputs) > 1 else None + if log_prob is not None: + verify_shape(context, log_prob, scores.shape) + functions.append(_log_prob_kernel(context, scores)) + statements.append( + call_kernel( + functions[-1].name, + [log_prob.expr, scores.expr, f"{rows}u", f"{classes}u", f"{inner}u"], + ) + ) + + kernel = _loss_kernel(context, scores, labels, weights, reduction) + arguments = [result.expr, scores.expr, labels.expr] + if weights is not None: + arguments.append(weights.expr) + ignore_index = context.attribute("ignore_index", None) + if ignore_index is not None: + arguments.append(scalar_literal(int(ignore_index), TensorProto.INT64)) + arguments += [f"{rows}u", f"{classes}u", f"{inner}u"] + functions.append(kernel) + statements.append(checked_call(context, kernel.name, arguments)) + return NodeEmission(functions=tuple(functions), statements=tuple(statements)) + + +def _log_prob_kernel(context: NodeContext, scores: TensorRef) -> CFunction: + name = kernel_name(context, "log_prob", numpy_dtype_name(scores.elem_type)) + return CFunction( + name, _fill(_LOG_PROB_TEMPLATE, name=name, elem_type=scores.elem_type) + ) + + +def _loss_kernel( + context: NodeContext, + scores: TensorRef, + labels: TensorRef, + weights: TensorRef | None, + reduction: str, +) -> CFunction: + ignore_index = context.attribute("ignore_index", None) + elem_type = scores.elem_type + ignored = "" + if ignore_index is not None: + template = _IGNORED_NONE if reduction == "none" else _IGNORED_REDUCED + ignored = _fill(template, elem_type=elem_type) + name = kernel_name( + context, + reduction, + "weighted" if weights is not None else "plain", + "ignoring" if ignore_index is not None else "all", + numpy_dtype_name(elem_type), + numpy_dtype_name(labels.elem_type), + ) + definition = _fill( + _LOSS_TEMPLATE, + name=name, + elem_type=elem_type, + label=c_type(labels.elem_type), + weight_parameter=( + f" const {c_type(elem_type)}* weights,\n" if weights is not None else "" + ), + ignore_parameter=( + " int64_t ignore_index,\n" if ignore_index is not None else "" + ), + ignored=ignored, + accumulators=( + "" if reduction == "none" else _fill(_ACCUMULATORS, elem_type=elem_type) + ), + weight=( + "weights[(size_t)label]" + if weights is not None + else scalar_literal(1, elem_type) + ), + write=_WRITE_ELEMENT if reduction == "none" else _WRITE_ACCUMULATED, + finish={"none": "", "mean": _FINISH_MEAN, "sum": _FINISH_SUM}[reduction], + ) + return CFunction(name, definition) + + +def _fill(template: str, *, elem_type: int, **fields: str) -> str: + """Substitute the element-type placeholders every template here shares, then `fields`.""" + return Template(template).substitute( + element=c_type(elem_type), + f=math_suffix(elem_type), + zero=scalar_literal(0, elem_type), + **fields, + ) + + +def _reduction(context: NodeContext) -> str: + reduction = context.attribute("reduction", b"mean") + name = reduction.decode() if isinstance(reduction, bytes) else str(reduction) + if name not in _REDUCTIONS: + raise CompileError( + f"Node `{context.label}`: `reduction` names `{name}`, which is not one of " + f"{', '.join(_REDUCTIONS)}." + ) + return name + + +def _partition( + context: NodeContext, scores: TensorRef, labels: TensorRef +) -> tuple[int, int, int]: + """The operand read as instances by classes by positions, checked against the labels.""" + if len(scores.shape) < 2: + raise CompileError( + f"Node `{context.label}`: SoftmaxCrossEntropyLoss reads `{scores.name}` as " + "instances by classes by any further axes, so it needs a rank of at least 2; " + f"this one has shape {list(scores.shape)}." + ) + expected = (scores.shape[0], *scores.shape[2:]) + if labels.shape != expected: + raise CompileError( + f"Node `{context.label}`: `{labels.name}` names one class per entry of " + f"`{scores.name}`, so ONNX gives it shape {list(expected)}; this model gives " + f"it {list(labels.shape)}." + ) + return scores.shape[0], scores.shape[1], math.prod(scores.shape[2:]) + + +register_kernel("", "SoftmaxCrossEntropyLoss", _VERSIONS, _softmax_cross_entropy_loss) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/ml.py b/src/python/fnnx/extras/compilers/c/onnx/ops/ml.py new file mode 100644 index 0000000..8d96a33 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/ml.py @@ -0,0 +1,722 @@ +"""The ONNX-ML preprocessing ops: the feature transforms a fitted pipeline is made of. + +They differ from the standard domain in where their parameters live: a scaler's offsets, an +encoder's categories and a label encoder's key/value pairs arrive as *attributes* rather than +as operands, so each one is emitted as `static const` data the shared kernel reads through a +pointer — which keeps one kernel per (op, element types) however many nodes run it, and keeps +the tables in the artifact's reported footprint rather than on the stack. + +Every one of these ops is at revision 1 of its schema and has been since ONNX-ML opset 1, +except `LabelEncoder`, whose revision 4 is the only one the reference evaluator implements +faithfully — the older two are left to the standard unsupported-version error. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from string import Template + +import numpy as np +from onnx import TensorProto +from onnx.numpy_helper import to_array + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + FLOAT_TYPES, + c_type, + element_type_name, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + constant_data, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.loader import ML_DOMAIN +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + checked_call, + kernel_name, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import Scalar, pointwise + +_INITIAL_VERSIONS = (1,) +_LABEL_ENCODER_VERSIONS = (4,) + +# What `Normalizer`'s reference implementation divides by when a row's norm underflows to +# zero, so that an all-zero row comes out as itself rather than as NaN. It is not in the +# op's prose, but it is what the executable form of the specification computes. +_NORM_FLOOR = "1e-30f" + +# The two lines each norm contributes to the row loop: how a value joins the running total, +# and what turns that total into the divisor. The largest magnitude is taken the way numpy's +# own `max` takes it, so one NaN in a row carries through to every element of it. +_NORMS = { + "MAX": ( + " const float magnitude = fabsf(value);\n" + " if (magnitude > norm || isnan(magnitude)) {\n" + " norm = magnitude;\n" + " }", + "", + ), + "L1": (" norm += fabsf(value);", ""), + "L2": (" norm += value * value;", " norm = sqrtf(norm);\n"), +} + +# What `OneHotEncoder` does about a value in no category, by whether `zeros` is set: leave +# the row at zero, or report the failure the schema prescribes through the status enum. +_ONE_HOT_MISS = " if (chosen < 0) {\n return 1;\n }\n" + +_SCALER_TEMPLATE = Template("""\ +static void $name( + float* out, + const $element* in, + const float* offset, + const float* scale, + size_t count, + size_t channels, + size_t offset_stride, + size_t scale_stride) +{ + size_t index; + for (index = 0; index < count; ++index) { + const size_t channel = index % channels; + out[index] = ((float)in[index] - offset[channel * offset_stride]) + * scale[channel * scale_stride]; + } +}""") + +_NORMALIZER_TEMPLATE = Template("""\ +static void $name( + float* out, + const $element* in, + size_t rows, + size_t columns) +{ + size_t row, column; + for (row = 0; row < rows; ++row) { + const size_t base = row * columns; + float norm = 0.0f; + for (column = 0; column < columns; ++column) { + const float value = (float)in[base + column]; +$accumulate + } +$finish if (norm < $floor) { + norm = $floor; + } + for (column = 0; column < columns; ++column) { + out[base + column] = (float)in[base + column] / norm; + } + } +}""") + +_IMPUTER_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + const $element* imputed, + size_t count, + size_t channels, + size_t imputed_stride$marker) +{ + size_t index; + for (index = 0; index < count; ++index) { + const $element value = in[index]; + out[index] = ($test) + ? imputed[(index % channels) * imputed_stride] + : value; + } +}""") + +_ONE_HOT_TEMPLATE = Template("""\ +static $status $name( + float* out, + const $element* in, + const $element* categories, + size_t count, + size_t category_count) +{ + size_t index, category; + for (index = 0; index < count; ++index) { + const $element value = in[index]; + ptrdiff_t chosen = -1; + for (category = 0; category < category_count; ++category) { + if (value == categories[category]) { + chosen = (ptrdiff_t)category; + } + } +$missing for (category = 0; category < category_count; ++category) { + out[index * category_count + category] = + (chosen == (ptrdiff_t)category) ? 1.0f : 0.0f; + } + } +$result}""") + +_LABEL_ENCODER_TEMPLATE = Template("""\ +static void $name( + $result* out, + const $element* in, + const $element* keys, + const $result* values, + size_t count, + size_t pair_count, + $result fallback) +{ + size_t index, pair; + for (index = 0; index < count; ++index) { + const $element value = in[index]; + $result mapped = fallback; + for (pair = 0; pair < pair_count; ++pair) { + if ($match) { + mapped = values[pair]; + } + } + out[index] = mapped; + } +}""") + +_EXTRACTOR_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const int64_t* indices, + size_t rows, + size_t extent, + size_t index_count) +{ + size_t row, chosen; + for (row = 0; row < rows; ++row) { + for (chosen = 0; chosen < index_count; ++chosen) { + ptrdiff_t position = (ptrdiff_t)indices[chosen]; + if (position < 0) { + position += (ptrdiff_t)extent; + } + if (position < 0 || position >= (ptrdiff_t)extent) { + return 1; + } + out[row * index_count + chosen] = in[row * extent + (size_t)position]; + } + } + return 0; +}""") + +_VECTORIZER_TEMPLATE = Template("""\ +static void $name( + float* out, + const $element* in, + size_t rows, + size_t columns, + size_t taken, + size_t width, + size_t stride, + size_t offset) +{ + size_t row, column; + for (row = 0; row < rows; ++row) { + for (column = 0; column < width; ++column) { + out[row * stride + offset + column] = + (column < taken) ? (float)in[row * columns + column] : 0.0f; + } + } +}""") + + +def _scaler(context: NodeContext) -> NodeEmission: + """`Y = (X - offset) * scale`, per feature, as the float32 the schema declares `Y`. + + Both coefficient lists are either one value shared by every feature or one per feature, + which is what their stride at the call site says. + """ + source = context.require_input(0) + result = _float_result(context, source.shape) + channels = _trailing_extent(source.shape) + offset = _coefficients(context, "offset", channels) + scale = _coefficients(context, "scale", channels) + + name = kernel_name(context, c_type(source.elem_type)) + definition = _SCALER_TEMPLATE.substitute( + name=name, element=c_type(source.elem_type) + ) + offset_data, offset_symbol = constant_data(context, "offset", offset) + scale_data, scale_symbol = constant_data(context, "scale", scale) + call = call_kernel( + name, + [ + result.expr, + source.expr, + offset_symbol, + scale_symbol, + f"{result.elem_count}u", + f"{max(1, channels)}u", + f"{int(len(offset) > 1)}u", + f"{int(len(scale) > 1)}u", + ], + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call,), + constants=(offset_data, scale_data), + ) + + +def _normalizer(context: NodeContext) -> NodeEmission: + """Each row divided by its own norm: the largest magnitude in it, its L1 or its L2.""" + source = context.require_input(0) + result = _float_result(context, source.shape) + if not source.shape: + raise CompileError( + f"Node `{context.label}`: Normalizer runs along the last axis of a `[C]` or " + "`[N,C]` tensor, and its input is a scalar." + ) + declared = context.attribute("norm", b"MAX") + norm = declared.decode() if isinstance(declared, bytes) else str(declared) + accumulate, finish = _NORMS.get(norm, (None, None)) + if accumulate is None: + raise CompileError( + f"Node `{context.label}`: `norm` is `{norm}`, which is none of the modes ONNX " + f"defines ({', '.join(sorted(_NORMS))})." + ) + + columns = source.shape[-1] + name = kernel_name(context, norm.lower(), c_type(source.elem_type)) + definition = _NORMALIZER_TEMPLATE.substitute( + name=name, + element=c_type(source.elem_type), + accumulate=accumulate, + finish=finish, + floor=_NORM_FLOOR, + ) + call = call_kernel( + name, + [ + result.expr, + source.expr, + f"{math.prod(source.shape[:-1])}u", + f"{columns}u", + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _imputer(context: NodeContext) -> NodeEmission: + """Every element equal to the marked value replaced by this feature's imputed one. + + ONNX carries the pair as two attribute families, one for the floating-point element + types and one for the integer ones; the float family is the one in force whenever it is + set, which is also how the reference evaluator picks between them. A marker that is NaN + is a test for NaN rather than a comparison, since nothing compares equal to it. + """ + source = context.require_input(0) + result = context.require_output(0) + verify_shape(context, result, source.shape) + if result.elem_type != source.elem_type: + raise CompileError( + f"Node `{context.label}`: Imputer leaves the element type alone, but its " + f"input is `{element_type_name(source.elem_type)}` and its output " + f"`{element_type_name(result.elem_type)}`." + ) + + floats = [float(value) for value in context.attribute("imputed_value_floats", [])] + integers = [int(value) for value in context.attribute("imputed_value_int64s", [])] + if bool(floats) == bool(integers): + raise CompileError( + f"Node `{context.label}`: Imputer must set exactly one of " + "`imputed_value_floats` and `imputed_value_int64s`." + ) + channels = _trailing_extent(source.shape) + values = np.array(floats or integers).astype(numpy_dtype_name(source.elem_type)) + if len(values) not in (1, channels): + raise CompileError( + f"Node `{context.label}`: Imputer was given {len(values)} imputed value(s) " + f"for a tensor whose last axis holds {channels}; it takes either one value " + "for every feature or a single value for all of them." + ) + + # The float family compares in double, which is what promoting an integer tensor to the + # attribute's own type does; the integer family compares in the tensor's own type, where + # a value too large for a double to hold exactly still has to match itself. + if floats: + marker, compare = context.float_attribute("replaced_value_float"), "double" + else: + marker, compare = ( + context.int_attribute("replaced_value_int64"), + c_type(source.elem_type), + ) + # Nothing compares equal to NaN, so a NaN marker is a test for it — and the marker then + # stops being a value the kernel is handed at all. + matches_nan = compare == "double" and math.isnan(marker) + test = "isnan((double)value)" if matches_nan else f"({compare})value == replaced" + name = kernel_name( + context, "nan" if matches_nan else compare, c_type(source.elem_type) + ) + definition = _IMPUTER_TEMPLATE.substitute( + name=name, + element=c_type(source.elem_type), + marker="" if matches_nan else f",\n {compare} replaced", + test=test, + ) + data, symbol = constant_data(context, "imputed", values) + arguments = [ + result.expr, + source.expr, + symbol, + f"{result.elem_count}u", + f"{max(1, channels)}u", + f"{int(len(values) > 1)}u", + ] + if not matches_nan: + arguments.append( + scalar_literal( + marker, TensorProto.DOUBLE if compare == "double" else source.elem_type + ) + ) + call = call_kernel(name, arguments) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call,), + constants=(data,), + ) + + +def _binarizer(context: NodeContext) -> NodeEmission: + """One where the value is strictly above the threshold, zero elsewhere. + + The comparison runs in double because that is what numpy's own promotion of the + tensor against the threshold — a value of the attribute's floating-point type — does. + """ + source = context.require_input(0) + verify_shape(context, context.require_output(0), source.shape) + return pointwise( + context, + "($element)((double)x0 > threshold)", + scalars=( + Scalar( + "threshold", TensorProto.DOUBLE, context.float_attribute("threshold") + ), + ), + ) + + +def _one_hot_encoder(context: NodeContext) -> NodeEmission: + """A row of indicators per element, one column per category, all zero for a miss. + + A repeated category is the last of its occurrences, which is the entry a mapping built + from the list in order ends up holding. With `zeros` cleared, a value in no category is + the failure the op's schema prescribes, reported through the status enum. + """ + source = context.require_input(0) + categories = [int(value) for value in context.attribute("cats_int64s", [])] + if not categories: + raise CompileError( + f"Node `{context.label}`: OneHotEncoder sets no `cats_int64s`; a numeric input " + "is encoded against the integer category list, `cats_strings` being reachable " + "only from the string input this compiler does not support." + ) + result = _float_result(context, (*source.shape, len(categories))) + + zeros = bool(context.int_attribute("zeros")) + name = kernel_name( + context, "zeros" if zeros else "strict", c_type(source.elem_type) + ) + definition = _ONE_HOT_TEMPLATE.substitute( + name=name, + element=c_type(source.elem_type), + status="void" if zeros else "int", + missing="" if zeros else _ONE_HOT_MISS, + result="" if zeros else " return 0;\n", + ) + data, symbol = constant_data( + context, + "categories", + np.array(categories).astype(numpy_dtype_name(source.elem_type)), + ) + arguments = [ + result.expr, + source.expr, + symbol, + f"{source.elem_count}u", + f"{len(categories)}u", + ] + call = ( + call_kernel(name, arguments) + if zeros + else checked_call(context, name, arguments) + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call,), + constants=(data,), + ) + + +def _label_encoder(context: NodeContext) -> NodeEmission: + """Each element looked up in the key list and replaced by the value beside it. + + A repeated key takes its last occurrence, as the op's schema states in as many words — + which is why the loop runs to the end rather than stopping at the first hit. The same + paragraph makes a NaN key match any NaN input, so a floating-point key list tests for + that as well as for equality. + """ + source = context.require_input(0) + result = context.require_output(0) + verify_shape(context, result, source.shape) + keys, values, fallback = _label_pairs(context, source, result) + + floating = source.elem_type in FLOAT_TYPES + match = ( + "value == keys[pair] || (isnan((double)value) && isnan((double)keys[pair]))" + if floating + else "value == keys[pair]" + ) + name = kernel_name(context, c_type(source.elem_type), c_type(result.elem_type)) + definition = _LABEL_ENCODER_TEMPLATE.substitute( + name=name, + element=c_type(source.elem_type), + result=c_type(result.elem_type), + match=match, + ) + key_data, key_symbol = constant_data(context, "keys", keys) + value_data, value_symbol = constant_data(context, "values", values) + call = call_kernel( + name, + [ + result.expr, + source.expr, + key_symbol, + value_symbol, + f"{result.elem_count}u", + f"{len(keys)}u", + scalar_literal(fallback, result.elem_type), + ], + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call,), + constants=(key_data, value_data), + ) + + +def _array_feature_extractor(context: NodeContext) -> NodeEmission: + """The columns the index operand names, taken from every row of the last axis. + + The indices come from the caller, so each one is normalized the way numpy's own + indexing — which the reference evaluator uses — defines it, and then bounds checked. + """ + source = context.require_input(0) + indices = context.require_input(1) + result = context.require_output(0) + if not source.shape: + raise CompileError( + f"Node `{context.label}`: ArrayFeatureExtractor selects along the last axis of " + "its input, and its input is a scalar." + ) + if indices.elem_type != TensorProto.INT64: + raise CompileError( + f"Node `{context.label}`: ArrayFeatureExtractor takes `int64` indices, not " + f"`{element_type_name(indices.elem_type)}`." + ) + rows = math.prod(source.shape[:-1]) + verify_shape( + context, + result, + (1, indices.elem_count) + if len(source.shape) == 1 + else (*source.shape[:-1], indices.elem_count), + ) + + name = kernel_name(context, c_type(source.elem_type)) + definition = _EXTRACTOR_TEMPLATE.substitute( + name=name, element=c_type(source.elem_type) + ) + call = checked_call( + context, + name, + [ + result.expr, + source.expr, + indices.expr, + f"{rows}u", + f"{source.shape[-1]}u", + f"{indices.elem_count}u", + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _feature_vectorizer(context: NodeContext) -> NodeEmission: + """Every input cut or zero-padded to its declared width, laid side by side as float32.""" + sources = [context.require_input(index) for index in range(len(context.node.input))] + result = _float_result(context, None) + widths = [int(value) for value in context.attribute("inputdimensions", [])] + if len(widths) != len(sources): + raise CompileError( + f"Node `{context.label}`: FeatureVectorizer declares {len(widths)} entry(s) in " + f"`inputdimensions` for {len(sources)} input(s); it takes one width per input." + ) + rows, columns = _vectorizer_layout(context, sources) + verify_shape(context, result, (rows, sum(widths))) + + functions: list[CFunction] = [] + statements: list[str] = [] + offset = 0 + for source, width, count in zip(sources, widths, columns): + name = kernel_name(context, c_type(source.elem_type)) + functions.append( + CFunction( + name, + _VECTORIZER_TEMPLATE.substitute( + name=name, element=c_type(source.elem_type) + ), + ) + ) + statements.append( + call_kernel( + name, + [ + result.expr, + source.expr, + f"{rows}u", + f"{count}u", + f"{min(count, width)}u", + f"{width}u", + f"{sum(widths)}u", + f"{offset}u", + ], + ) + ) + offset += width + return NodeEmission(functions=tuple(functions), statements=tuple(statements)) + + +def _vectorizer_layout( + context: NodeContext, sources: Sequence[TensorRef] +) -> tuple[int, tuple[int, ...]]: + """The shared row count and each input's own column count. + + A vector input is read as the single column a matrix of one column would be, which is + the shape the reference evaluator gives it before concatenating. + """ + layouts = [] + for source in sources: + if len(source.shape) not in (1, 2): + raise CompileError( + f"Node `{context.label}`: FeatureVectorizer takes inputs of rank 1 or 2, " + f"but `{source.name}` has shape {list(source.shape)}." + ) + layouts.append( + (source.shape[0], source.shape[1] if len(source.shape) == 2 else 1) + ) + rows = {row for row, _ in layouts} + if len(rows) != 1: + raise CompileError( + f"Node `{context.label}`: FeatureVectorizer lays its inputs side by side, so " + f"they must agree on their first dimension; these hold {sorted(rows)} rows." + ) + return layouts[0][0], tuple(count for _, count in layouts) + + +def _label_pairs( + context: NodeContext, source: TensorRef, result: TensorRef +) -> tuple[np.ndarray, np.ndarray, float | int]: + """`LabelEncoder`'s key/value tables and its default, in the attribute family in force. + + ONNX carries the mapping in one of three families and the schema picks whichever is set, + in the order the reference evaluator reads them; the pairs stop at the shorter of the two + lists. The string family never reaches here — ONNX's own type inference refuses a key type + that differs from the input's, and a string result is refused by the static verification — + so it is the numeric families or nothing. + """ + keys = _label_table(context, "keys", source.elem_type) + values = _label_table(context, "values", result.elem_type) + if keys is None or values is None: + raise CompileError( + f"Node `{context.label}`: LabelEncoder sets no numeric `keys_*`/`values_*` " + f"pair to map `{element_type_name(source.elem_type)}` through." + ) + pairs = min(len(keys), len(values)) + return keys[:pairs], values[:pairs], _label_default(context, result.elem_type) + + +def _label_table(context: NodeContext, role: str, elem_type: int) -> np.ndarray | None: + """The `role` list of the first attribute family the node sets, at `elem_type`.""" + for family in (f"{role}_floats", f"{role}_int64s"): + values = context.attribute(family, []) + if values: + return np.array(list(values)).astype(numpy_dtype_name(elem_type)) + tensor = context.attribute(f"{role}_tensor", None) + if tensor is not None: + return to_array(tensor).reshape(-1).astype(numpy_dtype_name(elem_type)) + return None + + +def _label_default(context: NodeContext, elem_type: int) -> float | int: + """The value an element matching no key takes, from the family the values came from.""" + tensor = context.attribute("default_tensor", None) + if tensor is not None: + return to_array(tensor).reshape(-1).tolist()[0] + if context.attribute("values_floats", []): + return context.float_attribute("default_float") + if context.attribute("values_int64s", []): + return context.int_attribute("default_int64") + # `values_tensor` with no `default_tensor`: the schema documents the default as -1 for an + # integral value type and -0 for a floating-point one, which are the two `default_*` + # attributes' own declared defaults. + return ( + context.float_attribute("default_float") + if elem_type in FLOAT_TYPES + else context.int_attribute("default_int64") + ) + + +def _coefficients(context: NodeContext, name: str, channels: int) -> np.ndarray: + """A `Scaler` coefficient list: one value per feature, or a single shared one.""" + values = [float(value) for value in context.attribute(name, [])] + if not values: + raise CompileError( + f"Node `{context.label}`: Scaler sets no `{name}`; it takes one value per " + "feature, or a single value for all of them." + ) + if len(values) not in (1, channels): + raise CompileError( + f"Node `{context.label}`: Scaler was given {len(values)} `{name}` value(s) for " + f"a tensor whose last axis holds {channels}; it takes either one value for " + "every feature or a single value for all of them." + ) + return np.array(values, dtype=np.float32) + + +def _float_result(context: NodeContext, shape: tuple[int, ...] | None) -> TensorRef: + """The node's result, checked to be the float32 tensor its ONNX schema declares.""" + result = context.require_output(0) + if result.elem_type != TensorProto.FLOAT: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` produces a `float` tensor, " + f"but its output `{result.name}` is declared " + f"`{element_type_name(result.elem_type)}`." + ) + if shape is not None: + verify_shape(context, result, shape) + return result + + +def _trailing_extent(shape: tuple[int, ...]) -> int: + """How many features a tensor holds per row: its last axis, or one for a scalar.""" + return shape[-1] if shape else 1 + + +register_kernel(ML_DOMAIN, "Scaler", _INITIAL_VERSIONS, _scaler) +register_kernel(ML_DOMAIN, "Normalizer", _INITIAL_VERSIONS, _normalizer) +register_kernel(ML_DOMAIN, "Imputer", _INITIAL_VERSIONS, _imputer) +register_kernel(ML_DOMAIN, "Binarizer", _INITIAL_VERSIONS, _binarizer) +register_kernel(ML_DOMAIN, "OneHotEncoder", _INITIAL_VERSIONS, _one_hot_encoder) +register_kernel(ML_DOMAIN, "LabelEncoder", _LABEL_ENCODER_VERSIONS, _label_encoder) +register_kernel( + ML_DOMAIN, "ArrayFeatureExtractor", _INITIAL_VERSIONS, _array_feature_extractor +) +register_kernel(ML_DOMAIN, "FeatureVectorizer", _INITIAL_VERSIONS, _feature_vectorizer) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/normalization.py b/src/python/fnnx/extras/compilers/c/onnx/ops/normalization.py new file mode 100644 index 0000000..1dbcf0e --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/normalization.py @@ -0,0 +1,936 @@ +"""The normalizations: standardizing a tensor by statistics taken from its own elements. + +LayerNormalization, InstanceNormalization, GroupNormalization, MeanVarianceNormalization and +BatchNormalization in training mode all do the same thing to different groups of elements: +take the group's mean and variance, subtract, divide, and — for all but the fourth — apply a +scale and a bias that vary along an axis. They share one loop nest, the group and element +loops `axes.py` builds, walked three times: once for the mean, once for the squared +deviations, once to write the result. Which axes form a group, where the epsilon guarding the +division goes, and how the affine operands are addressed is what each op supplies. +BatchNormalization at inference, RMSNormalization, LpNormalization and LRN take no group +statistics at all and carry loops of their own. + +The variance comes from the squared deviations rather than from `E[X^2] - E[X]^2`. The two +are the same quantity, and the deviation form is the one the reference evaluator and the +corpus's own expectations compute for every op here except GroupNormalization and +MeanVarianceNormalization, whose official function bodies subtract the squares — a +difference of one rounding error, on the side that cancellation cannot hurt. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass +from string import Template + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + FLOAT_TYPES, + c_type, + element_type_name, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + broadcast_strides, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + GROUP_PARAMETERS, + Grouping, + call_kernel, + group_axes, + kernel_name, + normalize_axes, + normalize_axis, + offset_helper, + verify_same_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents, math_suffix + +# BatchNormalization-15 only reworded 14's documentation and spelled out that the statistics +# of a float16 tensor accumulate in float. The revisions below 14 are a different op — they +# carry the `spatial` and `is_test` attributes and up to five outputs — and none is claimed. +_BATCH_VERSIONS = (14, 15) +_LAYER_VERSIONS = (17,) +_RMS_VERSIONS = (23,) +# GroupNormalization-18 took `scale` and `bias` per group rather than per channel and had no +# `stash_type`; it is a different op, so only the current revision is served. +_GROUP_VERSIONS = (21,) +# The remaining families changed only their type constraints across the revisions listed. +_INSTANCE_VERSIONS = (6, 22) +_LP_VERSIONS = (1, 22) +_MVN_VERSIONS = (9, 13) +_LRN_VERSIONS = (1, 13) + +# MeanVarianceNormalization's ONNX function body adds its epsilon to the standard deviation +# rather than to the variance, and standardizes these axes when the node names none. +_MVN_EPSILON = 1e-9 +_MVN_DEFAULT_AXES = (0, 2, 3) + +# ONNX defines LpNormalization for these orders only. +_SUPPORTED_ORDERS = (1, 2) + +# The three passes over a group: the mean, the squared deviations around it, and the write. +# A group with no elements at all divides zero by zero, which is the NaN numpy fills such a +# statistic with — the reference's own answer rather than a guess at what it should be. +_STANDARDIZE_TEMPLATE = Template("""\ +static void $name( +$parameters) +{ + size_t group, index; + for (group = 0; group < group_count; ++group) { + const size_t base = $offset(group, kept_rank, kept_shape, kept_strides); +$bases\ + $stash total = $stash_zero; + $stash spread = $stash_zero; + for (index = 0; index < group_size; ++index) { + total += ($stash)in[base + + $element_offset]; + } + const $stash mean = total / ($stash)group_size; + for (index = 0; index < group_size; ++index) { + const $stash deviation = ($stash)in[base + + $element_offset] - mean; + spread += deviation * deviation; + } + const $stash variance = spread / ($stash)group_size; + const $stash factor = $factor; + for (index = 0; index < group_size; ++index) { + const size_t position = base + + $element_offset; +$positions\ + const $stash centred = ($stash)in[position] - mean; + out[position] = ($element)( + $formula); + } +$statistics\ + } +}""") + +# Inference-mode BatchNormalization: the statistics arrive as operands, so there is nothing +# to reduce and an element's channel follows from where it sits in the buffer. +_BATCH_TEST_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + const $affine* scale, + const $affine* bias, + const $statistic* mean, + const $statistic* variance, + size_t count, + size_t inner, + size_t channels, + $stash epsilon) +{ + size_t index; + for (index = 0; index < count; ++index) { + const size_t channel = (index / inner) % channels; + out[index] = ($element)(($stash)scale[channel] + * (($stash)in[index] - ($stash)mean[channel]) + / sqrt$f(($stash)variance[channel] + epsilon) + + ($stash)bias[channel]); + } +}""") + +# RMSNormalization has no mean to centre on: each element is scaled by the reciprocal root of +# its group's mean square. The reference evaluator multiplies by that reciprocal rather than +# dividing by the root, and applies the linear coefficient afterwards, which is the rounding +# the corpus's own expectations were generated with. +_RMS_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + const $element* scale, +$parameters, + const size_t* scale_kept_strides, + const size_t* scale_reduced_strides, + $element epsilon) +{ + size_t group, index; + for (group = 0; group < group_count; ++group) { + const size_t base = $offset(group, kept_rank, kept_shape, kept_strides); + const size_t scale_base = + $offset(group, kept_rank, kept_shape, scale_kept_strides); + $element total = $zero; + for (index = 0; index < group_size; ++index) { + const $element x = in[base + $element_offset]; + total += x * x; + } + const $element factor = + $one / sqrt$f(total / ($element)group_size + epsilon); + for (index = 0; index < group_size; ++index) { + const size_t position = base + $element_offset; + const size_t scale_position = scale_base + + $offset(index, reduced_rank, reduced_shape, scale_reduced_strides); + out[position] = in[position] * factor * scale[scale_position]; + } + } +}""") + +# LpNormalization divides a group by its own norm, and answers a norm of zero with zero +# rather than with the NaN the division would give — as the reference evaluator does. +_LP_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, +$parameters) +{ + size_t group, index; + for (group = 0; group < group_count; ++group) { + const size_t base = $offset(group, kept_rank, kept_shape, kept_strides); + $element total = $zero; + for (index = 0; index < group_size; ++index) { + const $element x = in[base + + $element_offset]; + total += $term; + } + const $element norm = $norm; + for (index = 0; index < group_size; ++index) { + const size_t position = base + + $element_offset; + out[position] = (norm == $zero) ? $zero : in[position] / norm; + } + } +}""") + +# LRN's group is a window of channels around each element's own, so the channel axis is +# walked explicitly rather than through the grouping helpers. +_LRN_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + size_t batches, + size_t channels, + size_t inner, + size_t before, + size_t after, + $element bias, + $element scaled_alpha, + $element beta) +{ + size_t batch, channel, position, neighbour; + for (batch = 0; batch < batches; ++batch) { + for (channel = 0; channel < channels; ++channel) { + const size_t begin = (channel >= before) ? channel - before : 0; + const size_t last = channel + after + 1; + const size_t end = (last < channels) ? last : channels; + for (position = 0; position < inner; ++position) { + const size_t target = (batch * channels + channel) * inner + position; + $element square_sum = $zero; + for (neighbour = begin; neighbour < end; ++neighbour) { + const $element value = + in[(batch * channels + neighbour) * inner + position]; + square_sum += value * value; + } + out[target] = + in[target] / pow$f(bias + scaled_alpha * square_sum, beta); + } + } + } +}""") + + +@dataclass(frozen=True) +class _Operand: + """An operand read alongside the data: a scale, a bias, a running statistic. + + `strides` addresses it from the data's own coordinates, so a per-channel vector and a + tensor of the normalized shape are the same thing to the kernel — a stride per axis, zero + on every axis the operand does not vary along. `per_group` marks one that varies no + faster than the groups do and is read as `_base`, which is all a statistic blended + once per group needs; anything read per element carries `_position` as well. + """ + + name: str + ref: TensorRef + strides: tuple[int, ...] + per_group: bool = False + + +@dataclass(frozen=True) +class _Statistic: + """A per-group value the node takes as an output of its own. + + `value` is C over the locals the group loop has computed — `mean`, `variance`, `factor` — + and over any operand's `_base`, which for a per-group operand is its own index. + """ + + index: int + name: str + ref: TensorRef + value: str + + +@dataclass(frozen=True) +class _Scalar: + """An attribute the kernel reads as an argument rather than as an inlined literal.""" + + name: str + elem_type: int + value: float + + +def _standardize( + context: NodeContext, + *, + data: TensorRef, + result: TensorRef, + grouping: Grouping, + stash: int, + factor: str, + formula: str, + operands: Sequence[_Operand] = (), + statistics: Sequence[_Statistic] = (), + scalars: Sequence[_Scalar] = (), +) -> NodeEmission: + """Emit the group-statistics kernel and the call site addressing this node's buffers. + + `factor` is C over the group's `mean` and `variance`, and `formula` is C over `factor`, + the element's own `centred` deviation from the mean, and each operand's + `[_position]`; both are `$`-templated over the element and stash types. + """ + offset = offset_helper(context.prefix) + element_offset = ( + f"{offset.name}(index, reduced_rank, reduced_shape, reduced_strides)" + ) + element = c_type(result.elem_type) + + parameters = [f" {element}* out"] + parameters += [ + f" {c_type(entry.ref.elem_type)}* {entry.name}" for entry in statistics + ] + parameters.append(f" const {c_type(data.elem_type)}* in") + parameters += [ + f" const {c_type(operand.ref.elem_type)}* {operand.name}" + for operand in operands + ] + parameters.append(GROUP_PARAMETERS) + for operand in operands: + parameters.append(f" const size_t* {operand.name}_kept_strides") + if not operand.per_group: + parameters.append(f" const size_t* {operand.name}_reduced_strides") + parameters += [ + f" {c_type(scalar.elem_type)} {scalar.name}" for scalar in scalars + ] + + def expand(text: str) -> str: + return Template(text).substitute( + element=element, + stash=c_type(stash), + f=math_suffix(stash), + one=scalar_literal(1, stash), + zero=scalar_literal(0, stash), + ) + + name = _kernel_name( + context, + "", + [ + result.elem_type, + stash, + *(entry.ref.elem_type for entry in statistics), + *(operand.ref.elem_type for operand in operands), + ], + f"p{len(operands)}" + "".join(f"s{entry.index}" for entry in statistics), + ) + definition = _STANDARDIZE_TEMPLATE.substitute( + name=name, + parameters=",\n".join(parameters), + offset=offset.name, + element_offset=element_offset, + element=element, + stash=c_type(stash), + stash_zero=scalar_literal(0, stash), + factor=expand(factor), + formula=expand(formula), + bases="".join( + f" const size_t {operand.name}_base = {offset.name}(\n" + f" group, kept_rank, kept_shape, {operand.name}_kept_strides);\n" + for operand in operands + ), + positions="".join( + f" const size_t {operand.name}_position = {operand.name}_base\n" + f" + {offset.name}(index, reduced_rank, reduced_shape,\n" + f" {operand.name}_reduced_strides);\n" + for operand in operands + if not operand.per_group + ), + statistics="".join( + f" {entry.name}[group] = " + f"({c_type(entry.ref.elem_type)})({expand(entry.value)});\n" + for entry in statistics + ), + ) + + arguments = [result.expr] + arguments += [entry.ref.expr for entry in statistics] + arguments.append(data.expr) + arguments += [operand.ref.expr for operand in operands] + arguments += grouping.arguments + for operand in operands: + kept, reduced = _split_strides(grouping, operand.strides) + arguments += [kept] if operand.per_group else [kept, reduced] + arguments += [scalar_literal(scalar.value, scalar.elem_type) for scalar in scalars] + return NodeEmission( + functions=(offset, CFunction(name, definition)), + statements=(call_kernel(name, arguments),), + ) + + +def _split_strides(grouping: Grouping, strides: Sequence[int]) -> list[str]: + """An operand's per-axis strides, split the way the grouping splits the data's axes.""" + return [ + extents([strides[axis] for axis in grouping.kept_axes]), + extents([strides[axis] for axis in grouping.reduced_axes]), + ] + + +def _kernel_name( + context: NodeContext, variant: str, elem_types: Sequence[int], form: str +) -> str: + """A name encoding everything the emitted code depends on beyond the call-site literals. + + Two nodes running the same op share a kernel when their element types, the operands and + statistics they name, and the formula their attributes select all agree; anything else + would be two kernels colliding on one name. + """ + names = [numpy_dtype_name(elem_type) for elem_type in elem_types] + types = names[0] if len(set(names)) == 1 else "_".join(names) + return kernel_name(context, *(part for part in (variant, form, types) if part)) + + +# -------------------------------------------------------------------------------------- +# The ops +# -------------------------------------------------------------------------------------- + + +def _layer_normalization(context: NodeContext) -> NodeEmission: + """Standardize each row from `axis` on, then scale and shift it. + + `Scale` and `B` carry the normalized shape, but ONNX applies them by broadcasting, so + they are addressed through strides over the data's own axes rather than by position. + Stage one runs in the element type `stash_type` names, which is also the type ONNX gives + the mean and inverse deviation this op can report. + """ + data = context.require_input(0) + result = context.require_output(0) + verify_same_shape(context, data, result) + rank = len(data.shape) + axis = normalize_axis(context, context.int_attribute("axis"), rank) + grouping = group_axes(data.shape, tuple(range(axis, rank))) + stash = _stash_element(context) + + operands = [_broadcast_operand(context, "scale", context.require_input(1), data)] + formula = "($element)(centred * factor) * scale[scale_position]" + bias = context.optional_input(2) + if bias is not None: + operands.append(_broadcast_operand(context, "bias", bias, data)) + formula += " + bias[bias_position]" + + statistics = _statistics( + context, grouping, ((1, "mean_out", "mean"), (2, "inv_std_out", "factor")) + ) + for statistic in statistics: + if statistic.ref.elem_type != stash: + raise CompileError( + f"Node `{context.label}`: ONNX types `{statistic.ref.name}` by the " + f"`stash_type` attribute, which names `{element_type_name(stash)}`, but " + f"the graph gives it `{element_type_name(statistic.ref.elem_type)}`." + ) + return _standardize( + context, + data=data, + result=result, + grouping=grouping, + stash=stash, + factor="$one / sqrt$f(variance + epsilon)", + formula=formula, + operands=operands, + statistics=statistics, + scalars=(_Scalar("epsilon", stash, context.float_attribute("epsilon")),), + ) + + +def _rms_normalization(context: NodeContext) -> NodeEmission: + """Scale each row from `axis` on by the reciprocal root of its own mean square. + + ONNX's function body computes stage one in the element type `stash_type` names, casting + the data to it and back; the reference evaluator ignores the attribute, computes in the + data's own type, and refuses outright any value but the default. The kernel follows the + reference — it is what both suites compare against — and refuses the other values for the + same reason the reference's refusal gives them: nothing vouches for what they compute. + """ + data = context.require_input(0) + scale = context.require_input(1) + result = context.require_output(0) + verify_same_shape(context, data, result) + stash = context.int_attribute("stash_type") + if stash != TensorProto.FLOAT: + raise CompileError( + f"Node `{context.label}`: `stash_type` names element type " + f"`{element_type_name(stash)}`. The ONNX reference evaluator refuses every " + "value but the default here and takes RMSNormalization's statistics in the " + "data's own type, so nothing vouches for what another one computes." + ) + for operand in (scale, result): + # ONNX's own inference refuses a model whose scale and data disagree; a graph that + # reached here with one would have the kernel read a buffer at the wrong width, and + # this is where that stops rather than where it corrupts memory. + if operand.elem_type != data.elem_type: + raise CompileError( + f"Node `{context.label}`: RMSNormalization gives `{operand.name}` element " + f"type `{element_type_name(operand.elem_type)}` while `{data.name}` has " + f"`{element_type_name(data.elem_type)}`; the C compiler serves this op at " + "one element type only." + ) + + rank = len(data.shape) + axis = normalize_axis(context, context.int_attribute("axis"), rank) + grouping = group_axes(data.shape, tuple(range(axis, rank))) + offset = offset_helper(context.prefix) + elem_type = data.elem_type + name = _kernel_name(context, "", (elem_type,), "") + definition = _RMS_TEMPLATE.substitute( + name=name, + parameters=GROUP_PARAMETERS, + offset=offset.name, + element_offset=( + f"{offset.name}(index, reduced_rank, reduced_shape, reduced_strides)" + ), + element=c_type(elem_type), + f=math_suffix(elem_type), + one=scalar_literal(1, elem_type), + zero=scalar_literal(0, elem_type), + ) + strides = broadcast_strides(scale, data.shape, node_label=context.label) + arguments = [ + result.expr, + data.expr, + scale.expr, + *grouping.arguments, + *_split_strides(grouping, strides), + scalar_literal(context.float_attribute("epsilon"), elem_type), + ] + return NodeEmission( + functions=(offset, CFunction(name, definition)), + statements=(call_kernel(name, arguments),), + ) + + +def _instance_normalization(context: NodeContext) -> NodeEmission: + """Standardize each channel of each instance over its spatial axes.""" + data = context.require_input(0) + result = context.require_output(0) + verify_same_shape(context, data, result) + rank = _require_channel_axis(context, data) + grouping = group_axes(data.shape, tuple(range(2, rank))) + elem_type = result.elem_type + return _standardize( + context, + data=data, + result=result, + grouping=grouping, + stash=elem_type, + factor="sqrt$f(variance + epsilon)", + formula="scale[scale_position] * centred / factor + bias[bias_position]", + operands=_channel_operands(context, data, (("scale", 1), ("bias", 2))), + scalars=(_Scalar("epsilon", elem_type, context.float_attribute("epsilon")),), + ) + + +def _group_normalization(context: NodeContext) -> NodeEmission: + """Standardize each group of channels of each instance, then scale per channel. + + The data is read as though reshaped to `[N, num_groups, group_size, ...]`, which for a + contiguous row-major buffer is a change of coordinates and nothing more. + """ + data = context.require_input(0) + result = context.require_output(0) + verify_same_shape(context, data, result) + rank = _require_channel_axis(context, data) + channels = _channels(data) + groups = context.int_attribute("num_groups") + if groups <= 0 or channels % groups: + raise CompileError( + f"Node `{context.label}`: GroupNormalization splits {channels} channel(s) into " + f"`num_groups` = {groups} group(s), which does not divide them evenly." + ) + size = channels // groups + reshaped = (data.shape[0], groups, size, *data.shape[2:]) + grouping = group_axes(reshaped, tuple(range(2, len(reshaped)))) + # An element's channel is its group times the group's size plus its position within the + # group, which is what these strides — over the reshaped axes — add up to. + strides = (0, size, 1) + (0,) * (rank - 2) + stash = _stash_element(context) + return _standardize( + context, + data=data, + result=result, + grouping=grouping, + stash=stash, + factor="sqrt$f(variance + epsilon)", + formula=( + "($element)(centred / factor) * scale[scale_position] + bias[bias_position]" + ), + operands=[ + _Operand(name, _per_channel(context, index, channels), strides) + for name, index in (("scale", 1), ("bias", 2)) + ], + scalars=(_Scalar("epsilon", stash, context.float_attribute("epsilon")),), + ) + + +def _mean_variance_normalization(context: NodeContext) -> NodeEmission: + """Standardize over the named axes, guarding the division at the standard deviation. + + ONNX defines this op as a function whose epsilon is added to the deviation rather than + to the variance, which is what the reference evaluator and the corpus both compute. + """ + data = context.require_input(0) + result = context.require_output(0) + verify_same_shape(context, data, result) + axes = context.attribute("axes", list(_MVN_DEFAULT_AXES)) + grouping = group_axes(data.shape, normalize_axes(context, axes, len(data.shape))) + elem_type = result.elem_type + return _standardize( + context, + data=data, + result=result, + grouping=grouping, + stash=elem_type, + factor="sqrt$f(variance) + epsilon", + formula="centred / factor", + scalars=(_Scalar("epsilon", elem_type, _MVN_EPSILON),), + ) + + +def _batch_normalization(context: NodeContext) -> NodeEmission: + """Normalize per channel: by the statistics handed to it, or by the batch's own.""" + if context.int_attribute("training_mode"): + return _batch_training(context) + for index in (1, 2): + extra = context.outputs[index] if index < len(context.outputs) else None + if extra is not None: + raise CompileError( + f"Node `{context.label}`: BatchNormalization computes `{extra.name}` in " + "training mode only, and this node runs at inference, where ONNX leaves " + "the extra outputs undefined." + ) + return _batch_test(context) + + +def _batch_test(context: NodeContext) -> NodeEmission: + """Inference: the mean and variance are operands, so nothing is reduced.""" + data = context.require_input(0) + result = context.require_output(0) + verify_same_shape(context, data, result) + channels = _channels(data) + scale, bias, mean, variance = ( + _per_channel(context, index, channels) for index in (1, 2, 3, 4) + ) + + stash = _widest(data.elem_type, scale.elem_type, mean.elem_type) + name = _kernel_name( + context, "test", (result.elem_type, stash, scale.elem_type, mean.elem_type), "" + ) + definition = _BATCH_TEST_TEMPLATE.substitute( + name=name, + element=c_type(result.elem_type), + affine=c_type(scale.elem_type), + statistic=c_type(mean.elem_type), + stash=c_type(stash), + f=math_suffix(stash), + ) + arguments = [ + result.expr, + data.expr, + scale.expr, + bias.expr, + mean.expr, + variance.expr, + f"{result.elem_count}u", + f"{math.prod(data.shape[2:])}u", + f"{channels}u", + scalar_literal(context.float_attribute("epsilon"), stash), + ] + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call_kernel(name, arguments),), + ) + + +def _batch_training(context: NodeContext) -> NodeEmission: + """Training: the batch's own statistics normalize it, and carry the running ones on. + + The running statistics are read and written per channel, which is exactly one group + here, so the operands they blend are only passed when the node asks for them. + """ + data = context.require_input(0) + result = context.require_output(0) + verify_same_shape(context, data, result) + rank = len(data.shape) + channels = _channels(data) + grouping = group_axes(data.shape, tuple(axis for axis in range(rank) if axis != 1)) + strides = _channel_strides(rank) + + operands = _channel_operands(context, data, (("scale", 1), ("bias", 2))) + statistics = _statistics( + context, + grouping, + ( + ( + 1, + "running_mean", + "($stash)input_mean[input_mean_base] * momentum" + " + mean * ($one - momentum)", + ), + ( + 2, + "running_var", + "($stash)input_var[input_var_base] * momentum" + " + variance * ($one - momentum)", + ), + ), + ) + blended = {1: ("input_mean", 3), 2: ("input_var", 4)} + for statistic in statistics: + name, index = blended[statistic.index] + operands.append( + _Operand( + name, _per_channel(context, index, channels), strides, per_group=True + ) + ) + + stash = _widest( + data.elem_type, + context.require_input(1).elem_type, + context.require_input(3).elem_type, + ) + scalars = [_Scalar("epsilon", stash, context.float_attribute("epsilon"))] + if statistics: + # `momentum` blends the running statistics and is read nowhere else, so a node + # that reports neither takes an argument its kernel never touches — which the + # artifact's own `-Werror=unused-parameter` build refuses. + scalars.append(_Scalar("momentum", stash, context.float_attribute("momentum"))) + return _standardize( + context, + data=data, + result=result, + grouping=grouping, + stash=stash, + factor="sqrt$f(variance + epsilon)", + formula=( + "($stash)scale[scale_position] * centred / factor" + " + ($stash)bias[bias_position]" + ), + operands=operands, + statistics=statistics, + scalars=scalars, + ) + + +def _lp_normalization(context: NodeContext) -> NodeEmission: + """Divide each row along one axis by its own Lp norm. + + The norm sums absolute values, which is what ONNX defines and what the corpus's own + expectations compute; the reference evaluator raises the elements to the power `p` + without taking their absolute value, so for `p` = 1 it is an oracle on non-negative + operands only. + """ + data = context.require_input(0) + result = context.require_output(0) + verify_same_shape(context, data, result) + order = context.int_attribute("p") + if order not in _SUPPORTED_ORDERS: + raise CompileError( + f"Node `{context.label}`: ONNX defines LpNormalization for `p` in " + f"{list(_SUPPORTED_ORDERS)}, but this node asks for {order}." + ) + axis = normalize_axis(context, context.int_attribute("axis"), len(data.shape)) + grouping = group_axes(data.shape, (axis,)) + + offset = offset_helper(context.prefix) + elem_type = result.elem_type + suffix = math_suffix(elem_type) + name = _kernel_name(context, f"l{order}", (elem_type,), "") + definition = _LP_TEMPLATE.substitute( + name=name, + parameters=GROUP_PARAMETERS, + offset=offset.name, + element_offset=( + f"{offset.name}(index, reduced_rank, reduced_shape, reduced_strides)" + ), + element=c_type(elem_type), + zero=scalar_literal(0, elem_type), + term=f"fabs{suffix}(x)" if order == 1 else "x * x", + norm="total" if order == 1 else f"sqrt{suffix}(total)", + ) + return NodeEmission( + functions=(offset, CFunction(name, definition)), + statements=(call_kernel(name, [result.expr, data.expr, *grouping.arguments]),), + ) + + +def _local_response_normalization(context: NodeContext) -> NodeEmission: + """Divide each element by a power of the squared sum of its channel neighbourhood.""" + data = context.require_input(0) + result = context.require_output(0) + verify_same_shape(context, data, result) + _require_channel_axis(context, data) + size = context.int_attribute("size") + if size <= 0: + raise CompileError( + f"Node `{context.label}`: LRN sums over a window of `size` channels, which " + f"this node gives as {size}." + ) + elem_type = result.elem_type + name = _kernel_name(context, "", (elem_type,), "") + definition = _LRN_TEMPLATE.substitute( + name=name, + element=c_type(elem_type), + zero=scalar_literal(0, elem_type), + f=math_suffix(elem_type), + ) + arguments = [ + result.expr, + data.expr, + f"{data.shape[0]}u", + f"{_channels(data)}u", + f"{math.prod(data.shape[2:])}u", + # The window reaches `floor((size - 1) / 2)` channels back and + # `ceil((size - 1) / 2)` channels on, both clamped to the tensor. + f"{(size - 1) // 2}u", + f"{size // 2}u", + scalar_literal(context.float_attribute("bias"), elem_type), + scalar_literal(context.float_attribute("alpha") / size, elem_type), + scalar_literal(context.float_attribute("beta"), elem_type), + ] + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call_kernel(name, arguments),), + ) + + +# -------------------------------------------------------------------------------------- +# Operands, types and shapes the family shares +# -------------------------------------------------------------------------------------- + + +def _stash_element(context: NodeContext) -> int: + """The element type stage one computes in, as the node's `stash_type` names it.""" + stash = context.int_attribute("stash_type") + if stash not in FLOAT_TYPES: + raise CompileError( + f"Node `{context.label}`: `stash_type` names element type " + f"`{element_type_name(stash)}`, which the C compiler cannot take statistics " + "in; only FLOAT and DOUBLE are supported." + ) + return stash + + +def _statistics( + context: NodeContext, + grouping: Grouping, + entries: Sequence[tuple[int, str, str]], +) -> list[_Statistic]: + """The per-group outputs this node actually asks for. + + An output ONNX declares optional may be left out entirely or named as the empty string, + and one the node omits is not computed at all. + """ + statistics = [] + for index, name, value in entries: + ref = context.outputs[index] if index < len(context.outputs) else None + if ref is None: + continue + if ref.elem_count != grouping.group_count: + # The groups are counted from the operand's shape and the axes the node + # normalizes over, while the buffer is sized from the shape ONNX inferred; a + # disagreement is a compiler bug, and this is where it stops rather than where + # it corrupts memory. + raise CompileError( + f"Node `{context.label}`: normalizing leaves {grouping.group_count} " + f"group(s), but its output `{ref.name}` holds {ref.elem_count} element(s)." + ) + statistics.append(_Statistic(index, name, ref, value)) + return statistics + + +def _require_channel_axis(context: NodeContext, data: TensorRef) -> int: + """The data's rank, refusing one that has no channel axis to normalize per.""" + rank = len(data.shape) + if rank < 2: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads `{data.name}` as " + "instances by channels by spatial axes, so it needs a rank of at least 2; this " + f"one has shape {list(data.shape)}." + ) + return rank + + +def _channels(data: TensorRef) -> int: + """The extent of the channel axis; a tensor of rank 1 is a single channel to ONNX.""" + return data.shape[1] if len(data.shape) > 1 else 1 + + +def _channel_strides(rank: int) -> tuple[int, ...]: + """Strides addressing a per-channel operand from the data's coordinates.""" + return tuple(1 if axis == 1 else 0 for axis in range(rank)) + + +def _per_channel(context: NodeContext, index: int, channels: int) -> TensorRef: + """An operand ONNX gives one value per channel, checked against that shape.""" + operand = context.require_input(index) + if operand.shape != (channels,): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` takes `{operand.name}` per " + f"channel, so ONNX gives it shape [{channels}]; this model gives it " + f"{list(operand.shape)}." + ) + return operand + + +def _channel_operands( + context: NodeContext, data: TensorRef, entries: Sequence[tuple[str, int]] +) -> list[_Operand]: + channels = _channels(data) + strides = _channel_strides(len(data.shape)) + return [ + _Operand(name, _per_channel(context, index, channels), strides) + for name, index in entries + ] + + +def _broadcast_operand( + context: NodeContext, name: str, operand: TensorRef, data: TensorRef +) -> _Operand: + return _Operand( + name, operand, broadcast_strides(operand, data.shape, node_label=context.label) + ) + + +def _widest(*elem_types: int) -> int: + """The element type numpy's promotion would compute these operands in.""" + return TensorProto.DOUBLE if TensorProto.DOUBLE in elem_types else TensorProto.FLOAT + + +register_kernel("", "BatchNormalization", _BATCH_VERSIONS, _batch_normalization) +register_kernel("", "LayerNormalization", _LAYER_VERSIONS, _layer_normalization) +register_kernel("", "RMSNormalization", _RMS_VERSIONS, _rms_normalization) +register_kernel( + "", "InstanceNormalization", _INSTANCE_VERSIONS, _instance_normalization +) +register_kernel("", "GroupNormalization", _GROUP_VERSIONS, _group_normalization) +register_kernel("", "LpNormalization", _LP_VERSIONS, _lp_normalization) +register_kernel( + "", "MeanVarianceNormalization", _MVN_VERSIONS, _mean_variance_normalization +) +register_kernel("", "LRN", _LRN_VERSIONS, _local_response_normalization) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/pad.py b/src/python/fnnx/extras/compilers/c/onnx/ops/pad.py new file mode 100644 index 0000000..394121a --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/pad.py @@ -0,0 +1,292 @@ +"""Pad: the operand placed inside a larger result, and what fills the rest of it. + +Every mode is the same walk — each element of the result maps to the coordinate `i - begin` +of the operand along each axis — and they differ only in what that means once the coordinate +leaves the operand: a constant value, the nearest edge, the reflection back inside, or the +wrap around to the other end. So there is one kernel per mode and element type, taking the +per-axis pads as compile-time literals. + +The pads themselves have to be fixed at compile time: they place the operand inside the +result, which no shape can state on their behalf. A negative pad, which ONNX defines as +cropping instead, needs nothing of its own — it is the same coordinate map, shifted the +other way. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from functools import partial +from string import Template + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + kernel_name, + normalize_axis, + row_major_strides, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents + +# Pad-2 takes the pads as an attribute; from 11 they are an operand, 18 added the `axes` +# operand that says which axes they apply to, and 19 the `wrap` mode. The later revisions +# only widened the element types. Every form reads as the same per-axis pad vector, so one +# generator serves them all, told where to read the pads from. +# +# Pad-1, which spells the attribute `paddings`, is deliberately left out: ONNX's own shape +# inference derives nothing for that revision, so a node of it has no result shape to +# compile against and no oracle to prove one against either. +_ATTRIBUTE_VERSIONS = (2,) +_OPERAND_VERSIONS = (11, 13, 18, 19, 21, 23, 24, 25) + +# What each mode does with a coordinate that falls outside the operand. `constant` is not +# here: it reads nothing at all, so it is a template of its own. +_MAPPINGS = { + "edge": """\ + if (coordinate < 0) { + coordinate = 0; + } else if (coordinate >= extent) { + coordinate = extent - 1; + }""", + "reflect": """\ + if (extent > 1) { + const ptrdiff_t period = 2 * extent - 2; + coordinate = ((coordinate % period) + period) % period; + if (coordinate >= extent) { + coordinate = period - coordinate; + } + } else { + coordinate = 0; + }""", + "wrap": """\ + coordinate = ((coordinate % extent) + extent) % extent;""", +} + +_CONSTANT_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + $element value, + size_t count, + int rank, + const size_t* shape, + const size_t* limits, + const ptrdiff_t* pads, + const size_t* strides) +{ + size_t index; + for (index = 0; index < count; ++index) { + size_t remainder = index; + size_t source = 0; + int inside = 1; + int axis; + for (axis = rank - 1; axis >= 0; --axis) { + const ptrdiff_t coordinate = + (ptrdiff_t)(remainder % shape[axis]) - pads[axis]; + remainder /= shape[axis]; + if (coordinate < 0 || coordinate >= (ptrdiff_t)limits[axis]) { + inside = 0; + } else { + source += (size_t)coordinate * strides[axis]; + } + } + out[index] = inside ? in[source] : value; + } +}""") + +_MAPPED_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + size_t count, + int rank, + const size_t* shape, + const size_t* limits, + const ptrdiff_t* pads, + const size_t* strides) +{ + size_t index; + for (index = 0; index < count; ++index) { + size_t remainder = index; + size_t source = 0; + int axis; + for (axis = rank - 1; axis >= 0; --axis) { + const ptrdiff_t extent = (ptrdiff_t)limits[axis]; + ptrdiff_t coordinate = (ptrdiff_t)(remainder % shape[axis]) - pads[axis]; + remainder /= shape[axis]; +$mapping + source += (size_t)coordinate * strides[axis]; + } + out[index] = in[source]; + } +}""") + + +def _pad(context: NodeContext, *, attribute: str | None) -> NodeEmission: + source = context.require_input(0) + result = context.require_output(0) + mode = _mode(context) + begins, ends = _pads(context, attribute, rank=len(source.shape)) + verify_shape( + context, + result, + [ + extent + begin + end + for extent, begin, end in zip(source.shape, begins, ends) + ], + ) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + if mode != "constant": + _verify_readable(context, source.shape, result.shape, mode) + + element = c_type(result.elem_type) + name = kernel_name(context, mode, element) + arguments = [ + result.expr, + source.expr, + *([_fill(context, attribute)] if mode == "constant" else []), + f"{result.elem_count}u", + str(len(result.shape)), + extents(result.shape), + extents(source.shape), + _offsets(begins), + extents(row_major_strides(source.shape)), + ] + definition = ( + _CONSTANT_TEMPLATE.substitute(name=name, element=element) + if mode == "constant" + else _MAPPED_TEMPLATE.substitute( + name=name, element=element, mapping=_MAPPINGS[mode] + ) + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call_kernel(name, arguments),), + ) + + +def _mode(context: NodeContext) -> str: + value = context.attribute("mode", b"constant") + mode = value.decode() if isinstance(value, bytes) else str(value) + if mode == "constant" or mode in _MAPPINGS: + return mode + raise CompileError( + f"Node `{context.label}`: `Pad` mode `{mode}` is not one of the modes ONNX " + f"defines ({', '.join(['constant', *sorted(_MAPPINGS)])})." + ) + + +def _pads( + context: NodeContext, attribute: str | None, *, rank: int +) -> tuple[list[int], list[int]]: + """The pad before and after each axis, zero for the axes the node leaves out.""" + values = _pad_values(context, attribute) + axes = _axes(context, rank) + if len(values) != 2 * len(axes): + raise CompileError( + f"Node `{context.label}`: `Pad` was given {len(values)} pad(s) for " + f"{len(axes)} axis/axes; ONNX defines two — a begin and an end — per axis." + ) + begins = [0] * rank + ends = [0] * rank + seen: set[int] = set() + for position, axis in enumerate(axes): + if axis in seen: + raise CompileError( + f"Node `{context.label}`: `Pad` names axis {axis} of its operand more " + "than once." + ) + seen.add(axis) + begins[axis] = values[position] + ends[axis] = values[len(axes) + position] + return begins, ends + + +def _pad_values(context: NodeContext, attribute: str | None) -> list[int]: + if attribute is not None: + declared = context.attribute(attribute, None) + if declared is None: + raise CompileError( + f"Node `{context.label}`: `Pad` requires the `{attribute}` attribute at " + f"opset version {context.since_version}." + ) + return [int(value) for value in declared] + return _constant_operand(context, 1, "pads") + + +def _axes(context: NodeContext, rank: int) -> list[int]: + """Which axes the pads apply to: the ones the node names, or every one of them.""" + if context.optional_input(3) is None: + return list(range(rank)) + return [ + normalize_axis(context, axis, rank) + for axis in _constant_operand(context, 3, "axes") + ] + + +def _fill(context: NodeContext, attribute: str | None) -> str: + """What a constant pad is filled with, as a C expression. + + From 11 on it is an operand rather than an attribute, and one whose value decides + nothing about any shape — so it is read at run time and needs no folding. + """ + result = context.require_output(0) + if attribute is not None: + return scalar_literal(context.float_attribute("value"), result.elem_type) + operand = context.optional_input(2) + if operand is None: + return scalar_literal(0, result.elem_type) + if operand.elem_count != 1: + raise CompileError( + f"Node `{context.label}`: `Pad` fills with `{operand.name}`, which holds " + f"{operand.elem_count} values; ONNX defines it as a single one." + ) + return f"{operand.expr}[0]" + + +def _constant_operand(context: NodeContext, index: int, role: str) -> list[int]: + operand = context.require_input(index) + values = context.constant_input(index) + if values is None: + raise CompileError( + f"Node `{context.label}`: `Pad` takes its {role} from `{operand.name}`, which " + "is not known at compile time; where the operand sits inside the result — and " + "the shape of the result itself — then depends on input data, which the C " + "compiler cannot compile." + ) + return [int(value) for value in values.reshape(-1)] + + +def _verify_readable( + context: NodeContext, source: Sequence[int], result: Sequence[int], mode: str +) -> None: + """Refuse a mode that would have to read a value from an axis that holds none. + + Every mode but `constant` fills the result out of the operand, so an axis the operand is + empty along has nothing to fill a wider result with — as ONNX's own reference refuses too. + """ + for axis, (extent, wanted) in enumerate(zip(source, result)): + if extent == 0 and wanted > 0: + raise CompileError( + f"Node `{context.label}`: `Pad` in mode `{mode}` fills axis {axis} of its " + f"result, which measures {wanted}, from an operand that is empty along it." + ) + + +def _offsets(values: Sequence[int]) -> str: + """The per-axis pads as a compound literal; they are signed, since a pad may crop.""" + literals = ", ".join(str(value) for value in values) or "0" + return f"(const ptrdiff_t[]){{{literals}}}" + + +register_kernel("", "Pad", _ATTRIBUTE_VERSIONS, partial(_pad, attribute="pads")) +register_kernel("", "Pad", _OPERAND_VERSIONS, partial(_pad, attribute=None)) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/pool.py b/src/python/fnnx/extras/compilers/c/onnx/ops/pool.py new file mode 100644 index 0000000..86408af --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/pool.py @@ -0,0 +1,611 @@ +"""The poolings: the same sliding window as a convolution, folding rather than weighting. + +Every pooling walks the geometry `window.py` resolves — for each output position, each tap +reads the operand at `position * stride + tap * dilation - pad` along each spatial axis — and +folds what it finds into one value: the largest, the mean, or the Lp norm. So there is one +template, one kernel per fold and element type, and the geometry reaches it as call-site +literals. The `Global*` family is that same walk at a window the size of the operand's spatial +extent, which is why it shares those kernels outright. + +Two things only a pooling has. `ceil_mode` rounds the number of window positions up instead of +down, which lets the last window hang off the end of the operand. And a tap is *counted* +separately from being *read*: it counts when it lands inside the operand widened by the node's +own pads — which is what `count_include_pad` averages over — while only a tap inside the +operand itself is read. The padding `ceil_mode` implies is neither read nor counted. + +`MaxUnpool` runs a max pooling backwards: it scatters each element to the position that +pooling's `Indices` output recorded, leaving the rest at zero. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from string import Template + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + checked_call, + row_major_strides, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import combiner, expand, extents +from fnnx.extras.compilers.c.onnx.ops.reduce import extremum_identity, extremum_test +from fnnx.extras.compilers.c.onnx.ops.window import ( + auto_pad_mode, + declared_pads, + offsets, + resolve_pads, + spatial_attribute, + spatial_extents, +) + +# The parameters every pooling kernel takes after its buffers. `counted_shape` is the operand +# widened by the node's own pads: a tap at or past it belongs to the padding `ceil_mode` +# added, which no pooling reads or counts. +_GEOMETRY_PARAMETERS = """\ + size_t plane_count, + size_t input_size, + size_t output_size, + size_t window_size, + int spatial_rank, + const size_t* input_shape, + const size_t* counted_shape, + const size_t* output_shape, + const size_t* window_shape, + const size_t* strides, + const size_t* dilations, + const ptrdiff_t* pads""" + +# Where one tap of one window lands, as the offset into the operand plus the two flags a fold +# reads it through. +_WALK = Template("""\ + size_t remaining_position = position; + size_t remaining_tap = tap; + size_t offset = 0; + size_t stride = 1; + int inside = 1; + int counted = 1; + int axis; + for (axis = spatial_rank - 1; axis >= 0; --axis) { + const ptrdiff_t coordinate = + (ptrdiff_t)(remaining_position % output_shape[axis]) + * (ptrdiff_t)strides[axis] + + (ptrdiff_t)(remaining_tap % window_shape[axis]) + * (ptrdiff_t)dilations[axis] + - pads[axis]; + remaining_position /= output_shape[axis]; + remaining_tap /= window_shape[axis]; + if (coordinate >= (ptrdiff_t)counted_shape[axis]) { + counted = 0; + } + if (coordinate < 0 || coordinate >= (ptrdiff_t)input_shape[axis]) { + inside = 0; + } else { + offset += (size_t)coordinate * stride; + } + stride *= input_shape[axis]; + } +$combine""") + +_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, +$parameters$extra) +{ + size_t plane, position, tap; + for (plane = 0; plane < plane_count; ++plane) { + const $element* values = in + plane * input_size; + $element* result = out + plane * output_size; + for (position = 0; position < output_size; ++position) { + $state + for (tap = 0; tap < window_size; ++tap) { +$walk + } + result[position] = $finish; + } + } +}""") + +# The indexed max pooling reports where each maximum was read, as a flat index into the whole +# operand. `index_strides` is what `storage_order` chooses: the operand's own row-major +# strides, or the column-major ones ONNX defines the other order as. +_INDEXED_TEMPLATE = Template("""\ +static void $name( + $element* out, + int64_t* indices, + const $element* in, +$parameters, + const size_t* index_strides) +{ + size_t plane, position, tap; + (void)counted_shape; + for (plane = 0; plane < plane_count; ++plane) { + const $element* values = in + plane * input_size; + $element* result = out + plane * output_size; + int64_t* chosen = indices + plane * output_size; + for (position = 0; position < output_size; ++position) { + $element best = $identity; + size_t found = 0; + int seen = 0; + for (tap = 0; tap < window_size; ++tap) { + size_t remaining_position = position; + size_t remaining_tap = tap; + size_t offset = 0; + size_t reported = 0; + size_t stride = 1; + int inside = 1; + int axis; + for (axis = spatial_rank - 1; axis >= 0; --axis) { + const ptrdiff_t coordinate = + (ptrdiff_t)(remaining_position % output_shape[axis]) + * (ptrdiff_t)strides[axis] + + (ptrdiff_t)(remaining_tap % window_shape[axis]) + * (ptrdiff_t)dilations[axis] + - pads[axis]; + remaining_position /= output_shape[axis]; + remaining_tap /= window_shape[axis]; + if (coordinate < 0 || coordinate >= (ptrdiff_t)input_shape[axis]) { + inside = 0; + } else { + offset += (size_t)coordinate * stride; + reported += (size_t)coordinate * index_strides[axis]; + } + stride *= input_shape[axis]; + } + if (inside) { + const $element x = values[offset]; + if (!seen || ($better)) { + best = x; + found = reported; + seen = 1; + } + } + } + result[position] = best; + chosen[position] = (int64_t)(plane * input_size + found); + } + } +}""") + +_UNPOOL_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const int64_t* indices, + size_t count, + size_t limit) +{ + size_t index; + for (index = 0; index < limit; ++index) { + out[index] = $zero; + } + for (index = 0; index < count; ++index) { + const int64_t chosen = indices[index]; + /* ONNX leaves an index outside the result undefined; the artifact reports it + rather than writing past the buffer. */ + if (chosen < 0 || (size_t)chosen >= limit) { + return 1; + } + out[(size_t)chosen] = in[index]; + } + return 0; +}""") + +# Every pooling arrived at opset 1 and was revised repeatedly since — AveragePool gained +# `ceil_mode` at 10 and `dilations` at 19, MaxPool `Indices` at 8 and `dilations` at 10 — and +# 22 widened them all to bfloat16. Only 22 is claimed: it is the revision the reference +# evaluator is version-faithful for and the one every pooling test in the backend corpus +# imports, so it is the only one anything can vouch for. A model importing an older one gets +# the unsupported-version error. +_VERSIONS = (22,) + + +@dataclass(frozen=True) +class _Fold: + """How one pooling folds a window, as C over a tap's `inside` and `counted` flags. + + `state` opens the fold, `combine` runs per tap and `finish` is what the position is + written from; `parameters` are the kernel parameters only this fold takes and + `arguments` the literals its call sites pass for them. All three expressions are + `$`-templated over the element type. + """ + + name: str + state: str + combine: str + finish: str + parameters: tuple[str, ...] = () + arguments: tuple[str, ...] = () + helpers: tuple[CFunction, ...] = () + + +# What a pooling's fold is built from: the node, and the element type it folds at. +_Recipe = Callable[[NodeContext, int], _Fold] + + +@dataclass(frozen=True) +class _Geometry: + """A pooling's shape, resolved to the literals the kernel walks it with.""" + + batch_count: int + channels: int + input_shape: tuple[int, ...] + output_shape: tuple[int, ...] + window_shape: tuple[int, ...] + strides: tuple[int, ...] + dilations: tuple[int, ...] + pads: tuple[int, ...] + counted_shape: tuple[int, ...] + + @property + def result_shape(self) -> tuple[int, ...]: + return (self.batch_count, self.channels, *self.output_shape) + + @property + def arguments(self) -> list[str]: + """Call-site literals for the geometry parameters a pooling kernel takes.""" + return [ + f"{self.batch_count * self.channels}u", + f"{math.prod(self.input_shape)}u", + f"{math.prod(self.output_shape)}u", + f"{math.prod(self.window_shape)}u", + str(len(self.output_shape)), + extents(self.input_shape), + extents(self.counted_shape), + extents(self.output_shape), + extents(self.window_shape), + extents(self.strides), + extents(self.dilations), + offsets(self.pads), + ] + + def index_strides(self, storage_order: int) -> tuple[int, ...]: + """Strides turning a spatial coordinate into the index `storage_order` asks for.""" + if storage_order == 0: + return row_major_strides(self.input_shape) + return row_major_strides(self.input_shape[::-1])[::-1] + + +# -------------------------------------------------------------------------------------- +# The folds +# -------------------------------------------------------------------------------------- + + +def _average_fold(context: NodeContext, elem_type: int) -> _Fold: + """The mean over the taps a window covers, over as many of them as it counts. + + `count_include_pad` decides whether the padded positions are part of that count; the + value they contribute is zero either way, so only the divisor changes. A window counting + nothing at all divides by zero, which IEEE defines for the float families this op is + defined over. + + The default is stated here rather than read off the schema because `GlobalAveragePool` + folds through this too and has no such attribute — it pads nothing, so every tap it + covers is counted whichever way the flag would go. + """ + return _Fold( + name="average", + state="$element total = $zero; size_t inside_count = 0, counted_count = 0;", + combine="""\ + if (counted) { + ++counted_count; + if (inside) { + total += values[offset]; + ++inside_count; + } + }""", + finish="total / ($element)(include_pad ? counted_count : inside_count)", + parameters=("size_t include_pad",), + arguments=(f"{int(context.attribute('count_include_pad', 0) != 0)}u",), + ) + + +def _max_fold(context: NodeContext, elem_type: int) -> _Fold: + """The largest tap inside the operand; a padded position is not a candidate for it. + + ONNX says nothing about a NaN in the window, and the reference evaluator's two pooling + paths do not agree on one, so the fold takes numpy's `maximum`, as ReduceMax does. + """ + largest = combiner(context, elem_type, largest=True) + return _Fold( + name="max", + state=f"$element best = {extremum_identity(elem_type, largest=True)};", + combine=f"""\ + (void)counted; + if (inside) {{ + best = {largest.name}(best, values[offset]); + }}""", + finish="best", + helpers=(largest,), + ) + + +def _lp_fold(context: NodeContext, elem_type: int) -> _Fold: + """The Lp norm of the taps a window covers: a padded position contributes `|0|^p`. + + The norm itself, over however many taps the window turns out to hold. ONNX's reference + evaluator computes a window that `ceil_mode` clipped differently — it averages and scales + the result back up by the whole kernel's tap count, which its own source records as a + borrowed computation that differs from the spec's — where the backend corpus's stored + outputs and onnxruntime both take the plain norm, as this does. + """ + return _Fold( + name="lp", + state="$element total = $zero;", + combine="""\ + (void)counted; + if (inside) { + total += pow$f(fabs$f(values[offset]), ($element)order); + }""", + finish="pow$f(total, $one / ($element)order)", + parameters=("int order",), + arguments=(str(_lp_order(context)),), + ) + + +def _lp_order(context: NodeContext) -> int: + order = context.int_attribute("p") + if order < 1: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` takes the norm of order " + f"{order}; ONNX defines `p` as the order of an Lp norm, which is positive." + ) + return order + + +# -------------------------------------------------------------------------------------- +# Emitting a pooling +# -------------------------------------------------------------------------------------- + + +def _pool(context: NodeContext, recipe: _Recipe, geometry: _Geometry) -> NodeEmission: + source = context.require_input(0) + result = context.require_output(0) + fold = recipe(context, result.elem_type) + verify_shape(context, result, geometry.result_shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + name = f"{context.prefix}_pool_{fold.name}_{element}" + definition = _TEMPLATE.substitute( + name=name, + element=element, + parameters=_GEOMETRY_PARAMETERS, + extra="".join(f",\n {parameter}" for parameter in fold.parameters), + state=expand(fold.state, result.elem_type), + walk=_WALK.substitute(combine=expand(fold.combine, result.elem_type)), + finish=expand(fold.finish, result.elem_type), + ) + call = call_kernel( + name, [result.expr, source.expr, *geometry.arguments, *fold.arguments] + ) + return NodeEmission( + functions=(*fold.helpers, CFunction(name, definition)), statements=(call,) + ) + + +def _pooling(recipe: _Recipe) -> Callable[[NodeContext], NodeEmission]: + """A pooling over the window the node's own attributes describe.""" + return lambda context: _pool(context, recipe, _geometry(context)) + + +def _global_pooling(recipe: _Recipe) -> Callable[[NodeContext], NodeEmission]: + """A pooling over one window covering every spatial position of the operand.""" + return lambda context: _pool(context, recipe, _global_geometry(context)) + + +def _max_pool(context: NodeContext) -> NodeEmission: + """MaxPool, which also reports where each maximum came from when asked to.""" + geometry = _geometry(context) + indices = context.outputs[1] if len(context.outputs) > 1 else None + if indices is None: + return _pool(context, _max_fold, geometry) + + source = context.require_input(0) + result = context.require_output(0) + verify_shape(context, result, geometry.result_shape) + verify_shape(context, indices, geometry.result_shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + name = f"{context.prefix}_pool_max_indexed_{element}" + definition = _INDEXED_TEMPLATE.substitute( + name=name, + element=element, + parameters=_GEOMETRY_PARAMETERS, + identity=extremum_identity(result.elem_type, largest=True), + better=extremum_test(result.elem_type, largest=True, last=False), + ) + call = call_kernel( + name, + [ + result.expr, + indices.expr, + source.expr, + *geometry.arguments, + extents(geometry.index_strides(context.int_attribute("storage_order"))), + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _max_unpool(context: NodeContext) -> NodeEmission: + """MaxUnpool: every element written back where the pooling that chose it read it.""" + source = context.require_input(0) + positions = context.require_input(1) + result = context.require_output(0) + if positions.shape != source.shape: + raise CompileError( + f"Node `{context.label}`: `MaxUnpool` takes one index per value, but " + f"`{positions.name}` has shape {list(positions.shape)} against " + f"`{source.name}`'s {list(source.shape)}." + ) + verify_shape(context, result, _unpooled_shape(context, source)) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + name = f"{context.prefix}_maxunpool_{element}" + definition = _UNPOOL_TEMPLATE.substitute( + name=name, element=element, zero=scalar_literal(0, result.elem_type) + ) + call = checked_call( + context, + name, + [ + result.expr, + source.expr, + positions.expr, + f"{source.elem_count}u", + f"{result.elem_count}u", + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +# -------------------------------------------------------------------------------------- +# Reading the geometry off the node +# -------------------------------------------------------------------------------------- + + +def _geometry(context: NodeContext) -> _Geometry: + """The window a pooling node states, resolved to concrete extents and pads.""" + source = context.require_input(0) + rank = _spatial_rank(context, source) + window_shape = _kernel_shape(context, rank) + strides = spatial_attribute(context, "strides", rank, 1) + dilations = spatial_attribute(context, "dilations", rank, 1) + begins, ends = resolve_pads( + context, source.shape[2:], window_shape, dilations, strides + ) + return _Geometry( + batch_count=source.shape[0], + channels=source.shape[1], + input_shape=source.shape[2:], + output_shape=_window_positions( + context, source.shape[2:], window_shape, dilations, strides, begins, ends + ), + window_shape=window_shape, + strides=strides, + dilations=dilations, + pads=begins, + counted_shape=tuple( + extent + end for extent, end in zip(source.shape[2:], ends) + ), + ) + + +def _global_geometry(context: NodeContext) -> _Geometry: + """A `Global*` pooling: one window covering every spatial position of the operand.""" + source = context.require_input(0) + rank = _spatial_rank(context, source) + return _Geometry( + batch_count=source.shape[0], + channels=source.shape[1], + input_shape=source.shape[2:], + output_shape=(1,) * rank, + window_shape=source.shape[2:], + strides=(1,) * rank, + dilations=(1,) * rank, + pads=(0,) * rank, + counted_shape=source.shape[2:], + ) + + +def _spatial_rank(context: NodeContext, source: TensorRef) -> int: + if len(source.shape) < 3: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` pools a batch of " + f"multi-channel signals — a tensor of rank 3 or more — but `{source.name}` " + f"has shape {list(source.shape)}." + ) + return len(source.shape) - 2 + + +def _kernel_shape(context: NodeContext, rank: int) -> tuple[int, ...]: + window_shape = spatial_extents(context, "kernel_shape", rank) + if window_shape is None: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` states no `kernel_shape`, " + "which ONNX defines as a required attribute." + ) + return window_shape + + +def _window_positions( + context: NodeContext, + input_shape: Sequence[int], + window_shape: Sequence[int], + dilations: Sequence[int], + strides: Sequence[int], + begins: Sequence[int], + ends: Sequence[int], +) -> tuple[int, ...]: + """How many windows fit each axis, under `ceil_mode`'s rounding. + + Rounding up lets the last window reach past the operand's padded end, where the taps + beyond the pads are simply not read — unless it would *start* past the padding, in which + case the window would be nothing but padding and ONNX drops the position instead. + """ + ceil_mode = context.int_attribute("ceil_mode") != 0 + positions = [] + for extent, window, dilation, stride, begin, end in zip( + input_shape, window_shape, dilations, strides, begins, ends + ): + reach = extent + begin + end - (window - 1) * dilation - 1 + if not ceil_mode: + positions.append(reach // stride + 1) + continue + count = -(-reach // stride) + 1 + positions.append(count - 1 if (count - 1) * stride >= extent + begin else count) + return tuple(positions) + + +def _unpooled_shape(context: NodeContext, source: TensorRef) -> tuple[int, ...]: + """The extent the max pooling that produced this operand read, which is what is filled. + + ONNX also takes the result's shape from an optional `output_shape` operand. One the graph + does not fix makes that shape depend on input data, which the frontend rejects before any + kernel is reached; for one it does fix, ONNX's own shape inference derives nothing at all, + so there is no shape to compile against either way. + """ + rank = _spatial_rank(context, source) + window_shape = _kernel_shape(context, rank) + strides = spatial_attribute(context, "strides", rank, 1) + begins, ends = declared_pads(context, rank, auto_pad_mode(context)) or ( + (0,) * rank, + (0,) * rank, + ) + return ( + *source.shape[:2], + *( + (extent - 1) * stride - begin - end + window + for extent, stride, begin, end, window in zip( + source.shape[2:], strides, begins, ends, window_shape + ) + ), + ) + + +register_kernel("", "AveragePool", _VERSIONS, _pooling(_average_fold)) +register_kernel("", "LpPool", _VERSIONS, _pooling(_lp_fold)) +register_kernel("", "MaxPool", _VERSIONS, _max_pool) +register_kernel("", "GlobalAveragePool", _VERSIONS, _global_pooling(_average_fold)) +register_kernel("", "GlobalLpPool", _VERSIONS, _global_pooling(_lp_fold)) +register_kernel("", "GlobalMaxPool", _VERSIONS, _global_pooling(_max_fold)) +register_kernel("", "MaxUnpool", _VERSIONS, _max_unpool) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/quantize.py b/src/python/fnnx/extras/compilers/c/onnx/ops/quantize.py new file mode 100644 index 0000000..ff9b22b --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/quantize.py @@ -0,0 +1,974 @@ +"""Quantization: the affine grid a low-precision tensor stands on, and the ops over it. + +A quantized tensor is an integer grid standing for the reals `(value - zero_point) * scale`. +`DequantizeLinear` applies that map and `QuantizeLinear` inverts it — dividing by the scale, +rounding halves to even, and saturating to the target type's range, which is where the +information a low-precision tensor cannot carry is actually lost. Both read their scale and +zero point at one of the three granularities ONNX defines, and all three come to the same +addressing: a stride per axis into those operands' buffers, plus a divisor on the axis a +blocked scale is repeated along. + +The other four ops fold the map into a product. `MatMulInteger` and `ConvInteger` subtract the +zero points and accumulate in `int32`, leaving the result on a grid whose scale is the product +of the operands' — which is why they take no scale at all. `QLinearConv` and `QLinearMatMul` +run the same accumulation and then requantize onto a grid of their own, by the one factor +`a_scale * b_scale / y_scale` that product comes to. So the walk over the operands is the +convolution's and the matrix product's, taken from the kernels that run them unquantized, and +only the accumulation type, the zero-point offsets and the store differ. + +What is not served: a scale or zero point per row or per column of a matrix product. In the +form ONNX's own text describes it — an `M`-element vector against an `[M, K]` operand — its +reference evaluator stretches that vector along numpy's trailing axis instead, so no oracle +says what a kernel should compute there; the products are served at per-tensor granularity +alone rather than read one way in that form and another in the `[M, 1]` one. The +convolutions' per-output-channel `w_scale` and `w_zero_point` are served: on those the +evaluator and the backend corpus agree. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from string import Template + +import numpy as np +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + FLOAT_TYPES, + c_type, + element_type_name, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + broadcast_strides, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + kernel_name, + normalize_axis, + row_major_strides, + verify_same_shape, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents, math_suffix +from fnnx.extras.compilers.c.onnx.ops.conv import ( + WINDOW_PARAMETERS, + convolution_geometry, + verify_bias, +) +from fnnx.extras.compilers.c.onnx.ops.gemm import PRODUCT_PARAMETERS, matrix_product + +# The grids ONNX quantizes onto, of the ones the compiler's element types cover: a +# `QuantizeLinear` saturates to the range of one of these, and every quantized product reads +# its operands from one. `DequantizeLinear` reads `int32` as well — where an accumulated bias +# sits, which nothing ever quantizes *to*. +_GRID_TYPES = ( + TensorProto.INT8, + TensorProto.UINT8, + TensorProto.INT16, + TensorProto.UINT16, +) + +# QuantizeLinear arrived at 10, gained per-axis quantization at 13, float8 and `saturate` at +# 19, blocked quantization and the 4-bit types at 21, `precision` and float4 at 23, and more +# types at 24 and 25. Every revision but 13 is claimed: the reference evaluator is +# version-faithful for each of those — and the corpus's own tests import 11, which selects +# 10, and 25 — while nothing can vouch for 13, whose semantics the evaluator does not +# distinguish and which no corpus test imports. DequantizeLinear's history runs alongside it, +# vouched for from 19 on. +_QUANTIZE_VERSIONS = (10, 19, 21, 23, 24, 25) +_DEQUANTIZE_VERSIONS = (19, 21, 23, 24, 25) + +# QLinearMatMul is claimed at 21, the revision that widened its scales to a type parameter and +# the one the evaluator distinguishes; the integer products and QLinearConv have had one +# revision each, the one they arrived at. +_QLINEAR_MATMUL_VERSIONS = (21,) +_INTEGER_VERSIONS = (10,) + + +# -------------------------------------------------------------------------------------- +# Rounding onto a grid +# -------------------------------------------------------------------------------------- + +_SATURATE_TEMPLATE = Template("""\ +static $result $name(double value) +{ + /* The nearest integer, halves to even, saturated to the grid's range. A value that is + not a number has no nearest integer at all: ONNX leaves it undefined, and it lands at + the low end, which is at least the same end the reference evaluator's own conversion + leaves it at. */ + if (!(value > $low)) { + return ($result)$low; + } + if (value > $high) { + return ($result)$high; + } + return ($result)rint(value); +}""") + + +def _saturating_cast(context: NodeContext, elem_type: int) -> CFunction: + """The rounding-and-saturating store every op that writes a quantized grid ends in.""" + info = np.iinfo(np.dtype(numpy_dtype_name(elem_type))) + element = c_type(elem_type) + name = f"{context.prefix}_saturate_{element}" + return CFunction( + name, + _SATURATE_TEMPLATE.substitute( + name=name, + result=element, + low=f"{int(info.min)}.0", + high=f"{int(info.max)}.0", + ), + ) + + +def _verify_grid_type( + context: NodeContext, + operand: TensorRef, + role: str, + allowed: tuple[int, ...] = _GRID_TYPES, +) -> None: + """Refuse a tensor the op reads or writes as quantized whose type is no integer grid. + + A quantized tensor's saturation range is its own type's, so a kernel emitted for a type + ONNX does not quantize onto would round onto a grid that is not there — and a floating + one would be truncated into the accumulation without a word. + """ + if operand.elem_type not in allowed: + names = ", ".join(element_type_name(elem_type) for elem_type in allowed) + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` quantizes onto the integer " + f"grids {names}, but its {role} `{operand.name}` is " + f"`{element_type_name(operand.elem_type)}`." + ) + + +def _verify_zero_point( + context: NodeContext, zero_point: TensorRef | None, grid: TensorRef +) -> None: + """Refuse a zero point that does not sit on the same grid as the tensor it shifts.""" + if zero_point is not None and zero_point.elem_type != grid.elem_type: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads `{zero_point.name}` " + f"as `{element_type_name(zero_point.elem_type)}` against a quantized tensor of " + f"`{element_type_name(grid.elem_type)}`; ONNX defines the two as one type." + ) + + +# -------------------------------------------------------------------------------------- +# The affine map: QuantizeLinear and DequantizeLinear +# -------------------------------------------------------------------------------------- + +_AFFINE_TEMPLATE = Template("""\ +static void $name( + $result* out, + const $source* in, + const $scale* scale, + const $zero* zero_point, + size_t count, + int rank, + const size_t* shape, + int block_axis, + size_t block_size, + const size_t* scale_strides, + const size_t* zero_strides) +{ + size_t index; + for (index = 0; index < count; ++index) { + size_t remainder = index; + size_t scale_offset = 0; + size_t zero_offset = 0; + int axis; + for (axis = rank - 1; axis >= 0; --axis) { + size_t coordinate = remainder % shape[axis]; + remainder /= shape[axis]; + /* A blocked scale holds one value per `block_size` elements of this axis: the + repetition ONNX defines it by, read backwards. */ + if (axis == block_axis) { + coordinate /= block_size; + } + scale_offset += coordinate * scale_strides[axis]; + zero_offset += coordinate * zero_strides[axis]; + } +$body + } +}""") + +_QUANTIZE_BODY = Template("""\ + { + const $compute quotient = + ($compute)in[index] / ($compute)scale[scale_offset]; + const double zero = + (zero_point != NULL) ? (double)zero_point[zero_offset] : 0.0; + /* The quotient is rounded before the zero point shifts it, as ONNX rounds the + division alone: shifting first would send a half to the other neighbour. */ + out[index] = $saturate((double)$round(quotient) + zero); + }""") + +_DEQUANTIZE_BODY = Template("""\ + { + const $compute zero = (zero_point != NULL) + ? ($compute)zero_point[zero_offset] + : ($compute)0; + /* The grid is read at single precision whatever its width, which is where the + reference evaluator converts it too. */ + out[index] = ($result)(((float)in[index] - zero) + * ($compute)scale[scale_offset]); + }""") + + +@dataclass(frozen=True) +class _Granularity: + """How a scale and a zero point are addressed while the data's own shape is walked.""" + + shape: tuple[int, ...] + scale_strides: tuple[int, ...] + zero_strides: tuple[int, ...] + block_axis: int + block_size: int + + @property + def arguments(self) -> list[str]: + """Call-site literals for the addressing parameters an affine map's kernel takes.""" + return [ + f"{math.prod(self.shape)}u", + str(len(self.shape)), + extents(self.shape), + str(self.block_axis), + f"{self.block_size}u", + extents(self.scale_strides), + extents(self.zero_strides), + ] + + +def _quantize_linear(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + scale = context.require_input(1) + zero_point = context.optional_input(2) + result = context.require_output(0) + verify_same_shape(context, source, result) + _verify_grid_type(context, result, "output") + _verify_zero_point(context, zero_point, result) + granularity = _granularity(context, source, scale, zero_point) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + compute = _division_type(context, source, scale) + saturate = _saturating_cast(context, result.elem_type) + return _affine_emission( + context, + source=source, + scale=scale, + zero_point=zero_point, + grid=result, + result=result, + compute=compute, + granularity=granularity, + body=_QUANTIZE_BODY.substitute( + compute=c_type(compute), + round=f"rint{math_suffix(compute)}", + saturate=saturate.name, + ), + helpers=(saturate,), + ) + + +def _dequantize_linear(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + scale = context.require_input(1) + zero_point = context.optional_input(2) + result = context.require_output(0) + verify_same_shape(context, source, result) + # `int32` alone is read back from a grid without ever being quantized onto one: it is + # where an accumulated bias sits, whose scale is the product of the ones it is added to. + _verify_grid_type(context, source, "input", (*_GRID_TYPES, TensorProto.INT32)) + if result.elem_type not in FLOAT_TYPES: + raise CompileError( + f"Node `{context.label}`: `DequantizeLinear` reads a grid back as the reals it " + f"stands for, but its output `{result.name}` is " + f"`{element_type_name(result.elem_type)}`." + ) + _verify_zero_point(context, zero_point, source) + granularity = _granularity(context, source, scale, zero_point) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + compute = _difference_type(source, zero_point) + return _affine_emission( + context, + source=source, + scale=scale, + zero_point=zero_point, + grid=source, + result=result, + compute=compute, + granularity=granularity, + body=_DEQUANTIZE_BODY.substitute( + compute=c_type(compute), result=c_type(result.elem_type) + ), + ) + + +def _affine_emission( + context: NodeContext, + *, + source: TensorRef, + scale: TensorRef, + zero_point: TensorRef | None, + grid: TensorRef, + result: TensorRef, + compute: int, + granularity: _Granularity, + body: str, + helpers: tuple[CFunction, ...] = (), +) -> NodeEmission: + """The kernel and call site both affine maps share; only `body` differs between them. + + `grid` is the tensor the zero point sits on -- the result of a quantization and the + operand of a dequantization -- whose type the parameter keeps even where the node omits + the operand and the kernel is passed nothing at all. + """ + name = kernel_name( + context, + c_type(source.elem_type), + c_type(scale.elem_type), + c_type(result.elem_type), + c_type(compute), + *(("zp",) if zero_point is not None else ()), + ) + definition = _AFFINE_TEMPLATE.substitute( + name=name, + result=c_type(result.elem_type), + source=c_type(source.elem_type), + scale=c_type(scale.elem_type), + zero=c_type(grid.elem_type), + body=body, + ) + call = call_kernel( + name, + [ + result.expr, + source.expr, + scale.expr, + "NULL" if zero_point is None else zero_point.expr, + *granularity.arguments, + ], + ) + return NodeEmission( + functions=(*helpers, CFunction(name, definition)), statements=(call,) + ) + + +def _division_type(context: NodeContext, source: TensorRef, scale: TensorRef) -> int: + """The type `x / y_scale` is computed at. + + ONNX reads it off `y_scale`'s type, while its reference evaluator divides the two arrays + as numpy does — which promotes an `int32` operand against a `float32` one to `float64`. + The evaluator is what both suites compare against, so its promotion is what is emitted; + the two agree wherever both operands are floats, which is every model that quantizes one. + """ + precision = int(context.attribute("precision", 0)) + if precision: + raise CompileError( + f"Node `{context.label}`: `QuantizeLinear` states `precision` " + f"`{element_type_name(precision)}`, which the C compiler does not serve: the " + "newest revision ONNX's reference evaluator implements predates the attribute " + "and refuses a node carrying it, so nothing can vouch for what a kernel " + "dividing at that precision should produce. Drop the attribute to divide at " + "the type `y_scale` carries." + ) + single = {source.elem_type, scale.elem_type} == {TensorProto.FLOAT} + return TensorProto.FLOAT if single else TensorProto.DOUBLE + + +def _difference_type(source: TensorRef, zero_point: TensorRef | None) -> int: + """The type `x - x_zero_point` is computed at. + + The reference converts the grid to `float32` first and subtracts the zero point from + that, which numpy promotes to `float64` for an `int32` zero point and nothing narrower. + """ + if zero_point is None or source.elem_type != TensorProto.INT32: + return TensorProto.FLOAT + return TensorProto.DOUBLE + + +def _granularity( + context: NodeContext, + source: TensorRef, + scale: TensorRef, + zero_point: TensorRef | None, +) -> _Granularity: + block_size = int(context.attribute("block_size", 0)) + rank = len(source.shape) + scale_strides, blocked = _grid_strides(context, source, scale, block_size) + if ( + zero_point is not None + and zero_point.elem_count != 1 + and zero_point.shape != scale.shape + ): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads `{zero_point.name}` " + f"of shape {list(zero_point.shape)} against a scale of {list(scale.shape)}; " + "ONNX defines the two as one shape, which is what fixes the granularity." + ) + zero_strides = ( + (0,) * rank + if zero_point is None or zero_point.elem_count == 1 + else _grid_strides(context, source, zero_point, block_size)[0] + ) + return _Granularity( + shape=source.shape, + scale_strides=scale_strides, + zero_strides=zero_strides, + block_axis=_quantization_axis(context, source) if blocked else -1, + block_size=block_size if blocked else 1, + ) + + +def _grid_strides( + context: NodeContext, source: TensorRef, operand: TensorRef, block_size: int +) -> tuple[tuple[int, ...], bool]: + """Strides addressing `operand` as the data's coordinates are walked, and if it blocks. + + The three granularities are one addressing: a single-element scale is read at stride zero + on every axis, a per-axis vector at the stride its one axis carries, and a blocked tensor + at its own row-major strides — with the coordinate on the quantization axis divided by + the block size, which is the repetition ONNX defines blocking by. + """ + rank = len(source.shape) + if operand.elem_count == 1: + return (0,) * rank, False + axis = _quantization_axis(context, source) + if not block_size: + if len(operand.shape) != 1: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads `{operand.name}` " + f"of shape {list(operand.shape)}; with no `block_size`, ONNX defines a " + "scale of more than one element as one value per slice along the " + "quantization axis — a 1-D tensor." + ) + spread = tuple( + operand.shape[0] if position == axis else 1 for position in range(rank) + ) + strides = broadcast_strides( + replace(operand, shape=spread), source.shape, node_label=context.label + ) + return strides, False + _verify_blocks(context, source, operand, axis, block_size) + return row_major_strides(operand.shape), True + + +def _quantization_axis(context: NodeContext, source: TensorRef) -> int: + """The axis a per-axis or blocked scale runs along. + + `QuantizeLinear`-10 predates per-axis quantization and declares no `axis` attribute at + all; ONNX's reference reads the absent one as 1, which is the default every revision + that does declare it carries. + """ + return normalize_axis(context, int(context.attribute("axis", 1)), len(source.shape)) + + +def _verify_blocks( + context: NodeContext, + source: TensorRef, + operand: TensorRef, + axis: int, + block_size: int, +) -> None: + """Refuse a blocked scale the emitted addressing would read outside of. + + A blocked scale carries the data's own shape but for the quantization axis, where it + holds one value per block of `block_size` elements — so its extent there is what ONNX + states it: the block count the data's own extent comes to. + """ + if block_size < 1: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` states `block_size` " + f"{block_size}; ONNX defines it as a positive count of elements." + ) + blocks = -(-source.shape[axis] // block_size) + expected = tuple( + blocks if position == axis else extent + for position, extent in enumerate(source.shape) + ) + if operand.shape != expected: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` blocks `{source.name}` of " + f"shape {list(source.shape)} by {block_size} along axis {axis}, which takes a " + f"scale of {list(expected)}, but `{operand.name}` has shape " + f"{list(operand.shape)}." + ) + + +register_kernel("", "QuantizeLinear", _QUANTIZE_VERSIONS, _quantize_linear) +register_kernel("", "DequantizeLinear", _DEQUANTIZE_VERSIONS, _dequantize_linear) + + +# -------------------------------------------------------------------------------------- +# The quantized matrix products +# -------------------------------------------------------------------------------------- + +_PRODUCT_TEMPLATE = Template("""\ +static void $name( +$parameters) +{ +$locals + size_t batch, row, column, index; + for (batch = 0; batch < batch_count; ++batch) { + size_t left_base = 0; + size_t right_base = 0; + size_t remainder = batch; + int axis; + for (axis = batch_rank - 1; axis >= 0; --axis) { + const size_t coordinate = remainder % batch_shape[axis]; + remainder /= batch_shape[axis]; + left_base += coordinate * left_batch_strides[axis]; + right_base += coordinate * right_batch_strides[axis]; + } + for (row = 0; row < rows; ++row) { + for (column = 0; column < columns; ++column) { + int32_t sum = 0; + for (index = 0; index < inner; ++index) { + sum += ((int32_t)left[left_base + row * inner + index] - left_zero) + * ((int32_t)right[right_base + index * columns + column] + - right_zero); + } +$store + } + } + } +}""") + +_PRODUCT_STORE = " out[(batch * rows + row) * columns + column] =" + +_ZERO_LOCALS = Template("""\ + const int32_t left_zero = + (left_zero_point != NULL) ? (int32_t)left_zero_point[0] : 0; + const int32_t right_zero = + (right_zero_point != NULL) ? (int32_t)right_zero_point[0] : 0;""") + + +def _matmul_integer(context: NodeContext) -> NodeEmission: + left = context.require_input(0) + right = context.require_input(1) + left_zero = context.optional_input(2) + right_zero = context.optional_input(3) + result = context.require_output(0) + product = matrix_product(context, left, right) + verify_shape(context, result, product.result_shape) + _verify_per_tensor(context, left_zero, "zero point") + _verify_per_tensor(context, right_zero, "zero point") + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + name = kernel_name(context, c_type(left.elem_type), c_type(right.elem_type)) + definition = _PRODUCT_TEMPLATE.substitute( + name=name, + parameters=",\n".join( + [ + f" {c_type(result.elem_type)}* out", + *_product_operands(left, right), + PRODUCT_PARAMETERS, + ] + ), + locals=_ZERO_LOCALS.template, + store=f"{_PRODUCT_STORE} sum;", + ) + call = call_kernel( + name, + [ + result.expr, + left.expr, + right.expr, + _pointer(left_zero), + _pointer(right_zero), + *product.arguments, + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _qlinear_matmul(context: NodeContext) -> NodeEmission: + left = context.require_input(0) + left_scale = context.require_input(1) + left_zero = context.require_input(2) + right = context.require_input(3) + right_scale = context.require_input(4) + right_zero = context.require_input(5) + result_scale = context.require_input(6) + result_zero = context.require_input(7) + result = context.require_output(0) + product = matrix_product(context, left, right) + verify_shape(context, result, product.result_shape) + for operand in (left, right, result): + _verify_grid_type(context, operand, "operand") + _verify_zero_point(context, left_zero, left) + _verify_zero_point(context, right_zero, right) + _verify_zero_point(context, result_zero, result) + for operand, role in ( + (left_scale, "scale"), + (left_zero, "zero point"), + (right_scale, "scale"), + (right_zero, "zero point"), + (result_scale, "scale"), + (result_zero, "zero point"), + ): + _verify_per_tensor(context, operand, role) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + saturate = _saturating_cast(context, result.elem_type) + name = kernel_name( + context, + c_type(left.elem_type), + c_type(right.elem_type), + c_type(result.elem_type), + c_type(result_scale.elem_type), + ) + definition = _PRODUCT_TEMPLATE.substitute( + name=name, + parameters=",\n".join( + [ + f" {c_type(result.elem_type)}* out", + *_product_operands(left, right), + f" const {c_type(left_scale.elem_type)}* left_scale", + f" const {c_type(right_scale.elem_type)}* right_scale", + f" const {c_type(result_scale.elem_type)}* result_scale", + f" const {c_type(result.elem_type)}* result_zero_point", + PRODUCT_PARAMETERS, + ] + ), + locals="\n".join( + [ + _ZERO_LOCALS.template, + # The product of two grids stands on the product of their scales, so + # requantizing onto a third is this one factor, taken at the scales' own + # precision as ONNX's reference takes it. + f" const {c_type(result_scale.elem_type)} factor =", + " left_scale[0] * right_scale[0] / result_scale[0];", + " const double result_zero = (double)result_zero_point[0];", + ] + ), + store=( + f"{_PRODUCT_STORE}\n " + f"{saturate.name}((double)sum * (double)factor + result_zero);" + ), + ) + call = call_kernel( + name, + [ + result.expr, + left.expr, + right.expr, + left_zero.expr, + right_zero.expr, + left_scale.expr, + right_scale.expr, + result_scale.expr, + result_zero.expr, + *product.arguments, + ], + ) + return NodeEmission( + functions=(saturate, CFunction(name, definition)), statements=(call,) + ) + + +def _product_operands(left: TensorRef, right: TensorRef) -> list[str]: + """The operands and zero points both quantized products read, in call order.""" + return [ + f" const {c_type(left.elem_type)}* left", + f" const {c_type(right.elem_type)}* right", + f" const {c_type(left.elem_type)}* left_zero_point", + f" const {c_type(right.elem_type)}* right_zero_point", + ] + + +register_kernel("", "MatMulInteger", _INTEGER_VERSIONS, _matmul_integer) +register_kernel("", "QLinearMatMul", _QLINEAR_MATMUL_VERSIONS, _qlinear_matmul) + + +# -------------------------------------------------------------------------------------- +# The quantized convolutions +# -------------------------------------------------------------------------------------- + +_CONVOLUTION_TEMPLATE = Template("""\ +static void $name( +$parameters) +{ +$locals + size_t batch, group, filter, position, tap, channel; + for (batch = 0; batch < batch_count; ++batch) { + for (group = 0; group < groups; ++group) { + const $source* plane = + in + (batch * groups + group) * group_channels * input_size; + for (filter = 0; filter < group_filters; ++filter) { + const size_t channel_index = group * group_filters + filter; + const $weight* window = + weights + channel_index * group_channels * window_size; + $result* result = + out + (batch * groups * group_filters + channel_index) * output_size; +$channel + for (position = 0; position < output_size; ++position) { + int32_t sum = $initial; + for (tap = 0; tap < window_size; ++tap) { + size_t remaining_position = position; + size_t remaining_tap = tap; + size_t offset = 0; + size_t stride = 1; + int inside = 1; + int axis; + for (axis = spatial_rank - 1; axis >= 0; --axis) { + const ptrdiff_t coordinate = + (ptrdiff_t)(remaining_position % output_shape[axis]) + * (ptrdiff_t)strides[axis] + + (ptrdiff_t)(remaining_tap % window_shape[axis]) + * (ptrdiff_t)dilations[axis] + - pads[axis]; + remaining_position /= output_shape[axis]; + remaining_tap /= window_shape[axis]; + if (coordinate < 0 + || coordinate >= (ptrdiff_t)input_shape[axis]) { + inside = 0; + } else { + offset += (size_t)coordinate * stride; + } + stride *= input_shape[axis]; + } + if (inside) { + for (channel = 0; channel < group_channels; ++channel) { + sum += ((int32_t)plane[channel * input_size + offset] + - input_zero) + * ((int32_t)window[channel * window_size + tap] + - filter_zero); + } + } + } +$store + } + } + } + } +}""") + +_INPUT_ZERO_LOCAL = """\ + const int32_t input_zero = + (input_zero_point != NULL) ? (int32_t)input_zero_point[0] : 0;""" + +# The filter's zero point carries one value per output channel or one for the whole filter, +# which is a stride of one or of zero; either way it is read where the channel is known. +_FILTER_ZERO_CHANNEL = """\ + const int32_t filter_zero = (weight_zero_point != NULL) + ? (int32_t)weight_zero_point[channel_index * weight_zero_stride] + : 0;""" + + +def _conv_integer(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + weights = context.require_input(1) + source_zero = context.optional_input(2) + weight_zero = context.optional_input(3) + result = context.require_output(0) + geometry = convolution_geometry(context, source, weights) + channels = geometry.groups * geometry.group_filters + verify_shape(context, result, geometry.result_shape) + _verify_per_tensor(context, source_zero, "zero point") + weight_zero_stride = _channel_stride(context, weight_zero, channels, "zero point") + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + name = kernel_name(context, c_type(source.elem_type), c_type(weights.elem_type)) + definition = _CONVOLUTION_TEMPLATE.substitute( + name=name, + source=c_type(source.elem_type), + weight=c_type(weights.elem_type), + result=c_type(result.elem_type), + parameters=",\n".join( + [ + f" {c_type(result.elem_type)}* out", + *_convolution_operands(source, weights), + WINDOW_PARAMETERS, + ] + ), + locals=_INPUT_ZERO_LOCAL, + channel=_FILTER_ZERO_CHANNEL, + initial="0", + store=" result[position] = sum;", + ) + call = call_kernel( + name, + [ + result.expr, + source.expr, + weights.expr, + _pointer(source_zero), + _pointer(weight_zero), + f"{weight_zero_stride}u", + *geometry.arguments, + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _qlinear_conv(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + source_scale = context.require_input(1) + source_zero = context.require_input(2) + weights = context.require_input(3) + weight_scale = context.require_input(4) + weight_zero = context.require_input(5) + result_scale = context.require_input(6) + result_zero = context.require_input(7) + bias = context.optional_input(8) + result = context.require_output(0) + geometry = convolution_geometry(context, source, weights) + channels = geometry.groups * geometry.group_filters + verify_shape(context, result, geometry.result_shape) + verify_bias(context, bias, channels) + for operand in (source, weights, result): + _verify_grid_type(context, operand, "operand") + _verify_zero_point(context, source_zero, source) + _verify_zero_point(context, weight_zero, weights) + _verify_zero_point(context, result_zero, result) + _verify_per_tensor(context, source_scale, "scale") + _verify_per_tensor(context, source_zero, "zero point") + _verify_per_tensor(context, result_scale, "scale") + _verify_per_tensor(context, result_zero, "zero point") + weight_scale_stride = _channel_stride(context, weight_scale, channels, "scale") + weight_zero_stride = _channel_stride(context, weight_zero, channels, "zero point") + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + saturate = _saturating_cast(context, result.elem_type) + name = kernel_name( + context, + c_type(source.elem_type), + c_type(weights.elem_type), + c_type(result.elem_type), + c_type(result_scale.elem_type), + ) + definition = _CONVOLUTION_TEMPLATE.substitute( + name=name, + source=c_type(source.elem_type), + weight=c_type(weights.elem_type), + result=c_type(result.elem_type), + parameters=",\n".join( + [ + f" {c_type(result.elem_type)}* out", + *_convolution_operands(source, weights), + f" const {c_type(source_scale.elem_type)}* input_scale", + f" const {c_type(weight_scale.elem_type)}* weight_scale", + " size_t weight_scale_stride", + f" const {c_type(result_scale.elem_type)}* result_scale", + f" const {c_type(result.elem_type)}* result_zero_point", + " const int32_t* bias", + WINDOW_PARAMETERS, + ] + ), + locals="\n".join( + [ + _INPUT_ZERO_LOCAL, + " const double result_zero = (double)result_zero_point[0];", + ] + ), + channel="\n".join( + [ + _FILTER_ZERO_CHANNEL, + # One factor per output channel: the scales of the two grids the products + # come from, over the scale of the grid the result is written on. + f" const {c_type(result_scale.elem_type)} factor =", + " input_scale[0]", + " * weight_scale[channel_index" + " * weight_scale_stride]", + " / result_scale[0];", + ] + ), + initial="(bias != NULL) ? bias[channel_index] : 0", + store=( + " result[position] =\n" + f" {saturate.name}(" + "(double)sum * (double)factor + result_zero);" + ), + ) + call = call_kernel( + name, + [ + result.expr, + source.expr, + weights.expr, + source_zero.expr, + weight_zero.expr, + f"{weight_zero_stride}u", + source_scale.expr, + weight_scale.expr, + f"{weight_scale_stride}u", + result_scale.expr, + result_zero.expr, + _pointer(bias), + *geometry.arguments, + ], + ) + return NodeEmission( + functions=(saturate, CFunction(name, definition)), statements=(call,) + ) + + +def _convolution_operands(source: TensorRef, weights: TensorRef) -> list[str]: + """The operands and zero points both quantized convolutions read, in call order.""" + return [ + f" const {c_type(source.elem_type)}* in", + f" const {c_type(weights.elem_type)}* weights", + f" const {c_type(source.elem_type)}* input_zero_point", + f" const {c_type(weights.elem_type)}* weight_zero_point", + " size_t weight_zero_stride", + ] + + +def _verify_per_tensor( + context: NodeContext, operand: TensorRef | None, role: str +) -> None: + """Refuse a scale or zero point this op reads as the whole tensor's and that is not one. + + ONNX's reference evaluator stretches the per-row vector its own text describes along + numpy's trailing axis rather than the axis that text names, so nothing can vouch for what + a kernel should compute for it, and the whole granularity is left unserved rather than + read one way in that form and another in the one numpy does broadcast as written. + """ + if operand is not None and operand.elem_count != 1: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads `{operand.name}` as " + f"the one {role} of the whole tensor, but it has shape " + f"{list(operand.shape)}; the C compiler serves this op at per-tensor " + "granularity only." + ) + + +def _channel_stride( + context: NodeContext, operand: TensorRef | None, channels: int, role: str +) -> int: + """The stride reading `operand` per output channel: one per channel, or one for all.""" + if operand is None or operand.elem_count == 1: + return 0 + if operand.shape != (channels,): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads `{operand.name}` as " + f"one {role} per output channel — a 1-D tensor of {channels} — or as one for " + f"the whole filter, but it has shape {list(operand.shape)}." + ) + return 1 + + +def _pointer(operand: TensorRef | None) -> str: + return "NULL" if operand is None else operand.expr + + +register_kernel("", "ConvInteger", _INTEGER_VERSIONS, _conv_integer) +register_kernel("", "QLinearConv", _INTEGER_VERSIONS, _qlinear_conv) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/recurrent.py b/src/python/fnnx/extras/compilers/c/onnx/ops/recurrent.py new file mode 100644 index 0000000..8f82ac0 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/recurrent.py @@ -0,0 +1,968 @@ +"""The recurrent layers: one time-stepped kernel per op, driven by call-site geometry. + +An LSTM, a GRU and an RNN all walk a batch of sequences one time step at a time, carrying a +hidden state `H` — and, for the LSTM, a cell state `C` — from step to step. Everything that +varies between two nodes of the same op — how many steps, how wide the state is, which way +time runs, where the operands sit in memory — reaches the kernel as arguments, so one shared +`static` function per op and element type serves every node, and the states themselves live +in static scratch sized at compile time. + +Two things shape the emission. A **direction** is an independent pass over the sequence with +its own weights, so the kernel computes one direction and a bidirectional node calls it twice, +with every operand offset onto that direction's slice. And **layout** only permutes where the +operands' elements sit, so it never reaches the kernel as a flag: the call site passes the +strides that place a time step and a batch item, which layout 1 simply reorders. + +The three ops differ only in what one step computes, so they share a frame — the batch loop, +the per-item sequence length, the padding ONNX reports past a sequence's end, and the state +outputs — and each supplies the recurrence that runs inside it, along with the signature that +recurrence reads. That is what `_Layer` collects. + +The activations ONNX lets a node choose — `f` over the gates, `g` over the cell or hidden +candidate, `h` over the LSTM's cell state on the way out — are emitted as one small function +each, taking the alpha and beta that parameterize them, and reach the kernel as function +pointers. Only the ones a node actually names are emitted. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from functools import partial +from string import Template + +import onnx.defs + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + ScratchBuffer, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import checked_call, verify_shape +from fnnx.extras.compilers.c.onnx.ops.broadcast import expand + +# The recurrent ops arrived at opset 1 and were revised repeatedly since — 7 dropped the +# legacy broadcast attributes, 14 added `layout`, 22 widened the element types. Only 22 is +# claimed: it is the revision every recurrent test in the backend corpus imports, and the one +# the reference evaluator is version-faithful for across the whole family, so it is the only +# one anything can vouch for. A model importing an older one gets the unsupported-version +# error. +_VERSIONS = (22,) + +_DIRECTIONS = {"forward": 1, "reverse": 1, "bidirectional": 2} + +# The LSTM's peephole vector `P` carries only the three gates that read the cell state — +# input, output, forget — where `W`, `R` and `B` carry all four. +_PEEPHOLES = 3 + + +@dataclass(frozen=True) +class _Activation: + """One of the activations ONNX's recurrent ops may be told to run. + + `expression` is C over `x`, `alpha` and `beta`, as ONNX's LSTM specification writes the + function. `schema` names the ONNX operator whose own defaults for alpha and beta apply, + which is how a node that selects an activation without parameterizing it gets the values + ONNX's operator would have applied; the two activations ONNX defines no operator for have + no such default, so a node naming one has to state the parameters itself. + """ + + expression: str + alpha: bool = False + beta: bool = False + schema: str | None = None + + +# The activation functions ONNX's recurrent documentation lists, each as the formula written +# there. Sigmoid takes the numerically stable form the Sigmoid kernel uses — whichever branch +# is taken, the exponent is of a negative number, so it underflows to zero rather than +# overflowing — and Relu propagates NaN the way numpy's `maximum`, which ONNX's own Relu is +# defined through, does. +_ACTIVATIONS: dict[str, _Activation] = { + "Affine": _Activation("alpha * x + beta", alpha=True, beta=True), + "Elu": _Activation( + "(x >= $zero) ? x : alpha * (exp$f(x) - $one)", alpha=True, schema="Elu" + ), + "HardSigmoid": _Activation( + "fmin$f(fmax$f(alpha * x + beta, $zero), $one)", + alpha=True, + beta=True, + schema="HardSigmoid", + ), + "LeakyRelu": _Activation( + "(x >= $zero) ? x : alpha * x", alpha=True, schema="LeakyRelu" + ), + "Relu": _Activation("(x > $zero || isnan(x)) ? x : $zero"), + "ScaledTanh": _Activation("alpha * tanh$f(beta * x)", alpha=True, beta=True), + "Sigmoid": _Activation( + "x > $zero ? $one / ($one + exp$f(-x)) : exp$f(x) / ($one + exp$f(x))" + ), + "Softplus": _Activation("log$f($one + exp$f(x))"), + "Softsign": _Activation("x / ($one + fabs$f(x))"), + "Tanh": _Activation("tanh$f(x)"), + "ThresholdedRelu": _Activation( + "(x >= alpha) ? x : $zero", alpha=True, schema="ThresholdedRelu" + ), +} + +_ACTIVATION_TEMPLATE = Template("""\ +static $element $name($element x, $element alpha, $element beta) +{ + (void)alpha; + (void)beta; + return $expression; +}""") + +# The cell clip, as two comparisons rather than `fmin`/`fmax` so that a NaN — which fails both +# — reaches the activation instead of being replaced by a bound. +_CLIP_TEMPLATE = Template("""\ +static $element $name($element x, $element clip) +{ + return x < -clip ? -clip : (x > clip ? clip : x); +}""") + + +# -------------------------------------------------------------------------------------- +# The frame every recurrent kernel runs inside +# -------------------------------------------------------------------------------------- + +# One direction of one layer. The recurrence is per batch item — a sequence's gates read only +# its own step and its own state — so the batch is the outer loop and the state buffers hold +# one item's worth, which is also what lets each item stop at its own sequence length. What +# the op computes from a step sits in `$recurrence`, which reads `values` and the current +# `hidden` and leaves the new state in `hidden`. +_FRAME = Template("""\ +static int $name( +$parameters) +{ + const size_t gate_count = $gates * hidden_size; + size_t item, step, unit, term; + for (item = 0; item < batch_size; ++item) { + const $element* sequence = x + item * x_batch_stride; + const size_t state_base = item * state_batch_stride; + size_t length = seq_length; + if (lengths != NULL) { + /* ONNX defines a length as a position in the padded sequence; anything else + names a step that is not there. */ + if (lengths[item] < 0 || (size_t)lengths[item] > seq_length) { + return 1; + } + length = (size_t)lengths[item]; + } + for (unit = 0; unit < hidden_size; ++unit) { +$carried_in + } + for (step = 0; step < length; ++step) { + const size_t time = reverse ? length - 1 - step : step; + const $element* values = sequence + time * x_time_stride; +$recurrence + if (y != NULL) { + $element* written = y + item * y_batch_stride + time * y_time_stride; + for (unit = 0; unit < hidden_size; ++unit) { + written[unit] = hidden[unit]; + } + } + } + if (y != NULL) { + /* The steps past this sequence's own end are padding, which ONNX reports as + zeros rather than as state. */ + for (step = length; step < seq_length; ++step) { + $element* written = y + item * y_batch_stride + step * y_time_stride; + for (unit = 0; unit < hidden_size; ++unit) { + written[unit] = $zero; + } + } + } + for (unit = 0; unit < hidden_size; ++unit) { + /* A sequence of no steps carries no state to report: ONNX says nothing about + one, and this is what onnxruntime returns for it. */ +$carried_out + } + } + return 0; +}""") + +# The hidden state is the whole of what a GRU or an RNN carries between steps; the LSTM adds +# its cell state alongside, read and reported the same way. +_HIDDEN_IN = """\ + hidden[unit] = initial_h == NULL ? $zero : initial_h[state_base + unit];""" + +_HIDDEN_OUT = """\ + if (y_h != NULL) { + y_h[state_base + unit] = length == 0 ? $zero : hidden[unit]; + }""" + +_LSTM_CARRIED_IN = ( + _HIDDEN_IN + + """ + cell[unit] = initial_c == NULL ? $zero : initial_c[state_base + unit];""" +) + +_LSTM_CARRIED_OUT = ( + _HIDDEN_OUT + + """ + if (y_c != NULL) { + y_c[state_base + unit] = length == 0 ? $zero : cell[unit]; + }""" +) + + +def _frame( + parameters: str, carried_in: str, carried_out: str, recurrence: str +) -> Template: + """The frame with one op's pieces spliced in, still holding the per-node placeholders. + + Two passes: the pieces carry `$element` and `$zero` of their own, which a single + substitution would leave untouched inside the text it inserts. + """ + return Template( + _FRAME.safe_substitute( + parameters=parameters, + carried_in=carried_in, + carried_out=carried_out, + recurrence=recurrence, + ) + ) + + +# -------------------------------------------------------------------------------------- +# LSTM +# -------------------------------------------------------------------------------------- + +_LSTM_PARAMETERS = """\ + $element* y, + $element* y_h, + $element* y_c, + const $element* x, + const $element* w, + const $element* r, + const $element* bias, + const int32_t* lengths, + const $element* initial_h, + const $element* initial_c, + const $element* peepholes, + $element* hidden, + $element* cell, + $element* gates, + size_t seq_length, + size_t batch_size, + size_t input_size, + size_t hidden_size, + size_t x_time_stride, + size_t x_batch_stride, + size_t y_time_stride, + size_t y_batch_stride, + size_t state_batch_stride, + int reverse, + int coupled, + int clipped, + $element clip, + $element (*act_f)($element, $element, $element), + $element alpha_f, + $element beta_f, + $element (*act_g)($element, $element, $element), + $element alpha_g, + $element beta_g, + $element (*act_h)($element, $element, $element), + $element alpha_h, + $element beta_h""" + +# The four gates in the order ONNX concatenates them along `W`, `R` and `B`: input, output, +# forget, cell. Each unit's update reads only its own gates and its own cell, so the new +# hidden state can be written in place as it goes. +_LSTM_RECURRENCE = """\ + for (unit = 0; unit < gate_count; ++unit) { + $element total = + bias == NULL ? $zero : bias[unit] + bias[gate_count + unit]; + for (term = 0; term < input_size; ++term) { + total += w[unit * input_size + term] * values[term]; + } + for (term = 0; term < hidden_size; ++term) { + total += r[unit * hidden_size + term] * hidden[term]; + } + gates[unit] = total; + } + for (unit = 0; unit < hidden_size; ++unit) { + const $element state = cell[unit]; + $element input_gate = gates[unit]; + $element forget_gate = gates[2 * hidden_size + unit]; + $element candidate = gates[3 * hidden_size + unit]; + $element output_gate = gates[hidden_size + unit]; + $element updated; + if (peepholes != NULL) { + input_gate += peepholes[unit] * state; + forget_gate += peepholes[2 * hidden_size + unit] * state; + } + if (clipped) { + input_gate = $clip(input_gate, clip); + forget_gate = $clip(forget_gate, clip); + candidate = $clip(candidate, clip); + } + input_gate = act_f(input_gate, alpha_f, beta_f); + forget_gate = coupled + ? $one - input_gate + : act_f(forget_gate, alpha_f, beta_f); + updated = forget_gate * state + + input_gate * act_g(candidate, alpha_g, beta_g); + cell[unit] = updated; + if (peepholes != NULL) { + output_gate += peepholes[hidden_size + unit] * updated; + } + if (clipped) { + output_gate = $clip(output_gate, clip); + } + hidden[unit] = act_f(output_gate, alpha_f, beta_f) + * act_h(updated, alpha_h, beta_h); + }""" + + +# -------------------------------------------------------------------------------------- +# GRU +# -------------------------------------------------------------------------------------- + +_GRU_PARAMETERS = """\ + $element* y, + $element* y_h, + const $element* x, + const $element* w, + const $element* r, + const $element* bias, + const int32_t* lengths, + const $element* initial_h, + $element* hidden, + $element* gates, + size_t seq_length, + size_t batch_size, + size_t input_size, + size_t hidden_size, + size_t x_time_stride, + size_t x_batch_stride, + size_t y_time_stride, + size_t y_batch_stride, + size_t state_batch_stride, + int reverse, + int linear_before_reset, + int clipped, + $element clip, + $element (*act_f)($element, $element, $element), + $element alpha_f, + $element beta_f, + $element (*act_g)($element, $element, $element), + $element alpha_g, + $element beta_g""" + +# The three gates in ONNX's order: update, reset, hidden candidate. The candidate reads the +# reset gate, so it cannot share the first loop; and it reads the whole previous state +# through the recurrence weights, so the state is replaced only once every gate has read it. +# The candidate is parked in the third gate slot, which nothing else uses. +_GRU_RECURRENCE = """\ + for (unit = 0; unit < 2 * hidden_size; ++unit) { + $element total = + bias == NULL ? $zero : bias[unit] + bias[gate_count + unit]; + for (term = 0; term < input_size; ++term) { + total += w[unit * input_size + term] * values[term]; + } + for (term = 0; term < hidden_size; ++term) { + total += r[unit * hidden_size + term] * hidden[term]; + } + if (clipped) { + total = $clip(total, clip); + } + gates[unit] = act_f(total, alpha_f, beta_f); + } + for (unit = 0; unit < hidden_size; ++unit) { + const size_t gate = 2 * hidden_size + unit; + $element total = bias == NULL ? $zero : bias[gate]; + for (term = 0; term < input_size; ++term) { + total += w[gate * input_size + term] * values[term]; + } + if (linear_before_reset) { + /* This unit's own reset gate scales the whole recurrence, its bias + included. */ + $element carried = bias == NULL ? $zero : bias[gate_count + gate]; + for (term = 0; term < hidden_size; ++term) { + carried += r[gate * hidden_size + term] * hidden[term]; + } + total += gates[hidden_size + unit] * carried; + } else { + /* The reset gate scales the state, before the recurrence weights read + it, so each term is scaled by the gate of the unit it belongs to. */ + if (bias != NULL) { + total += bias[gate_count + gate]; + } + for (term = 0; term < hidden_size; ++term) { + total += r[gate * hidden_size + term] + * (gates[hidden_size + term] * hidden[term]); + } + } + if (clipped) { + total = $clip(total, clip); + } + gates[gate] = act_g(total, alpha_g, beta_g); + } + for (unit = 0; unit < hidden_size; ++unit) { + const $element update = gates[unit]; + hidden[unit] = ($one - update) * gates[2 * hidden_size + unit] + + update * hidden[unit]; + }""" + + +# -------------------------------------------------------------------------------------- +# RNN +# -------------------------------------------------------------------------------------- + +_RNN_PARAMETERS = """\ + $element* y, + $element* y_h, + const $element* x, + const $element* w, + const $element* r, + const $element* bias, + const int32_t* lengths, + const $element* initial_h, + $element* hidden, + $element* gates, + size_t seq_length, + size_t batch_size, + size_t input_size, + size_t hidden_size, + size_t x_time_stride, + size_t x_batch_stride, + size_t y_time_stride, + size_t y_batch_stride, + size_t state_batch_stride, + int reverse, + int clipped, + $element clip, + $element (*act_f)($element, $element, $element), + $element alpha_f, + $element beta_f""" + +# One gate, whose activation is the whole of the new state. Every unit reads the whole +# previous state, so the state is replaced only once every gate has been accumulated. +_RNN_RECURRENCE = """\ + for (unit = 0; unit < hidden_size; ++unit) { + $element total = + bias == NULL ? $zero : bias[unit] + bias[gate_count + unit]; + for (term = 0; term < input_size; ++term) { + total += w[unit * input_size + term] * values[term]; + } + for (term = 0; term < hidden_size; ++term) { + total += r[unit * hidden_size + term] * hidden[term]; + } + if (clipped) { + total = $clip(total, clip); + } + gates[unit] = total; + } + for (unit = 0; unit < hidden_size; ++unit) { + hidden[unit] = act_f(gates[unit], alpha_f, beta_f); + }""" + + +@dataclass(frozen=True) +class _Layer: + """What distinguishes one of ONNX's three recurrent ops from the other two. + + `gates` is how many gate rows `W`, `R` and `B` carry per direction. `defaults` is the + `(f, g, h)` — as many as the op runs — that a node naming no activations gets, which also + fixes how many a node that does name them has to name. `mode` is the single integer + attribute the op turns into a branch inside the kernel. `cell` marks the LSTM: the only + one carrying a second state between steps, and so the only one with `initial_c`, `Y_c` + and peephole weights. + """ + + op_type: str + gates: int + defaults: tuple[str, ...] + template: Template + scratch: tuple[str, ...] + mode: str | None = None + cell: bool = False + + @property + def symbol(self) -> str: + return self.op_type.lower() + + @property + def results(self) -> int: + return 3 if self.cell else 2 + + +_LAYERS = ( + _Layer( + op_type="LSTM", + gates=4, + defaults=("Sigmoid", "Tanh", "Tanh"), + template=_frame( + _LSTM_PARAMETERS, _LSTM_CARRIED_IN, _LSTM_CARRIED_OUT, _LSTM_RECURRENCE + ), + scratch=("hidden", "cell", "gates"), + mode="input_forget", + cell=True, + ), + _Layer( + op_type="GRU", + gates=3, + defaults=("Sigmoid", "Tanh"), + template=_frame(_GRU_PARAMETERS, _HIDDEN_IN, _HIDDEN_OUT, _GRU_RECURRENCE), + scratch=("hidden", "gates"), + mode="linear_before_reset", + ), + _Layer( + op_type="RNN", + gates=1, + defaults=("Tanh",), + template=_frame(_RNN_PARAMETERS, _HIDDEN_IN, _HIDDEN_OUT, _RNN_RECURRENCE), + scratch=("hidden", "gates"), + ), +) + + +@dataclass(frozen=True) +class _Geometry: + """A node's shape, and where each operand's elements sit under its layout. + + The strides are in elements. Layout 1 packs the batch outermost rather than time, which + changes nothing the kernel computes — only how far apart two steps of one sequence are. + """ + + seq_length: int + batch_size: int + input_size: int + hidden_size: int + direction: str + layout: int + + @property + def directions(self) -> int: + return _DIRECTIONS[self.direction] + + @property + def x_shape(self) -> tuple[int, ...]: + if self.layout: + return (self.batch_size, self.seq_length, self.input_size) + return (self.seq_length, self.batch_size, self.input_size) + + @property + def y_shape(self) -> tuple[int, ...]: + if self.layout: + return (self.batch_size, self.seq_length, self.directions, self.hidden_size) + return (self.seq_length, self.directions, self.batch_size, self.hidden_size) + + @property + def state_shape(self) -> tuple[int, ...]: + if self.layout: + return (self.batch_size, self.directions, self.hidden_size) + return (self.directions, self.batch_size, self.hidden_size) + + @property + def strides(self) -> tuple[int, ...]: + """`(x_time, x_batch, y_time, y_batch, state_batch)`, in elements.""" + if self.layout: + return ( + self.input_size, + self.seq_length * self.input_size, + self.directions * self.hidden_size, + self.seq_length * self.directions * self.hidden_size, + self.directions * self.hidden_size, + ) + return ( + self.batch_size * self.input_size, + self.input_size, + self.directions * self.batch_size * self.hidden_size, + self.hidden_size, + self.hidden_size, + ) + + def state_offset(self, direction: int) -> int: + """Where this direction's slice of `Y`, of a state output and of an initial one begins. + + One offset serves all of them: under either layout the direction axis sits directly + outside the hidden units in every one. + """ + if self.layout: + return direction * self.hidden_size + return direction * self.batch_size * self.hidden_size + + def runs_backwards(self, direction: int) -> bool: + return self.direction == "reverse" or ( + self.direction == "bidirectional" and direction == 1 + ) + + +def _recurrent(context: NodeContext, layer: _Layer) -> NodeEmission: + geometry = _geometry(context, layer) + activations = _activations(context, layer, geometry.directions) + _verify_operands(context, layer, geometry) + outputs = tuple( + context.outputs[index] if index < len(context.outputs) else None + for index in range(layer.results) + ) + if all(result is None or result.elem_count == 0 for result in outputs): + return NodeEmission(functions=(), statements=()) + + elem_type = context.require_input(0).elem_type + element = c_type(elem_type) + name = f"{context.prefix}_{layer.symbol}_{element}" + bound = _clip_function(context, elem_type) + definition = layer.template.substitute( + name=name, + element=element, + gates=layer.gates, + clip=bound.name, + one=scalar_literal(1, elem_type), + zero=scalar_literal(0, elem_type), + ) + scratch = _scratch(context, layer, geometry, elem_type) + helpers = { + function.name: function + for selection in activations + for function in selection.functions(context, elem_type) + } + return NodeEmission( + functions=(bound, *helpers.values(), CFunction(name, definition)), + statements=tuple( + checked_call( + context, + name, + _arguments( + context, layer, geometry, outputs, scratch, activations, direction + ), + ) + for direction in range(geometry.directions) + ), + scratch=scratch, + ) + + +# -------------------------------------------------------------------------------------- +# The activations a node selects +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _Selection: + """The activations one direction runs, each with the alpha and beta it is given.""" + + names: tuple[str, ...] + alphas: tuple[float, ...] + betas: tuple[float, ...] + + def functions(self, context: NodeContext, elem_type: int) -> list[CFunction]: + return [_activation_function(context, name, elem_type) for name in self.names] + + def arguments(self, context: NodeContext, elem_type: int) -> list[str]: + """The function pointer and the two parameters each activation is called through.""" + arguments = [] + for name, alpha, beta in zip(self.names, self.alphas, self.betas): + arguments.append(_activation_function(context, name, elem_type).name) + arguments.append(scalar_literal(alpha, elem_type)) + arguments.append(scalar_literal(beta, elem_type)) + return arguments + + +def _activations( + context: NodeContext, layer: _Layer, directions: int +) -> tuple[_Selection, ...]: + """The activations each direction runs, defaulting to the op's own. + + ONNX takes the names as one flat list per direction, and their parameters as two more + lists consumed in the same order. + """ + count = len(layer.defaults) + names = [name.decode() for name in context.attribute("activations", [])] + if not names: + names = list(layer.defaults) * directions + if len(names) != count * directions: + expected = count * directions + raise CompileError( + f"Node `{context.label}`: a `{_direction(context, layer)}` `{layer.op_type}` " + f"runs {expected} activation{'' if expected == 1 else 's'} — {count} per " + f"direction — but this node names {len(names)}: {', '.join(names) or 'none'}." + ) + alphas = _parameters(context, layer, names, "alpha") + betas = _parameters(context, layer, names, "beta") + return tuple( + _Selection( + names=tuple(names[count * index : count * (index + 1)]), + alphas=alphas[count * index : count * (index + 1)], + betas=betas[count * index : count * (index + 1)], + ) + for index in range(directions) + ) + + +def _parameters( + context: NodeContext, layer: _Layer, names: Sequence[str], role: str +) -> tuple[float, ...]: + """The alpha or beta each named activation runs with, else ONNX's own default. + + ONNX's list carries a value for the activations that take one and for no others: they + consume it in the order the node names them, over both directions of a bidirectional + node together, so a parameterized activation reads the next value rather than the one at + its own position. + """ + given = [float(value) for value in context.attribute(f"activation_{role}", [])] + resolved: list[float] = [] + consumed = 0 + for name in names: + activation = _activation(context, layer, name) + if not getattr(activation, role): + resolved.append(0.0) + continue + if consumed < len(given): + resolved.append(given[consumed]) + elif activation.schema is None: + raise CompileError( + f"Node `{context.label}`: `{layer.op_type}` runs `{name}`, which ONNX " + f"defines no operator for, so there is no default `activation_{role}` to " + f"fall back on; the attribute carries one value per activation that takes " + f"one, and `{name}` reads number {consumed + 1}." + ) + else: + resolved.append( + float( + onnx.defs.get_schema(activation.schema) + .attributes[role] + .default_value.f + ) + ) + consumed += 1 + return tuple(resolved) + + +def _activation(context: NodeContext, layer: _Layer, name: str) -> _Activation: + activation = _ACTIVATIONS.get(name) + if activation is None: + raise CompileError( + f"Node `{context.label}`: `{layer.op_type}` names the activation `{name}`, " + f"which is not one of the functions ONNX defines it over: " + f"{', '.join(_ACTIVATIONS)}." + ) + return activation + + +def _activation_function(context: NodeContext, name: str, elem_type: int) -> CFunction: + symbol = f"{context.prefix}_rnnact_{name.lower()}_{c_type(elem_type)}" + return CFunction( + symbol, + _ACTIVATION_TEMPLATE.substitute( + name=symbol, + element=c_type(elem_type), + expression=expand(_ACTIVATIONS[name].expression, elem_type), + ), + ) + + +def _clip_function(context: NodeContext, elem_type: int) -> CFunction: + name = f"{context.prefix}_rnnclip_{c_type(elem_type)}" + return CFunction( + name, + _CLIP_TEMPLATE.substitute(name=name, element=c_type(elem_type)), + ) + + +# -------------------------------------------------------------------------------------- +# Reading the geometry off the node, and placing the call +# -------------------------------------------------------------------------------------- + + +def _geometry(context: NodeContext, layer: _Layer) -> _Geometry: + source = context.require_input(0) + recurrence = context.require_input(2) + if len(source.shape) != 3: + raise CompileError( + f"Node `{context.label}`: `{layer.op_type}` reads a batch of sequences — a " + f"tensor of rank 3 — but `{source.name}` has shape {list(source.shape)}." + ) + if len(recurrence.shape) != 3: + raise CompileError( + f"Node `{context.label}`: `{layer.op_type}` reads its recurrence weights as " + f"rank 3, but `{recurrence.name}` has shape {list(recurrence.shape)}." + ) + layout = _layout(context, layer) + steps, items = source.shape[0], source.shape[1] + return _Geometry( + seq_length=items if layout else steps, + batch_size=steps if layout else items, + input_size=source.shape[2], + hidden_size=_hidden_size(context, layer, recurrence), + direction=_direction(context, layer), + layout=layout, + ) + + +def _layout(context: NodeContext, layer: _Layer) -> int: + layout = context.int_attribute("layout") + if layout not in (0, 1): + raise CompileError( + f"Node `{context.label}`: `{layer.op_type}`'s `layout` is {layout}, but ONNX " + "defines only 0 (time first) and 1 (batch first)." + ) + return layout + + +def _direction(context: NodeContext, layer: _Layer) -> str: + name = context.attribute("direction", b"forward").decode() + if name not in _DIRECTIONS: + raise CompileError( + f"Node `{context.label}`: `{layer.op_type}`'s `direction` is `{name}`, which " + f"is not one of the directions ONNX defines: {', '.join(_DIRECTIONS)}." + ) + return name + + +def _hidden_size(context: NodeContext, layer: _Layer, recurrence: TensorRef) -> int: + """The width of the state, which the recurrence weights carry. + + ONNX states it as an attribute too; a node whose attribute disagrees with the weights it + is handed describes two different layers, and there is no telling which one it meant. + """ + hidden_size = recurrence.shape[2] + declared = context.attribute("hidden_size", None) + if declared is not None and int(declared) != hidden_size: + raise CompileError( + f"Node `{context.label}`: `{layer.op_type}` states a `hidden_size` of " + f"{int(declared)}, but its recurrence weights `{recurrence.name}` are shaped " + f"for {hidden_size}." + ) + return hidden_size + + +def _operands( + layer: _Layer, geometry: _Geometry +) -> dict[int, tuple[str, tuple[int, ...]]]: + """Each optional operand past `X`, by input index: what it is, and the shape it must be.""" + hidden, directions = geometry.hidden_size, geometry.directions + row = layer.gates * hidden + operands = { + 1: ("its input weights", (directions, row, geometry.input_size)), + 2: ("its recurrence weights", (directions, row, hidden)), + 3: ("its biases", (directions, 2 * row)), + 4: ("its sequence lengths", (geometry.batch_size,)), + 5: ("an initial hidden state", geometry.state_shape), + } + if layer.cell: + operands[6] = ("an initial cell state", geometry.state_shape) + operands[7] = ("its peephole weights", (directions, _PEEPHOLES * hidden)) + return operands + + +def _verify_operands(context: NodeContext, layer: _Layer, geometry: _Geometry) -> None: + """Refuse to emit a kernel whose addressing disagrees with the buffers it is handed. + + ONNX's own shape inference derives most of this, but a graph may declare the shapes it + would infer rather than have them inferred, so the extents the kernel walks are checked + here against the buffers themselves. + """ + for index, (role, shape) in _operands(layer, geometry).items(): + operand = context.optional_input(index) + if operand is not None and operand.shape != shape: + raise CompileError( + f"Node `{context.label}`: `{layer.op_type}` reads `{operand.name}` as " + f"{role} of shape {list(shape)}, but it has shape {list(operand.shape)}." + ) + results = (geometry.y_shape, *(geometry.state_shape,) * (layer.results - 1)) + for index, shape in enumerate(results): + result = context.outputs[index] if index < len(context.outputs) else None + if result is not None: + verify_shape(context, result, shape) + + +def _scratch( + context: NodeContext, layer: _Layer, geometry: _Geometry, elem_type: int +) -> tuple[ScratchBuffer, ...]: + """The state of one batch item, and the gate values of one step, reserved statically. + + A step reads the whole previous state while writing the new one, so neither state can be + computed in place in a caller's buffer; the artifact allocates nothing, so the space is + reserved at compile time and counted in the reported footprint like every other buffer. + """ + element = c_type(elem_type) + counts = { + "hidden": geometry.hidden_size, + "cell": geometry.hidden_size, + "gates": layer.gates * geometry.hidden_size, + } + return tuple( + ScratchBuffer( + f"{context.prefix}_{layer.symbol}_{role}_{element}", elem_type, counts[role] + ) + for role in layer.scratch + ) + + +def _arguments( + context: NodeContext, + layer: _Layer, + geometry: _Geometry, + outputs: tuple[TensorRef | None, ...], + scratch: tuple[ScratchBuffer, ...], + activations: tuple[_Selection, ...], + direction: int, +) -> list[str]: + """The call site for one direction: every operand offset onto that direction's slice.""" + elem_type = context.require_input(0).elem_type + hidden, state = geometry.hidden_size, geometry.state_offset(direction) + row = layer.gates * hidden + offsets = { + 1: direction * row * geometry.input_size, + 2: direction * row * hidden, + 3: direction * 2 * row, + 4: 0, + 5: state, + 6: state, + 7: direction * _PEEPHOLES * hidden, + } + clip = context.attribute("clip", None) + arguments = [_operand(result, state) for result in outputs] + arguments.append(_operand(context.require_input(0), 0)) + arguments += [ + _operand(context.optional_input(index), offsets[index]) + for index in sorted(_operands(layer, geometry)) + ] + arguments += [buffer.symbol for buffer in scratch] + arguments += [ + f"{geometry.seq_length}u", + f"{geometry.batch_size}u", + f"{geometry.input_size}u", + f"{hidden}u", + *(f"{stride}u" for stride in geometry.strides), + str(int(geometry.runs_backwards(direction))), + ] + if layer.mode is not None: + arguments.append(str(int(context.int_attribute(layer.mode) != 0))) + arguments += [ + str(int(clip is not None)), + scalar_literal(_clip_threshold(context, layer, clip), elem_type), + *activations[direction].arguments(context, elem_type), + ] + return arguments + + +def _operand(ref: TensorRef | None, offset: int) -> str: + if ref is None: + return "NULL" + return ref.expr if offset == 0 else f"{ref.expr} + {offset}" + + +def _clip_threshold(context: NodeContext, layer: _Layer, clip: float | None) -> float: + if clip is None: + return 0.0 + if not clip >= 0.0: + raise CompileError( + f"Node `{context.label}`: `{layer.op_type}`'s `clip` is {clip}, but ONNX " + "defines it as the threshold a cell is bounded to, which is not negative." + ) + return float(clip) + + +for _layer in _LAYERS: + register_kernel("", _layer.op_type, _VERSIONS, partial(_recurrent, layer=_layer)) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/reduce.py b/src/python/fnnx/extras/compilers/c/onnx/ops/reduce.py new file mode 100644 index 0000000..2fcd55b --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/reduce.py @@ -0,0 +1,439 @@ +"""The reductions: the Reduce* family over any set of axes, and ArgMax/ArgMin. + +Every one of them folds the elements of a group into a single value, so they share one loop +nest — the group loop and the element loop the named axes describe — and differ only in what +the fold starts from, what it does per element, and what it makes of the accumulator. ONNX +moved `axes` from an attribute to an input partway through the family's history; both +conventions reach the same emitter once the axes are resolved. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from functools import partial +from string import Template + +import numpy as np +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + FLOAT_TYPES, + UNSIGNED_TYPES, + c_type, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + GROUP_PARAMETERS, + call_kernel, + group_axes, + kernel_name, + normalize_axes, + normalize_axis, + offset_helper, + verify_group_count, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import combiner, expand, math_suffix + +# The revisions that take `axes` as an attribute, and the ones that take it as an input. +# ReduceSum moved at 13 and the rest of the family at 18; the revisions listed alongside each +# move only widened type constraints — 12 added the int8 families to ReduceMax/ReduceMin and +# 20 the boolean ones — which leaves the emitted code unchanged. +_ATTRIBUTE_VERSIONS = (1, 11, 13) +_INPUT_VERSIONS = (18,) +_SUM_ATTRIBUTE_VERSIONS = (1, 11) +_SUM_INPUT_VERSIONS = (13,) +_EXTREMUM_ATTRIBUTE_VERSIONS = (1, 11, 12, 13) +_EXTREMUM_INPUT_VERSIONS = (18, 20) + +# ArgMax/ArgMin-12 added `select_last_index`, whose default is what the earlier revisions +# compute; 11 allowed a negative axis and 13 widened the type constraints. +_ARG_VERSIONS = (1, 11, 12, 13) + +_REDUCE_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, +$parameters) +{ + size_t group, index; + for (group = 0; group < group_count; ++group) { + const size_t base = $offset(group, kept_rank, kept_shape, kept_strides); + $element total = $identity; + for (index = 0; index < group_size; ++index) { + const $element x = in[base + + $offset(index, reduced_rank, reduced_shape, reduced_strides)]; + total = ($element)($combine); + } + out[group] = $result; + } +}""") + +# LogSumExp is the one reduction that reads its group twice: the largest element is +# subtracted from every exponent so that no term overflows, and added back afterwards, which +# is how the reference evaluator computes it too. +_LOG_SUM_EXP_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, +$parameters) +{ + size_t group, index; + for (group = 0; group < group_count; ++group) { + const size_t base = $offset(group, kept_rank, kept_shape, kept_strides); + $element largest = $identity; + $accumulator total = $accumulator_zero; + for (index = 0; index < group_size; ++index) { + largest = ($element)$maximum(largest, in[base + + $offset(index, reduced_rank, reduced_shape, reduced_strides)]); + } + for (index = 0; index < group_size; ++index) { + const $element x = in[base + + $offset(index, reduced_rank, reduced_shape, reduced_strides)]; + total += exp$suffix(($accumulator)(x - largest)); + } + out[group] = ($element)(log$suffix(total) + ($accumulator)largest); + } +}""") + +_ARG_TEMPLATE = Template("""\ +static void $name( + int64_t* out, + const $element* in, +$parameters) +{ + size_t group, index; + for (group = 0; group < group_count; ++group) { + const size_t base = $offset(group, kept_rank, kept_shape, kept_strides); + $element best = $zero; + int64_t chosen = 0; + for (index = 0; index < group_size; ++index) { + const $element x = in[base + + $offset(index, reduced_rank, reduced_shape, reduced_strides)]; + if (index == 0 || ($better)) { + best = x; + chosen = (int64_t)index; + } + } + out[group] = chosen; + } +}""") + + +@dataclass(frozen=True) +class Fold: + """How one reduction folds a group, as C over `total`, `x` and `group_size`. + + `identity` is where the fold starts, which is also what a group with no elements at all + yields — the value the reference evaluator fills such a reduction with. All three + expressions are `$`-templated over the element type. + """ + + identity: str + combine: str + result: str = "total" + helpers: tuple[CFunction, ...] = () + + +# What a kernel builder hands back: the kernel to call, and the functions it calls in turn. +Kernel = tuple[CFunction, tuple[CFunction, ...]] +KernelBuilder = Callable[[NodeContext], Kernel] +Recipe = Callable[[NodeContext], Fold] + + +def _sum_fold(context: NodeContext) -> Fold: + return Fold(identity="$zero", combine="total + x") + + +def _product_fold(context: NodeContext) -> Fold: + return Fold(identity="$one", combine="total * x") + + +def _mean_fold(context: NodeContext) -> Fold: + """The sum over the group's size, taken in the element type for the float families. + + numpy divides an integer sum in double and casts back, truncating toward zero; a group + with no elements makes that a 0/0 whose cast numpy leaves undefined, so the integer form + yields zero there rather than emitting a conversion C does not define either. + """ + if context.require_output(0).elem_type in FLOAT_TYPES: + return Fold( + identity="$zero", + combine="total + x", + result="total / ($element)group_size", + ) + return Fold( + identity="$zero", + combine="total + x", + result="($element)(group_size ? (double)total / (double)group_size : 0.0)", + ) + + +def _absolute_sum_fold(context: NodeContext) -> Fold: + return Fold(identity="$zero", combine=f"total + {_absolute(context)}") + + +def _square_sum_fold(context: NodeContext) -> Fold: + return Fold(identity="$zero", combine="total + x * x") + + +def _euclidean_fold(context: NodeContext) -> Fold: + return Fold( + identity="$zero", combine="total + x * x", result=_libm(context, "sqrt") + ) + + +def _log_sum_fold(context: NodeContext) -> Fold: + return Fold(identity="$zero", combine="total + x", result=_libm(context, "log")) + + +def _extremum_fold(context: NodeContext, *, largest: bool) -> Fold: + elem_type = context.require_output(0).elem_type + helper = combiner(context, elem_type, largest=largest) + return Fold( + identity=extremum_identity(elem_type, largest=largest), + combine=f"{helper.name}(total, x)", + helpers=(helper,), + ) + + +def _absolute(context: NodeContext) -> str: + """`|x|`, which the integer families take off a comparison rather than from libm.""" + elem_type = context.require_output(0).elem_type + if elem_type in FLOAT_TYPES: + return "fabs$f(x)" + if elem_type in UNSIGNED_TYPES: + return "x" + return "((x < $zero) ? ($element)-x : x)" + + +def _libm(context: NodeContext, function: str) -> str: + """A libm call on the accumulator, taken in double for the integer families. + + numpy evaluates these in floating point whatever the tensor holds and casts the result + back, so an integer reduction rounds once, on the way out. + """ + elem_type = context.require_output(0).elem_type + if elem_type in FLOAT_TYPES: + return f"{function}{math_suffix(elem_type)}(total)" + return f"($element){function}((double)total)" + + +def extremum_identity(elem_type: int, *, largest: bool) -> str: + """The neutral element of a max or min fold at this element type. + + Shared with the poolings, whose window is a max fold over part of a tensor. + """ + if elem_type in FLOAT_TYPES: + return "-INFINITY" if largest else "INFINITY" + if elem_type == TensorProto.BOOL: + return "0" if largest else "1" + info = np.iinfo(numpy_dtype_name(elem_type)) + return scalar_literal(info.min if largest else info.max, elem_type) + + +def _fold_kernel(context: NodeContext, *, recipe: Recipe) -> Kernel: + result = context.require_output(0) + fold = recipe(context) + offset = offset_helper(context.prefix) + name = kernel_name(context, numpy_dtype_name(result.elem_type)) + definition = _REDUCE_TEMPLATE.substitute( + name=name, + element=c_type(result.elem_type), + parameters=GROUP_PARAMETERS, + offset=offset.name, + identity=expand(fold.identity, result.elem_type), + combine=expand(fold.combine, result.elem_type), + result=expand(fold.result, result.elem_type), + ) + return CFunction(name, definition), (offset, *fold.helpers) + + +def _log_sum_exp_kernel(context: NodeContext) -> Kernel: + elem_type = context.require_output(0).elem_type + floating = elem_type in FLOAT_TYPES + accumulator = c_type(elem_type) if floating else "double" + largest = combiner(context, elem_type, largest=True) + offset = offset_helper(context.prefix) + name = kernel_name(context, numpy_dtype_name(elem_type)) + definition = _LOG_SUM_EXP_TEMPLATE.substitute( + name=name, + element=c_type(elem_type), + parameters=GROUP_PARAMETERS, + offset=offset.name, + identity=extremum_identity(elem_type, largest=True), + maximum=largest.name, + accumulator=accumulator, + accumulator_zero=scalar_literal( + 0, elem_type if floating else TensorProto.DOUBLE + ), + suffix=math_suffix(elem_type) if floating else "", + ) + return CFunction(name, definition), (offset, largest) + + +def _from_attribute(context: NodeContext, *, build: KernelBuilder) -> NodeEmission: + """A Reduce* revision whose axes are an attribute; without them it reduces every axis.""" + axes = context.attribute("axes", None) + return _reduce(context, build, tuple(axes) if axes else None) + + +def _from_input(context: NodeContext, *, build: KernelBuilder) -> NodeEmission: + """A Reduce* revision whose axes are an input, which it may also name none through. + + Without axes the op reduces every one of them, or — under `noop_with_empty_axes` — none. + Reducing none is not the identity: every element becomes a group of its own and still + goes through the fold, so ReduceL1 over no axes is an absolute value. + """ + axes = _axes_operand(context) + if axes is None and context.int_attribute("noop_with_empty_axes"): + axes = () + return _reduce(context, build, axes) + + +def _axes_operand(context: NodeContext) -> tuple[int, ...] | None: + """The axes the node's second operand names, None when it names none at all. + + An operand with no elements names no axes whatever it holds at run time, which is what + makes the corpus's `default_axes` models compilable; one carrying values the graph does + not fix is rejected by the frontend before any kernel is reached. + """ + operand = context.optional_input(1) + if operand is None or operand.elem_count == 0: + return None + values = context.constant_input(1) + if values is None: + raise CompileError( + f"Node `{context.label}`: the axes of `{context.node.op_type}` come from " + f"`{operand.name}`, which is not known at compile time; the shape of the " + "result then depends on input data, which the C compiler cannot compile." + ) + return tuple(int(axis) for axis in values.reshape(-1)) + + +def _reduce( + context: NodeContext, build: KernelBuilder, axes: Sequence[int] | None +) -> NodeEmission: + """Emit the fold over `axes`, where None stands for every axis and `()` for none.""" + source = context.require_input(0) + result = context.require_output(0) + rank = len(source.shape) + selected = ( + tuple(range(rank)) if axes is None else normalize_axes(context, axes, rank) + ) + grouping = group_axes(source.shape, selected) + verify_group_count(context, grouping, result) + + kernel, helpers = build(context) + return NodeEmission( + functions=(*helpers, kernel), + statements=( + call_kernel(kernel.name, [result.expr, source.expr, *grouping.arguments]), + ), + ) + + +def _arg_extremum(context: NodeContext, *, largest: bool) -> NodeEmission: + """ArgMax or ArgMin: the index of a group's extreme element, as numpy chooses it. + + A NaN is the extreme of any group it appears in, and the first one there wins, since + nothing that follows can better it; `select_last_index` reverses which of several equal + extremes is reported, exactly as the reference evaluator's own flip does. + """ + source = context.require_input(0) + result = context.require_output(0) + last = bool(context.attribute("select_last_index", 0)) + axis = normalize_axis(context, context.int_attribute("axis"), len(source.shape)) + grouping = group_axes(source.shape, (axis,)) + verify_group_count(context, grouping, result) + + offset = offset_helper(context.prefix) + name = kernel_name( + context, "last" if last else "first", numpy_dtype_name(source.elem_type) + ) + definition = _ARG_TEMPLATE.substitute( + name=name, + element=c_type(source.elem_type), + parameters=GROUP_PARAMETERS, + offset=offset.name, + zero=scalar_literal(0, source.elem_type), + better=extremum_test(source.elem_type, largest=largest, last=last), + ) + return NodeEmission( + functions=(offset, CFunction(name, definition)), + statements=( + call_kernel(name, [result.expr, source.expr, *grouping.arguments]), + ), + ) + + +def extremum_test(elem_type: int, *, largest: bool, last: bool) -> str: + """Whether `x` replaces `best`, under numpy's NaN-aware ordering. + + Shared with Hardmax, which is an ArgMax that writes a one-hot group rather than an index. + """ + comparison = (">" if largest else "<") + ("=" if last else "") + if elem_type not in FLOAT_TYPES: + return f"x {comparison} best" + if last: + # A NaN betters anything, itself included, so the last of them is what is reported. + return f"isnan(x) || (!isnan(best) && x {comparison} best)" + return f"!isnan(best) && (isnan(x) || x {comparison} best)" + + +def _register_reduction( + op_type: str, + build: KernelBuilder, + attribute_versions: tuple[int, ...], + input_versions: tuple[int, ...], +) -> None: + """Both axes conventions of one reduction, each at the revisions that take it.""" + register_kernel( + "", op_type, attribute_versions, partial(_from_attribute, build=build) + ) + register_kernel("", op_type, input_versions, partial(_from_input, build=build)) + + +# The family: how each op folds a group, and the revisions of both axes conventions. +_REDUCTIONS: tuple[tuple[str, Recipe, tuple[int, ...], tuple[int, ...]], ...] = ( + ("ReduceSum", _sum_fold, _SUM_ATTRIBUTE_VERSIONS, _SUM_INPUT_VERSIONS), + ("ReduceMean", _mean_fold, _ATTRIBUTE_VERSIONS, _INPUT_VERSIONS), + ("ReduceProd", _product_fold, _ATTRIBUTE_VERSIONS, _INPUT_VERSIONS), + ("ReduceL1", _absolute_sum_fold, _ATTRIBUTE_VERSIONS, _INPUT_VERSIONS), + ("ReduceL2", _euclidean_fold, _ATTRIBUTE_VERSIONS, _INPUT_VERSIONS), + ("ReduceLogSum", _log_sum_fold, _ATTRIBUTE_VERSIONS, _INPUT_VERSIONS), + ("ReduceSumSquare", _square_sum_fold, _ATTRIBUTE_VERSIONS, _INPUT_VERSIONS), + ( + "ReduceMax", + partial(_extremum_fold, largest=True), + _EXTREMUM_ATTRIBUTE_VERSIONS, + _EXTREMUM_INPUT_VERSIONS, + ), + ( + "ReduceMin", + partial(_extremum_fold, largest=False), + _EXTREMUM_ATTRIBUTE_VERSIONS, + _EXTREMUM_INPUT_VERSIONS, + ), +) + +for _op_type, _recipe, _attribute_versions, _input_versions in _REDUCTIONS: + _register_reduction( + _op_type, + partial(_fold_kernel, recipe=_recipe), + _attribute_versions, + _input_versions, + ) +_register_reduction( + "ReduceLogSumExp", _log_sum_exp_kernel, _ATTRIBUTE_VERSIONS, _INPUT_VERSIONS +) +register_kernel("", "ArgMax", _ARG_VERSIONS, partial(_arg_extremum, largest=True)) +register_kernel("", "ArgMin", _ARG_VERSIONS, partial(_arg_extremum, largest=False)) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/resize.py b/src/python/fnnx/extras/compilers/c/onnx/ops/resize.py new file mode 100644 index 0000000..5a6215c --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/resize.py @@ -0,0 +1,824 @@ +"""Resize: every output position read from the source coordinate its scale maps it to. + +One resize is a chain of one-dimensional ones — the interpolation is separable, so an N-d +result is the operand resized along one axis, then along the next — and that is how this is +emitted: one pass per named axis, through a pair of `double` working buffers, exactly as +ONNX's reference evaluator computes it. Within a pass every output position maps to a source +coordinate (that is what `coordinate_transformation_mode` names), takes a fixed number of +neighbouring elements around it, and weights them by the coefficients `mode` names — +`antialias` widening that footprint when the axis is being shrunk, `exclude_outside` dropping +the neighbours that fall off the end and renormalizing what is left. + +The geometry is read at run time, not baked in. `scales`, `sizes` and `roi` are operands, and +a model that computes one — the shape the corpus's own Resize tests are written in — makes +the result's extent a function of input data. What the artifact is compiled for is the result +shape ONNX inferred; the kernel derives the extents the operands ask for and refuses, through +the status enum, to write a result of any other shape. So the buffers stay static and a value +the artifact was not compiled for is reported rather than silently computed. + +`Upsample`, which ONNX deprecated in favour of `Resize`, is the same walk at the settings its +successor spells out: nearest neighbours taken at the floor of an asymmetrically mapped +coordinate. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from string import Template + +import numpy as np +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + c_type, + element_type_name, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + ScratchBuffer, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import checked_call, normalize_axis +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents + +# Only Resize-19 and Upsample-10 are claimed. They are the revisions the reference evaluator +# is version-faithful for, and 19 is the one every Resize test in the backend corpus imports; +# Upsample-9, which the corpus's own test imports, is claimed alongside 10 because ONNX +# revised nothing between them that this kernel reads. A model importing an older revision +# gets the unsupported-version error. +_RESIZE_VERSIONS = (19,) +_UPSAMPLE_VERSIONS = (9, 10) + +# The C values the kernel switches on, in the order its `switch` cases are written. +_MODES = {"nearest": 0, "linear": 1, "cubic": 2} +_NEAREST_MODES = { + "round_prefer_floor": 0, + "round_prefer_ceil": 1, + "floor": 2, + "ceil": 3, +} +_TRANSFORMS = { + "half_pixel": 0, + "half_pixel_symmetric": 1, + "pytorch_half_pixel": 2, + "align_corners": 3, + "asymmetric": 4, + "tf_crop_and_resize": 5, +} +_POLICIES = {"stretch": 0, "not_larger": 1, "not_smaller": 2} + +# The two modes Upsample carries over into its successor; the rest arrived after it. +_UPSAMPLE_MODES = ("nearest", "linear") + +# How far the largest extent the geometry may derive is allowed to reach before the kernel +# refuses it: comfortably past any buffer an artifact can hold, and well inside the range a +# `double` converts to an integer type exactly. +_EXTENT_LIMIT = 9.0e15 + +_TAPS_TEMPLATE = Template("""\ +static int $name(int mode, int antialias, double scale) +{ + double reach; + if (mode == $nearest) { + return 2; + } + if (!antialias) { + return (mode == $linear) ? 2 : 4; + } + /* Antialiasing widens the filter over the elements a shrinking axis merges together: + its footprint reaches `reach / scale` to each side of the coordinate instead of + `reach`. Growing an axis leaves it at its own width. */ + reach = (mode == $linear) ? 1.0 : 2.0; + return 2 - 2 * ((int)floor(-reach / ((scale < 1.0) ? scale : 1.0)) + 1); +}""") + +_COEFFICIENT_TEMPLATE = Template("""\ +static double $name( + int mode, + int nearest_mode, + int antialias, + double ratio, + double scale, + double cubic_a, + int tap, + int taps) +{ + double x; + double squared; + if (mode == $nearest) { + /* A coordinate that lands on an element takes that element, whichever way the + rounding rule would break a tie: the ratio is 1 there, and no rule applies. */ + if (ratio == 1.0) { + return (tap == 0) ? 0.0 : 1.0; + } + switch (nearest_mode) { + case $round_prefer_ceil: + return (tap == 0) ? (double)(ratio < 0.5) : (double)(ratio >= 0.5); + case $floor: + return (tap == 0) ? 1.0 : 0.0; + case $ceil: + return (tap == 0) ? 0.0 : 1.0; + default: + return (tap == 0) ? (double)(ratio <= 0.5) : (double)(ratio > 0.5); + } + } + if (!antialias) { + if (mode == $linear) { + return (tap == 0) ? 1.0 - ratio : ratio; + } + switch (tap) { + case 0: + x = ratio + 1.0; + return ((cubic_a * x - 5.0 * cubic_a) * x + 8.0 * cubic_a) * x + - 4.0 * cubic_a; + case 1: + return ((cubic_a + 2.0) * ratio - (cubic_a + 3.0)) * ratio * ratio + 1.0; + case 2: + x = 1.0 - ratio; + return ((cubic_a + 2.0) * x - (cubic_a + 3.0)) * x * x + 1.0; + default: + x = 2.0 - ratio; + return ((cubic_a * x - 5.0 * cubic_a) * x + 8.0 * cubic_a) * x + - 4.0 * cubic_a; + } + } + /* The antialiased filter is sampled at the operand's own spacing scaled down to the + result's, which is what spreads it over the elements being merged. */ + x = ((scale < 1.0) ? scale : 1.0) * ((double)(1 - taps / 2 + tap) - ratio); + if (mode == $linear) { + x = 1.0 - fabs(x); + return (x < 0.0) ? 0.0 : ((x > 1.0) ? 1.0 : x); + } + x = fabs(x); + squared = x * x; + if (x <= 1.0) { + return (cubic_a + 2.0) * (x * squared) - (cubic_a + 3.0) * squared + 1.0; + } + if (x < 2.0) { + return cubic_a * (x * squared) - 5.0 * cubic_a * squared + 8.0 * cubic_a * x + - 4.0 * cubic_a; + } + return 0.0; +}""") + +_COORDINATE_TEMPLATE = Template("""\ +static double $name( + int transform, + double position, + double scale, + double width, + double positions, + double extent, + double span, + double shift) +{ + switch (transform) { + case $half_pixel_symmetric: + return (width / 2.0) * (1.0 - extent / positions) + + (position + 0.5) / scale - 0.5; + case $pytorch_half_pixel: + return (positions == 1.0) ? -0.5 : (position + 0.5) / scale - 0.5; + case $align_corners: + return (positions == 1.0) + ? 0.0 + : position * (width - 1.0) / (positions - 1.0); + case $asymmetric: + return position / scale; + case $tf_crop_and_resize: + return ((positions == 1.0) + ? span * (width - 1.0) / 2.0 + : position * span * (width - 1.0) / (positions - 1.0)) + shift; + default: + return (position + 0.5) / scale - 0.5; + } +}""") + +# The extent one axis is asked for, and the scale that maps its coordinates. `scales` states +# the scale and the extent follows from it; `sizes` states the extent and the scale follows — +# unless a `keep_aspect_ratio_policy` overrides both with one scale shared across the axes. +_GEOMETRY_TEMPLATE = Template("""\ +static int $name( + const float* scales, + const int64_t* sizes, + size_t width, + int index, + int policy, + double policy_scale, + double* scale, + size_t* extent) +{ + double positions; + if (scales != NULL) { + *scale = (double)scales[index]; + positions = trunc(*scale * (double)width); + } else if (policy == $stretch) { + *scale = (double)sizes[index] / (double)width; + positions = (double)sizes[index]; + } else { + *scale = policy_scale; + positions = trunc(policy_scale * (double)width + 0.5); + } + /* Rejects a negative, infinite or absent-minded extent before the conversion below, + which C leaves undefined for anything outside `size_t`. */ + if (!(positions >= 0.0 && positions <= $limit)) { + return 1; + } + *extent = (size_t)positions; + return 0; +}""") + +_CLOSE_TEMPLATE = Template("""\ +static int $name(double left, double right) +{ + /* Python's `math.isclose` at its default tolerances, which is what the identity test + below is written against. */ + const double difference = fabs(left - right); + const double largest = (fabs(left) > fabs(right)) ? fabs(left) : fabs(right); + return difference <= 1e-9 * largest; +}""") + +_INTEGER_STORE_TEMPLATE = Template("""\ +static $element $name(double value) +{ + /* An interpolated value is rounded to the nearest even and then held inside the + element type's range. A NaN, which no comparison below admits, lands on the minimum + rather than in the undefined behaviour converting it would be. */ + const double rounded = rint(value); + if (!(rounded > $lower)) { + return $minimum; + } + if (rounded >= $upper) { + return $maximum; + } + return ($element)rounded; +}""") + +_FLOAT_STORE_TEMPLATE = Template("""\ +static float $name(double value) +{ + /* Saturating rather than overflowing: converting a value outside `float`'s range is + undefined in C, and the reference clips such a value instead of returning an + infinity. */ + if (value < $lower) { + return $minimum; + } + if (value > $upper) { + return $maximum; + } + return (float)value; +}""") + +_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const $roi_element* roi, + const float* scales, + const int64_t* sizes, + double* work, + double* spare, + size_t input_count, + size_t output_count, + int rank, + const size_t* input_shape, + const size_t* output_shape, + int axis_count, + const size_t* axes, + int mode, + int nearest_mode, + int transform, + int antialias, + int exclude_outside, + int policy, + double cubic_a, + double extrapolation_value) +{ + double* source = work; + double* target = spare; + double policy_scale = 0.0; + size_t index; + int step; + + /* One scale for every named axis, taken from the one that shrinks or grows the + operand most, is what a policy other than `stretch` asks for. */ + if (sizes != NULL && policy != $stretch) { + for (step = 0; step < axis_count; ++step) { + const double ratio = + (double)sizes[step] / (double)input_shape[axes[step]]; + const int closer = (policy == $not_larger) + ? (ratio < policy_scale) + : (ratio > policy_scale); + if (step == 0 || closer) { + policy_scale = ratio; + } + } + } + for (index = 0; index < input_count; ++index) { + source[index] = (double)in[index]; + } + + for (step = 0; step < axis_count; ++step) { + const size_t axis = axes[step]; + const size_t width = input_shape[axis]; + const size_t extent = output_shape[axis]; + const double edge = (double)width - 1.0; + double scale = 1.0; + double span = 1.0; + double shift = 0.0; + double positions; + size_t derived = 0; + size_t outer = 1; + size_t inner = 1; + size_t position; + int other; + int taps; + + if ($geometry( + scales, sizes, width, step, policy, policy_scale, &scale, &derived) != 0 + || derived != extent) { + return 1; + } + if (roi != NULL) { + /* In the region's own element type, as the reference reads it: a nearest + neighbour lands on the other side of a tie otherwise. */ + span = (double)(roi[axis_count + step] - roi[step]); + shift = (double)(roi[step] * ($roi_element)edge); + } + /* An axis at its own scale over its own region is the identity, and the reference + passes it through rather than mapping coordinates that would agree with it + everywhere but at the ends. */ + if ($close(scale, 1.0) && extent == width + && (roi == NULL + || (roi[step] == 0.0 + && $close((double)roi[axis_count + step], 1.0)))) { + continue; + } + + /* What lies around this axis at this point of the walk: the axes already resized + carry their result extent, the rest still carry the operand's. */ + for (other = 0; other < rank; ++other) { + size_t current = input_shape[other]; + int earlier; + for (earlier = 0; earlier < step; ++earlier) { + if (axes[earlier] == (size_t)other) { + current = output_shape[other]; + } + } + if ((size_t)other < axis) { + outer *= current; + } else if ((size_t)other > axis) { + inner *= current; + } + } + + taps = $taps(mode, antialias, scale); + positions = scale * (double)width; + for (position = 0; position < extent; ++position) { + const double placed = $coordinate( + transform, (double)position, scale, (double)width, positions, + (double)extent, span, shift); + size_t plane; + size_t offset; + int tap; + + /* A region mapping outside the operand carries the extrapolation value; the + test admits no coordinate a source position could not be taken from. */ + if (transform == $tf_crop_and_resize && !(placed >= 0.0 && placed <= edge)) { + for (plane = 0; plane < outer; ++plane) { + double* row = target + (plane * extent + position) * inner; + for (offset = 0; offset < inner; ++offset) { + row[offset] = extrapolation_value; + } + } + continue; + } + { + const double floored = floor(placed); + const double fraction = placed - floored; + /* The element the coordinate lands on is the one to its left, so a + coordinate that lands exactly on one carries a whole ratio, not none. */ + const double ratio = (fraction == 0.0) ? 1.0 : fraction; + const ptrdiff_t start = + (ptrdiff_t)floored - taps / 2 + ((fraction == 0.0) ? 0 : 1); + double denominator = 1.0; + + if (exclude_outside || antialias) { + double total = 0.0; + for (tap = 0; tap < taps; ++tap) { + const ptrdiff_t sampled = start + tap; + if (exclude_outside + && (sampled < 0 || sampled >= (ptrdiff_t)width)) { + continue; + } + total += $coefficient( + mode, nearest_mode, antialias, ratio, scale, cubic_a, + tap, taps); + } + denominator = (total == 0.0) ? 1.0 : total; + } + for (plane = 0; plane < outer; ++plane) { + double* row = target + (plane * extent + position) * inner; + for (offset = 0; offset < inner; ++offset) { + row[offset] = 0.0; + } + } + for (tap = 0; tap < taps; ++tap) { + const ptrdiff_t sampled = start + tap; + const int inside = (sampled >= 0) && (sampled < (ptrdiff_t)width); + /* Off the end, the operand's own edge element is read, which is what + padding it by repetition amounts to. */ + const size_t clamped = + inside ? (size_t)sampled : ((sampled < 0) ? 0 : width - 1); + const double raw = (exclude_outside && !inside) + ? 0.0 + : $coefficient( + mode, nearest_mode, antialias, ratio, scale, cubic_a, + tap, taps); + const double weight = raw / denominator; + for (plane = 0; plane < outer; ++plane) { + double* row = target + (plane * extent + position) * inner; + const double* taken = + source + (plane * width + clamped) * inner; + for (offset = 0; offset < inner; ++offset) { + row[offset] += weight * taken[offset]; + } + } + } + } + } + { + double* swapped = source; + source = target; + target = swapped; + } + } + for (index = 0; index < output_count; ++index) { + out[index] = $store; + } + return 0; +}""") + + +@dataclass(frozen=True) +class _Helpers: + """The shared functions the kernel calls, named so its call sites can reach them.""" + + taps: CFunction + coefficient: CFunction + coordinate: CFunction + geometry: CFunction + close: CFunction + store: CFunction | None + + @property + def functions(self) -> tuple[CFunction, ...]: + listed = ( + self.taps, + self.coefficient, + self.coordinate, + self.geometry, + self.close, + ) + return listed if self.store is None else (*listed, self.store) + + @property + def store_expression(self) -> str: + """How one interpolated value reaches the result buffer's element type.""" + return ( + "source[index]" + if self.store is None + else f"{self.store.name}(source[index])" + ) + + +@dataclass(frozen=True) +class _Options: + """What the node asks for, as the values the kernel switches on.""" + + mode: int + nearest_mode: int + transform: int + antialias: int + exclude_outside: int + policy: int + cubic_a: float + extrapolation_value: float + + @property + def arguments(self) -> list[str]: + return [ + str(self.mode), + str(self.nearest_mode), + str(self.transform), + str(self.antialias), + str(self.exclude_outside), + str(self.policy), + scalar_literal(self.cubic_a, TensorProto.DOUBLE), + scalar_literal(self.extrapolation_value, TensorProto.DOUBLE), + ] + + +def _resize(context: NodeContext) -> NodeEmission: + options = _Options( + mode=_choice(context, "mode", "nearest", _MODES), + nearest_mode=_choice( + context, "nearest_mode", "round_prefer_floor", _NEAREST_MODES + ), + transform=_choice( + context, "coordinate_transformation_mode", "half_pixel", _TRANSFORMS + ), + antialias=int(context.int_attribute("antialias") != 0), + exclude_outside=int(context.int_attribute("exclude_outside") != 0), + policy=_choice(context, "keep_aspect_ratio_policy", "stretch", _POLICIES), + cubic_a=context.float_attribute("cubic_coeff_a"), + extrapolation_value=context.float_attribute("extrapolation_value"), + ) + return _emit( + context, + options, + roi=_operand(context, 1), + scales=_operand(context, 2), + sizes=_operand(context, 3), + ) + + +def _upsample(context: NodeContext) -> NodeEmission: + """Upsample, which ONNX deprecated in favour of the Resize settings it spells out. + + Its successor's own specification records the equivalence: `asymmetric` is described + there as the coordinate mapping Resize-10 — the revision Upsample became — applies, and + `floor` as the neighbour it takes. + """ + options = _Options( + mode=_choice( + context, "mode", "nearest", {name: _MODES[name] for name in _UPSAMPLE_MODES} + ), + nearest_mode=_NEAREST_MODES["floor"], + transform=_TRANSFORMS["asymmetric"], + antialias=0, + exclude_outside=0, + policy=_POLICIES["stretch"], + cubic_a=0.0, + extrapolation_value=0.0, + ) + return _emit(context, options, roi=None, scales=_operand(context, 1), sizes=None) + + +def _emit( + context: NodeContext, + options: _Options, + *, + roi: TensorRef | None, + scales: TensorRef | None, + sizes: TensorRef | None, +) -> NodeEmission: + source = context.require_input(0) + result = context.require_output(0) + if len(result.shape) != len(source.shape): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` resizes `{source.name}` of " + f"shape {list(source.shape)} into `{result.name}` of shape " + f"{list(result.shape)}; a resize leaves the rank alone." + ) + if result.elem_type == TensorProto.BOOL: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` of a " + f"`{element_type_name(result.elem_type)}` tensor is not supported by the C " + "compiler; interpolating between two truth values has no defined result." + ) + if options.antialias and options.mode == _MODES["nearest"]: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` asks for `antialias` in " + "`nearest` mode, which ONNX defines only for the interpolating modes." + ) + + axes = _axes(context, len(source.shape)) + _verify_operands(context, source, result, axes, roi, scales, sizes) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + region = c_type(roi.elem_type if roi is not None else TensorProto.FLOAT) + name = f"{context.prefix}_resize_{element}_{region}" + helpers = _helpers(context, result.elem_type) + definition = _TEMPLATE.substitute( + name=name, + element=element, + roi_element=region, + store=helpers.store_expression, + geometry=helpers.geometry.name, + close=helpers.close.name, + taps=helpers.taps.name, + coefficient=helpers.coefficient.name, + coordinate=helpers.coordinate.name, + stretch=_POLICIES["stretch"], + not_larger=_POLICIES["not_larger"], + tf_crop_and_resize=_TRANSFORMS["tf_crop_and_resize"], + ) + work, spare = _scratch(context, source, result) + call = checked_call( + context, + name, + [ + result.expr, + source.expr, + "NULL" if roi is None else roi.expr, + "NULL" if scales is None else scales.expr, + "NULL" if sizes is None else sizes.expr, + work.symbol, + spare.symbol, + f"{source.elem_count}u", + f"{result.elem_count}u", + str(len(source.shape)), + extents(source.shape), + extents(result.shape), + str(len(axes)), + extents(axes), + *options.arguments, + ], + ) + return NodeEmission( + functions=(*helpers.functions, CFunction(name, definition)), + statements=(call,), + scratch=(work, spare), + ) + + +def _helpers(context: NodeContext, elem_type: int) -> _Helpers: + prefix = f"{context.prefix}_resize" + return _Helpers( + taps=CFunction( + f"{prefix}_taps", + _TAPS_TEMPLATE.substitute( + name=f"{prefix}_taps", + nearest=_MODES["nearest"], + linear=_MODES["linear"], + ), + ), + coefficient=CFunction( + f"{prefix}_coefficient", + _COEFFICIENT_TEMPLATE.substitute( + name=f"{prefix}_coefficient", + nearest=_MODES["nearest"], + linear=_MODES["linear"], + round_prefer_ceil=_NEAREST_MODES["round_prefer_ceil"], + floor=_NEAREST_MODES["floor"], + ceil=_NEAREST_MODES["ceil"], + ), + ), + coordinate=CFunction( + f"{prefix}_coordinate", + _COORDINATE_TEMPLATE.substitute(name=f"{prefix}_coordinate", **_TRANSFORMS), + ), + geometry=CFunction( + f"{prefix}_geometry", + _GEOMETRY_TEMPLATE.substitute( + name=f"{prefix}_geometry", + stretch=_POLICIES["stretch"], + limit=scalar_literal(_EXTENT_LIMIT, TensorProto.DOUBLE), + ), + ), + close=CFunction( + f"{prefix}_close", _CLOSE_TEMPLATE.substitute(name=f"{prefix}_close") + ), + store=_store_function(prefix, elem_type), + ) + + +def _store_function(prefix: str, elem_type: int) -> CFunction | None: + """How an interpolated value reaches the result's element type, where it has to narrow.""" + if elem_type == TensorProto.DOUBLE: + return None + name = f"{prefix}_store_{c_type(elem_type)}" + if elem_type == TensorProto.FLOAT: + limit = float(np.finfo("float32").max) + return CFunction( + name, + _FLOAT_STORE_TEMPLATE.substitute( + name=name, + lower=scalar_literal(-limit, TensorProto.DOUBLE), + upper=scalar_literal(limit, TensorProto.DOUBLE), + minimum=scalar_literal(-limit, TensorProto.FLOAT), + maximum=scalar_literal(limit, TensorProto.FLOAT), + ), + ) + info = np.iinfo(numpy_dtype_name(elem_type)) + return CFunction( + name, + _INTEGER_STORE_TEMPLATE.substitute( + name=name, + element=c_type(elem_type), + lower=scalar_literal(float(info.min), TensorProto.DOUBLE), + upper=scalar_literal(float(info.max), TensorProto.DOUBLE), + minimum=scalar_literal(info.min, elem_type), + maximum=scalar_literal(info.max, elem_type), + ), + ) + + +def _scratch( + context: NodeContext, source: TensorRef, result: TensorRef +) -> tuple[ScratchBuffer, ...]: + """The two working buffers one pass reads and the next writes. + + A pass over one axis reads the whole result of the pass before it, so neither can be + computed in place; the artifact allocates nothing, so the space is reserved at compile + time and counted in the reported footprint like every other buffer. Every intermediate + fits an operand whose every axis carries the larger of its two extents. + """ + count = math.prod( + max(before, after) for before, after in zip(source.shape, result.shape) + ) + return tuple( + ScratchBuffer(f"{context.prefix}_resize_{role}", TensorProto.DOUBLE, count) + for role in ("work", "spare") + ) + + +def _operand(context: NodeContext, index: int) -> TensorRef | None: + """Operand `index`, where one holding nothing at all reads as one left out. + + ONNX's own evaluator reads an empty `roi` that way, and exporters routinely pass an + empty `scales` alongside a `sizes` rather than leaving the position blank. + """ + operand = context.optional_input(index) + return operand if operand is not None and operand.elem_count > 0 else None + + +def _axes(context: NodeContext, rank: int) -> tuple[int, ...]: + """The axes the node resizes, in the order its operands describe them.""" + declared = context.attribute("axes", None) + if declared is None: + return tuple(range(rank)) + axes = tuple(normalize_axis(context, int(axis), rank) for axis in declared) + if len(set(axes)) != len(axes): + raise CompileError( + f"Node `{context.label}`: `axes` {[int(axis) for axis in declared]} names the " + "same dimension more than once." + ) + return axes + + +def _verify_operands( + context: NodeContext, + source: TensorRef, + result: TensorRef, + axes: tuple[int, ...], + roi: TensorRef | None, + scales: TensorRef | None, + sizes: TensorRef | None, +) -> None: + """Everything about this node the operands' shapes settle before it runs.""" + label = f"Node `{context.label}`: `{context.node.op_type}`" + if (scales is None) == (sizes is None): + raise CompileError( + f"{label} needs exactly one of `scales` and `sizes` to say what to resize to; " + f"this node passes {'both' if scales is not None else 'neither'}." + ) + for operand, role in ((scales, "scales"), (sizes, "sizes"), (roi, "roi")): + if operand is None: + continue + expected = 2 * len(axes) if role == "roi" else len(axes) + if operand.shape != (expected,): + raise CompileError( + f"{label} takes `{role}` as {expected} value(s) for the {len(axes)} " + f"axis/axes it resizes, but `{operand.name}` has shape " + f"{list(operand.shape)}." + ) + named = set(axes) + for axis, (before, after) in enumerate(zip(source.shape, result.shape)): + if axis not in named and before != after: + raise CompileError( + f"{label} does not resize axis {axis}, but `{source.name}` has " + f"{before} element(s) there against `{result.name}`'s {after}." + ) + if axis in named and before == 0 and after != 0: + raise CompileError( + f"{label} resizes axis {axis} of `{source.name}`, which holds no " + f"elements, into {after} element(s); there is nothing to interpolate." + ) + + +def _choice( + context: NodeContext, attribute: str, default: str, choices: dict[str, int] +) -> int: + value = context.attribute(attribute, default) + name = value.decode() if isinstance(value, bytes) else str(value) + if name not in choices: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` asks for `{attribute}` " + f"`{name}`, which is not one of the values ONNX defines for it " + f"({', '.join(f'`{choice}`' for choice in choices)})." + ) + return choices[name] + + +register_kernel("", "Resize", _RESIZE_VERSIONS, _resize) +register_kernel("", "Upsample", _UPSAMPLE_VERSIONS, _upsample) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/sampling.py b/src/python/fnnx/extras/compilers/c/onnx/ops/sampling.py new file mode 100644 index 0000000..bb70e7c --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/sampling.py @@ -0,0 +1,1054 @@ +"""The samplers: the ops that read an image at positions that fall between its elements. + +`GridSample` is the general one — every output position carries its own sampling coordinate, +normalized to [-1, 1] — and the other two are it at fixed geometries. `RoiAlign` samples a +regular grid inside each region of interest and folds what it reads into one value per bin; +`MaxRoiPool`, its predecessor, rounds each region to whole elements instead and takes the +largest in each bin, so it interpolates nothing. `AffineGrid` computes no samples at all: it +builds the coordinates a `GridSample` is then fed, which is why it sits here. + +Three things are shared, and emitted once per artifact whatever mixture of the four a model +holds: the reflection that folds a coordinate back between two borders, the resolution of one +tap's index under the padding mode, and the cubic weights. Everything else is per element +type, with the geometry reaching the kernels as call-site literals. + +A coordinate arrives at run time, so every conversion of one into an index is bounded first: +a value far outside the operand — or one that is not a number at all — is pulled to where the +sampling reads nothing rather than converted to whatever an out-of-range cast would give. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from string import Template + +import numpy as np +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + FLOAT_TYPES, + c_type, + element_type_name, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + checked_call, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents, math_suffix + +# GridSample arrived at 16, 20 generalized it past two spatial axes and renamed its modes, +# and 22 widened the element types; RoiAlign arrived at 10, 16 added the coordinate +# transformation and 22 widened the types; MaxRoiPool has 1 and 22. Only the newest revision +# of each is claimed: it is the one the reference evaluator is version-faithful for and the +# one every corpus test of these ops imports, so it is the only one anything can vouch for. +# A model importing an older one gets the unsupported-version error. AffineGrid has had a +# single revision since it arrived. +_VERSIONS = (22,) +_AFFINE_GRID_VERSIONS = (20,) + +# The C values the kernels switch on, in the order their branches are written. +_MODES = {"nearest": 0, "linear": 1, "cubic": 2} +_PADDING_MODES = {"zeros": 0, "border": 1, "reflection": 2} +_ROI_TRANSFORMS = {"half_pixel": 0, "output_half_pixel": 1} +_ROI_MODES = {"avg": 0, "max": 1} + +# How far a coordinate read at run time may reach before a kernel refuses it. Comfortably +# past any extent an artifact can hold, and small enough that the products the geometry +# forms out of it stay inside a 64-bit `ptrdiff_t`. +_COORDINATE_LIMIT = 1.0e9 + +_REFLECT_TEMPLATE = Template("""\ +static double $name(double value, double lower, double upper) +{ + double range = upper - lower; + double excess, periods, remainder; + /* A border with no room between its ends -- one element, measured corner to corner -- + leaves the single position it names. */ + if (!(range > 0.0)) { + return lower; + } + if (value < lower) { + excess = lower - value; + } else if (value > upper) { + excess = value - upper; + } else { + return value; + } + /* Every whole range crossed flips the direction the excess is measured in; the count is + kept in floating point so that a coordinate of any magnitude folds without a + conversion that would not be defined for it. */ + periods = floor(excess / range); + remainder = excess - periods * range; + if (fmod(periods, 2.0) != 0.0) { + return (value < lower) ? upper - remainder : lower + remainder; + } + return (value < lower) ? lower + remainder : upper - remainder; +}""") + +# Where one tap lands, once the padding mode has had its say. -1 is the answer `zeros` gives +# for a tap outside the operand: it reads nothing and contributes nothing. +_INDEX_TEMPLATE = Template("""\ +static ptrdiff_t $name( + ptrdiff_t index, + size_t extent, + int padding_mode, + int align_corners) +{ + /* An axis with no elements has nothing to clamp or reflect onto, whatever the mode. */ + if (extent == 0) { + return -1; + } + if (padding_mode == $zeros) { + return (index >= 0 && index < (ptrdiff_t)extent) ? index : -1; + } + if (padding_mode == $reflection) { + index = (ptrdiff_t)$reflect( + (double)index, + align_corners ? 0.0 : -0.5, + align_corners ? (double)extent - 1.0 : (double)extent - 0.5); + } + /* What `border` does, and what a reflected index needs anyway: the borders it folds + between reach half an element past the operand when the corners are not aligned. */ + if (index < 0) { + return 0; + } + return (index >= (ptrdiff_t)extent) ? (ptrdiff_t)extent - 1 : index; +}""") + +# The sampling coordinate one normalized grid value names, in the operand's own units. +_LOCATE_TEMPLATE = Template("""\ +static $coord $name( + $coord value, + size_t extent, + int align_corners, + int padding_mode, + int nearest) +{ + const double lower = align_corners ? 0.0 : -0.5; + const double upper = + align_corners ? (double)extent - 1.0 : (double)extent - 0.5; + const $coord reach = ($coord)extent + ($coord)8; + /* [-1, 1] spans the whole axis either corner to corner or edge to edge. */ + $coord x = align_corners + ? (value + $one) / ($coord)2 * (($coord)extent - $one) + : ((value + $one) * ($coord)extent - $one) / ($coord)2; + if (nearest) { + x = rint$f(x); + } + if ((double)x < lower || (double)x > upper) { + if (padding_mode == $border) { + x = (x < $zero) ? $zero : x; + x = (x > ($coord)extent - $one) ? ($coord)extent - $one : x; + } else if (padding_mode == $reflection) { + x = ($coord)$reflect((double)x, lower, upper); + } + } + /* `zeros` leaves the coordinate wherever it fell. Every tap around one this far out is + outside the operand whatever its exact value, so it is pulled to where that is still + true and the conversion to an index is defined -- as is a coordinate that is not a + number, which no comparison above holds for. */ + if (!(x > -reach)) { + return -reach; + } + return (x < reach) ? x : reach; +}""") + +# Keys' cubic convolution at a = -0.75, which GridSample fixes where Resize takes it as an +# attribute. The taps sit at -1, 0, 1 and 2 around the coordinate's floor. +_COEFFICIENT_TEMPLATE = Template("""\ +static $coord $name(int mode, $coord ratio, int tap) +{ + $coord x; + if (mode == $linear) { + return (tap == 0) ? $one - ratio : ratio; + } + switch (tap) { + case 0: + x = ratio + $one; + return ((($coord)-0.75 * x + ($coord)3.75) * x - ($coord)6) * x + ($coord)3; + case 1: + return (($coord)1.25 * ratio - ($coord)2.25) * ratio * ratio + $one; + case 2: + x = $one - ratio; + return (($coord)1.25 * x - ($coord)2.25) * x * x + $one; + default: + x = ($coord)2 - ratio; + return ((($coord)-0.75 * x + ($coord)3.75) * x - ($coord)6) * x + ($coord)3; + } +}""") + +# The parameters both GridSample kernels take after their buffers. +_GRID_PARAMETERS = """\ + size_t batch_count, + size_t channels, + size_t input_size, + size_t output_size, + int spatial_rank, + const size_t* input_shape, + int padding_mode, + int align_corners""" + +# Reading the grid at one output position: the coordinates are stored in the reverse of the +# operand's axis order, which is what the innermost subscript undoes. +_GRID_COORDINATE = """\ + const size_t extent = input_shape[axis]; + const $coord x = $locate( + coordinates[position * (size_t)spatial_rank + + (size_t)(spatial_rank - 1 - axis)], + extent, + align_corners, + padding_mode, + $nearest);""" + +_GRID_NEAREST_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + const $coord* grid, +$parameters) +{ + size_t batch, channel, position; + int axis; + for (batch = 0; batch < batch_count; ++batch) { + const $coord* coordinates = + grid + batch * output_size * (size_t)spatial_rank; + for (channel = 0; channel < channels; ++channel) { + const $element* plane = in + (batch * channels + channel) * input_size; + $element* result = out + (batch * channels + channel) * output_size; + for (position = 0; position < output_size; ++position) { + size_t offset = 0; + size_t stride = 1; + int outside = 0; + for (axis = spatial_rank - 1; axis >= 0; --axis) { +$coordinate + const ptrdiff_t resolved = + $index((ptrdiff_t)x, extent, padding_mode, align_corners); + if (resolved < 0) { + outside = 1; + } else { + offset += (size_t)resolved * stride; + } + stride *= extent; + } + result[position] = outside ? $zero : plane[offset]; + } + } + } +}""") + +_GRID_INTERPOLATE_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + const $coord* grid, +$parameters, + int mode) +{ + const int taps_per_axis = (mode == $linear) ? 2 : 4; + size_t batch, channel, position, tap, tap_count = 1; + int axis; + for (axis = 0; axis < spatial_rank; ++axis) { + tap_count *= (size_t)taps_per_axis; + } + for (batch = 0; batch < batch_count; ++batch) { + const $coord* coordinates = + grid + batch * output_size * (size_t)spatial_rank; + for (channel = 0; channel < channels; ++channel) { + const $element* plane = in + (batch * channels + channel) * input_size; + $element* result = out + (batch * channels + channel) * output_size; + for (position = 0; position < output_size; ++position) { + $element total = $zero; + for (tap = 0; tap < tap_count; ++tap) { + size_t remaining = tap; + size_t offset = 0; + size_t stride = 1; + $coord weight = $unit; + int outside = 0; + for (axis = spatial_rank - 1; axis >= 0; --axis) { +$coordinate + const $coord base = floor$f(x); + const int within = (int)(remaining % (size_t)taps_per_axis); + const ptrdiff_t resolved = $index( + (ptrdiff_t)base + within - ((taps_per_axis == 4) ? 1 : 0), + extent, + padding_mode, + align_corners); + remaining /= (size_t)taps_per_axis; + weight *= $coefficient(mode, x - base, within); + if (resolved < 0) { + outside = 1; + } else { + offset += (size_t)resolved * stride; + } + stride *= extent; + } + if (!outside) { + total += ($element)weight * plane[offset]; + } + } + result[position] = total; + } + } + } +}""") + +_AFFINE_COORDINATE_TEMPLATE = Template("""\ +static double $name(size_t index, size_t extent, int align_corners) +{ + double step; + if (align_corners) { + /* One position covers the whole axis, so there is no step to take along it. */ + if (extent < 2) { + return -1.0; + } + return -1.0 + (double)index * (2.0 / (double)(extent - 1)); + } + step = 2.0 / (double)extent; + return (-1.0 + step / 2.0) + (double)index * step; +}""") + +_AFFINE_GRID_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* theta, + size_t batch_count, + size_t output_size, + int spatial_rank, + const size_t* spatial_shape, + int align_corners) +{ + size_t batch, position; + int axis, row; + for (batch = 0; batch < batch_count; ++batch) { + for (position = 0; position < output_size; ++position) { + for (row = 0; row < spatial_rank; ++row) { + const $element* weights = + theta + ((batch * (size_t)spatial_rank) + (size_t)row) + * (size_t)(spatial_rank + 1); + size_t remainder = position; + double sum = 0.0; + /* The homogeneous coordinate runs the axes backwards -- the last spatial + axis first -- and closes with the constant one. */ + for (axis = spatial_rank - 1; axis >= 0; --axis) { + const size_t extent = spatial_shape[axis]; + const size_t index = remainder % extent; + remainder /= extent; + sum += (double)weights[spatial_rank - 1 - axis] + * $coordinate(index, extent, align_corners); + } + sum += (double)weights[spatial_rank]; + out[(batch * output_size + position) * (size_t)spatial_rank + (size_t)row] = + ($element)sum; + } + } + } +}""") + +_ROI_ALIGN_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const $element* rois, + const int64_t* batch_indices, + size_t roi_count, + size_t batch_count, + size_t channels, + size_t height, + size_t width, + size_t pooled_height, + size_t pooled_width, + int sampling_ratio, + int half_pixel, + int max_mode, + $element spatial_scale) +{ + const size_t plane_size = height * width; + size_t roi, channel, row, column; + for (roi = 0; roi < roi_count; ++roi) { + const $element* box = rois + roi * 4; + const int64_t batch = batch_indices[roi]; + const $element shift = half_pixel ? ($element)0.5 : $zero; + const $element start_column = box[0] * spatial_scale - shift; + const $element start_row = box[1] * spatial_scale - shift; + $element box_width = box[2] * spatial_scale - shift - start_column; + $element box_height = box[3] * spatial_scale - shift - start_row; + $element bin_width, bin_height, count; + ptrdiff_t grid_rows, grid_columns, down, right; + if (batch < 0 || (uint64_t)batch >= (uint64_t)batch_count) { + return 1; + } + if (!half_pixel) { + /* A region that measures less than one element is widened to one. */ + box_width = (box_width < $one) ? $one : box_width; + box_height = (box_height < $one) ? $one : box_height; + } + bin_height = box_height / ($element)pooled_height; + bin_width = box_width / ($element)pooled_width; + if (sampling_ratio > 0) { + grid_rows = sampling_ratio; + grid_columns = sampling_ratio; + } else { + /* Enough samples per bin to cover every element it spans. */ + const double rows = ceil((double)bin_height); + const double columns = ceil((double)bin_width); + if (!(rows >= -$limit && rows <= $limit) + || !(columns >= -$limit && columns <= $limit)) { + return 1; + } + grid_rows = (ptrdiff_t)rows; + grid_columns = (ptrdiff_t)columns; + } + count = ($element)((grid_rows * grid_columns > 1) ? grid_rows * grid_columns : 1); + for (channel = 0; channel < channels; ++channel) { + const $element* plane = + in + ((size_t)batch * channels + channel) * plane_size; + $element* result = + out + (roi * channels + channel) * pooled_height * pooled_width; + for (row = 0; row < pooled_height; ++row) { + for (column = 0; column < pooled_width; ++column) { + $element total = $zero; + int seen = 0; + for (down = 0; down < grid_rows; ++down) { + const $element y = start_row + ($element)row * bin_height + + (($element)down + ($element)0.5) * bin_height + / ($element)grid_rows; + for (right = 0; right < grid_columns; ++right) { + const $element x = start_column + ($element)column * bin_width + + (($element)right + ($element)0.5) * bin_width + / ($element)grid_columns; + ptrdiff_t low_row = 0, high_row = 0; + ptrdiff_t low_column = 0, high_column = 0; + $element w1 = $zero, w2 = $zero, w3 = $zero, w4 = $zero; + $element p1, p2, p3, p4; + /* A sample more than one element outside the feature map reads + nothing; so does one that is not a number at all. */ + if (y >= -$one && y <= ($element)height + && x >= -$one && x <= ($element)width) { + $element sampled_row = (y < $zero) ? $zero : y; + $element sampled_column = (x < $zero) ? $zero : x; + $element ly, lx; + low_row = (ptrdiff_t)sampled_row; + low_column = (ptrdiff_t)sampled_column; + if (low_row >= (ptrdiff_t)height - 1) { + high_row = low_row = (ptrdiff_t)height - 1; + sampled_row = ($element)low_row; + } else { + high_row = low_row + 1; + } + if (low_column >= (ptrdiff_t)width - 1) { + high_column = low_column = (ptrdiff_t)width - 1; + sampled_column = ($element)low_column; + } else { + high_column = low_column + 1; + } + ly = sampled_row - ($element)low_row; + lx = sampled_column - ($element)low_column; + w1 = ($one - ly) * ($one - lx); + w2 = ($one - ly) * lx; + w3 = ly * ($one - lx); + w4 = ly * lx; + } + p1 = w1 * plane[(size_t)low_row * width + (size_t)low_column]; + p2 = w2 * plane[(size_t)low_row * width + (size_t)high_column]; + p3 = w3 * plane[(size_t)high_row * width + (size_t)low_column]; + p4 = w4 * plane[(size_t)high_row * width + (size_t)high_column]; + if (max_mode) { + $element best = p1; + best = (p2 > best) ? p2 : best; + best = (p3 > best) ? p3 : best; + best = (p4 > best) ? p4 : best; + if (!seen || best > total) { + total = best; + seen = 1; + } + } else { + total += p1 + p2 + p3 + p4; + } + } + } + result[row * pooled_width + column] = + max_mode ? total : total / count; + } + } + } + } + return 0; +}""") + +_MAX_ROI_POOL_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* in, + const $element* rois, + size_t roi_count, + size_t batch_count, + size_t channels, + size_t height, + size_t width, + size_t pooled_height, + size_t pooled_width, + $element spatial_scale) +{ + const size_t plane_size = height * width; + size_t roi, channel, row, column; + for (roi = 0; roi < roi_count; ++roi) { + const $element* box = rois + roi * 5; + const double batch = (double)box[0]; + double corners[4]; + ptrdiff_t start_column, start_row, box_width, box_height; + $element bin_height, bin_width; + int corner; + if (!(batch >= 0.0 && batch < (double)batch_count)) { + return 1; + } + /* The region is rounded to whole elements before anything is pooled, which is what + separates this op from the RoiAlign that replaced it. */ + for (corner = 0; corner < 4; ++corner) { + corners[corner] = (double)round$f(box[corner + 1] * spatial_scale); + if (!(corners[corner] >= -$limit && corners[corner] <= $limit)) { + return 1; + } + } + start_column = (ptrdiff_t)corners[0]; + start_row = (ptrdiff_t)corners[1]; + box_width = (ptrdiff_t)corners[2] - start_column + 1; + box_height = (ptrdiff_t)corners[3] - start_row + 1; + box_width = (box_width > 1) ? box_width : 1; + box_height = (box_height > 1) ? box_height : 1; + bin_height = ($element)box_height / ($element)pooled_height; + bin_width = ($element)box_width / ($element)pooled_width; + for (channel = 0; channel < channels; ++channel) { + const $element* plane = + in + ((size_t)batch * channels + channel) * plane_size; + $element* result = + out + (roi * channels + channel) * pooled_height * pooled_width; + for (row = 0; row < pooled_height; ++row) { + for (column = 0; column < pooled_width; ++column) { + const ptrdiff_t first_row = $clamp( + (ptrdiff_t)floor$f(($element)row * bin_height) + start_row, + (ptrdiff_t)height); + const ptrdiff_t last_row = $clamp( + (ptrdiff_t)ceil$f(($element)(row + 1) * bin_height) + start_row, + (ptrdiff_t)height); + const ptrdiff_t first_column = $clamp( + (ptrdiff_t)floor$f(($element)column * bin_width) + start_column, + (ptrdiff_t)width); + const ptrdiff_t last_column = $clamp( + (ptrdiff_t)ceil$f(($element)(column + 1) * bin_width) + + start_column, + (ptrdiff_t)width); + /* A bin the region does not reach pools nothing and stays at zero. One + that does starts from the lowest finite value of its type rather than + from negative infinity: that is the floor Caffe's ROI pooling -- which + is what ONNX inherited this op from, and what onnxruntime, its only + implementation, still computes -- pools down to. */ + $element best = $lowest; + ptrdiff_t sampled_row, sampled_column; + if (last_row <= first_row || last_column <= first_column) { + result[row * pooled_width + column] = $zero; + continue; + } + for (sampled_row = first_row; sampled_row < last_row; ++sampled_row) { + for (sampled_column = first_column; + sampled_column < last_column; + ++sampled_column) { + const $element value = + plane[(size_t)sampled_row * width + + (size_t)sampled_column]; + /* The later of the two wins unless it is strictly smaller, + which is what `std::max` -- and so onnxruntime, the only + implementation ONNX has for this op -- folds a window with. + It is observable wherever one holds a NaN, which compares + smaller than nothing. */ + best = (value < best) ? best : value; + } + } + result[row * pooled_width + column] = best; + } + } + } + } + return 0; +}""") + +_CLAMP_TEMPLATE = Template("""\ +static ptrdiff_t $name(ptrdiff_t value, ptrdiff_t extent) +{ + if (value < 0) { + return 0; + } + return (value > extent) ? extent : value; +}""") + + +# -------------------------------------------------------------------------------------- +# GridSample +# -------------------------------------------------------------------------------------- + + +def _grid_sample(context: NodeContext) -> NodeEmission: + """GridSample: every output position read at the coordinate the grid names for it.""" + source = context.require_input(0) + grid = context.require_input(1) + result = context.require_output(0) + rank = _grid_geometry(context, source, grid, result) + mode = _mode(context) + padding_mode = _choice(context, "padding_mode", "zeros", _PADDING_MODES) + align_corners = int(context.int_attribute("align_corners") != 0) + _require_float(context, grid) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + coord = c_type(grid.elem_type) + padding = _padding_helpers(context, grid.elem_type) + arguments = [ + result.expr, + source.expr, + grid.expr, + f"{source.shape[0]}u", + f"{source.shape[1]}u", + f"{math.prod(source.shape[2:])}u", + f"{math.prod(result.shape[2:])}u", + str(rank), + extents(source.shape[2:]), + str(padding_mode), + str(align_corners), + ] + if mode == _MODES["nearest"]: + name = f"{context.prefix}_gridsample_nearest_{element}_{coord}" + definition = _GRID_NEAREST_TEMPLATE.substitute( + name=name, + element=element, + coord=coord, + parameters=_GRID_PARAMETERS, + coordinate=padding.coordinate(nearest=1), + index=padding.index.name, + zero=scalar_literal(0, result.elem_type), + ) + return NodeEmission( + functions=(*padding.functions, CFunction(name, definition)), + statements=(call_kernel(name, arguments),), + ) + + coefficient = _coefficient(context, grid.elem_type) + name = f"{context.prefix}_gridsample_{element}_{coord}" + definition = _GRID_INTERPOLATE_TEMPLATE.substitute( + name=name, + element=element, + coord=coord, + parameters=_GRID_PARAMETERS, + coordinate=padding.coordinate(nearest=0), + index=padding.index.name, + coefficient=coefficient.name, + linear=_MODES["linear"], + f=math_suffix(grid.elem_type), + zero=scalar_literal(0, result.elem_type), + unit=scalar_literal(1, grid.elem_type), + ) + return NodeEmission( + functions=(*padding.functions, coefficient, CFunction(name, definition)), + statements=(call_kernel(name, [*arguments, str(mode)]),), + ) + + +def _mode(context: NodeContext) -> int: + """The interpolation GridSample asks for, refused where its element type cannot hold it. + + `nearest` reads one element and computes nothing, so it serves every type the schema + allows; the other two weight the elements around a coordinate, which is arithmetic an + integer tensor cannot carry — ONNX's own reference truncates the weights to integers + there rather than defining anything usable. + """ + mode = _choice(context, "mode", "linear", _MODES) + elem_type = context.require_output(0).elem_type + if mode == _MODES["nearest"] or elem_type in FLOAT_TYPES: + return mode + name = next(choice for choice, value in _MODES.items() if value == mode) + raise CompileError( + f"Node `{context.label}`: `GridSample` in `{name}` mode weights the elements " + f"around each coordinate, which a `{element_type_name(elem_type)}` tensor cannot " + "hold; only `nearest` mode is supported for it." + ) + + +def _grid_geometry( + context: NodeContext, source: TensorRef, grid: TensorRef, result: TensorRef +) -> int: + """The number of axes the grid samples along, once every operand agrees on it.""" + rank = len(source.shape) - 2 + if rank < 1: + raise CompileError( + f"Node `{context.label}`: `GridSample` samples a batch of multi-channel " + f"signals — a tensor of rank 3 or more — but `{source.name}` has shape " + f"{list(source.shape)}." + ) + if len(grid.shape) != rank + 2 or grid.shape[0] != source.shape[0]: + raise CompileError( + f"Node `{context.label}`: `GridSample` reads `{grid.name}` of shape " + f"{list(grid.shape)} as one coordinate per sampled position of " + f"`{source.name}` of shape {list(source.shape)}; ONNX defines it as the batch " + f"of `{source.name}`, one axis per sampled position, and a trailing axis of " + f"{rank} coordinate(s)." + ) + if grid.shape[-1] != rank: + raise CompileError( + f"Node `{context.label}`: `GridSample` samples {rank} spatial axis/axes of " + f"`{source.name}`, but `{grid.name}` carries {grid.shape[-1]} coordinate(s) " + "per position." + ) + verify_shape(context, result, (*source.shape[:2], *grid.shape[1:-1])) + return rank + + +# -------------------------------------------------------------------------------------- +# AffineGrid +# -------------------------------------------------------------------------------------- + + +def _affine_grid(context: NodeContext) -> NodeEmission: + """AffineGrid: the coordinates an affine transform maps a regular grid onto.""" + theta = context.require_input(0) + result = context.require_output(0) + size = context.constant_input(1) + if size is None: + raise CompileError( + f"Node `{context.label}`: `AffineGrid` takes the grid's shape from " + f"`{context.require_input(1).name}`, which is not known at compile time; the " + "shape of the result then depends on input data, which the C compiler cannot " + "compile." + ) + _require_float(context, theta) + spatial = tuple(int(extent) for extent in size.reshape(-1))[2:] + rank = len(spatial) + if rank < 1 or any(extent < 0 for extent in spatial): + raise CompileError( + f"Node `{context.label}`: `AffineGrid` was given a size of " + f"{[int(extent) for extent in size.reshape(-1)]}; ONNX defines it as a batch, " + "a channel count and one nonnegative extent per spatial axis." + ) + batch = int(size.reshape(-1)[0]) + if theta.shape != (batch, rank, rank + 1): + raise CompileError( + f"Node `{context.label}`: `AffineGrid` maps {rank} spatial axis/axes of a " + f"batch of {batch}, which ONNX defines as a transform of shape " + f"{[batch, rank, rank + 1]}, but `{theta.name}` has shape " + f"{list(theta.shape)}." + ) + verify_shape(context, result, (batch, *spatial, rank)) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + coordinate = CFunction( + f"{context.prefix}_affine_coordinate", + _AFFINE_COORDINATE_TEMPLATE.substitute( + name=f"{context.prefix}_affine_coordinate" + ), + ) + name = f"{context.prefix}_affinegrid_{element}" + definition = _AFFINE_GRID_TEMPLATE.substitute( + name=name, element=element, coordinate=coordinate.name + ) + call = call_kernel( + name, + [ + result.expr, + theta.expr, + f"{batch}u", + f"{math.prod(spatial)}u", + str(rank), + extents(spatial), + str(int(context.int_attribute("align_corners") != 0)), + ], + ) + return NodeEmission( + functions=(coordinate, CFunction(name, definition)), statements=(call,) + ) + + +# -------------------------------------------------------------------------------------- +# The region-of-interest poolings +# -------------------------------------------------------------------------------------- + + +def _roi_align(context: NodeContext) -> NodeEmission: + """RoiAlign: each region divided into bins, each bin folding a grid of samples.""" + source = context.require_input(0) + rois = context.require_input(1) + indices = context.require_input(2) + result = context.require_output(0) + height, width = _feature_map(context, source) + pooled = ( + context.int_attribute("output_height"), + context.int_attribute("output_width"), + ) + _require_float(context, rois) + _verify_regions(context, rois, indices, columns=4) + verify_shape(context, result, (rois.shape[0], source.shape[1], *pooled)) + sampling_ratio = context.int_attribute("sampling_ratio") + mode = _choice(context, "mode", "avg", _ROI_MODES) + transform = _choice( + context, "coordinate_transformation_mode", "half_pixel", _ROI_TRANSFORMS + ) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + _require_samples(context, source) + + element = c_type(result.elem_type) + name = f"{context.prefix}_roialign_{element}" + definition = _ROI_ALIGN_TEMPLATE.substitute( + name=name, + element=element, + one=scalar_literal(1, result.elem_type), + zero=scalar_literal(0, result.elem_type), + limit=scalar_literal(_COORDINATE_LIMIT, TensorProto.DOUBLE), + ) + call = checked_call( + context, + name, + [ + result.expr, + source.expr, + rois.expr, + indices.expr, + f"{rois.shape[0]}u", + f"{source.shape[0]}u", + f"{source.shape[1]}u", + f"{height}u", + f"{width}u", + f"{pooled[0]}u", + f"{pooled[1]}u", + str(sampling_ratio), + str(int(transform == _ROI_TRANSFORMS["half_pixel"])), + str(mode), + scalar_literal(context.float_attribute("spatial_scale"), result.elem_type), + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _max_roi_pool(context: NodeContext) -> NodeEmission: + """MaxRoiPool: each region rounded to whole elements, then max-pooled into its bins.""" + source = context.require_input(0) + rois = context.require_input(1) + result = context.require_output(0) + height, width = _feature_map(context, source) + pooled = context.attribute("pooled_shape", None) + if pooled is None: + raise CompileError( + f"Node `{context.label}`: `MaxRoiPool` states no `pooled_shape`, which ONNX " + "defines as a required attribute." + ) + pooled = tuple(int(extent) for extent in pooled) + if len(pooled) != 2 or any(extent < 1 for extent in pooled): + raise CompileError( + f"Node `{context.label}`: `MaxRoiPool` was given `pooled_shape` " + f"{list(pooled)}; ONNX defines it as a positive height and width." + ) + _require_float(context, rois) + _verify_regions(context, rois, None, columns=5) + verify_shape(context, result, (rois.shape[0], source.shape[1], *pooled)) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + _require_samples(context, source) + + element = c_type(result.elem_type) + clamp = CFunction( + f"{context.prefix}_roi_clamp", + _CLAMP_TEMPLATE.substitute(name=f"{context.prefix}_roi_clamp"), + ) + name = f"{context.prefix}_maxroipool_{element}" + definition = _MAX_ROI_POOL_TEMPLATE.substitute( + name=name, + element=element, + clamp=clamp.name, + f=math_suffix(result.elem_type), + zero=scalar_literal(0, result.elem_type), + lowest=scalar_literal( + -float(np.finfo(numpy_dtype_name(result.elem_type)).max), result.elem_type + ), + limit=scalar_literal(_COORDINATE_LIMIT, TensorProto.DOUBLE), + ) + call = checked_call( + context, + name, + [ + result.expr, + source.expr, + rois.expr, + f"{rois.shape[0]}u", + f"{source.shape[0]}u", + f"{source.shape[1]}u", + f"{height}u", + f"{width}u", + f"{pooled[0]}u", + f"{pooled[1]}u", + scalar_literal(context.float_attribute("spatial_scale"), result.elem_type), + ], + ) + return NodeEmission( + functions=(clamp, CFunction(name, definition)), statements=(call,) + ) + + +def _feature_map(context: NodeContext, source: TensorRef) -> tuple[int, int]: + if len(source.shape) != 4: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` pools regions of a batch of " + f"images — a tensor of rank 4 — but `{source.name}` has shape " + f"{list(source.shape)}." + ) + return source.shape[2], source.shape[3] + + +def _verify_regions( + context: NodeContext, rois: TensorRef, indices: TensorRef | None, *, columns: int +) -> None: + if len(rois.shape) != 2 or rois.shape[1] != columns: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads `{rois.name}` of " + f"shape {list(rois.shape)} as its regions; ONNX defines them as one row of " + f"{columns} value(s) per region." + ) + if indices is not None and indices.shape != (rois.shape[0],): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` takes one batch index per " + f"region, but `{indices.name}` has shape {list(indices.shape)} against " + f"{rois.shape[0]} region(s)." + ) + + +def _require_samples(context: NodeContext, source: TensorRef) -> None: + """Refuse to pool regions of a feature map that holds no elements to sample. + + Every region samples the map whatever it holds, and ONNX's own reference implementation + indexes past the end of an empty one rather than defining a value for it, so there is + nothing to compile against. + """ + if math.prod(source.shape[2:]) == 0: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` pools regions of " + f"`{source.name}` of shape {list(source.shape)}, which holds no elements to " + "sample." + ) + + +# -------------------------------------------------------------------------------------- +# Shared emission +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _Padding: + """What turns a grid value into an element of the operand, under the padding mode.""" + + reflect: CFunction + locate: CFunction + index: CFunction + coord: str + + @property + def functions(self) -> tuple[CFunction, ...]: + return (self.reflect, self.locate, self.index) + + def coordinate(self, *, nearest: int) -> str: + """Reading one output position's coordinate along one axis, as kernel body text.""" + return Template(_GRID_COORDINATE).substitute( + locate=self.locate.name, nearest=nearest, coord=self.coord + ) + + +def _padding_helpers(context: NodeContext, elem_type: int) -> _Padding: + coord = c_type(elem_type) + reflect = CFunction( + f"{context.prefix}_sample_reflect", + _REFLECT_TEMPLATE.substitute(name=f"{context.prefix}_sample_reflect"), + ) + locate = f"{context.prefix}_sample_locate_{coord}" + index = f"{context.prefix}_sample_index" + return _Padding( + reflect=reflect, + locate=CFunction( + locate, + _LOCATE_TEMPLATE.substitute( + name=locate, + coord=coord, + reflect=reflect.name, + border=_PADDING_MODES["border"], + reflection=_PADDING_MODES["reflection"], + f=math_suffix(elem_type), + one=scalar_literal(1, elem_type), + zero=scalar_literal(0, elem_type), + ), + ), + index=CFunction( + index, + _INDEX_TEMPLATE.substitute( + name=index, + reflect=reflect.name, + zeros=_PADDING_MODES["zeros"], + reflection=_PADDING_MODES["reflection"], + ), + ), + coord=coord, + ) + + +def _coefficient(context: NodeContext, elem_type: int) -> CFunction: + name = f"{context.prefix}_sample_coefficient_{c_type(elem_type)}" + return CFunction( + name, + _COEFFICIENT_TEMPLATE.substitute( + name=name, + coord=c_type(elem_type), + linear=_MODES["linear"], + one=scalar_literal(1, elem_type), + ), + ) + + +def _require_float(context: NodeContext, operand: TensorRef) -> None: + """Refuse an operand of coordinates the emitted arithmetic would not compute in. + + ONNX defines every one of these as floating-point; a model declaring one otherwise is + rejected rather than served with integer arithmetic that would silently truncate. + """ + if operand.elem_type not in FLOAT_TYPES: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads `{operand.name}` as " + f"coordinates, which ONNX defines as floating-point, but it holds " + f"`{element_type_name(operand.elem_type)}`." + ) + + +def _choice( + context: NodeContext, attribute: str, default: str, choices: dict[str, int] +) -> int: + value = context.attribute(attribute, default) + name = value.decode() if isinstance(value, bytes) else str(value) + if name not in choices: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` asks for `{attribute}` " + f"`{name}`, which is not one of the values ONNX defines for it " + f"({', '.join(f'`{choice}`' for choice in choices)})." + ) + return choices[name] + + +register_kernel("", "GridSample", _VERSIONS, _grid_sample) +register_kernel("", "AffineGrid", _AFFINE_GRID_VERSIONS, _affine_grid) +register_kernel("", "RoiAlign", _VERSIONS, _roi_align) +register_kernel("", "MaxRoiPool", _VERSIONS, _max_roi_pool) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/scatter.py b/src/python/fnnx/extras/compilers/c/onnx/ops/scatter.py new file mode 100644 index 0000000..960fbdd --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/scatter.py @@ -0,0 +1,542 @@ +"""The ops that write into a copy of a tensor at positions decided at run time. + +ScatterElements, ScatterND and TensorScatter all do the same shape of work: the result is the +first operand copied, and then some of its elements are overwritten out of a second operand, +at positions a third names. The positions are values rather than shapes, so the addressing is +a loop; what stays static is the result's shape, which is the copied operand's own. Each op +is emitted as a `memcpy` of the whole operand followed by one kernel walking the updates. + +An index comes from the caller, so every one of them is normalized the way ONNX defines it — +a negative index counted back from the end of the axis — and then bounds checked: a kernel +returns nonzero for an index outside its axis and the entrypoint passes that on as an +argument error, rather than writing past a buffer. + +`reduction` says what an update does to the element already in the result, and ONNX's two +families disagree about one case of it. `ScatterElements` is defined by a reference that +folds with Python's own `max`/`min`, which keep the value already in the result whenever a +comparison against a NaN comes out false; `ScatterND`'s folds with `np.maximum`/`np.minimum`, +which propagate a NaN from either side. Both op documents say only "max" and "min", so each +kernel follows the reference implementation of its own op — the only thing that states what +these two do with a NaN at all. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from string import Template + +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + copy_tensor, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + checked_call, + kernel_name, + normalize_axis, + row_major_strides, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import combiner, extents + +# Only the newest revision of ScatterElements and ScatterND is claimed: it is the one the +# reference evaluator is version-faithful for and the one every corpus test of them imports, +# so it is the only one anything can vouch for. Both arrived at 11, took `reduction` with +# `add` and `mul` at 16 and gained `max` and `min` at 18. TensorScatter has had a single +# revision since it arrived at 24. A model importing an older one gets the +# unsupported-version error. +_SCATTER_ELEMENTS_VERSIONS = (18,) +_SCATTER_ND_VERSIONS = (18,) +_TENSOR_SCATTER_VERSIONS = (24,) + +# Scatter, which ONNX deprecated in favour of ScatterElements, is that op before `reduction` +# was added to it, so the same generator serves it. 9 is the revision the corpus's own tests +# select — they import opset 10 — and 11 is the deprecating revision, which changed nothing +# else about the op; `test_the_two_scatter_revisions_are_one_op` compares the two schemas +# rather than taking that on trust. +_SCATTER_VERSIONS = (9, 11) + +# The opset each `reduction` value arrived at, by the revision that added it. +_REDUCTION_VERSIONS = ((16, ("add", "mul")), (18, ("max", "min"))) + +_TENSOR_SCATTER_MODES = ("linear", "circular") + + +@dataclass(frozen=True) +class _Fold: + """How an update is combined with the element already in the result.""" + + expression: str + helpers: tuple[CFunction, ...] = () + + +_SCATTER_ELEMENTS_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* updates, + const $index* indices, + size_t count, + int rank, + const size_t* shape, + const size_t* strides, + int axis, + size_t extent) +{ + size_t index; + for (index = 0; index < count; ++index) { + ptrdiff_t position = (ptrdiff_t)indices[index]; + size_t remainder = index; + size_t offset = 0; + int walked; + for (walked = rank - 1; walked >= 0; --walked) { + const size_t coordinate = remainder % shape[walked]; + remainder /= shape[walked]; + if (walked != axis) { + offset += coordinate * strides[walked]; + } + } + if (position < 0) { + position += (ptrdiff_t)extent; + } + if (position < 0 || position >= (ptrdiff_t)extent) { + return 1; + } + offset += (size_t)position * strides[axis]; + out[offset] = $fold; + } + return 0; +}""") + +_SCATTER_ND_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* updates, + const $index* indices, + size_t rows, + size_t depth, + size_t slice_size, + const size_t* extents, + const size_t* strides) +{ + size_t row, level, element; + for (row = 0; row < rows; ++row) { + size_t offset = 0; + for (level = 0; level < depth; ++level) { + ptrdiff_t position = (ptrdiff_t)indices[row * depth + level]; + if (position < 0) { + position += (ptrdiff_t)extents[level]; + } + if (position < 0 || position >= (ptrdiff_t)extents[level]) { + return 1; + } + offset += (size_t)position * strides[level]; + } + for (element = 0; element < slice_size; ++element) { + out[offset + element] = $fold; + } + } + return 0; +}""") + +# The write index is validated once per sample rather than per step: with it inside the axis, +# no `written_at + step` the loop forms can leave `ptrdiff_t`, and running off the end of the +# axis part-way through the update is what the per-step check catches. +_TENSOR_SCATTER_LINEAR_TEMPLATE = Template("""\ +static int $name( + $element* out, + const $element* update,$index_parameters + size_t prefix_count, + size_t sequence_length, + size_t max_sequence_length, + size_t block) +{ + size_t prefix, step; + for (prefix = 0; prefix < prefix_count; ++prefix) { + const ptrdiff_t written_at = $written_at; + if (written_at < 0 || written_at >= (ptrdiff_t)max_sequence_length) { + return 1; + } + for (step = 0; step < sequence_length; ++step) { + const size_t position = (size_t)written_at + step; + if (position >= max_sequence_length) { + return 1; + } + memcpy( + out + (prefix * max_sequence_length + position) * block, + update + (prefix * sequence_length + step) * block, + block * sizeof(*out)); + } + } + return 0; +}""") + +# ONNX defines the circular mode by taking the whole cache coordinate modulo the sequence +# capacity — `np.mod(np.asarray(cache_idx), max_sequence_length)` in the op's own pseudocode, +# which is what its reference implementation runs — so the coordinates before the sequence +# axis wrap along with the write index. Hence the destination is recomposed from the wrapped +# coordinates rather than being the sample's own base offset. Wrapping only ever lowers a +# coordinate, so it always lands inside the axis it addresses and nothing is checked here. +_TENSOR_SCATTER_CIRCULAR_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* update,$index_parameters + size_t prefix_count, + size_t sequence_length, + size_t max_sequence_length, + size_t block, + int prefix_rank, + const size_t* prefix_shape, + const size_t* prefix_strides) +{ + size_t prefix, step; + for (prefix = 0; prefix < prefix_count; ++prefix) { + const ptrdiff_t written_at = $written_at; + const ptrdiff_t capacity = (ptrdiff_t)max_sequence_length; + const size_t start = (size_t)(((written_at % capacity) + capacity) % capacity); + size_t remainder = prefix; + size_t base = 0; + int walked; + for (walked = prefix_rank - 1; walked >= 0; --walked) { + const size_t coordinate = remainder % prefix_shape[walked]; + remainder /= prefix_shape[walked]; + base += (coordinate % max_sequence_length) * prefix_strides[walked]; + } + for (step = 0; step < sequence_length; ++step) { + memcpy( + out + base + ((start + step) % max_sequence_length) * block, + update + (prefix * sequence_length + step) * block, + block * sizeof(*out)); + } + } +}""") + + +def _scatter_elements(context: NodeContext) -> NodeEmission: + """ScatterElements: one element written per update, at the index's own coordinates.""" + data = context.require_input(0) + indices = context.require_input(1) + updates = context.require_input(2) + result = context.require_output(0) + rank = len(data.shape) + verify_shape(context, result, data.shape) + if rank == 0 or len(indices.shape) != rank: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` writes into `{data.name}` " + f"of rank {rank} through `{indices.name}` of rank {len(indices.shape)}; ONNX " + "defines the two as having the same rank, and at least one axis." + ) + if updates.shape != indices.shape: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` takes one update per index, " + f"but `{updates.name}` has shape {list(updates.shape)} and `{indices.name}` " + f"has shape {list(indices.shape)}." + ) + axis = normalize_axis(context, context.int_attribute("axis"), rank) + for other in range(rank): + if other != axis and indices.shape[other] > data.shape[other]: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` writes along axis " + f"{axis}, so `{indices.name}` of shape {list(indices.shape)} addresses " + f"`{data.name}` of shape {list(data.shape)} on every other axis, and " + f"reaches past it on axis {other}." + ) + + statements = list(copy_tensor(data, result).statements) + if updates.elem_count == 0: + return NodeEmission(functions=(), statements=tuple(statements)) + + reduction = _reduction(context) + fold = _fold( + context, + reduction, + data.elem_type, + current="out[offset]", + update="updates[index]", + numpy_extremum=False, + ) + name = kernel_name( + context, reduction, c_type(data.elem_type), c_type(indices.elem_type) + ) + statements.append( + checked_call( + context, + name, + [ + result.expr, + updates.expr, + indices.expr, + f"{updates.elem_count}u", + str(rank), + extents(indices.shape), + extents(row_major_strides(data.shape)), + str(axis), + f"{data.shape[axis]}u", + ], + ) + ) + return NodeEmission( + functions=( + *fold.helpers, + CFunction( + name, + _SCATTER_ELEMENTS_TEMPLATE.substitute( + name=name, + element=c_type(data.elem_type), + index=c_type(indices.elem_type), + fold=fold.expression, + ), + ), + ), + statements=tuple(statements), + ) + + +def _scatter_nd(context: NodeContext) -> NodeEmission: + """ScatterND: a slice written per index tuple, into the axes the tuple names.""" + data = context.require_input(0) + indices = context.require_input(1) + updates = context.require_input(2) + result = context.require_output(0) + rank = len(data.shape) + verify_shape(context, result, data.shape) + if not indices.shape: + raise CompileError( + f"Node `{context.label}`: `ScatterND` takes its index tuples from the last " + f"axis of `{indices.name}`, which is a scalar and has none." + ) + depth = indices.shape[-1] + if depth > rank: + raise CompileError( + f"Node `{context.label}`: `ScatterND` addresses {depth} dimension(s) of " + f"`{data.name}`, which has {rank}." + ) + expected = (*indices.shape[:-1], *data.shape[depth:]) + if updates.shape != expected: + raise CompileError( + f"Node `{context.label}`: `ScatterND` writes one slice per index tuple, so " + f"`{updates.name}` has shape {list(expected)}; it has " + f"{list(updates.shape)}." + ) + + statements = list(copy_tensor(data, result).statements) + if updates.elem_count == 0: + return NodeEmission(functions=(), statements=tuple(statements)) + + slice_size = math.prod(data.shape[depth:]) + reduction = _reduction(context) + fold = _fold( + context, + reduction, + data.elem_type, + current="out[offset + element]", + update="updates[row * slice_size + element]", + numpy_extremum=True, + ) + name = kernel_name( + context, reduction, c_type(data.elem_type), c_type(indices.elem_type) + ) + statements.append( + checked_call( + context, + name, + [ + result.expr, + updates.expr, + indices.expr, + f"{math.prod(indices.shape[:-1])}u", + f"{depth}u", + f"{slice_size}u", + extents(data.shape[:depth]), + extents(row_major_strides(data.shape)[:depth]), + ], + ) + ) + return NodeEmission( + functions=( + *fold.helpers, + CFunction( + name, + _SCATTER_ND_TEMPLATE.substitute( + name=name, + element=c_type(data.elem_type), + index=c_type(indices.elem_type), + fold=fold.expression, + ), + ), + ), + statements=tuple(statements), + ) + + +def _tensor_scatter(context: NodeContext) -> NodeEmission: + """TensorScatter: the update written into each sample's cache at that sample's index.""" + cache = context.require_input(0) + update = context.require_input(1) + written_at = context.optional_input(2) + result = context.require_output(0) + rank = len(cache.shape) + verify_shape(context, result, cache.shape) + axis = _sequence_axis(context, rank) + if len(update.shape) != rank or any( + extent != cached + for position, (extent, cached) in enumerate(zip(update.shape, cache.shape)) + if position != axis + ): + raise CompileError( + f"Node `{context.label}`: `TensorScatter` writes `{update.name}` of shape " + f"{list(update.shape)} into `{cache.name}` of shape {list(cache.shape)}; ONNX " + f"defines the two as differing on axis {axis} alone." + ) + if update.shape[axis] > cache.shape[axis]: + raise CompileError( + f"Node `{context.label}`: `TensorScatter` writes {update.shape[axis]} " + f"position(s) into axis {axis} of `{cache.name}`, which holds " + f"{cache.shape[axis]}." + ) + if written_at is not None and written_at.shape != (cache.shape[0],): + raise CompileError( + f"Node `{context.label}`: `TensorScatter` reads one write index per sample of " + f"the batch, so `{written_at.name}` has shape {[cache.shape[0]]}; it has " + f"{list(written_at.shape)}." + ) + + statements = list(copy_tensor(cache, result).statements) + if update.elem_count == 0: + return NodeEmission(functions=(), statements=tuple(statements)) + + indexed = written_at is not None + circular = _mode(context) == "circular" + prefix_shape = cache.shape[:axis] + block = math.prod(cache.shape[axis + 1 :]) + arguments = [result.expr, update.expr] + if written_at is not None: + arguments += [written_at.expr, f"{math.prod(prefix_shape[1:])}u"] + arguments += [ + f"{math.prod(prefix_shape)}u", + f"{update.shape[axis]}u", + f"{cache.shape[axis]}u", + f"{block}u", + ] + if circular: + arguments += [ + str(axis), + extents(prefix_shape), + extents(row_major_strides(cache.shape)[:axis]), + ] + name = kernel_name( + context, + "circular" if circular else "linear", + "indexed" if indexed else "appended", + c_type(cache.elem_type), + ) + template = ( + _TENSOR_SCATTER_CIRCULAR_TEMPLATE + if circular + else _TENSOR_SCATTER_LINEAR_TEMPLATE + ) + definition = template.substitute( + name=name, + element=c_type(cache.elem_type), + index_parameters=( + f"\n const {c_type(written_at.elem_type)}* write_indices," + "\n size_t batch_stride," + if written_at is not None + else "" + ), + written_at=( + "(ptrdiff_t)write_indices[prefix / batch_stride]" if indexed else "0" + ), + ) + statements.append( + call_kernel(name, arguments) + if circular + else checked_call(context, name, arguments) + ) + return NodeEmission( + functions=(CFunction(name, definition),), statements=tuple(statements) + ) + + +def _fold( + context: NodeContext, + reduction: str, + elem_type: int, + *, + current: str, + update: str, + numpy_extremum: bool, +) -> _Fold: + """The C expression writing `update` over `current`, as `reduction` combines the two. + + `numpy_extremum` selects between the two readings of `max` and `min` the ONNX reference + implementations of these ops carry — see the module docstring. + """ + if reduction == "none": + return _Fold(update) + if reduction == "add": + # numpy adds two booleans as their disjunction, which is what the reference folds + # with; a `+` on the byte a boolean is emitted as would leave a 2 behind. + operator = "|" if elem_type == TensorProto.BOOL else "+" + return _Fold(f"{current} {operator} {update}") + if reduction == "mul": + return _Fold(f"{current} * {update}") + largest = reduction == "max" + if numpy_extremum: + helper = combiner(context, elem_type, largest=largest) + return _Fold(f"{helper.name}({current}, {update})", (helper,)) + comparison = ">" if largest else "<" + return _Fold(f"({update} {comparison} {current}) ? {update} : {current}") + + +def _reduction(context: NodeContext) -> str: + """How the node folds an update into the result, of the values its revision defines.""" + allowed = ["none"] + for version, added in _REDUCTION_VERSIONS: + if context.since_version >= version: + allowed += added + value = context.attribute("reduction", b"none") + reduction = value.decode() if isinstance(value, bytes) else str(value) + if reduction not in allowed: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reduction `{reduction}` is " + f"not one of the reductions ONNX defines at opset version " + f"{context.since_version} ({', '.join(allowed)})." + ) + return reduction + + +def _mode(context: NodeContext) -> str: + value = context.attribute("mode", b"linear") + mode = value.decode() if isinstance(value, bytes) else str(value) + if mode not in _TENSOR_SCATTER_MODES: + raise CompileError( + f"Node `{context.label}`: `TensorScatter` mode `{mode}` is not one of the " + f"modes ONNX defines ({', '.join(_TENSOR_SCATTER_MODES)})." + ) + return mode + + +def _sequence_axis(context: NodeContext, rank: int) -> int: + """The axis TensorScatter writes along, which is never the batch it reads indices by.""" + axis = normalize_axis(context, context.int_attribute("axis"), rank) + if axis == 0: + raise CompileError( + f"Node `{context.label}`: `TensorScatter` writes along axis 0 of " + f"`{context.require_input(0).name}`, which is the batch it takes one write " + "index per sample of; ONNX defines the sequence axis as a later one." + ) + return axis + + +register_kernel("", "Scatter", _SCATTER_VERSIONS, _scatter_elements) +register_kernel("", "ScatterElements", _SCATTER_ELEMENTS_VERSIONS, _scatter_elements) +register_kernel("", "ScatterND", _SCATTER_ND_VERSIONS, _scatter_nd) +register_kernel("", "TensorScatter", _TENSOR_SCATTER_VERSIONS, _tensor_scatter) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/scores.py b/src/python/fnnx/extras/compilers/c/onnx/ops/scores.py new file mode 100644 index 0000000..22367ad --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/scores.py @@ -0,0 +1,360 @@ +"""What an ONNX-ML predictor does with a row of scores once it has computed one. + +The tree ensembles, the support vector machines and the linear models differ entirely in how +they arrive at those scores and not at all in what they do next: the five `post_transform` +values ONNX-ML defines, the second column a single-score binary classifier is paired with, +and the `argmax` that turns a row into a class label. All three are emitted from here, so a +graph running several kinds of predictor shares one kernel per element type rather than one +per op. + +Each is the ONNX reference implementation's own arithmetic rather than the textbook form: +`SOFTMAX_ZERO`'s threshold, `PROBIT`'s rational approximation of `erfinv` and the constants +either side of it, and the `LOGISTIC` that folds a negative argument onto a positive one are +what `onnx.reference.ops.aionnxml._common_classifier` computes, and that is the oracle every +predictor here is compared against. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from string import Template +from typing import TypeVar + +import numpy as np +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type, element_type_name +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + constant_data, +) +from fnnx.extras.compilers.c.onnx.ops.axes import call_kernel +from fnnx.extras.compilers.c.onnx.ops.broadcast import math_suffix + +# How an op names the transform or the aggregation it applies: the ONNX-ML ops of opset 1 +# spell them out as strings, the opset-5 `TreeEnsemble` numbers them. +Declared = TypeVar("Declared", str, int) + +# The score transforms, in the numbering opset 5 gave them; every op that names them as +# strings names these same five. +NONE, SOFTMAX, LOGISTIC, SOFTMAX_ZERO, PROBIT = range(5) +NAMED_TRANSFORMS = { + "NONE": NONE, + "SOFTMAX": SOFTMAX, + "LOGISTIC": LOGISTIC, + "SOFTMAX_ZERO": SOFTMAX_ZERO, + "PROBIT": PROBIT, +} +_TRANSFORM_NAMES = { + SOFTMAX: "softmax", + LOGISTIC: "logistic", + SOFTMAX_ZERO: "softmax_zero", + PROBIT: "probit", +} + +# What a classifier makes of the winning column of a row. Naming the class that column stands +# for is what every predictor here does bar one: `SVMClassifier` carries two readings of its +# own, which `set_score_svm` applies to an ensemble whose `rho` is a single value — the second +# class outright once the winning vote reaches one half, and the sign of the winning value as +# a label of its own where the ensemble does not have exactly two classes. +ARGMAX_LABEL, POSITIVE_CLASS, SIGN_LABEL = range(3) +# The transforms that read a whole row at once, against the ones that map value to value. +_ROW_TRANSFORMS = (SOFTMAX, SOFTMAX_ZERO) + +# `SOFTMAX_ZERO` leaves a value this close to zero out of the exponential and scales it +# instead; the constant is the reference implementation's own. +_ZERO_THRESHOLD = 1e-7 + +# `PROBIT` is `sqrt(2) * erfinv(2p - 1)` with `erfinv` the rational approximation the +# reference implementation spells out; these are the two constants that approximation is +# built from, either side of `0.5 * log((1 - x) * (1 + x))`. +_PROBIT_FIRST = 2.0 / (math.pi * 0.147) +_PROBIT_SECOND = 1.0 / 0.147 +_PROBIT_SCALE = 1.41421356 + +_SOFTMAX_TEMPLATE = Template("""\ +static void $name($result* scores, size_t rows, size_t columns) +{ + size_t row, index; + if (columns == 0) { + return; + } + for (row = 0; row < rows; ++row) { + $result* out = scores + row * columns; + $result largest = out[0]; + $result total = ($result)0; + for (index = 1; index < columns; ++index) { + if (out[index] > largest || isnan(out[index])) { + largest = out[index]; + } + } + for (index = 0; index < columns; ++index) { + out[index] = exp$suffix(out[index] - largest); + total += out[index]; + } + for (index = 0; index < columns; ++index) { + out[index] /= total; + } + } +}""") + +_SOFTMAX_ZERO_TEMPLATE = Template("""\ +static void $name($result* scores, size_t rows, size_t columns) +{ + size_t row, index; + if (columns == 0) { + return; + } + for (row = 0; row < rows; ++row) { + $result* out = scores + row * columns; + $result largest = out[0]; + $result total = ($result)0; + $result scale; + for (index = 1; index < columns; ++index) { + if (out[index] > largest || isnan(out[index])) { + largest = out[index]; + } + } + scale = exp$suffix(-largest); + for (index = 0; index < columns; ++index) { + const $result value = out[index]; + out[index] = (value > $threshold || value < -$threshold) + ? exp$suffix(value - largest) + : value * scale; + total += out[index]; + } + for (index = 0; index < columns; ++index) { + out[index] = (total == ($result)0) ? ($result)0.5 : out[index] / total; + } + } +}""") + +_LOGISTIC_TEMPLATE = Template("""\ +static void $name($result* scores, size_t count) +{ + size_t index; + for (index = 0; index < count; ++index) { + const $result value = scores[index]; + const $result mapped = + ($result)1 / (($result)1 + exp$suffix(-fabs$suffix(value))); + scores[index] = (value < 0) ? (($result)1 - mapped) : mapped; + } +}""") + +_PROBIT_TEMPLATE = Template("""\ +static void $name($result* scores, size_t count) +{ + size_t index; + for (index = 0; index < count; ++index) { + const $result value = scores[index] * ($result)2 - ($result)1; + const $result inner = (($result)1 - value) * (($result)1 + value); + $result mapped = ($result)0; + if (inner != ($result)0) { + const $result logarithm = log$suffix(inner); + const $result first = $constant + ($result)0.5 * logarithm; + const $result second = $reciprocal * logarithm; + const $result root = -first + sqrt$suffix(first * first - second); + mapped = ((value < 0) ? ($result)-1 : ($result)1) * sqrt$suffix(root); + } + scores[index] = $scale * mapped; + } +}""") + +_BINARY_TEMPLATE = Template("""\ +static void $name($result* scores, size_t rows, size_t columns, int complement) +{ + size_t row; + for (row = 0; row < rows; ++row) { + $result* out = scores + row * columns; + out[1] = out[0]; + out[0] = complement ? (($result)1 - out[1]) : -out[1]; + } +}""") + +_ARGMAX_TEMPLATE = Template("""\ +static void $name( + int64_t* labels, + const $result* values, + const int64_t* classes, + size_t rows, + size_t columns, + int rule) +{ + size_t row, index; + for (row = 0; row < rows; ++row) { + const $result* out = values + row * columns; + size_t best = 0; + for (index = 1; index < columns; ++index) { + /* The column carrying the first value that is not a number wins outright, which + is what the `argmax` this stands for returns. */ + if (isnan(out[best])) { + break; + } + if (out[index] > out[best] || isnan(out[index])) { + best = index; + } + } + if (rule == $sign) { + labels[row] = (out[best] > ($result)0) ? 1 : 0; + } else if (rule == $positive && out[best] >= ($result)0.5) { + labels[row] = classes[1]; + } else { + labels[row] = classes[best]; + } + } +}""") + +_TRANSFORM_TEMPLATES = { + SOFTMAX: _SOFTMAX_TEMPLATE, + SOFTMAX_ZERO: _SOFTMAX_ZERO_TEMPLATE, + LOGISTIC: _LOGISTIC_TEMPLATE, + PROBIT: _PROBIT_TEMPLATE, +} + + +def named_transform(context: NodeContext) -> int: + """The `post_transform` an op names as a string, as one of the five constants above.""" + return choice( + context, + "post_transform", + context.string_attribute("post_transform"), + NAMED_TRANSFORMS, + ) + + +def post_transform( + context: NodeContext, scores: TensorRef, transform: int, rows: int, columns: int +) -> NodeEmission | None: + """The transform the node names, applied over the scores in place.""" + template = _TRANSFORM_TEMPLATES.get(transform) + if template is None: + return None + element = c_type(scores.elem_type) + name = f"{context.prefix}_ml_{_TRANSFORM_NAMES[transform]}_{element}" + definition = template.substitute( + name=name, + result=element, + suffix=math_suffix(scores.elem_type), + threshold=scalar_literal(_ZERO_THRESHOLD, scores.elem_type), + constant=scalar_literal(_PROBIT_FIRST, scores.elem_type), + reciprocal=scalar_literal(_PROBIT_SECOND, scores.elem_type), + scale=scalar_literal(_PROBIT_SCALE, scores.elem_type), + ) + arguments = ( + [scores.expr, f"{rows}u", f"{columns}u"] + if transform in _ROW_TRANSFORMS + else [scores.expr, f"{rows * columns}u"] + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call_kernel(name, arguments),), + ) + + +def binary_scores( + context: NodeContext, + scores: TensorRef, + rows: int, + columns: int, + *, + complement: bool, +) -> NodeEmission: + """The second column a classifier scoring one value per row pairs that value with. + + The score is expected in the first column of each row of `scores`, which the caller has + already sized for both. Which second column it gets is the classifier's own rule: the + complement of the score, or its negation. + """ + element = c_type(scores.elem_type) + name = f"{context.prefix}_ml_binary_{element}" + definition = _BINARY_TEMPLATE.substitute(name=name, result=element) + call = call_kernel( + name, [scores.expr, f"{rows}u", f"{columns}u", str(int(complement))] + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def argmax_labels( + context: NodeContext, + labels: TensorRef, + values: str, + elem_type: int, + classes: Sequence[int], + rows: int, + columns: int, + rule: int = ARGMAX_LABEL, +) -> NodeEmission: + """Each row labelled from its winning column of `values`, by the rule the op applies. + + `values` is what the classes are ranked by, which is the scores themselves for every + predictor but the support-vector one, where the classes are ranked by their votes. + """ + element = c_type(elem_type) + name = f"{context.prefix}_ml_argmax_{element}" + definition = _ARGMAX_TEMPLATE.substitute( + name=name, result=element, sign=SIGN_LABEL, positive=POSITIVE_CLASS + ) + data, symbol = constant_data(context, "classes", np.array(classes, np.int64)) + call = call_kernel( + name, [labels.expr, values, symbol, f"{rows}u", f"{columns}u", str(rule)] + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call,), + constants=(data,), + ) + + +def extend(emission: NodeEmission, addition: NodeEmission | None) -> NodeEmission: + if addition is None: + return emission + return NodeEmission( + functions=emission.functions + addition.functions, + statements=emission.statements + addition.statements, + scratch=emission.scratch + addition.scratch, + constants=emission.constants + addition.constants, + ) + + +def choice( + context: NodeContext, + name: str, + value: Declared, + choices: Mapping[Declared, int], +) -> int: + if value not in choices: + raise CompileError( + f"Node `{context.label}`: `{name}` is `{value}`, which is none of the values " + f"ONNX defines for it ({', '.join(str(key) for key in choices)})." + ) + return choices[value] + + +def float_output(context: NodeContext, index: int) -> TensorRef: + """The scores, checked to be the float32 tensor the ONNX-ML schemas declare.""" + result = context.require_output(index) + if result.elem_type != TensorProto.FLOAT: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` scores in `float`, but its " + f"output `{result.name}` is declared " + f"`{element_type_name(result.elem_type)}`." + ) + return result + + +def label_output(context: NodeContext, index: int) -> TensorRef: + """The predicted classes, checked to be the `int64` tensor a numeric label table needs.""" + result = context.require_output(index) + if result.elem_type != TensorProto.INT64: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` labels each row with an " + f"`int64` class value, but its output `{result.name}` is declared " + f"`{element_type_name(result.elem_type)}`." + ) + return result diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/signal.py b/src/python/fnnx/extras/compilers/c/onnx/ops/signal.py new file mode 100644 index 0000000..a538e89 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/signal.py @@ -0,0 +1,536 @@ +"""DFT and STFT: the transform written out as the sum ONNX defines it to be. + +A discrete Fourier transform is one sum per output bin over the samples of one axis, and a +short-time transform is that same sum over a window slid along the signal. Both are emitted +as exactly that — no factorization, no twiddle tables — so the code is the definition and its +size does not depend on the transform's length. The transform is evaluated in `double` +whatever the tensor holds, which is what numpy's own FFT does before rounding back to the +input's type, so a `float` model is not compared against an oracle computed to a different +precision. + +What varies between calls is addressing, not code: the transformed axis is read as an outer +block count, an inner stride and an extent, so one kernel per element type serves every axis, +rank and length. The axis itself may be named at run time — ONNX passes it as an operand from +opset 20 — but only where it cannot change the shape of the result; `onesided` and +`dft_length` both resize the axis they land on, and which axis that is has to be known before +a buffer can be sized. The window functions and MelWeightMatrix have no kernels at all: their +operands are their result's own shape, so a model that fixes them is folded away through the +reference evaluator before dispatch, and one that does not is refused by the frontend. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass +from string import Template + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type +from fnnx.extras.compilers.c.onnx.emit import INVALID_ARGUMENT_STATUS +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import call_kernel, kernel_name, verify_shape + +# DFT-20 moved the axis from an attribute to an operand; the two revisions are otherwise the +# same transform. STFT has had a single revision since opset 17. +_AXIS_AS_OPERAND = 20 +_DFT_VERSIONS = (17, _AXIS_AS_OPERAND) +_STFT_VERSIONS = (17,) + +# Where that operand sits, and what the node transforms when it is left out: the last signal +# axis, the one before the real/imaginary pair. Revision 17's attribute defaults to the first +# signal axis instead, which its schema states and `int_attribute` reads. +_AXIS_INPUT = 2 +_DEFAULT_AXIS = -2 + +# The last axis of a signal tensor holds a real value alone or a real and an imaginary part. +_SIGNAL_COMPONENTS = (1, 2) + +# 2*pi to more digits than a double carries, so the constant is the nearest representable +# one however the C compiler parses it. +_TURN = "6.283185307179586476925286766559" + +_DFT_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + size_t outer, + size_t inner, + size_t samples, + size_t bins, + size_t length, + size_t components, + size_t written, + int inverse, + int mirrored) +{ + const double turn = $turn; + const size_t terms = mirrored ? length : (samples < length ? samples : length); + size_t lead, trail, bin, term; + for (lead = 0; lead < outer; ++lead) { + for (trail = 0; trail < inner; ++trail) { + const size_t from = (lead * samples * inner + trail) * components; + const size_t to = (lead * bins * inner + trail) * written; + for (bin = 0; bin < bins; ++bin) { + double real = 0.0; + double imaginary = 0.0; + for (term = 0; term < terms; ++term) { + const int conjugated = mirrored && term > length / 2; + const size_t source = conjugated ? length - term : term; + if (source < samples) { + const size_t at = from + source * inner * components; + const double angle = + turn * (double)(term * bin % length) / (double)length; + const double cosine = cos(angle); + const double sine = inverse ? sin(angle) : -sin(angle); + const double re = (double)in[at]; + double im = components == 2 ? (double)in[at + 1] : 0.0; + if (conjugated) { + im = -im; + } + real += re * cosine - im * sine; + imaginary += re * sine + im * cosine; + } + } + if (inverse) { + real /= (double)length; + imaginary /= (double)length; + } + out[to + bin * inner * written] = ($element)real; + if (written == 2) { + out[to + bin * inner * written + 1] = ($element)imaginary; + } + } + } + } +}""") + +_STFT_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + const $element* window, + size_t batch, + size_t signal, + size_t components, + size_t frames, + size_t step, + size_t length, + size_t bins) +{ + const double turn = $turn; + size_t index, frame, bin, term; + for (index = 0; index < batch; ++index) { + for (frame = 0; frame < frames; ++frame) { + for (bin = 0; bin < bins; ++bin) { + double real = 0.0; + double imaginary = 0.0; + for (term = 0; term < length; ++term) { + const size_t sample = frame * step + term; + if (sample < signal) { + const size_t at = (index * signal + sample) * components; + const double weight = + window != NULL ? (double)window[term] : 1.0; + const double angle = + turn * (double)(term * bin % length) / (double)length; + const double cosine = cos(angle); + const double sine = -sin(angle); + const double re = weight * (double)in[at]; + const double im = + components == 2 ? weight * (double)in[at + 1] : 0.0; + real += re * cosine - im * sine; + imaginary += re * sine + im * cosine; + } + } + out[((index * frames + frame) * bins + bin) * 2] = ($element)real; + out[((index * frames + frame) * bins + bin) * 2 + 1] = + ($element)imaginary; + } + } + } +}""") + +_NORMALIZE_TEMPLATE = Template("""\ +static int64_t $name(int64_t axis, int64_t rank) +{ + return (axis < 0) ? (axis + rank) : axis; +}""") + + +@dataclass(frozen=True) +class _Transform: + """One DFT call site: the shape it writes, and the addressing it walks to write it. + + `leading` and `trailing` are the operand's extents before and after the transformed axis, + the trailing ones less the real/imaginary axis: a block of transforms per coordinate of + the first, one stride apart for each coordinate of the second. `mirrored` marks the + inverse one-sided transform, whose operand holds only the non-redundant half of a + spectrum the conjugate symmetry fills back in — which is also the one transform writing + a real result rather than a complex one. + """ + + leading: tuple[int, ...] + trailing: tuple[int, ...] + samples: int + bins: int + length: int + components: int + mirrored: bool + inverse: bool + + @property + def written(self) -> int: + return 1 if self.mirrored else 2 + + @property + def shape(self) -> tuple[int, ...]: + return (*self.leading, self.bins, *self.trailing, self.written) + + def arguments(self, result: TensorRef, source: TensorRef) -> list[str]: + return [ + result.expr, + source.expr, + f"{math.prod(self.leading)}u", + f"{math.prod(self.trailing)}u", + f"{self.samples}u", + f"{self.bins}u", + f"{self.length}u", + f"{self.components}u", + f"{self.written}u", + str(int(self.inverse)), + str(int(self.mirrored)), + ] + + +def _dft(context: NodeContext) -> NodeEmission: + source = context.require_input(0) + result = context.require_output(0) + rank = _signal_rank(context, source) + inverse = context.int_attribute("inverse") != 0 + onesided = context.int_attribute("onesided") != 0 + length = _dft_length(context) + + fixed = _fixed_axis(context, rank) + axes = ( + (fixed,) + if fixed is not None + else _runtime_axes(context, rank, onesided, length) + ) + plans = { + axis: _plan( + source.shape, axis, inverse=inverse, onesided=onesided, length=length + ) + for axis in axes + } + for plan in plans.values(): + verify_shape(context, result, plan.shape) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + for plan in plans.values(): + _verify_length(context, plan.length) + + element = c_type(result.elem_type) + name = kernel_name(context, element) + kernel = CFunction( + name, _DFT_TEMPLATE.substitute(name=name, element=element, turn=_TURN) + ) + calls = { + axis: call_kernel(name, plan.arguments(result, source)) + for axis, plan in plans.items() + } + if fixed is not None: + return NodeEmission(functions=(kernel,), statements=(calls[fixed],)) + normalize = _normalize_helper(context.prefix) + operand = context.require_input(_AXIS_INPUT) + return NodeEmission( + functions=(normalize, kernel), + statements=(_dispatch(context, operand, rank, calls, normalize),), + ) + + +def _stft(context: NodeContext) -> NodeEmission: + signal = context.require_input(0) + result = context.require_output(0) + if len(signal.shape) != 3: + raise CompileError( + f"Node `{context.label}`: `STFT` was given `{signal.name}` of shape " + f"{list(signal.shape)}; ONNX defines its signal as " + "[batch_size][signal_length][1 or 2]." + ) + batch, samples, components = signal.shape + _verify_components(context, signal, components) + window = context.optional_input(2) + step = _frame_step(context) + length = _frame_length(context, window) + _verify_length(context, length) + onesided = context.int_attribute("onesided") != 0 + bins = length // 2 + 1 if onesided else length + frames = 1 + (samples - length) // step + if frames < 0: + raise CompileError( + f"Node `{context.label}`: `STFT` frames of {length} sample(s) do not fit a " + f"signal of {samples}; not one whole frame can be taken." + ) + verify_shape(context, result, (batch, frames, bins, 2)) + if result.elem_count == 0: + return NodeEmission(functions=(), statements=()) + + element = c_type(result.elem_type) + name = kernel_name(context, element) + definition = _STFT_TEMPLATE.substitute(name=name, element=element, turn=_TURN) + arguments = [ + result.expr, + signal.expr, + "NULL" if window is None else window.expr, + f"{batch}u", + f"{samples}u", + f"{components}u", + f"{frames}u", + f"{step}u", + f"{length}u", + f"{bins}u", + ] + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call_kernel(name, arguments),), + ) + + +def _plan( + shape: Sequence[int], + axis: int, + *, + inverse: bool, + onesided: bool, + length: int | None, +) -> _Transform: + """How the transform along `axis` addresses the tensor, and the result it writes. + + An absent `dft_length` is the extent of the axis itself, except for the inverse one-sided + transform, whose operand holds only the non-redundant half of an even-length spectrum. + """ + samples = shape[axis] + mirrored = inverse and onesided + if length is None: + length = 2 * (samples - 1) if mirrored else samples + return _Transform( + leading=tuple(shape[:axis]), + trailing=tuple(shape[axis + 1 : -1]), + samples=samples, + bins=length // 2 + 1 if onesided and not inverse else length, + length=length, + components=shape[-1], + mirrored=mirrored, + inverse=inverse, + ) + + +def _signal_rank(context: NodeContext, source: TensorRef) -> int: + """The operand's rank, once it is one a transform is defined over.""" + rank = len(source.shape) + if rank < 2: + raise CompileError( + f"Node `{context.label}`: `DFT` was given `{source.name}` of shape " + f"{list(source.shape)}; ONNX defines its input as at least one signal axis " + "followed by the axis holding the real and imaginary parts." + ) + _verify_components(context, source, source.shape[-1]) + return rank + + +def _verify_components( + context: NodeContext, source: TensorRef, components: int +) -> None: + if components not in _SIGNAL_COMPONENTS: + raise CompileError( + f"Node `{context.label}`: the last axis of `{source.name}` measures " + f"{components}; ONNX defines it as 1 for a real signal and 2 for a complex one." + ) + + +def _verify_length(context: NodeContext, length: int) -> None: + if length < 1: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` was given a transform " + f"length of {length}; ONNX defines nothing for a transform over no samples " + "at all." + ) + + +def _fixed_axis(context: NodeContext, rank: int) -> int | None: + """The axis this node transforms, or None where it names one only at run time.""" + if context.since_version < _AXIS_AS_OPERAND: + return _checked_axis(context, context.int_attribute("axis"), rank) + if context.optional_input(_AXIS_INPUT) is None: + return _checked_axis(context, _DEFAULT_AXIS, rank) + fixed = context.constant_input(_AXIS_INPUT) + if fixed is None: + return None + if fixed.size != 1: + raise CompileError( + f"Node `{context.label}`: the axis of `DFT` comes from " + f"`{context.require_input(_AXIS_INPUT).name}`, which holds {fixed.size} " + "values; ONNX defines it as a single one." + ) + return _checked_axis(context, int(fixed.reshape(-1)[0]), rank) + + +def _checked_axis(context: NodeContext, axis: int, rank: int) -> int: + """An ONNX DFT axis, which may count from the end, as an index into the operand. + + The last axis holds the real and imaginary parts rather than a signal, so it is the one + axis of the operand no transform may be taken along. + """ + if not -rank <= axis < rank - 1 or axis == -1: + raise CompileError( + f"Node `{context.label}`: `DFT` transforms axis {axis} of a rank-{rank} " + "operand; ONNX defines the axis over [-rank, -2] and [0, rank - 2], the last " + "axis being the real and imaginary parts rather than a signal." + ) + return axis + rank if axis < 0 else axis + + +def _runtime_axes( + context: NodeContext, rank: int, onesided: bool, length: int | None +) -> tuple[int, ...]: + """Every axis a run-time operand could name, where naming one cannot resize the result. + + A one-sided transform returns half its length and a stated `dft_length` replaces it, so + either makes the extent of the transformed axis differ from the operand's — and which + axis that is then has to be known to size a buffer at all. With neither, every axis + leaves the operand's own extents, so the result is one shape whichever the operand names + and the choice is a run-time switch over compile-time call sites. + """ + if onesided or length is not None: + raise CompileError( + f"Node `{context.label}`: `DFT` takes its axis from " + f"`{context.require_input(_AXIS_INPUT).name}`, which no initializer or constant " + f"folding fixes, while {'`onesided`' if onesided else '`dft_length`'} resizes " + "the axis it transforms; the shape of the result then depends on input data, " + "which the C compiler requires to be known at compile time." + ) + return tuple(range(rank - 1)) + + +def _dft_length(context: NodeContext) -> int | None: + """The transform's length where the node states one, as a compile-time value.""" + operand = context.optional_input(1) + if operand is None: + return None + fixed = context.constant_input(1) + if fixed is None or fixed.size != 1: + raise CompileError( + f"Node `{context.label}`: `DFT` takes its length from `{operand.name}`, which " + f"holds {'no single value' if fixed is not None else 'no value'} known at " + "compile time; the shape of the result then depends on input data, which the C " + "compiler cannot compile." + ) + return int(fixed.reshape(-1)[0]) + + +def _frame_step(context: NodeContext) -> int: + operand = context.require_input(1) + fixed = context.constant_input(1) + if fixed is None or fixed.size != 1: + raise CompileError( + f"Node `{context.label}`: `STFT` takes its frame step from `{operand.name}`, " + f"which holds {'no single value' if fixed is not None else 'no value'} known " + "at compile time; how many frames the signal yields then depends on input " + "data, which the C compiler cannot compile." + ) + step = int(fixed.reshape(-1)[0]) + if step < 1: + raise CompileError( + f"Node `{context.label}`: `STFT` steps {step} sample(s) between frames; ONNX " + "defines the step as the samples to advance by, which is positive." + ) + return step + + +def _frame_length(context: NodeContext, window: TensorRef | None) -> int: + """How many samples one frame holds, from whichever operand the node states it with. + + ONNX takes it from `frame_length`, and from the window's own extent where the node + passes a window instead; a node passing both states the same number twice. + """ + stated = context.optional_input(3) + fixed = context.constant_input(3) if stated is not None else None + if stated is not None and (fixed is None or fixed.size != 1): + raise CompileError( + f"Node `{context.label}`: `STFT` takes its frame length from `{stated.name}`, " + f"which holds {'no single value' if fixed is not None else 'no value'} known " + "at compile time; the shape of the result then depends on input data, which " + "the C compiler cannot compile." + ) + spanned = _window_span(context, window) + if fixed is None: + if spanned is None: + raise CompileError( + f"Node `{context.label}`: `STFT` states neither a window nor a frame " + "length; ONNX defines the frame from one of the two." + ) + return spanned + length = int(fixed.reshape(-1)[0]) + if spanned is not None and spanned != length: + raise CompileError( + f"Node `{context.label}`: `STFT` was given a window of {spanned} sample(s) and " + f"a frame length of {length}; ONNX defines a node stating both as stating the " + "same number twice." + ) + return length + + +def _window_span(context: NodeContext, window: TensorRef | None) -> int | None: + if window is None: + return None + if len(window.shape) != 1: + raise CompileError( + f"Node `{context.label}`: `STFT` was given the window `{window.name}` of shape " + f"{list(window.shape)}; ONNX defines it as a single sequence of weights." + ) + return window.shape[0] + + +def _normalize_helper(prefix: str) -> CFunction: + """An axis counted from the end, resolved against a rank, both known only as values.""" + name = f"{prefix}_normalized_axis" + return CFunction(name, _NORMALIZE_TEMPLATE.substitute(name=name)) + + +def _dispatch( + context: NodeContext, + operand: TensorRef, + rank: int, + calls: dict[int, str], + normalize: CFunction, +) -> str: + """The transform for whichever axis the operand names at run time, or an argument error. + + The axis is resolved through a function rather than into a local, so that the statement + introduces no identifier of its own — one would shadow the entrypoint parameter a tensor + of the same name is emitted as. + """ + cases = [] + for axis, call in sorted(calls.items()): + cases.append(f"case {axis}:") + cases.extend(f" {line}" if line else "" for line in call.splitlines()) + cases.append(" break;") + return "\n".join( + [ + f"switch ({normalize.name}((int64_t){operand.expr}[0], {rank})) {{", + *cases, + "default:", + f" return {context.prefix.upper()}_{INVALID_ARGUMENT_STATUS};", + "}", + ] + ) + + +register_kernel("", "DFT", _DFT_VERSIONS, _dft) +register_kernel("", "STFT", _STFT_VERSIONS, _stft) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/softmax.py b/src/python/fnnx/extras/compilers/c/onnx/ops/softmax.py new file mode 100644 index 0000000..0a07868 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/softmax.py @@ -0,0 +1,153 @@ +"""Softmax, LogSoftmax and Hardmax: the normalizations that run along one axis. + +Each writes a group in place of itself — the result carries the operand's shape — so one +grouping addresses both buffers. ONNX revised all three at opset 13: up to then they +flattened every axis from `axis` on into a single one, and nothing can vouch for those +revisions (the reference evaluator applies the current semantics to them, and the backend +corpus has no test at an older opset), so only the current revision is served and an older +import gets the standard unsupported-version error. +""" + +from __future__ import annotations + +from functools import partial +from string import Template + +from fnnx.extras.compilers.c.onnx.dtypes import c_type, numpy_dtype_name +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + GROUP_PARAMETERS, + call_kernel, + group_axes, + kernel_name, + normalize_axis, + offset_helper, + verify_same_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import combiner, expand, math_suffix +from fnnx.extras.compilers.c.onnx.ops.reduce import extremum_test + +_VERSIONS = (13,) + +# The group's largest element is subtracted from every exponent, so nothing overflows and +# the result is unchanged; the reference evaluator normalizes the same way, and LogSoftmax +# is the logarithm of what Softmax computes rather than a formula of its own — which is what +# ONNX defines it as, down to where the underflow to zero puts an infinity. +_SOFTMAX_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, +$parameters) +{ + size_t group, index; + for (group = 0; group < group_count; ++group) { + const size_t base = $offset(group, kept_rank, kept_shape, kept_strides); + $element largest = -INFINITY; + $element total = $zero; + for (index = 0; index < group_size; ++index) { + largest = ($element)$maximum(largest, in[base + + $offset(index, reduced_rank, reduced_shape, reduced_strides)]); + } + for (index = 0; index < group_size; ++index) { + total += exp$f(in[base + + $offset(index, reduced_rank, reduced_shape, reduced_strides)] + - largest); + } + for (index = 0; index < group_size; ++index) { + const size_t position = base + + $offset(index, reduced_rank, reduced_shape, reduced_strides); + out[position] = $result; + } + } +}""") + +_HARDMAX_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, +$parameters) +{ + size_t group, index; + for (group = 0; group < group_count; ++group) { + const size_t base = $offset(group, kept_rank, kept_shape, kept_strides); + $element best = $zero; + int64_t chosen = 0; + for (index = 0; index < group_size; ++index) { + const $element x = in[base + + $offset(index, reduced_rank, reduced_shape, reduced_strides)]; + if (index == 0 || ($better)) { + best = x; + chosen = (int64_t)index; + } + } + for (index = 0; index < group_size; ++index) { + out[base + $offset(index, reduced_rank, reduced_shape, reduced_strides)] = + (index == (size_t)chosen) ? $one : $zero; + } + } +}""") + +_NORMALIZED = "exp$f(in[position] - largest) / total" + + +def _softmax(context: NodeContext, *, logarithmic: bool) -> NodeEmission: + elem_type = context.require_output(0).elem_type + largest = combiner(context, elem_type, largest=True) + offset = offset_helper(context.prefix) + name = kernel_name(context, numpy_dtype_name(elem_type)) + result = f"log$f({_NORMALIZED})" if logarithmic else _NORMALIZED + definition = _SOFTMAX_TEMPLATE.substitute( + name=name, + element=c_type(elem_type), + parameters=GROUP_PARAMETERS, + offset=offset.name, + maximum=largest.name, + zero=scalar_literal(0, elem_type), + f=math_suffix(elem_type), + result=expand(result, elem_type), + ) + return _emit(context, CFunction(name, definition), (offset, largest)) + + +def _hardmax(context: NodeContext) -> NodeEmission: + """One where the group's largest element is, zero everywhere else, ties going first.""" + elem_type = context.require_output(0).elem_type + offset = offset_helper(context.prefix) + name = kernel_name(context, numpy_dtype_name(elem_type)) + definition = _HARDMAX_TEMPLATE.substitute( + name=name, + element=c_type(elem_type), + parameters=GROUP_PARAMETERS, + offset=offset.name, + zero=scalar_literal(0, elem_type), + one=scalar_literal(1, elem_type), + better=extremum_test(elem_type, largest=True, last=False), + ) + return _emit(context, CFunction(name, definition), (offset,)) + + +def _emit( + context: NodeContext, kernel: CFunction, helpers: tuple[CFunction, ...] +) -> NodeEmission: + source = context.require_input(0) + result = context.require_output(0) + verify_same_shape(context, source, result) + axis = normalize_axis(context, context.int_attribute("axis"), len(source.shape)) + grouping = group_axes(source.shape, (axis,)) + return NodeEmission( + functions=(*helpers, kernel), + statements=( + call_kernel(kernel.name, [result.expr, source.expr, *grouping.arguments]), + ), + ) + + +register_kernel("", "Softmax", _VERSIONS, partial(_softmax, logarithmic=False)) +register_kernel("", "LogSoftmax", _VERSIONS, partial(_softmax, logarithmic=True)) +register_kernel("", "Hardmax", _VERSIONS, _hardmax) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/svm.py b/src/python/fnnx/extras/compilers/c/onnx/ops/svm.py new file mode 100644 index 0000000..99f5c2e --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/svm.py @@ -0,0 +1,1093 @@ +"""Support vector machines and the linear models beside them. + +`LinearRegressor`, `LinearClassifier`, `SVMRegressor` and `SVMClassifier` are ONNX-ML's four +non-forest predictors, and three of the four are one computation: a row of `X` is scored +against a coefficient per feature per output, offset, and — for a classifier — turned into a +label. Only `SVMClassifier`'s support-vector mode differs, and that is libsvm's one-against-one +scheme: every pair of classes votes, and the winner of the vote names the row's class. All +four carry their parameters as *attributes*, so every table below is emitted as `static const` +data a shared kernel reads through a pointer, and each scores in float32 whatever the element +type of `X` — which is what their schemas declare `Y` and `Z` to be. + +Each op is at revision 1 of its schema and has been since ONNX-ML opset 1. + +Where this follows ONNX's reference implementation rather than the prose, because the +reference is the oracle these ops are compared against: + +* A `LinearRegressor` or `LinearClassifier` with no `intercepts` is refused. The reference + reads the missing attribute as a NaN and adds it to every score; onnxruntime treats it as + zero. Nothing can vouch for either, so the model is refused rather than silently taking a + side. +* A `LinearClassifier` scoring one column against two class labels pairs the score with its + own negation — `[-s, s]`, the *whole* score including the intercept, which is where the + reference and onnxruntime part company; converters emit one coefficient row per class, so + the case does not arise for them. +* `SVMClassifier`'s label is not always the class its winning column names: with a single + `rho`, two class labels, no probabilities and no negative coefficient, a winning vote of at + least 0.5 names the *second* class outright, and with a single `rho` and any other number + of labels the label is 1 or 0 by the sign of the winning value. Both are `set_score_svm`'s. +* A row of one score is widened to two only when a second class is called for, and is left + *untransformed* when it is not — `write_scores` returns such a row before it reaches the + transform, `PROBIT` alone excepted. + +What the compiler refuses: `prob_a`/`prob_b` on an ensemble of more than two classes. The +pairwise probabilities are coupled by an iterative solver whose matrix the reference builds +with a broadcast where libsvm — and onnxruntime with it — writes one entry per pair, so the +two disagree for three classes and up; they agree exactly for two, which is the case +scikit-learn's binary `SVC(probability=True)` emits. +""" + +from __future__ import annotations + +from string import Template +from typing import NamedTuple + +import numpy as np +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + ScratchBuffer, + TensorRef, + constant_data, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.loader import ML_DOMAIN +from fnnx.extras.compilers.c.onnx.ops.axes import call_kernel, verify_shape +from fnnx.extras.compilers.c.onnx.ops.scores import ( + ARGMAX_LABEL, + NONE, + POSITIVE_CLASS, + PROBIT, + SIGN_LABEL, + argmax_labels, + binary_scores, + choice, + extend, + float_output, + label_output, + named_transform, + post_transform, +) + +_VERSIONS = (1,) + +# The kernel functions a support vector machine compares a row against, numbered for the one +# switch the emitted helper is written as. ONNX names them as strings, which the reference +# implementation lowercases before it matches — so the compiler does too. +_LINEAR, _POLY, _RBF, _SIGMOID = range(4) +_KERNEL_TYPES = {"linear": _LINEAR, "poly": _POLY, "rbf": _RBF, "sigmoid": _SIGMOID} + +# `gamma`, `coef0` and `degree`, which ONNX packs into one attribute in that order. +_KERNEL_PARAMETERS = 3 + +# What a pairwise probability is clamped to either side, and the iteration that couples the +# pairs: the tolerance, and the cap on the number of passes. All four are the reference +# implementation's own, with the tolerance divided by the two classes this compiler serves. +_SMALLEST_PROBABILITY = 1e-7 +_LARGEST_PROBABILITY = 1 - 1e-7 +_COUPLING_TOLERANCE = 0.005 / 2 +_COUPLING_ITERATIONS = 100 + +_KERNEL_TEMPLATE = Template("""\ +static float $name( + const $element* row, + const float* support, + size_t width, + int kind, + float gamma, + float coef0, + float degree) +{ + size_t index; + float total = 0.0f; + if (kind == $rbf) { + for (index = 0; index < width; ++index) { + const float difference = (float)row[index] - support[index]; + total += difference * difference; + } + return expf(-gamma * total); + } + for (index = 0; index < width; ++index) { + total += (float)row[index] * support[index]; + } + if (kind == $poly) { + return powf(total * gamma + coef0, degree); + } + if (kind == $sigmoid) { + return tanhf(total * gamma + coef0); + } + return total; +}""") + +_SCORES_TEMPLATE = Template("""\ +static void $name( + float* out, + const $element* in, + const float* coefficients, + const float* bias, + size_t rows, + size_t width, + size_t columns, + size_t stride, + size_t bias_stride) +{ + size_t row, column, feature; + for (row = 0; row < rows; ++row) { + for (column = 0; column < columns; ++column) { + float total = 0.0f; + for (feature = 0; feature < width; ++feature) { + total += (float)in[row * width + feature] + * coefficients[column * width + feature]; + } + out[row * stride + column] = total + bias[column * bias_stride]; + } + } +}""") + +_SUPPORT_TEMPLATE = Template("""\ +static void $name( + float* out, + const $element* in, + const float* support_vectors, + const float* coefficients, + size_t rows, + size_t width, + size_t supports, + int kind, + float gamma, + float coef0, + float degree, + float rho) +{ + size_t row, index; + for (row = 0; row < rows; ++row) { + float total = 0.0f; + for (index = 0; index < supports; ++index) { + total += coefficients[index] * $kernel( + in + row * width, + support_vectors + index * width, + width, + kind, + gamma, + coef0, + degree); + } + out[row] = total + rho; + } +}""") + +_PAIRWISE_TEMPLATE = Template("""\ +static void $name( + float* out, + float* votes, + float* values, + const $element* in, + const float* support_vectors, + const float* coefficients, + const float* rho, + const int32_t* starts, + const int32_t* counts, + size_t rows, + size_t width, + size_t vectors, + size_t classes, + size_t stride, + int kind, + float gamma, + float coef0, + float degree) +{ + size_t row, index, first, second, taken, evaluated; + for (row = 0; row < rows; ++row) { + for (index = 0; index < vectors; ++index) { + values[index] = $kernel( + in + row * width, + support_vectors + index * width, + width, + kind, + gamma, + coef0, + degree); + } + for (index = 0; index < classes; ++index) { + votes[row * classes + index] = 0.0f; + } + evaluated = 0; + for (first = 0; first < classes; ++first) { + for (second = first + 1; second < classes; ++second) { + float total = rho[evaluated]; + float side = 0.0f; + for (taken = 0; taken < (size_t)counts[first]; ++taken) { + const size_t at = (size_t)starts[first] + taken; + side += coefficients[(second - 1) * vectors + at] * values[at]; + } + total += side; + side = 0.0f; + for (taken = 0; taken < (size_t)counts[second]; ++taken) { + const size_t at = (size_t)starts[second] + taken; + side += coefficients[first * vectors + at] * values[at]; + } + total += side; + out[row * stride + evaluated] = total; + votes[row * classes + ((total > 0.0f) ? first : second)] += 1.0f; + ++evaluated; + } + } + } +}""") + +# The two-class case of libsvm's pairwise coupling, which is what turns one decision value +# into a pair of probabilities. `probability` starts uniform and the loop drives the residual +# `Q*p - p'Qp` to zero; comparing that residual and the clamp bounds in double is what the +# reference's own mixed float32/Python-float arithmetic does. +_PROBABILITY_TEMPLATE = Template("""\ +static void $name( + float* out, + const float* scores, + size_t rows, + float prob_a, + float prob_b) +{ + size_t row, iteration, first, second; + for (row = 0; row < rows; ++row) { + const float raw = scores[row] * prob_a + prob_b; + const float mapped = 1.0f / (1.0f + expf(-fabsf(raw))); + float pair = 1.0f - ((raw < 0) ? (1.0f - mapped) : mapped); + float coupling[2][2]; + float probability[2]; + float product[2]; + float complement, total, largest; + if ((double)pair < $smallest_test) { + pair = $smallest; + } + if ((double)pair > $largest_test) { + pair = $largest; + } + complement = 1.0f - pair; + coupling[0][0] = complement * complement; + coupling[0][1] = -complement * pair; + coupling[1][0] = coupling[0][1]; + coupling[1][1] = pair * pair; + probability[0] = 0.5f; + probability[1] = 0.5f; + for (iteration = 0; iteration < $iterations; ++iteration) { + for (first = 0; first < 2; ++first) { + product[first] = coupling[first][0] * probability[0] + + coupling[first][1] * probability[1]; + } + total = probability[0] * product[0] + probability[1] * product[1]; + largest = 0.0f; + for (first = 0; first < 2; ++first) { + const float error = fabsf(product[first] - total); + /* `max(error, largest)` the way Python's own `max` takes it, which keeps a + value that is not a number rather than discarding it. */ + largest = (largest > error) ? largest : error; + } + if ((double)largest < $tolerance) { + break; + } + for (first = 0; first < 2; ++first) { + const float step = + (-product[first] + total) / coupling[first][first]; + const float scale = 1.0f + step; + probability[first] += step; + total = (total + + step * (step * coupling[first][first] + 2.0f * product[first])) + / (scale * scale); + probability[0] /= scale; + probability[1] /= scale; + for (second = 0; second < 2; ++second) { + product[second] = + (product[second] + step * coupling[first][second]) / scale; + } + } + } + out[row * 2] = probability[0]; + out[row * 2 + 1] = probability[1]; + } +}""") + +_ONE_CLASS_TEMPLATE = Template("""\ +static void $name(float* scores, size_t count) +{ + size_t index; + for (index = 0; index < count; ++index) { + scores[index] = (scores[index] > 0.0f) ? 1.0f : -1.0f; + } +}""") + +_THRESHOLD_TEMPLATE = Template("""\ +static void $name( + int64_t* labels, + const float* scores, + size_t rows, + int64_t positive, + float threshold) +{ + size_t row; + for (row = 0; row < rows; ++row) { + labels[row] = (scores[row] >= threshold) ? positive : 0; + } +}""") + + +# -------------------------------------------------------------------------------------- +# The four ops +# -------------------------------------------------------------------------------------- + + +def _linear_regressor(context: NodeContext) -> NodeEmission: + """`Y = X * coefficients' + intercepts`, one column per target.""" + source = context.require_input(0) + result = float_output(context, 0) + rows, width = _rows_and_width(context, source) + targets = context.int_attribute("targets") + if targets < 1: + raise CompileError( + f"Node `{context.label}`: LinearRegressor scores {targets} target(s); it takes " + "at least one." + ) + verify_shape(context, result, (rows, targets)) + + coefficients = _coefficient_matrix(context, targets, width) + emission = _scores( + context, + source, + result.expr, + coefficients, + _intercepts(context, targets), + rows=rows, + width=width, + columns=targets, + stride=targets, + ) + transform = named_transform(context) + return extend(emission, post_transform(context, result, transform, rows, targets)) + + +def _linear_classifier(context: NodeContext) -> NodeEmission: + """The same scores, transformed, and the class the winning column names. + + A single score column against two class labels is the binary case: the score is paired + with its own negation before the transform, and the pair is then read like any other row. + A single column that is *not* paired names its class by a threshold instead — zero on + untransformed scores, one half on anything the transform has mapped into `[0, 1]`. + """ + source = context.require_input(0) + labels = label_output(context, 0) + scores = float_output(context, 1) + rows, width = _rows_and_width(context, source) + classes = _class_labels(context, required=False) + coefficients = _required_floats(context, "coefficients") + if width == 0 or len(coefficients) % width: + raise CompileError( + f"Node `{context.label}`: LinearClassifier holds {len(coefficients)} " + f"coefficient(s) for an input of {width} feature(s); it takes one per feature " + "per class." + ) + produced = len(coefficients) // width + columns = 2 if produced == 1 and len(classes) == 2 else produced + # Read before the buffers are checked: ONNX's own inference sizes `Z` from the intercepts + # as much as from the class labels, so a node that sets none reaches this first. + intercepts = _intercepts(context, produced) + verify_shape(context, labels, (rows,)) + verify_shape(context, scores, (rows, columns)) + + transform = named_transform(context) + emission = _scores( + context, + source, + scores.expr, + coefficients.reshape(produced, width), + intercepts, + rows=rows, + width=width, + columns=produced, + stride=columns, + ) + if columns != produced: + emission = extend( + emission, binary_scores(context, scores, rows, columns, complement=False) + ) + emission = extend( + emission, post_transform(context, scores, transform, rows, columns) + ) + if columns > 1: + if classes and len(classes) != columns: + raise CompileError( + f"Node `{context.label}`: LinearClassifier scores {columns} column(s) " + f"against {len(classes)} class label(s); the winning column names a class, " + "so it takes one label per column." + ) + return extend( + emission, + argmax_labels( + context, + labels, + scores.expr, + scores.elem_type, + classes or tuple(range(columns)), + rows, + columns, + ), + ) + return extend( + emission, _threshold_labels(context, labels, scores, classes, transform, rows) + ) + + +def _svm_regressor(context: NodeContext) -> NodeEmission: + """One score per row: a kernel against every support vector, or one plain dot product.""" + source = context.require_input(0) + result = float_output(context, 0) + rows, width = _rows_and_width(context, source) + verify_shape(context, result, (rows, 1)) + + coefficients = _required_floats(context, "coefficients") + rho = _required_floats(context, "rho") + supports = context.int_attribute("n_supports") + if supports > 0: + emission = _support_scores( + context, + source, + result.expr, + _support_vectors(context, supports, width), + # The reference reads one coefficient per support vector and ignores the rest. + coefficients[:supports], + rho[0], + rows=rows, + width=width, + supports=supports, + ) + else: + if len(coefficients) != width: + raise CompileError( + f"Node `{context.label}`: SVMRegressor holds {len(coefficients)} " + f"coefficient(s) for an input of {width} feature(s); with no support " + "vectors it scores one plain dot product, which takes one per feature." + ) + emission = _scores( + context, + source, + result.expr, + coefficients.reshape(1, width), + rho[:1], + rows=rows, + width=width, + columns=1, + stride=1, + ) + if context.int_attribute("one_class"): + emission = extend(emission, _one_class(context, result, rows)) + transform = named_transform(context) + return extend(emission, post_transform(context, result, transform, rows, 1)) + + +def _svm_classifier(context: NodeContext) -> NodeEmission: + """Pairwise votes over support vectors, or one score per class, and then a label.""" + source = context.require_input(0) + labels = label_output(context, 0) + scores = float_output(context, 1) + rows, width = _rows_and_width(context, source) + classes = _class_labels(context, required=True) + coefficients = _required_floats(context, "coefficients") + rho = _required_floats(context, "rho") + counts = _vector_counts(context) + vectors = sum(counts) + transform = named_transform(context) + probabilities = _probability_pair(context, classes) if vectors > 0 else None + + # One score per class in the linear mode and one per class *pair* over support vectors, + # unless those pairs are coupled into a probability per class. + produced = ( + len(classes) + if vectors == 0 or probabilities is not None + else _pair_count(len(classes)) + ) + # A row of one score is paired with a second only where a second class is called for: + # one `rho`, two class labels, and a transform that is not the one `write_scores` answers + # before it ever reaches the pairing. + paired = ( + produced == 1 and len(rho) == 1 and len(classes) == 2 and transform != PROBIT + ) + columns = 2 if paired else produced + verify_shape(context, labels, (rows,)) + verify_shape(context, scores, (rows, columns)) + + emission, ranking = _svm_scores( + context, + source, + scores, + classes, + coefficients, + rho, + counts, + probabilities, + rows=rows, + width=width, + vectors=vectors, + columns=columns, + ) + emission = extend( + emission, + argmax_labels( + context, + labels, + ranking.expr, + TensorProto.FLOAT, + classes, + rows, + ranking.columns, + _label_rule(classes, rho, coefficients, probabilities is not None), + ), + ) + if paired: + emission = extend( + emission, binary_scores(context, scores, rows, columns, complement=False) + ) + # A single score the pairing left alone never reaches the transform at all, which is + # where `write_scores` returns it — unless the transform is the one it answers first. + if columns > 1 or transform == PROBIT: + emission = extend( + emission, post_transform(context, scores, transform, rows, columns) + ) + return emission + + +# -------------------------------------------------------------------------------------- +# Emission +# -------------------------------------------------------------------------------------- + + +class _Ranking(NamedTuple): + """The values a classifier's label is decided from, and how wide a row of them is. + + The support-vector mode ranks the classes by their votes and the linear mode by the + scores themselves, so this is a votes buffer in the first case and the score tensor in + the second — read before anything transforms it, which is the order the reference + computes the two in. + """ + + expr: str + columns: int + + +def _scores( + context: NodeContext, + source: TensorRef, + destination: str, + coefficients: np.ndarray, + bias: np.ndarray, + *, + rows: int, + width: int, + columns: int, + stride: int, +) -> NodeEmission: + """One dot product per output column, offset by the bias that column carries. + + A bias of one value covers every column, which is how the reference broadcasts a single + intercept — or the single `rho` a support vector machine shares — over the score matrix. + """ + element = c_type(source.elem_type) + name = f"{context.prefix}_ml_scores_{element}" + definition = _SCORES_TEMPLATE.substitute(name=name, element=element) + coefficient_data, coefficient_symbol = constant_data( + context, "coefficients", coefficients + ) + bias_data, bias_symbol = constant_data(context, "bias", bias) + call = call_kernel( + name, + [ + destination, + source.expr, + coefficient_symbol, + bias_symbol, + f"{rows}u", + f"{width}u", + f"{columns}u", + f"{stride}u", + f"{int(len(bias) > 1)}u", + ], + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call,), + constants=(coefficient_data, bias_data), + ) + + +def _support_scores( + context: NodeContext, + source: TensorRef, + destination: str, + vectors: np.ndarray, + coefficients: np.ndarray, + rho: float, + *, + rows: int, + width: int, + supports: int, +) -> NodeEmission: + """Every support vector's kernel value against the row, weighted and summed.""" + element = c_type(source.elem_type) + kernel = _kernel_function(context, element) + name = f"{context.prefix}_svm_supports_{element}" + definition = _SUPPORT_TEMPLATE.substitute( + name=name, element=element, kernel=kernel.name + ) + vector_data, vector_symbol = constant_data(context, "support_vectors", vectors) + coefficient_data, coefficient_symbol = constant_data( + context, "coefficients", coefficients + ) + call = call_kernel( + name, + [ + destination, + source.expr, + vector_symbol, + coefficient_symbol, + f"{rows}u", + f"{width}u", + f"{supports}u", + *_kernel_arguments(context), + scalar_literal(rho, TensorProto.FLOAT), + ], + ) + return NodeEmission( + functions=(kernel, CFunction(name, definition)), + statements=(call,), + constants=(vector_data, coefficient_data), + ) + + +def _svm_scores( + context: NodeContext, + source: TensorRef, + scores: TensorRef, + classes: tuple[int, ...], + coefficients: np.ndarray, + rho: np.ndarray, + counts: list[int], + probabilities: tuple[float, float] | None, + *, + rows: int, + width: int, + vectors: int, + columns: int, +) -> tuple[NodeEmission, _Ranking]: + """The scores a classifier's row holds, and what its label is then ranked from.""" + if vectors == 0: + if len(coefficients) != len(classes) * width: + raise CompileError( + f"Node `{context.label}`: SVMClassifier holds {len(coefficients)} " + f"coefficient(s) for {len(classes)} class(es) over {width} feature(s); with " + "no support vectors it scores one dot product per class, which takes one " + "coefficient per feature per class." + ) + emission = _scores( + context, + source, + scores.expr, + coefficients.reshape(len(classes), width), + rho[:1], + rows=rows, + width=width, + columns=len(classes), + stride=columns, + ) + return emission, _Ranking(scores.expr, len(classes)) + + votes = ScratchBuffer( + f"{context.prefix}_svm_votes", TensorProto.FLOAT, rows * len(classes) + ) + # A node computing probabilities scores the class pairs into working storage first: the + # result holds one column per class rather than one per pair. + pairs = ( + None + if probabilities is None + else ScratchBuffer(f"{context.prefix}_svm_pairs", TensorProto.FLOAT, rows) + ) + emission = _pairwise_scores( + context, + source, + scores.expr if pairs is None else pairs.symbol, + classes, + coefficients, + rho, + counts, + votes, + rows=rows, + width=width, + vectors=vectors, + stride=columns if pairs is None else 1, + ) + if probabilities is not None and pairs is not None: + emission = extend( + emission, _probabilities(context, scores, pairs, probabilities, rows) + ) + return emission, _Ranking(votes.symbol, len(classes)) + + +def _pairwise_scores( + context: NodeContext, + source: TensorRef, + destination: str, + classes: tuple[int, ...], + coefficients: np.ndarray, + rho: np.ndarray, + counts: list[int], + votes: ScratchBuffer, + *, + rows: int, + width: int, + vectors: int, + stride: int, +) -> NodeEmission: + """libsvm's one-against-one scoring: a decision value and a vote for every class pair.""" + if len(classes) < 2: + raise CompileError( + f"Node `{context.label}`: SVMClassifier declares {len(classes)} class(es) over " + "support vectors; the pairwise scheme it scores them with takes at least two, " + "and ONNX's own reference implementation refuses anything less." + ) + if len(counts) < len(classes): + raise CompileError( + f"Node `{context.label}`: `vectors_per_class` holds {len(counts)} entry(s) for " + f"{len(classes)} class(es); it takes one per class." + ) + if len(coefficients) % vectors: + raise CompileError( + f"Node `{context.label}`: SVMClassifier holds {len(coefficients)} " + f"coefficient(s) over {vectors} support vector(s); it takes a whole number of " + "rows of them." + ) + if len(coefficients) // vectors < len(classes) - 1: + raise CompileError( + f"Node `{context.label}`: SVMClassifier holds {len(coefficients) // vectors} " + f"row(s) of coefficients for {len(classes)} class(es); the pairwise scheme " + f"reads {len(classes) - 1} of them." + ) + pairs = _pair_count(len(classes)) + if len(rho) < pairs: + raise CompileError( + f"Node `{context.label}`: `rho` holds {len(rho)} value(s) for the {pairs} class " + "pair(s) this node scores; it takes one per pair." + ) + + element = c_type(source.elem_type) + kernel = _kernel_function(context, element) + name = f"{context.prefix}_svm_pairwise_{element}" + definition = _PAIRWISE_TEMPLATE.substitute( + name=name, element=element, kernel=kernel.name + ) + values = ScratchBuffer(f"{context.prefix}_svm_values", TensorProto.FLOAT, vectors) + starts = np.cumsum([0, *counts[: len(classes) - 1]], dtype=np.int32) + tables = [ + constant_data( + context, "support_vectors", _support_vectors(context, vectors, width) + ), + constant_data(context, "coefficients", coefficients), + constant_data(context, "rho", rho[:pairs]), + constant_data(context, "starts", starts), + constant_data(context, "counts", np.array(counts[: len(classes)], np.int32)), + ] + call = call_kernel( + name, + [ + destination, + votes.symbol, + values.symbol, + source.expr, + *(symbol for _, symbol in tables), + f"{rows}u", + f"{width}u", + f"{vectors}u", + f"{len(classes)}u", + f"{stride}u", + *_kernel_arguments(context), + ], + ) + return NodeEmission( + functions=(kernel, CFunction(name, definition)), + statements=(call,), + scratch=(votes, values), + constants=tuple(data for data, _ in tables), + ) + + +def _probabilities( + context: NodeContext, + scores: TensorRef, + pairs: ScratchBuffer, + pair: tuple[float, float], + rows: int, +) -> NodeEmission: + """The two class probabilities libsvm's Platt scaling turns one decision value into. + + `pairs` is the working storage the pairwise scoring wrote its one decision value per row + into, which is where this reads them from. + """ + name = f"{context.prefix}_svm_probabilities" + definition = _PROBABILITY_TEMPLATE.substitute( + name=name, + smallest_test=scalar_literal(_SMALLEST_PROBABILITY, TensorProto.DOUBLE), + smallest=scalar_literal(_SMALLEST_PROBABILITY, TensorProto.FLOAT), + largest_test=scalar_literal(_LARGEST_PROBABILITY, TensorProto.DOUBLE), + largest=scalar_literal(_LARGEST_PROBABILITY, TensorProto.FLOAT), + tolerance=scalar_literal(_COUPLING_TOLERANCE, TensorProto.DOUBLE), + iterations=f"{_COUPLING_ITERATIONS}u", + ) + call = call_kernel( + name, + [ + scores.expr, + pairs.symbol, + f"{rows}u", + scalar_literal(pair[0], TensorProto.FLOAT), + scalar_literal(pair[1], TensorProto.FLOAT), + ], + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call,), + scratch=(pairs,), + ) + + +def _one_class(context: NodeContext, result: TensorRef, rows: int) -> NodeEmission: + """Each score replaced by which side of zero it falls on.""" + name = f"{context.prefix}_svm_one_class" + definition = _ONE_CLASS_TEMPLATE.substitute(name=name) + call = call_kernel(name, [result.expr, f"{rows}u"]) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _threshold_labels( + context: NodeContext, + labels: TensorRef, + scores: TensorRef, + classes: tuple[int, ...], + transform: int, + rows: int, +) -> NodeEmission: + """A single score column labelled by which side of a threshold it falls on. + + The threshold is zero while the scores are the raw ones and one half once a transform has + mapped them onto a probability, and a row below it is labelled 0 whatever the class table + says — both of which are the reference implementation's. + """ + name = f"{context.prefix}_ml_threshold" + definition = _THRESHOLD_TEMPLATE.substitute(name=name) + call = call_kernel( + name, + [ + labels.expr, + scores.expr, + f"{rows}u", + scalar_literal(classes[0] if classes else 1, TensorProto.INT64), + scalar_literal(0.0 if transform == NONE else 0.5, TensorProto.FLOAT), + ], + ) + return NodeEmission(functions=(CFunction(name, definition),), statements=(call,)) + + +def _kernel_function(context: NodeContext, element: str) -> CFunction: + name = f"{context.prefix}_svm_kernel_{element}" + return CFunction( + name, + _KERNEL_TEMPLATE.substitute( + name=name, element=element, rbf=_RBF, poly=_POLY, sigmoid=_SIGMOID + ), + ) + + +def _kernel_arguments(context: NodeContext) -> list[str]: + """The kernel function the node names and the three parameters it reads.""" + declared = context.string_attribute("kernel_type") + kind = choice(context, "kernel_type", declared.lower(), _KERNEL_TYPES) + parameters = [float(value) for value in context.attribute("kernel_params", [])] + if parameters and len(parameters) < _KERNEL_PARAMETERS: + raise CompileError( + f"Node `{context.label}`: `kernel_params` holds {len(parameters)} value(s); it " + f"takes the {_KERNEL_PARAMETERS} its reference implementation reads — gamma, " + "coef0 and degree — or none at all." + ) + # No `kernel_params` leaves all three at zero, which is what the reference falls back on. + gamma, coef0, degree = parameters[:_KERNEL_PARAMETERS] or [0.0, 0.0, 0.0] + return [ + str(kind), + scalar_literal(gamma, TensorProto.FLOAT), + scalar_literal(coef0, TensorProto.FLOAT), + # The exponent is read as a whole number and applied as one. + scalar_literal(float(int(degree)), TensorProto.FLOAT), + ] + + +# -------------------------------------------------------------------------------------- +# Reading the attribute tables +# -------------------------------------------------------------------------------------- + + +def _pair_count(classes: int) -> int: + return classes * (classes - 1) // 2 + + +def _label_rule( + classes: tuple[int, ...], + rho: np.ndarray, + coefficients: np.ndarray, + probabilities: bool, +) -> int: + """Which of `set_score_svm`'s three readings of the winning column this node takes.""" + if len(rho) != 1: + return ARGMAX_LABEL + if len(classes) != 2: + return SIGN_LABEL + positive = bool(coefficients.size) and float(coefficients.min()) >= 0 + return POSITIVE_CLASS if positive and not probabilities else ARGMAX_LABEL + + +def _probability_pair( + context: NodeContext, classes: tuple[int, ...] +) -> tuple[float, float] | None: + """`prob_a`/`prob_b`, or None where the node carries no probabilities. + + Only the two-class ensemble is served: the coupling the reference solves for three + classes and up is built from a matrix it fills by broadcast where libsvm writes one + entry per pair, so its answer and onnxruntime's differ and neither can vouch for a + kernel. The two agree exactly for a single pair. + """ + first = [float(value) for value in context.attribute("prob_a", [])] + second = [float(value) for value in context.attribute("prob_b", [])] + if not first: + return None + if len(classes) != 2: + raise CompileError( + f"Node `{context.label}`: `prob_a` couples the {_pair_count(len(classes))} " + f"pairwise score(s) of {len(classes)} classes into probabilities, which the C " + "compiler supports for two classes only — ONNX's own reference implementation " + "and onnxruntime disagree on the coupling beyond that, so nothing can vouch for " + "a kernel; re-export the model without `prob_a`/`prob_b`." + ) + if not second: + raise CompileError( + f"Node `{context.label}`: SVMClassifier sets `prob_a` and no `prob_b`; the " + "probability of a class pair is read from one of each." + ) + return first[0], second[0] + + +def _vector_counts(context: NodeContext) -> list[int]: + """`vectors_per_class`: how many of the support vectors each class brought with it. + + Every count is a length the emitted loops run to and an offset they start from, so a + negative one would send them off both ends of the tables. The reference implementation + reads it as an empty slice and scores the pair as zero, which is not a reading anything + can vouch for. + """ + counts = [int(value) for value in context.attribute("vectors_per_class", [])] + if any(count < 0 for count in counts): + raise CompileError( + f"Node `{context.label}`: `vectors_per_class` holds a negative count " + f"({', '.join(str(count) for count in counts)}); each entry is how many of the " + "support vectors belong to one class." + ) + return counts + + +def _class_labels(context: NodeContext, *, required: bool) -> tuple[int, ...]: + """The class values a classifier labels its rows with.""" + if list(context.attribute("classlabels_strings", [])): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` labels its rows with the " + "strings in `classlabels_strings`, and a tensor of STRING at run time is not " + "something the C compiler supports." + ) + integers = tuple(int(value) for value in context.attribute("classlabels_ints", [])) + if not integers and required: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` sets no `classlabels_ints`; " + "it labels each row with one of them." + ) + return integers + + +def _coefficient_matrix(context: NodeContext, columns: int, width: int) -> np.ndarray: + """`coefficients` as the `[columns, width]` matrix the op reads it as.""" + values = _required_floats(context, "coefficients") + if len(values) != columns * width: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` holds {len(values)} " + f"coefficient(s) for {columns} target(s) over {width} feature(s); it takes one " + "per feature per target." + ) + return values.reshape(columns, width) + + +def _intercepts(context: NodeContext, columns: int) -> np.ndarray: + """`intercepts`, which every one of these models has to carry. + + ONNX's own reference implementation reads a missing `intercepts` as a NaN and adds it to + every score, while onnxruntime reads it as no offset at all; a model that sets none is + refused rather than compiled to one of the two. + """ + values = [float(value) for value in context.attribute("intercepts", [])] + if not values: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` sets no `intercepts`. ONNX's " + "own reference implementation scores every element of such a model as NaN and " + "onnxruntime ignores the attribute, so nothing can vouch for a kernel built " + "without them; re-export the model with explicit intercepts." + ) + if len(values) not in (1, columns): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` holds {len(values)} " + f"intercept(s) for {columns} score column(s); it takes either one per column or " + "a single one for all of them." + ) + return np.array(values, np.float32) + + +def _support_vectors(context: NodeContext, supports: int, width: int) -> np.ndarray: + """`support_vectors` as the `[supports, width]` matrix the kernel compares against.""" + values = _required_floats(context, "support_vectors") + if len(values) != supports * width: + raise CompileError( + f"Node `{context.label}`: `support_vectors` holds {len(values)} value(s) for " + f"{supports} support vector(s) over {width} feature(s); it takes one per " + "feature per vector." + ) + return values.reshape(supports, width) + + +def _required_floats(context: NodeContext, name: str) -> np.ndarray: + values = [float(value) for value in context.attribute(name, [])] + if not values: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` requires the `{name}` " + "attribute." + ) + return np.array(values, np.float32) + + +def _rows_and_width(context: NodeContext, source: TensorRef) -> tuple[int, int]: + """How many rows the node scores, and how many features each of them holds.""" + if len(source.shape) != 2: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads an `[N, F]` matrix of " + f"features, and `{source.name}` has shape {list(source.shape)}." + ) + return source.shape + + +register_kernel(ML_DOMAIN, "LinearRegressor", _VERSIONS, _linear_regressor) +register_kernel(ML_DOMAIN, "LinearClassifier", _VERSIONS, _linear_classifier) +register_kernel(ML_DOMAIN, "SVMRegressor", _VERSIONS, _svm_regressor) +register_kernel(ML_DOMAIN, "SVMClassifier", _VERSIONS, _svm_classifier) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/tfidf.py b/src/python/fnnx/extras/compilers/c/onnx/ops/tfidf.py new file mode 100644 index 0000000..99eebcd --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/tfidf.py @@ -0,0 +1,355 @@ +"""`TfIdfVectorizer`: counting the n-grams of a pool in a sequence of tokens. + +Everything the op matches against — the pool, which slice of it holds the n-grams of each +length, and the column each n-gram counts into — arrives in attributes, so the search +structure is built at compile time and emitted as `static const` tables. The reference +implementation walks a trie of pool entries, one level per token of an n-gram; the same trie +is flattened here into a node table (the column a node counts into, and the slice of the edge +table its outgoing edges occupy) and an edge table (the token an edge is taken on, and the +node it leads to). The kernel is then the reference's own walk over those tables, at the skip +distances and gram lengths the node's attributes ask for. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass, field +from string import Template + +import numpy as np +from onnx import TensorProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type, element_type_name +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + ScratchBuffer, + TensorRef, + constant_data, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + kernel_name, + verify_shape, +) + +_VERSIONS = (9,) + +# What each weighting mode makes of a column's count, given the weights the node carries. +# ONNX defines the three over the same counts, differing in this expression alone, so the +# mode is a compile-time choice rather than a branch the kernel takes per element. Without +# weights `TFIDF` is the count itself, which is what the reference computes for it. +_WEIGHTINGS = { + ("TF", False): "(float)count", + ("TF", True): "(float)count", + ("IDF", False): "count > 0 ? 1.0f : 0.0f", + ("IDF", True): "count > 0 ? weights[column] : 0.0f", + ("TFIDF", False): "(float)count", + ("TFIDF", True): "weights[column] * (float)count", +} + +_TEMPLATE = Template("""\ +static void $name( + float* out, + const $element* in, + int64_t* counts, + size_t rows, + size_t columns, + size_t width, + const int64_t* tokens, + const int32_t* targets, + const int32_t* first_edge, + const int32_t* edge_count, + const int32_t* counted_column, +$weights int min_gram, + int max_gram, + int max_skip) +{ + size_t row, column, index; + for (index = 0; index < rows * width; ++index) { + counts[index] = 0; + } + for (row = 0; row < rows; ++row) { + int start_gram = min_gram; + int skip; + for (skip = 1; skip <= max_skip + 1; ++skip) { + size_t start; + for (start = 0; start < columns; ++start) { + int32_t node = 0; + int gram = 1; + size_t item = start; + if (start + (size_t)skip * (size_t)(start_gram - 1) >= columns) { + break; + } + while (edge_count[node] > 0 && gram <= max_gram && item < columns) { + const int64_t token = (int64_t)in[row * columns + item]; + const int32_t last = first_edge[node] + edge_count[node]; + int32_t edge, taken = -1; + for (edge = first_edge[node]; edge < last; ++edge) { + if (tokens[edge] == token) { + taken = targets[edge]; + break; + } + } + if (taken < 0) { + break; + } + if (gram >= start_gram && counted_column[taken] >= 0) { + counts[row * width + (size_t)counted_column[taken]] += 1; + } + node = taken; + ++gram; + item += skip; + } + } + if (start_gram == 1) { + start_gram = 2; + if (start_gram > max_gram) { + break; + } + } + } + } + for (row = 0; row < rows; ++row) { + for (column = 0; column < width; ++column) { + const int64_t count = counts[row * width + column]; + out[row * width + column] = $weighting; + } + } +}""") + + +@dataclass +class _Trie: + """The pool's n-grams as a trie, flattened into the tables the kernel indexes. + + One entry per node in `first_edge`, `edge_count` and `counted_column`; one per edge in + `tokens` and `targets`. `counted_column` is -1 for a node no pool n-gram ends at, which + is how a prefix that is only on the way to a longer n-gram counts nothing of its own. + Node 0 is the root, and exists even for a pool the node registers nothing from. + """ + + counted_column: list[int] = field(default_factory=lambda: [-1]) + edges: list[dict[int, int]] = field(default_factory=lambda: [{}]) + + def child(self, node: int, token: int) -> int: + """The node `token` leads to from `node`, added if this is the first n-gram to use it.""" + existing = self.edges[node].get(token) + if existing is not None: + return existing + self.edges[node][token] = len(self.edges) + self.edges.append({}) + self.counted_column.append(-1) + return len(self.edges) - 1 + + @property + def tokens(self) -> list[int]: + return [token for edges in self.edges for token in edges] + + @property + def targets(self) -> list[int]: + return [target for edges in self.edges for target in edges.values()] + + @property + def first_edge(self) -> list[int]: + starts, total = [], 0 + for edges in self.edges: + starts.append(total) + total += len(edges) + return starts + + @property + def edge_count(self) -> list[int]: + return [len(edges) for edges in self.edges] + + +def _tf_idf_vectorizer(context: NodeContext) -> NodeEmission: + """One count per pool n-gram per row, weighted as the node's `mode` prescribes. + + The op reads its input as a batch of token sequences — a matrix is one sequence per row + and anything of lower rank a single sequence — which is the reading the shape of its + result follows from. + """ + source = context.require_input(0) + result = context.require_output(0) + if source.elem_type not in (TensorProto.INT32, TensorProto.INT64): + raise CompileError( + f"Node `{context.label}`: TfIdfVectorizer takes `int32` or `int64` tokens, not " + f"`{element_type_name(source.elem_type)}`; a string pool is matched against a " + "run-time string tensor, which the C compiler does not support." + ) + if len(source.shape) > 2: + raise CompileError( + f"Node `{context.label}`: TfIdfVectorizer takes one token sequence or a batch " + f"of them, but `{source.name}` has shape {list(source.shape)}." + ) + # A scalar is one sequence of one token, its shape being the empty product; a vector is + # one sequence of as many tokens as it holds, an empty one included. + batched = len(source.shape) == 2 + rows = source.shape[0] if batched else 1 + columns = source.shape[1] if batched else math.prod(source.shape) + + indexes = [int(value) for value in context.attribute("ngram_indexes", [])] + if not indexes or min(indexes) < 0: + raise CompileError( + f"Node `{context.label}`: TfIdfVectorizer needs a non-negative `ngram_indexes` " + "entry for every n-gram of its pool." + ) + width = max(indexes) + 1 + verify_shape(context, result, (rows, width) if batched else (width,)) + + weights = [float(value) for value in context.attribute("weights", [])] + if weights and len(weights) != width: + raise CompileError( + f"Node `{context.label}`: TfIdfVectorizer carries {len(weights)} weight(s) for " + f"{width} output column(s); the op weights one per column." + ) + return _emit( + context, source, result, _trie(context, indexes), weights, rows, columns, width + ) + + +def _trie(context: NodeContext, indexes: Sequence[int]) -> _Trie: + """The pool's n-grams, level by level, in the order ONNX numbers them. + + `ngram_counts` splits the pool by n-gram length: entry `i` is where the n-grams of length + `i + 1` start, and they run to the next entry. Every n-gram of the pool takes the next + identifier whether or not its length is one this node counts, so the identifiers — and + with them the `ngram_indexes` entry each n-gram counts into — stay aligned to the pool + however `min_gram_length` and `max_gram_length` narrow it. + """ + minimum = context.int_attribute("min_gram_length") + maximum = context.int_attribute("max_gram_length") + if minimum < 1 or maximum < minimum: + raise CompileError( + f"Node `{context.label}`: TfIdfVectorizer needs 1 <= `min_gram_length` <= " + f"`max_gram_length`, but they are {minimum} and {maximum}." + ) + if context.attribute("pool_strings", []): + raise CompileError( + f"Node `{context.label}`: TfIdfVectorizer sets `pool_strings`, which matches " + "against a run-time string tensor; the C compiler supports the `pool_int64s` " + "form only." + ) + pool = [int(value) for value in context.attribute("pool_int64s", [])] + counts = [int(value) for value in context.attribute("ngram_counts", [])] + + trie = _Trie() + identifier = 1 + for length, start in enumerate(counts, start=1): + end = counts[length] if length < len(counts) else len(pool) + available = (end - start) // length if end > start else 0 + if minimum <= length <= maximum: + identifier = _register( + context, trie, pool[start:end], length, available, identifier, indexes + ) + else: + identifier += available + return trie + + +def _register( + context: NodeContext, + trie: _Trie, + entries: Sequence[int], + length: int, + available: int, + identifier: int, + indexes: Sequence[int], +) -> int: + """Add `available` n-grams of `length` consecutive tokens, returning the next identifier. + + An n-gram the pool lists twice keeps the last of its `ngram_indexes` entries, which is + what re-walking the same path and overwriting the column at its end leaves behind. + """ + position = 0 + for _ in range(available): + node = 0 + for taken in range(1, length + 1): + if position >= len(entries): + break + node = trie.child(node, entries[position]) + position += 1 + if taken == length: + if identifier > len(indexes): + raise CompileError( + f"Node `{context.label}`: TfIdfVectorizer's pool holds more n-grams " + f"than its {len(indexes)} `ngram_indexes` entries account for." + ) + trie.counted_column[node] = indexes[identifier - 1] + identifier += 1 + return identifier + + +def _emit( + context: NodeContext, + source: TensorRef, + result: TensorRef, + trie: _Trie, + weights: Sequence[float], + rows: int, + columns: int, + width: int, +) -> NodeEmission: + mode = context.string_attribute("mode") + weighting = _WEIGHTINGS.get((mode, bool(weights))) + if weighting is None: + raise CompileError( + f"Node `{context.label}`: TfIdfVectorizer sets `mode` to `{mode}`; ONNX defines " + "`TF`, `IDF` and `TFIDF`." + ) + # `TF` ignores the weights a node may still carry, and so does `TFIDF` where there are + # none to apply; passing a table the expression never reads would be an unused parameter, + # which the artifact's `-Werror` build contract refuses. + weighted = "weights[" in weighting + element = c_type(source.elem_type) + name = kernel_name( + context, mode.lower(), "weighted" if weighted else "flat", element + ) + definition = _TEMPLATE.substitute( + name=name, + element=element, + weights=" const float* weights,\n" if weighted else "", + weighting=weighting, + ) + + tables = [ + constant_data(context, role, np.array(values, dtype=dtype)) + for role, values, dtype in ( + ("tokens", trie.tokens, np.int64), + ("targets", trie.targets, np.int32), + ("first_edge", trie.first_edge, np.int32), + ("edge_count", trie.edge_count, np.int32), + ("counted_column", trie.counted_column, np.int32), + *((("weights", weights, np.float32),) if weighted else ()), + ) + ] + counts = ScratchBuffer(f"{name}_counts", TensorProto.INT64, rows * width) + call = call_kernel( + name, + [ + result.expr, + source.expr, + counts.symbol, + f"{rows}u", + f"{columns}u", + f"{width}u", + *(symbol for _, symbol in tables), + str(context.int_attribute("min_gram_length")), + str(context.int_attribute("max_gram_length")), + str(context.int_attribute("max_skip_count")), + ], + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call,), + scratch=(counts,), + constants=tuple(data for data, _ in tables), + ) + + +register_kernel("", "TfIdfVectorizer", _VERSIONS, _tf_idf_vectorizer) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/tree.py b/src/python/fnnx/extras/compilers/c/onnx/ops/tree.py new file mode 100644 index 0000000..3ddad20 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/tree.py @@ -0,0 +1,1027 @@ +"""Decision forests: ONNX-ML's two legacy ensembles and the opset-5 op that supersedes them. + +`TreeEnsembleRegressor`, `TreeEnsembleClassifier` and `TreeEnsemble` are three encodings of +one computation: every row of `X` falls through a set of binary trees, and the weights of the +leaves it lands in are aggregated per target. What differs is the layout. The legacy pair name +a node by a `(tree, node)` pair, mark leaves with a `LEAF` mode among the branch tests, and +carry the leaf weights in a second family of attributes keyed by that same pair; opset 5 +splits interior nodes from leaves into two families indexed directly, and moves the tests, the +splits and the weights into tensors. Both are normalized here into one flat form — nodes, +leaves, `(target, weight)` entries and roots — so a single walker serves all three ops, and +everything it reads is `static const` data laid out at compile time. + +Three places the emitted code follows ONNX's reference implementation rather than the prose, +because the reference is the only oracle these ops have: + +* A `BRANCH_NEQ` node of a *legacy* ensemble sends a NaN feature down the branch its + `missing_value_tracks_true` flag names, while opset 5's `NEQ` sends it down the true branch + outright — `NaN != split` being true. The two differ by one table entry, so the legacy test + is emitted as a branch test of its own rather than the walker taking a flag. +* A set test reads its members off `membership_values` in the order the reference builds the + trees in — depth first from each root, true branch before false — and a member that is zero + ends the set early, which is what its `while (m := next(it)) and not isnan(m)` does. +* `TreeEnsembleClassifier` widens a one-class ensemble's scores to two columns and derives the + first from the second; that binary rule, and the `argmax` over the result, are the + reference's. + +What the compiler refuses: the `*_as_tensor` attribute families opset 3 added for +double-precision tables. The reference evaluator stores them and then reads the float32 +families anyway, so nothing could vouch for what a kernel built from them computes. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from string import Template + +import numpy as np +from onnx.numpy_helper import to_array + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import ( + c_type, + element_type_name, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.emit import scalar_literal +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + ConstantData, + NodeContext, + NodeEmission, + TensorRef, + constant_data, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.loader import ML_DOMAIN +from fnnx.extras.compilers.c.onnx.ops.axes import call_kernel, verify_shape +from fnnx.extras.compilers.c.onnx.ops.scores import ( + NONE, + PROBIT, + argmax_labels, + binary_scores, + choice, + extend, + float_output, + label_output, + named_transform, + post_transform, +) + +# The legacy pair changed twice after opset 1: at 3, which added the `*_as_tensor` families, +# and at 5, which deprecated them in favour of `TreeEnsemble` and changed nothing else. One +# generator therefore claims all three revisions, a claim +# `test_the_legacy_ensemble_revisions_are_one_op` holds to the schemas themselves. +_LEGACY_VERSIONS = (1, 3, 5) +_ENSEMBLE_VERSIONS = (5,) + +# The branch tests, in the encoding opset 5 numbers them with. The legacy families name them +# as strings and add `LEAF`, which this compiler turns into a leaf rather than a test. +_BRANCH_TESTS = { + "BRANCH_LEQ": 0, + "BRANCH_LT": 1, + "BRANCH_GTE": 2, + "BRANCH_GT": 3, + "BRANCH_EQ": 4, + "BRANCH_NEQ": 5, +} +_LEAF = "LEAF" +_MEMBER_TEST = 6 +# `BRANCH_NEQ` under the legacy NaN rule; see the module docstring. +_LEGACY_NEQ_TEST = 7 + +# Columns of one row of the emitted node table: the feature it tests, the test, the flags +# below, the two children, and the range of `membership_values` a set test reads. +_NODE_FIELDS = 7 +_TRUE_IS_LEAF = 1 +_FALSE_IS_LEAF = 2 +_MISSING_TRACKS_TRUE = 4 + +# What the walker does with a leaf's weight. AVERAGE differs from SUM only in what is divided +# by the number of trees, which the two divisors carry. +_ACCUMULATE = 0 +_MINIMUM = 1 +_MAXIMUM = 2 +_LEGACY_AGGREGATES = { + "SUM": _ACCUMULATE, + "AVERAGE": _ACCUMULATE, + "MIN": _MINIMUM, + "MAX": _MAXIMUM, +} +_ENSEMBLE_AGGREGATES = {0: _ACCUMULATE, 1: _ACCUMULATE, 2: _MINIMUM, 3: _MAXIMUM} + +# Opset 5 numbers the transforms the way `ops.scores` does, so its lookup exists to reject a +# value ONNX does not define rather than to translate one. +_ENSEMBLE_TRANSFORMS = {value: value for value in range(5)} + +# What is added to every target once the aggregation is done, where the node offsets none. +# Negative zero rather than positive is what makes the addition an identity for *every* value: +# `x + (-0.0) == x` holds for the negative zero a MIN fold can leave behind, where `x + 0.0` +# would flip its sign. +_NO_OFFSET = -0.0 + +_AGGREGATE_TEMPLATE = Template("""\ +static void $name( + $result* scores, + const $element* features, + const int32_t* nodes, + const double* splits, + const double* members, + const int32_t* leaves, + const int32_t* leaf_targets, + const $result* leaf_weights, + const int32_t* roots, + const $result* initial, + const $result* offset, + size_t rows, + size_t width, + size_t tree_count, + size_t target_count, + int aggregate, + $result weight_divisor, + $result total_divisor) +{ + size_t row, tree, target; + for (row = 0; row < rows; ++row) { + $result* out = scores + row * target_count; + const $element* in = features + row * width; + for (target = 0; target < target_count; ++target) { + out[target] = initial[target]; + } + for (tree = 0; tree < tree_count; ++tree) { + size_t index = (size_t)roots[tree * 2]; + int is_leaf = roots[tree * 2 + 1]; + size_t entry, last, member; + while (!is_leaf) { + const int32_t* node = nodes + index * $fields; + const double value = (double)in[node[0]]; + const double split = splits[index]; + int taken = 0; + switch (node[1]) { + case 0: taken = value <= split; break; + case 1: taken = value < split; break; + case 2: taken = value >= split; break; + case 3: taken = value > split; break; + case 4: taken = value == split; break; + case 5: taken = value != split; break; + case $member: + for (member = 0; member < (size_t)node[6]; ++member) { + if (value == members[(size_t)node[5] + member]) { + taken = 1; + } + } + break; + default: + /* The legacy encoding's BRANCH_NEQ, the one test that leaves a value + that is not a number to the missing-value rule below. */ + taken = !isnan(value) && value != split; + break; + } + if (!taken && (node[2] & $missing) != 0 && isnan(value)) { + taken = 1; + } + is_leaf = taken ? (node[2] & $true_leaf) : (node[2] & $false_leaf); + index = (size_t)(taken ? node[3] : node[4]); + } + entry = (size_t)leaves[index * 2]; + last = entry + (size_t)leaves[index * 2 + 1]; + for (; entry < last; ++entry) { + const size_t chosen = (size_t)leaf_targets[entry]; + const double weight = (double)(leaf_weights[entry] / weight_divisor); + if (aggregate == $accumulate) { + out[chosen] = ($result)((double)out[chosen] + weight); + } else if (aggregate == $minimum) { + if (weight < (double)out[chosen]) { + out[chosen] = ($result)weight; + } + } else if (weight > (double)out[chosen]) { + out[chosen] = ($result)weight; + } + } + } + for (target = 0; target < target_count; ++target) { + out[target] = out[target] / total_divisor + offset[target]; + } + } +}""") + + +@dataclass(frozen=True) +class _Ensemble: + """A forest in the one flat form the walker reads, whatever encoding it arrived in. + + `nodes` holds `_NODE_FIELDS` columns per interior node and `roots` an `(index, is_leaf)` + pair per tree; a child index addresses `nodes` or `leaves` according to the node's flags. + Each leaf names a range of the `targets`/`weights` pair list, which is what lets one leaf + of a multi-target legacy ensemble contribute to several targets at once. + """ + + nodes: np.ndarray + splits: np.ndarray + members: np.ndarray + leaves: np.ndarray + targets: np.ndarray + weights: np.ndarray + roots: np.ndarray + + @property + def tree_count(self) -> int: + return len(self.roots) + + +@dataclass(frozen=True) +class _Aggregation: + """How the walker folds the leaves a row reaches into the score of each target.""" + + mode: int + weight_divisor: float + total_divisor: float + initial: np.ndarray + offset: np.ndarray + + +# -------------------------------------------------------------------------------------- +# The three ops +# -------------------------------------------------------------------------------------- + + +def _tree_ensemble_regressor(context: NodeContext) -> NodeEmission: + """Leaf weights aggregated per target, then averaged, offset and transformed.""" + source = context.require_input(0) + result = float_output(context, 0) + rows, width = _rows_and_width(context, source) + targets = context.int_attribute("n_targets") + verify_shape(context, result, (rows, targets)) + _refuse_tensor_tables(context, ("base_values", "nodes_values", "target_weights")) + + ensemble = _legacy_ensemble(context, "target", width, targets) + declared = context.string_attribute("aggregate_function") + mode = choice(context, "aggregate_function", declared, _LEGACY_AGGREGATES) + base = _base_values(context, targets, result.elem_type) + aggregation = _Aggregation( + mode=mode, + weight_divisor=1.0, + # The reference divides the accumulated total by the number of trees and adds the + # base values only afterwards, which is why they offset rather than seed the scores. + total_divisor=float(ensemble.tree_count) if declared == "AVERAGE" else 1.0, + initial=_seed(mode, targets, result.elem_type), + offset=_filled(targets, _NO_OFFSET, result.elem_type) if base is None else base, + ) + emission = _aggregate(context, source, result, ensemble, aggregation, rows, width) + transform = named_transform(context) + return extend(emission, post_transform(context, result, transform, rows, targets)) + + +def _tree_ensemble_classifier(context: NodeContext) -> NodeEmission: + """The same aggregation, then the class the winning column names.""" + source = context.require_input(0) + labels = label_output(context, 0) + scores = float_output(context, 1) + rows, width = _rows_and_width(context, source) + classes = _class_labels(context) + binary = len({int(value) for value in context.attribute("class_ids", [])}) == 1 + # A one-class ensemble's scores are widened to the two columns the binary rule fills, and + # the winning column is then the label itself. + columns = 2 if binary and len(classes) == 1 else len(classes) + table = (0, 1)[:columns] if len(classes) == 1 else classes + verify_shape(context, labels, (rows,)) + verify_shape(context, scores, (rows, columns)) + _refuse_tensor_tables(context, ("base_values", "nodes_values", "class_weights")) + + ensemble = _legacy_ensemble(context, "class", width, len(classes)) + base = _base_values(context, len(classes), scores.elem_type) + initial = _filled(columns, 0.0, scores.elem_type) + if base is not None: + initial[: len(classes)] = base + aggregation = _Aggregation( + mode=_ACCUMULATE, + weight_divisor=1.0, + total_divisor=1.0, + initial=initial, + offset=_filled(columns, _NO_OFFSET, scores.elem_type), + ) + transform = named_transform(context) + emission = _aggregate(context, source, scores, ensemble, aggregation, rows, width) + if binary: + emission = extend( + emission, _binary_scores(context, scores, transform, rows, columns) + ) + emission = extend( + emission, post_transform(context, scores, transform, rows, columns) + ) + return extend( + emission, + argmax_labels( + context, labels, scores.expr, scores.elem_type, table, rows, columns + ), + ) + + +def _tree_ensemble(context: NodeContext) -> NodeEmission: + """The opset-5 op: the same forest, with interior nodes and leaves indexed apart.""" + source = context.require_input(0) + result = context.require_output(0) + rows, width = _rows_and_width(context, source) + targets = context.int_attribute("n_targets") + if result.elem_type != source.elem_type: + raise CompileError( + f"Node `{context.label}`: TreeEnsemble scores in the element type of its input, " + f"but `{source.name}` is `{element_type_name(source.elem_type)}` and its output " + f"`{result.name}` is `{element_type_name(result.elem_type)}`." + ) + verify_shape(context, result, (rows, targets)) + + ensemble = _opset5_ensemble(context, width, targets) + declared = context.int_attribute("aggregate_function") + mode = choice(context, "aggregate_function", declared, _ENSEMBLE_AGGREGATES) + aggregation = _Aggregation( + mode=mode, + # Each weight is divided by the number of trees before it is added, which does not + # round the way dividing the total once would. + weight_divisor=float(ensemble.tree_count) if declared == 0 else 1.0, + total_divisor=1.0, + initial=_seed(mode, targets, result.elem_type), + offset=_filled(targets, _NO_OFFSET, result.elem_type), + ) + emission = _aggregate(context, source, result, ensemble, aggregation, rows, width) + transform = choice( + context, + "post_transform", + context.int_attribute("post_transform"), + _ENSEMBLE_TRANSFORMS, + ) + return extend(emission, post_transform(context, result, transform, rows, targets)) + + +# -------------------------------------------------------------------------------------- +# Emission +# -------------------------------------------------------------------------------------- + + +def _aggregate( + context: NodeContext, + source: TensorRef, + result: TensorRef, + ensemble: _Ensemble, + aggregation: _Aggregation, + rows: int, + width: int, +) -> NodeEmission: + """The one call that walks every tree and folds the leaves it reaches into `result`.""" + element, value = c_type(source.elem_type), c_type(result.elem_type) + name = f"{context.prefix}_tree_aggregate_{element}_{value}" + definition = _AGGREGATE_TEMPLATE.substitute( + name=name, + element=element, + result=value, + fields=_NODE_FIELDS, + member=_MEMBER_TEST, + missing=_MISSING_TRACKS_TRUE, + true_leaf=_TRUE_IS_LEAF, + false_leaf=_FALSE_IS_LEAF, + accumulate=_ACCUMULATE, + minimum=_MINIMUM, + ) + tables, symbols = _tables(context, ensemble, aggregation, result.elem_type) + call = call_kernel( + name, + [ + result.expr, + source.expr, + *symbols, + f"{rows}u", + f"{width}u", + f"{ensemble.tree_count}u", + f"{len(aggregation.initial)}u", + str(aggregation.mode), + scalar_literal(aggregation.weight_divisor, result.elem_type), + scalar_literal(aggregation.total_divisor, result.elem_type), + ], + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(call,), + constants=tables, + ) + + +def _tables( + context: NodeContext, + ensemble: _Ensemble, + aggregation: _Aggregation, + elem_type: int, +) -> tuple[tuple[ConstantData, ...], tuple[str, ...]]: + """Every table the walker reads, as static constant data, in call-site order. + + The weights are accumulated in the element type the scores are held in, so that is what + they are laid out as. Nothing is lost by it: the legacy families carry float32 weights + into a float32 result, and opset 5's are required by ONNX's own type inference to carry + the element type of `X`, which is what its result is scored in. + """ + built = [ + constant_data(context, role, values) + for role, values in ( + ("nodes", ensemble.nodes), + ("splits", ensemble.splits), + ("members", ensemble.members), + ("leaves", ensemble.leaves), + ("leaf_targets", ensemble.targets), + ("leaf_weights", ensemble.weights.astype(numpy_dtype_name(elem_type))), + ("roots", ensemble.roots), + ("initial", aggregation.initial), + ("offset", aggregation.offset), + ) + ] + return tuple(data for data, _ in built), tuple(symbol for _, symbol in built) + + +def _binary_scores( + context: NodeContext, scores: TensorRef, transform: int, rows: int, columns: int +) -> NodeEmission: + """The second column a single-class-weight ensemble's score is paired with. + + Which one it gets is the reference implementation's rule: the complement of the score + where the transform leaves the scale of a probability alone, its negation otherwise. + """ + return binary_scores( + context, scores, rows, columns, complement=transform in (NONE, PROBIT) + ) + + +# -------------------------------------------------------------------------------------- +# Reading a legacy ensemble +# -------------------------------------------------------------------------------------- + + +def _legacy_ensemble( + context: NodeContext, role: str, width: int, targets: int +) -> _Ensemble: + """The `(tree, node)`-keyed encoding, flattened. + + A node is addressed by its position in the `nodes_*` families, while `nodes_truenodeids` + names a node *id* within the same tree, which is resolved here into that position; a + tree's root is its first node in those families rather than the node whose id is zero. + Leaves are given an index space of their own, and the `role_*` families — `target_*` for + the regressor, `class_*` for the classifier — say what each of them contributes. + """ + tests = [ + _branch_test(context, value) for value in context.attribute("nodes_modes", []) + ] + count = len(tests) + tree_ids = _integers(context, "nodes_treeids", count) + node_ids = _integers(context, "nodes_nodeids", count) + features = _integers(context, "nodes_featureids", count) + true_ids = _integers(context, "nodes_truenodeids", count) + false_ids = _integers(context, "nodes_falsenodeids", count) + missing = _integers( + context, "nodes_missing_value_tracks_true", count, optional=True + ) + splits = _floats(context, "nodes_values", count, optional=True) + + positions = {pair: index for index, pair in enumerate(zip(tree_ids, node_ids))} + slots = _index_spaces(tests) + entries = _leaf_entries(context, role, targets) + + nodes = [] + for index, test in enumerate(tests): + if test == _LEAF: + continue + true_child, true_leaf = _resolve( + context, positions, tests, slots, tree_ids[index], true_ids[index] + ) + false_child, false_leaf = _resolve( + context, positions, tests, slots, tree_ids[index], false_ids[index] + ) + nodes.append( + ( + _feature(context, features[index], width), + _LEGACY_NEQ_TEST if test == "BRANCH_NEQ" else _BRANCH_TESTS[test], + (_TRUE_IS_LEAF if true_leaf else 0) + | (_FALSE_IS_LEAF if false_leaf else 0) + | (_MISSING_TRACKS_TRUE if missing[index] else 0), + true_child, + false_child, + 0, + 0, + ) + ) + + leaves: list[tuple[int, int]] = [] + leaf_targets: list[int] = [] + leaf_weights: list[float] = [] + for index, test in enumerate(tests): + if test != _LEAF: + continue + pairs = entries.get((tree_ids[index], node_ids[index]), ()) + leaves.append((len(leaf_targets), len(pairs))) + leaf_targets.extend(target for target, _ in pairs) + leaf_weights.extend(weight for _, weight in pairs) + + roots = [] + for tree_id in sorted(set(tree_ids)): + position = tree_ids.index(tree_id) + leaf = tests[position] == _LEAF + roots.append((slots[position], int(leaf))) + + ensemble = _Ensemble( + nodes=_table(nodes, np.int32, _NODE_FIELDS), + splits=np.array( + [splits[index] for index, test in enumerate(tests) if test != _LEAF], + np.float64, + ), + members=np.zeros(0, np.float64), + leaves=_table(leaves, np.int32, 2), + targets=np.array(leaf_targets, np.int32), + # The reference reads these tables as the float32 the attribute stores them in. + weights=np.array(leaf_weights, np.float32), + roots=_table(roots, np.int32, 2), + ) + _verify_acyclic(context, ensemble) + return ensemble + + +def _index_spaces(tests: Sequence[str]) -> list[int]: + """Where each flat node lands once interior nodes and leaves are indexed apart.""" + slots = [] + interior = leaves = 0 + for test in tests: + if test == _LEAF: + slots.append(leaves) + leaves += 1 + else: + slots.append(interior) + interior += 1 + return slots + + +def _resolve( + context: NodeContext, + positions: Mapping[tuple[int, int], int], + tests: Sequence[str], + slots: Sequence[int], + tree_id: int, + node_id: int, +) -> tuple[int, bool]: + position = positions.get((tree_id, node_id)) + if position is None: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` names node {node_id} of tree " + f"{tree_id} as a child, which none of its `nodes_*` entries defines." + ) + return slots[position], tests[position] == _LEAF + + +def _leaf_entries( + context: NodeContext, role: str, targets: int +) -> dict[tuple[int, int], list[tuple[int, float]]]: + """The `(target, weight)` pairs each leaf contributes, keyed by `(tree, node)`.""" + tree_ids = _integers(context, f"{role}_treeids", None) + node_ids = _integers(context, f"{role}_nodeids", len(tree_ids)) + ids = _integers(context, f"{role}_ids", len(tree_ids)) + weights = _floats(context, f"{role}_weights", len(tree_ids)) + entries: dict[tuple[int, int], list[tuple[int, float]]] = {} + for tree_id, node_id, target, weight in zip(tree_ids, node_ids, ids, weights): + if not 0 <= target < targets: + raise CompileError( + f"Node `{context.label}`: `{role}_ids` names {target}, which is outside the " + f"{targets} target(s) this node scores." + ) + entries.setdefault((tree_id, node_id), []).append((target, weight)) + return entries + + +# -------------------------------------------------------------------------------------- +# Reading an opset-5 ensemble +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _Opset5Nodes: + """The interior-node families of an opset-5 ensemble, read and length-checked.""" + + tests: list[int] + features: list[int] + true_ids: list[int] + false_ids: list[int] + true_leafs: list[int] + false_leafs: list[int] + missing: list[int] + splits: np.ndarray + + def __len__(self) -> int: + return len(self.tests) + + def is_bare_leaf(self, root: int) -> bool: + """Whether the tree rooted here is a single leaf, as the reference reads that. + + Both children being leaves *and* being the same one is what marks the degenerate + tree, and the root's own position is then read as an index into the leaf families. + """ + return bool( + self.true_leafs[root] + and self.false_leafs[root] + and self.true_ids[root] == self.false_ids[root] + ) + + +def _opset5_ensemble(context: NodeContext, width: int, targets: int) -> _Ensemble: + """The opset-5 encoding, which is already the flat form bar the membership ranges.""" + parsed = _opset5_nodes(context) + weights = _required_tensor(context, "leaf_weights") + leaf_targets = _integers(context, "leaf_targetids", len(weights)) + for target in leaf_targets: + if not 0 <= target < targets: + raise CompileError( + f"Node `{context.label}`: `leaf_targetids` names {target}, which is outside " + f"the {targets} target(s) this node scores." + ) + roots = _roots(context, parsed, len(weights)) + members = _membership(context, parsed.tests) + ranges = _membership_ranges(context, parsed, roots, members, len(weights)) + + nodes = [ + ( + _feature(context, parsed.features[index], width), + parsed.tests[index], + (_TRUE_IS_LEAF if parsed.true_leafs[index] else 0) + | (_FALSE_IS_LEAF if parsed.false_leafs[index] else 0) + | (_MISSING_TRACKS_TRUE if parsed.missing[index] else 0), + _child( + context, + parsed.true_ids[index], + parsed.true_leafs[index], + parsed, + len(weights), + ), + _child( + context, + parsed.false_ids[index], + parsed.false_leafs[index], + parsed, + len(weights), + ), + *ranges.get(index, (0, 0)), + ) + for index in range(len(parsed)) + ] + ensemble = _Ensemble( + nodes=_table(nodes, np.int32, _NODE_FIELDS), + splits=parsed.splits.astype(np.float64), + members=np.array(members, np.float64), + leaves=_table([(index, 1) for index in range(len(weights))], np.int32, 2), + targets=np.array(leaf_targets, np.int32), + weights=weights, + roots=_table(roots, np.int32, 2), + ) + _verify_acyclic(context, ensemble) + return ensemble + + +def _opset5_nodes(context: NodeContext) -> _Opset5Nodes: + splits = _required_tensor(context, "nodes_splits") + tests = [ + _numbered_test(context, int(value)) + for value in _required_tensor(context, "nodes_modes") + ] + count = len(tests) + if len(splits) != count: + raise CompileError( + f"Node `{context.label}`: `nodes_splits` holds {len(splits)} entry(s) where " + f"{count} are described by the attributes beside it." + ) + return _Opset5Nodes( + tests=tests, + features=_integers(context, "nodes_featureids", count), + true_ids=_integers(context, "nodes_truenodeids", count), + false_ids=_integers(context, "nodes_falsenodeids", count), + true_leafs=_integers(context, "nodes_trueleafs", count), + false_leafs=_integers(context, "nodes_falseleafs", count), + missing=_integers( + context, "nodes_missing_value_tracks_true", count, optional=True + ), + splits=splits, + ) + + +def _roots( + context: NodeContext, parsed: _Opset5Nodes, leaves: int +) -> list[tuple[int, int]]: + """Each tree's root, and whether it addresses the leaf families rather than the nodes.""" + resolved = [] + for root in _integers(context, "tree_roots", None): + if not 0 <= root < len(parsed): + raise CompileError( + f"Node `{context.label}`: `tree_roots` names node {root}, which is outside " + f"the {len(parsed)} node(s) this ensemble defines." + ) + leaf = parsed.is_bare_leaf(root) + if leaf and root >= leaves: + raise CompileError( + f"Node `{context.label}`: the tree rooted at node {root} is a single leaf, " + f"which ONNX reads at that same position among the leaves — and this " + f"ensemble defines {leaves} of them." + ) + resolved.append((root, int(leaf))) + return resolved + + +def _child( + context: NodeContext, child: int, leaf: int, parsed: _Opset5Nodes, leaves: int +) -> int: + limit = leaves if leaf else len(parsed) + if not 0 <= child < limit: + raise CompileError( + f"Node `{context.label}`: TreeEnsemble names {'leaf' if leaf else 'node'} " + f"{child} as a child, which is outside the {limit} it defines." + ) + return child + + +def _membership(context: NodeContext, tests: Sequence[int]) -> list[float]: + """`membership_values`, checked against the set tests it is supposed to describe.""" + tensor = context.attribute("membership_values", None) + sets = sum(1 for test in tests if test == _MEMBER_TEST) + if tensor is None: + if sets: + raise CompileError( + f"Node `{context.label}`: TreeEnsemble has {sets} set test(s) and no " + "`membership_values` saying what they test against." + ) + return [] + values = [float(value) for value in to_array(tensor).reshape(-1)] + terminators = sum(1 for value in values if math.isnan(value)) + if terminators != sets: + raise CompileError( + f"Node `{context.label}`: `membership_values` holds {terminators} " + f"NaN-terminated set(s) for {sets} set test(s)." + ) + return values + + +def _membership_ranges( + context: NodeContext, + parsed: _Opset5Nodes, + roots: Sequence[tuple[int, int]], + members: Sequence[float], + leaves: int, +) -> dict[int, tuple[int, int]]: + """Which slice of `membership_values` each set test reads. + + The sets are laid out in the order the reference implementation builds the trees in — + depth first from each root, true branch before false — rather than in node order, so that + is the traversal here. A set ends at the first NaN *or zero*, which is where the + reference's own loop condition stops. + + This is the traversal `_verify_acyclic` makes over the normalized form, and it runs first, + so it carries the same guards: without them a cycle would be an endless walk here, and a + child naming a node the ensemble does not define an `IndexError`, rather than the compile + errors they are once the nodes below are built. + """ + ranges: dict[int, tuple[int, int]] = {} + consumed = 0 + visited: set[int] = set() + for root, leaf in roots: + stack = [] if leaf else [root] + while stack: + index = stack.pop() + if index in visited: + raise _revisited(context, index) + visited.add(index) + if parsed.tests[index] == _MEMBER_TEST: + start = consumed + consumed = _end_of_set(context, members, consumed, index) + ranges[index] = (start, consumed - start - 1) + if not parsed.false_leafs[index]: + stack.append( + _child(context, parsed.false_ids[index], 0, parsed, leaves) + ) + if not parsed.true_leafs[index]: + stack.append(_child(context, parsed.true_ids[index], 0, parsed, leaves)) + return ranges + + +def _end_of_set( + context: NodeContext, members: Sequence[float], start: int, index: int +) -> int: + """One past the terminator of the set beginning at `start`.""" + position = start + while True: + if position >= len(members): + raise CompileError( + f"Node `{context.label}`: `membership_values` runs out before the set node " + f"{index} tests against is terminated." + ) + value = members[position] + position += 1 + if value == 0.0 or math.isnan(value): + return position + + +# -------------------------------------------------------------------------------------- +# Shared reading and validation +# -------------------------------------------------------------------------------------- + + +def _verify_acyclic(context: NodeContext, ensemble: _Ensemble) -> None: + """Refuse a forest whose nodes are not a tree, which the emitted walker would loop on. + + A node two parents reach is refused along with one that reaches itself: the walk would be + the same, but the membership ranges an opset-5 node carries would not, since the reference + reads a set of its own every time it builds that node. + """ + nodes = ensemble.nodes.reshape(-1, _NODE_FIELDS).tolist() + visited: set[int] = set() + for root, leaf in ensemble.roots.reshape(-1, 2).tolist(): + stack = [] if leaf else [root] + while stack: + index = stack.pop() + if index in visited: + raise _revisited(context, index) + visited.add(index) + row = nodes[index] + for child, flag in ((row[3], _TRUE_IS_LEAF), (row[4], _FALSE_IS_LEAF)): + if not row[2] & flag: + stack.append(child) + + +def _revisited(context: NodeContext, index: int) -> CompileError: + return CompileError( + f"Node `{context.label}`: node {index} is reachable more than once; the C compiler " + "serves ensembles whose nodes form a tree." + ) + + +def _feature(context: NodeContext, feature: int, width: int) -> int: + if not 0 <= feature < width: + raise CompileError( + f"Node `{context.label}`: `nodes_featureids` names feature {feature}, which is " + f"outside the {width} column(s) of its input." + ) + return feature + + +def _class_labels(context: NodeContext) -> tuple[int, ...]: + """The class values a classifier labels its rows with, of the two families it may set.""" + integers = tuple( + int(value) for value in context.attribute("classlabels_int64s", []) + ) + strings = list(context.attribute("classlabels_strings", [])) + if bool(integers) == bool(strings): + raise CompileError( + f"Node `{context.label}`: TreeEnsembleClassifier must set exactly one of " + "`classlabels_int64s` and `classlabels_strings`." + ) + if strings: + raise CompileError( + f"Node `{context.label}`: TreeEnsembleClassifier labels its rows with the " + "strings in `classlabels_strings`, and a tensor of STRING at run time is not " + "something the C compiler supports." + ) + if integers == (1,): + return integers + if len(integers) == 1: + raise CompileError( + f"Node `{context.label}`: TreeEnsembleClassifier declares the single class " + f"{integers[0]}, which ONNX's own reference implementation refuses for any value " + "but 1; nothing says what such a row should be labelled." + ) + return integers + + +def _base_values( + context: NodeContext, targets: int, elem_type: int +) -> np.ndarray | None: + """`base_values` stretched over the targets it applies to, or None where none are set. + + A single value covers every target, which is how the reference broadcasts the attribute + over the score matrix. + """ + values = [float(value) for value in context.attribute("base_values", [])] + if not values: + return None + if len(values) not in (1, targets): + raise CompileError( + f"Node `{context.label}`: `base_values` holds {len(values)} value(s) for " + f"{targets} target(s); it takes either one value per target or a single value " + "for all of them." + ) + dtype = numpy_dtype_name(elem_type) + return np.broadcast_to(np.array(values, dtype), (targets,)).astype(dtype) + + +def _seed(mode: int, targets: int, elem_type: int) -> np.ndarray: + """What each target's score starts at: zero, or the extreme a fold runs down from.""" + if mode == _ACCUMULATE: + return _filled(targets, 0.0, elem_type) + info = np.finfo(numpy_dtype_name(elem_type)) + return _filled( + targets, float(info.max if mode == _MINIMUM else info.min), elem_type + ) + + +def _filled(count: int, value: float, elem_type: int) -> np.ndarray: + return np.full(count, value, numpy_dtype_name(elem_type)) + + +def _refuse_tensor_tables(context: NodeContext, families: Sequence[str]) -> None: + """Refuse the `*_as_tensor` families, which ONNX's reference implementation never reads. + + Opset 3 added them so that a double-precision ensemble need not round its tables to + float32. The reference evaluator stores them and goes on reading the float32 families, so + a kernel built from them would be answerable to nothing; the model is refused instead. + """ + named = [ + family + for family in families + if context.attribute(f"{family}_as_tensor", None) is not None + ] + if named: + raise CompileError( + f"Node `{context.label}`: `{named[0]}_as_tensor` is not supported by the C " + "compiler — ONNX's own reference implementation ignores the `*_as_tensor` " + f"families and reads `{named[0]}` instead, so nothing can vouch for a kernel " + "built from them; re-export the model with the float32 tables." + ) + + +def _rows_and_width(context: NodeContext, source: TensorRef) -> tuple[int, int]: + """How many rows an ensemble reads from `X`, and how many features each of them holds.""" + if len(source.shape) == 2: + return source.shape[0], source.shape[1] + if len(source.shape) == 1: + return 1, source.shape[0] + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads an `[N, F]` matrix of " + f"features, and `{source.name}` has shape {list(source.shape)}." + ) + + +def _branch_test(context: NodeContext, value: object) -> str: + test = value.decode("utf-8") if isinstance(value, bytes) else str(value) + if test != _LEAF and test not in _BRANCH_TESTS: + raise CompileError( + f"Node `{context.label}`: `nodes_modes` holds `{test}`, which is none of the " + f"branch tests ONNX defines ({', '.join(sorted(_BRANCH_TESTS))}, {_LEAF})." + ) + return test + + +def _numbered_test(context: NodeContext, test: int) -> int: + if not 0 <= test <= _MEMBER_TEST: + raise CompileError( + f"Node `{context.label}`: `nodes_modes` holds {test}, which is none of the " + f"branch tests ONNX numbers 0 to {_MEMBER_TEST}." + ) + return test + + +def _integers( + context: NodeContext, name: str, count: int | None, *, optional: bool = False +) -> list[int]: + values = [int(value) for value in context.attribute(name, [])] + if optional and not values: + return [0] * (count or 0) + if count is not None and len(values) != count: + raise CompileError( + f"Node `{context.label}`: `{name}` holds {len(values)} entry(s) where {count} " + "are described by the attributes beside it." + ) + return values + + +def _floats( + context: NodeContext, name: str, count: int, *, optional: bool = False +) -> list[float]: + values = [float(value) for value in context.attribute(name, [])] + if optional and not values: + return [0.0] * count + if len(values) != count: + raise CompileError( + f"Node `{context.label}`: `{name}` holds {len(values)} entry(s) where {count} " + "are described by the attributes beside it." + ) + return values + + +def _required_tensor(context: NodeContext, name: str) -> np.ndarray: + tensor = context.attribute(name, None) + if tensor is None: + raise CompileError( + f"Node `{context.label}`: TreeEnsemble requires the `{name}` attribute." + ) + return np.ascontiguousarray(to_array(tensor).reshape(-1)) + + +def _table(rows: Sequence[Sequence[int]], dtype: type, columns: int) -> np.ndarray: + return np.array(list(rows), dtype).reshape(len(rows), columns) + + +register_kernel( + ML_DOMAIN, "TreeEnsembleRegressor", _LEGACY_VERSIONS, _tree_ensemble_regressor +) +register_kernel( + ML_DOMAIN, "TreeEnsembleClassifier", _LEGACY_VERSIONS, _tree_ensemble_classifier +) +register_kernel(ML_DOMAIN, "TreeEnsemble", _ENSEMBLE_VERSIONS, _tree_ensemble) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/views.py b/src/python/fnnx/extras/compilers/c/onnx/ops/views.py new file mode 100644 index 0000000..bdea9c4 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/views.py @@ -0,0 +1,637 @@ +"""The views: the ops that rearrange a tensor's elements without computing any. + +Reshape, Flatten, Squeeze and Unsqueeze only relabel the axes of a row-major buffer, so each +of them is a copy of it. Transpose, Concat, Split, Slice, Expand, Tile and the two block +shuffles — DepthToSpace and SpaceToDepth — do reorder the elements, and all of them the same +way: every element of the result is read from one element of the operand, at an offset that +is a fixed base plus a stride per axis times that axis's coordinate. One shared kernel walks +that addressing, and the ops differ only in the strides and bases they hand it — compile-time +literals, all of them. Where both sides come out contiguous the kernel is skipped and the +move is a single `memcpy`. + +`Identity` is emitted alongside the elementwise family, and `Shape` and `Size` need no kernel +at all: their output follows from a shape the compiler has already made static, so constant +folding resolves them long before dispatch. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Sequence +from functools import partial +from string import Template + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import c_type +from fnnx.extras.compilers.c.onnx.kernels import ( + CFunction, + NodeContext, + NodeEmission, + TensorRef, + broadcast_strides, + copy_tensor, + register_kernel, +) +from fnnx.extras.compilers.c.onnx.ops.axes import ( + call_kernel, + normalize_axis, + row_major_strides, + verify_shape, +) +from fnnx.extras.compilers.c.onnx.ops.broadcast import extents + +# None of the ops below read the operand that decides their result's shape — the shape ONNX +# inferred is what every buffer and every stride is derived from — so a revision that only +# moved that operand between an attribute and an input, or widened the types it accepts, +# emits the identical code and is served by the same generator. +# +# Reshape moved `shape` to an input at 5 and gained `allowzero` at 14; Squeeze and Unsqueeze +# moved `axes` to an input at 13; Flatten, Transpose and Concat kept their attributes +# throughout. Concat-1 is the one revision left out: it defaulted `axis` to 1 while its +# schema reports no default at all, so the value a node omitting the attribute means cannot +# be read off ONNX itself. From 4 on the attribute is required. +_RESHAPE_VERSIONS = (1, 5, 13, 14, 19, 21, 23, 24, 25) +_FLATTEN_VERSIONS = (1, 9, 11, 13, 21, 23, 24, 25) +_SQUEEZE_VERSIONS = (1, 11, 13, 21, 23, 24, 25) +_UNSQUEEZE_VERSIONS = (1, 11, 13, 21, 23, 24, 25) +_TRANSPOSE_VERSIONS = (1, 13, 21, 23, 24, 25) +_CONCAT_VERSIONS = (4, 11, 13) +_SPLIT_VERSIONS = (1, 2, 11, 13, 18) +_EXPAND_VERSIONS = (8, 13) + +# Slice-1 takes its bounds as attributes and has no `steps`; from 10 on they are all +# operands. That is two ways of reading the same slice, so it is two generators. +_SLICE_ATTRIBUTE_VERSIONS = (1,) +_SLICE_OPERAND_VERSIONS = (10, 11, 13) + +# Tile-1 is deliberately absent: its second and third operands are a repeat count and the +# single axis to apply it to, not the per-axis repeat vector every revision since 6 takes. +_TILE_VERSIONS = (6, 13) + +# DepthToSpace gained its `mode` at 11 and SpaceToDepth never changed; 13 widened the types +# of both. Only 13 is claimed for either: it is the revision the reference evaluator is +# version-faithful for and the one both corpus tests import, so it is the only one anything +# can vouch for. A model importing an older one gets the unsupported-version error. +_BLOCK_VERSIONS = (13,) + +_COPY_TEMPLATE = Template("""\ +static void $name( + $element* out, + const $element* in, + size_t count, + int rank, + const size_t* shape, + const ptrdiff_t* out_strides, + const ptrdiff_t* in_strides, + ptrdiff_t out_base, + ptrdiff_t in_base) +{ + size_t index; + for (index = 0; index < count; ++index) { + size_t remainder = index; + ptrdiff_t source = in_base; + ptrdiff_t target = out_base; + int axis; + for (axis = rank - 1; axis >= 0; --axis) { + const size_t coordinate = remainder % shape[axis]; + remainder /= shape[axis]; + source += (ptrdiff_t)coordinate * in_strides[axis]; + target += (ptrdiff_t)coordinate * out_strides[axis]; + } + out[target] = in[source]; + } +}""") + + +def copy_elements( + context: NodeContext, + *, + source: TensorRef, + result: TensorRef, + shape: Sequence[int], + source_strides: Sequence[int], + result_strides: Sequence[int] | None = None, + source_base: int = 0, + result_base: int = 0, +) -> NodeEmission: + """Move a block of `shape` elements, each side addressed by its own strides and base. + + `result_strides` defaults to the row-major strides of `shape`, which is what an op + writing the whole of its result in order needs; Concat, which writes a slice of one, + passes the result's own strides instead. A move that comes out contiguous on both sides + is emitted as a `memcpy`, since the kernel would then walk it element by element to no + end. + """ + count = math.prod(shape) + if count == 0: + return NodeEmission(functions=(), statements=()) + strides = row_major_strides(shape) if result_strides is None else result_strides + if _is_contiguous(shape, source_strides) and _is_contiguous(shape, strides): + return NodeEmission( + functions=(), + statements=( + f"memcpy({_at(result.expr, result_base)}, " + f"{_at(source.expr, source_base)}, " + f"{count}u * sizeof(*{result.expr}));", + ), + ) + element = c_type(result.elem_type) + name = f"{context.prefix}_copy_{element}" + return NodeEmission( + functions=( + CFunction(name, _COPY_TEMPLATE.substitute(name=name, element=element)), + ), + statements=( + call_kernel( + name, + [ + result.expr, + source.expr, + f"{count}u", + str(len(shape)), + extents(shape), + _offsets(strides), + _offsets(source_strides), + str(result_base), + str(source_base), + ], + ), + ), + ) + + +def _is_contiguous(shape: Sequence[int], strides: Sequence[int]) -> bool: + """Whether walking `shape` under `strides` visits one unbroken run, in order. + + An axis of a single element is skipped: its stride multiplies a coordinate that is only + ever zero, so whatever it holds cannot break the run. + """ + expected = 1 + for extent, stride in zip(reversed(shape), reversed(strides)): + if extent != 1 and stride != expected: + return False + expected *= extent + return True + + +def _at(expr: str, offset: int) -> str: + return expr if offset == 0 else f"{expr} + {offset}" + + +def _offsets(values: Sequence[int]) -> str: + """Strides as a compound literal; they are signed, since a slice may walk backwards.""" + literals = ", ".join(str(value) for value in values) or "0" + return f"(const ptrdiff_t[]){{{literals}}}" + + +def _combined(emissions: Sequence[NodeEmission]) -> NodeEmission: + """One emission out of the several an op writing block by block contributes.""" + functions = { + function.name: function + for emission in emissions + for function in emission.functions + } + return NodeEmission( + functions=tuple(functions.values()), + statements=tuple( + statement for emission in emissions for statement in emission.statements + ), + ) + + +# -------------------------------------------------------------------------------------- +# The ops +# -------------------------------------------------------------------------------------- + + +def _relabel(context: NodeContext) -> NodeEmission: + """Reshape, Flatten, Squeeze, Unsqueeze: the same row-major buffer under other axes.""" + source = context.require_input(0) + result = context.require_output(0) + if source.elem_count != result.elem_count: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` reads " + f"{source.elem_count} element(s) from `{source.name}` but writes " + f"{result.elem_count} to `{result.name}`; this op only relabels axes." + ) + return copy_tensor(source, result) + + +def _transpose(context: NodeContext) -> NodeEmission: + """Transpose: the operand read along permuted axes, written out in order.""" + source = context.require_input(0) + result = context.require_output(0) + rank = len(source.shape) + declared = context.attribute("perm", None) + perm = ( + tuple(reversed(range(rank))) + if declared is None + else tuple(int(axis) for axis in declared) + ) + if sorted(perm) != list(range(rank)): + raise CompileError( + f"Node `{context.label}`: `perm` {list(perm)} is not a permutation of the " + f"{rank} axes of `{source.name}`." + ) + verify_shape(context, result, [source.shape[axis] for axis in perm]) + strides = row_major_strides(source.shape) + return copy_elements( + context, + source=source, + result=result, + shape=result.shape, + source_strides=[strides[axis] for axis in perm], + ) + + +def _concat(context: NodeContext) -> NodeEmission: + """Concat: each operand written into its own band of the result along one axis.""" + result = context.require_output(0) + operands = [ + context.require_input(index) for index in range(len(context.node.input)) + ] + axis = normalize_axis(context, context.int_attribute("axis"), len(result.shape)) + _verify_bands(context, operands, result, axis) + strides = row_major_strides(result.shape) + emissions = [] + offset = 0 + for operand in operands: + emissions.append( + copy_elements( + context, + source=operand, + result=result, + shape=operand.shape, + source_strides=row_major_strides(operand.shape), + result_strides=strides, + result_base=offset * strides[axis], + ) + ) + offset += operand.shape[axis] + return _combined(emissions) + + +def _split(context: NodeContext) -> NodeEmission: + """Split: consecutive bands of the operand along one axis, each its own result. + + How wide each band is comes from the shape ONNX inferred for it — which is what the + `split` operand, the `num_outputs` attribute and the equal division of neither all + ultimately say — so every revision of the op is read the same way here. + """ + source = context.require_input(0) + results = [ + context.require_output(index) for index in range(len(context.node.output)) + ] + axis = normalize_axis(context, context.int_attribute("axis"), len(source.shape)) + _verify_bands(context, results, source, axis) + strides = row_major_strides(source.shape) + emissions = [] + offset = 0 + for result in results: + emissions.append( + copy_elements( + context, + source=source, + result=result, + shape=result.shape, + source_strides=strides, + source_base=offset * strides[axis], + ) + ) + offset += result.shape[axis] + return _combined(emissions) + + +def _verify_bands( + context: NodeContext, + bands: Sequence[TensorRef], + joined: TensorRef, + axis: int, +) -> None: + """Refuse to emit a split or a join whose bands do not tile the whole tensor. + + Every band has to match the joined tensor on every axis but `axis`, and their extents + along `axis` have to add up to its own; anything else is a compiler bug that would read + or write outside a buffer. + """ + for band in bands: + expected = list(joined.shape) + if len(band.shape) == len(expected): + expected[axis] = band.shape[axis] + if band.shape != tuple(expected): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` joins `{band.name}` of " + f"shape {list(band.shape)} along axis {axis} of a tensor of shape " + f"{list(joined.shape)}; the two have to agree on every other axis." + ) + total = sum(band.shape[axis] for band in bands) + if total != joined.shape[axis]: + raise CompileError( + f"Node `{context.label}`: the bands of `{context.node.op_type}` measure " + f"{total} along axis {axis}, but `{joined.name}` measures " + f"{joined.shape[axis]}." + ) + + +# What a Slice revision reads its bounds from: the starts, ends, axes and steps, where None +# stands for a list the node leaves out. +SliceBounds = tuple[ + tuple[int, ...], tuple[int, ...], tuple[int, ...] | None, tuple[int, ...] | None +] + + +def _attribute_bounds(context: NodeContext) -> SliceBounds: + """Slice-1, whose bounds are attributes and which has no steps at all.""" + return ( + _required_attribute(context, "starts"), + _required_attribute(context, "ends"), + _optional_attribute(context, "axes"), + None, + ) + + +def _operand_bounds(context: NodeContext) -> SliceBounds: + """Slice-10 and later, whose bounds are operands the graph has to fix.""" + return ( + _constant_operand(context, 1), + _constant_operand(context, 2), + _optional_operand(context, 3), + _optional_operand(context, 4), + ) + + +def _required_attribute(context: NodeContext, name: str) -> tuple[int, ...]: + values = context.attribute(name, None) + if values is None: + raise CompileError( + f"Node `{context.label}`: `Slice` requires the `{name}` attribute at opset " + f"version {context.since_version}." + ) + return tuple(int(value) for value in values) + + +def _optional_attribute(context: NodeContext, name: str) -> tuple[int, ...] | None: + values = context.attribute(name, None) + return None if values is None else tuple(int(value) for value in values) + + +def _constant_operand(context: NodeContext, index: int) -> tuple[int, ...]: + operand = context.require_input(index) + values = context.constant_input(index) + if values is None: + raise CompileError( + f"Node `{context.label}`: `Slice` takes its bounds from `{operand.name}`, " + "which is not known at compile time; the shape of the result then depends on " + "input data, which the C compiler cannot compile." + ) + return tuple(int(value) for value in values.reshape(-1)) + + +def _optional_operand(context: NodeContext, index: int) -> tuple[int, ...] | None: + operand = context.optional_input(index) + return None if operand is None else _constant_operand(context, index) + + +def _slice( + context: NodeContext, *, bounds: Callable[[NodeContext], SliceBounds] +) -> NodeEmission: + """Slice: the operand walked from a per-axis start, by a per-axis step. + + A step is a stride multiplier and a start an offset into the operand, so the whole of the + op is addressing. The bounds are clamped the way Python's own slice does it, which is + what numpy — and through it the ONNX reference evaluator — applies; the extents that come + out are checked against the shape ONNX inferred rather than trusted. + """ + source = context.require_input(0) + result = context.require_output(0) + rank = len(source.shape) + starts, ends, axes, steps = bounds(context) + if axes is None: + axes = tuple(range(len(starts))) + if steps is None: + steps = (1,) * len(starts) + if not len(starts) == len(ends) == len(axes) == len(steps): + raise CompileError( + f"Node `{context.label}`: `Slice` was given {len(starts)} start(s), " + f"{len(ends)} end(s), {len(axes)} axis/axes and {len(steps)} step(s); ONNX " + "defines one of each per sliced axis." + ) + + strides = row_major_strides(source.shape) + walk = list(strides) + sliced = list(source.shape) + base = 0 + seen: set[int] = set() + for start, end, step, axis in zip(starts, ends, steps, axes): + resolved = normalize_axis(context, axis, rank) + if resolved in seen: + raise CompileError( + f"Node `{context.label}`: `Slice` names axis {resolved} of " + f"`{source.name}` more than once." + ) + seen.add(resolved) + if step == 0: + raise CompileError( + f"Node `{context.label}`: `Slice` steps by 0 along axis {resolved}, which " + "ONNX does not define." + ) + first, stop, stride = slice(start, end, step).indices(source.shape[resolved]) + sliced[resolved] = len(range(first, stop, stride)) + walk[resolved] = strides[resolved] * stride + base += first * strides[resolved] + + verify_shape(context, result, sliced) + return copy_elements( + context, + source=source, + result=result, + shape=result.shape, + source_strides=walk, + source_base=base, + ) + + +def _expand(context: NodeContext) -> NodeEmission: + """Expand: the operand stretched onto the shape it broadcasts with. + + A stretched axis is a stride of zero, so every coordinate along it reads the same + element; which axes those are follows from the shape ONNX inferred for the result, the + same shape the `shape` operand had to be constant to produce. + """ + source = context.require_input(0) + result = context.require_output(0) + return copy_elements( + context, + source=source, + result=result, + shape=result.shape, + source_strides=broadcast_strides( + source, result.shape, node_label=context.label + ), + ) + + +def _tile(context: NodeContext) -> NodeEmission: + """Tile: the operand repeated a given number of times along each of its axes. + + A tiling is a broadcast in disguise. Splitting every result axis into its repeat count + and the operand's own extent gives a tensor of rank 2n whose row-major order is exactly + the result's, and over which the operand is simply stretched along the repeat axes — so + the same strided move serves, with a stride of zero on each of them. + """ + source = context.require_input(0) + result = context.require_output(0) + repeats = context.constant_input(1) + if repeats is None: + raise CompileError( + f"Node `{context.label}`: `Tile` takes its repeat counts from " + f"`{context.require_input(1).name}`, which is not known at compile time; the " + "shape of the result then depends on input data, which the C compiler cannot " + "compile." + ) + counts = tuple(int(count) for count in repeats.reshape(-1)) + if len(counts) != len(source.shape): + raise CompileError( + f"Node `{context.label}`: `Tile` was given {len(counts)} repeat count(s) for " + f"the {len(source.shape)} axes of `{source.name}`; ONNX defines one per axis." + ) + verify_shape( + context, + result, + [extent * count for extent, count in zip(source.shape, counts)], + ) + + strides = row_major_strides(source.shape) + interleaved: list[int] = [] + source_strides: list[int] = [] + for axis, count in enumerate(counts): + interleaved += [count, source.shape[axis]] + source_strides += [0, strides[axis]] + return copy_elements( + context, + source=source, + result=result, + shape=interleaved, + source_strides=source_strides, + ) + + +def _block_shuffle( + context: NodeContext, view: Sequence[int], perm: Sequence[int] +) -> NodeEmission: + """The operand split into blocks by `view`, then transposed by `perm`. + + Both `DepthToSpace` and `SpaceToDepth` are defined as exactly that — a reshape, a + transpose and a reshape back — and the two reshapes are free: `view`'s row-major order is + the operand's own, and the transposed extents' row-major order is the result's. So only + the transpose is emitted, through the same strided move every view op runs. + """ + return copy_elements( + context, + source=context.require_input(0), + result=context.require_output(0), + shape=[view[axis] for axis in perm], + source_strides=[row_major_strides(view)[axis] for axis in perm], + ) + + +def _depth_to_space(context: NodeContext) -> NodeEmission: + """DepthToSpace: each channel of a block spread over one position of a spatial block. + + `mode` says how the channels are grouped before they are spread: `DCR` reads the block's + rows and columns as the outermost channel axes and the surviving depth as the innermost, + `CRD` the other way round. + """ + batch, channels, rows, columns = _image(context) + block = _blocksize(context) + if channels % (block * block) != 0: + raise CompileError( + f"Node `{context.label}`: `DepthToSpace` spreads {channels} channel(s) over " + f"{block}x{block} positions, which does not divide them evenly." + ) + depth = channels // (block * block) + verify_shape( + context, + context.require_output(0), + (batch, depth, rows * block, columns * block), + ) + mode = context.attribute("mode", b"DCR") + mode = mode.decode() if isinstance(mode, bytes) else str(mode) + if mode not in ("DCR", "CRD"): + raise CompileError( + f"Node `{context.label}`: `DepthToSpace` asks for `mode` `{mode}`, which is " + "not one of the modes ONNX defines (`DCR`, `CRD`)." + ) + if mode == "DCR": + return _block_shuffle( + context, (batch, block, block, depth, rows, columns), (0, 3, 4, 1, 5, 2) + ) + return _block_shuffle( + context, (batch, depth, block, block, rows, columns), (0, 1, 4, 2, 5, 3) + ) + + +def _space_to_depth(context: NodeContext) -> NodeEmission: + """SpaceToDepth: each position of a spatial block moved into a channel of its own.""" + batch, channels, rows, columns = _image(context) + block = _blocksize(context) + if rows % block != 0 or columns % block != 0: + raise CompileError( + f"Node `{context.label}`: `SpaceToDepth` splits a {rows}x{columns} image into " + f"{block}x{block} blocks, which do not tile it evenly." + ) + verify_shape( + context, + context.require_output(0), + (batch, channels * block * block, rows // block, columns // block), + ) + return _block_shuffle( + context, + (batch, channels, rows // block, block, columns // block, block), + (0, 3, 5, 1, 2, 4), + ) + + +def _image(context: NodeContext) -> tuple[int, int, int, int]: + source = context.require_input(0) + if len(source.shape) != 4: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` takes a batch of images — a " + f"tensor of rank 4 — but `{source.name}` has shape {list(source.shape)}." + ) + batch, channels, rows, columns = source.shape + return batch, channels, rows, columns + + +def _blocksize(context: NodeContext) -> int: + block = context.attribute("blocksize", None) + if block is None: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` states no `blocksize`, " + "which ONNX defines as a required attribute." + ) + if int(block) < 1: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` was given `blocksize` " + f"{int(block)}; ONNX defines it as the extent of a block, which is positive." + ) + return int(block) + + +register_kernel("", "Reshape", _RESHAPE_VERSIONS, _relabel) +register_kernel("", "Flatten", _FLATTEN_VERSIONS, _relabel) +register_kernel("", "Squeeze", _SQUEEZE_VERSIONS, _relabel) +register_kernel("", "Unsqueeze", _UNSQUEEZE_VERSIONS, _relabel) +register_kernel("", "Transpose", _TRANSPOSE_VERSIONS, _transpose) +register_kernel("", "Concat", _CONCAT_VERSIONS, _concat) +register_kernel("", "Split", _SPLIT_VERSIONS, _split) +register_kernel( + "", "Slice", _SLICE_ATTRIBUTE_VERSIONS, partial(_slice, bounds=_attribute_bounds) +) +register_kernel( + "", "Slice", _SLICE_OPERAND_VERSIONS, partial(_slice, bounds=_operand_bounds) +) +register_kernel("", "Expand", _EXPAND_VERSIONS, _expand) +register_kernel("", "Tile", _TILE_VERSIONS, _tile) +register_kernel("", "DepthToSpace", _BLOCK_VERSIONS, _depth_to_space) +register_kernel("", "SpaceToDepth", _BLOCK_VERSIONS, _space_to_depth) diff --git a/src/python/fnnx/extras/compilers/c/onnx/ops/window.py b/src/python/fnnx/extras/compilers/c/onnx/ops/window.py new file mode 100644 index 0000000..3e75224 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/ops/window.py @@ -0,0 +1,135 @@ +"""The sliding window's geometry, shared by every op that slides one. + +A convolution and a pooling place the same window over the same axes: `strides` move it, +`dilations` spread its taps, and `pads` or `auto_pad` place it against the operand's edges. +So the attributes are read and resolved once here — into the extents, steps and pads a kernel +walks with — and each family layers on what only it has: a filter's shape and the backward +walk for the convolutions, a `kernel_shape` and a `ceil_mode` for the poolings. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.kernels import NodeContext + +AUTO_PAD_MODES = ("NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID") + + +def auto_pad_mode(context: NodeContext) -> str: + value = context.attribute("auto_pad", b"NOTSET") + mode = value.decode() if isinstance(value, bytes) else str(value) + if mode not in AUTO_PAD_MODES: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` `auto_pad` `{mode}` is not " + f"one of the modes ONNX defines ({', '.join(AUTO_PAD_MODES)})." + ) + return mode + + +def spatial_extents( + context: NodeContext, name: str, rank: int, *, minimum: int = 1 +) -> tuple[int, ...] | None: + """The node's `name` attribute as one value per spatial axis, or None when absent.""" + declared = context.attribute(name, None) + if declared is None: + return None + values = tuple(int(value) for value in declared) + if len(values) != rank: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` was given {len(values)} " + f"`{name}` for {rank} spatial axis/axes." + ) + if any(value < minimum for value in values): + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` was given `{name}` " + f"{list(values)}; ONNX defines them as " + f"{'positive' if minimum > 0 else 'nonnegative'}." + ) + return values + + +def spatial_attribute( + context: NodeContext, name: str, rank: int, default: int, *, minimum: int = 1 +) -> tuple[int, ...]: + values = spatial_extents(context, name, rank, minimum=minimum) + return (default,) * rank if values is None else values + + +def declared_pads( + context: NodeContext, rank: int, mode: str +) -> tuple[tuple[int, ...], tuple[int, ...]] | None: + """The `pads` attribute split into begins and ends, or None when the node omits it.""" + declared = context.attribute("pads", None) + if declared is None: + return None + if mode != "NOTSET": + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` states both `auto_pad` " + f"`{mode}` and explicit `pads`, which ONNX defines as mutually exclusive." + ) + values = tuple(int(value) for value in declared) + if len(values) != 2 * rank: + raise CompileError( + f"Node `{context.label}`: `{context.node.op_type}` was given {len(values)} " + f"pad(s) for {rank} spatial axis/axes; ONNX defines two — a begin and an end " + "— per axis." + ) + return values[:rank], values[rank:] + + +def resolve_pads( + context: NodeContext, + input_shape: Sequence[int], + window_shape: Sequence[int], + dilations: Sequence[int], + strides: Sequence[int], +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """The pad before and after each spatial axis, after resolving `auto_pad`.""" + mode = auto_pad_mode(context) + rank = len(input_shape) + declared = declared_pads(context, rank, mode) + if declared is not None: + return declared + if mode in ("NOTSET", "VALID"): + return (0,) * rank, (0,) * rank + + begins, ends = [], [] + for extent, window, dilation, stride in zip( + input_shape, window_shape, dilations, strides + ): + # ONNX pads so the result measures `ceil(extent / stride)`, which puts the last + # window's start one stride before the end of the axis — or `extent % stride` + # before it, where the stride does not divide the extent — and pads whatever of + # the window's dilated reach then hangs off the end. + residual = extent % stride + reach = (window - 1) * dilation + 1 + total = max(reach - (stride if residual == 0 else residual), 0) + smaller, larger = total // 2, total - total // 2 + begins.append(smaller if mode == "SAME_UPPER" else larger) + ends.append(larger if mode == "SAME_UPPER" else smaller) + return tuple(begins), tuple(ends) + + +def output_extents( + input_shape: Sequence[int], + window_shape: Sequence[int], + dilations: Sequence[int], + strides: Sequence[int], + begins: Sequence[int], + ends: Sequence[int], +) -> tuple[int, ...]: + """Result extents of a forward walk: the window positions that fit each axis.""" + return tuple( + (extent + begin + end - (window - 1) * dilation - 1) // stride + 1 + for extent, begin, end, window, dilation, stride in zip( + input_shape, begins, ends, window_shape, dilations, strides + ) + ) + + +def offsets(values: Sequence[int]) -> str: + """The per-axis pads as a compound literal; they are signed, since a pad may crop.""" + literals = ", ".join(str(value) for value in values) or "0" + return f"(const ptrdiff_t[]){{{literals}}}" diff --git a/src/python/fnnx/extras/compilers/c/onnx/registry.py b/src/python/fnnx/extras/compilers/c/onnx/registry.py new file mode 100644 index 0000000..7fd5a40 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/registry.py @@ -0,0 +1,130 @@ +"""Kernel registry and opset-driven dispatch.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Generic, TypeVar + +import onnx.defs + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.loader import display_domain, normalize_domain + +G = TypeVar("G") + + +@dataclass(frozen=True) +class KernelSpec(Generic[G]): + domain: str + op_type: str + since_version: int + generator: G + + +def latest_semantic_revision( + domain: str, op_type: str, opset_version: int +) -> int | None: + """`since_version` of the ONNX schema in effect for the op at `opset_version`. + + None when the installed `onnx` package defines no schema for the op at that version. + """ + try: + schema = onnx.defs.get_schema(op_type, opset_version, normalize_domain(domain)) + except onnx.defs.SchemaError: + return None + return schema.since_version + + +class KernelRegistry(Generic[G]): + """Maps `(domain, op_type)` to version-keyed generators, mirroring ONNX's op versioning.""" + + def __init__(self) -> None: + self._specs: dict[tuple[str, str], dict[int, KernelSpec[G]]] = {} + + def register( + self, domain: str, op_type: str, since_version: int, generator: G + ) -> None: + normalized = normalize_domain(domain) + if latest_semantic_revision(normalized, op_type, since_version) is None: + raise ValueError( + f"ONNX defines no schema for `{op_type}` " + f"(domain `{display_domain(normalized)}`) at opset version {since_version}." + ) + versions = self._specs.setdefault((normalized, op_type), {}) + if since_version in versions: + raise ValueError( + f"A kernel for `{op_type}` (domain `{display_domain(normalized)}`) is " + f"already registered at since_version {since_version}." + ) + versions[since_version] = KernelSpec( + domain=normalized, + op_type=op_type, + since_version=since_version, + generator=generator, + ) + + def registered_ops(self) -> list[tuple[str, str]]: + """Every `(domain, op_type)` a kernel is registered for, in a stable order.""" + return sorted(self._specs) + + def registered_versions(self, domain: str, op_type: str) -> list[int]: + return sorted(self._specs.get((normalize_domain(domain), op_type), {})) + + def select( + self, domain: str, op_type: str, opset_version: int + ) -> KernelSpec[G] | None: + """Highest-versioned kernel valid at `opset_version`, or None if none can be vouched for. + + The semantic-revision guard rejects an otherwise applicable kernel when ONNX revised + the op's spec after the kernel's `since_version` and at or below the requested + version: old semantics are never silently applied to a newer opset. + """ + normalized = normalize_domain(domain) + versions = self._specs.get((normalized, op_type), {}) + applicable = [version for version in versions if version <= opset_version] + if not applicable: + return None + spec = versions[max(applicable)] + revision = latest_semantic_revision(normalized, op_type, opset_version) + if revision is None or revision > spec.since_version: + return None + return spec + + def unsupported_op_error( + self, + domain: str, + op_type: str, + opset_version: int, + *, + node_name: str | None = None, + ) -> CompileError: + """Build — but do not raise — the error for an op no registered kernel can serve. + + Callers fall through to function expansion before raising it. + """ + normalized = normalize_domain(domain) + versions = self.registered_versions(normalized, op_type) + prefix = f"Node `{node_name}`: " if node_name else "" + message = ( + f"{prefix}op `{op_type}` (domain `{display_domain(normalized)}`) is not " + f"supported at opset version {opset_version}" + ) + if not versions: + return CompileError(f"{message}: no kernel is registered for this op.") + nearest = min( + versions, key=lambda version: (abs(version - opset_version), version) + ) + if all(version > opset_version for version in versions): + reason = "every registered kernel targets a newer opset version" + else: + revision = latest_semantic_revision(normalized, op_type, opset_version) + if revision is None: + reason = "ONNX defines no schema for this op at that version" + else: + reason = ( + f"ONNX revised this op at opset version {revision} and no " + "registered kernel covers that revision" + ) + return CompileError( + f"{message}: {reason}. Nearest supported version: {nearest}." + ) diff --git a/src/python/fnnx/extras/compilers/c/onnx/runtime_dims.py b/src/python/fnnx/extras/compilers/c/onnx/runtime_dims.py new file mode 100644 index 0000000..13647e5 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/runtime_dims.py @@ -0,0 +1,80 @@ +"""Dimensions the caller sizes per call, within a maximum fixed at compile time.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.emit import sanitize_identifier + + +@dataclass(frozen=True) +class RuntimeDim: + """A symbolic dimension compiled for the whole family of sizes `[1, maximum]`. + + `identifier` is the dimension's name sanitized to a C identifier; the entrypoint + parameter carrying the actual value and the macro publishing the maximum are both + derived from it, so that two dimensions can never name the same one. + """ + + name: str + maximum: int + identifier: str + + @property + def c_name(self) -> str: + return f"dim_{self.identifier}" + + def macro(self, prefix: str) -> str: + return f"{prefix.upper()}_DIM_{self.identifier.upper()}_MAX" + + +@dataclass(frozen=True) +class ShapeTerm: + """One axis of a tensor: a constant extent, or a multiple of a runtime dimension. + + `size` is the extent at the dimension's maximum — the capacity the artifact's buffers + and macros are sized for — while `extent` gives the one a particular call works at. + """ + + size: int + dim: str | None = None + coefficient: int = 0 + + def extent(self, values: Mapping[str, int]) -> int: + return self.size if self.dim is None else self.coefficient * values[self.dim] + + +def resolve_runtime_dims( + runtime_dims: Mapping[str, int] | None, dim_bindings: Mapping[str, int] | None +) -> tuple[RuntimeDim, ...]: + """Validate the requested runtime dimensions, in the order they were declared.""" + identifiers: dict[str, str] = {} + resolved = [] + for name, maximum in (runtime_dims or {}).items(): + if not isinstance(maximum, int) or isinstance(maximum, bool) or maximum < 1: + raise CompileError( + f"Runtime dimension `{name}` needs a maximum of at least 1, got " + f"{maximum!r}." + ) + if name in (dim_bindings or {}): + raise CompileError( + f"Dimension `{name}` is both bound to " + f"{(dim_bindings or {})[name]} and declared runtime; a dimension is " + "either fixed at compile time or sized per call." + ) + identifier = sanitize_identifier(name, fallback="") + if not identifier: + raise CompileError( + f"Runtime dimension `{name}` has no C identifier to derive its " + "entrypoint parameter from; rename it." + ) + if identifier in identifiers: + raise CompileError( + f"Runtime dimensions `{identifiers[identifier]}` and `{name}` both " + f"sanitize to the C identifier `{identifier}`; rename one of them." + ) + identifiers[identifier] = name + resolved.append(RuntimeDim(name, maximum, identifier)) + return tuple(resolved) diff --git a/src/python/fnnx/extras/compilers/c/onnx/shapes.py b/src/python/fnnx/extras/compilers/c/onnx/shapes.py new file mode 100644 index 0000000..4e11fb4 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/shapes.py @@ -0,0 +1,699 @@ +"""Symbolic-dimension binding, shape inference, and tensor type/shape lookup.""" + +from __future__ import annotations + +import math +from collections.abc import Container, Iterator, Mapping + +import onnx.defs +import onnx.shape_inference +from onnx import ( + AttributeProto, + GraphProto, + ModelProto, + NodeProto, + TensorProto, + TensorShapeProto, + TypeProto, + ValueInfoProto, + helper, +) + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.loader import ( + ML_DOMAIN, + STANDARD_DOMAIN, + normalize_domain, +) + +UNBOUND_DIM_DEFAULT = 1 + +# Standard-domain ops whose ONNX schema states the output carries the input's shape and +# whose shape ONNX's own inference nevertheless does not derive; see +# `_propagate_preserved_shapes`. +SHAPE_PRESERVING_OPS = frozenset({"GroupNormalization"}) + +# `ai.onnx.ml` ops the installed `onnx` package ships no type-and-shape inference function +# for at all, plus the ones it has a function for that stop short of a case; see +# `_infer_ml_types`. +ML_UNINFERRED_OPS = frozenset( + { + "ArrayFeatureExtractor", + "FeatureVectorizer", + "Imputer", + "LinearRegressor", + "Normalizer", + "SVMClassifier", + "SVMRegressor", + "Scaler", + "TreeEnsembleClassifier", + "TreeEnsembleRegressor", + } +) + +# The Reduce* family, whose `axes` operand ONNX defines as naming every axis when it is +# absent — the one reading of an empty operand `drop_empty_shape_operands` rests on. +REDUCTIONS = ( + "ReduceL1", + "ReduceL2", + "ReduceLogSum", + "ReduceLogSumExp", + "ReduceMax", + "ReduceMean", + "ReduceMin", + "ReduceProd", + "ReduceSum", + "ReduceSumSquare", +) + +# Standard-domain ops with operands whose *values* decide the shape of their output, by the +# positions of those operands. ONNX passes them as tensors rather than as attributes, so a +# graph may compute one at run time; folding resolves every one a model fixes, and what is +# left makes the output shape a function of input data, which no binding can make static. +SHAPE_DEFINING_INPUTS: Mapping[str, tuple[int, ...]] = { + # `size` is the result's shape, spatial axes and all, less its trailing coordinate axis. + "AffineGrid": (1,), + # A window is `size` samples long and nothing else; a window op is a compile-time value + # rather than a kernel, so a `size` the graph does not fix leaves nothing to compile. + "BlackmanWindow": (0,), + # `shape` is the extent of every axis `axes` names, cropped or padded to reach it. + "CenterCropPad": (1,), + # `image_shape` is the result's spatial extent and `block_shape` decides how many + # channels the columns hold. + "Col2Im": (1, 2), + "ConstantOfShape": (0,), + # `dft_length` is the extent of the transformed axis, and `axis` says which axis takes + # it — the second only where the transform resizes an axis; see `_shape_defining_inputs`. + "DFT": (1, 2), + "Expand": (1,), + "HammingWindow": (0,), + "HannWindow": (0,), + "MaxUnpool": (2,), + # `num_mel_bins` and `dft_length` are the two extents of the matrix. + "MelWeightMatrix": (0, 1), + "OneHot": (1,), + # `pads` and `axes`: the two together say how much longer each axis of the result is. + "Pad": (1, 3), + "Range": (0, 1, 2), + "Reshape": (1,), + "Slice": (1, 2, 3, 4), + "Split": (1,), + "Squeeze": (1,), + # `frame_step` decides how many frames the signal yields and `frame_length` how long + # each transform is; the window is read for its values alone, so it may stay run-time. + "STFT": (1, 3), + "Tile": (1,), + "TopK": (1,), + "Unsqueeze": (1,), + **{op_type: (1,) for op_type in REDUCTIONS}, +} + + +def runtime_shape_operand( + node: NodeProto, + constants: Container[str], + types: Mapping[str, TypeProto], +) -> str | None: + """The first shape-deciding operand this node reads at run time, if it has any. + + An operand with no elements does not count: it names nothing whatever the data holds, so + a reduction over an empty axes tensor is as static as one with no axes operand at all. + """ + if normalize_domain(node.domain) != STANDARD_DOMAIN: + return None + indices = _shape_defining_inputs(node) + if indices is None: + return None + for index in indices: + name = node.input[index] if index < len(node.input) else "" + if not name or name in constants: + continue + shape = static_shape(types.get(name)) + if shape is None or math.prod(shape) != 0: + return name + return None + + +def _shape_defining_inputs(node: NodeProto) -> tuple[int, ...] | None: + """Which of the op's shape-deciding operands decide *this* node's result shape. + + DFT is the one op whose own attributes settle that: a one-sided transform returns half + the length of the axis it lands on, and a stated `dft_length` replaces that extent + outright, so with either of them which axis is transformed is part of the result's + shape. With neither, the result carries the operand's own extents whichever axis the + transform takes, and the axis is a value a kernel can switch on at run time. + """ + indices = SHAPE_DEFINING_INPUTS.get(node.op_type) + if indices is None or node.op_type != "DFT": + return indices + onesided = any( + attribute.name == "onesided" and attribute.i for attribute in node.attribute + ) + resized = onesided or (len(node.input) > 1 and bool(node.input[1])) + return indices if resized else indices[:1] + + +def drop_empty_shape_operands(model: ModelProto) -> None: + """Leave out a reduction's `axes` operand when it holds no elements. + + An empty axes tensor names no axes, which is what leaving the operand out means — ONNX + defines the two the same way — but ONNX's shape inference does not reason about the + values of an operand it cannot see, so it types the result of a reduction over an empty + `axes` *input* as being of unknown rank. Dropping the operand keeps the meaning and lets + inference derive the shape; the tensor stays a graph input, unread. + + The reductions alone: every other op reading an axis list defines an empty one as naming + nothing rather than everything, so dropping it there would change what the node computes. + """ + graph = model.graph + types = tensor_types(graph) + for node in graph.node: + if node.op_type not in REDUCTIONS: + continue + if normalize_domain(node.domain) != STANDARD_DOMAIN: + continue + name = node.input[1] if len(node.input) > 1 else "" + shape = static_shape(types.get(name)) if name else None + if shape is not None and math.prod(shape) == 0: + node.input[1] = "" + + +def state_stft_onesided(model: ModelProto) -> None: + """Write STFT's own default for `onesided` into the nodes that leave it out. + + ONNX's schema declares the default as 1 — and the reference evaluator applies it, so a + one-sided spectrum is what the op computes — while ONNX's shape inference falls back to + 0 and sizes the result at the whole frame length rather than at its non-redundant half. + Stating the default the schema itself declares changes nothing about what the node + computes and leaves inference deriving the shape the op actually produces. + """ + for node in model.graph.node: + if node.op_type != "STFT" or normalize_domain(node.domain) != STANDARD_DOMAIN: + continue + if any(attribute.name == "onesided" for attribute in node.attribute): + continue + declared = onnx.defs.get_schema("STFT").attributes["onesided"].default_value + node.attribute.append(helper.make_attribute("onesided", declared.i)) + + +def bind_dims(model: ModelProto, dim_bindings: Mapping[str, int]) -> dict[str, int]: + """Give every graph input a concrete shape and return the bindings that were applied. + + Symbolic (and unnamed unknown) input dimensions take their value from `dim_bindings`, + defaulting to `UNBOUND_DIM_DEFAULT`. Symbolic dimensions on graph outputs are dropped + rather than bound, and stale intermediate shapes are discarded, so that every shape + downstream follows from the computation instead of from a declaration the graph does + not actually guarantee. An output the graph does not compute is the exception: it just + aliases the input of that name, whose bound shape it therefore takes. Bindings naming a + dimension the model does not use are ignored; dimension names are global across a + bundle, and a node need not use all of them. + """ + for name, size in dim_bindings.items(): + if not isinstance(size, int) or isinstance(size, bool) or size < 0: + raise CompileError( + f"Dimension binding `{name}` must be a non-negative integer, got {size!r}." + ) + + graph = model.graph + initializer_names = {initializer.name for initializer in graph.initializer} + applied: dict[str, int] = {} + for value_info in graph.input: + if value_info.name in initializer_names: + continue + for dim in _dims(value_info): + _bind_dim(dim, dim_bindings, applied) + + bound_inputs = {value_info.name: value_info.type for value_info in graph.input} + for value_info in graph.output: + aliased = bound_inputs.get(value_info.name) + if aliased is not None: + value_info.type.CopyFrom(aliased) + continue + for dim in _dims(value_info): + if dim.WhichOneof("value") == "dim_param": + dim.ClearField("dim_param") + del graph.value_info[:] + return applied + + +def declared_output_types(graph: GraphProto) -> dict[str, TypeProto]: + """Snapshot of the graph outputs' declared types, taken before `bind_dims` clears them.""" + declared = {} + for value_info in graph.output: + if value_info.HasField("type"): + copied = TypeProto() + copied.CopyFrom(value_info.type) + declared[value_info.name] = copied + return declared + + +def apply_declared_output_shapes( + model: ModelProto, + declared: Mapping[str, TypeProto], + dim_bindings: Mapping[str, int], +) -> dict[str, int]: + """Fall back to a graph output's declared shape where inference could not derive one. + + Only outputs produced outside the standard domain qualify: `ai.onnx.ml` inference does + not propagate batch dimensions, so a classifier's declared output shape is the only + static description of it there is. A standard-domain op that inference cannot resolve is + one whose shape depends on runtime data, and its declaration is deliberately not + trusted — compilation fails instead of emitting code for a shape the graph never + guarantees. Returns the bindings the accepted declarations applied. + """ + producers = { + name: node for node in model.graph.node for name in node.output if name + } + applied: dict[str, int] = {} + for value_info in model.graph.output: + if static_shape(value_info.type) is not None: + continue + producer = producers.get(value_info.name) + declared_type = declared.get(value_info.name) + if producer is None or declared_type is None: + continue + if normalize_domain(producer.domain) == STANDARD_DOMAIN: + continue + candidate = TypeProto() + candidate.CopyFrom(declared_type) + bound: dict[str, int] = {} + for dim in _shape_dims(candidate): + _bind_dim(dim, dim_bindings, bound) + if static_shape(candidate) is None: + continue + value_info.type.tensor_type.shape.CopyFrom(candidate.tensor_type.shape) + applied.update(bound) + return applied + + +def drop_shadowed_inputs(model: ModelProto) -> None: + """Remove graph inputs that an initializer also defines. + + Pre-IR-4 models list every initializer as an input with the initializer as its default. + The C compiler embeds initializers as static weights, so those entries are dropped and + the remaining inputs are exactly the tensors a caller must provide. That is IR 4's rule, + and shape inference applies the older one literally — ignoring initializers that are not + inputs — so the model moves to the IR version whose semantics it is now compiled under. + """ + graph = model.graph + initializer_names = {initializer.name for initializer in graph.initializer} + kept = [entry for entry in graph.input if entry.name not in initializer_names] + if len(kept) == len(graph.input): + return + del graph.input[:] + graph.input.extend(kept) + model.ir_version = max(model.ir_version, 4) + + +def infer_shapes(model: ModelProto) -> ModelProto: + """Infer every tensor's type and shape, strictly wherever ONNX can be strict. + + Strict mode turns an inference error into a compile error naming the node, rather than + into a tensor the compiler discovers has no type much later. It also recurses into the + bodies ONNX defines for function ops, and a few of those raise on a node the model is + not answerable for — MeanVarianceNormalization's body builds its `axes` from a Constant + that carries nothing at all unless the node sets the attribute. A relaxed retry is + therefore accepted, but only when it leaves every tensor typed; anything less and the + strict diagnostic is the one worth reporting. + """ + try: + inferred = _infer(model, strict=True) + except CompileError as strict_error: + try: + inferred = _infer(model, strict=False) + except CompileError: + raise strict_error from None + if _has_untyped_tensor(inferred.graph): + raise strict_error from None + return inferred + + +def _infer(model: ModelProto, *, strict: bool) -> ModelProto: + """Run ONNX's inference, filling in the shapes it stops at until it stops adding any. + + Each round can only give a shape to a tensor that had none, so the loop shrinks a finite + set and ends; a graph with nothing to fill in runs inference exactly once. Entries left + shapeless by a previous run are discarded first: folding is what gives inference the + values it was missing, and a stale entry would shadow the shape this run derives. The + ONNX-ML results ONNX derives nothing for are filled in before the first round too, since + strict inference stops at the first node whose operand it cannot type. + """ + _drop_shapeless_value_info(model.graph) + _infer_ml_types(model.graph) + inferred = _run_inference(model, strict=strict) + while _fill_underived_shapes(inferred.graph): + _drop_shapeless_value_info(inferred.graph) + inferred = _run_inference(inferred, strict=strict) + return inferred + + +def _fill_underived_shapes(graph: GraphProto) -> bool: + """Both passes over what ONNX's inference left untyped, neither short-circuiting.""" + preserved = _propagate_preserved_shapes(graph) + return _infer_ml_types(graph) or preserved + + +def _drop_shapeless_value_info(graph: GraphProto) -> None: + """Discard the intermediate entries inference could not give a shape. + + They state an element type and nothing more, and ONNX leaves them behind when it runs + over a graph it has already seen — where a re-run derives the shape, the stale entry + would shadow it, since `tensor_types` reads `value_info` after the graph's outputs. + """ + kept = [entry for entry in graph.value_info if static_shape(entry.type) is not None] + if len(kept) != len(graph.value_info): + del graph.value_info[:] + graph.value_info.extend(kept) + + +def _run_inference(model: ModelProto, *, strict: bool) -> ModelProto: + try: + return onnx.shape_inference.infer_shapes( + model, check_type=True, strict_mode=strict, data_prop=True + ) + except Exception as exc: + raise CompileError( + f"ONNX shape inference failed for graph `{graph_label(model.graph)}`: {exc}" + ) from exc + + +def _propagate_preserved_shapes(graph: GraphProto) -> bool: + """Give an output the shape its op's ONNX schema states it takes from its input. + + ONNX's own inference derives a shape for every op this compiler serves but one: + GroupNormalization is defined as a function whose body reshapes through shapes it + computes, and inference stops at the first of them, leaving a rank it never states — + though the schema says in as many words that `Y` has the shape of `X`. Without this the + op would be compilable only in a model that declares that shape itself, and every tensor + downstream of one would lose its own. + """ + filled = False + types = tensor_types(graph) + for node in graph.node: + if node.op_type not in SHAPE_PRESERVING_OPS: + continue + if normalize_domain(node.domain) != STANDARD_DOMAIN: + continue + source = types.get(node.input[0]) if node.input else None + produced = node.output[0] if node.output else "" + if source is None or static_shape(source) is None or not produced: + continue + declared = types.get(produced) + if declared is None or declared.WhichOneof("value") != "tensor_type": + continue + if static_shape(declared) is not None: + continue + # Every entry the tensor has, since inference leaves a graph output described in + # both `output` and `value_info` and a stale one would shadow the other. + for entry in (*graph.output, *graph.value_info): + if entry.name == produced: + entry.type.tensor_type.shape.CopyFrom(source.tensor_type.shape) + filled = True + return filled + + +def _infer_ml_types(graph: GraphProto) -> bool: + """Type the `ai.onnx.ml` results ONNX's own inference derives nothing for. + + The installed `onnx` package registers no inference function for `Scaler`, `Normalizer`, + `Imputer` or `FeatureVectorizer`; the one it registers for `ArrayFeatureExtractor` stops + short of a rank-1 `X`; and the ones it registers for the two tree ensembles derive an + element type but no shape at all. Every rule below is the op's own schema read literally — + the element type its results are declared as, and the shape its documentation states — so + a model whose intermediates carry no `value_info`, which is most of what ONNX-ML + converters emit, still has a static type for every tensor. Anything derived wrongly here + is a wrong result shape, which the differential sweep compares against the reference + evaluator. + """ + types = tensor_types(graph) + filled = False + for node in graph.node: + if normalize_domain(node.domain) != ML_DOMAIN: + continue + if node.op_type not in ML_UNINFERRED_OPS or not node.output: + continue + derived = _ml_result_types(node, types) + for produced, type_proto in zip(node.output, derived): + if not produced or type_proto is None: + continue + if static_shape(types.get(produced)) is not None: + continue + _set_tensor_type(graph, produced, type_proto) + filled = True + return filled + + +def _ml_result_types( + node: NodeProto, types: Mapping[str, TypeProto] +) -> tuple[TypeProto | None, ...]: + """The types of `node`'s results, or nothing while an operand of its own has none.""" + operands = [types.get(name) for name in node.input] + shapes = [static_shape(operand) for operand in operands] + if not operands or any(shape is None for shape in shapes): + return () + source, shape = operands[0], shapes[0] + assert source is not None and shape is not None + if node.op_type == "Imputer": + return ( + helper.make_tensor_type_proto(source.tensor_type.elem_type, list(shape)), + ) + if node.op_type in ("Normalizer", "Scaler"): + return (helper.make_tensor_type_proto(TensorProto.FLOAT, list(shape)),) + if node.op_type == "FeatureVectorizer": + return (_feature_vectorizer_type(node, shape),) + if node.op_type == "TreeEnsembleRegressor": + return (_tree_ensemble_type(node, shape),) + if node.op_type == "TreeEnsembleClassifier": + return _tree_ensemble_classifier_types(node, shape) + if node.op_type == "LinearRegressor": + return (_linear_regressor_type(node, shape),) + if node.op_type == "SVMRegressor": + return (_svm_regressor_type(shape),) + if node.op_type == "SVMClassifier": + return _svm_classifier_types(node, shape) + # ArrayFeatureExtractor, which takes as many columns as its index operand holds elements + # — one for a scalar, which is the case ONNX's own inference leaves as an unknown + # dimension. A vector `X` is the other case it stops short of, and its reference + # implementation documents that one as following onnxruntime rather than the + # specification: the result is the single row a matrix of one row would have. + indices = shapes[1] + if not shape or indices is None: + return () + taken = math.prod(indices) + return ( + helper.make_tensor_type_proto( + source.tensor_type.elem_type, + [1, taken] if len(shape) == 1 else [*shape[:-1], taken], + ), + ) + + +def _tree_ensemble_rows(shape: tuple[int, ...]) -> int | None: + """How many rows an ensemble scores: one per row of `X`, a vector being a single row.""" + if len(shape) == 2: + return shape[0] + return 1 if len(shape) == 1 else None + + +def _tree_ensemble_type(node: NodeProto, shape: tuple[int, ...]) -> TypeProto | None: + """`TreeEnsembleRegressor`'s `Y`: one float score per target per row.""" + rows = _tree_ensemble_rows(shape) + targets = _attribute(node, "n_targets") + if rows is None or targets is None: + return None + return helper.make_tensor_type_proto(TensorProto.FLOAT, [rows, targets.i]) + + +def _tree_ensemble_classifier_types( + node: NodeProto, shape: tuple[int, ...] +) -> tuple[TypeProto | None, ...]: + """`TreeEnsembleClassifier`'s label and score outputs. + + An ensemble whose leaves all weight one class scores a *pair* of columns whatever its one + class label says, which is the binary rule its reference implementation applies; that is + what decides the width of `Z` and cannot be read off the class labels alone. + """ + rows = _tree_ensemble_rows(shape) + integers = _attribute(node, "classlabels_int64s") + strings = _attribute(node, "classlabels_strings") + if rows is None or (integers is None and strings is None): + return () + classes = max( + len(integers.ints) if integers else 0, len(strings.strings) if strings else 0 + ) + identifiers = _attribute(node, "class_ids") + binary = ( + len({int(value) for value in identifiers.ints}) == 1 if identifiers else False + ) + return ( + helper.make_tensor_type_proto( + TensorProto.INT64 if integers else TensorProto.STRING, [rows] + ), + helper.make_tensor_type_proto( + TensorProto.FLOAT, [rows, 2 if binary and classes == 1 else classes] + ), + ) + + +def _linear_regressor_type(node: NodeProto, shape: tuple[int, ...]) -> TypeProto | None: + """`LinearRegressor`'s `Y`: one float score per target per row.""" + if len(shape) != 2: + return None + targets = _attribute(node, "targets") + return helper.make_tensor_type_proto( + TensorProto.FLOAT, [shape[0], targets.i if targets else 1] + ) + + +def _svm_regressor_type(shape: tuple[int, ...]) -> TypeProto | None: + """`SVMRegressor`'s `Y`: the single score a row is given, as a column of its own.""" + if len(shape) != 2: + return None + return helper.make_tensor_type_proto(TensorProto.FLOAT, [shape[0], 1]) + + +def _svm_classifier_types( + node: NodeProto, shape: tuple[int, ...] +) -> tuple[TypeProto | None, ...]: + """`SVMClassifier`'s label and score outputs. + + How wide a row of `Z` is depends on the whole shape of the node: an ensemble over support + vectors scores one value per *pair* of classes unless it couples them into probabilities, + a linear one scores one per class, and a lone score is paired with a second where a second + class is called for. The kernel derives the same width from the same attributes and + refuses to write a buffer that disagrees with it. + """ + if len(shape) != 2: + return () + integers = _attribute(node, "classlabels_ints") + strings = _attribute(node, "classlabels_strings") + if integers is None and strings is None: + return () + classes = max( + len(integers.ints) if integers else 0, len(strings.strings) if strings else 0 + ) + counts = _attribute(node, "vectors_per_class") + vectors = sum(counts.ints) if counts else 0 + coupled = vectors > 0 and bool(_floats(node, "prob_a")) + scored = ( + max(classes, 1) if vectors == 0 or coupled else classes * (classes - 1) // 2 + ) + transform = _attribute(node, "post_transform") + paired = ( + scored == 1 + and classes == 2 + and len(_floats(node, "rho")) == 1 + and (transform is None or transform.s != b"PROBIT") + ) + return ( + helper.make_tensor_type_proto( + TensorProto.INT64 if integers else TensorProto.STRING, [shape[0]] + ), + helper.make_tensor_type_proto( + TensorProto.FLOAT, [shape[0], 2 if paired else scored] + ), + ) + + +def _floats(node: NodeProto, name: str) -> list[float]: + attribute = _attribute(node, name) + return list(attribute.floats) if attribute else [] + + +def _attribute(node: NodeProto, name: str) -> AttributeProto | None: + return next( + (entry for entry in node.attribute if entry.name == name), + None, + ) + + +def _feature_vectorizer_type( + node: NodeProto, first: tuple[int, ...] +) -> TypeProto | None: + """`FeatureVectorizer`'s result: one row per input row, `inputdimensions` wide in total.""" + widths = _attribute(node, "inputdimensions") + if widths is None or len(first) not in (1, 2): + return None + return helper.make_tensor_type_proto( + TensorProto.FLOAT, [first[0], sum(widths.ints)] + ) + + +def _set_tensor_type(graph: GraphProto, name: str, type_proto: TypeProto) -> None: + """Record `name`'s type, on every entry describing it and on a new one if there is none.""" + entries = [ + entry for entry in (*graph.output, *graph.value_info) if entry.name == name + ] + if not entries: + entries = [graph.value_info.add()] + entries[0].name = name + for entry in entries: + entry.type.CopyFrom(type_proto) + + +def _has_untyped_tensor(graph: GraphProto) -> bool: + types = tensor_types(graph) + produced = (name for node in graph.node for name in node.output if name) + return any( + static_shape(types.get(name)) is None + for name in (*produced, *(entry.name for entry in graph.output)) + ) + + +def tensor_types(graph: GraphProto) -> dict[str, TypeProto]: + """Type of every tensor the graph names, initializers included.""" + types = { + initializer.name: helper.make_tensor_type_proto( + initializer.data_type, list(initializer.dims) + ) + for initializer in graph.initializer + } + for value_info in (*graph.input, *graph.output, *graph.value_info): + if value_info.HasField("type"): + types[value_info.name] = value_info.type + return types + + +def static_shape(type_proto: TypeProto | None) -> tuple[int, ...] | None: + """Concrete dimensions of a tensor type, or None if any of them is not static. + + A zero-sized dimension is static: zero-element tensors are legal in ONNX. + """ + if type_proto is None or type_proto.WhichOneof("value") != "tensor_type": + return None + tensor_type = type_proto.tensor_type + if not tensor_type.HasField("shape"): + return None + dims: list[int] = [] + for dim in tensor_type.shape.dim: + if dim.WhichOneof("value") != "dim_value" or dim.dim_value < 0: + return None + dims.append(dim.dim_value) + return tuple(dims) + + +def graph_label(graph: GraphProto) -> str: + return graph.name or "" + + +def _bind_dim( + dim: TensorShapeProto.Dimension, + dim_bindings: Mapping[str, int], + applied: dict[str, int], +) -> None: + if dim.WhichOneof("value") == "dim_value": + return + name = dim.dim_param + dim.dim_value = dim_bindings.get(name, UNBOUND_DIM_DEFAULT) + if name: + applied[name] = dim.dim_value + + +def _dims(value_info: ValueInfoProto) -> Iterator[TensorShapeProto.Dimension]: + yield from _shape_dims(value_info.type) + + +def _shape_dims(type_proto: TypeProto) -> Iterator[TensorShapeProto.Dimension]: + if type_proto.WhichOneof("value") != "tensor_type": + return + yield from type_proto.tensor_type.shape.dim diff --git a/src/python/fnnx/extras/compilers/c/onnx/specialize.py b/src/python/fnnx/extras/compilers/c/onnx/specialize.py new file mode 100644 index 0000000..3c07b4b --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/specialize.py @@ -0,0 +1,609 @@ +"""Compiling one artifact for a whole family of shapes, not for a single one. + +A runtime dimension is compiled by compiling the model *repeatedly*, at a spread of values +for that dimension, and reading the family off the results: a tensor axis that is the same +at every value is a constant, one that is the dimension's value times a fixed factor is that +multiple, and anything else is a shape the artifact cannot be one piece of code for. The +same reading turns the compile-time literals in the emitted call sites into expressions over +the dimension's entrypoint parameter, so kernels loop to the size a call actually asks for +while every buffer stays sized for the maximum. + +Probing rather than symbolic inference is deliberate: the frontend that derives the shapes, +and the kernels that turn them into code, stay exactly the ones a fixed-shape compilation +uses — there is no second, symbolic implementation of either that could disagree with it. +What the probes do not agree on is rejected, never guessed at, so the artifact is either +correct across the whole family or refused at compile time. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, replace + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.codegen import ( + IOTensor, + NodeEntry, + Program, + StaticBuffer, +) +from fnnx.extras.compilers.c.onnx.runtime_dims import RuntimeDim, ShapeTerm + +Builder = Callable[[Mapping[str, int]], Program] + +# A decimal integer standing on its own: not part of an identifier, and not the fraction or +# the exponent of a floating-point literal. +_LITERAL = re.compile(r"(? Program: + """Compile through `build` for every size in the runtime dimensions' ranges. + + `build` compiles the model at one concrete set of dimension values, exactly as a + fixed-shape compilation would. The returned program is the one built at the maxima — + so every buffer is sized for the capacity — with its emitted code rewritten to work at + the value each call passes. + """ + probes = _probe_bindings(dims) + programs = [_build_at(build, dim_bindings, probe, dims) for probe in probes] + liftable = _liftable_probes(probes, dims) + base = programs[0] + + _check_signatures(programs, dims) + _check_definitions(programs, liftable, dims) + nodes = tuple( + replace( + entry, + inputs=_shaped( + entry.inputs, + [p.nodes[index].inputs for p in programs], + probes, + dims, + entry.id, + ), + outputs=_shaped( + entry.outputs, + [p.nodes[index].outputs for p in programs], + probes, + dims, + entry.id, + ), + body=_lift_body( + [p.nodes[index] for p in programs], liftable, probes, dims, entry.id + ), + ) + for index, entry in enumerate(base.nodes) + ) + return replace( + base, + inputs=_shaped(base.inputs, [p.inputs for p in programs], probes, dims, ""), + outputs=_shaped(base.outputs, [p.outputs for p in programs], probes, dims, ""), + body=_lift_body(programs, liftable, probes, dims, ""), + nodes=nodes, + runtime_dims=tuple(dims), + # The probe values are how the family was explored, not bindings the artifact was + # compiled under; only the dimensions genuinely fixed at compile time are reported. + dim_bindings={ + name: value + for name, value in base.dim_bindings.items() + if name not in {dim.name for dim in dims} + }, + ) + + +# -------------------------------------------------------------------------------------- +# The values the model is compiled at +# -------------------------------------------------------------------------------------- + + +def _build_at( + build: Builder, + dim_bindings: Mapping[str, int], + probe: Mapping[str, int], + dims: Sequence[RuntimeDim], +) -> Program: + """Compile at one probe, reporting a failure as one of the whole family. + + A model that only compiles at some sizes — a reshape into a fixed extent the smallest + size cannot fill, say — is one this artifact cannot be, and the size it broke at is the + most useful thing to say about it. + """ + try: + return build({**dim_bindings, **probe}) + except CompileError as error: + raise _untrackable( + f"compiling it at {_probe_label(probe, dims)} failed: {error}", "", dims + ) from error + + +def probe_values(dim: RuntimeDim) -> tuple[int, ...]: + """The sizes one dimension is probed at. + + 1, 2 and 3 pin down the small end of the range — where a shape that clamps, saturates + or rounds parts company with a linear one — and the two largest values anchor the other + end, so that a fit through the small values has to hold at the capacity as well. + """ + candidates = {1, 2, 3, dim.maximum - 1, dim.maximum} + return tuple(sorted(value for value in candidates if 1 <= value <= dim.maximum)) + + +def _probe_bindings(dims: Sequence[RuntimeDim]) -> tuple[dict[str, int], ...]: + """One probe per (dimension, size), every other dimension held at its maximum. + + The first probe holds every dimension at its maximum: it is the one the artifact's + buffers, macros and reported footprint are taken from. A last probe moves every + dimension away from its maximum at once, because the ones above move one at a time — a + size that is the *product* of two dimensions agrees with a linear reading along each of + them separately, and parts company with it only where both move together. + """ + maxima = {dim.name: dim.maximum for dim in dims} + probes = [dict(maxima)] + probes += [ + {**maxima, dim.name: value} + for dim in dims + for value in probe_values(dim) + if value != dim.maximum + ] + if len(dims) > 1: + probes.append({dim.name: min(2, dim.maximum) for dim in dims}) + seen: set[tuple[tuple[str, int], ...]] = set() + distinct = [] + for probe in probes: + key = tuple(sorted(probe.items())) + if key not in seen: + seen.add(key) + distinct.append(probe) + return tuple(distinct) + + +def _liftable_probes( + probes: Sequence[Mapping[str, int]], dims: Sequence[RuntimeDim] +) -> tuple[int, ...]: + """The probes whose emitted code is compared against each other, by index. + + A dimension of 1 makes a tensor's axis vanish: an operand stops broadcasting, a + concatenation becomes contiguous, and kernels legitimately emit a different — faster, + equally correct — form for it. Comparing those forms against the general one would + reject a model the general form serves perfectly well, so the code is read off the + sizes above 1 while the *shapes* are still checked at 1 as well. + + Leaving 1 out is only safe while the sizes that remain outnumber what an affine reading + of a literal has free parameters: two sizes determine a slope and an intercept, so a + quantity that is quadratic in the dimension would fit them and be lifted wrongly. A + dimension whose maximum is too small to spare three sizes above 1 is therefore read at + 1 as well — its whole range then appears among the probes, which makes the reading exact + rather than fitted, at the price of rejecting a model whose kernels take a size-1 form. + """ + return tuple( + index + for index, probe in enumerate(probes) + if all( + probe[dim.name] > 1 + for dim in dims + if dim.maximum >= _SMALLEST_GENERAL_MAXIMUM + ) + ) + + +# -------------------------------------------------------------------------------------- +# Shape families +# -------------------------------------------------------------------------------------- + + +def _shaped( + tensors: Sequence[IOTensor], + across: Sequence[Sequence[IOTensor]], + probes: Sequence[Mapping[str, int]], + dims: Sequence[RuntimeDim], + owner: str, +) -> tuple[IOTensor, ...]: + return tuple( + replace( + tensor, + runtime_shape=_shape_terms( + [variant[index].shape for variant in across], + probes, + dims, + subject=f"tensor `{tensor.name}`", + owner=tensor.owner or owner, + ), + ) + for index, tensor in enumerate(tensors) + ) + + +def _shape_terms( + shapes: Sequence[tuple[int, ...]], + probes: Sequence[Mapping[str, int]], + dims: Sequence[RuntimeDim], + *, + subject: str, + owner: str, +) -> tuple[ShapeTerm, ...]: + """Read one tensor's shape family off the shapes it took at each probe.""" + base = shapes[0] + if any(len(shape) != len(base) for shape in shapes): + raise _untrackable( + f"the rank of {subject} changes with the dimension's value", owner, dims + ) + return tuple( + _axis_term( + [shape[axis] for shape in shapes], + probes, + dims, + subject=subject, + subject_owner=owner, + axis=axis, + ) + for axis in range(len(base)) + ) + + +def _axis_term( + sizes: Sequence[int], + probes: Sequence[Mapping[str, int]], + dims: Sequence[RuntimeDim], + *, + subject: str, + subject_owner: str, + axis: int, +) -> ShapeTerm: + """One axis as a constant or as a fixed multiple of one runtime dimension.""" + if len(set(sizes)) == 1: + return ShapeTerm(sizes[0]) + for dim in dims: + coefficient, remainder = divmod(sizes[0], dim.maximum) + if remainder or coefficient < 1: + continue + if all( + size == coefficient * probe[dim.name] for size, probe in zip(sizes, probes) + ): + return ShapeTerm(sizes[0], dim.name, coefficient) + raise _untrackable( + f"axis {axis} of {subject} is neither a constant nor a constant multiple of a " + f"runtime dimension ({_observations(probes, sizes, dims)})", + subject_owner, + dims, + ) + + +# -------------------------------------------------------------------------------------- +# Lifting the emitted code +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _Fit: + """A literal read as `intercept + sum(slope * dimension value)`.""" + + intercept: int + slopes: tuple[int, ...] + + def at(self, probe: Mapping[str, int], dims: Sequence[RuntimeDim]) -> int: + return self.intercept + sum( + slope * probe[dim.name] for slope, dim in zip(self.slopes, dims) + ) + + def bounds(self, dims: Sequence[RuntimeDim]) -> tuple[int, int]: + """The smallest and largest values the fit takes anywhere in the dimensions' ranges.""" + low = high = self.intercept + for slope, dim in zip(self.slopes, dims): + ends = (slope, slope * dim.maximum) + low += min(ends) + high += max(ends) + return low, high + + def render(self, dims: Sequence[RuntimeDim]) -> str: + text = "" + for slope, dim in zip(self.slopes, dims): + if slope: + text += f" - {-slope}" if slope < 0 else f" + {slope}" + text += f" * {dim.c_name}" + if self.intercept or not text: + text += ( + f" - {-self.intercept}" + if self.intercept < 0 + else f" + {self.intercept}" + ) + return "(" + text.removeprefix(" + ").lstrip() + ")" + + +def _lift_body( + programs: Sequence[Program | NodeEntry], + liftable: Sequence[int], + probes: Sequence[Mapping[str, int]], + dims: Sequence[RuntimeDim], + entry: str, +) -> tuple[str, ...]: + """One entrypoint's statements, with every size-dependent literal made an expression.""" + bodies = [programs[index].body for index in liftable] + at = [probes[index] for index in liftable] + owners = programs[0].body_owners + where = f"node entrypoint `{entry}`" if entry else "the model entrypoint" + if len({len(body) for body in bodies}) != 1: + raise _untrackable( + f"the code emitted for {where} changes shape with the dimension's value", + "", + dims, + ) + return tuple( + _lift_statement( + [body[index] for body in bodies], + at, + dims, + owner=owners[index] if index < len(owners) else "", + ) + for index in range(len(bodies[0])) + ) + + +def _lift_statement( + variants: Sequence[str], + probes: Sequence[Mapping[str, int]], + dims: Sequence[RuntimeDim], + *, + owner: str, +) -> str: + separators, tokens = zip(*(_split(variant) for variant in variants)) + if len(set(separators)) != 1: + raise _untrackable( + "the code emitted for it is not the same at every size", owner, dims + ) + pieces = list(separators[0]) + lifted = [ + _lift_literal( + [token[index] for token in tokens], + probes, + dims, + owner=owner, + # A literal an expression cannot stand in for: a macro pastes its argument onto + # a suffix, and one inside a string literal is text rather than a size. + pasted=bool(_MACRO_ARGUMENT.search(pieces[index])) + or '"' in "".join(pieces[: index + 1]), + ) + for index in range(len(tokens[0])) + ] + text = pieces[0] + for literal, piece in zip(lifted, pieces[1:]): + text += literal + piece + return text + + +def _split(text: str) -> tuple[tuple[str, ...], tuple[str, ...]]: + """`text` as the runs between its integer literals, and the literals themselves.""" + separators, literals, position = [], [], 0 + for match in _LITERAL.finditer(text): + separators.append(text[position : match.start()]) + literals.append(match.group(0)) + position = match.end() + separators.append(text[position:]) + return tuple(separators), tuple(literals) + + +def _lift_literal( + literals: Sequence[str], + probes: Sequence[Mapping[str, int]], + dims: Sequence[RuntimeDim], + *, + owner: str, + pasted: bool, +) -> str: + if len(set(literals)) == 1: + return literals[0] + if pasted: + raise _untrackable( + "a value it pastes into a macro or a string depends on the dimension's value", + owner, + dims, + ) + values = [int(literal.rstrip("uU")) for literal in literals] + fit = _fit(values, probes, dims) + if fit is None: + raise _untrackable( + "a size in the code emitted for it does not scale linearly with the " + f"dimension's value ({_observations(probes, values, dims)})", + owner, + dims, + ) + low, high = fit.bounds(dims) + if low < 0 or high > _INT_MAX: + raise _untrackable( + f"a size in the code emitted for it runs to {low}..{high} across the " + "dimension's range, which no buffer of the artifact could be indexed by", + owner, + dims, + ) + return fit.render(dims) + + +def _fit( + values: Sequence[int], + probes: Sequence[Mapping[str, int]], + dims: Sequence[RuntimeDim], +) -> _Fit | None: + """The one affine reading of `values`, or None where they do not admit one. + + Every probe holds all but one dimension at its maximum, so each slope follows from the + pair of probes that differ in that dimension alone; the reading is then checked against + every probe, which is what rejects a size that is quadratic in a dimension, that + saturates, or that mixes two dimensions into a product. + """ + slopes = [] + for dim in dims: + pair = next( + ( + index + for index, probe in enumerate(probes) + if probe[dim.name] != dim.maximum + and all( + probe[other.name] == other.maximum + for other in dims + if other is not dim + ) + ), + None, + ) + if pair is None: + slopes.append(0) + continue + rise = values[0] - values[pair] + run = dim.maximum - probes[pair][dim.name] + if rise % run: + return None + slopes.append(rise // run) + fit = _Fit( + intercept=values[0] - sum(s * d.maximum for s, d in zip(slopes, dims)), + slopes=tuple(slopes), + ) + if any(fit.at(probe, dims) != value for probe, value in zip(probes, values)): + return None + return fit + + +# -------------------------------------------------------------------------------------- +# What every probe has to agree on +# -------------------------------------------------------------------------------------- + + +def _check_signatures(programs: Sequence[Program], dims: Sequence[RuntimeDim]) -> None: + """Which entrypoints the artifact publishes, and what each takes, is size-independent. + + Read at *every* probe, the degenerate ones included: the shapes checked against these + signatures are the model's own semantics, not a form some kernel happened to emit. + """ + base = programs[0] + shapes = [ + (len(program.inputs), len(program.outputs)) + + tuple( + (entry.id, entry.symbol, len(entry.inputs), len(entry.outputs)) + for entry in program.nodes + ) + for program in programs + ] + if len(set(shapes)) != 1: + raise _untrackable( + f"the entrypoints of `{base.prefix}` change with the dimension's value", + "", + dims, + ) + + +def _check_definitions( + programs: Sequence[Program], liftable: Sequence[int], dims: Sequence[RuntimeDim] +) -> None: + """Everything outside the entrypoint bodies has to be identical at every size. + + Weights and constant tables are the compile-time values the graph fixes, and a value + that moves with a dimension is a computation folded away that the artifact would have to + do at run time. Those are read at *every* probe, size 1 included: nothing about a folded + value is a form the emitter chose, and size 1 is where a clamped or degenerate fold most + often parts company with the rest of the range — exactly the probe the code comparison + leaves out. Which kernels get emitted, in contrast, *is* such a form, so it is read + where the code is, alongside the buffers those kernels ask for. + """ + base = programs[0] + for other in programs[1:]: + # Tensor by tensor first, so that the one that moved can be named — it is the only + # thread back to the computation the fold took out of the graph. + for one, another in zip(base.weights, other.weights): + _same(one, another, f"the constant `{one.name}` it embeds", dims) + _same(base.weights, other.weights, "the constant data it embeds", dims) + _same( + base.labels, + other.labels, + "the set of class-label tables it publishes", + dims, + ) + for index in liftable[1:]: + other = programs[index] + _same(base.functions, other.functions, "the set of kernels it emits", dims) + _check_capacity(base.scratch, other.scratch, dims) + + +def _same( + base: object, other: object, subject: str, dims: Sequence[RuntimeDim] +) -> None: + if base != other: + raise _untrackable(f"{subject} changes with the dimension's value", "", dims) + + +def _check_capacity( + base: Sequence[StaticBuffer], + other: Sequence[StaticBuffer], + dims: Sequence[RuntimeDim], +) -> None: + """The buffers planned at the maxima have to hold every smaller size as well.""" + sized = {buffer.symbol: buffer.declared_count for buffer in base} + for buffer in other: + capacity = sized.get(buffer.symbol) + if capacity is None: + raise _untrackable( + f"the buffer `{buffer.symbol}` it reserves depends on the dimension's " + "value", + "", + dims, + ) + if capacity < buffer.declared_count: + raise _untrackable( + f"the buffer `{buffer.symbol}` needs {buffer.declared_count} elements at " + f"a smaller size than the {capacity} it is given at the maximum", + "", + dims, + ) + + +# -------------------------------------------------------------------------------------- +# Errors +# -------------------------------------------------------------------------------------- + + +def _untrackable(detail: str, owner: str, dims: Sequence[RuntimeDim]) -> CompileError: + named = ", ".join(f"`{dim.name}`" for dim in dims) + where = f"Node `{owner}`: " if owner else "" + # With such a maximum, size 1 is one of the sizes the emitted code is read at (see + # `_liftable_probes`), and a kernel emits a different — still correct — form there, + # which is a failure of the schedule rather than of the model. + hint = ( + f" A maximum below {_SMALLEST_GENERAL_MAXIMUM} leaves too few sizes above 1 to " + "read the code at; try a larger one." + if any(dim.maximum < _SMALLEST_GENERAL_MAXIMUM for dim in dims) + else "" + ) + return CompileError( + f"{where}{detail}. Runtime dimension(s) {named} cannot be tracked through this " + f"model; pin them via `dim_bindings` to compile it.{hint}" + ) + + +def _probe_label(probe: Mapping[str, int], dims: Sequence[RuntimeDim]) -> str: + return ", ".join(f"{dim.name}={probe[dim.name]}" for dim in dims) + + +def _observations( + probes: Sequence[Mapping[str, int]], + values: Sequence[int], + dims: Sequence[RuntimeDim], +) -> str: + return ", ".join( + f"{_probe_label(probe, dims)} -> {value}" + for probe, value in zip(probes, values) + ) diff --git a/src/python/fnnx/extras/compilers/c/onnx/verify.py b/src/python/fnnx/extras/compilers/c/onnx/verify.py new file mode 100644 index 0000000..7e1fd5f --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/verify.py @@ -0,0 +1,141 @@ +"""Verification that a prepared graph is fully static and uses only supported types.""" + +from __future__ import annotations + +from collections.abc import Container, Mapping + +from onnx import GraphProto, ModelProto, NodeProto, TypeProto + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.dtypes import element_type_name, is_supported +from fnnx.extras.compilers.c.onnx.loader import display_domain, normalize_domain +from fnnx.extras.compilers.c.onnx.shapes import ( + graph_label, + runtime_shape_operand, + static_shape, + tensor_types, +) + +CONTROL_FLOW_OPS = frozenset({"If", "Loop", "Scan"}) + +# Ops whose output shape is a function of input values, not of input shapes: no binding or +# folding can make them static, so they are rejected by name rather than through the +# generic shape check, which could only report the symptom. +DATA_DEPENDENT_SHAPE_OPS = frozenset( + {"NonZero", "Unique", "Compress", "NonMaxSuppression"} +) + + +def verify_static(model: ModelProto) -> None: + """Raise a `CompileError` unless every tensor of the graph is statically compilable.""" + graph = model.graph + types = tensor_types(graph) + constants = {initializer.name for initializer in graph.initializer} + for node in graph.node: + _verify_node(graph, node, constants, types) + if graph.sparse_initializer: + names = ", ".join( + f"`{sparse.values.name}`" for sparse in graph.sparse_initializer + ) + raise CompileError( + f"Graph `{graph_label(graph)}`: sparse initializers ({names}) are not " + "supported by the C compiler." + ) + + for entry in graph.input: + _verify_tensor( + graph, entry.name, types.get(entry.name), f"input `{entry.name}`" + ) + for initializer in graph.initializer: + _verify_tensor( + graph, + initializer.name, + types.get(initializer.name), + f"initializer `{initializer.name}`", + ) + for node in graph.node: + for index, name in enumerate(node.output): + if not name: + continue + _verify_tensor( + graph, + name, + types.get(name), + f"output {index} of node `{_node_label(node)}`", + ) + + +def _verify_node( + graph: GraphProto, + node: NodeProto, + constants: Container[str], + types: Mapping[str, TypeProto], +) -> None: + domain = normalize_domain(node.domain) + if domain != "": + return + operand = runtime_shape_operand(node, constants, types) + if operand is not None: + raise CompileError( + f"Graph `{graph_label(graph)}`: node `{_node_label(node)}` takes the shape of " + f"its `{node.op_type}` output from `{operand}`, which no initializer or " + "constant folding fixes; that shape then depends on input data, which the C " + "compiler requires to be known at compile time." + ) + if node.op_type in CONTROL_FLOW_OPS: + raise CompileError( + f"Graph `{graph_label(graph)}`: node `{_node_label(node)}` uses control flow " + f"op `{node.op_type}`, which the C compiler supports only when constant " + "folding can resolve it away; that needs its inputs and everything its " + "subgraphs read from this graph to be known at compile time." + ) + if node.op_type in DATA_DEPENDENT_SHAPE_OPS: + raise CompileError( + f"Graph `{graph_label(graph)}`: node `{_node_label(node)}` uses op " + f"`{node.op_type}` (domain `{display_domain(domain)}`), whose output shape " + "depends on input data; the C compiler requires every shape to be known at " + "compile time." + ) + + +def _verify_tensor( + graph: GraphProto, name: str, type_proto: TypeProto | None, role: str +) -> None: + label = f"Graph `{graph_label(graph)}`: {role}" + if type_proto is None or not type_proto.WhichOneof("value"): + raise CompileError( + f"{label} has no type; the C compiler could not infer one for tensor `{name}`." + ) + kind = type_proto.WhichOneof("value") + if kind != "tensor_type": + raise CompileError( + f"{label} has type `{kind}`, which the C compiler does not support; only " + "tensors can be compiled." + ) + elem_type = type_proto.tensor_type.elem_type + if not is_supported(elem_type): + raise CompileError( + f"{label} has element type `{element_type_name(elem_type)}`, which the C " + "compiler does not support." + ) + if static_shape(type_proto) is None: + raise CompileError( + f"{label} has shape `{_shape_label(type_proto)}`, which is not static; bind " + "its symbolic dimensions with `dim_bindings` or remove the data-dependent " + "computation that produces it." + ) + + +def _shape_label(type_proto: TypeProto) -> str: + tensor_type = type_proto.tensor_type + if not tensor_type.HasField("shape"): + return "" + dims = [] + for dim in tensor_type.shape.dim: + kind = dim.WhichOneof("value") + dims.append(str(dim.dim_value) if kind == "dim_value" else dim.dim_param or "?") + return f"[{', '.join(dims)}]" + + +def _node_label(node: NodeProto) -> str: + return node.name or f"" diff --git a/src/python/fnnx/extras/compilers/c/onnx/zipmap.py b/src/python/fnnx/extras/compilers/c/onnx/zipmap.py new file mode 100644 index 0000000..b044195 --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/onnx/zipmap.py @@ -0,0 +1,123 @@ +"""The graph pass that turns an ONNX-ML classifier's map output back into a tensor. + +`ZipMap` exists only to pair a classifier's probability tensor with its class names, and its +result is a sequence of maps — not a tensor, so not something the C compiler can hand a +caller a buffer for. The pass removes a trailing one, promotes the probability tensor it was +reading to the graph output in its place, and hands the class names on as metadata the header +publishes alongside that output. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from onnx import GraphProto, ModelProto, NodeProto, TensorProto, ValueInfoProto, helper + +from fnnx.extras.compilers.c.errors import CompileError +from fnnx.extras.compilers.c.onnx.loader import ML_DOMAIN, normalize_domain +from fnnx.extras.compilers.c.onnx.shapes import graph_label + +ZIP_MAP = "ZipMap" + + +@dataclass(frozen=True) +class ClassLabels: + """The class names a removed `ZipMap` was keying its map output by. + + `tensor` is the graph output the labels describe, one label per element of its trailing + axis. `elem_type` is `STRING` or `INT64` — the two key types ONNX-ML defines. + """ + + tensor: str + elem_type: int + values: tuple[str, ...] | tuple[int, ...] + + +def remove_zipmap(model: ModelProto) -> tuple[ClassLabels, ...]: + """Drop every trailing `ZipMap` node, returning the class labels each one carried. + + A `ZipMap` anywhere but at the end of the graph is a compile error: its result is a + sequence of maps, and nothing this compiler emits can consume or produce one. + """ + graph = model.graph + consumed = {name for node in graph.node for name in node.input if name} + outputs = {entry.name for entry in graph.output} + labels: list[ClassLabels] = [] + kept: list[NodeProto] = [] + for node in graph.node: + if node.op_type != ZIP_MAP or normalize_domain(node.domain) != ML_DOMAIN: + kept.append(node) + continue + produced = node.output[0] if node.output else "" + source = node.input[0] if node.input else "" + if ( + not produced + or not source + or produced in consumed + or produced not in outputs + ): + raise CompileError( + f"Graph `{graph_label(graph)}`: node `{_label(node)}` produces a sequence " + "of maps, which the C compiler supports only as the last node on a graph " + "output, where it is removed and its class labels are published as header " + "metadata." + ) + elem_type, values = _class_labels(graph, node) + labels.append(ClassLabels(source, elem_type, values)) + _promote(graph, produced, source) + if len(kept) != len(graph.node): + del graph.node[:] + graph.node.extend(kept) + return tuple(labels) + + +def _class_labels( + graph: GraphProto, node: NodeProto +) -> tuple[int, tuple[str, ...] | tuple[int, ...]]: + strings = _attribute(node, "classlabels_strings") or () + integers = _attribute(node, "classlabels_int64s") or () + if bool(strings) == bool(integers): + raise CompileError( + f"Graph `{graph_label(graph)}`: node `{_label(node)}` must set exactly one of " + "`classlabels_strings` and `classlabels_int64s`." + ) + if strings: + return TensorProto.STRING, tuple( + value.decode("utf-8") if isinstance(value, bytes) else str(value) + for value in strings + ) + return TensorProto.INT64, tuple(int(value) for value in integers) + + +def _promote(graph: GraphProto, produced: str, source: str) -> None: + """Make `source` the graph output `produced` was. + + The declared type goes with the node: it described a sequence of maps, while the tensor + taking its place is typed by shape inference from the classifier that produces it. An + output the graph already exposes under its own name keeps that single entry. + """ + replacement = ValueInfoProto() + replacement.name = source + already_exposed = any(entry.name == source for entry in graph.output) + kept = [entry for entry in graph.output if entry.name != produced] + if not already_exposed: + kept.insert(_position(graph, produced), replacement) + del graph.output[:] + graph.output.extend(kept) + + +def _position(graph: GraphProto, name: str) -> int: + """Where `name` sits among the graph outputs, so the replacement keeps its place.""" + return next(index for index, entry in enumerate(graph.output) if entry.name == name) + + +def _attribute(node: NodeProto, name: str) -> Any: + for entry in node.attribute: + if entry.name == name: + return helper.get_attribute_value(entry) + return None + + +def _label(node: NodeProto) -> str: + return node.name or f"" diff --git a/src/python/fnnx/extras/compilers/c/result.py b/src/python/fnnx/extras/compilers/c/result.py new file mode 100644 index 0000000..c3c98be --- /dev/null +++ b/src/python/fnnx/extras/compilers/c/result.py @@ -0,0 +1,25 @@ +"""The object every `compile_*` entrypoint returns.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from fnnx.extras.compilers.c.harness import CompiledModel + + +@dataclass(frozen=True) +class CompileResult: + header_path: Path + report_path: Path + report: dict[str, Any] + + def load(self, *, compiler: str | None = None) -> CompiledModel: + """Build this artifact into a shared library and bind it for execution.""" + # Imported here: driving an artifact needs numpy and a C compiler, which merely + # compiling one does not. + from fnnx.extras.compilers.c.harness import load_compiled + + return load_compiled(self.report_path, compiler=compiler) diff --git a/src/python/pyproject.toml b/src/python/pyproject.toml index 44f17f6..1200e10 100644 --- a/src/python/pyproject.toml +++ b/src/python/pyproject.toml @@ -17,6 +17,9 @@ classifiers = [ core = ["numpy>=1.0.0,<3.0.0", "onnxruntime>=1.0.0,<2.0.0"] extras = ["pydantic>=2.0.0,<3.0.0", "pyyaml>=5.1"] mlflow = ["mlflow>=3.4"] +# Ahead-of-time compiler to C (fnnx.extras.compilers.c); the core runtime never +# depends on onnx. The installed version bounds the opsets that can be compiled. +compiler = ["onnx>=1.16,<2.0"] # Essential test deps: only the pure FNNX core (format/runtime/validators), # no mlflow converter and no ML frameworks. Drives the 3.10 / 3.14 CI legs, # which skip the extra suite (the `test_extra_*` files); see the test workflow. diff --git a/src/python/tests/conformance/dispositions.json b/src/python/tests/conformance/dispositions.json new file mode 100644 index 0000000..9efdefe --- /dev/null +++ b/src/python/tests/conformance/dispositions.json @@ -0,0 +1,744 @@ +{ + "ai.onnx": { + "Abs": { + "disposition": "native-kernel" + }, + "Acos": { + "disposition": "native-kernel" + }, + "Acosh": { + "disposition": "native-kernel" + }, + "Add": { + "disposition": "native-kernel" + }, + "AffineGrid": { + "disposition": "native-kernel" + }, + "And": { + "disposition": "native-kernel" + }, + "ArgMax": { + "disposition": "native-kernel" + }, + "ArgMin": { + "disposition": "native-kernel" + }, + "Asin": { + "disposition": "native-kernel" + }, + "Asinh": { + "disposition": "native-kernel" + }, + "Atan": { + "disposition": "native-kernel" + }, + "Atanh": { + "disposition": "native-kernel" + }, + "Attention": { + "disposition": "native-kernel" + }, + "AveragePool": { + "disposition": "native-kernel" + }, + "BatchNormalization": { + "disposition": "native-kernel" + }, + "Bernoulli": { + "disposition": "unsupported", + "reason": "random-op", + "note": "Draws from a Bernoulli distribution, so its result is not a function of its inputs; folding one would bake a single draw into the artifact and compiling one would need a generator the artifact does not carry." + }, + "BitCast": { + "disposition": "native-kernel" + }, + "BitShift": { + "disposition": "native-kernel" + }, + "BitwiseAnd": { + "disposition": "native-kernel" + }, + "BitwiseNot": { + "disposition": "native-kernel" + }, + "BitwiseOr": { + "disposition": "native-kernel" + }, + "BitwiseXor": { + "disposition": "native-kernel" + }, + "BlackmanWindow": { + "disposition": "folding-or-graph-pass", + "note": "A window is the `size` operand and nothing else, so the folding pass resolves every node whose `size` the graph fixes into an initializer; a `size` that is only known at run time makes the window's own extent data-dependent." + }, + "Cast": { + "disposition": "native-kernel" + }, + "CastLike": { + "disposition": "function-expansion" + }, + "CausalConvWithState": { + "disposition": "function-expansion" + }, + "Ceil": { + "disposition": "native-kernel" + }, + "Celu": { + "disposition": "native-kernel" + }, + "CenterCropPad": { + "disposition": "function-expansion" + }, + "Clip": { + "disposition": "native-kernel" + }, + "Col2Im": { + "disposition": "native-kernel" + }, + "Compress": { + "disposition": "unsupported", + "reason": "data-dependent-shape", + "note": "Keeps the elements a run-time condition selects, so how many elements its result holds is a function of input values rather than of input shapes." + }, + "Concat": { + "disposition": "native-kernel" + }, + "ConcatFromSequence": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Reads a sequence, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "Constant": { + "disposition": "folding-or-graph-pass", + "note": "Carries its value in an attribute, which the folding pass turns into an initializer before dispatch is reached." + }, + "ConstantOfShape": { + "disposition": "folding-or-graph-pass", + "note": "Reads the extent of its result out of its operand, so a node the graph fixes that operand for folds into an initializer and one that does not has a data-dependent shape." + }, + "Conv": { + "disposition": "native-kernel" + }, + "ConvInteger": { + "disposition": "native-kernel" + }, + "ConvTranspose": { + "disposition": "native-kernel" + }, + "Cos": { + "disposition": "native-kernel" + }, + "Cosh": { + "disposition": "native-kernel" + }, + "CumProd": { + "disposition": "native-kernel" + }, + "CumSum": { + "disposition": "native-kernel" + }, + "DFT": { + "disposition": "native-kernel" + }, + "DeformConv": { + "disposition": "native-kernel" + }, + "DepthToSpace": { + "disposition": "native-kernel" + }, + "DequantizeLinear": { + "disposition": "native-kernel" + }, + "Det": { + "disposition": "native-kernel" + }, + "Div": { + "disposition": "native-kernel" + }, + "Dropout": { + "disposition": "native-kernel" + }, + "DynamicQuantizeLinear": { + "disposition": "function-expansion" + }, + "Einsum": { + "disposition": "native-kernel" + }, + "Elu": { + "disposition": "native-kernel" + }, + "Equal": { + "disposition": "native-kernel" + }, + "Erf": { + "disposition": "native-kernel" + }, + "Exp": { + "disposition": "native-kernel" + }, + "Expand": { + "disposition": "native-kernel" + }, + "EyeLike": { + "disposition": "native-kernel" + }, + "Flatten": { + "disposition": "native-kernel" + }, + "Floor": { + "disposition": "native-kernel" + }, + "GRU": { + "disposition": "native-kernel" + }, + "Gather": { + "disposition": "native-kernel" + }, + "GatherElements": { + "disposition": "native-kernel" + }, + "GatherND": { + "disposition": "native-kernel" + }, + "Gelu": { + "disposition": "native-kernel" + }, + "Gemm": { + "disposition": "native-kernel" + }, + "GlobalAveragePool": { + "disposition": "native-kernel" + }, + "GlobalLpPool": { + "disposition": "native-kernel" + }, + "GlobalMaxPool": { + "disposition": "native-kernel" + }, + "Greater": { + "disposition": "native-kernel" + }, + "GreaterOrEqual": { + "disposition": "native-kernel" + }, + "GridSample": { + "disposition": "native-kernel" + }, + "GroupNormalization": { + "disposition": "native-kernel" + }, + "HammingWindow": { + "disposition": "folding-or-graph-pass", + "note": "As `BlackmanWindow`: the whole window follows from the `size` operand the folding pass resolves." + }, + "HannWindow": { + "disposition": "folding-or-graph-pass", + "note": "As `BlackmanWindow`: the whole window follows from the `size` operand the folding pass resolves." + }, + "HardSigmoid": { + "disposition": "native-kernel" + }, + "HardSwish": { + "disposition": "function-expansion" + }, + "Hardmax": { + "disposition": "native-kernel" + }, + "Identity": { + "disposition": "native-kernel" + }, + "If": { + "disposition": "unsupported", + "reason": "control-flow", + "note": "Control flow the compiler serves only where constant folding resolves it away; a branch that survives folding is one whose condition is a run-time value, and the two branches need not agree on the shapes they produce." + }, + "ImageDecoder": { + "disposition": "unsupported", + "reason": "external-codec", + "note": "Decodes an encoded image, which needs an image codec; the generated artifact depends on libm at most, and the extent of the decoded image is a function of the bytes it is handed." + }, + "InstanceNormalization": { + "disposition": "native-kernel" + }, + "IsInf": { + "disposition": "native-kernel" + }, + "IsNaN": { + "disposition": "native-kernel" + }, + "LRN": { + "disposition": "native-kernel" + }, + "LSTM": { + "disposition": "native-kernel" + }, + "LayerNormalization": { + "disposition": "native-kernel" + }, + "LeakyRelu": { + "disposition": "native-kernel" + }, + "Less": { + "disposition": "native-kernel" + }, + "LessOrEqual": { + "disposition": "native-kernel" + }, + "LinearAttention": { + "disposition": "native-kernel" + }, + "Log": { + "disposition": "native-kernel" + }, + "LogSoftmax": { + "disposition": "native-kernel" + }, + "Loop": { + "disposition": "unsupported", + "reason": "control-flow", + "note": "Control flow the compiler serves only where constant folding resolves it away; a loop that survives folding has a trip count no static buffer can be sized for." + }, + "LpNormalization": { + "disposition": "native-kernel" + }, + "LpPool": { + "disposition": "native-kernel" + }, + "MatMul": { + "disposition": "native-kernel" + }, + "MatMulInteger": { + "disposition": "native-kernel" + }, + "Max": { + "disposition": "native-kernel" + }, + "MaxPool": { + "disposition": "native-kernel" + }, + "MaxRoiPool": { + "disposition": "native-kernel" + }, + "MaxUnpool": { + "disposition": "native-kernel" + }, + "Mean": { + "disposition": "native-kernel" + }, + "MeanVarianceNormalization": { + "disposition": "native-kernel" + }, + "MelWeightMatrix": { + "disposition": "folding-or-graph-pass", + "note": "Both extents of the matrix come from operands, so the folding pass resolves the nodes whose operands the graph fixes and the rest have a data-dependent shape." + }, + "Min": { + "disposition": "native-kernel" + }, + "Mish": { + "disposition": "function-expansion" + }, + "Mod": { + "disposition": "native-kernel" + }, + "Mul": { + "disposition": "native-kernel" + }, + "Multinomial": { + "disposition": "unsupported", + "reason": "random-op", + "note": "Draws a categorical sample per row, so its result is not a function of its inputs." + }, + "Neg": { + "disposition": "native-kernel" + }, + "NegativeLogLikelihoodLoss": { + "disposition": "function-expansion" + }, + "NonMaxSuppression": { + "disposition": "unsupported", + "reason": "data-dependent-shape", + "note": "Returns the boxes that survive suppression, so the extent of its result is a function of input values." + }, + "NonZero": { + "disposition": "unsupported", + "reason": "data-dependent-shape", + "note": "Reports the coordinates of the non-zero elements, so the extent of its result is a function of input values." + }, + "Not": { + "disposition": "native-kernel" + }, + "OneHot": { + "disposition": "native-kernel" + }, + "Optional": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Produces an optional value, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "OptionalGetElement": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Unwraps an optional or a sequence, neither of which the compiler holds; the pass-through of a plain tensor that ONNX also allows at opset 18 is the degenerate case of an op whose purpose the compiler does not serve, and it has no kernel of its own." + }, + "OptionalHasElement": { + "disposition": "folding-or-graph-pass", + "note": "Whether an optional value is present is a compile-time property here: the node ONNX writes with no operand at all folds to `false`, and an operand of optional type is refused by the type check as non-tensor I/O." + }, + "Or": { + "disposition": "native-kernel" + }, + "PRelu": { + "disposition": "native-kernel" + }, + "Pad": { + "disposition": "native-kernel" + }, + "Pow": { + "disposition": "native-kernel" + }, + "QLinearConv": { + "disposition": "native-kernel" + }, + "QLinearMatMul": { + "disposition": "native-kernel" + }, + "QuantizeLinear": { + "disposition": "native-kernel" + }, + "RMSNormalization": { + "disposition": "native-kernel" + }, + "RNN": { + "disposition": "native-kernel" + }, + "RandomNormal": { + "disposition": "unsupported", + "reason": "random-op", + "note": "Draws from a normal distribution, so its result is not a function of its inputs." + }, + "RandomNormalLike": { + "disposition": "unsupported", + "reason": "random-op", + "note": "Draws from a normal distribution, so its result is not a function of its inputs." + }, + "RandomUniform": { + "disposition": "unsupported", + "reason": "random-op", + "note": "Draws from a uniform distribution, so its result is not a function of its inputs." + }, + "RandomUniformLike": { + "disposition": "unsupported", + "reason": "random-op", + "note": "Draws from a uniform distribution, so its result is not a function of its inputs." + }, + "Range": { + "disposition": "folding-or-graph-pass", + "note": "Reads its result, extent and values alike, out of three operands, so a node the graph fixes them for folds into an initializer. Below opset 27 the folding pass declines -- ONNX revised the op there, and the reference evaluator it folds through cannot be vouched for at the older revision -- and such a node is refused rather than served with semantics nothing can confirm." + }, + "Reciprocal": { + "disposition": "native-kernel" + }, + "ReduceL1": { + "disposition": "native-kernel" + }, + "ReduceL2": { + "disposition": "native-kernel" + }, + "ReduceLogSum": { + "disposition": "native-kernel" + }, + "ReduceLogSumExp": { + "disposition": "native-kernel" + }, + "ReduceMax": { + "disposition": "native-kernel" + }, + "ReduceMean": { + "disposition": "native-kernel" + }, + "ReduceMin": { + "disposition": "native-kernel" + }, + "ReduceProd": { + "disposition": "native-kernel" + }, + "ReduceSum": { + "disposition": "native-kernel" + }, + "ReduceSumSquare": { + "disposition": "native-kernel" + }, + "RegexFullMatch": { + "disposition": "unsupported", + "reason": "runtime-strings", + "note": "Matches a run-time string tensor against a pattern; the compiler holds strings only as the compile-time constants ONNX-ML attribute tables carry." + }, + "Relu": { + "disposition": "native-kernel" + }, + "Reshape": { + "disposition": "native-kernel" + }, + "Resize": { + "disposition": "native-kernel" + }, + "ReverseSequence": { + "disposition": "native-kernel" + }, + "RoiAlign": { + "disposition": "native-kernel" + }, + "RotaryEmbedding": { + "disposition": "native-kernel" + }, + "Round": { + "disposition": "native-kernel" + }, + "STFT": { + "disposition": "native-kernel" + }, + "Scan": { + "disposition": "unsupported", + "reason": "control-flow", + "note": "Control flow the compiler serves only where constant folding resolves it away; a scan that survives folding carries state across a run-time number of iterations." + }, + "Scatter": { + "disposition": "native-kernel" + }, + "ScatterElements": { + "disposition": "native-kernel" + }, + "ScatterND": { + "disposition": "native-kernel" + }, + "Selu": { + "disposition": "native-kernel" + }, + "SequenceAt": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Reads a sequence, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "SequenceConstruct": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Produces a sequence, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "SequenceEmpty": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Produces a sequence, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "SequenceErase": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Produces a sequence, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "SequenceInsert": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Produces a sequence, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "SequenceLength": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Reads a sequence, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "SequenceMap": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Maps a subgraph over a sequence: its operands are sequences, and the body runs a run-time number of times." + }, + "Shape": { + "disposition": "folding-or-graph-pass", + "note": "A function of its operand's shape alone, which the folding pass resolves from a static shape without needing the operand's values." + }, + "Shrink": { + "disposition": "native-kernel" + }, + "Sigmoid": { + "disposition": "native-kernel" + }, + "Sign": { + "disposition": "native-kernel" + }, + "Sin": { + "disposition": "native-kernel" + }, + "Sinh": { + "disposition": "native-kernel" + }, + "Size": { + "disposition": "folding-or-graph-pass", + "note": "As `Shape`: resolved by the folding pass from the operand's static shape." + }, + "Slice": { + "disposition": "native-kernel" + }, + "Softmax": { + "disposition": "native-kernel" + }, + "SoftmaxCrossEntropyLoss": { + "disposition": "native-kernel" + }, + "Softplus": { + "disposition": "native-kernel" + }, + "Softsign": { + "disposition": "native-kernel" + }, + "SpaceToDepth": { + "disposition": "native-kernel" + }, + "Split": { + "disposition": "native-kernel" + }, + "SplitToSequence": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Produces a sequence, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "Sqrt": { + "disposition": "native-kernel" + }, + "Squeeze": { + "disposition": "native-kernel" + }, + "StringConcat": { + "disposition": "unsupported", + "reason": "runtime-strings", + "note": "Concatenates run-time string tensors, whose elements have no fixed width to hold a buffer of." + }, + "StringNormalizer": { + "disposition": "unsupported", + "reason": "runtime-strings", + "note": "Case-folds and filters a run-time string tensor, and drops the elements it filters out, which also makes the extent of its result a function of its values." + }, + "StringSplit": { + "disposition": "unsupported", + "reason": "runtime-strings", + "note": "Splits a run-time string tensor, whose elements have no fixed width to hold a buffer of." + }, + "Sub": { + "disposition": "native-kernel" + }, + "Sum": { + "disposition": "native-kernel" + }, + "Swish": { + "disposition": "function-expansion" + }, + "Tan": { + "disposition": "native-kernel" + }, + "Tanh": { + "disposition": "native-kernel" + }, + "TensorScatter": { + "disposition": "native-kernel" + }, + "TfIdfVectorizer": { + "disposition": "native-kernel" + }, + "ThresholdedRelu": { + "disposition": "native-kernel" + }, + "Tile": { + "disposition": "native-kernel" + }, + "TopK": { + "disposition": "native-kernel" + }, + "Transpose": { + "disposition": "native-kernel" + }, + "Trilu": { + "disposition": "native-kernel" + }, + "Unique": { + "disposition": "unsupported", + "reason": "data-dependent-shape", + "note": "Returns the distinct elements, so the extent of its result is a function of input values." + }, + "Unsqueeze": { + "disposition": "native-kernel" + }, + "Upsample": { + "disposition": "native-kernel" + }, + "Where": { + "disposition": "native-kernel" + }, + "Xor": { + "disposition": "native-kernel" + } + }, + "ai.onnx.ml": { + "ArrayFeatureExtractor": { + "disposition": "native-kernel" + }, + "Binarizer": { + "disposition": "native-kernel" + }, + "CastMap": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Reads a map, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "CategoryMapper": { + "disposition": "unsupported", + "reason": "runtime-strings", + "note": "Maps integer categories to string ones and back, so one side of it is a run-time string tensor whichever direction a node runs in." + }, + "DictVectorizer": { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "Reads a map, which is not a tensor and has no buffer the entrypoint contract can hand a caller." + }, + "FeatureVectorizer": { + "disposition": "native-kernel" + }, + "Imputer": { + "disposition": "native-kernel" + }, + "LabelEncoder": { + "disposition": "native-kernel" + }, + "LinearClassifier": { + "disposition": "native-kernel" + }, + "LinearRegressor": { + "disposition": "native-kernel" + }, + "Normalizer": { + "disposition": "native-kernel" + }, + "OneHotEncoder": { + "disposition": "native-kernel" + }, + "SVMClassifier": { + "disposition": "native-kernel" + }, + "SVMRegressor": { + "disposition": "native-kernel" + }, + "Scaler": { + "disposition": "native-kernel" + }, + "TreeEnsemble": { + "disposition": "native-kernel" + }, + "TreeEnsembleClassifier": { + "disposition": "native-kernel" + }, + "TreeEnsembleRegressor": { + "disposition": "native-kernel" + }, + "ZipMap": { + "disposition": "folding-or-graph-pass", + "note": "Removed by the ZipMap graph pass, which promotes the probability tensor it reads to the graph output in its place and publishes its class labels as header metadata." + } + } +} diff --git a/src/python/tests/conformance/ledger.json b/src/python/tests/conformance/ledger.json new file mode 100644 index 0000000..b0f083a --- /dev/null +++ b/src/python/tests/conformance/ledger.json @@ -0,0 +1,721 @@ +{ + "test_adagrad": "out-of-scope-domain", + "test_adagrad_multiple": "out-of-scope-domain", + "test_adam": "out-of-scope-domain", + "test_adam_multiple": "out-of-scope-domain", + "test_affine_grid_2d": "data-dependent-shape", + "test_affine_grid_2d_align_corners": "data-dependent-shape", + "test_affine_grid_2d_align_corners_expanded": "control-flow", + "test_affine_grid_2d_expanded": "control-flow", + "test_affine_grid_3d": "data-dependent-shape", + "test_affine_grid_3d_align_corners": "data-dependent-shape", + "test_affine_grid_3d_align_corners_expanded": "control-flow", + "test_affine_grid_3d_expanded": "control-flow", + "test_ai_onnx_ml_label_encoder_string_int": "runtime-strings", + "test_ai_onnx_ml_label_encoder_string_int_no_default": "runtime-strings", + "test_ai_onnx_ml_label_encoder_tensor_mapping": "runtime-strings", + "test_ai_onnx_ml_label_encoder_tensor_value_only_mapping": "runtime-strings", + "test_attention_3d_attn_mask_expanded": "data-dependent-shape", + "test_attention_3d_causal_expanded": "data-dependent-shape", + "test_attention_3d_diff_heads_sizes_attn_mask_expanded": "data-dependent-shape", + "test_attention_3d_diff_heads_sizes_causal_expanded": "data-dependent-shape", + "test_attention_3d_diff_heads_sizes_expanded": "data-dependent-shape", + "test_attention_3d_diff_heads_sizes_scaled_expanded": "data-dependent-shape", + "test_attention_3d_diff_heads_sizes_softcap_expanded": "data-dependent-shape", + "test_attention_3d_diff_heads_with_past_and_present_expanded": "data-dependent-shape", + "test_attention_3d_expanded": "data-dependent-shape", + "test_attention_3d_gqa_attn_mask_expanded": "data-dependent-shape", + "test_attention_3d_gqa_causal_expanded": "data-dependent-shape", + "test_attention_3d_gqa_expanded": "data-dependent-shape", + "test_attention_3d_gqa_scaled_expanded": "data-dependent-shape", + "test_attention_3d_gqa_softcap_expanded": "data-dependent-shape", + "test_attention_3d_gqa_with_past_and_present_expanded": "data-dependent-shape", + "test_attention_3d_scaled_expanded": "data-dependent-shape", + "test_attention_3d_softcap_expanded": "data-dependent-shape", + "test_attention_3d_transpose_verification_expanded": "data-dependent-shape", + "test_attention_3d_with_past_and_present_expanded": "data-dependent-shape", + "test_attention_3d_with_past_and_present_qk_matmul_bias_expanded": "data-dependent-shape", + "test_attention_3d_with_past_and_present_qk_matmul_expanded": "data-dependent-shape", + "test_attention_3d_with_past_and_present_qk_matmul_softcap_expanded": "data-dependent-shape", + "test_attention_3d_with_past_and_present_qk_matmul_softmax_expanded": "data-dependent-shape", + "test_attention_4d_attn_mask_3d_causal_expanded": "data-dependent-shape", + "test_attention_4d_attn_mask_3d_expanded": "data-dependent-shape", + "test_attention_4d_attn_mask_4d_causal_expanded": "data-dependent-shape", + "test_attention_4d_attn_mask_4d_expanded": "data-dependent-shape", + "test_attention_4d_attn_mask_bool_4d_expanded": "data-dependent-shape", + "test_attention_4d_attn_mask_bool_expanded": "data-dependent-shape", + "test_attention_4d_attn_mask_expanded": "data-dependent-shape", + "test_attention_4d_causal_expanded": "data-dependent-shape", + "test_attention_4d_diff_heads_mask4d_padded_kv_expanded": "op-not-implemented", + "test_attention_4d_diff_heads_sizes_attn_mask_expanded": "data-dependent-shape", + "test_attention_4d_diff_heads_sizes_causal_expanded": "data-dependent-shape", + "test_attention_4d_diff_heads_sizes_expanded": "data-dependent-shape", + "test_attention_4d_diff_heads_sizes_scaled_expanded": "data-dependent-shape", + "test_attention_4d_diff_heads_sizes_softcap_expanded": "data-dependent-shape", + "test_attention_4d_diff_heads_with_past_and_present_expanded": "data-dependent-shape", + "test_attention_4d_diff_heads_with_past_and_present_mask3d_expanded": "data-dependent-shape", + "test_attention_4d_diff_heads_with_past_and_present_mask4d_expanded": "data-dependent-shape", + "test_attention_4d_expanded": "data-dependent-shape", + "test_attention_4d_fp16": "unsupported-dtype", + "test_attention_4d_fp16_expanded": "unsupported-dtype", + "test_attention_4d_gqa_attn_mask_expanded": "data-dependent-shape", + "test_attention_4d_gqa_causal_expanded": "data-dependent-shape", + "test_attention_4d_gqa_expanded": "data-dependent-shape", + "test_attention_4d_gqa_scaled_expanded": "data-dependent-shape", + "test_attention_4d_gqa_softcap_expanded": "data-dependent-shape", + "test_attention_4d_gqa_with_past_and_present_expanded": "data-dependent-shape", + "test_attention_4d_gqa_with_past_and_present_fp16": "unsupported-dtype", + "test_attention_4d_gqa_with_past_and_present_fp16_expanded": "unsupported-dtype", + "test_attention_4d_scaled_expanded": "data-dependent-shape", + "test_attention_4d_softcap_expanded": "data-dependent-shape", + "test_attention_4d_softcap_neginf_mask_expanded": "data-dependent-shape", + "test_attention_4d_softcap_neginf_mask_poison_expanded": "data-dependent-shape", + "test_attention_4d_with_past_and_present_expanded": "data-dependent-shape", + "test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask_causal_expanded": "data-dependent-shape", + "test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask_expanded": "data-dependent-shape", + "test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask_causal_expanded": "data-dependent-shape", + "test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask_expanded": "data-dependent-shape", + "test_attention_4d_with_past_and_present_qk_matmul_bias_expanded": "data-dependent-shape", + "test_attention_4d_with_past_and_present_qk_matmul_expanded": "data-dependent-shape", + "test_attention_4d_with_qk_matmul_bias_expanded": "data-dependent-shape", + "test_attention_4d_with_qk_matmul_expanded": "data-dependent-shape", + "test_attention_4d_with_qk_matmul_softcap_expanded": "data-dependent-shape", + "test_attention_4d_with_qk_matmul_softmax_expanded": "data-dependent-shape", + "test_bernoulli": "random-op", + "test_bernoulli_double": "random-op", + "test_bernoulli_double_expanded": "random-op", + "test_bernoulli_expanded": "random-op", + "test_bernoulli_seed": "random-op", + "test_bernoulli_seed_expanded": "random-op", + "test_blackmanwindow": "data-dependent-shape", + "test_blackmanwindow_expanded": "data-dependent-shape", + "test_blackmanwindow_symmetric": "data-dependent-shape", + "test_blackmanwindow_symmetric_expanded": "data-dependent-shape", + "test_cast_BFLOAT16_to_FLOAT": "unsupported-dtype", + "test_cast_DOUBLE_to_FLOAT16": "unsupported-dtype", + "test_cast_FLOAT16_to_DOUBLE": "unsupported-dtype", + "test_cast_FLOAT16_to_FLOAT": "unsupported-dtype", + "test_cast_FLOAT16_to_FLOAT4E2M1": "unsupported-dtype", + "test_cast_FLOAT16_to_FLOAT8E4M3FN": "unsupported-dtype", + "test_cast_FLOAT16_to_FLOAT8E4M3FNUZ": "unsupported-dtype", + "test_cast_FLOAT16_to_FLOAT8E5M2": "unsupported-dtype", + "test_cast_FLOAT16_to_FLOAT8E5M2FNUZ": "unsupported-dtype", + "test_cast_FLOAT16_to_INT2": "unsupported-dtype", + "test_cast_FLOAT16_to_INT4": "unsupported-dtype", + "test_cast_FLOAT16_to_UINT2": "unsupported-dtype", + "test_cast_FLOAT16_to_UINT4": "unsupported-dtype", + "test_cast_FLOAT4E2M1_to_FLOAT": "unsupported-dtype", + "test_cast_FLOAT4E2M1_to_FLOAT16": "unsupported-dtype", + "test_cast_FLOAT8E4M3FNUZ_to_FLOAT": "unsupported-dtype", + "test_cast_FLOAT8E4M3FNUZ_to_FLOAT16": "unsupported-dtype", + "test_cast_FLOAT8E4M3FN_to_FLOAT": "unsupported-dtype", + "test_cast_FLOAT8E4M3FN_to_FLOAT16": "unsupported-dtype", + "test_cast_FLOAT8E5M2FNUZ_to_FLOAT": "unsupported-dtype", + "test_cast_FLOAT8E5M2FNUZ_to_FLOAT16": "unsupported-dtype", + "test_cast_FLOAT8E5M2_to_FLOAT": "unsupported-dtype", + "test_cast_FLOAT8E5M2_to_FLOAT16": "unsupported-dtype", + "test_cast_FLOAT_to_BFLOAT16": "unsupported-dtype", + "test_cast_FLOAT_to_FLOAT16": "unsupported-dtype", + "test_cast_FLOAT_to_FLOAT4E2M1": "unsupported-dtype", + "test_cast_FLOAT_to_FLOAT8E4M3FN": "unsupported-dtype", + "test_cast_FLOAT_to_FLOAT8E4M3FNUZ": "unsupported-dtype", + "test_cast_FLOAT_to_FLOAT8E5M2": "unsupported-dtype", + "test_cast_FLOAT_to_FLOAT8E5M2FNUZ": "unsupported-dtype", + "test_cast_FLOAT_to_INT2": "unsupported-dtype", + "test_cast_FLOAT_to_INT4": "unsupported-dtype", + "test_cast_FLOAT_to_UINT2": "unsupported-dtype", + "test_cast_FLOAT_to_UINT4": "unsupported-dtype", + "test_cast_INT2_to_FLOAT": "unsupported-dtype", + "test_cast_INT2_to_FLOAT16": "unsupported-dtype", + "test_cast_INT2_to_INT8": "unsupported-dtype", + "test_cast_INT4_to_FLOAT": "unsupported-dtype", + "test_cast_INT4_to_FLOAT16": "unsupported-dtype", + "test_cast_INT4_to_INT8": "unsupported-dtype", + "test_cast_UINT2_to_FLOAT": "unsupported-dtype", + "test_cast_UINT2_to_FLOAT16": "unsupported-dtype", + "test_cast_UINT2_to_UINT8": "unsupported-dtype", + "test_cast_UINT4_to_FLOAT": "unsupported-dtype", + "test_cast_UINT4_to_FLOAT16": "unsupported-dtype", + "test_cast_UINT4_to_UINT8": "unsupported-dtype", + "test_cast_e8m0_FLOAT16_to_FLOAT8E8M0": "unsupported-dtype", + "test_cast_e8m0_FLOAT8E8M0_to_FLOAT": "unsupported-dtype", + "test_cast_e8m0_FLOAT8E8M0_to_FLOAT16": "unsupported-dtype", + "test_cast_e8m0_FLOAT_to_FLOAT8E8M0": "unsupported-dtype", + "test_cast_no_saturate_FLOAT16_to_FLOAT8E4M3FN": "unsupported-dtype", + "test_cast_no_saturate_FLOAT16_to_FLOAT8E4M3FNUZ": "unsupported-dtype", + "test_cast_no_saturate_FLOAT16_to_FLOAT8E5M2": "unsupported-dtype", + "test_cast_no_saturate_FLOAT16_to_FLOAT8E5M2FNUZ": "unsupported-dtype", + "test_cast_no_saturate_FLOAT_to_FLOAT8E4M3FN": "unsupported-dtype", + "test_cast_no_saturate_FLOAT_to_FLOAT8E4M3FNUZ": "unsupported-dtype", + "test_cast_no_saturate_FLOAT_to_FLOAT8E5M2": "unsupported-dtype", + "test_cast_no_saturate_FLOAT_to_FLOAT8E5M2FNUZ": "unsupported-dtype", + "test_castlike_BFLOAT16_to_FLOAT": "unsupported-dtype", + "test_castlike_BFLOAT16_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_DOUBLE_to_FLOAT16": "unsupported-dtype", + "test_castlike_DOUBLE_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_DOUBLE": "unsupported-dtype", + "test_castlike_FLOAT16_to_DOUBLE_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT4E2M1": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT4E2M1_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT8E4M3FN": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT8E4M3FNUZ": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT8E4M3FNUZ_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT8E4M3FN_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT8E5M2": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT8E5M2FNUZ": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT8E5M2FNUZ_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT8E5M2_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_INT2": "unsupported-dtype", + "test_castlike_FLOAT16_to_INT2_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_INT4": "unsupported-dtype", + "test_castlike_FLOAT16_to_INT4_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_UINT2": "unsupported-dtype", + "test_castlike_FLOAT16_to_UINT2_expanded": "unsupported-dtype", + "test_castlike_FLOAT16_to_UINT4": "unsupported-dtype", + "test_castlike_FLOAT16_to_UINT4_expanded": "unsupported-dtype", + "test_castlike_FLOAT4E2M1_to_FLOAT": "unsupported-dtype", + "test_castlike_FLOAT4E2M1_to_FLOAT16": "unsupported-dtype", + "test_castlike_FLOAT4E2M1_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_FLOAT4E2M1_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_FLOAT8E4M3FNUZ_to_FLOAT": "unsupported-dtype", + "test_castlike_FLOAT8E4M3FNUZ_to_FLOAT16": "unsupported-dtype", + "test_castlike_FLOAT8E4M3FNUZ_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_FLOAT8E4M3FNUZ_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_FLOAT8E4M3FN_to_FLOAT": "unsupported-dtype", + "test_castlike_FLOAT8E4M3FN_to_FLOAT16": "unsupported-dtype", + "test_castlike_FLOAT8E4M3FN_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_FLOAT8E4M3FN_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_FLOAT8E5M2FNUZ_to_FLOAT": "unsupported-dtype", + "test_castlike_FLOAT8E5M2FNUZ_to_FLOAT16": "unsupported-dtype", + "test_castlike_FLOAT8E5M2FNUZ_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_FLOAT8E5M2FNUZ_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_FLOAT8E5M2_to_FLOAT": "unsupported-dtype", + "test_castlike_FLOAT8E5M2_to_FLOAT16": "unsupported-dtype", + "test_castlike_FLOAT8E5M2_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_FLOAT8E5M2_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_BFLOAT16": "unsupported-dtype", + "test_castlike_FLOAT_to_BFLOAT16_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT16": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT4E2M1": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT4E2M1_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT8E4M3FN": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT8E4M3FNUZ": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT8E4M3FNUZ_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT8E4M3FN_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT8E5M2": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT8E5M2FNUZ": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT8E5M2FNUZ_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_FLOAT8E5M2_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_INT2": "unsupported-dtype", + "test_castlike_FLOAT_to_INT2_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_INT4": "unsupported-dtype", + "test_castlike_FLOAT_to_INT4_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_UINT2": "unsupported-dtype", + "test_castlike_FLOAT_to_UINT2_expanded": "unsupported-dtype", + "test_castlike_FLOAT_to_UINT4": "unsupported-dtype", + "test_castlike_FLOAT_to_UINT4_expanded": "unsupported-dtype", + "test_castlike_INT2_to_FLOAT": "unsupported-dtype", + "test_castlike_INT2_to_FLOAT16": "unsupported-dtype", + "test_castlike_INT2_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_INT2_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_INT2_to_INT8": "unsupported-dtype", + "test_castlike_INT2_to_INT8_expanded": "unsupported-dtype", + "test_castlike_INT4_to_FLOAT": "unsupported-dtype", + "test_castlike_INT4_to_FLOAT16": "unsupported-dtype", + "test_castlike_INT4_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_INT4_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_INT4_to_INT8": "unsupported-dtype", + "test_castlike_INT4_to_INT8_expanded": "unsupported-dtype", + "test_castlike_UINT2_to_FLOAT": "unsupported-dtype", + "test_castlike_UINT2_to_FLOAT16": "unsupported-dtype", + "test_castlike_UINT2_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_UINT2_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_UINT2_to_UINT8": "unsupported-dtype", + "test_castlike_UINT2_to_UINT8_expanded": "unsupported-dtype", + "test_castlike_UINT4_to_FLOAT": "unsupported-dtype", + "test_castlike_UINT4_to_FLOAT16": "unsupported-dtype", + "test_castlike_UINT4_to_FLOAT16_expanded": "unsupported-dtype", + "test_castlike_UINT4_to_FLOAT_expanded": "unsupported-dtype", + "test_castlike_UINT4_to_UINT8": "unsupported-dtype", + "test_castlike_UINT4_to_UINT8_expanded": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT16_to_FLOAT8E4M3FN": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT16_to_FLOAT8E4M3FNUZ": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT16_to_FLOAT8E4M3FNUZ_expanded": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT16_to_FLOAT8E4M3FN_expanded": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT16_to_FLOAT8E5M2": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT16_to_FLOAT8E5M2FNUZ": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT16_to_FLOAT8E5M2FNUZ_expanded": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT16_to_FLOAT8E5M2_expanded": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT_to_FLOAT8E4M3FN": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT_to_FLOAT8E4M3FNUZ": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT_to_FLOAT8E4M3FNUZ_expanded": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT_to_FLOAT8E4M3FN_expanded": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT_to_FLOAT8E5M2": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT_to_FLOAT8E5M2FNUZ": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT_to_FLOAT8E5M2FNUZ_expanded": "unsupported-dtype", + "test_castlike_no_saturate_FLOAT_to_FLOAT8E5M2_expanded": "unsupported-dtype", + "test_causal_conv_with_state_fp16": "unsupported-dtype", + "test_causal_conv_with_state_fp16_expanded": "unsupported-dtype", + "test_causal_conv_with_state_silu_fp16": "unsupported-dtype", + "test_causal_conv_with_state_silu_fp16_expanded": "unsupported-dtype", + "test_center_crop_pad_crop": "data-dependent-shape", + "test_center_crop_pad_crop_and_pad": "data-dependent-shape", + "test_center_crop_pad_crop_and_pad_expanded": "data-dependent-shape", + "test_center_crop_pad_crop_axes_chw": "data-dependent-shape", + "test_center_crop_pad_crop_axes_chw_expanded": "data-dependent-shape", + "test_center_crop_pad_crop_axes_hwc": "data-dependent-shape", + "test_center_crop_pad_crop_axes_hwc_expanded": "data-dependent-shape", + "test_center_crop_pad_crop_expanded": "data-dependent-shape", + "test_center_crop_pad_crop_negative_axes_hwc": "data-dependent-shape", + "test_center_crop_pad_crop_negative_axes_hwc_expanded": "data-dependent-shape", + "test_center_crop_pad_pad": "data-dependent-shape", + "test_center_crop_pad_pad_expanded": "data-dependent-shape", + "test_col2im": "data-dependent-shape", + "test_col2im_5d": "data-dependent-shape", + "test_col2im_dilations": "data-dependent-shape", + "test_col2im_pads": "data-dependent-shape", + "test_col2im_strides": "data-dependent-shape", + "test_compress_0": "data-dependent-shape", + "test_compress_1": "data-dependent-shape", + "test_compress_default_axis": "data-dependent-shape", + "test_compress_negative_axis": "data-dependent-shape", + "test_constant_pad": "data-dependent-shape", + "test_constant_pad_axes": "data-dependent-shape", + "test_constant_pad_negative_axes": "data-dependent-shape", + "test_constantofshape_float_ones": "data-dependent-shape", + "test_constantofshape_int_shape_zero": "data-dependent-shape", + "test_constantofshape_int_zeros": "data-dependent-shape", + "test_dequantizelinear_e4m3fn": "unsupported-dtype", + "test_dequantizelinear_e4m3fn_float16": "unsupported-dtype", + "test_dequantizelinear_e4m3fn_zero_point": "unsupported-dtype", + "test_dequantizelinear_e5m2": "unsupported-dtype", + "test_dequantizelinear_float4e2m1": "unsupported-dtype", + "test_dequantizelinear_int2": "unsupported-dtype", + "test_dequantizelinear_int4": "unsupported-dtype", + "test_dequantizelinear_uint2": "unsupported-dtype", + "test_dequantizelinear_uint4": "unsupported-dtype", + "test_dft_irfft": "data-dependent-shape", + "test_dft_rfft": "data-dependent-shape", + "test_edge_pad": "data-dependent-shape", + "test_equal_string": "runtime-strings", + "test_equal_string_broadcast": "runtime-strings", + "test_expand_dim_changed": "data-dependent-shape", + "test_expand_dim_unchanged": "data-dependent-shape", + "test_flexattention": "out-of-scope-domain", + "test_flexattention_causal_mask": "out-of-scope-domain", + "test_flexattention_causal_mask_expanded_ver26": "out-of-scope-domain", + "test_flexattention_diff_head_sizes": "out-of-scope-domain", + "test_flexattention_diff_head_sizes_expanded_ver26": "out-of-scope-domain", + "test_flexattention_double": "out-of-scope-domain", + "test_flexattention_double_expanded_ver26": "out-of-scope-domain", + "test_flexattention_expanded_ver26": "out-of-scope-domain", + "test_flexattention_fp16": "out-of-scope-domain", + "test_flexattention_fp16_expanded_ver26": "out-of-scope-domain", + "test_flexattention_gqa": "out-of-scope-domain", + "test_flexattention_gqa_expanded_ver26": "out-of-scope-domain", + "test_flexattention_prob_mod": "out-of-scope-domain", + "test_flexattention_prob_mod_expanded_ver26": "out-of-scope-domain", + "test_flexattention_relative_positional": "out-of-scope-domain", + "test_flexattention_relative_positional_expanded_ver26": "out-of-scope-domain", + "test_flexattention_scaled": "out-of-scope-domain", + "test_flexattention_scaled_expanded_ver26": "out-of-scope-domain", + "test_flexattention_score_mod": "out-of-scope-domain", + "test_flexattention_score_mod_expanded_ver26": "out-of-scope-domain", + "test_flexattention_soft_cap": "out-of-scope-domain", + "test_flexattention_soft_cap_expanded_ver26": "out-of-scope-domain", + "test_hammingwindow": "data-dependent-shape", + "test_hammingwindow_expanded": "data-dependent-shape", + "test_hammingwindow_symmetric": "data-dependent-shape", + "test_hammingwindow_symmetric_expanded": "data-dependent-shape", + "test_hannwindow": "data-dependent-shape", + "test_hannwindow_expanded": "data-dependent-shape", + "test_hannwindow_symmetric": "data-dependent-shape", + "test_hannwindow_symmetric_expanded": "data-dependent-shape", + "test_identity_opt": "non-tensor-io", + "test_identity_sequence": "non-tensor-io", + "test_if": "control-flow", + "test_if_opt": "non-tensor-io", + "test_if_seq": "non-tensor-io", + "test_image_decoder_decode_bmp_rgb": "external-codec", + "test_image_decoder_decode_jpeg2k_rgb": "external-codec", + "test_image_decoder_decode_jpeg_bgr": "external-codec", + "test_image_decoder_decode_jpeg_grayscale": "external-codec", + "test_image_decoder_decode_jpeg_rgb": "external-codec", + "test_image_decoder_decode_png_rgb": "external-codec", + "test_image_decoder_decode_pnm_rgb": "external-codec", + "test_image_decoder_decode_tiff_rgb": "external-codec", + "test_image_decoder_decode_webp_rgb": "external-codec", + "test_isinf_float16": "unsupported-dtype", + "test_isnan_float16": "unsupported-dtype", + "test_layer_normalization_2d_axis0_expanded": "data-dependent-shape", + "test_layer_normalization_2d_axis0_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_2d_axis1_expanded": "data-dependent-shape", + "test_layer_normalization_2d_axis1_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_2d_axis_negative_1_expanded": "data-dependent-shape", + "test_layer_normalization_2d_axis_negative_1_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_2d_axis_negative_2_expanded": "data-dependent-shape", + "test_layer_normalization_2d_axis_negative_2_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_3d_axis0_epsilon_expanded": "data-dependent-shape", + "test_layer_normalization_3d_axis0_epsilon_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_3d_axis1_epsilon_expanded": "data-dependent-shape", + "test_layer_normalization_3d_axis1_epsilon_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_3d_axis2_epsilon_expanded": "data-dependent-shape", + "test_layer_normalization_3d_axis2_epsilon_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_3d_axis_negative_1_epsilon_expanded": "data-dependent-shape", + "test_layer_normalization_3d_axis_negative_1_epsilon_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_3d_axis_negative_2_epsilon_expanded": "data-dependent-shape", + "test_layer_normalization_3d_axis_negative_2_epsilon_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_3d_axis_negative_3_epsilon_expanded": "data-dependent-shape", + "test_layer_normalization_3d_axis_negative_3_epsilon_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_4d_axis0_expanded": "data-dependent-shape", + "test_layer_normalization_4d_axis0_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_4d_axis1_expanded": "data-dependent-shape", + "test_layer_normalization_4d_axis1_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_4d_axis2_expanded": "data-dependent-shape", + "test_layer_normalization_4d_axis2_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_4d_axis3_expanded": "data-dependent-shape", + "test_layer_normalization_4d_axis3_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_4d_axis_negative_1_expanded": "data-dependent-shape", + "test_layer_normalization_4d_axis_negative_1_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_4d_axis_negative_2_expanded": "data-dependent-shape", + "test_layer_normalization_4d_axis_negative_2_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_4d_axis_negative_3_expanded": "data-dependent-shape", + "test_layer_normalization_4d_axis_negative_3_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_4d_axis_negative_4_expanded": "data-dependent-shape", + "test_layer_normalization_4d_axis_negative_4_expanded_ver18": "data-dependent-shape", + "test_layer_normalization_default_axis_expanded": "data-dependent-shape", + "test_layer_normalization_default_axis_expanded_ver18": "data-dependent-shape", + "test_linear_attention_decode_step_expanded": "control-flow", + "test_linear_attention_delta_expanded": "control-flow", + "test_linear_attention_explicit_scale_expanded": "control-flow", + "test_linear_attention_fp16": "unsupported-dtype", + "test_linear_attention_fp16_expanded": "unsupported-dtype", + "test_linear_attention_gated_delta_beta_scalar_expanded": "control-flow", + "test_linear_attention_gated_delta_expanded": "control-flow", + "test_linear_attention_gated_delta_gqa_expanded": "control-flow", + "test_linear_attention_gated_delta_mqa_expanded": "control-flow", + "test_linear_attention_gated_expanded": "control-flow", + "test_linear_attention_gated_per_head_decay_expanded": "control-flow", + "test_linear_attention_linear_expanded": "control-flow", + "test_linear_attention_linear_t1_no_past_expanded": "control-flow", + "test_linear_attention_no_past_explicit_zeros_expanded": "control-flow", + "test_linear_attention_prefill_with_past_expanded": "control-flow", + "test_loop11": "control-flow", + "test_loop13_seq": "non-tensor-io", + "test_loop16_seq_none": "non-tensor-io", + "test_max_float16": "unsupported-dtype", + "test_maxunpool_export_with_output_shape": "data-dependent-shape", + "test_melweightmatrix": "data-dependent-shape", + "test_min_float16": "unsupported-dtype", + "test_mod_mixed_sign_float16": "unsupported-dtype", + "test_momentum": "out-of-scope-domain", + "test_momentum_multiple": "out-of-scope-domain", + "test_nesterov_momentum": "out-of-scope-domain", + "test_nonmaxsuppression_center_point_box_format": "data-dependent-shape", + "test_nonmaxsuppression_flipped_coordinates": "data-dependent-shape", + "test_nonmaxsuppression_identical_boxes": "data-dependent-shape", + "test_nonmaxsuppression_iou_threshold_boundary": "data-dependent-shape", + "test_nonmaxsuppression_limit_output_size": "data-dependent-shape", + "test_nonmaxsuppression_single_box": "data-dependent-shape", + "test_nonmaxsuppression_suppress_by_IOU": "data-dependent-shape", + "test_nonmaxsuppression_suppress_by_IOU_and_scores": "data-dependent-shape", + "test_nonmaxsuppression_two_batches": "data-dependent-shape", + "test_nonmaxsuppression_two_classes": "data-dependent-shape", + "test_nonzero_example": "data-dependent-shape", + "test_onehot_negative_indices": "data-dependent-shape", + "test_onehot_with_axis": "data-dependent-shape", + "test_onehot_with_negative_axis": "data-dependent-shape", + "test_onehot_without_axis": "data-dependent-shape", + "test_optional_get_element_optional_sequence": "non-tensor-io", + "test_optional_get_element_optional_tensor": "non-tensor-io", + "test_optional_get_element_sequence": "non-tensor-io", + "test_optional_get_element_tensor": "op-not-implemented", + "test_optional_has_element_empty_optional_input": "non-tensor-io", + "test_optional_has_element_optional_input": "non-tensor-io", + "test_optional_has_element_tensor_input": "non-tensor-io", + "test_qlinearmatmul_2D_int8_float16": "unsupported-dtype", + "test_qlinearmatmul_2D_uint8_float16": "unsupported-dtype", + "test_qlinearmatmul_3D_int8_float16": "unsupported-dtype", + "test_qlinearmatmul_3D_uint8_float16": "unsupported-dtype", + "test_quantizelinear_e4m3fn": "unsupported-dtype", + "test_quantizelinear_e5m2": "unsupported-dtype", + "test_quantizelinear_float4e2m1": "unsupported-dtype", + "test_quantizelinear_int2": "unsupported-dtype", + "test_quantizelinear_int4": "unsupported-dtype", + "test_quantizelinear_uint2": "unsupported-dtype", + "test_quantizelinear_uint4": "unsupported-dtype", + "test_range_bfloat16_type_positive_delta": "unsupported-dtype", + "test_range_bfloat16_type_positive_delta_expanded": "unsupported-dtype", + "test_range_float16_type_positive_delta": "unsupported-dtype", + "test_range_float16_type_positive_delta_expanded": "unsupported-dtype", + "test_range_float_type_positive_delta": "data-dependent-shape", + "test_range_float_type_positive_delta_expanded": "control-flow", + "test_range_int32_type_negative_delta": "data-dependent-shape", + "test_range_int32_type_negative_delta_expanded": "control-flow", + "test_reduce_l1_do_not_keepdims_example": "data-dependent-shape", + "test_reduce_l1_do_not_keepdims_example_expanded": "data-dependent-shape", + "test_reduce_l1_do_not_keepdims_random": "data-dependent-shape", + "test_reduce_l1_do_not_keepdims_random_expanded": "data-dependent-shape", + "test_reduce_l1_empty_set": "data-dependent-shape", + "test_reduce_l1_empty_set_expanded": "data-dependent-shape", + "test_reduce_l1_keep_dims_example": "data-dependent-shape", + "test_reduce_l1_keep_dims_example_expanded": "data-dependent-shape", + "test_reduce_l1_keep_dims_random": "data-dependent-shape", + "test_reduce_l1_keep_dims_random_expanded": "data-dependent-shape", + "test_reduce_l1_negative_axes_keep_dims_example": "data-dependent-shape", + "test_reduce_l1_negative_axes_keep_dims_example_expanded": "data-dependent-shape", + "test_reduce_l1_negative_axes_keep_dims_random": "data-dependent-shape", + "test_reduce_l1_negative_axes_keep_dims_random_expanded": "data-dependent-shape", + "test_reduce_l2_do_not_keepdims_example": "data-dependent-shape", + "test_reduce_l2_do_not_keepdims_example_expanded": "data-dependent-shape", + "test_reduce_l2_do_not_keepdims_random": "data-dependent-shape", + "test_reduce_l2_do_not_keepdims_random_expanded": "data-dependent-shape", + "test_reduce_l2_empty_set": "data-dependent-shape", + "test_reduce_l2_empty_set_expanded": "data-dependent-shape", + "test_reduce_l2_keep_dims_example": "data-dependent-shape", + "test_reduce_l2_keep_dims_example_expanded": "data-dependent-shape", + "test_reduce_l2_keep_dims_random": "data-dependent-shape", + "test_reduce_l2_keep_dims_random_expanded": "data-dependent-shape", + "test_reduce_l2_negative_axes_keep_dims_example": "data-dependent-shape", + "test_reduce_l2_negative_axes_keep_dims_example_expanded": "data-dependent-shape", + "test_reduce_l2_negative_axes_keep_dims_random": "data-dependent-shape", + "test_reduce_l2_negative_axes_keep_dims_random_expanded": "data-dependent-shape", + "test_reduce_log_sum_asc_axes": "data-dependent-shape", + "test_reduce_log_sum_asc_axes_expanded": "data-dependent-shape", + "test_reduce_log_sum_desc_axes": "data-dependent-shape", + "test_reduce_log_sum_desc_axes_expanded": "data-dependent-shape", + "test_reduce_log_sum_empty_set": "data-dependent-shape", + "test_reduce_log_sum_empty_set_expanded": "data-dependent-shape", + "test_reduce_log_sum_exp_do_not_keepdims_example": "data-dependent-shape", + "test_reduce_log_sum_exp_do_not_keepdims_example_expanded": "data-dependent-shape", + "test_reduce_log_sum_exp_do_not_keepdims_random": "data-dependent-shape", + "test_reduce_log_sum_exp_do_not_keepdims_random_expanded": "data-dependent-shape", + "test_reduce_log_sum_exp_empty_set": "data-dependent-shape", + "test_reduce_log_sum_exp_empty_set_expanded": "data-dependent-shape", + "test_reduce_log_sum_exp_keepdims_example": "data-dependent-shape", + "test_reduce_log_sum_exp_keepdims_example_expanded": "data-dependent-shape", + "test_reduce_log_sum_exp_keepdims_random": "data-dependent-shape", + "test_reduce_log_sum_exp_keepdims_random_expanded": "data-dependent-shape", + "test_reduce_log_sum_exp_negative_axes_keepdims_example": "data-dependent-shape", + "test_reduce_log_sum_exp_negative_axes_keepdims_example_expanded": "data-dependent-shape", + "test_reduce_log_sum_exp_negative_axes_keepdims_random": "data-dependent-shape", + "test_reduce_log_sum_exp_negative_axes_keepdims_random_expanded": "data-dependent-shape", + "test_reduce_log_sum_negative_axes": "data-dependent-shape", + "test_reduce_log_sum_negative_axes_expanded": "data-dependent-shape", + "test_reduce_max_bool_inputs": "data-dependent-shape", + "test_reduce_max_do_not_keepdims_example": "data-dependent-shape", + "test_reduce_max_do_not_keepdims_random": "data-dependent-shape", + "test_reduce_max_empty_set": "data-dependent-shape", + "test_reduce_max_keepdims_example": "data-dependent-shape", + "test_reduce_max_keepdims_random": "data-dependent-shape", + "test_reduce_max_negative_axes_keepdims_example": "data-dependent-shape", + "test_reduce_max_negative_axes_keepdims_random": "data-dependent-shape", + "test_reduce_mean_do_not_keepdims_example": "data-dependent-shape", + "test_reduce_mean_do_not_keepdims_random": "data-dependent-shape", + "test_reduce_mean_keepdims_example": "data-dependent-shape", + "test_reduce_mean_keepdims_random": "data-dependent-shape", + "test_reduce_mean_negative_axes_keepdims_example": "data-dependent-shape", + "test_reduce_mean_negative_axes_keepdims_random": "data-dependent-shape", + "test_reduce_min_bool_inputs": "data-dependent-shape", + "test_reduce_min_do_not_keepdims_example": "data-dependent-shape", + "test_reduce_min_do_not_keepdims_random": "data-dependent-shape", + "test_reduce_min_empty_set": "data-dependent-shape", + "test_reduce_min_keepdims_example": "data-dependent-shape", + "test_reduce_min_keepdims_random": "data-dependent-shape", + "test_reduce_min_negative_axes_keepdims_example": "data-dependent-shape", + "test_reduce_min_negative_axes_keepdims_random": "data-dependent-shape", + "test_reduce_prod_do_not_keepdims_example": "data-dependent-shape", + "test_reduce_prod_do_not_keepdims_random": "data-dependent-shape", + "test_reduce_prod_empty_set": "data-dependent-shape", + "test_reduce_prod_keepdims_example": "data-dependent-shape", + "test_reduce_prod_keepdims_random": "data-dependent-shape", + "test_reduce_prod_negative_axes_keepdims_example": "data-dependent-shape", + "test_reduce_prod_negative_axes_keepdims_random": "data-dependent-shape", + "test_reduce_sum_do_not_keepdims_example": "data-dependent-shape", + "test_reduce_sum_do_not_keepdims_random": "data-dependent-shape", + "test_reduce_sum_empty_set": "data-dependent-shape", + "test_reduce_sum_empty_set_non_reduced_axis_zero": "data-dependent-shape", + "test_reduce_sum_keepdims_example": "data-dependent-shape", + "test_reduce_sum_keepdims_random": "data-dependent-shape", + "test_reduce_sum_negative_axes_keepdims_example": "data-dependent-shape", + "test_reduce_sum_negative_axes_keepdims_random": "data-dependent-shape", + "test_reduce_sum_square_do_not_keepdims_example": "data-dependent-shape", + "test_reduce_sum_square_do_not_keepdims_example_expanded": "data-dependent-shape", + "test_reduce_sum_square_do_not_keepdims_random": "data-dependent-shape", + "test_reduce_sum_square_do_not_keepdims_random_expanded": "data-dependent-shape", + "test_reduce_sum_square_empty_set": "data-dependent-shape", + "test_reduce_sum_square_empty_set_expanded": "data-dependent-shape", + "test_reduce_sum_square_keepdims_example": "data-dependent-shape", + "test_reduce_sum_square_keepdims_example_expanded": "data-dependent-shape", + "test_reduce_sum_square_keepdims_random": "data-dependent-shape", + "test_reduce_sum_square_keepdims_random_expanded": "data-dependent-shape", + "test_reduce_sum_square_negative_axes_keepdims_example": "data-dependent-shape", + "test_reduce_sum_square_negative_axes_keepdims_example_expanded": "data-dependent-shape", + "test_reduce_sum_square_negative_axes_keepdims_random": "data-dependent-shape", + "test_reduce_sum_square_negative_axes_keepdims_random_expanded": "data-dependent-shape", + "test_reflect_pad": "data-dependent-shape", + "test_regex_full_match_basic": "runtime-strings", + "test_regex_full_match_email_domain": "runtime-strings", + "test_regex_full_match_empty": "runtime-strings", + "test_reshape_allowzero_reordered": "data-dependent-shape", + "test_reshape_extended_dims": "data-dependent-shape", + "test_reshape_negative_dim": "data-dependent-shape", + "test_reshape_negative_extended_dims": "data-dependent-shape", + "test_reshape_one_dim": "data-dependent-shape", + "test_reshape_reduced_dims": "data-dependent-shape", + "test_reshape_reordered_all_dims": "data-dependent-shape", + "test_reshape_reordered_last_dims": "data-dependent-shape", + "test_reshape_zero_and_negative_dim": "data-dependent-shape", + "test_reshape_zero_dim": "data-dependent-shape", + "test_rms_normalization_2d_axis0_expanded": "data-dependent-shape", + "test_rms_normalization_2d_axis1_expanded": "data-dependent-shape", + "test_rms_normalization_2d_axis_negative_1_expanded": "data-dependent-shape", + "test_rms_normalization_2d_axis_negative_2_expanded": "data-dependent-shape", + "test_rms_normalization_3d_axis0_epsilon_expanded": "data-dependent-shape", + "test_rms_normalization_3d_axis1_epsilon_expanded": "data-dependent-shape", + "test_rms_normalization_3d_axis2_epsilon_expanded": "data-dependent-shape", + "test_rms_normalization_3d_axis_negative_1_epsilon_expanded": "data-dependent-shape", + "test_rms_normalization_3d_axis_negative_2_epsilon_expanded": "data-dependent-shape", + "test_rms_normalization_3d_axis_negative_3_epsilon_expanded": "data-dependent-shape", + "test_rms_normalization_4d_axis0_expanded": "data-dependent-shape", + "test_rms_normalization_4d_axis1_expanded": "data-dependent-shape", + "test_rms_normalization_4d_axis2_expanded": "data-dependent-shape", + "test_rms_normalization_4d_axis3_expanded": "data-dependent-shape", + "test_rms_normalization_4d_axis_negative_1_expanded": "data-dependent-shape", + "test_rms_normalization_4d_axis_negative_2_expanded": "data-dependent-shape", + "test_rms_normalization_4d_axis_negative_3_expanded": "data-dependent-shape", + "test_rms_normalization_4d_axis_negative_4_expanded": "data-dependent-shape", + "test_rms_normalization_default_axis_expanded": "data-dependent-shape", + "test_rotary_embedding_3d_input_expanded": "data-dependent-shape", + "test_rotary_embedding_expanded": "data-dependent-shape", + "test_rotary_embedding_interleaved_expanded": "data-dependent-shape", + "test_rotary_embedding_no_position_ids_expanded": "data-dependent-shape", + "test_rotary_embedding_no_position_ids_interleaved_expanded": "data-dependent-shape", + "test_rotary_embedding_no_position_ids_rotary_dim_expanded": "data-dependent-shape", + "test_rotary_embedding_with_interleaved_rotary_dim_expanded": "data-dependent-shape", + "test_rotary_embedding_with_rotary_dim_expanded": "data-dependent-shape", + "test_scan9_multi_state": "control-flow", + "test_scan9_scalar": "control-flow", + "test_scan9_sum": "control-flow", + "test_scan_sum": "control-flow", + "test_sce_NCd1_mean_weight_negative_ii_expanded": "data-dependent-shape", + "test_sce_NCd1_mean_weight_negative_ii_log_prob_expanded": "data-dependent-shape", + "test_sce_NCd1d2d3_none_no_weight_negative_ii_expanded": "data-dependent-shape", + "test_sce_NCd1d2d3_none_no_weight_negative_ii_log_prob_expanded": "data-dependent-shape", + "test_sce_NCd1d2d3_sum_weight_high_ii_expanded": "data-dependent-shape", + "test_sce_NCd1d2d3_sum_weight_high_ii_log_prob_expanded": "data-dependent-shape", + "test_sce_NCd1d2d3d4d5_mean_weight_expanded": "data-dependent-shape", + "test_sce_NCd1d2d3d4d5_mean_weight_log_prob_expanded": "data-dependent-shape", + "test_sce_NCd1d2d3d4d5_none_no_weight_expanded": "data-dependent-shape", + "test_sce_NCd1d2d3d4d5_none_no_weight_log_prob_expanded": "data-dependent-shape", + "test_sce_mean_3d_expanded": "data-dependent-shape", + "test_sce_mean_3d_log_prob_expanded": "data-dependent-shape", + "test_sce_mean_expanded": "data-dependent-shape", + "test_sce_mean_log_prob_expanded": "data-dependent-shape", + "test_sce_mean_no_weight_ii_3d_expanded": "data-dependent-shape", + "test_sce_mean_no_weight_ii_3d_log_prob_expanded": "data-dependent-shape", + "test_sce_mean_no_weight_ii_4d_expanded": "data-dependent-shape", + "test_sce_mean_no_weight_ii_4d_log_prob_expanded": "data-dependent-shape", + "test_sce_mean_no_weight_ii_expanded": "data-dependent-shape", + "test_sce_mean_no_weight_ii_log_prob_expanded": "data-dependent-shape", + "test_sce_mean_weight_expanded": "data-dependent-shape", + "test_sce_mean_weight_ii_3d_expanded": "data-dependent-shape", + "test_sce_mean_weight_ii_3d_log_prob_expanded": "data-dependent-shape", + "test_sce_mean_weight_ii_4d_expanded": "data-dependent-shape", + "test_sce_mean_weight_ii_4d_log_prob_expanded": "data-dependent-shape", + "test_sce_mean_weight_ii_expanded": "data-dependent-shape", + "test_sce_mean_weight_ii_log_prob_expanded": "data-dependent-shape", + "test_sce_mean_weight_log_prob_expanded": "data-dependent-shape", + "test_sce_none_expanded": "data-dependent-shape", + "test_sce_none_log_prob_expanded": "data-dependent-shape", + "test_sce_none_weights_expanded": "data-dependent-shape", + "test_sce_none_weights_log_prob_expanded": "data-dependent-shape", + "test_sce_sum_expanded": "data-dependent-shape", + "test_sce_sum_log_prob_expanded": "data-dependent-shape", + "test_sequence_insert_at_back": "non-tensor-io", + "test_sequence_insert_at_front": "non-tensor-io", + "test_sequence_map_add_1_sequence_1_tensor": "non-tensor-io", + "test_sequence_map_add_1_sequence_1_tensor_expanded": "non-tensor-io", + "test_sequence_map_add_2_sequences": "non-tensor-io", + "test_sequence_map_add_2_sequences_expanded": "non-tensor-io", + "test_sequence_map_extract_shapes": "non-tensor-io", + "test_sequence_map_extract_shapes_expanded": "non-tensor-io", + "test_sequence_map_identity_1_sequence": "non-tensor-io", + "test_sequence_map_identity_1_sequence_1_tensor": "non-tensor-io", + "test_sequence_map_identity_1_sequence_1_tensor_expanded": "non-tensor-io", + "test_sequence_map_identity_1_sequence_expanded": "non-tensor-io", + "test_sequence_map_identity_2_sequences": "non-tensor-io", + "test_sequence_map_identity_2_sequences_expanded": "non-tensor-io", + "test_slice": "data-dependent-shape", + "test_slice_default_axes": "data-dependent-shape", + "test_slice_default_steps": "data-dependent-shape", + "test_slice_end_out_of_bounds": "data-dependent-shape", + "test_slice_neg": "data-dependent-shape", + "test_slice_neg_steps": "data-dependent-shape", + "test_slice_negative_axes": "data-dependent-shape", + "test_slice_start_out_of_bounds": "data-dependent-shape", + "test_split_to_sequence_1": "non-tensor-io", + "test_split_to_sequence_2": "non-tensor-io", + "test_split_to_sequence_nokeepdims": "non-tensor-io", + "test_split_variable_parts_1d_opset13": "data-dependent-shape", + "test_split_variable_parts_1d_opset18": "data-dependent-shape", + "test_split_variable_parts_2d_opset13": "data-dependent-shape", + "test_split_variable_parts_2d_opset18": "data-dependent-shape", + "test_split_variable_parts_default_axis_opset13": "data-dependent-shape", + "test_split_variable_parts_default_axis_opset18": "data-dependent-shape", + "test_split_zero_size_splits_opset13": "data-dependent-shape", + "test_split_zero_size_splits_opset18": "data-dependent-shape", + "test_squeeze": "data-dependent-shape", + "test_squeeze_negative_axes": "data-dependent-shape", + "test_stft": "data-dependent-shape", + "test_stft_with_window": "data-dependent-shape", + "test_string_concat": "runtime-strings", + "test_string_concat_broadcasting": "runtime-strings", + "test_string_concat_empty_string": "runtime-strings", + "test_string_concat_utf8": "runtime-strings", + "test_string_concat_zero_dimensional": "runtime-strings", + "test_string_split_basic": "runtime-strings", + "test_string_split_consecutive_delimiters": "runtime-strings", + "test_string_split_empty_string_delimiter": "runtime-strings", + "test_string_split_empty_tensor": "runtime-strings", + "test_string_split_maxsplit": "runtime-strings", + "test_string_split_no_delimiter": "runtime-strings", + "test_strnormalizer_export_monday_casesensintive_lower": "runtime-strings", + "test_strnormalizer_export_monday_casesensintive_nochangecase": "runtime-strings", + "test_strnormalizer_export_monday_casesensintive_upper": "runtime-strings", + "test_strnormalizer_export_monday_empty_output": "runtime-strings", + "test_strnormalizer_export_monday_insensintive_upper_twodim": "runtime-strings", + "test_strnormalizer_nostopwords_nochangecase": "runtime-strings", + "test_tile": "data-dependent-shape", + "test_tile_precomputed": "data-dependent-shape", + "test_top_k": "data-dependent-shape", + "test_top_k_negative_axis": "data-dependent-shape", + "test_top_k_same_values": "data-dependent-shape", + "test_top_k_same_values_2d": "data-dependent-shape", + "test_top_k_same_values_largest": "data-dependent-shape", + "test_top_k_smallest": "data-dependent-shape", + "test_top_k_uint64": "data-dependent-shape", + "test_training_dropout": "random-op", + "test_training_dropout_default": "random-op", + "test_training_dropout_default_mask": "random-op", + "test_training_dropout_mask": "random-op", + "test_training_dropout_zero_ratio": "random-op", + "test_training_dropout_zero_ratio_mask": "random-op", + "test_unique_length_1": "data-dependent-shape", + "test_unique_not_sorted_without_axis": "data-dependent-shape", + "test_unique_sorted_with_axis": "data-dependent-shape", + "test_unique_sorted_with_axis_3d": "data-dependent-shape", + "test_unique_sorted_with_negative_axis": "data-dependent-shape", + "test_unique_sorted_without_axis": "data-dependent-shape", + "test_unsqueeze_axis_0": "data-dependent-shape", + "test_unsqueeze_axis_1": "data-dependent-shape", + "test_unsqueeze_axis_2": "data-dependent-shape", + "test_unsqueeze_negative_axes": "data-dependent-shape", + "test_unsqueeze_three_axes": "data-dependent-shape", + "test_unsqueeze_two_axes": "data-dependent-shape", + "test_unsqueeze_unsorted_axes": "data-dependent-shape", + "test_wrap_pad": "data-dependent-shape" +} diff --git a/src/python/tests/conformance/passing.txt b/src/python/tests/conformance/passing.txt new file mode 100644 index 0000000..1f0c3a8 --- /dev/null +++ b/src/python/tests/conformance/passing.txt @@ -0,0 +1,1050 @@ +# ONNX backend conformance pass-list ratchet. +# Every test listed here compiles, builds, runs and matches the corpus outputs. +# The list only ever grows: a kernel milestone adds its tests here and removes +# them from conformance/ledger.json in the same change. +test_abs +test_acos +test_acos_example +test_acosh +test_acosh_example +test_add +test_add_bcast +test_add_int16 +test_add_int8 +test_add_uint16 +test_add_uint32 +test_add_uint64 +test_add_uint8 +test_ai_onnx_ml_array_feature_extractor +test_ai_onnx_ml_binarizer +test_ai_onnx_ml_tree_ensemble_set_membership +test_ai_onnx_ml_tree_ensemble_single_tree +test_and2d +test_and3d +test_and4d +test_and_bcast3v1d +test_and_bcast3v2d +test_and_bcast4v2d +test_and_bcast4v3d +test_and_bcast4v4d +test_argmax_default_axis_example +test_argmax_default_axis_example_select_last_index +test_argmax_default_axis_random +test_argmax_default_axis_random_select_last_index +test_argmax_keepdims_example +test_argmax_keepdims_example_select_last_index +test_argmax_keepdims_random +test_argmax_keepdims_random_select_last_index +test_argmax_negative_axis_keepdims_example +test_argmax_negative_axis_keepdims_example_select_last_index +test_argmax_negative_axis_keepdims_random +test_argmax_negative_axis_keepdims_random_select_last_index +test_argmax_no_keepdims_example +test_argmax_no_keepdims_example_select_last_index +test_argmax_no_keepdims_random +test_argmax_no_keepdims_random_select_last_index +test_argmin_default_axis_example +test_argmin_default_axis_example_select_last_index +test_argmin_default_axis_random +test_argmin_default_axis_random_select_last_index +test_argmin_keepdims_example +test_argmin_keepdims_example_select_last_index +test_argmin_keepdims_random +test_argmin_keepdims_random_select_last_index +test_argmin_negative_axis_keepdims_example +test_argmin_negative_axis_keepdims_example_select_last_index +test_argmin_negative_axis_keepdims_random +test_argmin_negative_axis_keepdims_random_select_last_index +test_argmin_no_keepdims_example +test_argmin_no_keepdims_example_select_last_index +test_argmin_no_keepdims_random +test_argmin_no_keepdims_random_select_last_index +test_asin +test_asin_example +test_asinh +test_asinh_example +test_atan +test_atan_example +test_atanh +test_atanh_example +test_attention_3d +test_attention_3d_attn_mask +test_attention_3d_causal +test_attention_3d_diff_heads_sizes +test_attention_3d_diff_heads_sizes_attn_mask +test_attention_3d_diff_heads_sizes_causal +test_attention_3d_diff_heads_sizes_scaled +test_attention_3d_diff_heads_sizes_softcap +test_attention_3d_diff_heads_with_past_and_present +test_attention_3d_gqa +test_attention_3d_gqa_attn_mask +test_attention_3d_gqa_causal +test_attention_3d_gqa_scaled +test_attention_3d_gqa_softcap +test_attention_3d_gqa_with_past_and_present +test_attention_3d_scaled +test_attention_3d_softcap +test_attention_3d_transpose_verification +test_attention_3d_with_past_and_present +test_attention_3d_with_past_and_present_qk_matmul +test_attention_3d_with_past_and_present_qk_matmul_bias +test_attention_3d_with_past_and_present_qk_matmul_softcap +test_attention_3d_with_past_and_present_qk_matmul_softmax +test_attention_4d +test_attention_4d_attn_mask +test_attention_4d_attn_mask_3d +test_attention_4d_attn_mask_3d_causal +test_attention_4d_attn_mask_4d +test_attention_4d_attn_mask_4d_causal +test_attention_4d_attn_mask_bool +test_attention_4d_attn_mask_bool_4d +test_attention_4d_causal +test_attention_4d_diff_heads_mask4d_padded_kv +test_attention_4d_diff_heads_sizes +test_attention_4d_diff_heads_sizes_attn_mask +test_attention_4d_diff_heads_sizes_causal +test_attention_4d_diff_heads_sizes_scaled +test_attention_4d_diff_heads_sizes_softcap +test_attention_4d_diff_heads_with_past_and_present +test_attention_4d_diff_heads_with_past_and_present_mask3d +test_attention_4d_diff_heads_with_past_and_present_mask4d +test_attention_4d_gqa +test_attention_4d_gqa_attn_mask +test_attention_4d_gqa_causal +test_attention_4d_gqa_scaled +test_attention_4d_gqa_softcap +test_attention_4d_gqa_with_past_and_present +test_attention_4d_scaled +test_attention_4d_softcap +test_attention_4d_softcap_neginf_mask +test_attention_4d_softcap_neginf_mask_poison +test_attention_4d_with_past_and_present +test_attention_4d_with_past_and_present_qk_matmul +test_attention_4d_with_past_and_present_qk_matmul_bias +test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask +test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask_causal +test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask +test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask_causal +test_attention_4d_with_qk_matmul +test_attention_4d_with_qk_matmul_bias +test_attention_4d_with_qk_matmul_softcap +test_attention_4d_with_qk_matmul_softmax +test_averagepool_1d_default +test_averagepool_2d_ceil +test_averagepool_2d_ceil_last_window_starts_on_pad +test_averagepool_2d_default +test_averagepool_2d_dilations +test_averagepool_2d_pads +test_averagepool_2d_pads_count_include_pad +test_averagepool_2d_precomputed_pads +test_averagepool_2d_precomputed_pads_count_include_pad +test_averagepool_2d_precomputed_same_upper +test_averagepool_2d_precomputed_strides +test_averagepool_2d_same_lower +test_averagepool_2d_same_upper +test_averagepool_2d_strides +test_averagepool_3d_default +test_averagepool_3d_dilations_large_count_include_pad_is_0_ceil_mode_is_False +test_averagepool_3d_dilations_large_count_include_pad_is_0_ceil_mode_is_True +test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_False +test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True +test_averagepool_3d_dilations_small +test_basic_conv_with_padding +test_basic_conv_without_padding +test_basic_deform_conv_with_padding +test_basic_deform_conv_without_padding +test_batchnorm_epsilon +test_batchnorm_epsilon_training_mode +test_batchnorm_example +test_batchnorm_example_training_mode +test_bitcast_2d_float32_to_int32 +test_bitcast_bool_to_uint8 +test_bitcast_float32_to_int32 +test_bitcast_float64_to_int64 +test_bitcast_int32_to_float32 +test_bitcast_int64_to_float64 +test_bitcast_int8_to_uint8 +test_bitcast_scalar_float32_to_int32 +test_bitcast_uint16_to_int16 +test_bitcast_uint32_to_int32 +test_bitshift_left_uint16 +test_bitshift_left_uint32 +test_bitshift_left_uint64 +test_bitshift_left_uint8 +test_bitshift_right_uint16 +test_bitshift_right_uint32 +test_bitshift_right_uint64 +test_bitshift_right_uint8 +test_bitwise_and_i16_3d +test_bitwise_and_i32_2d +test_bitwise_and_ui64_bcast_3v1d +test_bitwise_and_ui8_bcast_4v3d +test_bitwise_not_2d +test_bitwise_not_3d +test_bitwise_not_4d +test_bitwise_or_i16_4d +test_bitwise_or_i32_2d +test_bitwise_or_ui64_bcast_3v1d +test_bitwise_or_ui8_bcast_4v3d +test_bitwise_xor_i16_3d +test_bitwise_xor_i32_2d +test_bitwise_xor_ui64_bcast_3v1d +test_bitwise_xor_ui8_bcast_4v3d +test_cast_DOUBLE_to_FLOAT +test_cast_FLOAT_to_DOUBLE +test_castlike_DOUBLE_to_FLOAT +test_castlike_DOUBLE_to_FLOAT_expanded +test_castlike_FLOAT_to_DOUBLE +test_castlike_FLOAT_to_DOUBLE_expanded +test_causal_conv_with_state_b1_c1_degenerate +test_causal_conv_with_state_b1_c1_degenerate_expanded +test_causal_conv_with_state_basic +test_causal_conv_with_state_basic_expanded +test_causal_conv_with_state_decode_step +test_causal_conv_with_state_decode_step_expanded +test_causal_conv_with_state_kernel_size_one +test_causal_conv_with_state_kernel_size_one_expanded +test_causal_conv_with_state_short_input_no_past_state +test_causal_conv_with_state_short_input_no_past_state_expanded +test_causal_conv_with_state_silu +test_causal_conv_with_state_silu_expanded +test_causal_conv_with_state_silu_with_past_state +test_causal_conv_with_state_silu_with_past_state_expanded +test_causal_conv_with_state_swish_alias +test_causal_conv_with_state_swish_alias_expanded +test_causal_conv_with_state_with_bias +test_causal_conv_with_state_with_bias_and_past_state +test_causal_conv_with_state_with_bias_and_past_state_expanded +test_causal_conv_with_state_with_bias_expanded +test_causal_conv_with_state_with_past_state +test_causal_conv_with_state_with_past_state_expanded +test_ceil +test_ceil_example +test_celu +test_celu_expanded +test_clip +test_clip_default_inbounds +test_clip_default_inbounds_expanded +test_clip_default_int8_inbounds +test_clip_default_int8_inbounds_expanded +test_clip_default_int8_max +test_clip_default_int8_max_expanded +test_clip_default_int8_min +test_clip_default_int8_min_expanded +test_clip_default_max +test_clip_default_max_expanded +test_clip_default_min +test_clip_default_min_expanded +test_clip_example +test_clip_example_expanded +test_clip_expanded +test_clip_inbounds +test_clip_inbounds_expanded +test_clip_min_greater_than_max +test_clip_min_greater_than_max_expanded +test_clip_outbounds +test_clip_outbounds_expanded +test_clip_splitbounds +test_clip_splitbounds_expanded +test_concat_1d_axis_0 +test_concat_1d_axis_negative_1 +test_concat_2d_axis_0 +test_concat_2d_axis_1 +test_concat_2d_axis_negative_1 +test_concat_2d_axis_negative_2 +test_concat_3d_axis_0 +test_concat_3d_axis_1 +test_concat_3d_axis_2 +test_concat_3d_axis_negative_1 +test_concat_3d_axis_negative_2 +test_concat_3d_axis_negative_3 +test_constant +test_conv_with_autopad_same +test_conv_with_strides_and_asymmetric_padding +test_conv_with_strides_no_padding +test_conv_with_strides_padding +test_convinteger_with_padding +test_convinteger_without_padding +test_convtranspose +test_convtranspose_1d +test_convtranspose_3d +test_convtranspose_autopad_same +test_convtranspose_dilations +test_convtranspose_group_2 +test_convtranspose_group_2_image_3 +test_convtranspose_kernel_shape +test_convtranspose_output_shape +test_convtranspose_pad +test_convtranspose_pads +test_cos +test_cos_example +test_cosh +test_cosh_example +test_cumprod_1d +test_cumprod_1d_exclusive +test_cumprod_1d_int32_exclusive +test_cumprod_1d_reverse +test_cumprod_1d_reverse_exclusive +test_cumprod_2d_axis_0 +test_cumprod_2d_axis_1 +test_cumprod_2d_int32 +test_cumprod_2d_negative_axis +test_cumsum_1d +test_cumsum_1d_exclusive +test_cumsum_1d_int32_exclusive +test_cumsum_1d_reverse +test_cumsum_1d_reverse_exclusive +test_cumsum_2d_axis_0 +test_cumsum_2d_axis_1 +test_cumsum_2d_int32 +test_cumsum_2d_negative_axis +test_deform_conv_with_mask_bias +test_deform_conv_with_multiple_offset_groups +test_depthtospace_crd_mode_example +test_depthtospace_example +test_dequantizelinear +test_dequantizelinear_axis +test_dequantizelinear_blocked +test_dequantizelinear_int16 +test_dequantizelinear_uint16 +test_det_2d +test_det_nd +test_dft +test_dft_axis +test_dft_axis_opset19 +test_dft_inverse +test_dft_inverse_opset19 +test_dft_irfft_opset19 +test_dft_opset19 +test_dft_rfft_opset19 +test_div +test_div_bcast +test_div_example +test_div_int16 +test_div_int32_trunc +test_div_int8 +test_div_uint16 +test_div_uint32 +test_div_uint64 +test_div_uint8 +test_dropout_default +test_dropout_default_mask +test_dropout_default_mask_ratio +test_dropout_default_old +test_dropout_default_ratio +test_dropout_random_old +test_dynamicquantizelinear +test_dynamicquantizelinear_expanded +test_dynamicquantizelinear_max_adjusted +test_dynamicquantizelinear_max_adjusted_expanded +test_dynamicquantizelinear_min_adjusted +test_dynamicquantizelinear_min_adjusted_expanded +test_einsum_batch_diagonal +test_einsum_batch_matmul +test_einsum_inner_prod +test_einsum_scalar +test_einsum_sum +test_einsum_transpose +test_elu +test_elu_default +test_elu_default_expanded_ver18 +test_elu_example +test_elu_example_expanded_ver18 +test_elu_expanded_ver18 +test_equal +test_equal_bcast +test_equal_int16 +test_equal_int8 +test_equal_uint16 +test_equal_uint32 +test_equal_uint64 +test_equal_uint8 +test_erf +test_exp +test_exp_example +test_eyelike_populate_off_main_diagonal +test_eyelike_with_dtype +test_eyelike_without_dtype +test_flatten_axis0 +test_flatten_axis1 +test_flatten_axis2 +test_flatten_axis3 +test_flatten_default_axis +test_flatten_negative_axis1 +test_flatten_negative_axis2 +test_flatten_negative_axis3 +test_flatten_negative_axis4 +test_floor +test_floor_example +test_gather_0 +test_gather_1 +test_gather_2d_indices +test_gather_elements_0 +test_gather_elements_1 +test_gather_elements_negative_indices +test_gather_negative_indices +test_gathernd_example_float32 +test_gathernd_example_int32 +test_gathernd_example_int32_batch_dim1 +test_gelu_default_1 +test_gelu_default_1_expanded +test_gelu_default_2 +test_gelu_default_2_expanded +test_gelu_tanh_1 +test_gelu_tanh_1_expanded +test_gelu_tanh_2 +test_gelu_tanh_2_expanded +test_gemm_all_attributes +test_gemm_alpha +test_gemm_beta +test_gemm_default_matrix_bias +test_gemm_default_no_bias +test_gemm_default_scalar_bias +test_gemm_default_single_elem_vector_bias +test_gemm_default_vector_bias +test_gemm_default_zero_bias +test_gemm_transposeA +test_gemm_transposeB +test_globalaveragepool +test_globalaveragepool_precomputed +test_globalmaxpool +test_globalmaxpool_precomputed +test_greater +test_greater_bcast +test_greater_equal +test_greater_equal_bcast +test_greater_equal_bcast_expanded +test_greater_equal_expanded +test_greater_equal_int16 +test_greater_equal_int16_expanded +test_greater_equal_int8 +test_greater_equal_int8_expanded +test_greater_equal_uint16 +test_greater_equal_uint16_expanded +test_greater_equal_uint32 +test_greater_equal_uint32_expanded +test_greater_equal_uint64 +test_greater_equal_uint64_expanded +test_greater_equal_uint8 +test_greater_equal_uint8_expanded +test_greater_int16 +test_greater_int8 +test_greater_uint16 +test_greater_uint32 +test_greater_uint64 +test_greater_uint8 +test_gridsample +test_gridsample_aligncorners_true +test_gridsample_bicubic +test_gridsample_bicubic_align_corners_0_additional_1 +test_gridsample_bicubic_align_corners_1_additional_1 +test_gridsample_bilinear +test_gridsample_bilinear_align_corners_0_additional_1 +test_gridsample_bilinear_align_corners_1_additional_1 +test_gridsample_border_padding +test_gridsample_nearest +test_gridsample_nearest_align_corners_0_additional_1 +test_gridsample_nearest_align_corners_1_additional_1 +test_gridsample_reflection_padding +test_gridsample_volumetric_bilinear_align_corners_0 +test_gridsample_volumetric_bilinear_align_corners_1 +test_gridsample_volumetric_nearest_align_corners_0 +test_gridsample_volumetric_nearest_align_corners_1 +test_gridsample_zeros_padding +test_group_normalization_epsilon +test_group_normalization_epsilon_expanded +test_group_normalization_example +test_group_normalization_example_expanded +test_gru_batchwise +test_gru_defaults +test_gru_seq_length +test_gru_with_initial_bias +test_hardmax_axis_0 +test_hardmax_axis_1 +test_hardmax_axis_2 +test_hardmax_default_axis +test_hardmax_example +test_hardmax_negative_axis +test_hardmax_one_hot +test_hardsigmoid +test_hardsigmoid_default +test_hardsigmoid_default_expanded_ver18 +test_hardsigmoid_example +test_hardsigmoid_example_expanded_ver18 +test_hardsigmoid_expanded_ver18 +test_hardswish +test_hardswish_expanded +test_identity +test_instancenorm_epsilon +test_instancenorm_example +test_isinf +test_isinf_negative +test_isinf_positive +test_isnan +test_l1normalization_axis_0 +test_l1normalization_axis_1 +test_l1normalization_axis_last +test_l2normalization_axis_0 +test_l2normalization_axis_1 +test_layer_normalization_2d_axis0 +test_layer_normalization_2d_axis1 +test_layer_normalization_2d_axis_negative_1 +test_layer_normalization_2d_axis_negative_2 +test_layer_normalization_3d_axis0_epsilon +test_layer_normalization_3d_axis1_epsilon +test_layer_normalization_3d_axis2_epsilon +test_layer_normalization_3d_axis_negative_1_epsilon +test_layer_normalization_3d_axis_negative_2_epsilon +test_layer_normalization_3d_axis_negative_3_epsilon +test_layer_normalization_4d_axis0 +test_layer_normalization_4d_axis1 +test_layer_normalization_4d_axis2 +test_layer_normalization_4d_axis3 +test_layer_normalization_4d_axis_negative_1 +test_layer_normalization_4d_axis_negative_2 +test_layer_normalization_4d_axis_negative_3 +test_layer_normalization_4d_axis_negative_4 +test_layer_normalization_default_axis +test_leakyrelu +test_leakyrelu_default +test_leakyrelu_default_expanded +test_leakyrelu_example +test_leakyrelu_example_expanded +test_leakyrelu_expanded +test_less +test_less_bcast +test_less_equal +test_less_equal_bcast +test_less_equal_bcast_expanded +test_less_equal_expanded +test_less_equal_int16 +test_less_equal_int16_expanded +test_less_equal_int8 +test_less_equal_int8_expanded +test_less_equal_uint16 +test_less_equal_uint16_expanded +test_less_equal_uint32 +test_less_equal_uint32_expanded +test_less_equal_uint64 +test_less_equal_uint64_expanded +test_less_equal_uint8 +test_less_equal_uint8_expanded +test_less_int16 +test_less_int8 +test_less_uint16 +test_less_uint32 +test_less_uint64 +test_less_uint8 +test_linear_attention_decode_step +test_linear_attention_delta +test_linear_attention_explicit_scale +test_linear_attention_gated +test_linear_attention_gated_delta +test_linear_attention_gated_delta_beta_scalar +test_linear_attention_gated_delta_gqa +test_linear_attention_gated_delta_mqa +test_linear_attention_gated_per_head_decay +test_linear_attention_linear +test_linear_attention_linear_t1_no_past +test_linear_attention_no_past_explicit_zeros +test_linear_attention_prefill_with_past +test_log +test_log_example +test_logsoftmax_axis_0 +test_logsoftmax_axis_0_expanded +test_logsoftmax_axis_0_expanded_ver18 +test_logsoftmax_axis_1 +test_logsoftmax_axis_1_expanded +test_logsoftmax_axis_1_expanded_ver18 +test_logsoftmax_axis_2 +test_logsoftmax_axis_2_expanded +test_logsoftmax_axis_2_expanded_ver18 +test_logsoftmax_default_axis +test_logsoftmax_default_axis_expanded +test_logsoftmax_default_axis_expanded_ver18 +test_logsoftmax_example_1 +test_logsoftmax_example_1_expanded +test_logsoftmax_example_1_expanded_ver18 +test_logsoftmax_large_number +test_logsoftmax_large_number_expanded +test_logsoftmax_large_number_expanded_ver18 +test_logsoftmax_negative_axis +test_logsoftmax_negative_axis_expanded +test_logsoftmax_negative_axis_expanded_ver18 +test_lpnormalization_default +test_lppool_1d_default +test_lppool_2d_default +test_lppool_2d_dilations +test_lppool_2d_pads +test_lppool_2d_same_lower +test_lppool_2d_same_upper +test_lppool_2d_strides +test_lppool_3d_default +test_lrn +test_lrn_default +test_lstm_batchwise +test_lstm_defaults +test_lstm_with_initial_bias +test_lstm_with_peepholes +test_matmul_1d_1d +test_matmul_1d_3d +test_matmul_2d +test_matmul_3d +test_matmul_4d +test_matmul_4d_1d +test_matmul_bcast +test_matmulinteger +test_max_example +test_max_float32 +test_max_float64 +test_max_int16 +test_max_int32 +test_max_int64 +test_max_int8 +test_max_one_input +test_max_two_inputs +test_max_uint16 +test_max_uint32 +test_max_uint64 +test_max_uint8 +test_maxpool_1d_default +test_maxpool_2d_ceil +test_maxpool_2d_ceil_output_size_reduce_by_one +test_maxpool_2d_default +test_maxpool_2d_dilations +test_maxpool_2d_pads +test_maxpool_2d_precomputed_pads +test_maxpool_2d_precomputed_same_upper +test_maxpool_2d_precomputed_strides +test_maxpool_2d_same_lower +test_maxpool_2d_same_upper +test_maxpool_2d_strides +test_maxpool_2d_uint8 +test_maxpool_3d_default +test_maxpool_3d_dilations +test_maxpool_3d_dilations_use_ref_impl +test_maxpool_3d_dilations_use_ref_impl_large +test_maxpool_with_argmax_2d_precomputed_pads +test_maxpool_with_argmax_2d_precomputed_strides +test_maxunpool_export_without_output_shape +test_mean_example +test_mean_one_input +test_mean_two_inputs +test_min_example +test_min_float32 +test_min_float64 +test_min_int16 +test_min_int32 +test_min_int64 +test_min_int8 +test_min_one_input +test_min_two_inputs +test_min_uint16 +test_min_uint32 +test_min_uint64 +test_min_uint8 +test_mish +test_mish_expanded +test_mod_broadcast +test_mod_int64_fmod +test_mod_mixed_sign_float32 +test_mod_mixed_sign_float64 +test_mod_mixed_sign_int16 +test_mod_mixed_sign_int32 +test_mod_mixed_sign_int64 +test_mod_mixed_sign_int8 +test_mod_uint16 +test_mod_uint32 +test_mod_uint64 +test_mod_uint8 +test_mul +test_mul_bcast +test_mul_example +test_mul_int16 +test_mul_int8 +test_mul_uint16 +test_mul_uint32 +test_mul_uint64 +test_mul_uint8 +test_mvn +test_mvn_expanded +test_mvn_expanded_ver18 +test_neg +test_neg_example +test_nllloss_NC +test_nllloss_NC_expanded +test_nllloss_NCd1 +test_nllloss_NCd1_expanded +test_nllloss_NCd1_ii +test_nllloss_NCd1_ii_expanded +test_nllloss_NCd1_mean_weight_negative_ii +test_nllloss_NCd1_mean_weight_negative_ii_expanded +test_nllloss_NCd1_weight +test_nllloss_NCd1_weight_expanded +test_nllloss_NCd1_weight_ii +test_nllloss_NCd1_weight_ii_expanded +test_nllloss_NCd1d2 +test_nllloss_NCd1d2_expanded +test_nllloss_NCd1d2_no_weight_reduction_mean_ii +test_nllloss_NCd1d2_no_weight_reduction_mean_ii_expanded +test_nllloss_NCd1d2_reduction_mean +test_nllloss_NCd1d2_reduction_mean_expanded +test_nllloss_NCd1d2_reduction_sum +test_nllloss_NCd1d2_reduction_sum_expanded +test_nllloss_NCd1d2_with_weight +test_nllloss_NCd1d2_with_weight_expanded +test_nllloss_NCd1d2_with_weight_reduction_mean +test_nllloss_NCd1d2_with_weight_reduction_mean_expanded +test_nllloss_NCd1d2_with_weight_reduction_sum +test_nllloss_NCd1d2_with_weight_reduction_sum_expanded +test_nllloss_NCd1d2_with_weight_reduction_sum_ii +test_nllloss_NCd1d2_with_weight_reduction_sum_ii_expanded +test_nllloss_NCd1d2d3_none_no_weight_negative_ii +test_nllloss_NCd1d2d3_none_no_weight_negative_ii_expanded +test_nllloss_NCd1d2d3_sum_weight_high_ii +test_nllloss_NCd1d2d3_sum_weight_high_ii_expanded +test_nllloss_NCd1d2d3d4d5_mean_weight +test_nllloss_NCd1d2d3d4d5_mean_weight_expanded +test_nllloss_NCd1d2d3d4d5_none_no_weight +test_nllloss_NCd1d2d3d4d5_none_no_weight_expanded +test_not_2d +test_not_3d +test_not_4d +test_optional_has_element_empty_no_input_name_optional_input +test_optional_has_element_empty_no_input_name_tensor_input +test_optional_has_element_empty_no_input_optional_input +test_optional_has_element_empty_no_input_tensor_input +test_or2d +test_or3d +test_or4d +test_or_bcast3v1d +test_or_bcast3v2d +test_or_bcast4v2d +test_or_bcast4v3d +test_or_bcast4v4d +test_pow +test_pow_bcast_array +test_pow_bcast_scalar +test_pow_example +test_pow_types_float32_int32 +test_pow_types_float32_int64 +test_pow_types_float32_uint32 +test_pow_types_float32_uint64 +test_pow_types_int32_float32 +test_pow_types_int32_int32 +test_pow_types_int64_float32 +test_pow_types_int64_int64 +test_prelu_broadcast +test_prelu_broadcast_expanded +test_prelu_example +test_prelu_example_expanded +test_qlinearconv +test_qlinearmatmul_2D_int8_float32 +test_qlinearmatmul_2D_uint8_float32 +test_qlinearmatmul_3D_int8_float32 +test_qlinearmatmul_3D_uint8_float32 +test_quantizelinear +test_quantizelinear_axis +test_quantizelinear_blocked_asymmetric +test_quantizelinear_blocked_symmetric +test_quantizelinear_int16 +test_quantizelinear_uint16 +test_reciprocal +test_reciprocal_example +test_reduce_l1_default_axes_keepdims_example +test_reduce_l1_default_axes_keepdims_example_expanded +test_reduce_l1_default_axes_keepdims_random +test_reduce_l1_default_axes_keepdims_random_expanded +test_reduce_l2_default_axes_keepdims_example +test_reduce_l2_default_axes_keepdims_example_expanded +test_reduce_l2_default_axes_keepdims_random +test_reduce_l2_default_axes_keepdims_random_expanded +test_reduce_log_sum_default +test_reduce_log_sum_default_expanded +test_reduce_log_sum_exp_default_axes_keepdims_example +test_reduce_log_sum_exp_default_axes_keepdims_example_expanded +test_reduce_log_sum_exp_default_axes_keepdims_random +test_reduce_log_sum_exp_default_axes_keepdims_random_expanded +test_reduce_max_default_axes_keepdim_example +test_reduce_max_default_axes_keepdims_random +test_reduce_mean_default_axes_keepdims_example +test_reduce_mean_default_axes_keepdims_random +test_reduce_min_default_axes_keepdims_example +test_reduce_min_default_axes_keepdims_random +test_reduce_prod_default_axes_keepdims_example +test_reduce_prod_default_axes_keepdims_random +test_reduce_sum_default_axes_keepdims_example +test_reduce_sum_default_axes_keepdims_random +test_reduce_sum_empty_axes_input_noop +test_reduce_sum_empty_axes_input_noop_example +test_reduce_sum_square_default_axes_keepdims_example +test_reduce_sum_square_default_axes_keepdims_example_expanded +test_reduce_sum_square_default_axes_keepdims_random +test_reduce_sum_square_default_axes_keepdims_random_expanded +test_relu +test_relu_expanded_ver18 +test_resize_downsample_scales_cubic +test_resize_downsample_scales_cubic_A_n0p5_exclude_outside +test_resize_downsample_scales_cubic_align_corners +test_resize_downsample_scales_cubic_antialias +test_resize_downsample_scales_linear +test_resize_downsample_scales_linear_align_corners +test_resize_downsample_scales_linear_antialias +test_resize_downsample_scales_linear_half_pixel_symmetric +test_resize_downsample_scales_nearest +test_resize_downsample_sizes_cubic +test_resize_downsample_sizes_cubic_antialias +test_resize_downsample_sizes_linear_antialias +test_resize_downsample_sizes_linear_pytorch_half_pixel +test_resize_downsample_sizes_nearest +test_resize_downsample_sizes_nearest_not_larger +test_resize_downsample_sizes_nearest_not_smaller +test_resize_tf_crop_and_resize +test_resize_tf_crop_and_resize_axes_2_3 +test_resize_tf_crop_and_resize_axes_3_2 +test_resize_tf_crop_and_resize_extrapolation_value +test_resize_upsample_scales_cubic +test_resize_upsample_scales_cubic_A_n0p5_exclude_outside +test_resize_upsample_scales_cubic_align_corners +test_resize_upsample_scales_cubic_asymmetric +test_resize_upsample_scales_linear +test_resize_upsample_scales_linear_align_corners +test_resize_upsample_scales_linear_half_pixel_symmetric +test_resize_upsample_scales_nearest +test_resize_upsample_scales_nearest_axes_2_3 +test_resize_upsample_scales_nearest_axes_3_2 +test_resize_upsample_sizes_cubic +test_resize_upsample_sizes_nearest +test_resize_upsample_sizes_nearest_axes_2_3 +test_resize_upsample_sizes_nearest_axes_3_2 +test_resize_upsample_sizes_nearest_ceil_half_pixel +test_resize_upsample_sizes_nearest_floor_align_corners +test_resize_upsample_sizes_nearest_not_larger +test_resize_upsample_sizes_nearest_not_smaller +test_resize_upsample_sizes_nearest_round_prefer_ceil_asymmetric +test_reversesequence_batch +test_reversesequence_time +test_rms_normalization_2d_axis0 +test_rms_normalization_2d_axis1 +test_rms_normalization_2d_axis_negative_1 +test_rms_normalization_2d_axis_negative_2 +test_rms_normalization_3d_axis0_epsilon +test_rms_normalization_3d_axis1_epsilon +test_rms_normalization_3d_axis2_epsilon +test_rms_normalization_3d_axis_negative_1_epsilon +test_rms_normalization_3d_axis_negative_2_epsilon +test_rms_normalization_3d_axis_negative_3_epsilon +test_rms_normalization_4d_axis0 +test_rms_normalization_4d_axis1 +test_rms_normalization_4d_axis2 +test_rms_normalization_4d_axis3 +test_rms_normalization_4d_axis_negative_1 +test_rms_normalization_4d_axis_negative_2 +test_rms_normalization_4d_axis_negative_3 +test_rms_normalization_4d_axis_negative_4 +test_rms_normalization_default_axis +test_rnn_seq_length +test_roialign_aligned_false +test_roialign_aligned_true +test_roialign_mode_max +test_rotary_embedding +test_rotary_embedding_3d_input +test_rotary_embedding_interleaved +test_rotary_embedding_no_position_ids +test_rotary_embedding_no_position_ids_interleaved +test_rotary_embedding_no_position_ids_rotary_dim +test_rotary_embedding_with_interleaved_rotary_dim +test_rotary_embedding_with_rotary_dim +test_round +test_scatter_elements_with_axis +test_scatter_elements_with_duplicate_indices +test_scatter_elements_with_negative_indices +test_scatter_elements_with_reduction_max +test_scatter_elements_with_reduction_min +test_scatter_elements_with_reduction_mul +test_scatter_elements_without_axis +test_scatter_with_axis +test_scatter_without_axis +test_scatternd +test_scatternd_add +test_scatternd_max +test_scatternd_min +test_scatternd_multiply +test_sce_NCd1_mean_weight_negative_ii +test_sce_NCd1_mean_weight_negative_ii_log_prob +test_sce_NCd1d2d3_none_no_weight_negative_ii +test_sce_NCd1d2d3_none_no_weight_negative_ii_log_prob +test_sce_NCd1d2d3_sum_weight_high_ii +test_sce_NCd1d2d3_sum_weight_high_ii_log_prob +test_sce_NCd1d2d3d4d5_mean_weight +test_sce_NCd1d2d3d4d5_mean_weight_log_prob +test_sce_NCd1d2d3d4d5_none_no_weight +test_sce_NCd1d2d3d4d5_none_no_weight_log_prob +test_sce_mean +test_sce_mean_3d +test_sce_mean_3d_log_prob +test_sce_mean_log_prob +test_sce_mean_no_weight_ii +test_sce_mean_no_weight_ii_3d +test_sce_mean_no_weight_ii_3d_log_prob +test_sce_mean_no_weight_ii_4d +test_sce_mean_no_weight_ii_4d_log_prob +test_sce_mean_no_weight_ii_log_prob +test_sce_mean_weight +test_sce_mean_weight_ii +test_sce_mean_weight_ii_3d +test_sce_mean_weight_ii_3d_log_prob +test_sce_mean_weight_ii_4d +test_sce_mean_weight_ii_4d_log_prob +test_sce_mean_weight_ii_log_prob +test_sce_mean_weight_log_prob +test_sce_none +test_sce_none_log_prob +test_sce_none_weights +test_sce_none_weights_log_prob +test_sce_sum +test_sce_sum_log_prob +test_selu +test_selu_default +test_selu_default_expanded_ver18 +test_selu_example +test_selu_example_expanded_ver18 +test_selu_expanded_ver18 +test_shape +test_shape_clip_end +test_shape_clip_start +test_shape_end_1 +test_shape_end_negative_1 +test_shape_example +test_shape_start_1 +test_shape_start_1_end_2 +test_shape_start_1_end_negative_1 +test_shape_start_greater_than_end +test_shape_start_negative_1 +test_shrink_hard +test_shrink_hard_expanded_ver18 +test_shrink_soft +test_shrink_soft_expanded_ver18 +test_sigmoid +test_sigmoid_example +test_sign +test_simple_rnn_batchwise +test_simple_rnn_defaults +test_simple_rnn_with_initial_bias +test_sin +test_sin_example +test_sinh +test_sinh_example +test_size +test_size_example +test_softmax_axis_0 +test_softmax_axis_0_expanded +test_softmax_axis_0_expanded_ver18 +test_softmax_axis_1 +test_softmax_axis_1_expanded +test_softmax_axis_1_expanded_ver18 +test_softmax_axis_2 +test_softmax_axis_2_expanded +test_softmax_axis_2_expanded_ver18 +test_softmax_default_axis +test_softmax_default_axis_expanded +test_softmax_default_axis_expanded_ver18 +test_softmax_example +test_softmax_example_expanded +test_softmax_example_expanded_ver18 +test_softmax_large_number +test_softmax_large_number_expanded +test_softmax_large_number_expanded_ver18 +test_softmax_negative_axis +test_softmax_negative_axis_expanded +test_softmax_negative_axis_expanded_ver18 +test_softplus +test_softplus_example +test_softplus_example_expanded_ver18 +test_softplus_expanded_ver18 +test_softsign +test_softsign_example +test_softsign_example_expanded_ver18 +test_softsign_expanded_ver18 +test_spacetodepth +test_spacetodepth_example +test_split_1d_uneven_split_opset18 +test_split_2d_uneven_split_opset18 +test_split_equal_parts_1d_opset13 +test_split_equal_parts_1d_opset18 +test_split_equal_parts_2d +test_split_equal_parts_2d_opset13 +test_split_equal_parts_default_axis_opset13 +test_split_equal_parts_default_axis_opset18 +test_sqrt +test_sqrt_example +test_sub +test_sub_bcast +test_sub_example +test_sub_int16 +test_sub_int8 +test_sub_uint16 +test_sub_uint32 +test_sub_uint64 +test_sub_uint8 +test_sum_example +test_sum_one_input +test_sum_two_inputs +test_swish +test_swish_expanded +test_tan +test_tan_example +test_tanh +test_tanh_example +test_tensorscatter +test_tensorscatter_3d +test_tensorscatter_circular +test_tfidfvectorizer_tf_batch_onlybigrams_skip0 +test_tfidfvectorizer_tf_batch_onlybigrams_skip5 +test_tfidfvectorizer_tf_batch_uniandbigrams_skip5 +test_tfidfvectorizer_tf_only_bigrams_skip0 +test_tfidfvectorizer_tf_onlybigrams_levelempty +test_tfidfvectorizer_tf_onlybigrams_skip5 +test_tfidfvectorizer_tf_uniandbigrams_skip5 +test_thresholdedrelu +test_thresholdedrelu_default +test_thresholdedrelu_default_expanded_ver18 +test_thresholdedrelu_example +test_thresholdedrelu_example_expanded_ver18 +test_thresholdedrelu_expanded_ver18 +test_transpose_all_permutations_0 +test_transpose_all_permutations_1 +test_transpose_all_permutations_2 +test_transpose_all_permutations_3 +test_transpose_all_permutations_4 +test_transpose_all_permutations_5 +test_transpose_default +test_tril +test_tril_neg +test_tril_one_row_neg +test_tril_out_neg +test_tril_out_pos +test_tril_pos +test_tril_square +test_tril_square_neg +test_tril_zero +test_triu +test_triu_neg +test_triu_one_row +test_triu_out_neg_out +test_triu_out_pos +test_triu_pos +test_triu_square +test_triu_square_neg +test_triu_zero +test_upsample_nearest +test_where_example +test_where_long_example +test_xor2d +test_xor3d +test_xor4d +test_xor_bcast3v1d +test_xor_bcast3v2d +test_xor_bcast4v2d +test_xor_bcast4v3d +test_xor_bcast4v4d diff --git a/src/python/tests/conformance/tolerances.json b/src/python/tests/conformance/tolerances.json new file mode 100644 index 0000000..f836f1c --- /dev/null +++ b/src/python/tests/conformance/tolerances.json @@ -0,0 +1,18 @@ +{ + "documentation": [ + "Per-op tolerance overrides for the ONNX backend conformance suite.", + "The defaults are ONNX's own (rtol 1e-3, atol 1e-7, or whatever a test's data.json", + "states); an entry here loosens them for every corpus test containing that op.", + "An override must carry a written numerical justification -- libm accuracy,", + "summation order -- and must stay within MAX_TOLERANCE_FACTOR of the ONNX default.", + "Anything that needs more than that is a wrong kernel, not a tolerance problem.", + "Format: \"\": {\"rtol\": , \"atol\": , \"justification\": \"...\"}." + ], + "overrides": { + "DFT": { + "rtol": 0.001, + "atol": 1e-06, + "justification": "The expected values come from numpy's FFT run at the tensor's own precision -- a single-precision butterfly for a float32 signal -- so a bin whose exact value cancels to 0 is stored as ~1 ULP of the largest term instead. test_dft_inverse and test_dft_inverse_opset19 are that case: an all-real ramp up to 90 whose Nyquist bin cancels exactly, stored as 1.907e-07 where the kernel, summing in double, leaves 3e-15. That is the oracle's noise floor, not the kernel's error: 1.906e-07 remains after the corpus rtol, so atol is loosened to ten times the ONNX default -- a tenth of the bound -- which clears it with margin and still catches a wrong bin." + } + } +} diff --git a/src/python/tests/test_extra_compiler_bundle.py b/src/python/tests/test_extra_compiler_bundle.py new file mode 100644 index 0000000..69cd76b --- /dev/null +++ b/src/python/tests/test_extra_compiler_bundle.py @@ -0,0 +1,1193 @@ +"""The FNNX bundle layer: reading a pipeline bundle, node entrypoints, and pipeline glue. + +The oracle for a compiled pipeline is the FNNX `Runtime` — onnxruntime executing the same +bundle — and, for a single node, the ONNX reference evaluator. Nothing here states an +expected output of its own. +""" + +from __future__ import annotations + +import json +import re +import shutil +import tarfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +pytest.importorskip("onnxruntime") +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") +bundle_module = pytest.importorskip("fnnx.extras.compilers.c.bundle") +codegen = pytest.importorskip("fnnx.extras.compilers.c.onnx.codegen") +kernels = pytest.importorskip("fnnx.extras.compilers.c.onnx.kernels") + +from fnnx.extras.compilers.c import compile_bundle # noqa: E402 +from fnnx.runtime import Runtime # noqa: E402 +from onnx import TensorProto, helper # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +OPSET = 21 +SEED = 20260726 + +MODELS = Path(__file__).parent / "models" +PIPELINE_BUNDLE = MODELS / "onnx_pipeline.fnnx" +PIPELINE_TAR = MODELS / "onnx_pipeline.fnnx.tar" + +# Compiling is only half of what these tests assert; the other half runs the artifact. +pytestmark = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + + +# -------------------------------------------------------------------------------------- +# Bundle fixtures +# -------------------------------------------------------------------------------------- + + +def _values(shape, *, seed: int = SEED): + return np.random.default_rng(seed).normal(size=shape).astype(np.float32) + + +def _spec(shape, dtype: str = "float32") -> dict[str, Any]: + return {"dtype": f"Array[{dtype}]", "shape": list(shape)} + + +def _manifest_tensor(name: str, shape=(), dtype: str = "float32") -> dict[str, Any]: + return { + "name": name, + "content_type": "NDJSON", + "dtype": f"Array[{dtype}]", + "shape": list(shape), + } + + +@dataclass +class _Node: + """One pipeline node of a bundle written for a test.""" + + id: str + model: Any + inputs: tuple[str, ...] + outputs: tuple[str, ...] + input_specs: tuple[dict[str, Any], ...] + output_specs: tuple[dict[str, Any], ...] + dynamic_attributes: dict[str, Any] = field(default_factory=dict) + extra_dynattrs: dict[str, str] = field(default_factory=dict) + op: str = "ONNX_v1" + opset: int = OPSET + + +def _affine_model(weight: np.ndarray, bias: float, *, name: str): + """`y = x * weight + bias` over a symbolic batch, with weights of its own.""" + width = int(weight.size) + return helper.make_model( + helper.make_graph( + [ + helper.make_node("Mul", ["x", "w"], ["scaled"], name="mul"), + helper.make_node("Add", ["scaled", "b"], ["y"], name="add"), + ], + name, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["batch", width])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, ["batch", width])], + initializer=[ + onnx.numpy_helper.from_array(weight.astype(np.float32), "w"), + onnx.numpy_helper.from_array(np.array([bias], dtype=np.float32), "b"), + ], + ), + opset_imports=[helper.make_opsetid("", OPSET)], + ) + + +def _difference_model(name: str): + """`y = a - b` over a symbolic batch: the join of a diamond. + + Subtraction rather than addition so that the two operands are not interchangeable: a + join wired to its inputs in the wrong order has to show up in the result. + """ + return helper.make_model( + helper.make_graph( + [helper.make_node("Sub", ["a", "b"], ["y"], name="sub")], + name, + [ + helper.make_tensor_value_info("a", TensorProto.FLOAT, ["batch", 3]), + helper.make_tensor_value_info("b", TensorProto.FLOAT, ["batch", 3]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, ["batch", 3])], + ), + opset_imports=[helper.make_opsetid("", OPSET)], + ) + + +def _write_bundle( + directory: Path, + nodes, + inputs, + outputs, + *, + variant: str = "pipeline", + manifest_dynamic_attributes=(), +) -> Path: + """Write a bundle the runtime can load and the compiler can read.""" + directory.mkdir(parents=True, exist_ok=True) + manifest = { + "variant": variant, + "producer_name": "tests", + "producer_version": "0.0.0", + "producer_tags": [], + "inputs": list(inputs), + "outputs": list(outputs), + "dynamic_attributes": list(manifest_dynamic_attributes), + "env_vars": [], + } + ops = [ + { + "id": node.id, + "op": node.op, + "inputs": list(node.input_specs), + "outputs": list(node.output_specs), + "attributes": { + "opsets": [{"domain": "ai.onnx", "version": node.opset}], + "requires_ort_extensions": False, + "has_external_data": False, + "onnx_ir_version": 10, + }, + "dynamic_attributes": dict(node.dynamic_attributes), + } + for node in nodes + ] + variant_config = { + "nodes": [ + { + "op_instance_id": node.id, + "inputs": list(node.inputs), + "outputs": list(node.outputs), + "extra_dynattrs": dict(node.extra_dynattrs), + } + for node in nodes + ] + } + for name, document in ( + ("manifest.json", manifest), + ("ops.json", ops), + ("variant_config.json", variant_config), + ("dtypes.json", {}), + ): + (directory / name).write_text(json.dumps(document, indent=2), encoding="utf-8") + for node in nodes: + artifacts = directory / "ops_artifacts" / node.id + artifacts.mkdir(parents=True, exist_ok=True) + if node.model is not None: + onnx.save_model(node.model, str(artifacts / "model.onnx")) + return directory + + +def _diamond_nodes(): + """`head` fans out to `left` and `right`, which `join` sums; `head` is also an output. + + Every node carries weights of its own, so a fan-out edge routed to the wrong buffer — + or a weight shared between nodes that should not share one — changes the result. + """ + tensor = _spec(["batch", 3]) + return [ + _Node( + id="head", + model=_affine_model(np.array([1.0, 2.0, 3.0]), 0.5, name="head"), + inputs=("x",), + outputs=("h",), + input_specs=(tensor,), + output_specs=(tensor,), + ), + _Node( + id="left", + model=_affine_model(np.array([-1.0, 0.25, 4.0]), -2.0, name="left"), + inputs=("h",), + outputs=("l",), + input_specs=(tensor,), + output_specs=(tensor,), + ), + _Node( + id="right", + model=_affine_model(np.array([7.0, -0.5, 0.125]), 3.0, name="right"), + inputs=("h",), + outputs=("r",), + input_specs=(tensor,), + output_specs=(tensor,), + ), + _Node( + id="join", + model=_difference_model("join"), + inputs=("l", "r"), + outputs=("out",), + input_specs=(tensor, tensor), + output_specs=(tensor,), + ), + ] + + +@pytest.fixture +def diamond_bundle(tmp_path): + return _write_bundle( + tmp_path / "diamond.fnnx", + _diamond_nodes(), + [_manifest_tensor("x")], + [_manifest_tensor("h"), _manifest_tensor("out")], + ) + + +def _runtime_outputs(bundle: Path, feeds: dict[str, Any]) -> dict[str, Any]: + return Runtime(str(bundle)).compute(dict(feeds), {}) + + +def _assert_matches(actual: dict[str, Any], expected: dict[str, Any]) -> None: + assert sorted(actual) == sorted(expected) + for name, want in expected.items(): + got = actual[name] + assert got.dtype == want.dtype, name + assert got.shape == want.shape, name + np.testing.assert_allclose(got, want, rtol=1e-6, atol=1e-6, err_msg=name) + + +# -------------------------------------------------------------------------------------- +# Compiling the pipeline test bundle +# -------------------------------------------------------------------------------------- + + +def test_the_pipeline_bundle_compiles_to_one_header_and_report(tmp_path): + result = compile_bundle(PIPELINE_BUNDLE, tmp_path) + header = result.header_path.read_text() + + assert sorted(path.name for path in tmp_path.iterdir()) == [ + "fnnx_model.h", + "fnnx_model_report.json", + ] + assert "int fnnx_model_run(" in header + for node_id in ("linreg", "linreg2", "linreg3", "concat_reduce"): + assert f"int fnnx_model_node_{node_id}_run(" in header + # The preamble points a reader at the per-node entrypoints the pipeline exposes. + assert "fnnx_model_node__run" in header.split("*/", 1)[0] + assert result.report["dim_bindings"] == {"batch": 1} + assert [node["id"] for node in result.report["nodes"]] == [ + "linreg", + "linreg2", + "linreg3", + "concat_reduce", + ] + # Builds the header under `-std=c99 -Wall -Wextra -Werror -Werror=vla`. + result.load() + + +def test_the_compiled_pipeline_matches_the_python_runtime(tmp_path): + compiled = compile_bundle(PIPELINE_BUNDLE, tmp_path).load() + feeds = {"x": _values((1, 3))} + + _assert_matches(compiled.run(feeds), _runtime_outputs(PIPELINE_BUNDLE, feeds)) + + +def test_a_dimension_binding_sizes_every_buffer_of_the_pipeline(tmp_path): + result = compile_bundle(PIPELINE_BUNDLE, tmp_path, dim_bindings={"batch": 4}) + header = result.header_path.read_text() + feeds = {"x": _values((4, 3))} + + outputs = result.load().run(feeds) + + assert result.report["dim_bindings"] == {"batch": 4} + assert result.report["options"]["dim_bindings"] == {"batch": 4} + assert result.report["entrypoint"]["inputs"][0]["shape"] == [4, 3] + assert "#define FNNX_MODEL_INPUT_X_DIM_0 4" in header + assert "#define FNNX_MODEL_OUTPUT_Y4_DIM_0 4" in header + assert "#define FNNX_MODEL_NODE_LINREG_INPUT_FLOAT_INPUT_DIM_0 4" in header + _assert_matches(outputs, _runtime_outputs(PIPELINE_BUNDLE, feeds)) + + +def test_a_tar_packaged_bundle_compiles_like_its_directory(tmp_path): + from_directory = compile_bundle(PIPELINE_BUNDLE, tmp_path / "directory") + from_tar = compile_bundle(PIPELINE_TAR, tmp_path / "tar") + feeds = {"x": _values((1, 3))} + + outputs = from_tar.load().run(feeds) + + # The two artifacts differ only in the file the preamble names as their source. + for key in ("prefix", "entrypoint", "nodes", "memory", "dim_bindings"): + assert from_tar.report[key] == from_directory.report[key], key + _assert_matches(outputs, _runtime_outputs(PIPELINE_BUNDLE, feeds)) + + +def test_repeated_runs_carry_no_state_between_them(tmp_path): + compiled = compile_bundle(PIPELINE_BUNDLE, tmp_path).load() + first = {"x": _values((1, 3), seed=1)} + second = {"x": _values((1, 3), seed=2)} + + first_outputs = compiled.run(first) + second_outputs = compiled.run(second) + repeated = compiled.run(first) + + _assert_matches(first_outputs, _runtime_outputs(PIPELINE_BUNDLE, first)) + _assert_matches(second_outputs, _runtime_outputs(PIPELINE_BUNDLE, second)) + _assert_matches(repeated, _runtime_outputs(PIPELINE_BUNDLE, first)) + + +def test_a_node_entrypoint_computes_what_the_node_computes(tmp_path): + """The `linreg` entrypoint alone, against the reference evaluator on its own model.""" + compiled = compile_bundle(PIPELINE_BUNDLE, tmp_path).load() + model = onnx.load(str(PIPELINE_BUNDLE / "ops_artifacts" / "linreg" / "model.onnx")) + feeds = {"float_input": _values((1, 3))} + + outputs = compiled.run_node("linreg", feeds) + + expected = ReferenceEvaluator(model).run(None, feeds) + _assert_matches(outputs, {"variable": expected[0]}) + + +def test_compiling_a_bundle_twice_is_byte_identical(tmp_path): + first = compile_bundle(PIPELINE_BUNDLE, tmp_path / "first") + second = compile_bundle(PIPELINE_BUNDLE, tmp_path / "second") + + assert first.header_path.read_bytes() == second.header_path.read_bytes() + assert first.report_path.read_bytes() == second.report_path.read_bytes() + + +def test_the_pipeline_glue_allocates_nothing(tmp_path): + header = compile_bundle(PIPELINE_BUNDLE, tmp_path).header_path.read_text() + + for token in ("malloc", "calloc", "realloc", "free", "alloca"): + assert not re.search(rf"\b{token}\b", header), token + + +def test_a_kernel_two_nodes_share_is_emitted_once(tmp_path): + """`linreg`, `linreg2` and `linreg3` are the same op at the same types.""" + result = compile_bundle(PIPELINE_BUNDLE, tmp_path) + header = result.header_path.read_text() + + kernels = result.report["kernels"] + + assert len(kernels) == len(set(kernels)) + scorers = [name for name in kernels if "ml_scores" in name] + assert len(scorers) == 1 + assert header.count(f"static void {scorers[0]}(") == 1 + + +def test_the_prefix_defaults_to_the_manifest_name(tmp_path, diamond_bundle): + """`onnx_pipeline.fnnx` carries no name, so only a bundle that does shows this.""" + manifest = json.loads((diamond_bundle / "manifest.json").read_text()) + manifest["name"] = "my diamond" + (diamond_bundle / "manifest.json").write_text( + json.dumps(manifest), encoding="utf-8" + ) + + result = compile_bundle(diamond_bundle, tmp_path / "out") + + assert result.report["prefix"] == "my_diamond" + assert result.header_path.name == "my_diamond.h" + assert "int my_diamond_run(" in result.header_path.read_text() + assert result.report["nodes"][0]["symbol"] == "my_diamond_node_head_run" + result.load() + + +def test_an_explicit_prefix_renames_every_entrypoint(tmp_path): + result = compile_bundle(PIPELINE_BUNDLE, tmp_path, prefix="my model") + header = result.header_path.read_text() + + assert result.report["prefix"] == "my_model" + assert "int my_model_run(" in header + assert "int my_model_node_linreg_run(" in header + result.load() + + +# -------------------------------------------------------------------------------------- +# Pipeline wiring +# -------------------------------------------------------------------------------------- + + +def test_a_fan_out_edge_that_is_also_an_output_reaches_every_consumer( + tmp_path, diamond_bundle +): + compiled = compile_bundle(diamond_bundle, tmp_path / "out").load() + feeds = {"x": _values((1, 3))} + + outputs = compiled.run(feeds) + + _assert_matches(outputs, _runtime_outputs(diamond_bundle, feeds)) + + +def test_a_node_entrypoint_matches_what_the_pipeline_computes(tmp_path, diamond_bundle): + """`head`'s output is a manifest output too, so both routes are observable.""" + compiled = compile_bundle(diamond_bundle, tmp_path / "out").load() + feeds = {"x": _values((1, 3))} + + inside = compiled.run(feeds)["h"] + alone = compiled.run_node("head", {"x": feeds["x"]})["y"] + + np.testing.assert_array_equal(alone, inside) + + +def test_nodes_are_ordered_by_their_edges_not_by_their_declaration( + tmp_path, diamond_bundle +): + """Declared back to front, the nodes still compile — and run — in dependency order.""" + nodes = _diamond_nodes() + shuffled = _write_bundle( + tmp_path / "shuffled.fnnx", + [nodes[3], nodes[2], nodes[1], nodes[0]], + [_manifest_tensor("x")], + [_manifest_tensor("h"), _manifest_tensor("out")], + ) + feeds = {"x": _values((1, 3))} + + result = compile_bundle(shuffled, tmp_path / "out") + + assert [node["id"] for node in result.report["nodes"]] == [ + "head", + "right", + "left", + "join", + ] + _assert_matches(result.load().run(feeds), _runtime_outputs(diamond_bundle, feeds)) + + +def test_a_manifest_shape_that_agrees_with_the_nodes_is_accepted(tmp_path): + bundle = _write_bundle( + tmp_path / "declared.fnnx", + _diamond_nodes(), + [_manifest_tensor("x", ["batch", 3])], + [_manifest_tensor("h", ["batch", 3]), _manifest_tensor("out", ["batch", 3])], + ) + feeds = {"x": _values((2, 3))} + + result = compile_bundle(bundle, tmp_path / "out", dim_bindings={"batch": 2}) + + _assert_matches(result.load().run(feeds), _runtime_outputs(bundle, feeds)) + + +def test_one_op_instance_may_run_at_two_places_in_the_pipeline(tmp_path): + """Two pipeline nodes on one op instance share its entrypoint, called twice.""" + tensor = _spec(["batch", 3]) + scale = _Node( + id="scale", + model=_affine_model(np.array([2.0, 3.0, 4.0]), 1.0, name="scale"), + inputs=("x",), + outputs=("once",), + input_specs=(tensor,), + output_specs=(tensor,), + ) + again = _Node( + id="scale", + model=None, + inputs=("once",), + outputs=("twice",), + input_specs=(tensor,), + output_specs=(tensor,), + ) + bundle = _write_bundle( + tmp_path / "twice.fnnx", + [scale, again], + [_manifest_tensor("x")], + [_manifest_tensor("twice")], + ) + # `ops.json` holds the instance once; the variant config wires it twice. + ops = json.loads((bundle / "ops.json").read_text()) + (bundle / "ops.json").write_text(json.dumps(ops[:1]), encoding="utf-8") + feeds = {"x": _values((1, 3))} + + result = compile_bundle(bundle, tmp_path / "out") + + assert [node["id"] for node in result.report["nodes"]] == ["scale"] + _assert_matches(result.load().run(feeds), _runtime_outputs(bundle, feeds)) + + +DET_OPSET = 22 + + +def _determinant_model(order: int, *, name: str): + """`Det`, whose kernel works on scratch storage its call sites share.""" + return helper.make_model( + helper.make_graph( + [helper.make_node("Det", ["x"], ["y"], name="det")], + name, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, order, order])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1])], + ), + opset_imports=[helper.make_opsetid("", DET_OPSET)], + ) + + +def test_kernel_scratch_two_nodes_share_is_sized_for_the_larger(tmp_path): + """Two `Det` nodes share one kernel, and its working storage has to fit both. + + The smaller node is compiled first, so a merge that kept the first claim would leave + the larger node writing past the end of the buffer. + """ + nodes = [ + _Node( + id="small", + model=_determinant_model(2, name="small"), + inputs=("a",), + outputs=("da",), + input_specs=(_spec([1, 2, 2]),), + output_specs=(_spec([1]),), + opset=DET_OPSET, + ), + _Node( + id="large", + model=_determinant_model(5, name="large"), + inputs=("b",), + outputs=("db",), + input_specs=(_spec([1, 5, 5]),), + output_specs=(_spec([1]),), + opset=DET_OPSET, + ), + ] + bundle = _write_bundle( + tmp_path / "dets.fnnx", + nodes, + [_manifest_tensor("a"), _manifest_tensor("b")], + [_manifest_tensor("da"), _manifest_tensor("db")], + ) + feeds = {"a": _values((1, 2, 2)), "b": _values((1, 5, 5))} + + result = compile_bundle(bundle, tmp_path / "out") + + scratch = re.findall( + r"static float \w+_work\[(\d+)\];", result.header_path.read_text() + ) + assert scratch == ["25"] + _assert_matches(result.load().run(feeds), _runtime_outputs(bundle, feeds)) + + +def test_one_kernel_name_with_two_definitions_is_refused(): + """The invariant kernel sharing rests on: a name encodes everything its code needs. + + Within one graph the emitter enforces it; between the nodes of a pipeline, the merge + has to, or one node would silently run the other's code. + """ + + def program(definition: str) -> Any: + return codegen.Program( + prefix="p", + graph_name="g", + source="test", + opsets={}, + dim_bindings={}, + inputs=(), + outputs=(), + weights=(), + scratch=(), + functions=(kernels.CFunction("p_kernel", definition),), + body=(), + ) + + with pytest.raises(CompileError, match="emitted twice with different definitions"): + bundle_module._merged_functions( + [program("void a(void);"), program("void b(void);")] + ) + + +def test_a_node_weight_in_a_side_file_is_embedded_at_compile_time(tmp_path): + """`has_external_data` nodes: the side file is read while compiling, never afterwards.""" + nodes = _diamond_nodes() + weight = np.array([2.0, -3.0, 0.5], dtype=np.float32) + nodes[0].model = _affine_model(weight, 0.5, name="head") + for tensor in nodes[0].model.graph.initializer: + if tensor.name == "w": + onnx.external_data_helper.set_external_data(tensor, location="head_w.bin") + tensor.ClearField("raw_data") + bundle = _write_bundle( + tmp_path / "external.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + side_file = bundle / "ops_artifacts" / "head" / "head_w.bin" + side_file.write_bytes(weight.tobytes()) + feeds = {"x": _values((1, 3))} + expected = _runtime_outputs(bundle, feeds) + + compiled = compile_bundle(bundle, tmp_path / "out").load() + side_file.unlink() + + _assert_matches(compiled.run(feeds), expected) + + +def test_an_edge_no_node_reads_still_gets_a_buffer(tmp_path): + """`left` writes `l`, which nothing downstream reads and the manifest does not expose.""" + nodes = _diamond_nodes() + nodes[3].inputs = ("h", "r") + bundle = _write_bundle( + tmp_path / "deadend.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + feeds = {"x": _values((1, 3))} + + result = compile_bundle(bundle, tmp_path / "out") + + _assert_matches(result.load().run(feeds), _runtime_outputs(bundle, feeds)) + + +def test_a_pipeline_of_zero_element_tensors_runs(tmp_path, diamond_bundle): + """Binding the batch to 0 makes every edge empty; the buffers still have addresses.""" + result = compile_bundle(diamond_bundle, tmp_path / "out", dim_bindings={"batch": 0}) + feeds = {"x": np.zeros((0, 3), dtype=np.float32)} + + outputs = result.load().run(feeds) + + assert result.report["entrypoint"]["inputs"][0]["shape"] == [0, 3] + _assert_matches(outputs, _runtime_outputs(diamond_bundle, feeds)) + + +def test_an_input_no_node_reads_is_still_a_parameter(tmp_path): + """An unused pipeline input keeps its place in the signature and builds warning-free.""" + nodes = _diamond_nodes() + bundle = _write_bundle( + tmp_path / "unused.fnnx", + nodes, + [_manifest_tensor("x"), _manifest_tensor("spare", ["batch", 2])], + [_manifest_tensor("out")], + ) + + result = compile_bundle(bundle, tmp_path / "out") + + assert [tensor["name"] for tensor in result.report["entrypoint"]["inputs"]] == [ + "x", + "spare", + ] + assert "(void)spare;" in result.header_path.read_text() + result.load() + + +# -------------------------------------------------------------------------------------- +# Rejections +# -------------------------------------------------------------------------------------- + + +def _assert_nothing_written(output_dir: Path) -> None: + assert not output_dir.exists() or not list(output_dir.iterdir()) + + +def test_a_non_pipeline_bundle_is_rejected(tmp_path): + bundle = _write_bundle( + tmp_path / "pyfunc.fnnx", + _diamond_nodes(), + [_manifest_tensor("x")], + [_manifest_tensor("out")], + variant="pyfunc", + ) + output = tmp_path / "out" + + with pytest.raises(CompileError, match="compiles `pipeline` bundles.*`pyfunc`"): + compile_bundle(bundle, output) + _assert_nothing_written(output) + + +def test_a_node_op_the_compiler_cannot_compile_is_rejected(tmp_path): + nodes = _diamond_nodes() + nodes[1].op = "PyFunc_v1" + bundle = _write_bundle( + tmp_path / "custom.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + output = tmp_path / "out" + + with pytest.raises( + CompileError, match="`left` runs op `PyFunc_v1`.*no node compiler.*`ONNX_v1`" + ): + compile_bundle(bundle, output) + _assert_nothing_written(output) + + +@pytest.mark.parametrize( + ("mutate", "expected"), + [ + pytest.param( + lambda nodes, manifest: manifest.update( + dynamic_attributes=[{"name": "temperature", "description": "how hot"}] + ), + r"dynamic attributes.*the manifest declares `temperature`", + id="manifest", + ), + pytest.param( + lambda nodes, manifest: nodes[1].dynamic_attributes.update( + {"scale": {"name": "scale", "default_value": "1"}} + ), + r"dynamic attributes.*op instance `left` declares `scale`", + id="op-instance", + ), + pytest.param( + lambda nodes, manifest: nodes[2].extra_dynattrs.update({"scale": "other"}), + r"dynamic attributes.*pipeline node `right` declares `scale`", + id="pipeline-node", + ), + ], +) +def test_dynamic_attributes_are_rejected(tmp_path, mutate, expected): + nodes = _diamond_nodes() + manifest: dict[str, Any] = {} + mutate(nodes, manifest) + bundle = _write_bundle( + tmp_path / "dynattrs.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + manifest_dynamic_attributes=manifest.get("dynamic_attributes", ()), + ) + output = tmp_path / "out" + + with pytest.raises(CompileError, match=expected): + compile_bundle(bundle, output) + _assert_nothing_written(output) + + +@pytest.mark.parametrize( + ("tensor", "expected"), + [ + pytest.param( + {"name": "x", "content_type": "JSON", "dtype": "MyRecord"}, + "content type `JSON`", + id="json-content-type", + ), + pytest.param( + _manifest_tensor("x") | {"dtype": "NDContainer[MyRecord]"}, + r"dtype `NDContainer\[MyRecord\]`.*`Array\[\.\.\.\]`", + id="ndcontainer", + ), + pytest.param( + _manifest_tensor("x") | {"dtype": "Array[string]"}, + "element type `string`.*does not support", + id="runtime-strings", + ), + pytest.param( + _manifest_tensor("x") | {"dtype": "Array[float16]"}, + "element type `float16`.*does not support", + id="float16", + ), + ], +) +def test_non_tensor_pipeline_io_is_rejected(tmp_path, tensor, expected): + bundle = _write_bundle( + tmp_path / "io.fnnx", + _diamond_nodes(), + [tensor], + [_manifest_tensor("out")], + ) + output = tmp_path / "out" + + with pytest.raises(CompileError, match=f"Manifest input `x`.*{expected}"): + compile_bundle(bundle, output) + _assert_nothing_written(output) + + +def test_two_manifest_entries_sharing_a_name_are_rejected(tmp_path): + """Two parameters of one name: the second is wired, the first reaches nothing.""" + bundle = _write_bundle( + tmp_path / "duplicate.fnnx", + _diamond_nodes(), + [_manifest_tensor("x"), _manifest_tensor("x")], + [_manifest_tensor("out")], + ) + output = tmp_path / "out" + + with pytest.raises(CompileError, match="Manifest input `x` is declared twice"): + compile_bundle(bundle, output) + _assert_nothing_written(output) + + +def test_a_dimension_written_as_minus_one_is_rejected(tmp_path): + """`-1` is how ONNX writes an unknown dimension; here it names nothing to bind.""" + nodes = _diamond_nodes() + nodes[0].input_specs = (_spec([-1, 3]),) + bundle = _write_bundle( + tmp_path / "negative.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, match="`head` input 0 has shape entry -1, which is neither" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_node_spec_dtype_the_compiler_cannot_hold_is_rejected(tmp_path): + nodes = _diamond_nodes() + nodes[1].input_specs = (_spec(["batch", 3], dtype="string"),) + bundle = _write_bundle( + tmp_path / "strings.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, match="Op instance `left` input 0 has element type `string`" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_an_edge_no_node_produces_is_rejected(tmp_path): + nodes = _diamond_nodes() + nodes[3].inputs = ("l", "missing") + bundle = _write_bundle( + tmp_path / "dangling.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, + match="`join` reads `missing`, which no manifest input and no node", + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_cycle_between_nodes_is_rejected(tmp_path): + nodes = _diamond_nodes() + nodes[1].inputs = ("out",) + bundle = _write_bundle( + tmp_path / "cycle.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises(CompileError, match="cycle.*`left`.*`join`"): + compile_bundle(bundle, tmp_path / "out") + + +def test_two_nodes_writing_one_edge_are_rejected(tmp_path): + nodes = _diamond_nodes() + nodes[2].outputs = ("l",) + bundle = _write_bundle( + tmp_path / "double.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises(CompileError, match="`left` and `right` both write `l`"): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_node_writing_a_manifest_input_is_rejected(tmp_path): + """The caller's input buffer is `const`, and overwriting it would lose what it holds.""" + nodes = _diamond_nodes() + nodes[1].outputs = ("x",) + bundle = _write_bundle( + tmp_path / "overwrite.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, match="`left` writes `x`, which is also a manifest input" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_an_output_no_node_produces_is_rejected(tmp_path): + bundle = _write_bundle( + tmp_path / "unproduced.fnnx", + _diamond_nodes(), + [_manifest_tensor("x")], + [_manifest_tensor("out"), _manifest_tensor("elsewhere")], + ) + + with pytest.raises( + CompileError, match="output `elsewhere` is produced by no pipeline node" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_an_output_that_names_a_pipeline_input_is_rejected(tmp_path): + """A pass-through output would have the caller's output buffer take the input's name.""" + bundle = _write_bundle( + tmp_path / "passthrough.fnnx", + _diamond_nodes(), + [_manifest_tensor("x")], + [_manifest_tensor("out"), _manifest_tensor("x")], + ) + + with pytest.raises( + CompileError, match="output `x` is produced by no pipeline node" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_manifest_shape_that_contradicts_the_nodes_is_rejected(tmp_path): + bundle = _write_bundle( + tmp_path / "contradiction.fnnx", + _diamond_nodes(), + [_manifest_tensor("x", ["batch", 5])], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, match=r"input `x` declares shape \[1, 5\].*takes \[1, 3\]" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_manifest_dtype_that_contradicts_the_nodes_is_rejected(tmp_path): + bundle = _write_bundle( + tmp_path / "dtype.fnnx", + _diamond_nodes(), + [_manifest_tensor("x", dtype="int64")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, + match=r"input `x` is declared `Array\[int64\]`.*takes `Array\[float32\]`", + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_an_input_nothing_reads_and_nothing_sizes_is_rejected(tmp_path): + """Without a consumer and without a declared shape, nothing says how big `spare` is.""" + bundle = _write_bundle( + tmp_path / "unsized.fnnx", + _diamond_nodes(), + [_manifest_tensor("x"), _manifest_tensor("spare")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, + match="input `spare` is read by no pipeline node and declares no shape", + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_an_edge_the_two_ends_disagree_about_is_rejected(tmp_path): + """One buffer cannot be [batch, 3] where it is written and [batch, 2] where it is read. + + Each node here agrees with its own ONNX graph; what disagrees is the two ends of `h`. + """ + nodes = _diamond_nodes() + nodes[1].model = _affine_model(np.array([1.0, 2.0]), 0.0, name="left") + nodes[1].input_specs = (_spec(["batch", 2]),) + nodes[1].output_specs = (_spec(["batch", 2]),) + nodes[3].inputs = ("r", "r") + bundle = _write_bundle( + tmp_path / "mismatch.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, + match=r"edge `h` is float32\[1, 3\] as output 0 of node `head`, but " + r"float32\[1, 2\] as input 0 of node `left`", + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_node_spec_that_contradicts_its_onnx_graph_is_rejected(tmp_path): + nodes = _diamond_nodes() + nodes[0].input_specs = (_spec(["batch", 3]), _spec(["batch", 3])) + nodes[0].inputs = ("x", "x") + bundle = _write_bundle( + tmp_path / "arity.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, match="`head`: its op spec declares 2 input.*ONNX graph has 1" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_node_spec_dimension_that_contradicts_its_graph_is_rejected(tmp_path): + nodes = _diamond_nodes() + nodes[0].input_specs = (_spec(["batch", 4]),) + bundle = _write_bundle( + tmp_path / "dims.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, match="input 0 .*has size 3 on axis 1 in the ONNX graph, but 4" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_node_spec_dtype_that_contradicts_its_graph_is_rejected(tmp_path): + nodes = _diamond_nodes() + nodes[0].input_specs = (_spec(["batch", 3], dtype="int64"),) + bundle = _write_bundle( + tmp_path / "elemtype.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, + match="`head`: input 0 .*is `float32` in the ONNX graph, but `int64` in the op", + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_node_wired_to_the_wrong_number_of_edges_is_rejected(tmp_path): + nodes = _diamond_nodes() + nodes[3].inputs = ("l",) + bundle = _write_bundle( + tmp_path / "wiring.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises( + CompileError, match="`join` is wired to 1 input.*op spec declares 2" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_missing_node_model_is_rejected(tmp_path): + nodes = _diamond_nodes() + bundle = _write_bundle( + tmp_path / "incomplete.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + (bundle / "ops_artifacts" / "left" / "model.onnx").unlink() + + with pytest.raises(CompileError, match="`left`: `model.onnx` is missing"): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_node_the_ops_file_does_not_define_is_rejected(tmp_path): + bundle = _write_bundle( + tmp_path / "unknown.fnnx", + _diamond_nodes(), + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + config = json.loads((bundle / "variant_config.json").read_text()) + config["nodes"][1]["op_instance_id"] = "ghost" + (bundle / "variant_config.json").write_text(json.dumps(config), encoding="utf-8") + + with pytest.raises( + CompileError, match="op instance `ghost`, which `ops.json` does not define" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_bundle_path_that_does_not_exist_is_rejected(tmp_path): + with pytest.raises(CompileError, match="FNNX bundle not found"): + compile_bundle(tmp_path / "nowhere.fnnx", tmp_path / "out") + + +def test_a_malformed_bundle_file_is_rejected(tmp_path): + bundle = _write_bundle( + tmp_path / "broken.fnnx", + _diamond_nodes(), + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + (bundle / "manifest.json").write_text("{not json", encoding="utf-8") + + with pytest.raises(CompileError, match="Could not read `manifest.json`"): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_manifest_missing_a_required_field_is_rejected(tmp_path): + bundle = _write_bundle( + tmp_path / "invalid.fnnx", + _diamond_nodes(), + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + manifest = json.loads((bundle / "manifest.json").read_text()) + del manifest["producer_name"] + (bundle / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8") + + with pytest.raises(CompileError, match="`manifest.json` is not valid"): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_node_requiring_ort_extensions_is_rejected(tmp_path): + bundle = _write_bundle( + tmp_path / "extensions.fnnx", + _diamond_nodes(), + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + ops = json.loads((bundle / "ops.json").read_text()) + ops[1]["attributes"]["requires_ort_extensions"] = True + (bundle / "ops.json").write_text(json.dumps(ops), encoding="utf-8") + + with pytest.raises( + CompileError, match="`left`: the op requires the onnxruntime extensions" + ): + compile_bundle(bundle, tmp_path / "out") + + +def test_a_node_the_onnx_core_rejects_names_the_op_instance(tmp_path): + """A failure inside a node's graph still points at the node it came from.""" + nodes = _diamond_nodes() + nodes[1].model = helper.make_model( + helper.make_graph( + [helper.make_node("NonZero", ["x"], ["y"], name="nonzero")], + "left", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["batch", 3])], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2, "count"])], + ), + opset_imports=[helper.make_opsetid("", OPSET)], + ) + nodes[1].output_specs = (_spec([2, 3], dtype="int64"),) + nodes[2].inputs = ("h",) + nodes[3].input_specs = (_spec([2, 3], dtype="int64"), _spec(["batch", 3])) + bundle = _write_bundle( + tmp_path / "nonzero.fnnx", + nodes, + [_manifest_tensor("x")], + [_manifest_tensor("out")], + ) + + with pytest.raises(CompileError, match="Op instance `left`.*`nonzero`.*NonZero"): + compile_bundle(bundle, tmp_path / "out") + + +# -------------------------------------------------------------------------------------- +# Packaging +# -------------------------------------------------------------------------------------- + + +def test_a_tar_bundle_leaves_no_temporary_directory_behind(tmp_path, monkeypatch): + """The bundle is unpacked to compile it; what is unpacked has to be cleaned up.""" + unpacked: list[Path] = [] + original = bundle_module.unpack_model + + def record(path): + directory, temporary = original(path) + unpacked.append(Path(directory)) + return directory, temporary + + monkeypatch.setattr(bundle_module, "unpack_model", record) + packed = tmp_path / "packed.tar" + with tarfile.open(packed, "w") as archive: + for entry in PIPELINE_BUNDLE.iterdir(): + archive.add(entry, arcname=entry.name) + + compile_bundle(packed, tmp_path / "out") + + assert unpacked and not any(directory.exists() for directory in unpacked) diff --git a/src/python/tests/test_extra_compiler_cli.py b/src/python/tests/test_extra_compiler_cli.py new file mode 100644 index 0000000..719d94b --- /dev/null +++ b/src/python/tests/test_extra_compiler_cli.py @@ -0,0 +1,257 @@ +"""The command-line entry `python -m fnnx.extras.compilers.c` and the report it summarizes. + +What a compiled artifact computes is the business of the conformance, differential and +bundle suites; what is pinned here is the command line — which files land where, which +options reach the compiler, what the summary says, and what a failure looks like on stderr. +The one execution check compares against the ONNX reference evaluator, never against a +hand-written expectation. +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +import pytest + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from fnnx.extras.compilers.c.__main__ import main # noqa: E402 +from onnx import TensorProto, helper # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +OPSET = 21 +SEED = 20260726 + +MODELS = Path(__file__).parent / "models" +PIPELINE_BUNDLE = MODELS / "onnx_pipeline.fnnx" + +needs_c_compiler = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + + +def _relu_model( + path: Path, name: str = "tiny", shape: tuple[str | int, ...] = ("batch", 3) +) -> Path: + """A one-node model with a symbolic batch dimension, saved as a `.onnx` file.""" + graph = helper.make_graph( + [helper.make_node("Relu", ["x"], ["y"], name="relu")], + name, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, list(shape))], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, list(shape))], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", OPSET)]) + onnx.save(model, path) + return path + + +def _nonzero_model(path: Path) -> Path: + """`NonZero`'s output shape depends on the data, so this can never compile.""" + graph = helper.make_graph( + [helper.make_node("NonZero", ["x"], ["y"], name="find")], + "dynamic", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [4])], + [helper.make_tensor_value_info("y", TensorProto.INT64, [1, "n"])], + ) + onnx.save( + helper.make_model(graph, opset_imports=[helper.make_opsetid("", OPSET)]), path + ) + return path + + +def _summary_fields(stdout: str) -> dict[str, str]: + """The summary's `label: value` lines, keyed by label; the heading line dropped.""" + return { + label.strip(): value.strip() + for label, _, value in (line.partition(":") for line in stdout.splitlines()[1:]) + } + + +def _report(directory: Path, prefix: str) -> dict[str, Any]: + return json.loads((directory / f"{prefix}_report.json").read_text()) + + +def test_the_module_entry_compiles_a_bundle(tmp_path): + """The advertised invocation, run as its own process: `python -m ...`.""" + completed = subprocess.run( + [ + sys.executable, + "-m", + "fnnx.extras.compilers.c", + str(PIPELINE_BUNDLE), + "-o", + str(tmp_path), + ], + capture_output=True, + text=True, + ) + + assert completed.returncode == 0, completed.stderr + assert sorted(path.name for path in tmp_path.iterdir()) == [ + "fnnx_model.h", + "fnnx_model_report.json", + ] + assert str(tmp_path / "fnnx_model.h") in completed.stdout + assert str(tmp_path / "fnnx_model_report.json") in completed.stdout + + +@needs_c_compiler +def test_a_raw_onnx_model_compiles_and_runs(tmp_path): + model_path = _relu_model(tmp_path / "tiny.onnx") + + assert main([str(model_path), "-o", str(tmp_path / "out"), "--dim", "batch=2"]) == 0 + + compiled = harness.load_compiled(tmp_path / "out" / "tiny_report.json") + feeds = {"x": np.random.default_rng(SEED).normal(size=(2, 3)).astype(np.float32)} + reference = ReferenceEvaluator(str(model_path)) + expected = dict(zip(reference.output_names, reference.run(None, feeds))) + np.testing.assert_allclose( + compiled.run(feeds)["y"], expected["y"], rtol=1e-6, atol=1e-6 + ) + + +def test_the_dim_flag_binds_symbolic_dimensions(tmp_path): + assert main([str(PIPELINE_BUNDLE), "-o", str(tmp_path), "--dim", "batch=4"]) == 0 + + report = _report(tmp_path, "fnnx_model") + assert report["options"]["dim_bindings"] == {"batch": 4} + assert report["dim_bindings"] == {"batch": 4} + assert report["entrypoint"]["inputs"][0]["shape"] == [4, 3] + assert ( + "#define FNNX_MODEL_INPUT_X_DIM_0 4" in (tmp_path / "fnnx_model.h").read_text() + ) + + +def test_the_dim_flag_repeats_once_per_symbolic_dimension(tmp_path): + model_path = _relu_model(tmp_path / "pair.onnx", "pair", ("batch", "features")) + output = tmp_path / "out" + flags = ["--dim", "batch=2", "--dim", "features=5"] + + assert main([str(model_path), "-o", str(output), *flags]) == 0 + + report = _report(output, "pair") + assert report["dim_bindings"] == {"batch": 2, "features": 5} + assert report["entrypoint"]["inputs"][0]["shape"] == [2, 5] + + +def test_unbound_dimensions_default_to_one(tmp_path, capsys): + assert main([str(PIPELINE_BUNDLE), "-o", str(tmp_path)]) == 0 + + assert _report(tmp_path, "fnnx_model")["dim_bindings"] == {"batch": 1} + assert _summary_fields(capsys.readouterr().out)["dimensions"] == "batch=1" + + +def test_the_prefix_flag_names_the_files_and_the_symbols(tmp_path): + model_path = _relu_model(tmp_path / "tiny.onnx") + output = tmp_path / "out" + + assert main([str(model_path), "-o", str(output), "--prefix", "my model"]) == 0 + + assert sorted(path.name for path in output.iterdir()) == [ + "my_model.h", + "my_model_report.json", + ] + assert "int my_model_run(" in (output / "my_model.h").read_text() + assert _report(output, "my_model")["options"]["prefix"] == "my model" + + +def test_the_summary_reports_footprint_kernels_opsets_and_bindings(tmp_path, capsys): + assert main([str(PIPELINE_BUNDLE), "-o", str(tmp_path), "--dim", "batch=2"]) == 0 + + fields = _summary_fields(capsys.readouterr().out) + report = _report(tmp_path, "fnnx_model") + memory = report["memory"] + assert fields["header"] == str(tmp_path / "fnnx_model.h") + assert fields["report"] == str(tmp_path / "fnnx_model_report.json") + assert fields["entrypoint"] == f"{report['entrypoint']['symbol']}()" + assert fields["opsets"] == "ai.onnx=21, ai.onnx.ml=1" + assert fields["opsets"] == ", ".join( + f"{domain}={version}" for domain, version in report["opsets"].items() + ) + assert fields["dimensions"] == "batch=2" + assert fields["kernels"] == str(len(report["kernels"])) + assert fields["static memory"] == ( + f"{memory['static_bytes']} bytes " + f"(weights {memory['weights_bytes']}, arena {memory['arena_bytes']})" + ) + assert memory["static_bytes"] == memory["weights_bytes"] + memory["arena_bytes"] + + +@pytest.mark.parametrize("flag", ["batch", "batch=four", "batch=", "=4", "batch=4.0"]) +def test_a_malformed_dim_flag_is_rejected_before_compiling(tmp_path, flag, capsys): + output = tmp_path / "out" + + with pytest.raises(SystemExit) as failure: + main([str(PIPELINE_BUNDLE), "-o", str(output), "--dim", flag]) + + assert failure.value.code == 2 + assert "--dim" in capsys.readouterr().err + assert not output.exists() + + +def test_the_output_directory_is_required(tmp_path, capsys): + with pytest.raises(SystemExit) as failure: + main([str(PIPELINE_BUNDLE)]) + + assert failure.value.code == 2 + assert "--output-dir" in capsys.readouterr().err + + +def test_a_negative_dim_binding_is_reported_as_a_compile_error(tmp_path, capsys): + output = tmp_path / "out" + + assert main([str(PIPELINE_BUNDLE), "-o", str(output), "--dim", "batch=-1"]) == 1 + + streams = capsys.readouterr() + assert streams.out == "" + assert streams.err.startswith("error: ") + assert "batch" in streams.err + assert not output.exists() + + +def test_an_unsupported_model_fails_without_writing_anything(tmp_path, capsys): + model_path = _nonzero_model(tmp_path / "dynamic.onnx") + output = tmp_path / "out" + + assert main([str(model_path), "-o", str(output)]) == 1 + + streams = capsys.readouterr() + assert streams.out == "" + assert streams.err.startswith("error: ") + assert "find" in streams.err and "NonZero" in streams.err + assert not output.exists() + + +def test_an_output_path_that_is_not_a_directory_is_reported_cleanly(tmp_path, capsys): + occupied = tmp_path / "taken" + occupied.write_text("not a directory") + + assert main([str(PIPELINE_BUNDLE), "-o", str(occupied)]) == 1 + + streams = capsys.readouterr() + assert streams.out == "" + assert streams.err.startswith("error: ") + assert str(occupied) in streams.err + assert "Traceback" not in streams.err + assert occupied.read_text() == "not a directory" + + +@pytest.mark.parametrize("missing", ["absent.fnnx", "absent.onnx"]) +def test_a_missing_source_is_reported_cleanly(tmp_path, missing, capsys): + output = tmp_path / "out" + + assert main([str(tmp_path / missing), "-o", str(output)]) == 1 + + streams = capsys.readouterr() + assert streams.err.startswith("error: ") + assert missing in streams.err + assert "Traceback" not in streams.err + assert not output.exists() diff --git a/src/python/tests/test_extra_compiler_conformance.py b/src/python/tests/test_extra_compiler_conformance.py new file mode 100644 index 0000000..3e9a3a0 --- /dev/null +++ b/src/python/tests/test_extra_compiler_conformance.py @@ -0,0 +1,634 @@ +"""ONNX backend conformance: the shipped node corpus is the oracle. + +Every expected value comes from the test data the `onnx` package ships, and every +comparison goes through ONNX's own `Runner.assert_similar_outputs` under the tolerances +ONNX records per test. Nothing in this module decides what an op should compute. + +The suite is fail-closed. Each enumerated corpus test is either + +* **ratcheted** — compiled, built, executed and compared, where any error in that chain is + a failure, or +* **ledgered** — excluded for a reason the governance check re-derives from the model + itself, so an entry cannot be invented to silence a failing test, + +and never neither nor both. The only skip is environmental and takes the whole module with +it: no `onnx`, no C compiler, or an `onnx` other than the pinned one whose schema set and +corpus the ledger and pass list are keyed to. +""" + +from __future__ import annotations + +import json +import shutil +import tempfile +from collections.abc import Iterator, Mapping, Sequence +from functools import cache +from pathlib import Path +from typing import Any + +import pytest + +onnx = pytest.importorskip("onnx") +# The harness refuses to import without numpy, so this covers both dependencies. +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from onnx import ( # noqa: E402 + GraphProto, + ModelProto, + NodeProto, + TensorProto, + TypeProto, + ValueInfoProto, + helper, + numpy_helper, +) +from onnx.backend.test.loader import load_model_tests # noqa: E402 +from onnx.backend.test.runner import Runner # noqa: E402 + +from fnnx.extras.compilers.c import compile_onnx # noqa: E402 +from fnnx.extras.compilers.c.errors import CompileError # noqa: E402 +from fnnx.extras.compilers.c.onnx.dtypes import C_TYPES # noqa: E402 +from fnnx.extras.compilers.c.onnx.folding import fold_constants # noqa: E402 +from fnnx.extras.compilers.c.onnx.functions import ( # noqa: E402 + MAX_EXPANSION_DEPTH, + function_model, +) +from fnnx.extras.compilers.c.onnx.kernels import KERNELS # noqa: E402 +from fnnx.extras.compilers.c.onnx.loader import ( # noqa: E402 + SUPPORTED_DOMAINS, + normalize_domain, + resolve_opsets, +) +from fnnx.extras.compilers.c.onnx.shapes import ( # noqa: E402 + infer_shapes, + runtime_shape_operand, + tensor_types, +) +from fnnx.extras.compilers.c.onnx.verify import ( # noqa: E402 + CONTROL_FLOW_OPS, + DATA_DEPENDENT_SHAPE_OPS, +) + +# The ledger, the pass list and the disposition of every op are keyed to one exact `onnx` +# release: it defines both the schema set the compiler dispatches against and the corpus +# this suite enumerates. Running under another release is not a weaker run of this suite, +# it is a run of a different suite, so the module steps aside rather than reporting on it. +PINNED_ONNX = "1.22" + +CONFORMANCE_DIR = Path(__file__).parent / "conformance" +LEDGER_PATH = CONFORMANCE_DIR / "ledger.json" +RATCHET_PATH = CONFORMANCE_DIR / "passing.txt" +TOLERANCES_PATH = CONFORMANCE_DIR / "tolerances.json" + +# The closed set of reasons a corpus test may be excluded, in the order `_classify` applies +# them. All but the last are structural — properties of the model that no amount of kernel +# work changes — and cover the compiler's v1 unsupported surface. `op-not-implemented` is +# the milestone category: the kernel tasks drive it down, and the governance check keeps it +# honest by refusing it for any op the kernel registry serves. What is left of it is +# accounted for op by op in `conformance/dispositions.json`, whose own suite holds the table +# to these same categories and refuses an entry that would exempt an op the compiler serves. +LEDGER_CATEGORIES = ( + "out-of-scope-domain", + "non-tensor-io", + "runtime-strings", + "unsupported-dtype", + "control-flow", + "data-dependent-shape", + "random-op", + "external-codec", + "op-not-implemented", +) + +RANDOM_OPS = frozenset( + { + "Bernoulli", + "Multinomial", + "RandomNormal", + "RandomNormalLike", + "RandomUniform", + "RandomUniformLike", + } +) +CODEC_OPS = frozenset({"ImageDecoder"}) + +# What `onnx.backend.test.loader.load_model_tests` applies to a test that ships no +# `data.json`; a test asserts these are still the corpus-wide defaults. +ONNX_DEFAULT_RTOL = 1e-3 +ONNX_DEFAULT_ATOL = 1e-7 + +# How far a per-op override may loosen those defaults. Two orders of magnitude covers what +# summation order and libm accuracy can cost a float32 kernel — `atol` reaches 1e-5, still +# well inside float32's ~1e-7 epsilon times a few thousand accumulations. Anything needing +# more is a wrong kernel, not a tolerance problem. +MAX_TOLERANCE_FACTOR = 100 + +pytestmark = [ + pytest.mark.skipif( + not onnx.__version__.startswith(f"{PINNED_ONNX}."), + reason=( + f"the conformance corpus and ledger are pinned to onnx {PINNED_ONNX}.*, " + f"but onnx {onnx.__version__} is installed" + ), + ), + pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", + ), +] + + +# -------------------------------------------------------------------------------------- +# The corpus, and the category a model may be ledgered under +# -------------------------------------------------------------------------------------- + + +@cache +def _corpus() -> dict[str, Any]: + """Every node test the pinned `onnx` package ships, by name.""" + return {case.name: case for case in load_model_tests(kind="node")} + + +def _subgraphs(node: NodeProto) -> Iterator[GraphProto]: + for attribute in node.attribute: + if attribute.HasField("g"): + yield attribute.g + yield from attribute.graphs + + +def _nodes(graph: GraphProto) -> Iterator[NodeProto]: + for node in graph.node: + yield node + for subgraph in _subgraphs(node): + yield from _nodes(subgraph) + + +def _value_infos(graph: GraphProto) -> Iterator[ValueInfoProto]: + yield from graph.input + yield from graph.output + yield from graph.value_info + for node in graph.node: + for subgraph in _subgraphs(node): + yield from _value_infos(subgraph) + + +def _classify(model: ModelProto) -> str | None: + """The category this model may be ledgered under, or None if it has to compile and run. + + Everything is read off the model proto, the compiler's folding pass and the kernel + registry, never off a compilation attempt: a compiler bug must surface as a failing + test, not as a test that quietly becomes ledgerable. + """ + graph = model.graph + nodes = list(_nodes(graph)) + domains = {normalize_domain(imported.domain) for imported in model.opset_import} + domains |= {normalize_domain(node.domain) for node in nodes} + if not domains <= set(SUPPORTED_DOMAINS): + return "out-of-scope-domain" + + infos = list(_value_infos(graph)) + if graph.sparse_initializer or any( + info.type.WhichOneof("value") not in (None, "tensor_type") for info in infos + ): + return "non-tensor-io" + + declared = { + info.type.tensor_type.elem_type + for info in infos + if info.type.WhichOneof("value") == "tensor_type" + } + declared |= {initializer.data_type for initializer in graph.initializer} + if TensorProto.STRING in declared: + return "runtime-strings" + # UNDEFINED is a tensor whose element type the model does not state, not one whose + # element type is unsupported; the compiler's own verification rejects it by name. + if not declared <= set(C_TYPES) | {TensorProto.UNDEFINED}: + return "unsupported-dtype" + + op_types = {node.op_type for node in nodes} + for category, family in ( + ("control-flow", CONTROL_FLOW_OPS), + ("data-dependent-shape", DATA_DEPENDENT_SHAPE_OPS), + ): + if op_types & family: + return category + if any(_draws_at_random(node, graph) for node in nodes): + return "random-op" + if op_types & CODEC_OPS: + return "external-codec" + + try: + opsets = resolve_opsets(model) + except CompileError: + return "op-not-implemented" + folded = _folded_graph(model) + types = tensor_types(folded) + constants = {initializer.name for initializer in folded.initializer} + # An op that takes its output shape from an operand's values — a reduction's axes — has a + # data-dependent output shape unless the graph fixes that operand, whatever kernel serves + # the op. That is structural, so it is derived before the milestone category below. + if any( + runtime_shape_operand(node, constants, types) is not None + for node in _nodes(folded) + ): + return "data-dependent-shape" + if any(not _serviceable(node, opsets, types) for node in _nodes(folded)): + return "op-not-implemented" + return None + + +def _draws_at_random(node: NodeProto, graph: GraphProto) -> bool: + """Whether the node's output is a draw rather than a function of its inputs. + + `Dropout` is one only in training mode: without a `training_mode` operand, or with one + the graph pins to false, it passes its input through. + """ + if node.op_type in RANDOM_OPS: + return True + if node.op_type != "Dropout" or len(node.input) < 3 or not node.input[2]: + return False + for initializer in graph.initializer: + if initializer.name == node.input[2]: + return bool(numpy_helper.to_array(initializer).any()) + return True + + +def _serviceable( + node: NodeProto, + opsets: Mapping[str, int], + types: Mapping[str, TypeProto], + depth: int = 0, +) -> bool: + """Whether the compiler has a way to compile this node, short of trying. + + A registered kernel, or — the fallback dispatch takes — the ONNX function body defining + the op, provided the compiler can serve what that body is made of in turn. Deliberately + "has a kernel at all" rather than "has one at this model's opset": a version the registry + cannot vouch for is a gap to close, not a reason to exempt. + + The body is built and folded exactly as `codegen` builds and folds it, through the + compiler's own passes, so that the nodes weighed here are the ones dispatch would really + reach: a `Shape` the body computes its result from resolves by folding rather than by a + kernel, and a context-dependent body — `CastLike`'s — only exists once its operands carry + types. What is deliberately *not* run is the static verification: whether the body would + compile is the question a failing test answers, not one a ledger entry may. + """ + domain = normalize_domain(node.domain) + if KERNELS.registered_versions(domain, node.op_type): + return True + opset_version = opsets.get(domain) + if opset_version is None or depth >= MAX_EXPANSION_DEPTH: + return False + try: + built = function_model( + node, domain, opset_version, [types.get(name) for name in node.input] + ) + except CompileError: + return False + if built is None or not built.model.graph.node: + return False + body_opsets = { + normalize_domain(imported.domain): imported.version + for imported in built.model.opset_import + } + folded = _folded_graph(built.model) + body_types = tensor_types(folded) + return all( + _serviceable(inner, body_opsets, body_types, depth + 1) + for inner in _nodes(folded) + ) + + +def _folded_graph(model: ModelProto) -> GraphProto: + """What is left for kernels once the compiler's own folding pass has run. + + `Shape`, `Size` and constant subgraphs are resolved by folding rather than by a kernel — + a disposition of its own — so a test of one is not a test of an unimplemented op, and + ledgering it would hide a test that passes today. The pass run here is the compiler's, + never a reimplementation of it. When it cannot run at all (ONNX's own shape inference + rejects a handful of corpus models outright) nothing folds, which is what the unfolded + graph already says. + """ + folded = ModelProto() + folded.CopyFrom(model) + try: + opsets = resolve_opsets(folded) + folded = infer_shapes(folded) + while fold_constants(folded, opsets): + folded = infer_shapes(folded) + except CompileError: + return model.graph + return folded.graph + + +@cache +def _derived_categories() -> dict[str, str | None]: + return { + name: _classify(onnx.load(Path(case.model_dir) / "model.onnx")) + for name, case in _corpus().items() + } + + +# -------------------------------------------------------------------------------------- +# The checked-in ledger, pass list and tolerance overrides +# -------------------------------------------------------------------------------------- + + +def _ledger() -> dict[str, str]: + return json.loads(LEDGER_PATH.read_text(encoding="utf-8")) + + +def _ratchet() -> tuple[str, ...]: + lines = RATCHET_PATH.read_text(encoding="utf-8").splitlines() + return tuple( + line.strip() + for line in lines + if line.strip() and not line.lstrip().startswith("#") + ) + + +def _overrides() -> dict[str, dict[str, Any]]: + return json.loads(TOLERANCES_PATH.read_text(encoding="utf-8"))["overrides"] + + +def _governance_problems( + ledger: Mapping[str, str], ratchet: Sequence[str] +) -> list[str]: + """Everything wrong with the ledger and pass list as a pair, worst case all of it.""" + corpus = _corpus() + derived = _derived_categories() + problems = [] + for name, category in sorted(ledger.items()): + if name not in corpus: + problems.append(f"`{name}` is ledgered but the corpus has no such test.") + elif category not in LEDGER_CATEGORIES: + problems.append( + f"`{name}` is ledgered as `{category}`, which is not one of the " + f"categories {', '.join(LEDGER_CATEGORIES)}." + ) + elif derived[name] != category: + actual = derived[name] + problems.append( + f"`{name}` is ledgered as `{category}`, but its model is " + + ( + f"`{actual}`." + if actual is not None + else "compilable: a supported op's test cannot be ledgered." + ) + ) + for name in sorted(set(ratchet) - set(corpus)): + problems.append( + f"`{name}` is in the pass list but the corpus has no such test." + ) + for name in sorted(set(ratchet) & set(ledger)): + problems.append(f"`{name}` is both ledgered and in the pass list.") + for name in sorted(set(corpus) - set(ledger) - set(ratchet)): + problems.append( + f"`{name}` is in neither the ledger nor the pass list; it has to pass, " + f"or be ledgered as `{derived[name]}`." + ) + if list(ratchet) != sorted(set(ratchet)): + problems.append("The pass list is not sorted, or lists a test twice.") + return problems + + +def _tolerance_problems(overrides: Mapping[str, Mapping[str, Any]]) -> list[str]: + problems = [] + for op_type, override in sorted(overrides.items()): + if not str(override.get("justification", "")).strip(): + problems.append( + f"The `{op_type}` tolerance override carries no written justification." + ) + for field, default in ( + ("rtol", ONNX_DEFAULT_RTOL), + ("atol", ONNX_DEFAULT_ATOL), + ): + value = override.get(field) + if value is None: + problems.append(f"The `{op_type}` tolerance override has no `{field}`.") + elif value > default * MAX_TOLERANCE_FACTOR: + problems.append( + f"The `{op_type}` tolerance override loosens `{field}` to {value}, " + f"beyond {MAX_TOLERANCE_FACTOR}x the ONNX default {default}." + ) + return problems + + +# -------------------------------------------------------------------------------------- +# Compiling, executing and comparing one corpus test +# -------------------------------------------------------------------------------------- + + +def _read_tensor(path: Path) -> Any: + tensor = TensorProto() + tensor.ParseFromString(path.read_bytes()) + return numpy_helper.to_array(tensor) + + +def _load_data_set(directory: Path, model: ModelProto) -> tuple[dict[str, Any], list]: + """Inputs by graph position, outputs in graph order — as ONNX's own runner reads them.""" + feeds = {} + for index, entry in enumerate(model.graph.input): + path = directory / f"input_{index}.pb" + if path.is_file(): + feeds[entry.name] = _read_tensor(path) + expected = [ + _read_tensor(directory / f"output_{index}.pb") + for index in range(len(model.graph.output)) + ] + return feeds, expected + + +def _tolerances(model: ModelProto, rtol: float, atol: float) -> tuple[float, float]: + op_types = {node.op_type for node in _nodes(model.graph)} + for op_type, override in sorted(_overrides().items()): + if op_type in op_types: + rtol = max(rtol, override["rtol"]) + atol = max(atol, override["atol"]) + return rtol, atol + + +def _execute(model_dir: Path, *, rtol: float, atol: float) -> None: + """Compile a corpus test, run every data set it ships, compare against its outputs.""" + model_path = model_dir / "model.onnx" + model = onnx.load(model_path) + rtol, atol = _tolerances(model, rtol, atol) + data_sets = sorted( + path for path in model_dir.iterdir() if path.name.startswith("test_data_set") + ) + assert data_sets, f"`{model_dir}` ships no test data set." + + with tempfile.TemporaryDirectory() as artifact_dir: + compiled = compile_onnx(model_path, artifact_dir).load() + for data_set in data_sets: + feeds, expected = _load_data_set(data_set, model) + outputs = compiled.run( + {spec.name: feeds[spec.name] for spec in compiled.inputs} + ) + Runner.assert_similar_outputs( + expected, + [outputs[entry.name] for entry in model.graph.output], + rtol=rtol, + atol=atol, + model_dir=str(model_dir), + ) + + +# -------------------------------------------------------------------------------------- +# The suite +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("name", _ratchet()) +def test_ratcheted_conformance_test_passes(name): + """Every test the ratchet records as passing still compiles, runs and matches.""" + case = _corpus().get(name) + assert case is not None, f"the corpus has no test `{name}`" + + _execute(Path(case.model_dir), rtol=case.rtol, atol=case.atol) + + +def test_every_corpus_test_is_ledgered_or_ratcheted(): + problems = _governance_problems(_ledger(), _ratchet()) + + assert not problems, "\n".join(problems) + + +def test_the_ledger_cannot_exempt_a_supported_op(): + """A ledger entry for an op the registry serves is caught, whatever reason it claims.""" + ledger = {**_ledger(), "test_add": "op-not-implemented"} + ratchet = [name for name in _ratchet() if name != "test_add"] + + problems = _governance_problems(ledger, ratchet) + + assert [ + problem + for problem in problems + if "test_add" in problem and "cannot be ledgered" in problem + ] + + +def test_the_ledger_cannot_exempt_an_op_folding_resolves(): + """`Shape` never needs a kernel, so `op-not-implemented` cannot be claimed for it.""" + ledger = {**_ledger(), "test_shape": "op-not-implemented"} + ratchet = [name for name in _ratchet() if name != "test_shape"] + + problems = _governance_problems(ledger, ratchet) + + assert [ + problem + for problem in problems + if "`test_shape`" in problem and "cannot be ledgered" in problem + ] + + +def test_the_ledger_cannot_exempt_an_op_function_expansion_serves(): + """No kernel serves `HardSwish`; the body ONNX defines for it compiles all the same.""" + assert not KERNELS.registered_versions("", "HardSwish") + ledger = {**_ledger(), "test_hardswish": "op-not-implemented"} + ratchet = [name for name in _ratchet() if name != "test_hardswish"] + + problems = _governance_problems(ledger, ratchet) + + assert [ + problem + for problem in problems + if "`test_hardswish`" in problem and "cannot be ledgered" in problem + ] + + +def test_ledgering_a_ratcheted_test_fails_the_suite(): + """The ratchet only grows: a passing test cannot be silenced by ledgering it.""" + ledger = {**_ledger(), "test_relu": "unsupported-dtype"} + + problems = _governance_problems(ledger, _ratchet()) + + assert [ + problem + for problem in problems + if "test_relu" in problem and "both ledgered and in the pass list" in problem + ] + + +def test_an_unaccounted_corpus_test_fails_the_suite(): + """A corpus test in neither file — an `onnx` upgrade's new tests — is not silence.""" + ledger = dict(_ledger()) + dropped = sorted(set(ledger) - set(_ratchet()))[0] + del ledger[dropped] + + problems = _governance_problems(ledger, _ratchet()) + + assert [ + problem + for problem in problems + if dropped in problem and "neither the ledger nor the pass list" in problem + ] + + +def test_a_reason_outside_the_closed_set_is_rejected(): + ledger = {**_ledger(), sorted(_ledger())[0]: "not yet implemented"} + + problems = _governance_problems(ledger, _ratchet()) + + assert [problem for problem in problems if "not one of the categories" in problem] + + +def test_the_recorded_tolerance_defaults_match_the_corpus(): + """The defaults the override bound rests on are ONNX's, not this suite's.""" + defaults = {(case.rtol, case.atol) for case in _corpus().values()} + + assert (ONNX_DEFAULT_RTOL, ONNX_DEFAULT_ATOL) in defaults + + +def test_the_tolerance_overrides_are_bounded_and_justified(): + problems = _tolerance_problems(_overrides()) + + assert not problems, "\n".join(problems) + + +def test_an_unjustified_or_unbounded_override_is_rejected(): + problems = _tolerance_problems( + { + "Conv": {"rtol": ONNX_DEFAULT_RTOL, "atol": ONNX_DEFAULT_ATOL}, + "Gemm": { + "rtol": ONNX_DEFAULT_RTOL, + "atol": ONNX_DEFAULT_ATOL * MAX_TOLERANCE_FACTOR * 10, + "justification": "summation order", + }, + } + ) + + assert len(problems) == 2 + assert any("Conv" in problem and "justification" in problem for problem in problems) + assert any("Gemm" in problem and "beyond" in problem for problem in problems) + + +def test_a_wrong_expectation_fails_the_case(tmp_path): + """Fail-closed: the comparison is real, on the corpus's own data.""" + copied = tmp_path / "case" + shutil.copytree(_corpus()["test_add"].model_dir, copied) + output = copied / "test_data_set_0" / "output_0.pb" + output.write_bytes( + numpy_helper.from_array(_read_tensor(output) + 1.0).SerializeToString() + ) + + with pytest.raises(AssertionError): + _execute(copied, rtol=ONNX_DEFAULT_RTOL, atol=ONNX_DEFAULT_ATOL) + + +def test_a_model_the_compiler_rejects_fails_the_case(tmp_path): + """Fail-closed: a compile error is a failure, never a skip.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("NonZero", ["x"], ["y"], name="nonzero")], + "rejected", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2, 6])], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + case = tmp_path / "case" + (case / "test_data_set_0").mkdir(parents=True) + onnx.save_model(model, case / "model.onnx") + + with pytest.raises(CompileError, match="NonZero"): + _execute(case, rtol=ONNX_DEFAULT_RTOL, atol=ONNX_DEFAULT_ATOL) diff --git a/src/python/tests/test_extra_compiler_differential.py b/src/python/tests/test_extra_compiler_differential.py new file mode 100644 index 0000000..f003781 --- /dev/null +++ b/src/python/tests/test_extra_compiler_differential.py @@ -0,0 +1,7134 @@ +"""Differential sweep: every registered kernel against the ONNX reference evaluator. + +The backend corpus is example-based. It does not cover every attribute combination, every +dtype a schema allows, or the numerical edges, so this suite generates single-node models +systematically and takes every expected value from `onnx.reference.ReferenceEvaluator` -- +the executable form of the spec. Nothing here decides what an op should compute. + +**Oracle validity.** The evaluator carries a versioned implementation class only for the +revisions whose semantics it distinguishes; to every other opset it silently applies the +newest semantics it knows. A case is therefore generated only for an (op, version) pair the +evaluator is version-faithful for -- the same mechanical check the compiler's folding pass +applies, never an assumption about which revisions "did not really change". Registered +revisions outside that set rest on the backend corpus's frozen old-opset tests, whose +expected outputs are stored rather than computed. + +**Acceptance rule.** An op counts as implemented only when both suites pass for it. The +tests at the bottom fail if the kernel registry serves an op this sweep does not execute a +case for, or one that no test in the conformance pass list exercises. +""" + +from __future__ import annotations + +import json +import math +import shutil +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field, replace +from enum import Enum +from functools import cache, partial +from pathlib import Path +from typing import Any, NamedTuple + +import pytest + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +# The harness refuses to import without numpy, so this covers both dependencies. +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from onnx import ModelProto, TensorProto, helper, numpy_helper # noqa: E402 +from onnx.backend.test.loader import load_model_tests # noqa: E402 +from onnx.backend.test.runner import Runner # noqa: E402 +from onnx.defs import OpSchema, get_schema # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +from fnnx.extras.compilers.c import compile_onnx # noqa: E402 +from fnnx.extras.compilers.c.onnx.dtypes import ( # noqa: E402 + C_TYPES, + FLOAT_TYPES, + UNSIGNED_TYPES, + c_type, + element_size, + numpy_dtype_name, +) +from fnnx.extras.compilers.c.onnx.folding import ( # noqa: E402 + evaluator_is_version_faithful, +) +from fnnx.extras.compilers.c.onnx.kernels import ( # noqa: E402 + KERNELS, + CFunction, + NodeContext, + NodeEmission, +) +from fnnx.extras.compilers.c.onnx.loader import ( # noqa: E402 + ML_DOMAIN, + display_domain, + normalize_domain, +) +from fnnx.extras.compilers.c.onnx.registry import KernelSpec # noqa: E402 + +# Every draw below comes from `numpy.random.default_rng([SEED, operand index])`, so a +# reported case reproduces exactly. +SEED = 20260726 + +# ONNX's own backend-test defaults. The compiler is compared under one tolerance policy, and +# that policy is ONNX's; a kernel that needs more than the conformance suite's bounded +# per-op overrides allow is a wrong kernel, here as much as there. +RTOL = 1e-3 +ATOL = 1e-7 + +# The conformance suite's checked-in record of which corpus tests pass, and the `onnx` +# release it -- and the corpus it names -- are keyed to. +RATCHET_PATH = Path(__file__).parent / "conformance" / "passing.txt" +LEDGER_PATH = Path(__file__).parent / "conformance" / "ledger.json" +PINNED_ONNX = "1.22" + +pytestmark = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + + +# -------------------------------------------------------------------------------------- +# What the sweep covers +# -------------------------------------------------------------------------------------- + + +class Kind(Enum): + """What an op does to its operands, which is what bounds the values it can be fed. + + The oracle is the reference evaluator either way; this only decides which inputs a + disagreement with it would be attributable to the kernel. + """ + + # Nothing the op computes can leave the dtype's range, so it sees every special value + # the dtype has, including its extremes. + POINTWISE = "pointwise" + # Elementwise arithmetic: IEEE pins the float edges (Inf, NaN, signed zero, overflow to + # Inf) exactly, so floats still sweep them, but integers stay small -- ONNX does not + # define integer overflow and C's is undefined for the signed families, so a wrap-around + # difference would not be a divergence from the spec. + ARITHMETIC = "arithmetic" + # Sums many products, in an order the spec does not fix: with Inf or dtype extremes in + # one dot product the reference's summation order and the kernel's legitimately + # disagree. These cases run on finite values, and their edge behaviour rests on the + # backend corpus. + ACCUMULATING = "accumulating" + + +class Domain(Enum): + """A restriction on an operand's values, for what ONNX itself leaves undefined.""" + + # Integer division and remainder by zero: undefined by ONNX, undefined in C, and a trap + # on the platforms this compiler targets. Integer operands only -- the floating-point + # case is pinned by IEEE, so those operands keep their full range of special values. + NONZERO = "nonzero" + # A negative integer exponent, which numpy refuses to evaluate at all, and one large + # enough to overflow the dtype, which ONNX does not define. Integer operands only. + SMALL_EXPONENT = "small_exponent" + # A floating-point operand converted to an integer type, which ONNX and C alike leave + # undefined outside the target's range: the special values that cannot survive the + # conversion -- NaN, the infinities, the dtype's extremes -- are pulled into a range + # every signed integer type holds. Integer operands convert for every value, so they are + # left alone. + CONVERTIBLE = "convertible" + # The same, for an unsigned target, which holds no negative value however small. + CONVERTIBLE_UNSIGNED = "convertible_unsigned" + # A factor of a product, running or reduced: many of them multiply together, so integer + # operands are pulled into {-1, 0, 1}, where no product can leave the dtype's range — + # ONNX does not define integer overflow and C's is undefined for the signed families. + SMALL_FACTOR = "small_factor" + # An operand a square root or a fractional power is taken of. A negative one puts the + # whole result at NaN, which compares equal to itself and would leave the case asserting + # nothing at all; the reference evaluator's LpNormalization needs it for a second reason, + # recorded where it is asked for. + NONNEGATIVE = "nonnegative" + + +@dataclass(frozen=True) +class Variant: + """One attribute combination and the operand shapes it applies to. + + `None` leaves an optional operand out. A variant is generated only at the versions whose + schema takes the operands it declares, so one table serves every revision of an op; + `versions` and `elem_types` narrow that further, for attributes a revision does not have + and combinations a dtype does not allow. `values` pins an operand a variant treats as a + parameter rather than as data — Clip's bounds, a reduction's axes — where a seeded draw + would make the case's coverage an accident of the seed; a single value fills the operand + and a sequence gives it element by element. `domains` restricts an operand for this + variant alone, where what ONNX leaves undefined depends on the attributes — Cast's target + type. `operand_types` does the same for an operand's element type, where a schema leaves + a choice the sweep's own type operand does not range over — the integer width of a + gathering op's indices. `outputs` asks for the optional results an op computes alongside + its first — LayerNormalization's mean, BatchNormalization's running statistics — which + are then compared like any other output. + """ + + label: str + shapes: tuple[tuple[int, ...] | None, ...] + attributes: Mapping[str, Any] = field(default_factory=dict) + values: Mapping[int, float | Sequence[float]] = field(default_factory=dict) + domains: Mapping[int, Domain] = field(default_factory=dict) + operand_types: Mapping[int, int] = field(default_factory=dict) + versions: tuple[int, ...] | None = None + elem_types: tuple[int, ...] | None = None + outputs: int = 1 + + +@dataclass(frozen=True) +class Sweep: + """Everything generated for one op: its variants, and how its operands are fed. + + `type_operand` is the operand whose schema type constraint enumerates the element types + swept — the first, unless the op takes something else there (Where's condition). + `operand_types` fixes the element type of the operands the schema does not leave free. + `constant_operands` are carried in the model as initializers rather than fed at run time: + an op that reads an operand as configuration — a reduction's axes — needs it at compile + time, and a model that computes it is a model the compiler refuses by design. + `equivalent_model` builds what the oracle is run on instead, for an op ONNX defines as + equal to another one and implements either nowhere or not faithfully; the expected values + still come from the oracle, on the model ONNX's own specification says computes the same + thing. `oracle` replaces the reference evaluator itself, for the one op ONNX ships no + reference implementation for at all -- see the MaxRoiPool entry, which is the only user. + """ + + kind: Kind + variants: tuple[Variant, ...] + operand_domains: Mapping[int, Domain] = field(default_factory=dict) + type_operand: int = 0 + operand_types: Mapping[int, int] = field(default_factory=dict) + constant_operands: tuple[int, ...] = () + equivalent_model: Callable[[Case], ModelProto] | None = None + oracle: Callable[[ModelProto, Mapping[str, Any]], list[Any]] | None = None + + +# Broadcasting cases: equal shapes, trailing-axis alignment, both-ways broadcasting, a +# scalar operand, mismatched ranks, rank 0, and zero-element tensors on either axis. `wide` +# is the only one larger than the special-value list on both sides, so it is what carries +# every dtype edge into both operands; the rest are shape coverage. +_BROADCAST_VARIANTS = ( + Variant("wide", ((4, 8), (4, 8))), + Variant("same", ((2, 3), (2, 3))), + Variant("trailing_axis", ((2, 3), (3,))), + Variant("both_ways", ((1, 3), (2, 1))), + Variant("scalar_operand", ((2, 3), ())), + Variant("mixed_rank", ((2, 1, 3), (4, 3))), + Variant("rank_0", ((), ())), + Variant("empty_rows", ((0, 3), (3,))), + Variant("empty_axis", ((2, 0, 3), (3,))), +) + +# `wide` is larger than the special-value list, so seeded random draws reach these ops too. +_UNARY_VARIANTS = ( + Variant("matrix", ((2, 3),)), + Variant("rank_0", ((),)), + Variant("empty", ((0, 3),)), + Variant("wide", ((4, 8),)), +) + +# PRelu broadcasts its slope onto the data unidirectionally: the result keeps the data's +# shape, so only the second operand may be stretched. +_UNIDIRECTIONAL_VARIANTS = ( + Variant("wide", ((4, 8), (4, 8))), + Variant("trailing_axis", ((2, 3), (3,))), + Variant("scalar_operand", ((2, 3), ())), + Variant("rank_0", ((), ())), + Variant("empty_rows", ((0, 3), (3,))), +) + +# The variadic families take one operand or many, each broadcasting against the others. The +# broadcast case gives the first operand the result's own shape: the reference's Mean +# accumulates into a copy of it, so it cannot evaluate a case where a later operand widens +# the result, and there is no oracle for one. +_VARIADIC_VARIANTS = ( + Variant("wide", ((4, 8), (4, 8), (4, 8))), + Variant("single", ((2, 3),)), + Variant("pair", ((2, 3), (2, 3))), + Variant("broadcast", ((2, 4, 3), (4, 3), ())), + Variant("rank_0", ((), ())), + Variant("empty_rows", ((0, 3), (3,))), +) + +_FLOAT_ELEM_TYPES = tuple(sorted(FLOAT_TYPES)) + +# Clip's bounds are scalars the op treats as parameters, so they are pinned rather than +# drawn: an inverted pair (numpy applies the lower bound first, so the upper one wins) and a +# NaN bound (which wins outright) are edges a random draw would only reach by luck. +_CLIP_BOUND_VALUES = ( + ("both", {1: -1.0, 2: 1.0}, None), + ("inverted", {1: 1.0, 2: -1.0}, None), + ("nan_low", {1: float("nan"), 2: 1.0}, _FLOAT_ELEM_TYPES), +) +_CLIP_VARIANTS = ( + Variant("unbounded", ((4, 8),)), + Variant("unbounded_rank_0", ((),)), + Variant("unbounded_empty", ((0, 3),)), + Variant("low_only", ((4, 8), ()), values={1: 0.0}), + Variant("high_only", ((4, 8), None, ()), values={2: 0.0}), + *( + Variant(label, ((4, 8), (), ()), values=values, elem_types=elem_types) + for label, values, elem_types in _CLIP_BOUND_VALUES + ), + # Up to opset 10 the bounds are attributes instead, which is a kernel of its own. + Variant("attributes", ((4, 8),), {"min": -1.0, "max": 1.0}, versions=(6,)), + Variant("attribute_low", ((4, 8),), {"min": 0.0}, versions=(6,)), + Variant("attribute_high", ((4, 8),), {"max": 0.0}, versions=(6,)), +) + +# Dropout is the identity in inference mode, whatever ratio it is handed — including one +# that is not a number at all. +_DROPOUT_VARIANTS = ( + Variant("data", ((2, 3),)), + Variant("wide", ((4, 8),)), + Variant("empty", ((0, 3),)), + Variant("ratio", ((4, 8), ()), values={1: 0.75}), + Variant("ratio_zero", ((4, 8), ()), values={1: 0.0}), + Variant("ratio_nan", ((4, 8), ()), values={1: float("nan")}), +) + +# ONNX defines Mod's fmod=0 for the integer families only, so the floored formula is swept +# there alone; the broadcasting it shares with fmod=1 is covered by the shape family above. +_INTEGER_TYPES = tuple( + elem_type + for elem_type in sorted(C_TYPES) + if elem_type not in FLOAT_TYPES and elem_type != TensorProto.BOOL +) +_MOD_VARIANTS = tuple( + replace(variant, label=f"{variant.label}_fmod", attributes={"fmod": 1}) + for variant in _BROADCAST_VARIANTS +) + tuple( + replace(variant, label=f"{variant.label}_floored", elem_types=_INTEGER_TYPES) + for variant in _BROADCAST_VARIANTS + if variant.label in ("wide", "trailing_axis") +) + +# BitShift's direction is a required attribute, so every shape carries one of the two. +_BIT_SHIFT_VARIANTS = tuple( + replace( + variant, + label=f"{variant.label}_{direction.lower()}", + attributes={"direction": direction}, + ) + for direction in ("LEFT", "RIGHT") + for variant in _BROADCAST_VARIANTS +) + +# IsInf's attributes decide which infinities count, down to neither of them, where the +# operand goes unread. +_IS_INF_ATTRIBUTES: Mapping[str, Mapping[str, Any]] = { + "positive_only": {"detect_negative": 0}, + "negative_only": {"detect_positive": 0}, + "neither": {"detect_positive": 0, "detect_negative": 0}, +} + +# Where broadcasts all three operands against each other; the condition is boolean whatever +# the branches carry, which is what the sweep's `operand_types` pins. +_SELECT_VARIANTS = ( + Variant("wide", ((4, 8), (4, 8), (4, 8))), + Variant("same", ((2, 3), (2, 3), (2, 3))), + Variant("condition_broadcast", ((2, 1), (2, 3), (2, 3))), + Variant("branch_broadcast", ((2, 3), (3,), ())), + Variant("every_operand_stretched", ((2, 1, 1), (1, 3, 1), (1, 1, 4))), + Variant("rank_0", ((), (), ())), + Variant("empty_rows", ((0, 3), (3,), ())), +) + + +def _cast_variants() -> tuple[Variant, ...]: + """Every supported target type, plus the shapes a conversion does not vary over. + + The source types come from the schema, so each variant is generated once per source: the + pair matrix is the cross product of the two, the identity conversions included. + """ + variants = [ + Variant( + f"to_{numpy_dtype_name(target)}", + ((4, 8),), + {"to": target}, + # Only the conversions to another floating-point type, or to bool, are defined + # for every value a float operand can hold. + domains=( + {} + if target in FLOAT_TYPES or target == TensorProto.BOOL + else { + 0: ( + Domain.CONVERTIBLE_UNSIGNED + if target in UNSIGNED_TYPES + else Domain.CONVERTIBLE + ) + } + ), + ) + for target in sorted(C_TYPES) + ] + variants += [ + Variant(label, (shape,), {"to": TensorProto.DOUBLE}) + for label, shape in (("matrix", (2, 3)), ("rank_0", ()), ("empty", (0, 3))) + ] + return tuple(variants) + + +def _bitcast_variants() -> tuple[Variant, ...]: + """One variant per target type, offered to the source types of that same width. + + ONNX defines BitCast only between types of equal width, and this compiler refuses a + boolean target outright — its bytes are contractually 0 or 1, which arbitrary bits are + not — so neither is generated here; both are error-path tests instead. + """ + by_width: dict[int, list[int]] = {} + for elem_type in sorted(C_TYPES): + by_width.setdefault(element_size(elem_type), []).append(elem_type) + variants = [ + Variant( + f"to_{numpy_dtype_name(target)}", + ((4, 8),), + {"to": target}, + elem_types=tuple(members), + ) + for members in by_width.values() + for target in members + if target != TensorProto.BOOL + ] + variants += [ + Variant( + label, + (shape,), + {"to": TensorProto.INT32}, + elem_types=(TensorProto.FLOAT,), + ) + for label, shape in (("rank_0", ()), ("empty", (0, 3))) + ] + return tuple(variants) + + +# Gemm's attributes are swept as a full cross product rather than a chosen handful: which +# combinations interact is exactly what a hand-picked list would be guessing at. The bias +# shapes cover every way C broadcasts onto the result, including leaving it out. +_GEMM_BIASES = ( + ("vector", (4,)), + ("matrix", (2, 4)), + ("row", (1, 4)), + ("scalar", ()), + ("absent", None), +) + + +def _gemm_variants() -> tuple[Variant, ...]: + variants = [ + Variant( + f"transA{transpose_left}_transB{transpose_right}" + f"_alpha{alpha}_beta{beta}_bias_{bias_label}", + ( + (3, 2) if transpose_left else (2, 3), + (4, 3) if transpose_right else (3, 4), + bias, + ), + { + "transA": transpose_left, + "transB": transpose_right, + "alpha": alpha, + "beta": beta, + }, + ) + for transpose_left in (0, 1) + for transpose_right in (0, 1) + for alpha, beta in ((1.0, 1.0), (0.5, 2.0), (1.5, 0.0)) + for bias_label, bias in _GEMM_BIASES + ] + # Zero-element operands, which no attribute combination reaches: an empty result, and an + # empty contraction whose sum is over nothing at all. + variants += [ + Variant("empty_rows", ((0, 3), (3, 4), (4,))), + Variant("empty_inner", ((2, 0), (0, 4), (4,))), + ] + return tuple(variants) + + +_GEMM_VARIANTS = _gemm_variants() + +# MatMul is numpy's `matmul`, so the sweep is the shapes that convention distinguishes: a +# rank-1 operand on either side, which is promoted to the row or column that makes the product +# defined and dropped from the result again, and batch axes that stretch against each other or +# are absent from one operand entirely. The zero-element cases cover an empty result and an +# empty contraction, whose sum is over nothing at all. +_MATMUL_VARIANTS = ( + Variant("wide", ((4, 8), (8, 4))), + Variant("matrix", ((2, 3), (3, 4))), + Variant("vector_vector", ((3,), (3,))), + Variant("vector_matrix", ((3,), (3, 4))), + Variant("matrix_vector", ((2, 3), (3,))), + Variant("batched", ((2, 3, 4), (2, 4, 3))), + Variant("batched_vector", ((2, 3, 4), (4,))), + Variant("vector_batched", ((4,), (2, 4, 3))), + Variant("unbatched_left", ((2, 3), (2, 3, 4))), + Variant("unbatched_right", ((2, 3, 4), (4, 3))), + Variant("both_batches_stretched", ((3, 1, 2, 4), (1, 2, 4, 3))), + Variant("rank_4", ((1, 2, 3, 4), (1, 2, 4, 3))), + Variant("empty_rows", ((0, 3), (3, 4))), + Variant("empty_columns", ((2, 3), (3, 0))), + Variant("empty_inner", ((2, 0), (0, 4))), + Variant("empty_batch", ((0, 2, 3), (0, 3, 4))), + # An empty batch axis stretched against a batch of one, from either side: broadcasting a + # 1 against a 0 yields 0, which is the one place the rule is not "take the larger". + Variant("empty_batch_over_one", ((0, 2, 3), (1, 3, 4))), + Variant("one_batch_over_empty", ((1, 2, 3), (0, 3, 4))), +) + +# Det factorizes each of the trailing square matrices. Orders 1 through 4 cover the pivoting — +# the seeded draws put the largest element of a column off the diagonal often enough to swap +# rows — and the batch axes are what carries more than one matrix through the same kernel. A +# matrix of order 0 is square too, and its determinant is the empty product. +_DET_VARIANTS = ( + Variant("order_2", ((2, 2),)), + Variant("order_3", ((3, 3),)), + Variant("order_4", ((4, 4),)), + Variant("order_1", ((1, 1),)), + Variant("batched", ((3, 2, 2),)), + Variant("batched_rank_4", ((2, 3, 3, 3),)), + Variant("empty_batch", ((0, 2, 2),)), + Variant("order_0", ((2, 0, 0),)), +) + +# Einsum's surface is its equation, so the sweep is every reading an equation has and the +# shapes each one addresses: an output stated and one left implicit (where the labels are +# ordered alphabetically rather than as written), a label repeated inside a term at each +# position it can repeat in — which is a diagonal — a label two terms share, which is summed, +# an ellipsis standing for a leading, inner, trailing or absent block of axes, and the spaces +# numpy takes out before reading any of it. The zero-element cases cover an empty result and +# an empty contraction, whose sum is over nothing at all. +# +# The ellipsis broadcasting swept here is the one ONNX's shape inference derives: equal +# ellipsis ranks, an extent of 1 stretching against another operand's. An ellipsis standing +# for a different *number* of axes on two operands is deliberately absent — the pinned `onnx` +# release crashes outright inferring some of those models, so there is no compiling one to +# compare against anything. +_EINSUM_VARIANTS = ( + Variant("matmul", ((2, 3), (3, 4)), {"equation": "ij,jk->ik"}), + Variant("matmul_implicit", ((2, 3), (3, 4)), {"equation": "ij,jk"}), + Variant("batch_matmul", ((2, 3, 4), (2, 4, 5)), {"equation": "bij,bjk->bik"}), + Variant( + "spaced_terms", ((2, 3, 4), (2, 4, 5)), {"equation": "b i j, b j k -> b i k"} + ), + Variant("transpose", ((2, 3),), {"equation": "ij->ji"}), + Variant("identity_implicit", ((2, 3),), {"equation": "ij"}), + Variant("reordered_implicit", ((2, 3),), {"equation": "ji"}), + Variant("summed_axis", ((2, 3),), {"equation": "ij->i"}), + Variant("summed_all", ((2, 3),), {"equation": "ij->"}), + Variant("diagonal", ((4, 4),), {"equation": "ii->i"}), + Variant("trace_implicit", ((4, 4),), {"equation": "ii"}), + Variant("interleaved_diagonal", ((3, 4, 3),), {"equation": "iji->ij"}), + Variant("batch_diagonal", ((2, 4, 4),), {"equation": "...ii ->...i"}), + Variant("inner_product", ((5,), (5,)), {"equation": "i,i"}), + Variant("outer_product", ((3,), (4,)), {"equation": "i,j->ij"}), + Variant("hadamard", ((2, 3), (2, 3)), {"equation": "ij,ij->ij"}), + Variant("scalar", ((),), {"equation": "->"}), + Variant("scaled_by_scalar", ((2, 3), ()), {"equation": "ij,->ij"}), + Variant("three_terms", ((2, 3), (3, 4), (4, 5)), {"equation": "ij,jk,kl->il"}), + Variant( + "leading_ellipsis", ((2, 3, 4), (2, 4, 5)), {"equation": "...ij,...jk->...ik"} + ), + Variant( + "stretched_ellipsis", ((1, 2, 3), (5, 3, 4)), {"equation": "...ij,...jk->...ik"} + ), + Variant("inner_ellipsis", ((2, 3, 4, 5),), {"equation": "i...j->ji"}), + Variant("summed_ellipsis", ((2, 3),), {"equation": "...i->..."}), + Variant("dropped_ellipsis", ((2, 3),), {"equation": "...i->i"}), + Variant("implicit_ellipsis", ((2, 3),), {"equation": "...i"}), + Variant("stretched_label", ((1, 3), (2, 3)), {"equation": "ij,ij->j"}), + Variant("empty_contraction", ((2, 0), (0, 3)), {"equation": "ij,jk->ik"}), + Variant("empty_result", ((0, 3),), {"equation": "ij->ji"}), + Variant("empty_batch", ((0, 2, 3), (0, 3, 4)), {"equation": "bij,bjk->bik"}), +) + +# DFT is one sum per output bin, so the sweep is everything that decides which samples a bin +# runs over and what it writes: each of the four `inverse`/`onesided` combinations — the +# forward transform, the RFFT that keeps the non-redundant half of a real signal's spectrum, +# the inverse, and the IRFFT that mirrors such a half back into a real signal — over real and +# complex operands, at even and odd lengths, with a `dft_length` that leaves the axis alone, +# truncates it or zero-pads it. The axis is swept as the attribute revision 17 reads and as +# the operand revision 20 takes, stated, counted from the end, and left out — the two +# revisions default it differently, which a rank-4 operand tells apart. The transformed axis +# itself is never empty: numpy refuses a transform of no points at all, so there is no oracle +# for one, and an empty batch carries the zero-element case instead. +_DFT_SIGNALS: tuple[tuple[str, tuple[int, ...], Mapping[str, Any]], ...] = ( + ("real", (3, 8, 1), {}), + ("complex", (3, 8, 2), {}), + ("inverse", (3, 8, 2), {"inverse": 1}), + ("rfft", (3, 8, 1), {"onesided": 1}), + ("irfft", (3, 6, 2), {"inverse": 1, "onesided": 1}), + ("real_odd", (3, 7, 1), {}), + ("complex_odd", (3, 7, 2), {}), + ("rfft_odd", (3, 7, 1), {"onesided": 1}), + ("irfft_odd", (3, 5, 2), {"inverse": 1, "onesided": 1}), + ("empty_batch", (0, 8, 1), {}), + ("rank_4", (2, 3, 4, 1), {}), +) + +# What a stated `dft_length` does to an 8-sample axis, against leaving it out. +_DFT_LENGTHS = (("truncated", 5), ("padded", 12), ("same", 8)) + + +def _dft_variants() -> tuple[Variant, ...]: + variants = [ + Variant(label, (shape,), attributes) + for label, shape, attributes in _DFT_SIGNALS + ] + variants += [ + Variant( + f"{label}_{combination}", + ((3, 8, 1 if onesided and not inverse else 2), ()), + {"inverse": inverse, "onesided": onesided}, + values={1: length}, + ) + for label, length in _DFT_LENGTHS + for combination, inverse, onesided in ( + ("forward", 0, 0), + ("rfft", 0, 1), + ("inverse", 1, 0), + ("irfft", 1, 1), + ) + ] + # The axis as revision 17's attribute and as revision 20's operand, each at the positions + # ONNX defines: a leading axis, a trailing signal axis, and the same two counted from the + # end. A rank-2 operand has exactly one axis a transform is defined over, and neither + # revision's default names it, so it is reached only by stating the axis. + for label, shape, axis in ( + ("leading", (3, 8, 4, 1), 0), + ("trailing", (3, 8, 4, 1), 2), + ("from_the_end", (3, 8, 4, 1), -2), + ("far_from_the_end", (3, 8, 4, 1), -4), + ("rank_2", (8, 1), 0), + ): + variants.append( + Variant(f"axis_{label}_attribute", (shape,), {"axis": axis}, versions=(17,)) + ) + variants.append( + Variant(f"axis_{label}_operand", (shape, None, ()), values={2: axis}) + ) + return tuple(variants) + + +_DFT_VARIANTS = _dft_variants() + +# STFT is a DFT per frame of a window slid along the signal, so the sweep is how the frames +# are laid out — the step against the frame length, from a dense overlap to none at all — and +# where the length comes from: the `frame_length` operand, the window's own extent, or both +# stating it. The window is fed at run time rather than folded in, since its values reach the +# kernel and nothing else about it does. `onesided` is swept stated both ways and left out: +# ONNX's shape inference reads a default of 0 there where the schema declares 1, so the +# omitted case is the one that proves the compiler sizes the result the way the op computes it. +_STFT_VARIANTS = ( + Variant("frame_length", ((2, 32, 1), (), None, ()), values={1: 8, 3: 16}), + Variant("window", ((2, 32, 1), (), (16,)), values={1: 8}), + Variant( + "window_and_frame_length", ((2, 32, 1), (), (16,), ()), values={1: 8, 3: 16} + ), + Variant( + "onesided", + ((2, 32, 1), (), None, ()), + {"onesided": 1}, + values={1: 8, 3: 16}, + ), + Variant( + "twosided", + ((2, 32, 1), (), None, ()), + {"onesided": 0}, + values={1: 8, 3: 16}, + ), + Variant( + "complex_signal", + ((2, 32, 2), (), None, ()), + {"onesided": 0}, + values={1: 8, 3: 16}, + ), + Variant("dense_overlap", ((1, 16, 1), (), None, ()), values={1: 1, 3: 8}), + Variant("no_overlap", ((1, 16, 1), (), None, ()), values={1: 8, 3: 8}), + Variant("one_frame", ((1, 16, 1), (), None, ()), values={1: 8, 3: 16}), + Variant("odd_frame_length", ((1, 16, 1), (), None, ()), values={1: 4, 3: 7}), + Variant("empty_batch", ((0, 32, 1), (), None, ()), values={1: 8, 3: 16}), +) + +# Conv slides a filter over the spatial axes of a batch of multi-channel signals, so the sweep +# is its geometry: rank 1 through 3, the attributes that move the window (strides, dilations) +# and the ones that place it (pads, auto_pad), the channel groups that split the operand into +# independent stacks, and the bias. +# +# `auto_pad` is swept only at shapes whose batch and channel extents repeat the spatial ones. +# The reference resolves the mode by reading `X.shape[i]` where the spec reads `X.shape[i+2]`, +# so it computes the padding ONNX defines exactly when those coincide, and is no oracle +# elsewhere. `VALID` goes further: the reference puts it through the same SAME-style formula, +# so the only geometries it agrees with the spec on are the ones where SAME pads nothing +# either -- hence the stride-2 2x2 window here. That the compiler does *not* read `VALID` as +# SAME is settled instead by ONNX's own shape inference, in the kernel tests. +# +# No zero-element variant: the reference raises outright on an empty batch and on a filter +# with no input channels, so neither has an oracle here. That a Conv writing no elements emits +# no code at all is an emission-contract assertion, and lives with the kernel tests. +_CONV_VARIANTS = ( + Variant("spatial_1d", ((2, 2, 7), (3, 2, 3))), + Variant("spatial_1d_strided", ((2, 2, 7), (3, 2, 3)), {"strides": [2]}), + Variant("spatial_1d_padded", ((2, 2, 7), (3, 2, 3)), {"pads": [2, 1]}), + Variant("spatial_1d_dilated", ((2, 2, 7), (3, 2, 3)), {"dilations": [2]}), + Variant("spatial_2d", ((2, 3, 5, 4), (2, 3, 3, 2))), + Variant("pads", ((2, 3, 5, 4), (2, 3, 3, 2)), {"pads": [1, 1, 1, 1]}), + Variant("asymmetric_pads", ((2, 3, 5, 4), (2, 3, 3, 2)), {"pads": [2, 0, 1, 1]}), + Variant("strides", ((2, 3, 5, 4), (2, 3, 3, 2)), {"strides": [2, 2]}), + Variant("dilations", ((2, 3, 5, 4), (2, 3, 3, 2)), {"dilations": [2, 1]}), + Variant( + "strided_dilated_padded", + ((2, 3, 5, 4), (2, 3, 3, 2)), + {"strides": [2, 1], "dilations": [1, 2], "pads": [1, 2, 0, 1]}, + ), + Variant("bias", ((2, 3, 5, 4), (2, 3, 3, 2), (2,))), + Variant("unit_window", ((2, 3, 5, 4), (2, 3, 1, 1))), + # A window wider than the signal, which only the padding makes fit at all. + Variant("window_wider_than_input", ((1, 1, 2, 2), (1, 1, 3, 3)), {"pads": [2] * 4}), + Variant("groups", ((2, 4, 5, 4), (6, 2, 3, 2)), {"group": 2}), + Variant("depthwise", ((2, 4, 5, 4), (4, 1, 3, 2), (4,)), {"group": 4}), + Variant("spatial_3d", ((1, 2, 4, 3, 3), (2, 2, 2, 2, 2))), + Variant( + "spatial_3d_strided_padded", + ((1, 2, 4, 3, 3), (2, 2, 2, 2, 2)), + {"strides": [2, 1, 2], "pads": [1, 0, 1, 0, 1, 0]}, + ), + # An odd total pad, which is what tells the two SAME modes apart: the extra one goes at + # the end for SAME_UPPER and at the beginning for SAME_LOWER. + Variant( + "auto_pad_same_upper", + ((3, 4, 3, 4), (2, 4, 2, 3)), + {"auto_pad": "SAME_UPPER", "strides": [2, 2]}, + ), + Variant( + "auto_pad_same_lower", + ((3, 4, 3, 4), (2, 4, 2, 3)), + {"auto_pad": "SAME_LOWER", "strides": [2, 2]}, + ), + Variant( + "auto_pad_same_upper_unit_stride", + ((3, 4, 3, 4), (2, 4, 2, 3)), + {"auto_pad": "SAME_UPPER"}, + ), + Variant( + "auto_pad_same_lower_unit_stride", + ((3, 4, 3, 4), (2, 4, 2, 3)), + {"auto_pad": "SAME_LOWER"}, + ), + # SAME pads for the window's *dilated* reach, not its tap count. Shape inference only + # pins the total -- at stride 2 a whole family of totals rounds to the same result shape + # -- and says nothing at all about which end the odd one goes to, so the split under + # dilation is a value question, and only the evaluator answers it. The first axis takes + # a 2-tap window dilated to a reach of 4 over an extent of 5 at stride 2: three pads, + # which the two modes divide the other way round from each other. + Variant( + "auto_pad_same_upper_dilated", + ((5, 4, 5, 4), (2, 4, 2, 2)), + {"auto_pad": "SAME_UPPER", "dilations": [3, 2], "strides": [2, 1]}, + ), + Variant( + "auto_pad_same_lower_dilated", + ((5, 4, 5, 4), (2, 4, 2, 2)), + {"auto_pad": "SAME_LOWER", "dilations": [3, 2], "strides": [2, 1]}, + ), + Variant( + "auto_pad_valid", + ((4, 4, 4, 4), (2, 4, 2, 2)), + {"auto_pad": "VALID", "strides": [2, 2]}, + ), +) + +# ConvTranspose runs the same window backwards, so the sweep is Conv's geometry again -- +# rank, strides, dilations, pads -- plus what only the backward walk has: `output_padding` +# and `output_shape`, which name the result the stride leaves ambiguous, and a filter laid +# out per input channel, (C, M/group, ...). +# +# Two restrictions come from the oracle. The reference resolves an `output_shape` under +# NOTSET by padding nothing and then reading the operand back through a column count it +# derives from that shape instead of from the operand -- which only agree where the padding +# the spec's own equation gives is zero anyway, as it is in the corpus's own `output_shape` +# tests. The shapes here are those. And its grouped path slices `W` by output rather than +# input channels and hands every group the whole bias, so it can only evaluate a group at +# all where each holds exactly one channel of each -- hence the depthwise-shaped `groups` +# variants below, with no bias among them. What the general case computes is settled +# instead by decomposing it into per-group transposed convolutions the reference *can* +# evaluate, in the kernel tests. +_CONV_TRANSPOSE_VARIANTS = ( + Variant("spatial_1d", ((2, 2, 7), (2, 3, 3))), + Variant("spatial_1d_strided", ((2, 2, 7), (2, 3, 3)), {"strides": [2]}), + Variant("spatial_1d_padded", ((2, 2, 7), (2, 3, 3)), {"pads": [2, 1]}), + Variant("spatial_1d_dilated", ((2, 2, 7), (2, 3, 3)), {"dilations": [2]}), + Variant("spatial_2d", ((2, 3, 5, 4), (3, 2, 3, 2))), + Variant("pads", ((2, 3, 5, 4), (3, 2, 3, 2)), {"pads": [1, 1, 1, 1]}), + Variant("asymmetric_pads", ((2, 3, 5, 4), (3, 2, 3, 2)), {"pads": [2, 0, 1, 1]}), + Variant("strides", ((2, 3, 5, 4), (3, 2, 3, 2)), {"strides": [2, 2]}), + Variant("dilations", ((2, 3, 5, 4), (3, 2, 3, 2)), {"dilations": [2, 1]}), + Variant( + "strided_dilated_padded", + ((2, 3, 5, 4), (3, 2, 3, 2)), + {"strides": [3, 2], "dilations": [1, 2], "pads": [1, 2, 0, 1]}, + ), + Variant("bias", ((2, 3, 5, 4), (3, 2, 3, 2), (2,))), + Variant("unit_window", ((2, 3, 5, 4), (3, 2, 1, 1))), + # Padding wider than the operand's own reach, which crops the result to nothing but + # the overlap. + Variant("cropping_pads", ((1, 1, 3, 3), (1, 1, 2, 2)), {"pads": [1, 1, 1, 1]}), + Variant("groups", ((2, 3, 5, 4), (3, 1, 3, 2)), {"group": 3}), + Variant( + "groups_strided_padded", + ((2, 3, 5, 4), (3, 1, 3, 2)), + {"group": 3, "strides": [2, 1], "pads": [1, 0, 0, 1]}, + ), + # `output_padding` extends the result past the window's reach: the positions it adds + # take no tap at all, and carry the bias alone. + Variant( + "output_padding", + ((2, 3, 5, 4), (3, 2, 3, 2), (2,)), + {"strides": [3, 2], "output_padding": [1, 1]}, + ), + Variant( + "output_shape", + ((1, 1, 3, 3), (1, 2, 3, 3)), + {"strides": [3, 2], "output_shape": [10, 8]}, + ), + Variant( + "output_shape_below_the_reach", + ((1, 1, 3, 3), (1, 2, 3, 3)), + {"strides": [3, 2], "output_shape": [9, 7]}, + ), + Variant( + "output_shape_with_output_padding", + ((1, 1, 3, 3), (1, 2, 3, 3)), + {"strides": [3, 2], "output_shape": [10, 8], "output_padding": [1, 1]}, + ), + # An odd total pad, which is what tells the two SAME modes apart. Here a 3-tap window + # at stride 2 leaves one pad over; the modes put it at opposite ends. + Variant( + "auto_pad_same_upper", + ((2, 3, 5, 4), (3, 2, 3, 2)), + {"auto_pad": "SAME_UPPER", "strides": [2, 2]}, + ), + Variant( + "auto_pad_same_lower", + ((2, 3, 5, 4), (3, 2, 3, 2)), + {"auto_pad": "SAME_LOWER", "strides": [2, 2]}, + ), + Variant( + "auto_pad_same_upper_dilated", + ((2, 3, 5, 4), (3, 2, 3, 2)), + {"auto_pad": "SAME_UPPER", "strides": [3, 2], "dilations": [2, 1]}, + ), + Variant( + "auto_pad_same_lower_dilated", + ((2, 3, 5, 4), (3, 2, 3, 2)), + {"auto_pad": "SAME_LOWER", "strides": [3, 2], "dilations": [2, 1]}, + ), + # A SAME mode against an explicit `output_shape`: the mode no longer decides the result, + # only which end of it the odd pad goes to. + Variant( + "auto_pad_same_upper_output_shape", + ((1, 1, 3, 3), (1, 2, 3, 3)), + {"auto_pad": "SAME_UPPER", "strides": [2, 2], "output_shape": [6, 6]}, + ), + Variant( + "auto_pad_same_lower_output_shape", + ((1, 1, 3, 3), (1, 2, 3, 3)), + {"auto_pad": "SAME_LOWER", "strides": [2, 2], "output_shape": [6, 6]}, + ), + Variant("auto_pad_valid", ((2, 3, 5, 4), (3, 2, 3, 2)), {"auto_pad": "VALID"}), + Variant("spatial_3d", ((1, 2, 4, 3, 3), (2, 2, 2, 2, 2))), + Variant( + "spatial_3d_strided_padded", + ((1, 2, 4, 3, 3), (2, 2, 2, 2, 2)), + {"strides": [2, 1, 2], "pads": [1, 0, 1, 0, 1, 0]}, + ), +) + +# DeformConv shifts every tap by an offset it reads at run time, so its window lands between +# elements and is interpolated. The sweep is the geometry again — the offsets themselves are +# a seeded draw of ordinary magnitude, which puts sampling points inside the operand, across +# its border and well outside it — plus what only this op has: the `offset_group`s that +# choose which offsets a channel follows, and the `mask` that weights each tap. +# +# Every variant is two-dimensional: the reference evaluator raises outright on any other +# rank, and the compiler emits a kernel for no other rank either, for that reason. The +# `offset` shape is the geometry restated, two coordinates per tap per offset group at every +# result position, and `mask` is the same without the coordinate pair. +_DEFORM_CONV_VARIANTS = ( + Variant("basic", ((1, 1, 4, 4), (1, 1, 2, 2), (1, 8, 3, 3))), + Variant( + "pads", + ((1, 1, 4, 4), (1, 1, 2, 2), (1, 8, 5, 5)), + {"pads": [1, 1, 1, 1]}, + ), + Variant( + "asymmetric_pads", + ((1, 1, 4, 4), (1, 1, 2, 2), (1, 8, 4, 5)), + {"pads": [1, 1, 0, 1]}, + ), + Variant( + "strides", + ((2, 1, 5, 4), (3, 1, 2, 3), (2, 12, 2, 2)), + {"strides": [2, 1]}, + ), + Variant( + "dilations", + ((2, 2, 5, 4), (3, 2, 2, 3), (2, 12, 3, 2)), + {"dilations": [2, 1]}, + ), + Variant( + "strided_dilated_padded", + ((1, 1, 6, 5), (1, 1, 3, 2), (1, 12, 2, 3), (1,), (1, 6, 2, 3)), + {"strides": [2, 2], "dilations": [2, 1], "pads": [1, 1, 1, 1]}, + ), + Variant("bias", ((1, 1, 4, 4), (2, 1, 2, 2), (1, 8, 3, 3), (2,))), + Variant( + "mask", + ((1, 1, 4, 4), (2, 1, 2, 2), (1, 8, 3, 3), (2,), (1, 4, 3, 3)), + ), + # A mask without a bias: the operand ONNX puts between them is left out. + Variant( + "mask_without_bias", + ((1, 1, 4, 4), (2, 1, 2, 2), (1, 8, 3, 3), None, (1, 4, 3, 3)), + ), + Variant("unit_window", ((1, 1, 4, 4), (1, 1, 1, 1), (1, 2, 4, 4))), + Variant("groups", ((1, 2, 5, 5), (2, 1, 2, 2), (1, 8, 4, 4)), {"group": 2}), + Variant( + "offset_groups", + ((1, 2, 5, 5), (1, 2, 2, 2), (1, 16, 4, 4)), + {"offset_group": 2}, + ), + Variant( + "groups_and_offset_groups", + ((1, 4, 5, 5), (2, 2, 2, 2), (1, 16, 4, 4), (2,), (1, 8, 4, 4)), + {"group": 2, "offset_group": 2}, + ), + # One offset group per channel, which is the most deformation the op allows. + Variant( + "offset_group_per_channel", + ((1, 4, 5, 5), (2, 2, 2, 2), (1, 32, 4, 4)), + {"group": 2, "offset_group": 4}, + ), +) + + +# The quantization family. What the two affine maps have to be swept over is the granularity +# of their scale -- one per tensor, one per slice along the quantization axis, one per block +# of such a slice -- at each revision that has each, and the grid types they round onto, +# which is where saturation lives. The revisions that declare `axis`, `block_size` and +# `output_dtype` are named per variant, since a node cannot carry an attribute a revision +# does not declare. `precision`, which opset 23 added, is swept nowhere: the newest +# QuantizeLinear the reference evaluator implements is revision 21, which refuses a node +# carrying the attribute outright, so there is no oracle for it -- and the compiler refuses +# it for that reason, which the kernel tests assert. +_AXIS_VERSIONS = (19, 21, 23, 24, 25) +_BLOCKED_VERSIONS = (21, 23, 24, 25) + +# One scale per slice, a mix of exactly representable ones and ones that are not: a scale of +# 0.75 puts quotients on the halves that the round-to-even rule is the whole of. +_PER_AXIS_SCALES = ( + 0.5, 0.25, 2.0, 1.0, 0.125, 4.0, 0.75, 1.5, + 0.0625, 8.0, 3.0, 0.3125, 0.5, 16.0, 0.875, 2.5, +) # fmt: skip + +# QuantizeLinear's scale is pinned rather than drawn wherever the quotient it divides out has +# to stay where the oracle is defined: the reference converts `rint(x / y_scale)` to `int32` +# before it clips, which numpy leaves undefined outside that range, and a drawn scale near +# zero would put it there. `x` is pulled into the same range for the same reason, which is +# what `Domain.CONVERTIBLE` is. The zero point *is* drawn, so every grid's own extremes reach +# the addition that follows -- and that is what saturates the result at either end. +_QUANTIZE_VARIANTS = ( + Variant("per_tensor", ((4, 8), (), ()), values={1: 0.5}), + # Small enough that every grid's range is left at both ends, uint16's included. + Variant("saturating", ((4, 8), (), ()), values={1: 1e-4}), + # Halves throughout, which is the rule that rounds them to the even neighbour. + Variant("halves", ((4, 8), (), ()), values={0: 1.5, 1: 1.0}), + Variant("rank_0", ((), (), ()), values={1: 0.5}), + Variant("empty", ((0, 3), (), ()), values={1: 0.5}), + Variant("single_element_scale", ((4, 8), (1,), (1,)), values={1: 0.25}), + # No zero point: the grid is uint8 unless `output_dtype` names another, which is the + # variant below, so both run at that one type. + Variant( + "no_zero_point", + ((4, 8), ()), + values={1: 0.5}, + elem_types=(TensorProto.UINT8,), + ), + Variant( + "output_dtype", + ((4, 8), ()), + {"output_dtype": TensorProto.INT16}, + values={1: 0.5}, + elem_types=(TensorProto.UINT8,), + versions=_BLOCKED_VERSIONS, + ), + # An `int32` operand, which ONNX allows alongside the float ones and numpy promotes + # against a float scale. Only at the revisions where `y_scale` is a type of its own: + # 19 and 21 tie it to the operand's. + Variant( + "int32_data", + ((4, 8), (), ()), + values={1: 0.5}, + operand_types={0: TensorProto.INT32}, + elem_types=(TensorProto.UINT8,), + versions=(10, 23, 24, 25), + ), + Variant("per_axis", ((4, 16), (16,), (16,)), values={1: _PER_AXIS_SCALES}), + Variant( + "per_axis_first", + ((16, 4), (16,), (16,)), + {"axis": 0}, + values={1: _PER_AXIS_SCALES}, + versions=_AXIS_VERSIONS, + ), + Variant( + "per_axis_negative", + ((4, 16), (16,), (16,)), + {"axis": -1}, + values={1: _PER_AXIS_SCALES}, + versions=_AXIS_VERSIONS, + ), + Variant( + "blocked", + ((4, 8), (4, 2), (4, 2)), + {"axis": 1, "block_size": 4}, + values={1: (0.5, 0.25, 2.0, 1.0, 0.125, 4.0, 0.75, 1.5)}, + versions=_BLOCKED_VERSIONS, + ), + # A block the axis does not divide, whose last block covers what is left of it. + Variant( + "blocked_remainder", + ((4, 7), (4, 2), (4, 2)), + {"axis": 1, "block_size": 4}, + values={1: (0.5, 0.25, 2.0, 1.0, 0.125, 4.0, 0.75, 1.5)}, + versions=_BLOCKED_VERSIONS, + ), + Variant( + "blocked_first_axis", + ((8, 3), (2, 3), (2, 3)), + {"axis": 0, "block_size": 4}, + values={1: (0.5, 0.25, 2.0, 1.0, 0.125, 4.0)}, + versions=_BLOCKED_VERSIONS, + ), +) + +# DequantizeLinear only multiplies, so nothing it computes can leave the type it computes in: +# its scale is drawn like any other operand, special values and all, and the grid it reads +# sweeps its own extremes. The granularities are QuantizeLinear's, at the revisions that have +# them -- every revision claimed here declares `axis`. +_DEQUANTIZE_VARIANTS = ( + Variant("per_tensor", ((4, 8), (), ())), + Variant("rank_0", ((), (), ())), + Variant("empty", ((0, 3), (), ())), + Variant("single_element_scale", ((4, 8), (1,), (1,))), + Variant("no_zero_point", ((4, 8), ())), + Variant("per_axis", ((4, 16), (16,), (16,))), + Variant("per_axis_first", ((16, 4), (16,), (16,)), {"axis": 0}), + Variant("per_axis_negative", ((4, 16), (16,), (16,)), {"axis": -1}), + Variant( + "blocked", + ((4, 8), (4, 2), (4, 2)), + {"axis": 1, "block_size": 4}, + versions=_BLOCKED_VERSIONS, + ), + Variant( + "blocked_remainder", + ((4, 7), (4, 2), (4, 2)), + {"axis": 1, "block_size": 4}, + versions=_BLOCKED_VERSIONS, + ), + Variant( + "blocked_first_axis", + ((8, 3), (2, 3), (2, 3)), + {"axis": 0, "block_size": 4}, + versions=_BLOCKED_VERSIONS, + ), +) + +# The quantized products walk their operands exactly as MatMul and Conv do, so the shapes +# swept here are those tables' own, thinned to what the walk distinguishes; what is added is +# the zero points, present and absent, and the grids, whose types are free of one another. +_MATMUL_INTEGER_VARIANTS = ( + Variant("matrix", ((2, 3), (3, 4), (), ())), + Variant("wide", ((4, 8), (8, 4), (), ())), + Variant("no_zero_points", ((2, 3), (3, 4))), + Variant("left_zero_point_only", ((2, 3), (3, 4), ())), + Variant("single_element_zero_points", ((2, 3), (3, 4), (1,), (1,))), + Variant("batched", ((2, 3, 4), (2, 4, 3), (), ())), + Variant("unbatched_right", ((2, 3, 4), (4, 3), (), ())), + Variant("vector_matrix", ((3,), (3, 4), (), ())), + Variant("matrix_vector", ((2, 3), (3,), (), ())), + Variant("empty_rows", ((0, 3), (3, 4), (), ())), + Variant("empty_inner", ((2, 0), (0, 4), (), ())), + Variant( + "mixed_grids", + ((2, 3), (3, 4), (), ()), + operand_types={1: TensorProto.INT8, 3: TensorProto.INT8}, + ), +) + +# The scales of a quantized product are parameters, not data: the one factor the requantization +# comes to is `a_scale * b_scale / y_scale`, and pinning them puts the products it scales where +# both saturation and rounding are reached rather than leaving that to the draw. +# Both ops take them at the same three positions, so one table serves them both. +_PRODUCT_SCALES = {1: 0.02, 4: 0.03, 6: 0.01} + +# Every grid the second operand and the result can carry; the first operand's is the type the +# sweep itself ranges over, so the pairs below cover all eight combinations of the three. +_GRID_PAIRS = ( + ("int8_int8", TensorProto.INT8, TensorProto.INT8), + ("int8_uint8", TensorProto.INT8, TensorProto.UINT8), + ("uint8_int8", TensorProto.UINT8, TensorProto.INT8), + ("uint8_uint8", TensorProto.UINT8, TensorProto.UINT8), +) + +_QLINEAR_MATMUL_SHAPES = ( + ("matrix", (2, 3), (3, 4)), + ("wide", (4, 8), (8, 4)), + ("batched", (2, 3, 4), (2, 4, 3)), + ("unbatched_right", (2, 3, 4), (4, 3)), + ("vector_matrix", (3,), (3, 4)), + ("matrix_vector", (2, 3), (3,)), + ("empty_rows", (0, 3), (3, 4)), + ("empty_inner", (2, 0), (0, 4)), +) + + +def _qlinear_matmul_variants() -> tuple[Variant, ...]: + variants = [ + Variant(label, (left, (), (), right, (), (), (), ()), values=_PRODUCT_SCALES) + for label, left, right in _QLINEAR_MATMUL_SHAPES + ] + variants += [ + Variant( + f"grids_{label}", + ((2, 3), (), (), (3, 4), (), (), (), ()), + values=_PRODUCT_SCALES, + operand_types={3: right, 5: right, 7: result}, + ) + for label, right, result in _GRID_PAIRS + ] + variants.append( + Variant( + "single_element_parameters", + ((2, 3), (1,), (1,), (3, 4), (1,), (1,), (1,), (1,)), + values=_PRODUCT_SCALES, + ) + ) + return tuple(variants) + + +# The geometry both quantized convolutions are swept over: Conv's own table, thinned to the +# cases the walk distinguishes, and under `auto_pad` restricted to the shapes the reference +# resolves the mode correctly for -- the ones whose batch and channel extents repeat the +# spatial ones, for the reason Conv's table records. +_QUANTIZED_CONV_GEOMETRY: tuple[ + tuple[str, tuple[int, ...], tuple[int, ...], Mapping[str, Any]], ... +] = ( + ("spatial_1d", (2, 2, 7), (3, 2, 3), {}), + ("spatial_2d", (2, 3, 5, 4), (2, 3, 3, 2), {}), + ("pads", (2, 3, 5, 4), (2, 3, 3, 2), {"pads": [1, 1, 1, 1]}), + ("asymmetric_pads", (2, 3, 5, 4), (2, 3, 3, 2), {"pads": [2, 0, 1, 1]}), + ("strides", (2, 3, 5, 4), (2, 3, 3, 2), {"strides": [2, 2]}), + ("dilations", (2, 3, 5, 4), (2, 3, 3, 2), {"dilations": [2, 1]}), + ("groups", (2, 4, 5, 4), (6, 2, 3, 2), {"group": 2}), + ("depthwise", (2, 4, 5, 4), (4, 1, 3, 2), {"group": 4}), + ("spatial_3d", (1, 2, 4, 3, 3), (2, 2, 2, 2, 2), {}), + ( + "auto_pad_same_upper", + (3, 4, 3, 4), + (2, 4, 2, 3), + {"auto_pad": "SAME_UPPER", "strides": [2, 2]}, + ), + ( + "auto_pad_same_lower", + (3, 4, 3, 4), + (2, 4, 2, 3), + {"auto_pad": "SAME_LOWER", "strides": [2, 2]}, + ), +) + + +def _conv_integer_variants() -> tuple[Variant, ...]: + variants = [ + Variant(label, (source, weights, (), ()), attributes) + for label, source, weights, attributes in _QUANTIZED_CONV_GEOMETRY + ] + variants += [ + Variant("no_zero_points", ((2, 3, 5, 4), (2, 3, 3, 2))), + Variant("input_zero_point_only", ((2, 3, 5, 4), (2, 3, 3, 2), ())), + # One zero point per output channel. The reference stretches a filter's zero point + # over four axes whatever the operand's rank, so it is an oracle for this at two + # spatial axes and nowhere else; that the kernel addresses it by output channel at + # any rank is settled in the kernel tests. + Variant("per_channel_zero_point", ((2, 3, 5, 4), (2, 3, 3, 2), (), (2,))), + Variant( + "mixed_grids", + ((2, 3, 5, 4), (2, 3, 3, 2), (), ()), + operand_types={1: TensorProto.INT8, 3: TensorProto.INT8}, + ), + ] + return tuple(variants) + + +def _qlinear_conv_variants() -> tuple[Variant, ...]: + variants = [ + Variant( + label, + (source, (), (), weights, (), (), (), ()), + attributes, + values=_PRODUCT_SCALES, + ) + for label, source, weights, attributes in _QUANTIZED_CONV_GEOMETRY + ] + variants += [ + Variant( + "bias", + ((2, 3, 5, 4), (), (), (2, 3, 3, 2), (), (), (), (), (2,)), + values=_PRODUCT_SCALES, + ), + # One scale and one zero point per output channel, at the two spatial axes the + # reference stretches them over. + Variant( + "per_channel", + ((2, 3, 5, 4), (), (), (2, 3, 3, 2), (2,), (2,), (), ()), + values={**_PRODUCT_SCALES, 4: (0.03, 0.05)}, + ), + Variant( + "single_element_parameters", + ((2, 3, 5, 4), (1,), (1,), (2, 3, 3, 2), (1,), (1,), (1,), (1,)), + values=_PRODUCT_SCALES, + ), + ] + variants += [ + Variant( + f"grids_{label}", + ((2, 3, 5, 4), (), (), (2, 3, 3, 2), (), (), (), ()), + values=_PRODUCT_SCALES, + operand_types={3: weights, 5: weights, 7: result}, + ) + for label, weights, result in _GRID_PAIRS + ] + return tuple(variants) + + +# A pooling slides the same window as a convolution, so the sweep is that geometry again -- +# rank 1 through 3, the attributes that move the window and the ones that place it -- plus +# what only a pooling has: `ceil_mode`, which lets the last window hang off the end, and (for +# AveragePool) `count_include_pad`, which decides whether the padded positions are part of the +# divisor. The window is deliberately never wider than what the padding fills, since a window +# covering nothing at all averages over an empty array, which the reference raises on. +# +# Two restrictions come from the oracle, both about `auto_pad`. The reference resolves it +# through the *undilated* kernel, so the pads it derives are the ones ONNX defines only where +# the dilations are 1; and it refuses `ceil_mode` alongside it outright. What the compiler +# derives for the dilated case is settled instead by ONNX's own shape inference, in the kernel +# tests: a result shape the pads do not imply fails to compile at all. +_AVERAGE_POOL_VARIANTS = ( + Variant("spatial_1d", ((2, 2, 7),), {"kernel_shape": [3]}), + Variant("spatial_1d_strided", ((2, 2, 7),), {"kernel_shape": [3], "strides": [2]}), + Variant("spatial_1d_padded", ((2, 2, 7),), {"kernel_shape": [3], "pads": [2, 1]}), + Variant( + "spatial_1d_dilated", ((2, 2, 7),), {"kernel_shape": [3], "dilations": [2]} + ), + Variant("spatial_2d", ((2, 3, 5, 4),), {"kernel_shape": [3, 2]}), + Variant("pads", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "pads": [1, 1, 1, 1]}), + Variant( + "asymmetric_pads", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2], "pads": [2, 0, 1, 1]}, + ), + Variant("strides", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "strides": [2, 2]}), + Variant( + "dilations", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "dilations": [2, 1]} + ), + Variant( + "strided_dilated_padded", + ((2, 3, 5, 4),), + { + "kernel_shape": [3, 2], + "strides": [2, 1], + "dilations": [1, 2], + "pads": [1, 2, 0, 1], + }, + ), + # The padded positions counted, and not: the only thing that tells the two apart is a + # window the padding reaches into, so every one of these pads. + Variant( + "count_include_pad", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2], "pads": [1, 1, 1, 1], "count_include_pad": 1}, + ), + Variant( + "count_include_pad_asymmetric", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2], "pads": [2, 0, 1, 1], "count_include_pad": 1}, + ), + Variant( + "count_include_pad_strided", + ((2, 3, 5, 4),), + { + "kernel_shape": [3, 2], + "pads": [1, 1, 1, 1], + "strides": [2, 2], + "count_include_pad": 1, + }, + ), + # `ceil_mode` keeps a last window the floor would have dropped; the taps it reaches past + # the operand's own padding are read by nothing. + Variant( + "ceil_mode", + ((1, 1, 4, 4),), + {"kernel_shape": [3, 3], "strides": [2, 2], "ceil_mode": 1}, + ), + Variant( + "ceil_mode_dilated", + ((1, 1, 4, 4),), + { + "kernel_shape": [2, 2], + "strides": [1, 1], + "dilations": [2, 2], + "ceil_mode": 1, + }, + ), + # The position ONNX drops again: rounding up puts the last window's own start beyond the + # padding, where it would cover nothing but pad. + Variant( + "ceil_mode_last_window_starts_on_pad", + ((1, 3, 2, 2),), + { + "kernel_shape": [3, 3], + "pads": [1, 1, 1, 1], + "strides": [3, 3], + "ceil_mode": 1, + "count_include_pad": 1, + }, + ), + Variant("unit_window", ((2, 3, 5, 4),), {"kernel_shape": [1, 1]}), + Variant("spatial_3d", ((1, 2, 4, 3, 3),), {"kernel_shape": [2, 2, 2]}), + Variant( + "spatial_3d_strided_padded", + ((1, 2, 4, 3, 3),), + { + "kernel_shape": [2, 2, 2], + "strides": [2, 1, 2], + "pads": [1, 0, 1, 0, 1, 0], + }, + ), + # An odd total pad, which is what tells the two SAME modes apart: the extra one goes at + # the end for SAME_UPPER and at the beginning for SAME_LOWER. + Variant( + "auto_pad_same_upper", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "SAME_UPPER", "strides": [2, 2]}, + ), + Variant( + "auto_pad_same_lower", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "SAME_LOWER", "strides": [2, 2]}, + ), + Variant( + "auto_pad_same_upper_unit_stride", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "SAME_UPPER"}, + ), + Variant( + "auto_pad_same_lower_unit_stride", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "SAME_LOWER"}, + ), + Variant( + "auto_pad_valid", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "VALID", "strides": [2, 2]}, + ), + Variant("empty_batch", ((0, 2, 5, 4),), {"kernel_shape": [2, 2]}), +) + +# MaxPool is the same geometry, folded by comparison instead. The reference evaluator carries +# two implementations of it and picks between them by the attributes: a general one whenever +# any stride or dilation is other than 1, and — at unit strides and dilations — a padding-based +# one that re-pads the operand itself. That second path is no oracle for much: it unpacks 2-D +# `pads` in the wrong order, ignores them outright at any other rank, rounds `ceil_mode` up +# without dropping the window that then starts beyond the padding, and derives the `Indices` +# output through the wrong extents. So everything but the plainest unit-stride window carries a +# stride or a dilation, which is what puts the case on the general path. +# +# `auto_pad` narrows it once more: the general path reads SAME_LOWER as flooring the result and +# then splits the pad the SAME_UPPER way, which is neither of the two modes ONNX defines. The +# corpus's own `same_lower` tests are what vouch for that mode. +_MAX_POOL_VARIANTS = ( + # The one unit-stride, unit-dilation case, on the second path — and floating-point only, + # since that path pads the operand with a NaN no integer dtype can hold. + Variant( + "spatial_2d", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2]}, + elem_types=_FLOAT_ELEM_TYPES, + ), + Variant("spatial_1d_strided", ((2, 2, 7),), {"kernel_shape": [3], "strides": [2]}), + Variant( + "spatial_1d_padded", + ((2, 2, 7),), + {"kernel_shape": [3], "pads": [2, 1], "strides": [2]}, + ), + Variant( + "spatial_1d_dilated", ((2, 2, 7),), {"kernel_shape": [3], "dilations": [2]} + ), + Variant("strides", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "strides": [2, 2]}), + Variant( + "pads", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2], "pads": [1, 1, 1, 1], "strides": [2, 2]}, + ), + Variant( + "asymmetric_pads", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2], "pads": [2, 0, 1, 1], "strides": [2, 2]}, + ), + Variant( + "dilations", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "dilations": [2, 1]} + ), + Variant( + "strided_dilated_padded", + ((2, 3, 5, 4),), + { + "kernel_shape": [3, 2], + "strides": [2, 1], + "dilations": [1, 2], + "pads": [1, 2, 0, 1], + }, + ), + Variant( + "ceil_mode", + ((1, 1, 4, 4),), + {"kernel_shape": [3, 3], "strides": [2, 2], "ceil_mode": 1}, + ), + Variant( + "ceil_mode_last_window_starts_on_pad", + ((1, 3, 2, 2),), + { + "kernel_shape": [3, 3], + "pads": [1, 1, 1, 1], + "strides": [3, 3], + "ceil_mode": 1, + }, + ), + Variant( + "unit_window", ((2, 3, 5, 4),), {"kernel_shape": [1, 1], "strides": [2, 2]} + ), + Variant( + "spatial_3d", + ((1, 2, 4, 3, 3),), + {"kernel_shape": [2, 2, 2], "strides": [2, 1, 2]}, + ), + Variant( + "auto_pad_same_upper", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "SAME_UPPER", "strides": [2, 2]}, + ), + Variant( + "auto_pad_valid", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "VALID", "strides": [2, 2]}, + ), + Variant( + "empty_batch", ((0, 2, 5, 4),), {"kernel_shape": [2, 2], "strides": [2, 2]} + ), + # `Indices` reports where each maximum was read as one flat index into the whole operand, + # which `storage_order` lays out row-major or column-major. + Variant( + "indices", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2], "strides": [2, 2]}, + outputs=2, + ), + Variant( + "indices_column_major", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2], "strides": [2, 2], "storage_order": 1}, + outputs=2, + ), + Variant( + "indices_padded", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2], "pads": [1, 1, 1, 1], "strides": [2, 2]}, + outputs=2, + ), + Variant( + "indices_1d", + ((2, 2, 7),), + {"kernel_shape": [3], "strides": [2]}, + outputs=2, + ), + Variant( + "indices_3d_column_major", + ((1, 2, 4, 3, 3),), + {"kernel_shape": [2, 2, 2], "strides": [2, 1, 2], "storage_order": 1}, + outputs=2, + ), +) + +# LpPool is the same window under a norm. `ceil_mode` is left out of its sweep: for a window +# that rounding up clipped, the reference averages over the taps the window really holds and +# scales the result back up by the whole kernel's tap count, which its own source records as a +# computation borrowed from elsewhere that differs from the spec's -- and which the corpus's +# stored LpPool outputs, taken as the plain norm, disagree with. Only rounding up can clip a +# window, so everything below is geometry the two agree on. +_LP_POOL_VARIANTS = ( + Variant("spatial_1d", ((2, 2, 7),), {"kernel_shape": [3], "p": 3}), + Variant("spatial_2d", ((2, 3, 5, 4),), {"kernel_shape": [3, 2]}), + Variant("order_1", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "p": 1}), + Variant("order_3", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "p": 3}), + Variant( + "pads", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "pads": [1, 1, 1, 1], "p": 3} + ), + Variant( + "asymmetric_pads", + ((2, 3, 5, 4),), + {"kernel_shape": [3, 2], "pads": [2, 0, 1, 1]}, + ), + Variant("strides", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "strides": [2, 2]}), + Variant( + "dilations", ((2, 3, 5, 4),), {"kernel_shape": [3, 2], "dilations": [2, 1]} + ), + Variant( + "strided_dilated_padded", + ((2, 3, 5, 4),), + { + "kernel_shape": [3, 2], + "strides": [2, 1], + "dilations": [1, 2], + "pads": [1, 2, 0, 1], + "p": 3, + }, + ), + Variant("unit_window", ((2, 3, 5, 4),), {"kernel_shape": [1, 1]}), + Variant("spatial_3d", ((1, 2, 4, 3, 3),), {"kernel_shape": [2, 2, 2], "p": 1}), + Variant( + "auto_pad_same_upper", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "SAME_UPPER", "strides": [2, 2]}, + ), + Variant( + "auto_pad_same_lower", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "SAME_LOWER", "strides": [2, 2]}, + ), + Variant( + "auto_pad_valid", + ((2, 3, 5, 4),), + {"kernel_shape": [2, 3], "auto_pad": "VALID", "strides": [2, 2]}, + ), + Variant("empty_batch", ((0, 2, 5, 4),), {"kernel_shape": [2, 2]}), +) + +# The global poolings take one window the size of the operand's spatial extent, so there is +# no geometry to sweep — only the rank they take it over, and an operand with nothing in it. +# The mean of an empty batch is one the reference cannot take: it divides by a count it reads +# off the array, which is why only the two folds below carry that variant. +_GLOBAL_VARIANTS = ( + Variant("spatial_1d", ((2, 3, 7),)), + Variant("spatial_2d", ((2, 3, 5, 4),)), + Variant("spatial_3d", ((1, 2, 4, 3, 3),)), + Variant("single_position", ((2, 3, 1, 1),)), +) +# GlobalMaxPool and GlobalLpPool are compared against the windowed poolings they are defined +# to equal, which fold over the taps of a window rather than over an array — so the empty +# batch has an oracle here. +_GLOBAL_MAX_VARIANTS = (*_GLOBAL_VARIANTS, Variant("empty_batch", ((0, 2, 5, 4),))) +# The order of the norm is worth sweeping for GlobalLpPool, as it is for LpPool itself. +_GLOBAL_LP_VARIANTS = ( + *_GLOBAL_MAX_VARIANTS, + Variant("order_1", ((2, 3, 5, 4),), {"p": 1}), + Variant("order_3", ((2, 3, 5, 4),), {"p": 3}), +) + + +def _as_windowed_pooling(op_type: str) -> Callable[[Case], ModelProto]: + """The case's `Global*` pooling written as the windowed op its schema defines it to equal. + + Both schemas say it in the same words — "equivalent to [MaxPool/LpPool] with kernel size + equal to the spatial dimension of input tensor" — so the equivalent is that op over one + window covering the operand's whole spatial extent. + """ + + def build(case: Case) -> ModelProto: + spatial = (case.variant.shapes[0] or ())[2:] + return _model( + replace( + case, + op_type=op_type, + variant=replace( + case.variant, + attributes={ + **case.variant.attributes, + "kernel_shape": list(spatial), + }, + ), + ) + ) + + return build + + +# MaxUnpool scatters rather than folds, so every special value its dtype has goes in and has +# to come back out unchanged. Its indices are the op's own addressing — a draw would land +# outside the result, which ONNX leaves undefined and the artifact reports as an error status +# — so every variant pins them, including one that names the same position twice, where the +# last value written is what stays. +_MAX_UNPOOL_VARIANTS = ( + Variant( + "basic", + ((1, 1, 2, 2), (1, 1, 2, 2)), + {"kernel_shape": [2, 2], "strides": [2, 2]}, + values={1: (0, 3, 12, 15)}, + ), + Variant( + "wide", + ((1, 1, 4, 4), (1, 1, 4, 4)), + {"kernel_shape": [2, 2], "strides": [2, 2]}, + values={1: tuple(range(0, 64, 4))}, + ), + Variant( + "duplicate_positions", + ((1, 1, 2, 2), (1, 1, 2, 2)), + {"kernel_shape": [2, 2], "strides": [2, 2]}, + values={1: (5, 5, 5, 5)}, + ), + Variant( + "channels", + ((2, 3, 2, 2), (2, 3, 2, 2)), + {"kernel_shape": [2, 2], "strides": [2, 2]}, + values={1: tuple(range(0, 96, 4))}, + ), + Variant( + "padded", + ((1, 1, 2, 2), (1, 1, 2, 2)), + {"kernel_shape": [2, 2], "strides": [2, 2], "pads": [1, 1, 1, 1]}, + values={1: (0, 1, 2, 3)}, + ), + Variant( + "unit_stride", + ((1, 1, 2, 2), (1, 1, 2, 2)), + {"kernel_shape": [2, 2], "strides": [1, 1]}, + values={1: (0, 2, 6, 8)}, + ), + Variant( + "spatial_1d", + ((1, 1, 3), (1, 1, 3)), + {"kernel_shape": [2], "strides": [2]}, + values={1: (0, 3, 5)}, + ), + Variant( + "spatial_3d", + ((1, 1, 2, 2, 2), (1, 1, 2, 2, 2)), + {"kernel_shape": [2, 2, 2], "strides": [2, 2, 2]}, + values={1: (0, 9, 18, 27, 36, 45, 54, 63)}, + ), +) + + +# The reductions sweep both of ONNX's axes conventions: the attribute the family started +# with, and the input it moved to (ReduceSum at 13, the rest at 18). A variant is generated +# only at the revisions that take the form it declares. Axes are the op's configuration, so +# the input convention pins them — which is also what makes them compile-time constant, the +# only form the compiler can specialize a shape to. +_ATTRIBUTE_AXES_VARIANTS = ( + Variant("axis_1", ((2, 3, 4),), {"axes": [1]}), + Variant("axes_0_2_no_keepdims", ((2, 3, 4),), {"axes": [0, 2], "keepdims": 0}), + Variant("all_axes", ((2, 3, 4),)), + Variant("rank_1", ((5,),), {"axes": [0], "keepdims": 0}), + Variant("empty_result", ((0, 3),), {"axes": [1]}), +) +_INPUT_AXES_VARIANTS = ( + Variant("input_axis_1", ((2, 3, 4), (1,)), values={1: 1}), + Variant( + "input_axes_0_2_no_keepdims", + ((2, 3, 4), (2,)), + {"keepdims": 0}, + values={1: (0, 2)}, + ), + Variant("input_negative_axis", ((2, 3, 4), (1,)), values={1: -1}), + Variant("input_absent_axes", ((2, 3, 4),)), + Variant("input_empty_axes", ((2, 3, 4), (0,)), values={1: 0}), + Variant( + "input_noop_empty_axes", + ((2, 3, 4), (0,)), + {"noop_with_empty_axes": 1}, + values={1: 0}, + ), + Variant("input_rank_1", ((5,), (1,)), {"keepdims": 0}, values={1: 0}), + Variant("input_empty_result", ((0, 3), (1,)), values={1: 1}), +) +# Reducing a group with no elements at all, which every reduction answers with its identity. +_EMPTY_GROUP_VARIANTS = ( + Variant("empty_group", ((0, 3),), {"axes": [0]}), + Variant("input_empty_group", ((0, 3), (1,)), values={1: 0}), +) + +# The reference fills a reduction over no elements from ±inf, which `np.full(..., bool)` +# turns into True for a minimum as much as for a maximum — not a value to compare a kernel +# against, so the boolean families sit that one variant out. +_NON_BOOL_TYPES = tuple( + elem_type for elem_type in sorted(C_TYPES) if elem_type != TensorProto.BOOL +) + + +# The revisions of each convention: ReduceSum moved `axes` to an input at 13 and the rest of +# the family at 18, with ReduceMax and ReduceMin picking up the int8 families at 12 and the +# boolean ones at 20. +_REDUCTION_VERSIONS = ((1, 11, 13), (18,)) +_SUM_VERSIONS = ((1, 11), (13,)) +_EXTREMUM_VERSIONS = ((1, 11, 12, 13), (18, 20)) + + +def _reduction_sweep( + kind: Kind, + attribute_versions: tuple[int, ...], + input_versions: tuple[int, ...], + *, + elem_types: tuple[int, ...] | None = None, + empty_group_types: tuple[int, ...] | None = None, + factors: bool = False, +) -> Sweep: + """One reduction's shape and axes family, in both of ONNX's conventions. + + `elem_types` narrows the whole sweep and `empty_group_types` only the variant reducing a + group with no elements — each where the reference evaluator stops being an oracle. + `factors` marks a reduction that multiplies its group rather than adding it up. + """ + empty_group_types = empty_group_types if empty_group_types else elem_types + attribute_empty, given_empty = _EMPTY_GROUP_VARIANTS + return Sweep( + kind, + ( + *( + replace(variant, versions=attribute_versions, elem_types=elem_types) + for variant in _ATTRIBUTE_AXES_VARIANTS + ), + *( + replace(variant, versions=input_versions, elem_types=elem_types) + for variant in _INPUT_AXES_VARIANTS + ), + replace( + attribute_empty, + versions=attribute_versions, + elem_types=empty_group_types, + ), + replace(given_empty, versions=input_versions, elem_types=empty_group_types), + ), + operand_domains={0: Domain.SMALL_FACTOR} if factors else {}, + operand_types={1: TensorProto.INT64}, + constant_operands=(1,), + ) + + +# ArgMax and ArgMin take their axis as an attribute throughout; negative axes arrived at +# opset 11 and `select_last_index` at 12. +_ARG_VARIANTS = ( + Variant("axis_0", ((2, 3, 4),), {"axis": 0}), + Variant("axis_1_no_keepdims", ((2, 3, 4),), {"axis": 1, "keepdims": 0}), + Variant("rank_1", ((5,),), {"axis": 0}), + Variant("empty_result", ((0, 3),), {"axis": 1}), + Variant("negative_axis", ((2, 3, 4),), {"axis": -1}, versions=(11, 12, 13)), + Variant( + "select_last", + ((2, 3, 4),), + {"axis": 1, "select_last_index": 1}, + versions=(12, 13), + ), + Variant( + "select_last_negative_axis", + ((2, 3, 4),), + {"axis": -1, "select_last_index": 1}, + versions=(12, 13), + ), +) + +# Softmax, LogSoftmax and Hardmax normalize along one axis, leaving the shape alone. +_ALONG_AXIS_VARIANTS = ( + Variant("axis_0", ((2, 3, 4),), {"axis": 0}), + Variant("axis_1", ((2, 3, 4),), {"axis": 1}), + Variant("default_axis", ((2, 3, 4),)), + Variant("negative_axis", ((2, 3, 4),), {"axis": -2}), + Variant("rank_1", ((5,),), {"axis": 0}), + Variant("empty_groups", ((0, 3),), {"axis": 1}), + Variant("empty_group", ((0, 3),), {"axis": 0}), +) + +# The cumulative folds take their axis as an operand, which models — the backend corpus among +# them — pass at run time, so the sweep feeds it rather than pinning it into the model. +_CUMULATIVE_VARIANTS = ( + Variant("axis_0", ((2, 3, 4), ()), values={1: 0}), + Variant("axis_1_exclusive", ((2, 3, 4), ()), {"exclusive": 1}, values={1: 1}), + Variant("axis_2_reverse", ((2, 3, 4), ()), {"reverse": 1}, values={1: 2}), + Variant( + "negative_axis_exclusive_reverse", + ((2, 3, 4), ()), + {"exclusive": 1, "reverse": 1}, + values={1: -1}, + ), + Variant("rank_1", ((5,), ()), values={1: 0}), + Variant("empty", ((0, 3), ()), values={1: 1}), +) + + +# The normalizations. Each reduces a group of elements to a mean and a variance, so all of +# them carry ACCUMULATING: the summation order the spec leaves open is what a special value +# would make the comparison depend on. +# +# BatchNormalization's five operands are the data and four per-channel vectors; the variance +# it is handed at inference goes under a square root, so it is drawn non-negative. Training +# mode reduces every axis but the channel and reports the running statistics, which the +# `outputs` variants ask for. +_BATCH_DATA = (2, 3, 4, 5) +_BATCH_PARAMETER = (3,) + + +def _batch_variants() -> tuple[Variant, ...]: + def shapes(data: tuple[int, ...], channels: int) -> tuple[tuple[int, ...], ...]: + return (data, *((channels,),) * 4) + + ranks = ( + ("spatial_2d", _BATCH_DATA, 3), + ("spatial_1d", (2, 3, 4), 3), + ("no_spatial", (2, 3), 3), + # ONNX reads a tensor of rank 1 as a single channel of N instances. + ("rank_1", (4,), 1), + ("empty_instances", (0, 3, 2), 3), + ("empty_channels", (2, 0, 3), 0), + ("empty_spatial", (2, 3, 0), 3), + ) + variants = [ + Variant(label, shapes(data, channels)) for label, data, channels in ranks + ] + variants += [ + Variant("epsilon", shapes(_BATCH_DATA, 3), {"epsilon": 0.01}), + # ONNX's own shape inference refuses a training-mode node with anything but the + # three outputs, so every one of these reports the running statistics. + Variant("training", shapes(_BATCH_DATA, 3), {"training_mode": 1}, outputs=3), + Variant( + "training_momentum", + shapes(_BATCH_DATA, 3), + {"training_mode": 1, "momentum": 0.25, "epsilon": 0.01}, + outputs=3, + ), + Variant( + "training_empty_group", + shapes((0, 3, 2), 3), + {"training_mode": 1}, + outputs=3, + ), + Variant( + "training_empty_channels", + shapes((2, 0, 3), 0), + {"training_mode": 1}, + outputs=3, + ), + ] + return tuple(variants) + + +# LayerNormalization standardizes every row from `axis` on. The reference evaluator computes +# stage one in the tensor's own dtype and reports the mean and inverse deviation in it too, +# while ONNX derives both from `stash_type` — the two coincide at float32, which is therefore +# the only element type it is an oracle for here. +_LAYER_DATA = (2, 3, 4, 5) + + +def _layer_variants() -> tuple[Variant, ...]: + variants = [] + for axis in (0, 1, 2, 3, -1, -2, -4): + normalized = _LAYER_DATA[axis if axis >= 0 else axis + len(_LAYER_DATA) :] + label = f"axis_{axis}".replace("-", "negative_") + variants.append( + Variant(label, (_LAYER_DATA, normalized, normalized), {"axis": axis}) + ) + trailing = _LAYER_DATA[-1:] + variants += [ + Variant("default_axis", (_LAYER_DATA, trailing, trailing)), + Variant("no_bias", (_LAYER_DATA, trailing)), + Variant("epsilon", (_LAYER_DATA, trailing, trailing), {"epsilon": 0.01}), + Variant("stash_float", (_LAYER_DATA, trailing, trailing), {"stash_type": 1}), + Variant("statistics", (_LAYER_DATA, trailing, trailing), outputs=3), + Variant("statistics_no_bias", (_LAYER_DATA, trailing), outputs=3), + Variant("statistics_mean_only", (_LAYER_DATA, trailing, trailing), outputs=2), + # ONNX applies the scale and bias by broadcasting, so they need not carry the + # normalized shape at all. + Variant("broadcast_scale", (_LAYER_DATA, (), ()), {"axis": 3}), + Variant("rank_1", ((5,), (5,), (5,)), {"axis": 0}, outputs=3), + Variant("empty_groups", ((0, 3), (3,), (3,)), {"axis": 1}, outputs=3), + Variant("empty_group", ((2, 0), (0,), (0,)), {"axis": 1}, outputs=3), + ] + return tuple( + replace(variant, elem_types=(TensorProto.FLOAT,)) for variant in variants + ) + + +# InstanceNormalization standardizes each channel of each instance over its spatial axes, +# and GroupNormalization each group of channels; both take a per-channel scale and bias. +_INSTANCE_VARIANTS = ( + Variant("spatial_2d", ((2, 3, 4, 5), (3,), (3,))), + Variant("spatial_1d", ((2, 3, 4), (3,), (3,))), + Variant("no_spatial", ((2, 3), (3,), (3,))), + Variant("epsilon", ((2, 3, 4, 5), (3,), (3,)), {"epsilon": 0.01}), + Variant("empty_instances", ((0, 3, 2), (3,), (3,))), + Variant("empty_channels", ((2, 0, 3), (0,), (0,))), + Variant("empty_spatial", ((2, 3, 0), (3,), (3,))), +) + +_GROUP_DATA = (3, 4, 2, 2) +_GROUP_VARIANTS = ( + Variant("one_group", (_GROUP_DATA, (4,), (4,)), {"num_groups": 1}), + Variant("two_groups", (_GROUP_DATA, (4,), (4,)), {"num_groups": 2}), + # A group per channel, where the op is InstanceNormalization. + Variant("channel_groups", (_GROUP_DATA, (4,), (4,)), {"num_groups": 4}), + Variant("epsilon", (_GROUP_DATA, (4,), (4,)), {"num_groups": 2, "epsilon": 0.01}), + Variant( + "stash_float", (_GROUP_DATA, (4,), (4,)), {"num_groups": 2, "stash_type": 1} + ), + Variant("spatial_1d", ((2, 4, 3), (4,), (4,)), {"num_groups": 2}), + Variant("no_spatial", ((2, 4), (4,), (4,)), {"num_groups": 2}), + # No zero-element variant: the reference evaluates GroupNormalization's function body, + # whose `Reshape` to [0, 0, -1] numpy refuses outright on an empty tensor, so there is + # no oracle for one. The loops it would exercise are the ones the other members of this + # family carry empty variants for. +) + + +# RMSNormalization scales each row from `axis` on by the reciprocal root of its own mean +# square. `stash_type` is swept at its default alone: the reference refuses every other value +# outright, so there would be no oracle for one. +_RMS_DATA = (2, 3, 4, 5) + + +def _rms_variants() -> tuple[Variant, ...]: + variants = [] + for axis in (0, 1, 2, 3, -1, -2, -4): + normalized = _RMS_DATA[axis if axis >= 0 else axis + len(_RMS_DATA) :] + label = f"axis_{axis}".replace("-", "negative_") + variants.append(Variant(label, (_RMS_DATA, normalized), {"axis": axis})) + trailing = _RMS_DATA[-1:] + return ( + *variants, + Variant("default_axis", (_RMS_DATA, trailing)), + Variant("epsilon", (_RMS_DATA, trailing), {"epsilon": 0.01}), + Variant("stash_float", (_RMS_DATA, trailing), {"stash_type": 1}), + # ONNX applies the scale by broadcasting, so it need not carry the normalized shape. + Variant("broadcast_scale", (_RMS_DATA, ()), {"axis": 3}), + Variant("rank_1", ((5,), (5,)), {"axis": 0}), + Variant("empty_groups", ((0, 3), (3,)), {"axis": 1}), + Variant("empty_group", ((2, 0), (0,)), {"axis": 1}), + ) + + +# SoftmaxCrossEntropyLoss reads its logits as instances by classes by any further axes. The +# labels are a parameter rather than data — one outside the class axis is a read ONNX leaves +# undefined and the reference raises on — so every variant pins them, the `ignore_index` ones +# to values that exercise both the skipped and the counted branch. The weights are drawn +# non-negative so that the weighted mean's denominator cannot land on the zero that would make +# every expectation a NaN. +_SCE_2D = ((3, 5), (3,)) +_SCE_3D = ((3, 5, 2), (3, 2)) +_SCE_4D = ((2, 3, 2, 2), (2, 2, 2)) +_SCE_WEIGHTS = (5,) +_SCE_LABELS = { + _SCE_2D: (0, 4, 2), + _SCE_3D: (0, 1, 2, 3, 4, 0), + _SCE_4D: (0, 1, 2, 0, 2, 1, 0, 1), +} +_SCE_IGNORED = (0, -1, 2) + + +def _sce_variants() -> tuple[Variant, ...]: + def shaped( + label: str, + shapes: tuple[tuple[int, ...], ...], + attributes: Mapping[str, Any] | None = None, + **fields: Any, + ) -> Variant: + return Variant( + label, + shapes, + attributes or {}, + values={1: _SCE_LABELS[(shapes[0], shapes[1])]}, + domains={2: Domain.NONNEGATIVE}, + **fields, + ) + + variants = [ + shaped(f"{reduction}_{rank}d", shapes, {"reduction": reduction}) + for reduction in ("mean", "sum", "none") + for rank, shapes in ((2, _SCE_2D), (3, _SCE_3D), (4, _SCE_4D)) + ] + weighted = (*_SCE_2D, _SCE_WEIGHTS) + variants += [ + shaped("default_reduction", _SCE_2D), + shaped("weighted_mean", weighted, {"reduction": "mean"}), + shaped("weighted_sum", weighted, {"reduction": "sum"}), + shaped("weighted_none", weighted, {"reduction": "none"}), + shaped("weighted_3d", (*_SCE_3D, _SCE_WEIGHTS)), + shaped("log_prob", _SCE_2D, outputs=2), + shaped("log_prob_weighted", weighted, outputs=2), + shaped("log_prob_none", _SCE_2D, {"reduction": "none"}, outputs=2), + shaped("int32_labels", _SCE_2D, operand_types={1: TensorProto.INT32}), + # No zero-instance variant: the reference reshapes its log-softmax to `(N, C, -1)`, + # which numpy refuses on an empty array — it cannot infer the free axis from no + # elements at all — so there is no oracle for one. The kernel's own empty case is + # the `ignored_all` variant below, where nothing reaches either fold. + ] + variants += [ + Variant( + f"ignored_{label}", + shapes, + {"ignore_index": -1, **(attributes or {})}, + values={1: labels}, + domains={2: Domain.NONNEGATIVE}, + outputs=outputs, + ) + for label, shapes, labels, attributes, outputs in ( + ("mean", _SCE_2D, _SCE_IGNORED, None, 1), + ("sum", _SCE_2D, _SCE_IGNORED, {"reduction": "sum"}, 1), + ("none", _SCE_2D, _SCE_IGNORED, {"reduction": "none"}, 1), + ("weighted", weighted, _SCE_IGNORED, None, 1), + ("log_prob", _SCE_2D, _SCE_IGNORED, None, 2), + ("3d", _SCE_3D, (0, -1, 2, -1, 4, 0), None, 1), + # Every entry skipped: the weighted mean divides zero by zero. + ("all", _SCE_2D, (-1, -1, -1), None, 1), + ) + ] + return tuple(variants) + + +# LpNormalization divides each row along one axis by its own norm. The reference raises the +# elements to the power `p` without taking their absolute value, so at `p` = 1 it computes a +# signed sum rather than the norm ONNX defines: those variants feed non-negative operands, +# where the two agree. The zero-filled variants are the 0/0 the op answers with zero. +def _lp_variants() -> tuple[Variant, ...]: + shapes = ( + ("axis_0", (2, 3, 4), 0), + ("axis_1", (2, 3, 4), 1), + ("axis_negative_1", (2, 3, 4), -1), + ("rank_1", (5,), 0), + ("empty_rows", (0, 3), 1), + ("empty_group", (2, 0), 1), + ) + variants = [ + Variant(f"l2_{label}", (shape,), {"p": 2, "axis": axis}) + for label, shape, axis in shapes + ] + variants += [ + Variant( + f"l1_{label}", + (shape,), + {"p": 1, "axis": axis}, + domains={0: Domain.NONNEGATIVE}, + ) + for label, shape, axis in shapes + ] + variants += [ + Variant("l2_default", ((2, 3, 4),)), + Variant("l2_zeros", ((2, 3, 4),), {"p": 2, "axis": 1}, values={0: 0.0}), + Variant("l1_zeros", ((2, 3, 4),), {"p": 1, "axis": 1}, values={0: 0.0}), + ] + return tuple(variants) + + +# MeanVarianceNormalization is defined as an ONNX function, and the reference evaluates that +# body — whose epsilon constant is a float32, which makes the body itself ill-typed for a +# float64 tensor. float32 is therefore the only element type it can vouch for. +def _mvn_variants() -> tuple[Variant, ...]: + variants = [ + Variant("default_axes", ((2, 3, 4, 5),)), + Variant("axes_0", ((2, 3, 4, 5),), {"axes": [0]}), + Variant("axes_1", ((2, 3, 4, 5),), {"axes": [1]}), + Variant("axes_0_1", ((2, 3, 4, 5),), {"axes": [0, 1]}), + Variant("all_axes", ((2, 3, 4, 5),), {"axes": [0, 1, 2, 3]}), + Variant("negative_axes", ((2, 3, 4, 5),), {"axes": [-1]}), + Variant("rank_2", ((4, 3),), {"axes": [0]}), + Variant("empty_group", ((0, 3),), {"axes": [0]}), + Variant("empty_result", ((2, 0),), {"axes": [0]}), + ] + return tuple( + replace(variant, elem_types=(TensorProto.FLOAT,)) for variant in variants + ) + + +# LRN sums the squares of a window of channels around each element. The reference's channel +# loop is bounded by the *batch* extent rather than by the channel one, so it computes the +# whole result only where the two are equal; shapes that differ have no oracle here and rest +# on the corpus's own 5x5x5x5 expectations. +_LRN_VARIANTS = ( + Variant("window_3", ((3, 3, 2, 2),), {"size": 3}), + Variant("window_1", ((3, 3, 2, 2),), {"size": 1}), + # An even window reaches one further forward than back, and one wider than the tensor + # clamps at both ends. + Variant("window_2", ((3, 3, 2, 2),), {"size": 2}), + Variant("window_wider_than_channels", ((3, 3, 2, 2),), {"size": 5}), + Variant( + "scaled", + ((3, 3, 2, 2),), + {"size": 3, "alpha": 0.0002, "beta": 0.75, "bias": 2.0}, + ), + Variant("beta_one", ((3, 3, 2, 2),), {"size": 3, "beta": 1.0}), + Variant("singleton_spatial", ((2, 2, 1, 1),), {"size": 2}), + Variant("empty_spatial", ((2, 2, 0, 3),), {"size": 3}), + Variant("empty_channels", ((0, 0, 2, 2),), {"size": 3}), +) + + +# The views -- the ops that rearrange elements rather than compute them. Every one of them +# takes the shape of its result from an operand or an attribute, so the interesting axis is +# that description and not the values: the shapes and attributes below are the sweep, and the +# operand carrying them is pinned rather than drawn, since a random draw would be a shape. +# +# What a view kernel emits depends on the element type only through the C type it copies, so +# `_typed_at` sweeps the shape family at float32 and one variant per op at every type the +# schema allows -- crossing the two would repeat the same addressing once per dtype. +_FLOAT_ONLY = (TensorProto.FLOAT,) + + +def _typed_at(variants: tuple[Variant, ...], *labels: str) -> tuple[Variant, ...]: + """The variants named by `labels` at every element type, and the rest at float32.""" + return tuple( + variant + if variant.label in labels + else replace(variant, elem_types=variant.elem_types or _FLOAT_ONLY) + for variant in variants + ) + + +# Squeeze and Unsqueeze moved `axes` from an attribute to an operand at 13; a variant naming +# it one way does not apply to the revisions that take the other. +_AXES_ATTRIBUTE_VERSIONS = (1, 11) +_AXES_OPERAND_VERSIONS = (13, 21, 23, 24, 25) + +# Reshape's `allowzero` arrived at 14. Without it a zero in the shape copies the input's own +# extent; with it the zero is the extent. +_RESHAPE_ALLOWZERO_VERSIONS = (14, 19, 21, 23, 24, 25) + +_RESHAPE_VARIANTS = _typed_at( + ( + Variant("merge_dims", ((2, 3, 4), (2,)), values={1: (6, 4)}), + Variant("flatten", ((2, 3, 4), (1,)), values={1: 24}), + Variant("split_dims", ((2, 3, 4), (4,)), values={1: (2, 3, 2, 2)}), + Variant("negative_dim", ((2, 3, 4), (2,)), values={1: (2, -1)}), + # A zero copies the extent of the input's axis of that position. + Variant("zero_dim", ((2, 3, 4), (3,)), values={1: (2, 0, 4)}), + Variant("zero_and_negative_dim", ((2, 3, 4), (3,)), values={1: (2, 0, -1)}), + Variant("to_scalar", ((1,), (0,)), values={1: ()}), + Variant("from_scalar", ((), (1,)), values={1: 1}), + # A zero in the shape copies the input's own extent, so an empty result is asked for + # through the axis that carries it, or through the inferred one. + Variant("empty", ((0, 3), (2,)), values={1: (3, -1)}), + Variant( + "allowzero", + ((0, 3, 4), (3,)), + {"allowzero": 1}, + values={1: (3, 4, 0)}, + versions=_RESHAPE_ALLOWZERO_VERSIONS, + ), + ), + "merge_dims", +) + +_FLATTEN_VARIANTS = _typed_at( + ( + Variant("axis_2", ((2, 3, 4),), {"axis": 2}), + Variant("axis_0", ((2, 3, 4),), {"axis": 0}), + Variant("axis_1", ((2, 3, 4),), {"axis": 1}), + # An axis equal to the rank leaves every dimension in the first factor. + Variant("axis_rank", ((2, 3, 4),), {"axis": 3}), + Variant("default_axis", ((2, 3, 4),)), + Variant("negative_axis_1", ((2, 3, 4),), {"axis": -1}), + Variant("negative_axis_3", ((2, 3, 4),), {"axis": -3}), + Variant("rank_1", ((5,),), {"axis": 0}), + # The empty axis sits after `axis`: the reference reshapes to `(prod(shape[:axis]), + # -1)`, which numpy refuses outright when the first factor is itself zero, so there + # is no oracle for an empty dimension before it. + Variant("empty", ((2, 0, 4),), {"axis": 1}), + ), + "axis_2", +) + +# A variant naming no axes at all applies to both conventions: leaving the attribute out and +# leaving the operand out are the same instruction, and both mean every single dimension. +_SQUEEZE_VARIANTS = _typed_at( + ( + Variant("all_single_dims", ((1, 4, 1, 4),)), + Variant("nothing_to_squeeze", ((2, 3),)), + Variant("scalar", ((1,),)), + Variant( + "attribute_axis_0", + ((1, 3, 1, 4),), + {"axes": [0]}, + versions=_AXES_ATTRIBUTE_VERSIONS, + ), + Variant( + "attribute_axes_0_2", + ((1, 3, 1, 4),), + {"axes": [0, 2]}, + versions=_AXES_ATTRIBUTE_VERSIONS, + ), + Variant( + "operand_axis_0", + ((1, 3, 1, 4), (1,)), + values={1: 0}, + versions=_AXES_OPERAND_VERSIONS, + ), + Variant( + "operand_axes_0_2", + ((1, 3, 1, 4), (2,)), + values={1: (0, 2)}, + versions=_AXES_OPERAND_VERSIONS, + ), + Variant( + "operand_negative_axis", + ((1, 3, 1, 4), (1,)), + values={1: -2}, + versions=_AXES_OPERAND_VERSIONS, + ), + Variant( + "operand_no_axes", + ((2, 1, 3), (0,)), + values={1: ()}, + versions=_AXES_OPERAND_VERSIONS, + ), + Variant( + "empty", ((1, 0, 3),), {"axes": [0]}, versions=_AXES_ATTRIBUTE_VERSIONS + ), + Variant( + "operand_empty", + ((1, 0, 3), (1,)), + values={1: 0}, + versions=_AXES_OPERAND_VERSIONS, + ), + ), + "all_single_dims", +) + +_UNSQUEEZE_VARIANTS = _typed_at( + ( + Variant( + "attribute_axis_0", + ((3, 4),), + {"axes": [0]}, + versions=_AXES_ATTRIBUTE_VERSIONS, + ), + Variant( + "attribute_axes_0_3", + ((3, 4),), + {"axes": [0, 3]}, + versions=_AXES_ATTRIBUTE_VERSIONS, + ), + Variant( + "attribute_scalar", + ((),), + {"axes": [0]}, + versions=_AXES_ATTRIBUTE_VERSIONS, + ), + Variant("attribute_empty", ((0, 3),), {"axes": [1]}, versions=(11,)), + Variant( + "operand_axis_0", + ((4, 4), (1,)), + values={1: 0}, + versions=_AXES_OPERAND_VERSIONS, + ), + Variant( + "operand_axes_0_3", + ((3, 4), (2,)), + values={1: (0, 3)}, + versions=_AXES_OPERAND_VERSIONS, + ), + # ONNX resolves every axis against the *output's* rank, so an unsorted pair inserts + # the same dimensions a sorted one does. + Variant( + "operand_unsorted_axes", + ((3, 4), (2,)), + values={1: (3, 0)}, + versions=_AXES_OPERAND_VERSIONS, + ), + Variant( + "operand_negative_axes", + ((3, 4), (2,)), + values={1: (-1, -4)}, + versions=_AXES_OPERAND_VERSIONS, + ), + Variant( + "operand_scalar", ((), (1,)), values={1: 0}, versions=_AXES_OPERAND_VERSIONS + ), + Variant( + "operand_empty", + ((0, 3), (1,)), + values={1: 1}, + versions=_AXES_OPERAND_VERSIONS, + ), + ), + "operand_axis_0", +) + +_TRANSPOSE_VARIANTS = _typed_at( + ( + Variant("perm_2_0_1", ((2, 3, 4),), {"perm": [2, 0, 1]}), + Variant("perm_0_2_1", ((2, 3, 4),), {"perm": [0, 2, 1]}), + Variant("perm_1_0_2", ((2, 3, 4),), {"perm": [1, 0, 2]}), + Variant("identity_perm", ((2, 3, 4),), {"perm": [0, 1, 2]}), + Variant("default_perm", ((2, 3, 4),)), + Variant("rank_1", ((5,),), {"perm": [0]}), + Variant("rank_0", ((),)), + Variant("empty", ((0, 3, 4),), {"perm": [2, 0, 1]}), + ), + "perm_2_0_1", +) + +_CONCAT_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8), (4, 8), (4, 8)), {"axis": 0}), + Variant("axis_0", ((2, 3), (4, 3)), {"axis": 0}), + Variant("axis_1", ((2, 3), (2, 4), (2, 1)), {"axis": 1}), + Variant("negative_axis", ((2, 3, 4), (2, 3, 2)), {"axis": -1}), + Variant("inner_axis", ((2, 3, 4), (2, 1, 4)), {"axis": 1}), + Variant("single_operand", ((4, 8),), {"axis": 0}), + Variant("rank_1", ((5,), (3,)), {"axis": 0}), + Variant("empty_operand", ((0, 3), (2, 3)), {"axis": 0}), + Variant("empty_axis", ((2, 0), (2, 0)), {"axis": 1}), + ), + "wide", +) + +# Split's bands come from a `split` attribute up to 11, from an operand at 13, and from +# `num_outputs` -- or from an equal division of neither -- at 18. +_SPLIT_ATTRIBUTE_VERSIONS = (2, 11) +_SPLIT_OPERAND_VERSIONS = (13, 18) +_SPLIT_VARIANTS = _typed_at( + ( + Variant( + "operand_even", + ((6, 4), (2,)), + values={1: (3, 3)}, + outputs=2, + versions=_SPLIT_OPERAND_VERSIONS, + ), + Variant( + "operand_uneven", + ((6, 4), (3,)), + values={1: (1, 2, 3)}, + outputs=3, + versions=_SPLIT_OPERAND_VERSIONS, + ), + Variant( + "operand_zero_size_band", + ((6, 4), (3,)), + values={1: (0, 3, 3)}, + outputs=3, + versions=_SPLIT_OPERAND_VERSIONS, + ), + Variant( + "operand_axis_1", + ((4, 6), (2,)), + {"axis": 1}, + values={1: (2, 4)}, + outputs=2, + versions=_SPLIT_OPERAND_VERSIONS, + ), + Variant( + "operand_empty", + ((0, 4), (2,)), + values={1: (0, 0)}, + outputs=2, + versions=_SPLIT_OPERAND_VERSIONS, + ), + Variant( + "attribute_even", + ((6, 4),), + {"split": [3, 3]}, + outputs=2, + versions=_SPLIT_ATTRIBUTE_VERSIONS, + ), + Variant( + "attribute_uneven", + ((6, 4),), + {"split": [1, 2, 3]}, + outputs=3, + versions=_SPLIT_ATTRIBUTE_VERSIONS, + ), + Variant( + "attribute_axis_1", + ((4, 6),), + {"axis": 1, "split": [2, 4]}, + outputs=2, + versions=_SPLIT_ATTRIBUTE_VERSIONS, + ), + Variant( + "attribute_negative_axis", + ((4, 6),), + {"axis": -1, "split": [2, 4]}, + outputs=2, + versions=(11,), + ), + # Neither form: the axis is divided equally among the outputs the node declares. + Variant("equal_parts", ((6, 4),), outputs=2, versions=(2, 11, 13)), + Variant( + "num_outputs_even", ((6, 4),), {"num_outputs": 3}, outputs=3, versions=(18,) + ), + # An uneven division under `num_outputs` leaves the remainder in the last band. + Variant( + "num_outputs_uneven", + ((5, 4),), + {"num_outputs": 2}, + outputs=2, + versions=(18,), + ), + ), + "operand_even", + "attribute_even", +) + +# Slice-1 takes its bounds as attributes and has no steps; from 10 on they are operands. +_SLICE_ATTRIBUTE_VERSIONS = (1,) +_SLICE_OPERAND_VERSIONS = (10, 11, 13) +_SLICE_NEGATIVE_AXIS_VERSIONS = (11, 13) +_SLICE_VARIANTS = _typed_at( + ( + Variant( + "operand_default_axes", + ((2, 3, 4), (3,), (3,)), + values={1: (0, 1, 1), 2: (2, 3, 3)}, + versions=_SLICE_OPERAND_VERSIONS, + ), + Variant( + "operand_axes", + ((2, 3, 4), (2,), (2,), (2,)), + values={1: (1, 0), 2: (3, 3), 3: (1, 2)}, + versions=_SLICE_OPERAND_VERSIONS, + ), + Variant( + "operand_steps", + ((2, 3, 4), (2,), (2,), (2,), (2,)), + values={1: (0, 0), 2: (3, 4), 3: (1, 2), 4: (2, 3)}, + versions=_SLICE_OPERAND_VERSIONS, + ), + Variant( + "operand_negative_bounds", + ((2, 3, 4), (2,), (2,), (2,)), + values={1: (-2, -3), 2: (-1, -1), 3: (1, 2)}, + versions=_SLICE_OPERAND_VERSIONS, + ), + Variant( + "operand_negative_steps", + ((2, 3, 4), (2,), (2,), (2,), (2,)), + values={1: (-1, -1), 2: (-4, -5), 3: (1, 2), 4: (-1, -2)}, + versions=_SLICE_OPERAND_VERSIONS, + ), + # Bounds beyond the extent clamp to it, in both directions. + Variant( + "operand_out_of_bounds", + ((2, 3, 4), (2,), (2,), (2,)), + values={1: (-1000, 1000), 2: (1000, -1000), 3: (1, 2)}, + versions=_SLICE_OPERAND_VERSIONS, + ), + Variant( + "operand_empty_result", + ((2, 3, 4), (1,), (1,), (1,)), + values={1: 2, 2: 2, 3: 1}, + versions=_SLICE_OPERAND_VERSIONS, + ), + Variant( + "operand_empty_source", + ((0, 3), (1,), (1,), (1,)), + values={1: 0, 2: 2, 3: 1}, + versions=_SLICE_OPERAND_VERSIONS, + ), + Variant( + "operand_negative_axes", + ((2, 3, 4), (1,), (1,), (1,)), + values={1: 1, 2: 3, 3: -2}, + versions=_SLICE_NEGATIVE_AXIS_VERSIONS, + ), + Variant( + "attribute_default_axes", + ((2, 3, 4),), + {"starts": [0, 1], "ends": [2, 3]}, + versions=_SLICE_ATTRIBUTE_VERSIONS, + ), + Variant( + "attribute_axes", + ((2, 3, 4),), + {"starts": [1, 0], "ends": [3, 3], "axes": [1, 2]}, + versions=_SLICE_ATTRIBUTE_VERSIONS, + ), + # Only the end runs out of bounds here: ONNX's own shape inference gives the + # attribute form no static shape once the *start* leaves the extent, so a model + # with one is a compile error rather than a case. The operand form above sweeps + # both ends of the clamping. + Variant( + "attribute_out_of_bounds_end", + ((2, 3, 4),), + {"starts": [0], "ends": [1000], "axes": [1]}, + versions=_SLICE_ATTRIBUTE_VERSIONS, + ), + ), + "operand_default_axes", + "attribute_default_axes", +) + +_EXPAND_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8), (2,)), values={1: (4, 8)}), + Variant("stretch_axis", ((3, 1), (2,)), values={1: (3, 4)}), + Variant("new_axes", ((3, 1), (3,)), values={1: (2, 3, 6)}), + # The shape broadcasts against the operand rather than replacing it, so it may be + # shorter than the operand's own rank, and may stretch it on either side. + Variant("broadcast_shape", ((3, 1), (3,)), values={1: (1, 3, 4)}), + Variant("shorter_shape", ((2, 3, 1), (2,)), values={1: (3, 4)}), + Variant("scalar_source", ((), (2,)), values={1: (2, 3)}), + Variant("empty", ((0, 1), (2,)), values={1: (0, 4)}), + ), + "wide", +) + +_TILE_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8), (2,)), values={1: (2, 1)}), + Variant("every_axis", ((2, 3), (2,)), values={1: (2, 3)}), + Variant("one_axis", ((2, 3), (2,)), values={1: (1, 3)}), + Variant("no_repeats", ((2, 3), (2,)), values={1: (1, 1)}), + Variant("rank_3", ((2, 3, 4), (3,)), values={1: (2, 1, 2)}), + Variant("zero_repeat", ((2, 3), (2,)), values={1: (2, 0)}), + Variant("empty_source", ((0, 3), (2,)), values={1: (2, 2)}), + ), + "wide", +) + + +# -------------------------------------------------------------------------------------- +# Reading through an index, padding, and the ops built from their own coordinates +# -------------------------------------------------------------------------------------- + + +def _cycled(count: int, extent: int) -> tuple[int, ...]: + """`count` indices inside `[0, extent)`, striding so that every position is reached.""" + return tuple((step * 3) % extent for step in range(count)) + + +# Every index operand below is pinned rather than drawn. An index is a position into another +# operand, so a seeded draw would spend every case on the out-of-range value ONNX leaves +# undefined — which the artifact answers with an argument error and the reference with an +# exception. What is swept is what ONNX does define: positions at both ends of the axis, and +# the negative index counted back from it. +# +# An index tensor with no elements is left out for the same reason a wrong expectation would +# be: the reference returns it as shape `(0,)` whatever the axes around it measure, which is +# not the shape ONNX's own inference derives, so there is no oracle for that case. +_GATHER_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8), (3,)), {"axis": 0}, values={1: (3, 0, 2)}), + Variant("axis_1", ((2, 3, 4), (2,)), {"axis": 1}, values={1: (2, 0)}), + Variant("negative_axis", ((2, 3, 4), (2,)), {"axis": -1}, values={1: (3, 1)}), + Variant("default_axis", ((3, 4), (2,)), values={1: (2, 0)}), + Variant( + "negative_indices", ((2, 3, 4), (2,)), {"axis": 1}, values={1: (-1, -3)} + ), + Variant( + "index_matrix", ((3, 4), (2, 2)), {"axis": 0}, values={1: (2, 0, 1, 1)} + ), + Variant("scalar_index", ((3, 4), ()), {"axis": 1}, values={1: 2}), + Variant("rank_1", ((5,), (2,)), {"axis": 0}, values={1: (4, 0)}), + Variant( + "int32_indices", + ((3, 4), (2,)), + {"axis": 0}, + values={1: (2, 0)}, + operand_types={1: TensorProto.INT32}, + ), + # An axis with no elements is gathered along another one, which is a shape ONNX + # defines and a buffer the kernel must leave alone. + Variant("empty_source_axis", ((0, 3), (2,)), {"axis": 1}, values={1: (0, 2)}), + ), + "wide", +) + +_GATHER_ELEMENTS_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8), (4, 8)), {"axis": 1}, values={1: _cycled(32, 8)}), + Variant("axis_0", ((3, 4), (2, 4)), {"axis": 0}, values={1: _cycled(8, 3)}), + Variant("axis_1", ((3, 4), (3, 2)), {"axis": 1}, values={1: _cycled(6, 4)}), + # The reference resolves a negative axis only where both operands have the same + # shape: its own cross-section check reads `shape[dim + 1:]` unnormalized, and + # differing extents make it refuse the case rather than answer it. + Variant( + "negative_axis", ((3, 4), (3, 4)), {"axis": -1}, values={1: _cycled(12, 4)} + ), + Variant("negative_indices", ((3, 4), (3, 4)), {"axis": 1}, values={1: -1}), + Variant("rank_1", ((5,), (3,)), {"axis": 0}, values={1: (4, 0, 2)}), + Variant( + "rank_3", ((2, 3, 4), (2, 3, 2)), {"axis": 2}, values={1: _cycled(12, 4)} + ), + Variant( + "int32_indices", + ((3, 4), (3, 4)), + {"axis": 1}, + values={1: _cycled(12, 4)}, + operand_types={1: TensorProto.INT32}, + ), + ), + "wide", +) + +# GatherND's `batch_dims` arrived at 12. +_GATHER_ND_BATCH_VERSIONS = (12, 13) +_GATHER_ND_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8), (3, 1)), values={1: (0, 3, 2)}), + Variant("full_depth", ((3, 4), (2, 2)), values={1: (0, 1, 2, 3)}), + Variant("slices", ((2, 3, 4), (2, 2)), values={1: (0, 1, 1, 2)}), + Variant("index_rank_3", ((3, 4), (2, 2, 1)), values={1: (0, 1, 2, 0)}), + Variant("negative_indices", ((3, 4), (1, 2)), values={1: (-1, -2)}), + Variant( + "batch_dims_1", + ((2, 3, 4), (2, 2, 1)), + {"batch_dims": 1}, + values={1: (0, 2, 1, 1)}, + versions=_GATHER_ND_BATCH_VERSIONS, + ), + # Unlike Gather's, GatherND's reference reshapes an empty gather onto the shape ONNX + # infers for it, so the case has an oracle. + Variant("empty_indices", ((4, 8), (0, 1)), values={1: ()}), + ), + "wide", +) + +# Pad's pads are its configuration — where the operand sits inside the result — so they are +# carried as an initializer and pinned per variant, as are the axes they apply to. The fill +# value is pinned for the reason Clip's bounds are: one scalar cannot carry a dtype's edges, +# and which one it carried would be an accident of the seed. +# +# A negative pad, which ONNX defines as cropping, is absent: `np.pad` refuses one outright, +# so the reference evaluator is no oracle for it. +_PAD_ATTRIBUTE_MODES = ("constant", "reflect") +_PAD_AXES_VERSIONS = (18, 19, 21, 23, 24, 25) +_PAD_WRAP_VERSIONS = (19, 21, 23, 24, 25) +_PAD_MODE_VERSIONS = {"wrap": _PAD_WRAP_VERSIONS} +_PAD_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8), (4,), ()), values={1: (1, 2, 3, 0), 2: 5.0}), + *( + Variant( + mode, + ((2, 3, 4), (6,), ()), + {"mode": mode}, + values={1: (1, 0, 2, 1, 2, 0), 2: 3.5}, + versions=_PAD_MODE_VERSIONS.get(mode), + ) + for mode in ("constant", "edge", "reflect", "wrap") + ), + Variant("no_pads", ((2, 3), (4,), ()), values={1: (0, 0, 0, 0), 2: 1.0}), + Variant("one_sided", ((2, 3), (4,), ()), values={1: (0, 2, 3, 0), 2: 1.0}), + Variant("rank_1", ((5,), (2,), ()), values={1: (2, 3), 2: 1.0}), + Variant( + "fill_is_not_a_number", + ((2, 3), (4,), ()), + values={1: (1, 1, 1, 1), 2: float("nan")}, + elem_types=_FLOAT_ELEM_TYPES, + ), + # An operand with no elements has nothing to reflect or repeat, but a constant pad + # fills the whole result from the value alone. + Variant("empty_source", ((0, 3), (4,), ()), values={1: (1, 0, 1, 0), 2: 7.0}), + Variant( + "axes", + ((2, 3, 4), (2,), (), (1,)), + values={1: (1, 2), 2: 1.0, 3: (1,)}, + versions=_PAD_AXES_VERSIONS, + ), + Variant( + "negative_axes", + ((2, 3, 4), (4,), (), (2,)), + values={1: (1, 2, 0, 1), 2: 1.0, 3: (-1, -3)}, + versions=_PAD_AXES_VERSIONS, + ), + Variant( + "wrap_axes", + ((2, 3, 4), (2,), (), (1,)), + {"mode": "wrap"}, + values={1: (2, 3), 2: 1.0, 3: (2,)}, + versions=_PAD_WRAP_VERSIONS, + ), + # Up to opset 10 the pads and the fill are attributes instead, which is a kernel of + # its own. + *( + Variant( + f"attribute_{mode}", + ((2, 3, 4),), + {"pads": [1, 0, 2, 1, 2, 0], "value": 2.5, "mode": mode}, + versions=(2,), + ) + for mode in _PAD_ATTRIBUTE_MODES + ), + ), + "wide", +) + +# OneHot's depth is the extent of the axis it inserts, so it is an initializer; the two +# values it selects between are parameters, pinned for the reason Clip's bounds are — and +# pinned to ordinary numbers, since the reference reaches them through `off + (on - off) * m` +# rather than by selecting, which is the same value for every pair but an infinite one and no +# value at all for a boolean one, where numpy refuses to subtract. ONNX's own shape inference +# refuses indices of rank 0, so the smallest case here is rank 1. +_ONE_HOT_VALUES = (2, 7) +_ONE_HOT_VARIANTS = _typed_at( + ( + Variant( + "wide", + ((4, 8), (), (2,)), + values={0: _cycled(32, 5), 1: 5, 2: _ONE_HOT_VALUES}, + elem_types=_NON_BOOL_TYPES, + ), + Variant( + "default_axis", + ((3,), (), (2,)), + values={0: (0, 2, 1), 1: 3, 2: _ONE_HOT_VALUES}, + ), + Variant( + "axis_0", + ((3, 2), (), (2,)), + {"axis": 0}, + values={0: _cycled(6, 4), 1: 4, 2: _ONE_HOT_VALUES}, + ), + Variant( + "axis_1", + ((3, 2), (), (2,)), + {"axis": 1}, + values={0: _cycled(6, 4), 1: 4, 2: _ONE_HOT_VALUES}, + ), + Variant( + "negative_axis", + ((3, 2), (), (2,)), + {"axis": -2}, + values={0: _cycled(6, 4), 1: 4, 2: _ONE_HOT_VALUES}, + ), + Variant( + "negative_indices", + ((4,), (), (2,)), + values={0: (-1, -5, -3, 0), 1: 5, 2: _ONE_HOT_VALUES}, + ), + Variant( + "empty_indices", + ((0, 3), (), (2,)), + values={0: (), 1: 4, 2: _ONE_HOT_VALUES}, + ), + # ONNX types the indices and the depth as any numeric tensor, and folds an index into + # the depth whichever it is: a fractional one then matches no position at all. + Variant( + "float_indices", + ((4,), (), (2,)), + values={0: (1.5, 2.0, -1.0, 0.0), 1: 4, 2: _ONE_HOT_VALUES}, + operand_types={0: TensorProto.FLOAT}, + ), + Variant( + "int32_indices", + ((3,), (), (2,)), + values={0: (0, 2, 1), 1: 3, 2: _ONE_HOT_VALUES}, + operand_types={0: TensorProto.INT32}, + ), + Variant( + "float_depth", + ((3,), (), (2,)), + values={0: (0, 2, 1), 1: 3.0, 2: _ONE_HOT_VALUES}, + operand_types={1: TensorProto.FLOAT}, + ), + ), + "wide", +) + +_EYE_LIKE_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8),)), + Variant("square", ((5, 5),)), + Variant("tall", ((5, 3),)), + Variant("above_the_diagonal", ((4, 5),), {"k": 2}), + Variant("below_the_diagonal", ((4, 5),), {"k": -2}), + Variant("beyond_the_matrix", ((3, 3),), {"k": 5}), + Variant("to_double", ((3, 4),), {"dtype": TensorProto.DOUBLE}), + Variant("to_int64", ((3, 4),), {"dtype": TensorProto.INT64}), + Variant("empty", ((0, 3),)), + ), + "wide", +) + +# Trilu's diagonal decides no shape, so it stays a run-time operand — pinned all the same, +# since a drawn one would put every case on the far side of the matrix. +_TRILU_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8),)), + Variant("upper_default", ((4, 5),)), + Variant("lower", ((4, 5),), {"upper": 0}), + Variant("upper_offset", ((4, 5), ()), {"upper": 1}, values={1: 2}), + Variant("lower_offset", ((4, 5), ()), {"upper": 0}, values={1: -1}), + Variant("offset_zero", ((4, 5), ()), values={1: 0}), + Variant("offset_beyond_the_matrix", ((4, 5), ()), values={1: 9}), + Variant("batched", ((2, 3, 4),), {"upper": 0}), + Variant("one_row", ((1, 5),)), + Variant("one_column", ((5, 1),), {"upper": 0}), + Variant("empty", ((0, 3),)), + ), + "wide", +) + +# ONNX defines a sequence length as being in [1, s], and the artifact refuses anything else +# at run time rather than reading past its buffers, so the lengths are pinned inside it. +_REVERSE_SEQUENCE_VARIANTS = _typed_at( + ( + Variant( + "wide", + ((4, 8), (8,)), + {"batch_axis": 1, "time_axis": 0}, + values={1: (1, 2, 3, 4, 4, 1, 2, 3)}, + ), + Variant( + "batch_major", + ((3, 5), (3,)), + {"batch_axis": 0, "time_axis": 1}, + values={1: (1, 5, 3)}, + ), + Variant( + "time_major", + ((5, 3), (3,)), + {"batch_axis": 1, "time_axis": 0}, + values={1: (5, 1, 3)}, + ), + Variant("default_axes", ((4, 3), (3,)), values={1: (4, 2, 1)}), + Variant( + "full_length", + ((3, 4), (3,)), + {"batch_axis": 0, "time_axis": 1}, + values={1: 4}, + ), + Variant( + "single_step", + ((3, 4), (3,)), + {"batch_axis": 0, "time_axis": 1}, + values={1: 1}, + ), + Variant( + "rank_3", + ((2, 3, 4), (2,)), + {"batch_axis": 0, "time_axis": 1}, + values={1: (3, 2)}, + ), + Variant( + "empty_features", + ((2, 3, 0), (2,)), + {"batch_axis": 0, "time_axis": 1}, + values={1: (3, 1)}, + ), + Variant( + "empty_batch", + ((0, 3), (0,)), + {"batch_axis": 0, "time_axis": 1}, + values={1: ()}, + ), + ), + "wide", +) + +# `largest` and `sorted` arrived at 11, and `k` moved from an attribute to an operand at 10. +_TOP_K_RANKED_VERSIONS = (11, 24) +_TOP_K_VARIANTS = _typed_at( + ( + Variant("wide", ((4, 8), (1,)), {"axis": 1}, values={1: 3}, outputs=2), + Variant("axis_0", ((4, 3), (1,)), {"axis": 0}, values={1: 2}, outputs=2), + Variant("default_axis", ((3, 5), (1,)), values={1: 2}, outputs=2), + Variant( + "negative_axis", ((2, 3, 4), (1,)), {"axis": -2}, values={1: 2}, outputs=2 + ), + Variant("every_element", ((3, 4), (1,)), {"axis": 1}, values={1: 4}, outputs=2), + Variant("one_element", ((3, 4), (1,)), {"axis": 1}, values={1: 1}, outputs=2), + Variant("rank_1", ((5,), (1,)), {"axis": 0}, values={1: 3}, outputs=2), + Variant("empty_rows", ((0, 4), (1,)), {"axis": 1}, values={1: 2}, outputs=2), + # A group of equal values is what pins the tie-break: ONNX ranks the smaller index + # first, at either end of the ranking. + Variant("ties", ((3, 4), (1,)), {"axis": 1}, values={0: 1, 1: 2}, outputs=2), + Variant( + "smallest", + ((4, 8), (1,)), + {"axis": 1, "largest": 0}, + values={1: 3}, + outputs=2, + versions=_TOP_K_RANKED_VERSIONS, + ), + Variant( + "smallest_ties", + ((3, 4), (1,)), + {"axis": 1, "largest": 0}, + values={0: 1, 1: 2}, + outputs=2, + versions=_TOP_K_RANKED_VERSIONS, + ), + Variant( + "attribute_k", ((4, 8),), {"axis": 1, "k": 3}, outputs=2, versions=(1,) + ), + ), + "wide", +) + + +# The reference evaluator is an oracle for a slice of each recurrent op, and the slice is what +# is swept. All three implement one direction only (they raise outright for a bidirectional +# weight tensor), ignore `clip` and `sequence_lens`, and return two outputs at most. The LSTM +# additionally drops `activations` and `input_forget` on the floor -- its own source says +# "TODO: support overridden attributes" -- so `Y_c` never comes back either; the GRU does read +# `linear_before_reset`; and the RNN honours its activation, of the two it knows. Every case +# below therefore runs one forward direction, with any sequence length equal to the padded one. +# What the rest of the family computes is settled against onnxruntime in the kernel suite, and +# against the corpus's stored outputs in the conformance suite. +def _recurrent_shapes( + seq: int, + batch: int, + input_size: int, + *, + gates: int, + hidden: int, + through: int, +) -> tuple[tuple[int, ...] | None, ...]: + """`X, W, R` and the optional operands up to `through`, at one hidden size. + + The last two are the LSTM's alone: no other layer of the family carries a cell state, so + no other one reaches past `initial_h`. + """ + rows = gates * hidden + shapes: list[tuple[int, ...] | None] = [ + (seq, batch, input_size), + (1, rows, input_size), + (1, rows, hidden), + (1, 2 * rows), + (batch,), + (1, batch, hidden), + (1, batch, hidden), + (1, 3 * hidden), + ] + return tuple(shapes[: through + 1]) + + +_LSTM_HIDDEN = 5 +_LSTM_GATES = 4 * _LSTM_HIDDEN + +_lstm_shapes = partial(_recurrent_shapes, gates=4, hidden=_LSTM_HIDDEN) + +_LSTM_ATTRIBUTES: Mapping[str, Any] = {"hidden_size": _LSTM_HIDDEN} + +_LSTM_VARIANTS = ( + Variant("minimal", _lstm_shapes(3, 2, 4, through=2), _LSTM_ATTRIBUTES, outputs=2), + Variant("bias", _lstm_shapes(3, 2, 4, through=3), _LSTM_ATTRIBUTES, outputs=2), + # Every sequence runs to the padded length: the reference reads the operand and then + # ignores it, so it is an oracle for no other value. + Variant( + "sequence_lens", + _lstm_shapes(3, 2, 4, through=4), + _LSTM_ATTRIBUTES, + values={4: 3}, + outputs=2, + ), + Variant( + "initial_state", + ( + *_lstm_shapes(3, 2, 4, through=3), + None, + (1, 2, _LSTM_HIDDEN), + (1, 2, _LSTM_HIDDEN), + ), + _LSTM_ATTRIBUTES, + outputs=2, + ), + Variant( + "peepholes", + _lstm_shapes(3, 2, 4, through=7), + _LSTM_ATTRIBUTES, + values={4: 3}, + outputs=2, + ), + # `Y` alone, which is what the node's first output slot carries. + Variant("states_dropped", _lstm_shapes(3, 2, 4, through=2), _LSTM_ATTRIBUTES), + Variant( + "one_step", + _lstm_shapes(1, 3, 4, through=7), + _LSTM_ATTRIBUTES, + values={4: 1}, + outputs=2, + ), + Variant( + "one_item", + _lstm_shapes(4, 1, 1, through=7), + _LSTM_ATTRIBUTES, + values={4: 4}, + outputs=2, + ), + Variant( + "empty_batch", _lstm_shapes(3, 0, 4, through=2), _LSTM_ATTRIBUTES, outputs=2 + ), + # Layout 1 packs the batch outermost, so `X` reads as `[batch, seq, input]`. The + # reference's own layout path holds only for a single step -- it takes the batch size off + # the operand before transposing it, so a longer sequence broadcasts the initial state + # against the wrong axis, and it reports `Y_h` at a shape of its own -- which is also the + # shape the corpus's one batchwise test runs at. + Variant( + "batchwise", + ((3, 1, 4), (1, _LSTM_GATES, 4), (1, _LSTM_GATES, _LSTM_HIDDEN)), + {**_LSTM_ATTRIBUTES, "layout": 1}, + outputs=2, + ), +) + +_GRU_HIDDEN = 5 +_GRU_GATES = 3 * _GRU_HIDDEN + +_gru_shapes = partial(_recurrent_shapes, gates=3, hidden=_GRU_HIDDEN) + +_GRU_ATTRIBUTES: Mapping[str, Any] = {"hidden_size": _GRU_HIDDEN} + +_GRU_VARIANTS = ( + Variant("minimal", _gru_shapes(3, 2, 4, through=2), _GRU_ATTRIBUTES, outputs=2), + Variant("bias", _gru_shapes(3, 2, 4, through=3), _GRU_ATTRIBUTES, outputs=2), + # Every sequence runs to the padded length: the reference reads the operand and then + # ignores it, so it is an oracle for no other value. It also squeezes the operand's only + # axis before ignoring it, which is a batch of one or nothing at all. + Variant( + "sequence_lens", + _gru_shapes(3, 1, 4, through=4), + _GRU_ATTRIBUTES, + values={4: 3}, + outputs=2, + ), + Variant( + "initial_state", + (*_gru_shapes(3, 2, 4, through=3), None, (1, 2, _GRU_HIDDEN)), + _GRU_ATTRIBUTES, + outputs=2, + ), + # `linear_before_reset` is the one attribute of the family's the reference does read, and + # it moves both what the reset gate scales and where the candidate's recurrent bias is + # added -- which is why it is swept with a bias and without one. + Variant( + "linear_before_reset", + _gru_shapes(3, 2, 4, through=3), + {**_GRU_ATTRIBUTES, "linear_before_reset": 1}, + outputs=2, + ), + Variant( + "linear_before_reset_unbiased", + _gru_shapes(3, 2, 4, through=2), + {**_GRU_ATTRIBUTES, "linear_before_reset": 1}, + outputs=2, + ), + # `Y` alone, which is what the node's first output slot carries. + Variant("states_dropped", _gru_shapes(3, 2, 4, through=2), _GRU_ATTRIBUTES), + Variant("one_step", _gru_shapes(1, 3, 4, through=3), _GRU_ATTRIBUTES, outputs=2), + Variant( + "one_item", + _gru_shapes(4, 1, 1, through=5), + _GRU_ATTRIBUTES, + values={4: 4}, + outputs=2, + ), + Variant("empty_batch", _gru_shapes(3, 0, 4, through=2), _GRU_ATTRIBUTES, outputs=2), + # Layout 1 packs the batch outermost, so `X` reads as `[batch, seq, input]`. The + # reference's own layout path holds only for a single step, for the reason recorded on + # the LSTM's own batchwise case. + Variant( + "batchwise", + ((3, 1, 4), (1, _GRU_GATES, 4), (1, _GRU_GATES, _GRU_HIDDEN)), + {**_GRU_ATTRIBUTES, "layout": 1}, + outputs=2, + ), +) + +_RNN_HIDDEN = 5 + +_rnn_shapes = partial(_recurrent_shapes, gates=1, hidden=_RNN_HIDDEN) + +_RNN_ATTRIBUTES: Mapping[str, Any] = {"hidden_size": _RNN_HIDDEN} + +_RNN_VARIANTS = ( + Variant("minimal", _rnn_shapes(3, 2, 4, through=2), _RNN_ATTRIBUTES, outputs=2), + Variant("bias", _rnn_shapes(3, 2, 4, through=3), _RNN_ATTRIBUTES, outputs=2), + Variant( + "sequence_lens", + _rnn_shapes(3, 1, 4, through=4), + _RNN_ATTRIBUTES, + values={4: 3}, + outputs=2, + ), + Variant( + "initial_state", + (*_rnn_shapes(3, 2, 4, through=3), None, (1, 2, _RNN_HIDDEN)), + _RNN_ATTRIBUTES, + outputs=2, + ), + # The RNN is the one member of the family whose activation the reference reads, and + # `Affine` the one parameterized function it knows -- so it is also the only oracle in + # the suite for `activation_alpha` and `activation_beta` reaching the kernel at all. + Variant( + "affine", + _rnn_shapes(3, 2, 4, through=3), + { + **_RNN_ATTRIBUTES, + "activations": ["Affine"], + "activation_alpha": [0.5], + "activation_beta": [-0.25], + }, + outputs=2, + ), + Variant("states_dropped", _rnn_shapes(3, 2, 4, through=2), _RNN_ATTRIBUTES), + Variant("one_step", _rnn_shapes(1, 3, 4, through=3), _RNN_ATTRIBUTES, outputs=2), + Variant( + "one_item", + _rnn_shapes(4, 1, 1, through=5), + _RNN_ATTRIBUTES, + values={4: 4}, + outputs=2, + ), + Variant("empty_batch", _rnn_shapes(3, 0, 4, through=2), _RNN_ATTRIBUTES, outputs=2), + Variant( + "batchwise", + ((3, 1, 4), (1, _RNN_HIDDEN, 4), (1, _RNN_HIDDEN, _RNN_HIDDEN)), + {**_RNN_ATTRIBUTES, "layout": 1}, + outputs=2, + ), +) + + +# Resize maps every output position to a source coordinate and weights the elements around +# it, so the sweep is the three things that decide those two: `mode`, which chooses the +# weights, `coordinate_transformation_mode`, which chooses the map, and the shape family the +# scales themselves describe -- growing an axis, shrinking one, leaving one alone, and doing +# each at once. `mode` and the transformation are swept as a full cross product: which +# combinations interact is exactly what a hand-picked list would be guessing at. +# +# Its three operands are carried in the model as initializers rather than fed: they decide +# the shape of the result, and a model that computes one is a model the compiler refuses by +# design. The corpus's own Resize tests, which pass all three at run time against a declared +# result shape, are what exercise the other half -- the kernel reads them from their buffers +# either way. +_RESIZE_DATA = (1, 2, 4, 5) +_RESIZE_SCALES = (1.0, 1.0, 0.6, 2.5) + +# The element types the reference evaluator can be asked about: it casts its result through +# `saturate_cast`, which refuses a boolean outright, and the compiler refuses one too. +_RESIZE_ELEM_TYPES = tuple( + elem_type for elem_type in sorted(C_TYPES) if elem_type != TensorProto.BOOL +) + +# A region of interest per axis: the pair ONNX reads as the fraction of each axis the +# result is taken from, and one that reaches past the operand so that the extrapolation +# value is what lands there. +_RESIZE_ROI = (0.0, 0.0, 0.2, 0.1, 1.0, 1.0, 0.9, 0.8) +_RESIZE_ROI_OUTSIDE = (0.0, 0.0, -0.3, 0.1, 1.0, 1.0, 1.4, 0.8) + +_RESIZE_TRANSFORMS = ( + "half_pixel", + "half_pixel_symmetric", + "pytorch_half_pixel", + "align_corners", + "asymmetric", +) + + +def _resize_variants() -> tuple[Variant, ...]: + variants = [ + Variant( + f"{mode}_{transform}", + (_RESIZE_DATA, None, (4,)), + {"mode": mode, "coordinate_transformation_mode": transform}, + values={2: _RESIZE_SCALES}, + ) + for mode in ("nearest", "linear", "cubic") + for transform in _RESIZE_TRANSFORMS + ] + # The one transformation that reads the region of interest, and so takes one. + variants += [ + Variant( + f"{mode}_tf_crop_and_resize", + (_RESIZE_DATA, (8,), (4,)), + {"mode": mode, "coordinate_transformation_mode": "tf_crop_and_resize"}, + values={1: _RESIZE_ROI, 2: _RESIZE_SCALES}, + ) + for mode in ("nearest", "linear", "cubic") + ] + # The shape family: an axis grown, shrunk, left alone, and every rank the op serves. + variants += [ + Variant( + label, + (shape, None, (len(shape),)), + {"mode": "linear"}, + values={2: scales}, + ) + for label, shape, scales in ( + ("upsampled", _RESIZE_DATA, (1.0, 1.0, 2.0, 3.0)), + ("downsampled", _RESIZE_DATA, (1.0, 1.0, 0.6, 0.5)), + ("unchanged", _RESIZE_DATA, (1.0, 1.0, 1.0, 1.0)), + ("spatial_1d", (2, 3, 7), (1.0, 1.0, 1.7)), + ("spatial_3d", (1, 2, 3, 2, 4), (1.0, 1.0, 2.0, 0.5, 1.5)), + ("rank_1", (5,), (2.5,)), + ("single_element", (1, 1, 1, 1), (1.0, 1.0, 3.0, 1.0)), + # A scale that shrinks an axis to nothing at all, and an operand holding + # nothing to begin with. + ("empty_result", _RESIZE_DATA, (1.0, 1.0, 0.2, 1.0)), + ("empty_batch", (0, 2, 4, 5), (1.0, 1.0, 2.0, 0.5)), + ) + ] + # `sizes` states the result's extents instead, and a policy other than `stretch` takes + # one scale for every axis it names -- the smallest or the largest the sizes ask for. + variants += [ + Variant( + label, + (_RESIZE_DATA, None, None, (len(sizes),)), + {"mode": "linear", **attributes}, + values={3: sizes}, + ) + for label, sizes, attributes in ( + ("sizes", (1, 2, 3, 7), {}), + ("sizes_axes", (3, 7), {"axes": [2, 3]}), + ("sizes_axes_reversed", (7, 3), {"axes": [3, 2]}), + ( + "sizes_not_larger", + (3, 7), + {"axes": [2, 3], "keep_aspect_ratio_policy": "not_larger"}, + ), + ( + "sizes_not_smaller", + (3, 7), + {"axes": [2, 3], "keep_aspect_ratio_policy": "not_smaller"}, + ), + ( + "sizes_not_larger_every_axis", + (1, 2, 3, 7), + {"keep_aspect_ratio_policy": "not_larger"}, + ), + ) + ] + variants += [ + # `axes` names which axes the operands describe, in the order they describe them. + Variant( + "scales_axes", + (_RESIZE_DATA, None, (2,)), + {"mode": "cubic", "axes": [2, 3]}, + values={2: (0.6, 2.5)}, + ), + Variant( + "scales_axes_reversed", + (_RESIZE_DATA, None, (2,)), + {"mode": "cubic", "axes": [3, 2]}, + values={2: (2.5, 0.6)}, + ), + Variant( + "roi_axes", + (_RESIZE_DATA, (4,), None, (2,)), + { + "mode": "linear", + "axes": [2, 3], + "coordinate_transformation_mode": "tf_crop_and_resize", + }, + values={1: (0.2, 0.1, 0.9, 0.8), 3: (3, 7)}, + ), + # A region reaching past the operand, where the extrapolation value lands instead. + Variant( + "extrapolation", + (_RESIZE_DATA, (8,), (4,)), + { + "mode": "linear", + "coordinate_transformation_mode": "tf_crop_and_resize", + "extrapolation_value": 7.5, + }, + values={1: _RESIZE_ROI_OUTSIDE, 2: _RESIZE_SCALES}, + elem_types=_FLOAT_ELEM_TYPES, + ), + # ONNX allows the region in any floating-point type, and the arithmetic that reads + # it is that type's own. + Variant( + "roi_double", + (_RESIZE_DATA, (8,), (4,)), + {"mode": "linear", "coordinate_transformation_mode": "tf_crop_and_resize"}, + values={1: _RESIZE_ROI, 2: _RESIZE_SCALES}, + operand_types={1: TensorProto.DOUBLE}, + ), + # Every element type the schema allows and the oracle can be asked about, at one + # shape: what a resize does per element differs by type only in the narrowing at + # the end, and crossing the two would repeat the same walk once per type. + Variant( + "typed", + (_RESIZE_DATA, None, (4,)), + {"mode": "linear"}, + values={2: _RESIZE_SCALES}, + elem_types=_RESIZE_ELEM_TYPES, + ), + ] + # The neighbour a `nearest` resize rounds to, on scales that put coordinates on both + # sides of a half and exactly on an element. + variants += [ + Variant( + f"nearest_{nearest_mode}", + (_RESIZE_DATA, None, (4,)), + {"mode": "nearest", "nearest_mode": nearest_mode}, + values={2: scales}, + ) + for nearest_mode in ("round_prefer_floor", "round_prefer_ceil", "floor", "ceil") + for scales in (_RESIZE_SCALES,) + ] + variants += [ + Variant( + f"nearest_{nearest_mode}_asymmetric", + (_RESIZE_DATA, None, (4,)), + { + "mode": "nearest", + "nearest_mode": nearest_mode, + "coordinate_transformation_mode": "asymmetric", + }, + values={2: (1.0, 1.0, 3.0, 1.5)}, + ) + for nearest_mode in ("round_prefer_floor", "round_prefer_ceil", "floor", "ceil") + ] + # `antialias` widens the filter over the elements a shrinking axis merges, and + # `exclude_outside` drops the taps that fall off the end and renormalizes the rest. + variants += [ + Variant( + f"{mode}{'_antialias' if antialias else ''}" + f"{'_exclude_outside' if exclude_outside else ''}_{label}", + (_RESIZE_DATA, None, (4,)), + { + "mode": mode, + "antialias": antialias, + "exclude_outside": exclude_outside, + }, + values={2: scales}, + ) + for mode in ("linear", "cubic") + for antialias in (0, 1) + for exclude_outside in (0, 1) + for label, scales in ( + ("downsampled", (1.0, 1.0, 0.3, 0.5)), + ("upsampled", (1.0, 1.0, 2.0, 3.0)), + ) + if antialias or exclude_outside + ] + # The shape of the cubic filter itself, which `cubic_coeff_a` states. + variants += [ + Variant( + f"cubic_coeff_a_{index}", + (_RESIZE_DATA, None, (4,)), + {"mode": "cubic", "cubic_coeff_a": coefficient, "antialias": antialias}, + values={2: (1.0, 1.0, 0.4, 2.3)}, + ) + for index, (coefficient, antialias) in enumerate( + ((-0.5, 0), (-1.0, 0), (-0.5, 1)) + ) + ] + return _typed_at(tuple(variants), "typed") + + +_RESIZE_VARIANTS = _resize_variants() + +# Upsample is Resize's predecessor at a fixed set of its settings, so the sweep is the +# shapes and the two modes it carries; everything else about the walk is Resize's above. +_UPSAMPLE_VARIANTS = _typed_at( + tuple( + Variant( + f"{mode}_{label}", + (shape, (len(shape),)), + {"mode": mode}, + values={1: scales}, + ) + for mode in ("nearest", "linear") + for label, shape, scales in ( + ("integer_scales", (1, 2, 4, 5), (1.0, 1.0, 2.0, 3.0)), + ("fractional_scales", (1, 2, 4, 5), (1.0, 1.0, 1.7, 2.5)), + ("unit_scales", (1, 2, 4, 5), (1.0, 1.0, 1.0, 1.0)), + ("spatial_1d", (2, 3, 7), (1.0, 1.0, 1.5)), + ("rank_1", (5,), (2.5,)), + ("empty_batch", (0, 2, 4, 5), (1.0, 1.0, 2.0, 2.0)), + ) + ) + + ( + Variant( + "typed", + ((1, 2, 4, 5), (4,)), + {"mode": "nearest"}, + values={1: (1.0, 1.0, 2.0, 3.0)}, + elem_types=_RESIZE_ELEM_TYPES, + ), + ), + "typed", +) + + +def _as_resize(case: Case) -> ModelProto: + """The case's Upsample written as the Resize its successor's own spec defines it to equal. + + Resize's specification records the equivalence in as many words: `asymmetric` is + described there as the coordinate mapping "used by Resize-10 and Upsample", and `floor` + as the neighbour that revision takes. The reference evaluator's own Upsample is no + oracle -- it implements integer-scaled `nearest` and raises on everything else. + """ + variant = case.variant + return _model( + replace( + case, + op_type="Resize", + version=19, + variant=replace( + variant, + shapes=(variant.shapes[0], None, variant.shapes[1]), + values={2: variant.values[1]}, + attributes={ + **variant.attributes, + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + ), + ) + ) + + +# The two block shuffles move elements without computing any: their attributes decide the +# addressing and nothing else, so the sweep is the shape family crossed with the block, at +# every mode DepthToSpace groups its channels by. +_DEPTH_TO_SPACE_VARIANTS = _typed_at( + tuple( + Variant(f"{mode.lower()}_{label}", (shape,), {"blocksize": block, "mode": mode}) + for mode in ("DCR", "CRD") + for label, shape, block in ( + ("wide", (2, 8, 2, 3), 2), + # A block of one leaves every element where it is. + ("unit_block", (2, 3, 4, 5), 1), + ("cubed", (1, 9, 2, 2), 3), + ("single_position", (1, 4, 1, 1), 2), + ("empty_batch", (0, 8, 2, 3), 2), + ("empty_spatial", (2, 8, 0, 3), 2), + ) + ) + + (Variant("typed", ((2, 8, 2, 3),), {"blocksize": 2}),), + "typed", +) + +_SPACE_TO_DEPTH_VARIANTS = _typed_at( + tuple( + Variant(label, (shape,), {"blocksize": block}) + for label, shape, block in ( + ("wide", (2, 2, 6, 4), 2), + ("unit_block", (2, 3, 4, 5), 1), + ("cubed", (1, 2, 3, 6), 3), + ("single_position", (1, 3, 2, 2), 2), + ("empty_batch", (0, 2, 6, 4), 2), + ("empty_channels", (2, 0, 6, 4), 2), + ) + ) + + (Variant("typed", ((2, 2, 6, 4),), {"blocksize": 2}),), + "typed", +) + +# Col2Im folds a stack of blocks back into an image, so the sweep is the geometry that +# decides where each block sat -- the block's own extents, and the strides, dilations and +# pads that place it -- plus the overlap that makes an image position sum more than one. +# The extents are carried in the model: they are the shape of the result, and a model that +# computes one is a model the compiler refuses by design. +# +# ONNX's own reference returns nothing at all for an empty batch or an empty channel axis +# (its accumulator is never built), so there is no oracle for one; the emitted code's +# handling of those is asserted in the kernel tests instead. +# Summing truth values has no defined result, which is why the compiler refuses a boolean +# Col2Im outright; the error path is a kernel test of its own. +_COL2IM_ELEM_TYPES = tuple( + elem_type for elem_type in sorted(C_TYPES) if elem_type != TensorProto.BOOL +) + +_COL2IM_VARIANTS = _typed_at( + tuple( + Variant( + label, + (shape, (len(image),), (len(block),)), + attributes, + values={1: image, 2: block}, + ) + for label, shape, image, block, attributes in ( + ("blocks", (1, 5, 5), (5, 5), (1, 5), {}), + # Every interior position is reached by four blocks, which it sums. + ("overlapping", (1, 4, 16), (5, 5), (2, 2), {}), + ("pads", (1, 5, 15), (5, 5), (1, 5), {"pads": [0, 1, 0, 1]}), + # A pad that differs from the one at the other end of its own axis, and a stride + # that differs from the one on the other axis: only the beginnings shift where a + # block sat, so a geometry that read the ends instead computes this one wrong. + ( + "asymmetric", + (1, 6, 24), + (6, 5), + (3, 2), + {"strides": [1, 2], "pads": [2, 1, 0, 2]}, + ), + ("strides", (1, 9, 4), (5, 5), (3, 3), {"strides": [2, 2]}), + ("dilations", (1, 4, 5), (6, 6), (2, 2), {"dilations": [1, 5]}), + ("channels", (2, 12, 9), (4, 4), (2, 2), {}), + ("signal", (1, 3, 2), (6,), (3,), {"strides": [2]}), + ("volume", (1, 10, 12), (3, 4, 5), (1, 1, 5), {}), + ) + ) + + ( + Variant( + "typed", + ((1, 4, 16), (2,), (2,)), + values={1: (5, 5), 2: (2, 2)}, + elem_types=_COL2IM_ELEM_TYPES, + ), + ), + "typed", +) + +# GridSample is decided by three things: `mode`, which chooses the weights around a +# coordinate, `padding_mode`, which decides what a coordinate outside the operand reads, and +# `align_corners`, which decides what [-1, 1] spans. All three are swept as a full cross +# product: which combinations interact is exactly what a hand-picked list would be guessing +# at. The grid itself is drawn rather than pinned -- a standard normal puts roughly a third +# of its coordinates outside [-1, 1], which is what exercises the padding. +# +# The grid is float whatever the data holds: ONNX types it separately, and the reference +# denormalizes every coordinate through a float32 array whatever type it arrives in, so that +# is the width it is an oracle at. +_GRID_SAMPLE_DATA = (1, 2, 4, 5) +_GRID_SAMPLE_GRID = (1, 3, 6, 2) + +# `nearest` reads one element and computes nothing, so it serves every type the schema +# allows; the interpolating modes are floating-point only, and refused otherwise. +_GRID_SAMPLE_ELEM_TYPES = tuple(sorted(C_TYPES)) + + +def _grid_sample_variants() -> tuple[Variant, ...]: + variants = [ + Variant( + f"{mode}_{padding}_align{align}", + (_GRID_SAMPLE_DATA, _GRID_SAMPLE_GRID), + {"mode": mode, "padding_mode": padding, "align_corners": align}, + ) + for mode in ("nearest", "linear", "cubic") + for padding in ("zeros", "border", "reflection") + for align in (0, 1) + ] + # The rank family, and the shapes a sampling does not vary over. An empty result is + # absent because ONNX's reference returns a bare array for one rather than the tuple its + # own runner expects, which leaves nothing to compare against. + variants += [ + Variant(label, (data, grid), {"mode": mode}) + for label, data, grid, mode in ( + ("signal", (2, 3, 7), (2, 4, 1), "linear"), + ("volume", (1, 2, 3, 4, 5), (1, 2, 2, 2, 3), "linear"), + ("volume_cubic", (1, 1, 4, 4, 4), (1, 2, 2, 2, 3), "cubic"), + ("grown", (1, 2, 3, 3), (1, 7, 8, 2), "linear"), + ("shrunk", (1, 2, 7, 8), (1, 2, 2, 2), "cubic"), + ("single_element", (1, 1, 1, 1), (1, 2, 2, 2), "linear"), + # An axis with no elements at all: every coordinate falls outside it. + ("empty_spatial", (1, 1, 0, 4), (1, 2, 2, 2), "linear"), + ) + ] + return _typed_at( + tuple(variants) + + ( + Variant( + "typed", + (_GRID_SAMPLE_DATA, _GRID_SAMPLE_GRID), + {"mode": "nearest"}, + elem_types=_GRID_SAMPLE_ELEM_TYPES, + ), + ), + "typed", + ) + + +_GRID_SAMPLE_VARIANTS = _grid_sample_variants() + +# AffineGrid maps a regular grid through one transform per batch, so the sweep is the shape +# of that grid and the two spacings `align_corners` chooses between. `size` is carried in +# the model: it is the shape of the result. +# +# float32 alone: the reference casts its result to float32 whatever type the transform +# arrives in, so it is no oracle for a wider one. The double kernel is the same code at +# another type, and stands on the kernel tests' second oracle. +_AFFINE_GRID_VARIANTS = tuple( + Variant( + f"{label}_align{align}", + ((size[0], len(size) - 2, len(size) - 1), (len(size),)), + {"align_corners": align}, + values={1: size}, + elem_types=_FLOAT_ONLY, + ) + for align in (0, 1) + for label, size in ( + ("image", (2, 3, 5, 6)), + ("volume", (2, 1, 4, 5, 6)), + ("single_row", (1, 3, 1, 4)), + ("smallest", (3, 2, 2, 2)), + ("empty_batch", (0, 3, 5, 6)), + ) +) + +# RoiAlign divides each region into bins and folds a grid of bilinear samples in each, so +# the sweep is what decides the region (`spatial_scale`, the coordinate transformation), what +# decides the bins (`output_height`/`output_width`) and what decides the samples +# (`sampling_ratio`, `mode`), crossed. The regions themselves are drawn: a standard normal +# against a small feature map puts them inside it, across its edge and inverted. +# +# The batch index each region reads is pinned rather than drawn -- it is a choice of plane, +# not data, and one outside the batch is an argument the artifact rejects at run time. +_ROI_ALIGN_DATA = (2, 3, 6, 5) +_ROI_ALIGN_BATCHES = (0, 1, 1, 0) + + +def _roi_align_variants() -> tuple[Variant, ...]: + shapes = (_ROI_ALIGN_DATA, (len(_ROI_ALIGN_BATCHES), 4), (len(_ROI_ALIGN_BATCHES),)) + variants = [ + Variant( + f"{mode}_{transform}_ratio{ratio}", + shapes, + { + "mode": mode, + "coordinate_transformation_mode": transform, + "output_height": 2, + "output_width": 3, + "sampling_ratio": ratio, + }, + values={2: _ROI_ALIGN_BATCHES}, + ) + for mode in ("avg", "max") + for transform in ("half_pixel", "output_half_pixel") + # A ratio of zero takes as many samples per bin as the bin spans elements. + for ratio in (0, 1, 2, 3) + ] + variants += [ + Variant( + label, + (data, (rois, 4), (rois,)), + {**attributes, "sampling_ratio": 2}, + values={2: _ROI_ALIGN_BATCHES[:rois]}, + ) + for label, data, rois, attributes in ( + ("scaled_down", _ROI_ALIGN_DATA, 4, {"spatial_scale": 0.5}), + ("scaled_up", _ROI_ALIGN_DATA, 4, {"spatial_scale": 2.0}), + ("single_bin", _ROI_ALIGN_DATA, 4, {}), + ("tall_bins", _ROI_ALIGN_DATA, 4, {"output_height": 5, "output_width": 1}), + ("single_region", _ROI_ALIGN_DATA, 1, {"output_height": 3}), + ("no_regions", _ROI_ALIGN_DATA, 0, {"output_height": 2}), + ("empty_channels", (2, 0, 6, 5), 4, {"output_height": 2}), + ) + ] + return tuple(variants) + + +_ROI_ALIGN_VARIANTS = _roi_align_variants() + +# MaxRoiPool rounds each region to whole elements and takes the largest in each bin, so +# nothing it computes can leave the dtype's range and every special value goes in. Its +# regions are pinned: the first of the five columns is the plane the region is read from, +# which a draw would put outside the batch. The table covers a region inside the image, one +# reaching past it, one before it, an inverted one and one that names a single element. +_MAX_ROI_POOL_DATA = (2, 2, 6, 5) +_MAX_ROI_POOL_REGIONS = ( + (0, 0, 0, 3, 3), + (1, 1.5, 0.5, 4.5, 3.5), + (0, -3, -2, 2, 1), + (1, 2, 4, 8, 9), + (1, 3, 2, 1, 1), + (0, 2, 2, 2, 2), +) + +# float32 alone: onnxruntime, the only implementation ONNX has for this op and so the +# sweep's oracle, registers it for that type only. The double kernel is the same code at +# another type, and stands on the kernel tests' equivalence between the two. +_MAX_ROI_POOL_VARIANTS = tuple( + Variant( + label, + (data, (len(_MAX_ROI_POOL_REGIONS), 5)), + {"pooled_shape": pooled, **attributes}, + values={ + 1: tuple(value for region in _MAX_ROI_POOL_REGIONS for value in region) + }, + elem_types=_FLOAT_ONLY, + ) + for label, data, pooled, attributes in ( + ("bins", _MAX_ROI_POOL_DATA, [2, 2], {}), + ("single_bin", _MAX_ROI_POOL_DATA, [1, 1], {}), + ("tall_bins", _MAX_ROI_POOL_DATA, [5, 1], {}), + # More bins than the region has elements, so some of them pool nothing at all. + ("finer_than_the_region", _MAX_ROI_POOL_DATA, [7, 6], {}), + ("scaled_down", _MAX_ROI_POOL_DATA, [2, 2], {"spatial_scale": 0.5}), + ("scaled_up", _MAX_ROI_POOL_DATA, [2, 2], {"spatial_scale": 2.0}), + ("empty_channels", (2, 0, 6, 5), [2, 2], {}), + ) +) + +# The opset onnxruntime serves MaxRoiPool at. ONNX revised the op once, at 22, and only to +# widen its type constraints -- which `test_the_maxroipool_oracle_runs_the_same_op` checks +# against the schemas themselves rather than taking on trust. +_MAX_ROI_POOL_ORACLE_VERSION = 21 + + +def _onnxruntime_outputs(model: ModelProto, feeds: Mapping[str, Any]) -> list[Any]: + """The second oracle the compiler's parity test already stands on, run on `model`.""" + runtime = pytest.importorskip("onnxruntime") + runtime.set_default_logger_severity(3) + session = runtime.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + return list(session.run(None, dict(feeds))) + + +def _at_the_oracle_opset(case: Case) -> ModelProto: + """The case's node at the newest opset its oracle implements it at.""" + return _model(replace(case, version=_MAX_ROI_POOL_ORACLE_VERSION)) + + +# -------------------------------------------------------------------------------------- +# Writing through an index +# -------------------------------------------------------------------------------------- + +# Every index operand below is pinned rather than drawn, for the reason the gathering ops' +# are: a seeded draw would spend every case on the out-of-range value ONNX leaves undefined. +# What is swept is what ONNX does define — positions at both ends of the axis, the negative +# index counted back from it, and the duplicates a `reduction` folds together. +# +# The rank family stops at 4: ONNX's reference implements ScatterElements for indices of +# rank 1 through 4 and raises outright for anything deeper, so there is no oracle past it. +_SCATTER_ELEMENTS_SHAPES = ( + Variant("wide", ((4, 8), (4, 8), (4, 8)), {"axis": 1}, values={1: _cycled(32, 8)}), + # `indices` may be shorter than `data` on the axes it does not write along, which leaves + # the elements past it alone. + Variant("axis_0", ((3, 4), (2, 4), (2, 4)), {"axis": 0}, values={1: _cycled(8, 3)}), + Variant("default_axis", ((3, 4), (2, 4), (2, 4)), values={1: _cycled(8, 3)}), + Variant( + "negative_axis", + ((3, 4), (3, 2), (3, 2)), + {"axis": -1}, + values={1: _cycled(6, 4)}, + ), + Variant("negative_indices", ((3, 4), (3, 4), (3, 4)), {"axis": 1}, values={1: -1}), + Variant("rank_1", ((5,), (3,), (3,)), values={1: (4, 0, 2)}), + Variant( + "rank_3", + ((2, 3, 4), (2, 3, 2), (2, 3, 2)), + {"axis": 2}, + values={1: _cycled(12, 4)}, + ), + Variant( + "rank_4", + ((2, 2, 3, 4), (2, 2, 3, 2), (2, 2, 3, 2)), + {"axis": 3}, + values={1: _cycled(24, 4)}, + ), + Variant( + "int32_indices", + ((3, 4), (3, 4), (3, 4)), + {"axis": 1}, + values={1: _cycled(12, 4)}, + operand_types={1: TensorProto.INT32}, + ), + # Nothing to write at all, which still has to leave the operand's own values behind. + Variant("no_updates", ((3, 4), (0, 4), (0, 4)), values={1: ()}), +) + +# The folds, on indices that name the same element more than once — which is what tells a +# fold from a plain write, and what ONNX defines `reduction` for. Every element type goes +# through each of them: the fold is emitted per type, and two of them are the type's own +# (a boolean sum is a disjunction, an integer extremum has no NaN to reckon with). +# +# `add` and `mul` pull their integer operands into the range no sum or product of a handful +# of them can leave: ONNX does not define integer overflow and C's is undefined for the +# signed families, so a wrap-around difference would not be a divergence from the spec. +_FOLD_DOMAINS = {0: Domain.SMALL_FACTOR, 2: Domain.SMALL_FACTOR} +_SCATTER_ELEMENTS_REDUCTIONS = tuple( + Variant( + f"reduction_{reduction}", + ((4, 8), (4, 8), (4, 8)), + {"axis": 1, "reduction": reduction}, + values={1: _cycled(32, 4)}, + domains=_FOLD_DOMAINS if reduction in ("add", "mul") else {}, + ) + for reduction in ("add", "mul", "max", "min") +) + +_SCATTER_ELEMENTS_VARIANTS = ( + _typed_at(_SCATTER_ELEMENTS_SHAPES, "wide") + _SCATTER_ELEMENTS_REDUCTIONS +) + +# Scatter is compiled by the generator its successor is, so what its own sweep has to show +# is that the deprecated op dispatches and computes what ONNX replaced it with; the surface +# above is not swept a second time. +_SCATTER_VARIANTS = _typed_at( + tuple( + variant + for variant in _SCATTER_ELEMENTS_SHAPES + if variant.label in ("wide", "axis_0", "negative_indices", "rank_1") + ), + "wide", +) + + +def _as_scatter_elements(case: Case) -> ModelProto: + """The case's Scatter written as the ScatterElements ONNX replaced it with. + + Scatter's own document says it in as many words: "This operator is deprecated. Please use + ScatterElements, which provides the same functionality." The reference evaluator carries + no implementation of Scatter at all, so the successor's — on the same operands — is the + oracle. + """ + return _model(replace(case, op_type="ScatterElements", version=18)) + + +# ScatterND writes a slice per index tuple, so its sweep is the depth of those tuples — from +# one axis of the operand to all of them — crossed with the shapes around it. +_SCATTER_ND_SHAPES = ( + Variant("wide", ((4, 8), (3, 1), (3, 8)), values={1: (0, 3, 2)}), + Variant("full_depth", ((3, 4), (2, 2), (2,)), values={1: (0, 1, 2, 3)}), + Variant("slices", ((2, 3, 4), (2, 2), (2, 4)), values={1: (0, 1, 1, 2)}), + Variant("index_rank_3", ((3, 4), (2, 2, 1), (2, 2, 4)), values={1: (0, 1, 2, 0)}), + Variant("negative_indices", ((3, 4), (1, 2), (1,)), values={1: (-1, -2)}), + Variant("rank_1", ((5,), (2, 1), (2,)), values={1: (4, 0)}), + Variant("no_updates", ((4, 8), (0, 1), (0, 8)), values={1: ()}), +) + +# The folds again, on tuples naming the same slice twice. They stay at depth 1: the +# reference's `max` and `min` branch indexes its output with the index array rather than the +# tuple built from it, which numpy answers with the wrong shape from depth 2 on and refuses +# to store — so it is an oracle for those two at depth 1 only, and the four are kept together +# rather than sweeping the same fold at two depths for half of them. +_SCATTER_ND_REDUCTIONS = tuple( + Variant( + f"reduction_{reduction}", + ((4, 8), (4, 1), (4, 8)), + {"reduction": reduction}, + values={1: (0, 2, 0, 2)}, + domains=_FOLD_DOMAINS if reduction in ("add", "mul") else {}, + ) + for reduction in ("add", "mul", "max", "min") +) + +_SCATTER_ND_VARIANTS = ( + _typed_at(_SCATTER_ND_SHAPES, "wide") + + _SCATTER_ND_REDUCTIONS + # A fold onto single elements rather than onto slices. + + ( + Variant( + "reduction_add_full_depth", + ((3, 4), (4, 2), (4,)), + {"reduction": "add"}, + values={1: (0, 1, 0, 1, 2, 3, 2, 3)}, + domains=_FOLD_DOMAINS, + elem_types=_FLOAT_ONLY, + ), + ) +) + +# TensorScatter writes each sample's update into that sample's own place in a cache, so its +# sweep is what decides that place: the axis the sequence runs along, the index the write +# starts at, whether the op is given those indices at all, and the mode that says what +# happens once a write runs past the end of the axis. +_TENSOR_SCATTER_CACHE = (2, 1, 4, 5) +_TENSOR_SCATTER_UPDATE = (2, 1, 2, 5) + +_TENSOR_SCATTER_VARIANTS = _typed_at( + ( + Variant( + "wide", + (_TENSOR_SCATTER_CACHE, _TENSOR_SCATTER_UPDATE, (2,)), + values={2: (1, 2)}, + ), + # Without the operand every sample is written from the start of the axis. + Variant("appended", (_TENSOR_SCATTER_CACHE, _TENSOR_SCATTER_UPDATE, None)), + Variant( + "circular", + (_TENSOR_SCATTER_CACHE, _TENSOR_SCATTER_UPDATE, (2,)), + {"mode": "circular"}, + values={2: (3, 2)}, + ), + Variant( + "circular_negative", + (_TENSOR_SCATTER_CACHE, _TENSOR_SCATTER_UPDATE, (2,)), + {"mode": "circular"}, + values={2: (-1, -3)}, + ), + Variant( + "circular_appended", + (_TENSOR_SCATTER_CACHE, _TENSOR_SCATTER_UPDATE, None), + {"mode": "circular"}, + ), + Variant("rank_3", ((3, 4, 5), (3, 2, 5), (3,)), values={2: (0, 1, 2)}), + Variant( + "axis_1", ((3, 4, 5), (3, 2, 5), (3,)), {"axis": 1}, values={2: (2, 0, 1)} + ), + # The sequence axis is the last one, so each write moves a single element. + Variant( + "last_axis", ((3, 4), (3, 2), (3,)), {"axis": 1}, values={2: (0, 2, 1)} + ), + Variant( + "full_sequence", + (_TENSOR_SCATTER_CACHE, _TENSOR_SCATTER_CACHE, (2,)), + values={2: (0, 0)}, + ), + Variant( + "single_step", + (_TENSOR_SCATTER_CACHE, (2, 1, 1, 5), (2,)), + values={2: (3, 0)}, + ), + # A batch wider than the axis the mode wraps against: ONNX takes the whole cache + # coordinate modulo the capacity, so the sample being written wraps along with the + # position inside it. + Variant( + "batch_wraps", + ((5, 3, 1), (5, 1, 1), (5,)), + {"mode": "circular"}, + values={2: (2, 2, 2, 2, 2)}, + ), + Variant("empty_batch", ((0, 1, 4, 5), (0, 1, 2, 5), (0,)), values={2: ()}), + ), + "wide", +) + + +def _with_attributes( + base: tuple[Variant, ...], combinations: Mapping[str, Mapping[str, Any]] +) -> tuple[Variant, ...]: + """The shape family at the op's default attributes, plus each combination at `wide`. + + An attribute changes what a kernel computes per element, never how it walks its + operands, so crossing the two would only repeat the shape family. + """ + return base + tuple( + Variant(label, ((4, 8),), attributes) + for label, attributes in sorted(combinations.items()) + ) + + +# Attention. `B` batch items, `H` query heads over `KV` key/value heads, `Q` query positions +# against `KV_LEN` incoming key positions, at head size `D` and value head size `DV`. The 3-D +# layout packs the heads into the last axis and names their count in an attribute, which is +# what `_attention_3d` builds; the two layouts are the same node at two sets of strides, so +# every family below is swept at whichever one is not redundant for it. +_ATTENTION_BATCH = 2 +_ATTENTION_HEADS = 2 +_ATTENTION_Q = 2 +_ATTENTION_KV = 3 +_ATTENTION_SIZE = 4 + +_ATTENTION_3D_HEADS: Mapping[str, Any] = { + "q_num_heads": _ATTENTION_HEADS, + "kv_num_heads": _ATTENTION_HEADS, +} + + +def _attention_4d( + *, + batch: int = _ATTENTION_BATCH, + q_heads: int = _ATTENTION_HEADS, + kv_heads: int = _ATTENTION_HEADS, + q_seq: int = _ATTENTION_Q, + kv_seq: int = _ATTENTION_KV, + size: int = _ATTENTION_SIZE, + value_size: int | None = None, +) -> tuple[tuple[int, ...], ...]: + """`Q, K, V` at the `(batch, head, sequence, size)` layout.""" + return ( + (batch, q_heads, q_seq, size), + (batch, kv_heads, kv_seq, size), + (batch, kv_heads, kv_seq, value_size if value_size is not None else size), + ) + + +def _attention_3d( + *, + batch: int = _ATTENTION_BATCH, + q_heads: int = _ATTENTION_HEADS, + kv_heads: int = _ATTENTION_HEADS, + q_seq: int = _ATTENTION_Q, + kv_seq: int = _ATTENTION_KV, + size: int = _ATTENTION_SIZE, + value_size: int | None = None, +) -> tuple[tuple[int, ...], ...]: + """The same node at the `(batch, sequence, head * size)` layout.""" + return ( + (batch, q_seq, q_heads * size), + (batch, kv_seq, kv_heads * size), + (batch, kv_seq, kv_heads * (value_size if value_size is not None else size)), + ) + + +# The cache a node that carries one reads, and the total key length it leaves. +_ATTENTION_PAST = (_ATTENTION_BATCH, _ATTENTION_HEADS, 2, _ATTENTION_SIZE) +_ATTENTION_TOTAL = _ATTENTION_PAST[2] + _ATTENTION_KV + +# A mask row that is entirely -inf, which max-subtracting leaves as `-inf - -inf`: the +# reference's softmax reports the whole row as NaN, and so must the kernel. The other row +# mixes a finite bias with a single masked-out column. +_ATTENTION_MASK_VALUES = (0.5, -math.inf, 0.0, -math.inf, -math.inf, -math.inf) + +_ATTENTION_SHAPE_VARIANTS = ( + Variant("wide", _attention_4d()), + Variant("value_head_size", _attention_4d(value_size=6)), + Variant("grouped_query", _attention_4d(q_heads=4)), + Variant("multi_query", _attention_4d(q_heads=4, kv_heads=1)), + Variant("single", _attention_4d(batch=1, q_heads=1, kv_heads=1, q_seq=1, kv_seq=1)), + Variant("empty_batch", _attention_4d(batch=0)), + Variant("empty_head_size", _attention_4d(size=0, value_size=4)), + Variant("empty_query", _attention_4d(q_seq=0)), + Variant("wide_3d", _attention_3d(), _ATTENTION_3D_HEADS), + Variant("value_head_size_3d", _attention_3d(value_size=6), _ATTENTION_3D_HEADS), + Variant( + "grouped_query_3d", + _attention_3d(q_heads=4), + {**_ATTENTION_3D_HEADS, "q_num_heads": 4}, + ), +) + +# `scale` scales the product, and the reference applies its square root to each operand; the +# negative one is the edge numpy answers with a NaN rather than an error. `softcap` is only +# applied when it is positive, which is what the ignored one pins. +_ATTENTION_ATTRIBUTE_VARIANTS = ( + Variant("causal", _attention_4d(), {"is_causal": 1}), + Variant("causal_3d", _attention_3d(), {**_ATTENTION_3D_HEADS, "is_causal": 1}), + Variant("scaled", _attention_4d(), {"scale": 0.25}), + Variant("scale_negative", _attention_4d(), {"scale": -1.0}), + Variant("softcap", _attention_4d(), {"softcap": 2.0}), + Variant("softcap_ignored", _attention_4d(), {"softcap": -1.0}), + Variant("softcap_causal", _attention_4d(), {"softcap": 0.5, "is_causal": 1}), + # The softmax runs in float64 unless this narrows it, whatever the tensors hold. + Variant( + "softmax_precision_float", + _attention_4d(), + {"softmax_precision": TensorProto.FLOAT}, + ), + Variant( + "softmax_precision_double", + _attention_4d(), + {"softmax_precision": TensorProto.DOUBLE}, + ), +) + +# The mask families. A mask shorter than the key axis is padded out with -inf rather than +# broadcast, and every axis before that one broadcasts; a boolean mask is a different +# expression again, and a different one under `is_causal` than without it. +_ATTENTION_MASK_VARIANTS = ( + Variant("mask_2d", (*_attention_4d(), (_ATTENTION_Q, _ATTENTION_KV))), + Variant("mask_padded", (*_attention_4d(), (_ATTENTION_Q, 2))), + Variant( + "mask_3d", (*_attention_4d(), (_ATTENTION_HEADS, _ATTENTION_Q, _ATTENTION_KV)) + ), + Variant( + "mask_4d", + ( + *_attention_4d(), + (_ATTENTION_BATCH, _ATTENTION_HEADS, _ATTENTION_Q, _ATTENTION_KV), + ), + ), + Variant("mask_stretched", (*_attention_4d(), (1, 1, 1, _ATTENTION_KV))), + Variant( + "mask_infinite", + (*_attention_4d(), (_ATTENTION_Q, _ATTENTION_KV)), + values={3: _ATTENTION_MASK_VALUES}, + ), + Variant( + "mask_causal", + (*_attention_4d(), (_ATTENTION_Q, _ATTENTION_KV)), + {"is_causal": 1}, + ), + Variant( + "mask_bool", + (*_attention_4d(), (_ATTENTION_Q, _ATTENTION_KV)), + operand_types={3: TensorProto.BOOL}, + ), + Variant( + "mask_bool_4d", + ( + *_attention_4d(), + (_ATTENTION_BATCH, _ATTENTION_HEADS, _ATTENTION_Q, _ATTENTION_KV), + ), + operand_types={3: TensorProto.BOOL}, + ), + # A True entry becomes `0 * -inf` here, which is the NaN that poisons its whole row. + Variant( + "mask_bool_causal", + (*_attention_4d(), (_ATTENTION_Q, _ATTENTION_KV)), + {"is_causal": 1}, + operand_types={3: TensorProto.BOOL}, + ), + # The reference adds the causal triangle *into* the mask before broadcasting it, so a + # mask carrying one row on the query axis — the padding mask a decoder passes alongside + # `is_causal` — takes the triangle's first row for every query position rather than one + # row each. Every rank that can carry a singleton there is swept, since which axis the + # 1 sits on is what the addressing has to read it off. + Variant( + "mask_causal_one_row_2d", + (*_attention_4d(), (1, _ATTENTION_KV)), + {"is_causal": 1}, + ), + Variant( + "mask_causal_one_row_3d", + (*_attention_4d(), (1, 1, _ATTENTION_KV)), + {"is_causal": 1}, + ), + Variant( + "mask_causal_one_row_4d", + (*_attention_4d(), (_ATTENTION_BATCH, 1, 1, _ATTENTION_KV)), + {"is_causal": 1}, + ), + Variant( + "mask_causal_one_row_per_head", + (*_attention_4d(), (_ATTENTION_BATCH, _ATTENTION_HEADS, 1, _ATTENTION_KV)), + {"is_causal": 1}, + ), + Variant( + "mask_bool_causal_one_row", + (*_attention_4d(), (_ATTENTION_BATCH, 1, 1, _ATTENTION_KV)), + {"is_causal": 1}, + operand_types={3: TensorProto.BOOL}, + ), + # The same singleton against a cache, where the triangle's columns start at `past_seq`. + Variant( + "mask_causal_one_row_past", + ( + *_attention_4d(), + (_ATTENTION_BATCH, 1, 1, _ATTENTION_TOTAL), + _ATTENTION_PAST, + _ATTENTION_PAST, + ), + {"is_causal": 1}, + ), + Variant( + "mask_3d_layout", + (*_attention_3d(), (_ATTENTION_Q, _ATTENTION_KV)), + _ATTENTION_3D_HEADS, + ), +) + +# The cache families. `present_key`/`present_value` are the concatenation of the cache and +# the incoming keys or values, which is what the extra outputs carry; `is_causal` counts its +# diagonal from the cached length rather than from the start of the row. Every variant asking +# for them carries a cache, and one with incoming keys: ONNX's own shape inference reports +# nothing at all for the two present outputs otherwise — no `past_key`, or an empty incoming +# `K` — and a node whose result has no inferable shape is one the compiler refuses by design. +_ATTENTION_CACHE_VARIANTS = ( + Variant("past", (*_attention_4d(), None, _ATTENTION_PAST, _ATTENTION_PAST)), + Variant( + "past_present", + (*_attention_4d(), None, _ATTENTION_PAST, _ATTENTION_PAST), + outputs=3, + ), + Variant( + "past_present_causal", + (*_attention_4d(), None, _ATTENTION_PAST, _ATTENTION_PAST), + {"is_causal": 1}, + outputs=3, + ), + Variant( + "past_present_mask", + ( + *_attention_4d(), + (_ATTENTION_Q, _ATTENTION_TOTAL), + _ATTENTION_PAST, + _ATTENTION_PAST, + ), + outputs=3, + ), + Variant( + "past_present_3d", + (*_attention_3d(), None, _ATTENTION_PAST, _ATTENTION_PAST), + _ATTENTION_3D_HEADS, + outputs=3, + ), + # Nothing incoming: the whole of the keys and values is what the cache already holds. + Variant( + "past_only", + (*_attention_4d(kv_seq=0), None, _ATTENTION_PAST, _ATTENTION_PAST), + ), +) + +# The fourth output, and the four intermediates it may carry, each against a mask shorter +# than the key axis so that the padding is in the reported tensor too. Modes 0 and 1 are +# swept with a softcap because that is where ONNX's text distinguishes them and the +# reference does not. +_ATTENTION_REPORTED_VARIANTS = tuple( + Variant( + f"qk_mode_{mode}{label}", + ( + *_attention_4d(), + (_ATTENTION_Q, _ATTENTION_KV), + _ATTENTION_PAST, + _ATTENTION_PAST, + ), + {"qk_matmul_output_mode": mode, **attributes}, + outputs=4, + ) + for mode in (0, 1, 2, 3) + for label, attributes in (("", {}), ("_softcap", {"softcap": 2.0})) +) + +# `nonpad_kv_seqlen` masks off every key position at or past a batch item's own length. The +# reference adds that mask into an attention bias it has already reshaped to rank 4, and the +# addition is in place, so it can only evaluate a node whose bias carries the batch axis +# already — one batch item, or a 4-D mask. Those are the two forms swept. +_ATTENTION_NONPAD_VARIANTS = ( + Variant( + "nonpad_single_batch", + (*_attention_4d(batch=1), None, None, None, (1,)), + values={6: (2,)}, + ), + Variant( + "nonpad_mask_4d", + ( + *_attention_4d(), + (_ATTENTION_BATCH, _ATTENTION_HEADS, _ATTENTION_Q, _ATTENTION_KV), + None, + None, + (_ATTENTION_BATCH,), + ), + values={6: (1, 3)}, + ), + Variant( + "nonpad_causal", + ( + *_attention_4d(batch=1), + (1, _ATTENTION_HEADS, _ATTENTION_Q, _ATTENTION_KV), + None, + None, + (1,), + ), + {"is_causal": 1}, + values={6: (2,)}, + ), +) + +_ATTENTION_VARIANTS = ( + *_ATTENTION_SHAPE_VARIANTS, + *_ATTENTION_ATTRIBUTE_VARIANTS, + *_ATTENTION_MASK_VARIANTS, + *_ATTENTION_CACHE_VARIANTS, + *_ATTENTION_REPORTED_VARIANTS, + *_ATTENTION_NONPAD_VARIANTS, +) + +# RotaryEmbedding. `position_ids` indexes the caches, so it is pinned rather than drawn: a +# seeded draw would name rows that are not there, which ONNX leaves undefined. The negative +# one is numpy's own indexing, which the reference gathers with. +_ROTARY_POSITIONS = (0, 2, 1, 3, 7, 4) +_ROTARY_NEGATIVE_POSITIONS = (-1, -8, 0, 5, -3, 2) +# Wide enough that a cache of one angle per position holds the whole special-value list. +_ROTARY_ROWS = 8 + +_ROTARY_VARIANTS = ( + Variant( + "wide", + ((2, 2, 3, 4), (_ROTARY_ROWS, 2), (_ROTARY_ROWS, 2), (2, 3)), + values={3: _ROTARY_POSITIONS}, + ), + Variant( + "interleaved", + ((2, 2, 3, 4), (_ROTARY_ROWS, 2), (_ROTARY_ROWS, 2), (2, 3)), + {"interleaved": 1}, + values={3: _ROTARY_POSITIONS}, + ), + Variant( + "negative_positions", + ((2, 2, 3, 4), (_ROTARY_ROWS, 2), (_ROTARY_ROWS, 2), (2, 3)), + values={3: _ROTARY_NEGATIVE_POSITIONS}, + ), + # A partial rotation: the lanes past `rotary_embedding_dim` are copied through untouched. + Variant( + "rotary_dim", + ((2, 2, 3, 4), (_ROTARY_ROWS, 1), (_ROTARY_ROWS, 1), (2, 3)), + {"rotary_embedding_dim": 2}, + values={3: _ROTARY_POSITIONS}, + ), + Variant( + "rotary_dim_interleaved", + ((2, 2, 3, 4), (_ROTARY_ROWS, 1), (_ROTARY_ROWS, 1), (2, 3)), + {"rotary_embedding_dim": 2, "interleaved": 1}, + values={3: _ROTARY_POSITIONS}, + ), + Variant( + "hidden_axis", + ((2, 3, 8), (_ROTARY_ROWS, 2), (_ROTARY_ROWS, 2), (2, 3)), + {"num_heads": 2}, + values={3: _ROTARY_POSITIONS}, + ), + Variant( + "hidden_axis_interleaved", + ((2, 3, 8), (_ROTARY_ROWS, 2), (_ROTARY_ROWS, 2), (2, 3)), + {"num_heads": 2, "interleaved": 1}, + values={3: _ROTARY_POSITIONS}, + ), + Variant( + "hidden_axis_rotary_dim", + ((2, 3, 8), (_ROTARY_ROWS, 1), (_ROTARY_ROWS, 1), (2, 3)), + {"num_heads": 2, "rotary_embedding_dim": 2}, + values={3: _ROTARY_POSITIONS}, + ), + Variant( + "empty_batch", + ((0, 2, 3, 4), (_ROTARY_ROWS, 2), (_ROTARY_ROWS, 2), (0, 3)), + values={3: ()}, + ), + Variant( + "one_position", + ((1, 1, 1, 2), (1, 1), (1, 1), (1, 1)), + values={3: (0,)}, + ), + # Without `position_ids` the caches stand as they are, carrying the batch and the + # position themselves and stretching over the heads. + Variant("cached", ((2, 2, 3, 4), (2, 3, 2), (2, 3, 2))), + Variant( + "cached_interleaved", ((2, 2, 3, 4), (2, 3, 2), (2, 3, 2)), {"interleaved": 1} + ), + Variant( + "cached_rotary_dim", + ((2, 2, 3, 4), (2, 3, 1), (2, 3, 1)), + {"rotary_embedding_dim": 2}, + ), + Variant("cached_stretched", ((2, 2, 3, 4), (1, 3, 2), (1, 3, 2))), + Variant("cached_hidden_axis", ((2, 3, 8), (2, 3, 2), (2, 3, 2)), {"num_heads": 2}), +) + + +# Ops with no attributes and one operand, whose whole surface is the shape family. +_POINTWISE_UNARY_OPS = ( + "Acos", + "Acosh", + "Asin", + "Asinh", + "Atan", + "Atanh", + "Ceil", + "Cos", + "Cosh", + "Erf", + "Exp", + "Floor", + "Identity", + "Log", + "Reciprocal", + "Relu", + "Round", + "Sigmoid", + "Sign", + "Sin", + "Sinh", + "Softplus", + "Softsign", + "Sqrt", + "Tan", + "Tanh", +) + +# The zero-valued alpha the reference's `alpha or self.alpha` would silently replace with +# the schema default is deliberately absent from every combination below; a case built on it +# would be comparing against a value the reference does not claim to compute. +_ACTIVATION_ATTRIBUTES: dict[str, Mapping[str, Mapping[str, Any]]] = { + "Celu": {"alpha_half": {"alpha": 0.5}, "alpha_two": {"alpha": 2.0}}, + "Elu": {"alpha_quarter": {"alpha": 0.25}, "alpha_two": {"alpha": 2.0}}, + # `approximate` selects between two formulas rather than scaling one, so both are swept + # even though `none` is the default the shape family already runs. + "Gelu": {"erf": {"approximate": "none"}, "tanh": {"approximate": "tanh"}}, + "HardSigmoid": { + "shifted": {"alpha": 0.5, "beta": 0.25}, + "steep": {"alpha": 2.0, "beta": -1.0}, + }, + "LeakyRelu": {"steep": {"alpha": 0.5}, "negative": {"alpha": -0.25}}, + "Selu": { + "unit": {"alpha": 1.0, "gamma": 1.0}, + "scaled": {"alpha": 0.5, "gamma": 2.0}, + }, + "ThresholdedRelu": {"high": {"alpha": 2.0}, "negative": {"alpha": -1.0}}, +} + +# Shrink's bias and lambd stay non-negative: with either negative, an integer result can +# leave the dtype's range, where the reference's own float64-to-integer cast is undefined. +_SHRINK_ATTRIBUTES: Mapping[str, Mapping[str, Any]] = { + "hard": {"lambd": 1.5, "bias": 0.0}, + "soft": {"lambd": 1.5, "bias": 1.5}, + "wide_band": {"lambd": 3.0, "bias": 0.5}, +} + + +def _linear_attention_variant( + label: str, + *, + rule: str | None = None, + batch: int = 2, + steps: int = 3, + q_heads: int = 2, + kv_heads: int = 2, + d_k: int = 3, + d_v: int = 2, + past: bool = False, + decay: str | None = None, + beta: str | None = None, + **attributes: Any, +) -> Variant: + """One LinearAttention case, from the head counts and head widths it runs at. + + Every 3-D operand packs `H * D` into its last axis, so the shapes follow from those four + numbers. `decay` and `beta` name the granularity ONNX packs each of them at -- one value + per key dimension or one per head for the decay, one per head or one the heads share for + beta -- and leaving either out is how `update_rule` reaches the model: the rule forbids + the operand it does not read, so which operands a node passes is the rule. + """ + decay_width = kv_heads * (1 if decay == "per_head" else d_k) + beta_width = kv_heads if beta == "per_head" else 1 + shapes: list[tuple[int, ...] | None] = [ + (batch, steps, q_heads * d_k), + (batch, steps, kv_heads * d_k), + (batch, steps, kv_heads * d_v), + (batch, kv_heads, d_k, d_v) if past else None, + None if decay is None else (batch, steps, decay_width), + None if beta is None else (batch, steps, beta_width), + ] + while shapes and shapes[-1] is None: + shapes.pop() + return Variant( + label, + tuple(shapes), + { + "q_num_heads": q_heads, + "kv_num_heads": kv_heads, + **({} if rule is None else {"update_rule": rule}), + **attributes, + }, + outputs=2, + ) + + +_LINEAR_ATTENTION_VARIANTS = ( + # The four recurrences, each with exactly the gates it reads. + _linear_attention_variant("linear", rule="linear"), + _linear_attention_variant("gated", rule="gated", decay="per_key_dim"), + _linear_attention_variant("gated_per_head_decay", rule="gated", decay="per_head"), + _linear_attention_variant("delta", rule="delta", beta="per_head"), + _linear_attention_variant("delta_shared_beta", rule="delta", beta="shared"), + _linear_attention_variant( + "gated_delta", rule="gated_delta", decay="per_key_dim", beta="per_head" + ), + # The same rule, left to the schema's own default rather than named. + _linear_attention_variant("default_rule", decay="per_key_dim", beta="per_head"), + _linear_attention_variant( + "gated_delta_per_head_decay", decay="per_head", beta="shared" + ), + # `past_state` seeds the state the sequence starts from, which each rule carries forward + # its own way. + _linear_attention_variant("linear_with_past", rule="linear", past=True), + _linear_attention_variant( + "gated_with_past", rule="gated", decay="per_key_dim", past=True + ), + _linear_attention_variant( + "delta_with_past", rule="delta", beta="per_head", past=True + ), + _linear_attention_variant( + "gated_delta_with_past", decay="per_key_dim", beta="per_head", past=True + ), + # Grouped-query attention: one KV head, and the one state it carries, answers several + # query heads. A `kv_num_heads` of 1 is the multi-query special case. + _linear_attention_variant( + "grouped_query", q_heads=4, decay="per_key_dim", beta="per_head" + ), + _linear_attention_variant( + "multi_query", + q_heads=4, + kv_heads=1, + decay="per_key_dim", + beta="per_head", + past=True, + ), + # A `scale` of 0 -- the schema's default, which every variant above runs -- asks for + # `1/sqrt(d_k)`; anything else is taken as given. + _linear_attention_variant( + "explicit_scale", decay="per_key_dim", beta="per_head", scale=0.25 + ), + _linear_attention_variant( + "negative_scale", decay="per_key_dim", beta="per_head", scale=-1.5 + ), + # `chunk_size` is documented as a tuning hint that does not affect the output, which is + # only worth asserting where one chunk would not span the whole sequence anyway. + _linear_attention_variant( + "chunk_size_hint", decay="per_key_dim", beta="per_head", chunk_size=2 + ), + # One decode step, which is what the state inputs exist for, and the two head widths the + # other way round. + _linear_attention_variant( + "decode_step", steps=1, decay="per_key_dim", beta="per_head", past=True + ), + _linear_attention_variant( + "wide_values", d_k=2, d_v=5, decay="per_key_dim", beta="per_head" + ), + # A key width of one packs the decay's two granularities identically; ONNX reads the + # per-head one first, and here the two mean the same thing. + _linear_attention_variant( + "scalar_head", + q_heads=1, + kv_heads=1, + d_k=1, + d_v=1, + decay="per_key_dim", + beta="per_head", + ), + # The zero-element shapes the op admits. A key width of zero is the sharp one: the + # derived scale is `1/sqrt(0)`, so every answer is that infinity times an empty sum -- a + # NaN in the reference and here alike -- while an explicit scale leaves it a zero. + _linear_attention_variant( + "empty_batch", batch=0, decay="per_key_dim", beta="per_head", past=True + ), + _linear_attention_variant( + "empty_sequence", steps=0, decay="per_key_dim", beta="per_head", past=True + ), + _linear_attention_variant( + "empty_value_width", d_v=0, decay="per_key_dim", beta="per_head", past=True + ), + _linear_attention_variant( + "empty_key_width", d_k=0, decay="per_key_dim", beta="per_head", past=True + ), + _linear_attention_variant( + "empty_key_width_scaled", + d_k=0, + decay="per_key_dim", + beta="per_head", + scale=0.5, + ), +) + +# -------------------------------------------------------------------------------------- +# Counting n-grams +# -------------------------------------------------------------------------------------- + +# TfIdfVectorizer matches its input against a pool of n-grams its attributes carry, so what a +# case covers is decided by the pool and the tokens together: a draw over the whole of an +# integer dtype matches nothing at all and would compare one zero tensor against another. Most +# variants therefore pin the token sequence to values built out of the pool, and the widest +# draws freely -- which is what carries the dtype's extremes into an operand every value of +# which is compared against the pool rather than computed with. +# +# Each pool below splits differently by n-gram length: `ngram_counts` says where the entries +# of each length start, and the identifiers ONNX numbers them by run across all of them, which +# is what `ngram_indexes` is indexed by. `level_empty` is a pool whose unigram level holds +# nothing, `repeated` one listing the same n-gram twice -- the case ONNX defines as taking the +# last of the two -- and `unreached` one holding lengths outside the gram range every variant +# asks for, whose identifiers still have to advance past them. +_TFIDF_POOLS: Mapping[str, Mapping[str, Any]] = { + "unigrams": { + "ngram_counts": [0], + "ngram_indexes": [0, 1, 2, 3], + "pool_int64s": [2, 3, 5, 4], + }, + "level_empty": { + "ngram_counts": [0, 0], + "ngram_indexes": [0, 1, 2], + "pool_int64s": [5, 6, 7, 8, 6, 7], + }, + "mixed": { + "ngram_counts": [0, 4], + "ngram_indexes": [0, 1, 2, 3, 4, 5, 6], + "pool_int64s": [2, 3, 5, 4, 5, 6, 7, 8, 6, 7], + }, + "trigrams": { + "ngram_counts": [0, 2, 6], + "ngram_indexes": [0, 1, 2, 3, 4, 5], + "pool_int64s": [2, 3, 5, 4, 5, 6, 2, 3, 4, 4, 5, 6], + }, + "repeated": { + "ngram_counts": [0], + "ngram_indexes": [2, 0, 1], + "pool_int64s": [3, 3, 5], + }, + "unreached": { + "ngram_counts": [0, 1, 3], + "ngram_indexes": [0, 1, 2, 3], + "pool_int64s": [9, 5, 6, 2, 3, 5], + }, +} + +# A token sequence holding several of the pools' n-grams, adjacent and spread apart, so that +# both the skip distances and the gram lengths have something to find. +_TFIDF_TOKENS = (2, 3, 5, 4, 5, 6, 7, 8, 6, 7, 2, 3) + +_TFIDF_WEIGHTS = (0.5, 1.5, -2.0, 0.25, 3.0, 1.0, 0.125) + +_TFIDF_GRAM_RANGES = ((1, 1), (1, 2), (2, 2), (1, 3), (2, 3), (3, 3)) + + +def _tfidf_variant( + label: str, + pool: str, + *, + shape: tuple[int, ...] = (12,), + tokens: Sequence[int] | None = _TFIDF_TOKENS, + mode: str = "TF", + grams: tuple[int, int] = (1, 2), + skip: int = 1, + weighted: bool = False, +) -> Variant: + attributes = { + **_TFIDF_POOLS[pool], + "mode": mode, + "min_gram_length": grams[0], + "max_gram_length": grams[1], + "max_skip_count": skip, + } + if weighted: + width = max(attributes["ngram_indexes"]) + 1 + attributes["weights"] = list(_TFIDF_WEIGHTS[:width]) + values = ( + {} + if tokens is None + else { + 0: tuple(tokens[index % len(tokens)] for index in range(math.prod(shape))) + } + ) + return Variant(f"{label}_{pool}", (shape,), attributes, values=values) + + +def _with_unit_weights(case: Case) -> ModelProto: + """The case's node, with the weights it leaves out set to the identity of their product. + + The reference reads an absent `weights` as `None` rather than as an empty list, and every + mode that reaches for it -- `IDF`, which multiplies a truncated count by the weight of its + n-gram, and `TFIDF`, which multiplies the count itself -- raises before reaching the + branch its own code carries for a node that sets none. Running the oracle with those + weights set to 1 is that multiplication left as it was, and is exactly what the branch the + raise blocks computes. + """ + model = _model(case) + node = model.graph.node[0] + attributes = {entry.name for entry in node.attribute} + if "weights" in attributes or case.variant.attributes["mode"] == "TF": + return model + width = max(case.variant.attributes["ngram_indexes"]) + 1 + node.attribute.append(helper.make_attribute("weights", [1.0] * width)) + return model + + +_TFIDF_VARIANTS = ( + tuple( + _tfidf_variant( + f"{mode.lower()}{'_weighted' if weighted else ''}", + pool, + mode=mode, + weighted=weighted, + ) + for pool in _TFIDF_POOLS + for mode in ("TF", "IDF", "TFIDF") + for weighted in (False, True) + ) + + tuple( + _tfidf_variant( + f"grams_{low}_{high}_skip_{skip}", pool, grams=(low, high), skip=skip + ) + for pool in ("mixed", "trigrams") + for low, high in _TFIDF_GRAM_RANGES + for skip in (0, 1, 5) + ) + + ( + _tfidf_variant("batch", "mixed", shape=(2, 6)), + _tfidf_variant("batch_skipped", "mixed", shape=(3, 5), skip=3, grams=(2, 3)), + _tfidf_variant("single_row", "mixed", shape=(1, 12)), + _tfidf_variant("single_token", "unigrams", shape=(1,)), + # The zero-element shapes the reference reads: a sequence of no tokens, and a batch + # of them. A batch of *no sequences* is not one of them — it refuses any `[N, C]` + # with `N` below 1 outright — so that shape has no oracle here. + _tfidf_variant("empty_sequence", "mixed", shape=(0,)), + _tfidf_variant("empty_rows", "mixed", shape=(2, 0)), + # The one variant whose tokens are drawn rather than pinned, and the only one wide + # enough to carry the whole special-value list into the operand. + _tfidf_variant("wide", "mixed", shape=(4, 8), tokens=None), + ) +) + + +# -------------------------------------------------------------------------------------- +# ONNX-ML preprocessing +# -------------------------------------------------------------------------------------- + +# The standard-domain opset the helper nodes an ONNX-ML oracle is built from are imported at; +# any revision serves, since none of them changed what a cast or a softmax computes here. +_ML_STANDARD_OPSET = 21 + + +def _with_float_inputs(case: Case) -> ModelProto: + """The case's node with every operand cast to float32 in front of it. + + `Scaler`, `Normalizer`, `FeatureVectorizer` and the four predictors below are each + declared by their own schema to produce a `tensor(float)` whatever they are handed, while + the reference evaluator returns a result of the *input's* element type — so for anything + but a float32 input its output contradicts the op's stated output type, and comparing + against it would be comparing against the wrong type. The oracle is therefore run on the + model whose input already is that float, where the two agree; for a float32 case the cast + is an identity and this is the compiled model itself. + """ + return _cast_inputs_to_float(_model(case)) + + +def _cast_inputs_to_float(model: ModelProto) -> ModelProto: + """`model` with a Cast to float32 in front of every operand it is fed.""" + graph = model.graph + fed = [entry.name for entry in graph.input] + casts = [ + helper.make_node("Cast", [name], [f"{name}_float"], to=TensorProto.FLOAT) + for name in fed + ] + for node in graph.node: + for index, name in enumerate(node.input): + if name in fed: + node.input[index] = f"{name}_float" + nodes = casts + list(graph.node) + del graph.node[:] + graph.node.extend(nodes) + return _importing_standard_ops(model) + + +def _importing_standard_ops(model: ModelProto) -> ModelProto: + """`model` importing the standard domain, which its helper nodes are defined in.""" + if not any(entry.domain in ("", "ai.onnx") for entry in model.opset_import): + model.opset_import.append(helper.make_opsetid("", _ML_STANDARD_OPSET)) + return model + + +# Scaler addresses its coefficients along the last axis, which is what the shapes sweep: one +# coefficient per feature and a single one shared by all, at each rank. A rank-0 input is +# deliberately absent — the reference broadcasts the scalar against the coefficient list and +# returns the rank-1 result of that, which is not the shape the op's own schema gives it. +_SCALER_COEFFICIENTS = ( + ("per_feature", [1.0, -2.0, 0.5], [0.5, 2.0, -1.0]), + ("shared", [1.5], [-0.25]), + ("shared_offset", [0.0], [2.0, -3.0, 0.5]), + ("shared_scale", [1.0, -1.0, 0.0], [4.0]), +) +_SCALER_VARIANTS = tuple( + Variant(f"{label}_{shape_label}", (shape,), {"offset": offset, "scale": scale}) + for label, offset, scale in _SCALER_COEFFICIENTS + for shape_label, shape in ( + ("matrix", (4, 3)), + ("vector", (3,)), + ("rank_3", (2, 2, 3)), + ) +) + ( + Variant("wide", ((4, 8),), {"offset": [1.0] * 8, "scale": [0.5] * 8}), + Variant("empty_rows", ((0, 3),), {"offset": [1.0, 2.0, 3.0], "scale": [1.0] * 3}), + Variant("empty_features", ((2, 0),), {"offset": [1.0], "scale": [2.0]}), +) + +# Normalizer divides each row by a norm of that row, so every shape is the `[N,C]` matrix its +# schema describes — the reference reduces along axis 1 and has no other reading of the +# input, and it reshapes that reduction in a way no zero-element matrix survives, so those +# shapes have no oracle here. The pinned rows are the edges a draw would not reach: a row +# whose norm is zero, and hence divided by the floor the reference falls back on rather than +# by nothing; a row carrying a NaN, which its `max` and its sum both propagate; and one +# carrying an infinity, which leaves every element of the row at NaN or zero. +_NORMALIZER_EDGE_ROWS = ( + 0.0, + -0.0, + 0.0, + float("nan"), + 1.0, + 2.0, + float("inf"), + 1.0, + 2.0, + -3.0, + 0.5, + -0.25, +) +_NORMALIZER_VARIANTS = ( + tuple( + Variant(f"{label}_{norm.lower()}", (shape,), {"norm": norm}) + for norm in ("MAX", "L1", "L2") + for label, shape in ( + ("matrix", (4, 3)), + ("wide", (4, 8)), + ("single_row", (1, 5)), + ("single_column", (4, 1)), + ) + ) + + tuple( + Variant( + f"edge_rows_{norm.lower()}", + ((4, 3),), + {"norm": norm}, + values={0: _NORMALIZER_EDGE_ROWS}, + elem_types=_FLOAT_ELEM_TYPES, + ) + for norm in ("MAX", "L1", "L2") + ) + + ( + # The default `norm`, which the schema declares rather than the node. + Variant("default_norm", ((4, 3),)), + ) +) + +# Imputer's two attribute families are the floating-point and the integer marker; ONNX pairs +# each with the element types it describes, and the reference reads the float one first +# whenever it is set. Its shapes are the matrices the reference accepts — it refuses any other +# rank outright — with a value per column or a single one shared by all of them. +_INTEGER_ML_TYPES = (TensorProto.INT32, TensorProto.INT64) +_IMPUTER_VARIANTS = ( + Variant( + "per_column", + ((4, 3),), + {"imputed_value_floats": [1.5, -2.5, 0.0], "replaced_value_float": 0.0}, + elem_types=_FLOAT_ELEM_TYPES, + ), + Variant( + "shared_value", + ((4, 3),), + {"imputed_value_floats": [9.0], "replaced_value_float": 1.0}, + elem_types=_FLOAT_ELEM_TYPES, + ), + Variant( + "nan_marker", + ((4, 3),), + {"imputed_value_floats": [7.0], "replaced_value_float": float("nan")}, + elem_types=_FLOAT_ELEM_TYPES, + ), + Variant( + "wide", + ((4, 8),), + {"imputed_value_floats": [0.25] * 8, "replaced_value_float": -1.0}, + elem_types=_FLOAT_ELEM_TYPES, + ), + Variant( + "empty_rows", + ((0, 3),), + {"imputed_value_floats": [1.0], "replaced_value_float": 0.0}, + elem_types=_FLOAT_ELEM_TYPES, + ), + Variant( + "integer_per_column", + ((4, 3),), + {"imputed_value_int64s": [7, -8, 9], "replaced_value_int64": 0}, + elem_types=_INTEGER_ML_TYPES, + ), + Variant( + "integer_shared", + ((4, 3),), + {"imputed_value_int64s": [5], "replaced_value_int64": 1}, + elem_types=_INTEGER_ML_TYPES, + ), +) + +# Binarizer is a comparison against a threshold, so its shapes are the unary family and its +# thresholds include the schema's own default, which the node leaves out. +_BINARIZER_VARIANTS = tuple( + replace( + variant, + label=f"{variant.label}_{label}", + attributes={} if threshold is None else {"threshold": threshold}, + ) + for label, threshold in ( + ("default", None), + ("zero", 0.0), + ("positive", 1.5), + ("negative", -2.0), + ) + for variant in _UNARY_VARIANTS +) + +# OneHotEncoder's categories are the integer list, which is what a numeric input is matched +# against; the reference reads rank 1 and rank 2 and refuses anything deeper. `zeros` cleared +# makes a value in no category the failure the schema prescribes, so those cases pin every +# element to a category rather than drawing one. +# +# A category list with a repeat is deliberately absent: the reference sizes its result from +# the *distinct* categories while ONNX's own shape inference sizes it from the list, so the +# reference indexes past its own buffer and there is nothing to compare a kernel against. +_ENCODER_CATEGORIES = [0, 1, -1] +_IN_CATEGORY = tuple(_ENCODER_CATEGORIES[index % 3] for index in range(12)) +_ONE_HOT_ENCODER_VARIANTS = ( + Variant("matrix", ((4, 3),), {"cats_int64s": _ENCODER_CATEGORIES, "zeros": 1}), + Variant("wide", ((4, 4),), {"cats_int64s": _ENCODER_CATEGORIES, "zeros": 1}), + Variant("vector", ((5,),), {"cats_int64s": _ENCODER_CATEGORIES, "zeros": 1}), + Variant("single_category", ((4, 3),), {"cats_int64s": [1], "zeros": 1}), + Variant("empty_rows", ((0, 3),), {"cats_int64s": _ENCODER_CATEGORIES, "zeros": 1}), + Variant( + "strict", + ((4, 3),), + {"cats_int64s": _ENCODER_CATEGORIES, "zeros": 0}, + values={0: _IN_CATEGORY}, + ), +) + +# LabelEncoder pairs a key family with a value family. ONNX's own type inference rejects a +# key element type differing from the input's, so the `keys_int64s` and `keys_floats` families +# pair with exactly one input type each and the `keys_tensor` family is what reaches the rest. +# The keys themselves are values the generator is certain to feed — the specials every dtype's +# sweep starts with — so that hits and misses both occur. +_LABEL_KEYS_INT = [0, 1, -1] +_LABEL_KEYS_FLOAT = [0.0, 1.0, -1.0] +_LABEL_INT_MAPPING = { + "keys_int64s": _LABEL_KEYS_INT, + "values_int64s": [10, 20, 30], + "default_int64": -99, +} + + +def _label_tensor_mapping(elem_type: int) -> dict[str, Any]: + """A `keys_tensor`/`values_tensor` mapping, the one family that names its own types.""" + keys = np.array(_LABEL_KEYS_INT, numpy_dtype_name(elem_type)) + return { + "keys_tensor": numpy_helper.from_array(keys, "keys"), + "values_tensor": numpy_helper.from_array( + np.array([3, -4, 5], np.int32), "values" + ), + "default_tensor": numpy_helper.from_array(np.array([-77], np.int32), "default"), + } + + +_LABEL_ENCODER_VARIANTS = ( + Variant( + "int_keys_int_values", + ((4, 3),), + _LABEL_INT_MAPPING, + elem_types=(TensorProto.INT64,), + ), + Variant( + "int_keys_float_values", + ((4, 3),), + { + "keys_int64s": _LABEL_KEYS_INT, + "values_floats": [1.5, -2.5, 0.0], + "default_float": -0.5, + }, + elem_types=(TensorProto.INT64,), + ), + Variant( + "float_keys_float_values", + ((4, 3),), + { + "keys_floats": _LABEL_KEYS_FLOAT, + "values_floats": [1.5, -2.5, 0.0], + "default_float": -0.5, + }, + elem_types=(TensorProto.FLOAT,), + ), + Variant( + "float_keys_int_values", + ((4, 3),), + { + "keys_floats": _LABEL_KEYS_FLOAT, + "values_int64s": [10, 20, 30], + "default_int64": -99, + }, + elem_types=(TensorProto.FLOAT,), + ), + # A repeated key takes its last occurrence, as the schema states in as many words. + Variant( + "repeated_key", + ((4, 3),), + { + "keys_int64s": [1, 0, 1], + "values_int64s": [10, 20, 30], + "default_int64": -99, + }, + elem_types=(TensorProto.INT64,), + ), + Variant( + "default_only", + ((4, 3),), + {"keys_int64s": [7], "values_int64s": [8], "default_int64": 42}, + elem_types=(TensorProto.INT64,), + ), + *( + Variant(label, (shape,), _LABEL_INT_MAPPING, elem_types=(TensorProto.INT64,)) + for label, shape in (("wide", (4, 8)), ("rank_0", ()), ("empty", (0, 3))) + ), + # The tensor family, which is where the element types the other two cannot describe are + # swept, and which brings its own default along. + *( + Variant( + f"tensor_pair_{numpy_dtype_name(elem_type)}", + ((4, 3),), + _label_tensor_mapping(elem_type), + elem_types=(elem_type,), + ) + for elem_type in ( + TensorProto.DOUBLE, + TensorProto.FLOAT, + TensorProto.INT16, + TensorProto.INT32, + ) + ), +) + +# ArrayFeatureExtractor takes the columns its index operand names, so the indices are pinned +# rather than drawn, for the reason the gathering ops' are: a seeded draw would spend every +# case on the out-of-range value ONNX leaves undefined. A vector `X` is the one case the +# reference documents as following onnxruntime rather than the specification, returning the +# single row a one-row matrix would have. +_EXTRACTOR_VARIANTS = ( + Variant("matrix", ((3, 4), (2,)), values={1: (3, 0)}), + Variant("vector", ((4,), (2,)), values={1: (1, 3)}), + Variant("rank_3", ((2, 3, 4), (2,)), values={1: (0, 2)}), + Variant("row_of_indices", ((3, 4), (1, 2)), values={1: (2, 1)}), + Variant("negative_index", ((3, 4), (2,)), values={1: (-1, -4)}), + Variant("repeated_index", ((3, 4), (3,)), values={1: (1, 1, 2)}), + Variant("every_column", ((3, 4), (4,)), values={1: (0, 1, 2, 3)}), + Variant("wide", ((4, 8), (2,)), values={1: (7, 0)}), + Variant("empty_rows", ((0, 4), (2,)), values={1: (0, 3)}), + Variant("no_indices", ((3, 4), (0,)), values={1: ()}), +) + +# FeatureVectorizer lays its inputs side by side, each cut or zero-padded to its declared +# width, so the sweep is every relation a width can have to the input it describes. +_VECTORIZER_VARIANTS = ( + Variant("exact", ((4, 3), (4, 2)), {"inputdimensions": [3, 2]}), + Variant("padded", ((4, 3), (4, 2)), {"inputdimensions": [5, 2]}), + Variant("truncated", ((4, 3), (4, 2)), {"inputdimensions": [2, 1]}), + Variant("single", ((4, 3),), {"inputdimensions": [3]}), + Variant("vector_input", ((4,), (4, 2)), {"inputdimensions": [1, 2]}), + Variant("vector_padded", ((4,),), {"inputdimensions": [3]}), + Variant("three_inputs", ((2, 2), (2, 1), (2, 3)), {"inputdimensions": [2, 1, 3]}), + Variant("zero_width", ((4, 3), (4, 2)), {"inputdimensions": [0, 2]}), + Variant("empty_rows", ((0, 3), (0, 2)), {"inputdimensions": [3, 2]}), + # Wide enough on every input for the dtype's whole special-value list to reach it. + Variant("wide", ((4, 4), (4, 4), (4, 4)), {"inputdimensions": [4, 4, 4]}), +) + +# -------------------------------------------------------------------------------------- +# Tree ensembles +# -------------------------------------------------------------------------------------- + +# The forests below are written as trees and flattened into the two encodings ONNX-ML uses, +# so that one description covers the legacy `(tree, node)` families and opset 5's separate +# node and leaf families -- and so that what a variant is actually testing stays readable. + + +class _Split(NamedTuple): + """One interior node: the test it applies, and the two subtrees it chooses between. + + `members` holds the set a `MEMBER` test matches against, which only opset 5 defines. A + branch is either another `_Split` or a leaf, written as a list of `(target, weight)` + pairs -- more than one of which only the legacy encoding can express. + """ + + feature: int + test: str + value: float + true_branch: Any + false_branch: Any + missing: int = 0 + members: tuple[float, ...] = () + + +# The branch tests, named without the `BRANCH_` the legacy families spell them with; opset 5 +# numbers them by their position here. +_TESTS = ("LEQ", "LT", "GTE", "GT", "EQ", "NEQ", "MEMBER") + +_LEGACY_NODE_FAMILIES = ( + "treeids", + "nodeids", + "featureids", + "modes", + "values", + "truenodeids", + "falsenodeids", + "missing_value_tracks_true", +) +_LEGACY_LEAF_FAMILIES = ("treeids", "nodeids", "ids", "weights") +_ENSEMBLE_NODE_FAMILIES = ( + "featureids", + "truenodeids", + "falsenodeids", + "trueleafs", + "falseleafs", + "missing_value_tracks_true", +) + + +def _legacy_attributes(trees: Sequence[Any], role: str, **extra: Any) -> dict[str, Any]: + """A forest in the `(tree, node)`-keyed families the legacy ensembles are encoded in.""" + nodes: dict[str, list] = {name: [] for name in _LEGACY_NODE_FAMILIES} + leaves: dict[str, list] = {name: [] for name in _LEGACY_LEAF_FAMILIES} + for tree_id, tree in enumerate(trees): + _append_legacy(tree, tree_id, nodes, leaves) + return { + **{f"nodes_{name}": values for name, values in nodes.items()}, + **{f"{role}_{name}": values for name, values in leaves.items()}, + **extra, + } + + +def _append_legacy( + node: Any, tree_id: int, nodes: dict[str, list], leaves: dict[str, list] +) -> int: + """Append `node` and its subtrees, returning the node id it was given in its tree.""" + node_id = sum(1 for value in nodes["treeids"] if value == tree_id) + position = len(nodes["treeids"]) + nodes["treeids"].append(tree_id) + nodes["nodeids"].append(node_id) + split = node if isinstance(node, _Split) else None + nodes["featureids"].append(split.feature if split else 0) + nodes["modes"].append(f"BRANCH_{split.test}" if split else "LEAF") + nodes["values"].append(split.value if split else 0.0) + nodes["missing_value_tracks_true"].append(split.missing if split else 0) + nodes["truenodeids"].append(0) + nodes["falsenodeids"].append(0) + if split is None: + for target, weight in node: + leaves["treeids"].append(tree_id) + leaves["nodeids"].append(node_id) + leaves["ids"].append(target) + leaves["weights"].append(weight) + return node_id + nodes["truenodeids"][position] = _append_legacy( + split.true_branch, tree_id, nodes, leaves + ) + nodes["falsenodeids"][position] = _append_legacy( + split.false_branch, tree_id, nodes, leaves + ) + return node_id + + +def _ensemble_attributes( + trees: Sequence[Any], dtype: str = "float32", **extra: Any +) -> dict[str, Any]: + """The same forest in opset 5's families, where leaves are indexed apart from nodes.""" + nodes: dict[str, list] = {name: [] for name in _ENSEMBLE_NODE_FAMILIES} + tests: list[int] = [] + splits: list[float] = [] + targets: list[int] = [] + weights: list[float] = [] + members: list[float] = [] + roots = [] + for tree in trees: + index, leaf = _append_ensemble( + tree, nodes, tests, splits, targets, weights, members + ) + # A tree that is a single leaf is encoded at a position in *both* families at once, + # which is what `_bare_leaf_attributes` is written out for. + assert not leaf, "a bare-leaf tree cannot be built from a leaf alone" + roots.append(index) + attributes = { + **{f"nodes_{name}": values for name, values in nodes.items()}, + "nodes_modes": numpy_helper.from_array( + np.array(tests, np.uint8), "nodes_modes" + ), + "nodes_splits": numpy_helper.from_array( + np.array(splits, dtype), "nodes_splits" + ), + "leaf_targetids": targets, + "leaf_weights": numpy_helper.from_array( + np.array(weights, dtype), "leaf_weights" + ), + "tree_roots": roots, + **extra, + } + if members: + attributes["membership_values"] = numpy_helper.from_array( + np.array(members, dtype), "membership_values" + ) + return attributes + + +def _append_ensemble( + node: Any, + nodes: dict[str, list], + tests: list[int], + splits: list[float], + targets: list[int], + weights: list[float], + members: list[float], +) -> tuple[int, int]: + """Append `node`, returning its index and whether it landed in the leaf families.""" + if not isinstance(node, _Split): + ((target, weight),) = node + targets.append(target) + weights.append(weight) + return len(targets) - 1, 1 + index = len(tests) + tests.append(_TESTS.index(node.test)) + splits.append(node.value) + nodes["featureids"].append(node.feature) + nodes["missing_value_tracks_true"].append(node.missing) + for family in ("truenodeids", "falsenodeids", "trueleafs", "falseleafs"): + nodes[family].append(0) + # The sets are read in the order the reference builds the trees in, so they are appended + # where that traversal reaches them: before either branch. + if node.test == "MEMBER": + members.extend([*node.members, float("nan")]) + for branch, child, leaf in ( + (node.true_branch, "truenodeids", "trueleafs"), + (node.false_branch, "falsenodeids", "falseleafs"), + ): + nodes[child][index], nodes[leaf][index] = _append_ensemble( + branch, nodes, tests, splits, targets, weights, members + ) + return index, 0 + + +# One stump per branch test, so both branches of each are taken. The ordering tests split at +# 0.5, which the seeded draws fall either side of; the equality tests split at zero, which the +# special-value list carries as both of its signs. +_STUMPS = { + test: [ + _Split(0, test, 0.0 if test in ("EQ", "NEQ") else 0.5, [(0, 1.5)], [(0, -2.5)]) + ] + for test in _TESTS[:6] +} +# A tree deep enough that a row's path depends on more than one feature, and one whose +# branches route a missing feature the way the flag names rather than the test. +_DEEP = [ + _Split( + 0, + "LEQ", + 0.5, + _Split(1, "GT", -1.0, [(0, 1.0)], [(1, 2.0)]), + _Split(2, "LT", 2.0, [(1, -1.0)], [(0, 0.25)]), + ) +] +_MISSING = [ + _Split( + 0, + "LEQ", + 0.5, + [(0, 1.0)], + _Split(1, "GTE", 0.0, [(1, 2.0)], [(0, -3.0)], missing=1), + missing=1, + ) +] +_FOREST = [ + _Split(0, "LEQ", 0.5, [(0, 1.0)], [(1, 2.0)]), + _Split(1, "GT", 0.0, [(1, -0.5)], [(0, 3.0)]), + _Split(2, "LEQ", -1.0, [(0, 0.75)], [(1, -2.25)]), +] +# A leaf weighting two targets at once, which only the legacy encoding can express. +_MULTI_TARGET = [_Split(0, "LEQ", 0.5, [(0, 1.0), (1, -1.0)], [(1, 2.0)])] +# A forest whose every leaf weights both classes, so that no row's score reaches 0 or 1 — +# the two values a probit is not defined at, and the two the reference implementation's own +# `numpy.vectorize` returns a Python `int` for, which makes it return `float64` scores +# instead of the `tensor(float)` the schema declares. +_PROBABILITY_FOREST = [ + _Split(0, "LEQ", 0.5, [(0, 0.1), (1, 0.25)], [(0, 0.3), (1, 0.05)]), + _Split(1, "GT", 0.0, [(0, 0.2), (1, 0.15)], [(0, 0.05), (1, 0.3)]), + _Split(2, "LEQ", -1.0, [(0, 0.25), (1, 0.1)], [(0, 0.15), (1, 0.2)]), +] +# Set tests, which only opset 5 defines. The second set carries a zero, which the reference's +# own loop reads as the end of the set rather than as a member of it. +_MEMBERSHIP = [ + _Split( + 0, + "MEMBER", + 0.0, + [(0, 1.0)], + _Split(0, "MEMBER", 0.0, [(1, 2.0)], [(0, -3.0)], members=(0.0, 1.0)), + members=(1.0, -1.0, 2.0), + ) +] + +# `wide` is the only shape larger than the special-value list, so it is what carries every +# float edge -- NaN and the infinities included -- into the features a branch tests. +_LEGACY_SHAPES = ( + ("matrix", (4, 3)), + ("wide", (5, 3)), + ("vector", (3,)), + ("empty_rows", (0, 3)), +) + + +def _regressor_variants() -> tuple[Variant, ...]: + """`TreeEnsembleRegressor` over every branch test, aggregation and score transform. + + Only float32 inputs are swept: the reference implementation scores in the element type of + its *input*, so for a double `X` its result is not the `tensor(float)` the schema declares + and for an integer one it refuses to accumulate at all -- neither is an oracle for what + the op's own contract says the result should be. + """ + forests = [ + (f"stump_{test.lower()}", trees, 1) for test, trees in _STUMPS.items() + ] + [ + ("deep", _DEEP, 2), + ("missing", _MISSING, 2), + ("forest", _FOREST, 2), + ("multi_target", _MULTI_TARGET, 2), + ] + variants = [ + Variant( + label, + ((4, 3),), + _legacy_attributes(trees, "target", n_targets=targets), + elem_types=_FLOAT_ONLY, + ) + for label, trees, targets in forests + ] + variants += [ + Variant( + f"forest_{label}", + (shape,), + _legacy_attributes(_FOREST, "target", n_targets=2, **attributes), + elem_types=_FLOAT_ONLY, + ) + for label, shape, attributes in ( + *((label, shape, {}) for label, shape in _LEGACY_SHAPES[1:]), + ("sum", (4, 3), {"aggregate_function": "SUM"}), + ("average", (4, 3), {"aggregate_function": "AVERAGE"}), + ("minimum", (4, 3), {"aggregate_function": "MIN"}), + ("maximum", (4, 3), {"aggregate_function": "MAX"}), + ("base_values", (4, 3), {"base_values": [0.25, -0.5]}), + ("shared_base_value", (4, 3), {"base_values": [1.5]}), + ( + "base_values_average", + (4, 3), + {"base_values": [0.25, -0.5], "aggregate_function": "AVERAGE"}, + ), + ( + "base_values_minimum", + (4, 3), + {"base_values": [0.25, -0.5], "aggregate_function": "MIN"}, + ), + ("softmax", (4, 3), {"post_transform": "SOFTMAX"}), + ("explicit_none", (4, 3), {"post_transform": "NONE"}), + ) + ] + return tuple(variants) + + +def _classifier_variants() -> tuple[Variant, ...]: + """`TreeEnsembleClassifier` over the same forests, its binary rule and all five transforms. + + The scores are float32 and the labels int64 whatever `X` holds, so float32 and double are + both swept; the integer input types are not, since the reference rounds those to float32 + before it compares -- which is neither what the op's text says nor what the regressor's + own reference does with them. + """ + scored = (TensorProto.FLOAT, TensorProto.DOUBLE) + forests = [ + (f"stump_{test.lower()}", trees, [0, 1]) for test, trees in _STUMPS.items() + ] + [ + ("deep", _DEEP, [10, 20]), + ("missing", _MISSING, [10, 20]), + ("forest", _FOREST, [10, 20]), + ("three_classes", _MULTI_TARGET, [1, 2, 3]), + ] + variants = [ + Variant( + label, + ((4, 3),), + _legacy_attributes(trees, "class", classlabels_int64s=classes), + elem_types=scored, + outputs=2, + ) + for label, trees, classes in forests + ] + # An ensemble whose leaves all weight one class is the binary case: the reference pairs + # the score with a second column derived from it, which the transform then decides. + binary = [_Split(0, "LEQ", 0.5, [(0, 0.25)], [(0, 0.75)])] + variants += [ + Variant( + f"binary_{label}", + ((4, 3),), + _legacy_attributes( + binary, "class", classlabels_int64s=classes, **attributes + ), + elem_types=scored, + outputs=2, + ) + for label, classes, attributes in ( + ("pair", [0, 1], {}), + ("softmax", [0, 1], {"post_transform": "SOFTMAX"}), + ("logistic", [0, 1], {"post_transform": "LOGISTIC"}), + ("softmax_zero", [0, 1], {"post_transform": "SOFTMAX_ZERO"}), + ("probit", [0, 1], {"post_transform": "PROBIT"}), + ("three_labels", [7, 8, 9], {}), + # A single class label, where the reference widens the scores to two columns. + ("single_label", [1], {}), + ("single_label_logistic", [1], {"post_transform": "LOGISTIC"}), + ) + ] + variants += [ + Variant( + f"forest_{label}", + (shape,), + _legacy_attributes( + _FOREST, "class", classlabels_int64s=[10, 20], **attributes + ), + elem_types=scored, + outputs=2, + ) + for label, shape, attributes in ( + *((label, shape, {}) for label, shape in _LEGACY_SHAPES[1:]), + ("softmax", (4, 3), {"post_transform": "SOFTMAX"}), + ("logistic", (4, 3), {"post_transform": "LOGISTIC"}), + ("softmax_zero", (4, 3), {"post_transform": "SOFTMAX_ZERO"}), + ("base_values", (4, 3), {"base_values": [0.25, -0.5]}), + ( + "base_values_softmax", + (4, 3), + {"base_values": [0.25, -0.5], "post_transform": "SOFTMAX"}, + ), + ) + ] + # A probit maps a probability, so it is swept over the forest whose scores are ones. + variants += [ + Variant( + f"probabilities_{label}", + ((4, 3),), + _legacy_attributes( + _PROBABILITY_FOREST, + "class", + classlabels_int64s=[10, 20], + **attributes, + ), + elem_types=scored, + outputs=2, + ) + for label, attributes in ( + ("probit", {"post_transform": "PROBIT"}), + ("none", {}), + ) + ] + return tuple(variants) + + +def _bare_leaf_attributes(dtype: str) -> dict[str, Any]: + """A tree that is a single leaf, in the one shape its reference implementation reads. + + The root of such a tree indexes the *leaf* families rather than the node ones, so it is + only expressible where the two indices coincide -- which is why this one is written out + rather than built from a tree. + """ + return { + "nodes_featureids": [0], + "nodes_truenodeids": [0], + "nodes_falsenodeids": [0], + "nodes_trueleafs": [1], + "nodes_falseleafs": [1], + "nodes_modes": numpy_helper.from_array(np.array([0], np.uint8), "nodes_modes"), + "nodes_splits": numpy_helper.from_array(np.array([0.5], dtype), "nodes_splits"), + "leaf_targetids": [1], + "leaf_weights": numpy_helper.from_array(np.array([4.5], dtype), "leaf_weights"), + "tree_roots": [0], + "n_targets": 2, + } + + +def _tree_ensemble_variants() -> tuple[Variant, ...]: + """`TreeEnsemble` over the branch tests, the set tests, and every aggregation. + + Every case is generated once per element type rather than once for both: ONNX's own type + inference requires the splits, the weights and the set members to carry the element type + of `X`, so a forest's tables belong to the dtype they are swept at. + + A zero-row `X` is deliberately absent. The reference implementation drops each row through + the trees with `numpy.apply_along_axis`, which refuses an empty axis outright, so there is + no oracle for one; the legacy pair, which loop over the rows themselves, do sweep it. + """ + forests = [ + (f"stump_{test.lower()}", trees, 1) for test, trees in _STUMPS.items() + ] + [ + ("deep", _DEEP, 2), + ("missing", _MISSING, 2), + ("forest", _FOREST, 2), + ("membership", _MEMBERSHIP, 2), + ] + aggregations = ( + # Wider than the special-value list, so every float edge reaches the features. + ("wide", (5, 3), {}), + ("average", (4, 3), {"aggregate_function": 0}), + ("sum", (4, 3), {"aggregate_function": 1}), + ("minimum", (4, 3), {"aggregate_function": 2}), + ("maximum", (4, 3), {"aggregate_function": 3}), + ("softmax", (4, 3), {"post_transform": 1}), + ("explicit_none", (4, 3), {"post_transform": 0}), + ) + variants: list[Variant] = [] + for elem_type in (TensorProto.FLOAT, TensorProto.DOUBLE): + dtype = numpy_dtype_name(elem_type) + variants += [ + Variant( + f"{label}_{dtype}", + ((4, 3),), + _ensemble_attributes(trees, dtype, n_targets=targets), + elem_types=(elem_type,), + ) + for label, trees, targets in forests + ] + variants += [ + Variant( + f"forest_{label}_{dtype}", + (shape,), + _ensemble_attributes(_FOREST, dtype, n_targets=2, **attributes), + elem_types=(elem_type,), + ) + for label, shape, attributes in aggregations + ] + variants.append( + Variant( + f"bare_leaf_{dtype}", + ((4, 3),), + _bare_leaf_attributes(dtype), + elem_types=(elem_type,), + ) + ) + return tuple(variants) + + +# What ONNX's own standard-domain ops say a score transform computes, for the predictors +# whose reference implementation cannot apply one. +_TRANSFORM_EQUIVALENTS: Mapping[Any, tuple[str, Mapping[str, Any]]] = { + "SOFTMAX": ("Softmax", {"axis": 1}), + 1: ("Softmax", {"axis": 1}), +} + + +def _as_transformed_scores(case: Case) -> ModelProto: + """The case's predictor with its score transform spelled out as the op ONNX defines. + + `TreeEnsembleRegressor`'s reference implementation raises for every transform but `NONE`, + `TreeEnsemble`'s ignores the attribute outright, and the two regressors below raise like + the first, so none of them can be the oracle for a transformed model. Each is compared + instead against the untransformed op followed by the standard-domain op the transform is + defined to be -- which is also what the transform the classifiers' references *do* + implement computes, and those are compared against the reference directly. + """ + equivalent = _TRANSFORM_EQUIVALENTS.get( + case.variant.attributes.get("post_transform") + ) + if equivalent is None: + return _model(case) + untransformed = { + name: value + for name, value in case.variant.attributes.items() + if name != "post_transform" + } + model = _model( + replace(case, variant=replace(case.variant, attributes=untransformed)) + ) + op_type, attributes = equivalent + model.graph.node[0].output[0] = "scores" + model.graph.node.append( + helper.make_node(op_type, ["scores"], ["out0"], name="transform", **attributes) + ) + return _importing_standard_ops(model) + + +# -------------------------------------------------------------------------------------- +# Support vector machines and linear models +# -------------------------------------------------------------------------------------- + +# Every one of the four scores dot products, so their sweeps are `ACCUMULATING` for the reason +# Gemm's is: the summation order the reference takes is numpy's and the kernel's is a loop, +# and with an infinity or a dtype extreme among the products the two legitimately disagree. +# What the variants sweep instead is the *shape* of each model -- how many coefficient rows, +# how many class labels, which kernel function, and which of the readings of a row of scores +# each combination lands on. + +_PREDICTOR_SHAPES = (("matrix", (4, 3)), ("single_row", (1, 3)), ("empty_rows", (0, 3))) +# One coefficient per feature, and a second row of them for the two-output cases. +_FIRST_ROW = [1.0, 0.0, -1.0] +_SECOND_ROW = [-0.5, 0.5, 0.25] +_TWO_ROWS = _FIRST_ROW + _SECOND_ROW +# Three support vectors over the same three features, and the gamma/coef0/degree triple every +# kernel function reads from. +_SUPPORT_VECTORS = [1.0, 2.0, 3.0, 0.0, 0.0, 1.0, -1.0, 0.5, 2.0] +_KERNEL_PARAMS = [0.5, 1.0, 3.0] +_KERNEL_TYPES = ("LINEAR", "POLY", "RBF", "SIGMOID") +# Coefficients small enough, against an intercept of one half, that every score lands inside +# the `[0, 1]` a probit is defined over -- outside it the transform is NaN, which compares +# equal to itself and would leave the case asserting nothing about the arithmetic. +_SMALL_ROW = [0.01, 0.02, -0.01] + + +def _linear_regressor_variants() -> tuple[Variant, ...]: + """`LinearRegressor` over one and several targets, and over both intercept layouts.""" + single = {"coefficients": _FIRST_ROW, "intercepts": [0.5]} + two = {"coefficients": _TWO_ROWS, "intercepts": [0.5, -0.25], "targets": 2} + return ( + *(Variant(label, (shape,), single) for label, shape in _PREDICTOR_SHAPES), + Variant("two_targets", ((4, 3),), two), + Variant( + "shared_intercept", + ((4, 3),), + {"coefficients": _TWO_ROWS, "intercepts": [0.5], "targets": 2}, + ), + Variant("wide", ((4, 8),), {"coefficients": [0.25] * 8, "intercepts": [0.5]}), + Variant("explicit_none", ((4, 3),), {**single, "post_transform": "NONE"}), + Variant("softmax", ((4, 3),), {**two, "post_transform": "SOFTMAX"}), + ) + + +def _linear_classifier_variants() -> tuple[Variant, ...]: + """`LinearClassifier` over every reading of a row of scores its reference has. + + One coefficient row per class is what converters emit and what the transforms are swept + over; a single row against two labels is the paired case, where the score is set against + its own negation, and a single row against anything else is the thresholded one, where + the label is decided by which side of zero -- or of one half, once a transform has mapped + the score onto a probability -- the single column falls. + """ + binary = {"coefficients": _TWO_ROWS, "intercepts": [0.5, -0.25]} + labelled = {**binary, "classlabels_ints": [3, 7]} + single = {"coefficients": _FIRST_ROW, "intercepts": [0.5]} + three = { + "coefficients": _TWO_ROWS + [0.25, -0.75, 1.0], + "intercepts": [0.5, -0.25, 0.0], + "classlabels_ints": [1, 2, 3], + } + probabilities = { + "coefficients": _SMALL_ROW * 3, + "intercepts": [0.5, 0.5, 0.5], + "classlabels_ints": [1, 2, 3], + } + return ( + *( + Variant(label, (shape,), labelled, outputs=2) + for label, shape in _PREDICTOR_SHAPES + ), + *( + Variant( + f"binary_{transform.lower()}", + ((4, 3),), + {**labelled, "post_transform": transform}, + outputs=2, + ) + for transform in ("NONE", "LOGISTIC", "SOFTMAX", "SOFTMAX_ZERO") + ), + Variant( + "shared_intercept", + ((4, 3),), + {**labelled, "intercepts": [0.5]}, + outputs=2, + ), + Variant( + "wide", ((4, 8),), {**labelled, "coefficients": [0.25] * 16}, outputs=2 + ), + Variant("three_classes", ((4, 3),), three, outputs=2), + Variant( + "three_classes_softmax", + ((4, 3),), + {**three, "post_transform": "SOFTMAX"}, + outputs=2, + ), + Variant( + "probit", + ((4, 3),), + {**probabilities, "post_transform": "PROBIT"}, + outputs=2, + ), + # `multi_class` is an attribute the reference reads and then ignores. + Variant("multi_class", ((4, 3),), {**three, "multi_class": 1}, outputs=2), + Variant("paired", ((4, 3),), {**single, "classlabels_ints": [3, 7]}, outputs=2), + Variant( + "paired_logistic", + ((4, 3),), + {**single, "classlabels_ints": [3, 7], "post_transform": "LOGISTIC"}, + outputs=2, + ), + Variant( + "single_class", ((4, 3),), {**single, "classlabels_ints": [7]}, outputs=2 + ), + Variant( + "single_class_logistic", + ((4, 3),), + {**single, "classlabels_ints": [7], "post_transform": "LOGISTIC"}, + outputs=2, + ), + # No class labels at all, where a thresholded row is labelled 1 or 0. + Variant("no_labels", ((4, 3),), single, outputs=2), + ) + + +def _svm_regressor_variants() -> tuple[Variant, ...]: + """`SVMRegressor` in both its modes, over every kernel function ONNX defines.""" + linear = {"coefficients": _FIRST_ROW, "rho": [0.25]} + supports = { + "coefficients": [1.0, -0.5, 0.25], + "rho": [0.25], + "n_supports": 3, + "support_vectors": _SUPPORT_VECTORS, + "kernel_params": _KERNEL_PARAMS, + } + return ( + *(Variant(label, (shape,), linear) for label, shape in _PREDICTOR_SHAPES), + Variant("wide", ((4, 8),), {"coefficients": [0.25] * 8, "rho": [0.25]}), + Variant("one_class", ((4, 3),), {**linear, "one_class": 1}), + Variant("softmax", ((4, 3),), {**linear, "post_transform": "SOFTMAX"}), + *( + Variant( + f"supports_{kernel.lower()}", + ((4, 3),), + {**supports, "kernel_type": kernel}, + ) + for kernel in _KERNEL_TYPES + ), + # No `kernel_params` leaves gamma, coef0 and the degree at zero. + Variant( + "supports_no_params", + ((4, 3),), + { + name: value + for name, value in supports.items() + if name != "kernel_params" + }, + ), + Variant("supports_one_class", ((4, 3),), {**supports, "one_class": 1}), + Variant("supports_empty_rows", ((0, 3),), supports), + # The reference reads one coefficient per support vector and ignores the rest. + Variant( + "supports_spare_coefficients", + ((4, 3),), + {**supports, "coefficients": [1.0, -0.5, 0.25, 9.0]}, + ), + ) + + +def _svm_classifier_variants() -> tuple[Variant, ...]: + """`SVMClassifier` over both modes, all three label rules and the probability coupling. + + A zero-row `X` is deliberately absent: the reference sizes its score matrix from the first + row it computes, so with no rows it returns `None` for the scores and there is nothing to + compare against. The other three ops sweep that shape, over the same two kernels this one + writes its scores with. + """ + labels = {"classlabels_ints": [3, 7]} + linear = {"coefficients": _TWO_ROWS, "rho": [0.25], **labels} + three = { + "coefficients": _TWO_ROWS + [0.25, -0.75, 1.0], + "rho": [0.25], + "classlabels_ints": [1, 2, 3], + } + single = {"coefficients": _FIRST_ROW, "rho": [0.25], "classlabels_ints": [7]} + supports = { + "coefficients": [0.5, -0.25, 0.75], + "rho": [0.25], + "vectors_per_class": [2, 1], + "support_vectors": _SUPPORT_VECTORS, + "kernel_params": _KERNEL_PARAMS, + **labels, + } + trio = { + "coefficients": [0.5, -0.25, 0.75, 0.1, 0.2, -0.3], + "rho": [0.25, 0.1, -0.2], + "vectors_per_class": [1, 1, 1], + "support_vectors": _SUPPORT_VECTORS, + "kernel_params": _KERNEL_PARAMS, + "kernel_type": "LINEAR", + "classlabels_ints": [1, 2, 3], + } + return ( + Variant("linear_binary", ((4, 3),), linear, outputs=2), + Variant("linear_single_row", ((1, 3),), linear, outputs=2), + Variant( + "linear_wide", ((4, 8),), {**linear, "coefficients": [0.25] * 16}, outputs=2 + ), + *( + Variant( + f"linear_binary_{transform.lower()}", + ((4, 3),), + {**linear, "post_transform": transform}, + outputs=2, + ) + for transform in ("NONE", "LOGISTIC", "SOFTMAX", "SOFTMAX_ZERO", "PROBIT") + ), + Variant("linear_three_classes", ((4, 3),), three, outputs=2), + # More than one `rho` takes the plain reading of the winning column, whatever the + # class labels say. + Variant( + "linear_three_rho", + ((4, 3),), + {**three, "rho": [0.25, 0.1, -0.2]}, + outputs=2, + ), + # A single score against a single class label: the label is the sign of that score, + # and the row is returned before it ever reaches the transform. + Variant("linear_single_class", ((4, 3),), single, outputs=2), + Variant( + "linear_single_class_logistic", + ((4, 3),), + {**single, "post_transform": "LOGISTIC"}, + outputs=2, + ), + Variant( + "linear_single_class_probit", + ((4, 3),), + {**single, "post_transform": "PROBIT"}, + outputs=2, + ), + *( + Variant( + f"supports_{kernel.lower()}", + ((4, 3),), + {**supports, "kernel_type": kernel}, + outputs=2, + ) + for kernel in _KERNEL_TYPES + ), + Variant("supports_single_row", ((1, 3),), supports, outputs=2), + # No coefficient below zero is the one case where a winning vote of at least one half + # names the second class outright. + Variant( + "supports_all_positive", + ((4, 3),), + {**supports, "coefficients": [0.5, 0.25, 0.75]}, + outputs=2, + ), + *( + Variant( + f"supports_{transform.lower()}", + ((4, 3),), + {**supports, "post_transform": transform}, + outputs=2, + ) + for transform in ("LOGISTIC", "SOFTMAX", "SOFTMAX_ZERO", "PROBIT") + ), + Variant("supports_three_classes", ((4, 3),), trio, outputs=2), + Variant( + "supports_three_classes_softmax", + ((4, 3),), + {**trio, "post_transform": "SOFTMAX"}, + outputs=2, + ), + # Platt scaling, which turns the single decision value of a class pair into the two + # probabilities the row is then scored with. + Variant( + "supports_probabilities", + ((4, 3),), + {**supports, "prob_a": [-1.5], "prob_b": [0.25]}, + outputs=2, + ), + Variant( + "supports_probabilities_softmax", + ((4, 3),), + { + **supports, + "prob_a": [-1.5], + "prob_b": [0.25], + "post_transform": "SOFTMAX", + }, + outputs=2, + ), + ) + + +def _as_float_predictor(case: Case) -> ModelProto: + """A predictor's oracle: float32 operands, and a transform its reference cannot apply.""" + return _cast_inputs_to_float(_as_transformed_scores(case)) + + +SWEEP: dict[tuple[str, str], Sweep] = { + **{ + ("", op_type): Sweep(Kind.POINTWISE, _UNARY_VARIANTS) + for op_type in _POINTWISE_UNARY_OPS + }, + # The reference divides into a 0-d output array, which numpy refuses outright, so a + # rank-0 Softsign has no oracle to be compared against. + ("", "Softsign"): Sweep( + Kind.POINTWISE, + tuple(variant for variant in _UNARY_VARIANTS if variant.label != "rank_0"), + ), + # Negation and the bias are arithmetic, so the integer extremes — where C's overflow is + # undefined and ONNX's is unstated — stay out of these. + **{ + ("", op_type): Sweep(Kind.ARITHMETIC, _UNARY_VARIANTS) + for op_type in ("Abs", "Neg") + }, + **{ + ("", op_type): Sweep(Kind.ARITHMETIC, _BROADCAST_VARIANTS) + for op_type in ("Add", "Sub", "Mul") + }, + **{ + ("", op_type): Sweep(Kind.POINTWISE, _VARIADIC_VARIANTS) + for op_type in ("Max", "Min") + }, + **{ + ("", op_type): Sweep(Kind.ARITHMETIC, _VARIADIC_VARIANTS) + for op_type in ("Mean", "Sum") + }, + **{ + ("", op_type): Sweep( + Kind.POINTWISE, _with_attributes(_UNARY_VARIANTS, combinations) + ) + for op_type, combinations in _ACTIVATION_ATTRIBUTES.items() + }, + # Comparisons and boolean logic: nothing they compute can leave a dtype's range, so + # every operand sweeps its extremes. + **{ + ("", op_type): Sweep(Kind.POINTWISE, _BROADCAST_VARIANTS) + for op_type in ( + "And", + "BitwiseAnd", + "BitwiseOr", + "BitwiseXor", + "Equal", + "Greater", + "GreaterOrEqual", + "Less", + "LessOrEqual", + "Or", + "Xor", + ) + }, + **{ + ("", op_type): Sweep(Kind.POINTWISE, _UNARY_VARIANTS) + for op_type in ("BitwiseNot", "IsNaN", "Not") + }, + ("", "BitShift"): Sweep(Kind.POINTWISE, _BIT_SHIFT_VARIANTS), + ("", "BitCast"): Sweep(Kind.POINTWISE, _bitcast_variants()), + ("", "Cast"): Sweep(Kind.POINTWISE, _cast_variants()), + ("", "IsInf"): Sweep( + Kind.POINTWISE, _with_attributes(_UNARY_VARIANTS, _IS_INF_ATTRIBUTES) + ), + ("", "Where"): Sweep( + Kind.POINTWISE, + _SELECT_VARIANTS, + type_operand=1, + operand_types={0: TensorProto.BOOL}, + ), + ("", "Div"): Sweep(Kind.ARITHMETIC, _BROADCAST_VARIANTS, {1: Domain.NONZERO}), + ("", "Mod"): Sweep(Kind.ARITHMETIC, _MOD_VARIANTS, {1: Domain.NONZERO}), + ("", "Pow"): Sweep( + Kind.ARITHMETIC, _BROADCAST_VARIANTS, {1: Domain.SMALL_EXPONENT} + ), + ("", "PRelu"): Sweep(Kind.ARITHMETIC, _UNIDIRECTIONAL_VARIANTS), + ("", "Shrink"): Sweep( + Kind.ARITHMETIC, _with_attributes(_UNARY_VARIANTS, _SHRINK_ATTRIBUTES) + ), + ("", "Clip"): Sweep(Kind.POINTWISE, _CLIP_VARIANTS), + ("", "Dropout"): Sweep(Kind.POINTWISE, _DROPOUT_VARIANTS), + ("", "Gemm"): Sweep(Kind.ACCUMULATING, _GEMM_VARIANTS), + ("", "MatMul"): Sweep(Kind.ACCUMULATING, _MATMUL_VARIANTS), + # Det multiplies the pivots of a factorization, which is an accumulation of its own: a + # special value anywhere in a matrix decides the whole determinant, and by which pivots + # were chosen rather than by the spec. + ("", "Det"): Sweep(Kind.ACCUMULATING, _DET_VARIANTS), + # Einsum sums a product per result element in an order the equation does not fix, so it + # accumulates for the reason Gemm does. + ("", "Einsum"): Sweep(Kind.ACCUMULATING, _EINSUM_VARIANTS), + # The transforms sum every sample of an axis into every bin, so they accumulate for the + # reason Gemm does — and more so: the reference runs numpy's FFT at the operand's own + # precision, whose butterfly leaves a residue of its own where a bin cancels to zero. + # What each reads as configuration rather than as data — the length, the axis, the step — + # is carried as an initializer, since a compiler with no run-time shapes needs it fixed. + ("", "DFT"): Sweep( + Kind.ACCUMULATING, + _DFT_VARIANTS, + operand_types={1: TensorProto.INT64, 2: TensorProto.INT64}, + constant_operands=(1, 2), + ), + ("", "STFT"): Sweep( + Kind.ACCUMULATING, + _STFT_VARIANTS, + operand_types={1: TensorProto.INT64, 3: TensorProto.INT64}, + constant_operands=(1, 3), + ), + # A convolution sums a whole window of products per output element, in an order the spec + # leaves open, so it accumulates for the reason Gemm does. + ("", "Conv"): Sweep(Kind.ACCUMULATING, _CONV_VARIANTS), + ("", "ConvTranspose"): Sweep(Kind.ACCUMULATING, _CONV_TRANSPOSE_VARIANTS), + ("", "DeformConv"): Sweep(Kind.ACCUMULATING, _DEFORM_CONV_VARIANTS), + # The affine maps are elementwise arithmetic, so they sweep the float edges; what bounds + # them is only what the reference leaves undefined, which is the conversion at the end of + # a quantization -- hence the restriction on its operand and the pinned scales its table + # records. The type swept is the grid rather than the operand: for QuantizeLinear that is + # the zero point's, which is what decides the result's, and its operand and scale are + # pinned to the one float type ONNX allows them both to carry at every revision claimed. + ("", "QuantizeLinear"): Sweep( + Kind.ARITHMETIC, + _QUANTIZE_VARIANTS, + {0: Domain.CONVERTIBLE}, + type_operand=2, + operand_types={0: TensorProto.FLOAT, 1: TensorProto.FLOAT}, + ), + ("", "DequantizeLinear"): Sweep( + Kind.POINTWISE, + _DEQUANTIZE_VARIANTS, + operand_types={1: TensorProto.FLOAT}, + ), + # The quantized products sum a window or a row of products, so they accumulate for the + # reason Gemm does. Their bias is `int32` whatever grid the operands stand on. + ("", "MatMulInteger"): Sweep(Kind.ACCUMULATING, _MATMUL_INTEGER_VARIANTS), + ("", "QLinearMatMul"): Sweep( + Kind.ACCUMULATING, + _qlinear_matmul_variants(), + operand_types={ + 1: TensorProto.FLOAT, + 4: TensorProto.FLOAT, + 6: TensorProto.FLOAT, + }, + ), + ("", "ConvInteger"): Sweep(Kind.ACCUMULATING, _conv_integer_variants()), + ("", "QLinearConv"): Sweep( + Kind.ACCUMULATING, + _qlinear_conv_variants(), + operand_types={ + 1: TensorProto.FLOAT, + 4: TensorProto.FLOAT, + 6: TensorProto.FLOAT, + 8: TensorProto.INT32, + }, + ), + # AveragePool and LpPool sum a window, so they accumulate for the reason Gemm does. + # MaxPool and GlobalMaxPool only compare — but the reference evaluator runs them through + # whichever of its two pooling implementations the attributes select, and the two disagree + # about a NaN in the window (one drops it, the other lets it win where it comes first) and + # about which of a window's equal maxima is reported, which is what tells -0.0 and 0.0 + # apart. So the whole family is fed the finite values every path agrees on, and the + # special ones rest on the backend corpus. + ("", "AveragePool"): Sweep(Kind.ACCUMULATING, _AVERAGE_POOL_VARIANTS), + ("", "LpPool"): Sweep(Kind.ACCUMULATING, _LP_POOL_VARIANTS), + ("", "MaxPool"): Sweep(Kind.ACCUMULATING, _MAX_POOL_VARIANTS), + ("", "GlobalAveragePool"): Sweep(Kind.ACCUMULATING, _GLOBAL_VARIANTS), + # The other two global folds are run against the windowed op each is defined to equal. + # GlobalLpPool has to be: the reference evaluator implements it nowhere. GlobalMaxPool it + # has one for, but that one reduces `range(rank - 2, rank)` — the spatial axes only for a + # 4-D operand, and the wrong axes for every other rank — so it is no oracle for the ranks + # this kernel equally serves, and nothing else in the corpus is either. + ("", "GlobalMaxPool"): Sweep( + Kind.ACCUMULATING, + _GLOBAL_MAX_VARIANTS, + equivalent_model=_as_windowed_pooling("MaxPool"), + ), + ("", "GlobalLpPool"): Sweep( + Kind.ACCUMULATING, + _GLOBAL_LP_VARIANTS, + equivalent_model=_as_windowed_pooling("LpPool"), + ), + # MaxUnpool moves elements without computing any, so every special value goes in and has + # to come back out unchanged; the positions it moves them to are int64 whatever they hold. + ("", "MaxUnpool"): Sweep( + Kind.POINTWISE, _MAX_UNPOOL_VARIANTS, operand_types={1: TensorProto.INT64} + ), + # The reductions and the running folds accumulate a whole group in an order the spec does + # not fix, so they carry ACCUMULATING for the same reason Gemm does. ReduceMax, ReduceMin, + # ArgMax, ArgMin and Hardmax compare rather than accumulate — no order changes what they + # select — so those sweep every special value their dtype has. + ("", "ReduceSum"): _reduction_sweep(Kind.ACCUMULATING, *_SUM_VERSIONS), + ("", "ReduceMean"): _reduction_sweep( + Kind.ACCUMULATING, + *_REDUCTION_VERSIONS, + # numpy's mean of nothing is a 0/0 it casts unsafely to the element type, which is a + # value only the floating-point families have. + empty_group_types=_FLOAT_ELEM_TYPES, + ), + ("", "ReduceProd"): _reduction_sweep( + Kind.ACCUMULATING, *_REDUCTION_VERSIONS, factors=True + ), + ("", "ReduceL1"): _reduction_sweep(Kind.ACCUMULATING, *_REDUCTION_VERSIONS), + ("", "ReduceL2"): _reduction_sweep(Kind.ACCUMULATING, *_REDUCTION_VERSIONS), + # The reference evaluator is no oracle for these three on the integer families: it raises + # outright for the two logarithmic ones, and returns a ReduceSumSquare whose dtype + # disagrees with ONNX's own type inference. Their integer kernels are emitted the way the + # op is defined; the sweep stops where the oracle does. + ("", "ReduceLogSum"): _reduction_sweep( + Kind.ACCUMULATING, *_REDUCTION_VERSIONS, elem_types=_FLOAT_ELEM_TYPES + ), + ("", "ReduceLogSumExp"): _reduction_sweep( + Kind.ACCUMULATING, *_REDUCTION_VERSIONS, elem_types=_FLOAT_ELEM_TYPES + ), + ("", "ReduceSumSquare"): _reduction_sweep( + Kind.ACCUMULATING, *_REDUCTION_VERSIONS, elem_types=_FLOAT_ELEM_TYPES + ), + ("", "ReduceMax"): _reduction_sweep( + Kind.POINTWISE, *_EXTREMUM_VERSIONS, empty_group_types=_NON_BOOL_TYPES + ), + ("", "ReduceMin"): _reduction_sweep( + Kind.POINTWISE, *_EXTREMUM_VERSIONS, empty_group_types=_NON_BOOL_TYPES + ), + **{ + ("", op_type): Sweep(Kind.POINTWISE, _ARG_VARIANTS) + for op_type in ("ArgMax", "ArgMin") + }, + **{ + ("", op_type): Sweep(Kind.ACCUMULATING, _ALONG_AXIS_VARIANTS) + for op_type in ("LogSoftmax", "Softmax") + }, + ("", "Hardmax"): Sweep(Kind.POINTWISE, _ALONG_AXIS_VARIANTS), + ("", "CumSum"): Sweep( + Kind.ACCUMULATING, + _CUMULATIVE_VARIANTS, + operand_types={1: TensorProto.INT64}, + ), + ("", "CumProd"): Sweep( + Kind.ACCUMULATING, + _CUMULATIVE_VARIANTS, + {0: Domain.SMALL_FACTOR}, + operand_types={1: TensorProto.INT64}, + ), + ("", "BatchNormalization"): Sweep( + Kind.ACCUMULATING, _batch_variants(), {4: Domain.NONNEGATIVE} + ), + ("", "LayerNormalization"): Sweep(Kind.ACCUMULATING, _layer_variants()), + ("", "RMSNormalization"): Sweep(Kind.ACCUMULATING, _rms_variants()), + # The labels index the class axis, so they do not range over the element type the logits + # are swept at; the weights do, being read as one coefficient per class. + ("", "SoftmaxCrossEntropyLoss"): Sweep( + Kind.ACCUMULATING, _sce_variants(), operand_types={1: TensorProto.INT64} + ), + ("", "InstanceNormalization"): Sweep(Kind.ACCUMULATING, _INSTANCE_VARIANTS), + ("", "GroupNormalization"): Sweep(Kind.ACCUMULATING, _GROUP_VARIANTS), + ("", "LpNormalization"): Sweep(Kind.ACCUMULATING, _lp_variants()), + ("", "MeanVarianceNormalization"): Sweep(Kind.ACCUMULATING, _mvn_variants()), + ("", "LRN"): Sweep(Kind.ACCUMULATING, _LRN_VARIANTS), + # The views move elements without computing any, so every special value their dtype has + # goes in and has to come back out unchanged, signed zeros included. + ("", "Transpose"): Sweep(Kind.POINTWISE, _TRANSPOSE_VARIANTS), + ("", "Concat"): Sweep(Kind.POINTWISE, _CONCAT_VARIANTS), + ("", "Flatten"): Sweep(Kind.POINTWISE, _FLATTEN_VARIANTS), + # The operand describing the result's shape is carried in the model as an initializer: + # one a graph computes at run time makes the result's shape depend on input data, which + # the compiler refuses by design. + **{ + ("", op_type): Sweep( + Kind.POINTWISE, + variants, + operand_types={1: TensorProto.INT64}, + constant_operands=(1,), + ) + for op_type, variants in ( + ("Reshape", _RESHAPE_VARIANTS), + ("Squeeze", _SQUEEZE_VARIANTS), + ("Unsqueeze", _UNSQUEEZE_VARIANTS), + ("Split", _SPLIT_VARIANTS), + ("Expand", _EXPAND_VARIANTS), + ("Tile", _TILE_VARIANTS), + ) + }, + ("", "Slice"): Sweep( + Kind.POINTWISE, + _SLICE_VARIANTS, + operand_types=dict.fromkeys((1, 2, 3, 4), TensorProto.INT64), + constant_operands=(1, 2, 3, 4), + ), + # The gathering ops move elements without computing any, so every special value their + # dtype has goes in and has to come back out unchanged. Their indices are read at run + # time — the result's shape follows from the operands' shapes alone — so they are fed + # rather than carried in the model. + **{ + ("", op_type): Sweep( + Kind.POINTWISE, + variants, + operand_types={1: TensorProto.INT64}, + ) + for op_type, variants in ( + ("Gather", _GATHER_VARIANTS), + ("GatherElements", _GATHER_ELEMENTS_VARIANTS), + ("GatherND", _GATHER_ND_VARIANTS), + ) + }, + # The scattering ops write elements without computing any, so every special value their + # dtype has goes in and has to come back out unchanged — except where `reduction` folds + # an update into the element already there, which the variants carrying one restrict for. + # Their indices are read at run time: the result's shape is the operand's own, whatever + # they hold. + **{ + ("", op_type): Sweep( + Kind.POINTWISE, + variants, + operand_types={1: TensorProto.INT64}, + ) + for op_type, variants in ( + ("ScatterElements", _SCATTER_ELEMENTS_VARIANTS), + ("ScatterND", _SCATTER_ND_VARIANTS), + ) + }, + ("", "Scatter"): Sweep( + Kind.POINTWISE, + _SCATTER_VARIANTS, + operand_types={1: TensorProto.INT64}, + equivalent_model=_as_scatter_elements, + ), + ("", "TensorScatter"): Sweep( + Kind.POINTWISE, + _TENSOR_SCATTER_VARIANTS, + operand_types={2: TensorProto.INT64}, + ), + ("", "Pad"): Sweep( + Kind.POINTWISE, + _PAD_VARIANTS, + operand_types={1: TensorProto.INT64, 3: TensorProto.INT64}, + constant_operands=(1, 3), + ), + ("", "OneHot"): Sweep( + Kind.POINTWISE, + _ONE_HOT_VARIANTS, + # The output — and so the sweep's element types — is the type of the two values the + # op selects between; the indices and the depth are typed independently of it. + type_operand=2, + operand_types={0: TensorProto.INT64, 1: TensorProto.INT64}, + constant_operands=(1,), + ), + ("", "EyeLike"): Sweep(Kind.POINTWISE, _EYE_LIKE_VARIANTS), + ("", "Trilu"): Sweep( + Kind.POINTWISE, _TRILU_VARIANTS, operand_types={1: TensorProto.INT64} + ), + ("", "ReverseSequence"): Sweep( + Kind.POINTWISE, + _REVERSE_SEQUENCE_VARIANTS, + operand_types={1: TensorProto.INT64}, + ), + # A recurrent layer sums a whole row of products per gate and then carries the result + # forward through every remaining step, so it accumulates for the reason Gemm does, + # several times over. Its lengths are int32 by ONNX's own type constraint whatever the + # sweep's element type is. + ("", "LSTM"): Sweep( + Kind.ACCUMULATING, _LSTM_VARIANTS, operand_types={4: TensorProto.INT32} + ), + ("", "GRU"): Sweep( + Kind.ACCUMULATING, _GRU_VARIANTS, operand_types={4: TensorProto.INT32} + ), + ("", "RNN"): Sweep( + Kind.ACCUMULATING, _RNN_VARIANTS, operand_types={4: TensorProto.INT32} + ), + # LinearAttention sums a whole key dimension per state cell and per answer, and carries + # the state forward through every remaining token, so it accumulates for the reason the + # recurrent layers do. Its state operand is typed independently of its activations by + # ONNX, but the compiler serves the op at one element type and the sweep only has that + # one to offer either of them. + ("", "LinearAttention"): Sweep(Kind.ACCUMULATING, _LINEAR_ATTENTION_VARIANTS), + # A resize sums a filter's worth of products per output element, in an order the spec + # leaves open, so it accumulates for the reason Gemm does. Its three operands are typed + # by ONNX rather than by the sweep: the scales are float and the sizes int64 whatever + # the data holds, and the region defaults to float where a variant does not ask for + # another floating-point type. + ("", "Resize"): Sweep( + Kind.ACCUMULATING, + _RESIZE_VARIANTS, + operand_types={ + 1: TensorProto.FLOAT, + 2: TensorProto.FLOAT, + 3: TensorProto.INT64, + }, + constant_operands=(1, 2, 3), + ), + ("", "Upsample"): Sweep( + Kind.ACCUMULATING, + _UPSAMPLE_VARIANTS, + operand_types={1: TensorProto.FLOAT}, + constant_operands=(1,), + equivalent_model=_as_resize, + ), + # The block shuffles move elements without computing any, so every special value their + # dtype has goes in and has to come back out unchanged, signed zeros included. + ("", "DepthToSpace"): Sweep(Kind.POINTWISE, _DEPTH_TO_SPACE_VARIANTS), + ("", "SpaceToDepth"): Sweep(Kind.POINTWISE, _SPACE_TO_DEPTH_VARIANTS), + # Col2Im sums the blocks that reach an image position, in the order ONNX's own reference + # accumulates them, so the float edges are compared rather than left out: arithmetic, not + # accumulation. Its two extent operands are int64 whatever the data holds. + ("", "Col2Im"): Sweep( + Kind.ARITHMETIC, + _COL2IM_VARIANTS, + operand_types={1: TensorProto.INT64, 2: TensorProto.INT64}, + constant_operands=(1, 2), + ), + # The samplers weight the elements around a coordinate and sum them, in an order the + # spec leaves open, so they accumulate for the reason Gemm does. + ("", "GridSample"): Sweep( + Kind.ACCUMULATING, + _GRID_SAMPLE_VARIANTS, + operand_types={1: TensorProto.FLOAT}, + ), + ("", "AffineGrid"): Sweep( + Kind.ACCUMULATING, + _AFFINE_GRID_VARIANTS, + operand_types={1: TensorProto.INT64}, + constant_operands=(1,), + ), + ("", "RoiAlign"): Sweep( + Kind.ACCUMULATING, + _ROI_ALIGN_VARIANTS, + operand_types={2: TensorProto.INT64}, + ), + # MaxRoiPool compares rather than accumulates, so it sweeps every special value its dtype + # has. ONNX ships neither a reference implementation nor a node test for it — it is the + # one registered op with no ONNX-published oracle at all — so the expected values come + # from onnxruntime, the second oracle the compiler's parity test already stands on, run + # on the same node at the newest opset onnxruntime serves it at. + ("", "MaxRoiPool"): Sweep( + Kind.POINTWISE, + _MAX_ROI_POOL_VARIANTS, + equivalent_model=_at_the_oracle_opset, + oracle=_onnxruntime_outputs, + ), + # TopK compares rather than accumulates — no summation order changes what it selects — + # so it sweeps every special value its dtype has, NaN included. `k` is the extent of the + # result's axis, which makes it configuration the model has to carry. + ("", "TopK"): Sweep( + Kind.POINTWISE, + _TOP_K_VARIANTS, + operand_types={1: TensorProto.INT64}, + constant_operands=(1,), + ), + # Attention sums a head's worth of products per score and a whole key row per output + # element, so it accumulates for the reason Gemm does; the special values its mask paths + # turn on -- the -inf that masks a column out, and the NaN a boolean mask under + # `is_causal` poisons a row with -- are pinned by the variants that need them rather than + # drawn. Only revision 24 is generated: the evaluator is version-faithful for the newest + # revision alone, and 23 rests on the 63 corpus models that import it. Its key lengths + # are int64 by ONNX's own type constraint whatever the tensors hold. + ("", "Attention"): Sweep( + Kind.ACCUMULATING, + _ATTENTION_VARIANTS, + operand_types={6: TensorProto.INT64}, + ), + # RotaryEmbedding computes two products and one sum per rotated pair, in the order the + # reference writes them, so the float edges are compared rather than left out: arithmetic, + # not accumulation. Its positions index the caches, so they are pinned per variant. + ("", "RotaryEmbedding"): Sweep( + Kind.ARITHMETIC, + _ROTARY_VARIANTS, + operand_types={3: TensorProto.INT64}, + ), + # TfIdfVectorizer compares each token against a pool and counts the matches, so no value + # it reads is computed with and every dtype extreme goes into the one variant that draws. + ("", "TfIdfVectorizer"): Sweep( + Kind.POINTWISE, _TFIDF_VARIANTS, equivalent_model=_with_unit_weights + ), + # The ONNX-ML preprocessing ops. The three whose schema declares a `tensor(float)` result + # take the oracle on the model that casts their input to that float, for the reason + # `_with_float_inputs` records; the rest are compared against the evaluator directly. + # + # Scaler subtracts and multiplies once per element, so the float edges are compared rather + # than left out; no integer operand can overflow, since every one of them is converted to + # float before anything is computed with it. + (ML_DOMAIN, "Scaler"): Sweep( + Kind.POINTWISE, _SCALER_VARIANTS, equivalent_model=_with_float_inputs + ), + # Normalizer sums a row before dividing by it, which is the accumulation Gemm's sweep + # keeps the extremes out of; the rows that carry the special values are pinned instead. + (ML_DOMAIN, "Normalizer"): Sweep( + Kind.ACCUMULATING, _NORMALIZER_VARIANTS, equivalent_model=_with_float_inputs + ), + # FeatureVectorizer copies and pads, so nothing it computes can leave a dtype's range. + (ML_DOMAIN, "FeatureVectorizer"): Sweep( + Kind.POINTWISE, _VECTORIZER_VARIANTS, equivalent_model=_with_float_inputs + ), + **{ + (ML_DOMAIN, op_type): Sweep(Kind.POINTWISE, variants) + for op_type, variants in ( + ("Binarizer", _BINARIZER_VARIANTS), + ("Imputer", _IMPUTER_VARIANTS), + ("LabelEncoder", _LABEL_ENCODER_VARIANTS), + ("OneHotEncoder", _ONE_HOT_ENCODER_VARIANTS), + ) + }, + # The index operand is int64 by ONNX's own type constraint whatever the data holds. + (ML_DOMAIN, "ArrayFeatureExtractor"): Sweep( + Kind.POINTWISE, _EXTRACTOR_VARIANTS, operand_types={1: TensorProto.INT64} + ), + # The tree ensembles. Nothing they compute is a function of the input's magnitude -- a + # feature is compared against a split and the scores are sums of the weights the + # attributes carry -- so the whole special-value list goes into `X`, NaN included, where + # the missing-value flag decides which branch it takes. + (ML_DOMAIN, "TreeEnsembleRegressor"): Sweep( + Kind.POINTWISE, + _regressor_variants(), + equivalent_model=_as_transformed_scores, + ), + (ML_DOMAIN, "TreeEnsembleClassifier"): Sweep( + Kind.POINTWISE, _classifier_variants() + ), + (ML_DOMAIN, "TreeEnsemble"): Sweep( + Kind.POINTWISE, + _tree_ensemble_variants(), + equivalent_model=_as_transformed_scores, + ), + # The support vector machines and the linear models, which score in float32 whatever they + # are handed and are therefore compared on the model whose operands already are that + # float. The two regressors' references raise for any transform but `NONE`, so theirs is + # spelled out as the standard-domain op it is defined to be; the classifiers' references + # apply the transforms themselves — including the one case where a row is deliberately + # returned without one — and are compared directly. + **{ + (ML_DOMAIN, op_type): Sweep( + Kind.ACCUMULATING, variants, equivalent_model=_as_float_predictor + ) + for op_type, variants in ( + ("LinearRegressor", _linear_regressor_variants()), + ("SVMRegressor", _svm_regressor_variants()), + ) + }, + **{ + (ML_DOMAIN, op_type): Sweep( + Kind.ACCUMULATING, variants, equivalent_model=_with_float_inputs + ) + for op_type, variants in ( + ("LinearClassifier", _linear_classifier_variants()), + ("SVMClassifier", _svm_classifier_variants()), + ) + }, +} + + +@dataclass(frozen=True) +class Case: + """One generated model: an op at one revision, one element type, one variant.""" + + domain: str + op_type: str + version: int + elem_type: int + kind: Kind + variant: Variant + + @property + def dtype(self) -> Any: + return np.dtype(numpy_dtype_name(self.elem_type)) + + @property + def id(self) -> str: + return f"{self.op_type}-{self.version}-{self.dtype.name}-{self.variant.label}" + + def __str__(self) -> str: + shapes = ", ".join( + "omitted" if shape is None else str(list(shape)) + for shape in self.variant.shapes + ) + attributes = ( + ", ".join( + f"{name}={value}" + for name, value in sorted(self.variant.attributes.items()) + ) + or "none" + ) + return ( + f"`{self.op_type}` (domain `{display_domain(self.domain)}`) at opset version " + f"{self.version}, dtype `{self.dtype.name}`, variant `{self.variant.label}`, " + f"operand shapes {shapes}, attributes {attributes}, seed {SEED}" + ) + + +def _faithful_revisions(domain: str, op_type: str) -> tuple[int, ...]: + """Registered revisions of the op the reference evaluator is a valid oracle for.""" + return tuple( + version + for version in KERNELS.registered_versions(domain, op_type) + if evaluator_is_version_faithful(domain, op_type, version) + ) + + +@cache +def _cases() -> tuple[Case, ...]: + cases: list[Case] = [] + for domain, op_type in KERNELS.registered_ops(): + sweep = SWEEP.get((domain, op_type)) + if sweep is None: + # The acceptance-rule test reports this; skipping keeps that one failure from + # multiplying into a collection error per missing case. + continue + for version in _faithful_revisions(domain, op_type): + schema = get_schema(op_type, version, domain) + for elem_type in _swept_element_types(schema, sweep.type_operand): + cases.extend( + Case(domain, op_type, version, elem_type, sweep.kind, variant) + for variant in sweep.variants + if _variant_applies(schema, version, elem_type, variant) + ) + return tuple(cases) + + +def _variant_applies( + schema: OpSchema, version: int, elem_type: int, variant: Variant +) -> bool: + if variant.versions is not None and version not in variant.versions: + return False + if variant.elem_types is not None and elem_type not in variant.elem_types: + return False + return _schema_takes(schema, variant) + + +def _swept_element_types(schema: OpSchema, operand: int) -> tuple[int, ...]: + """Element types the schema allows for the sweep's type operand and the compiler supports. + + Read off the schema's type constraints rather than listed by hand, so a dtype ONNX adds + to an op is swept from the moment the installed package defines it. + """ + allowed = _allowed_type_strings(schema, operand) + return tuple( + elem_type for elem_type in sorted(C_TYPES) if _type_string(elem_type) in allowed + ) + + +def _allowed_type_strings(schema: OpSchema, operand: int) -> frozenset[str]: + type_str = schema.inputs[operand].type_str + for constraint in schema.type_constraints: + if constraint.type_param_str == type_str: + return frozenset(constraint.allowed_type_strs) + return frozenset({type_str}) + + +def _type_string(elem_type: int) -> str: + return f"tensor({TensorProto.DataType.Name(elem_type).lower()})" + + +def _schema_takes(schema: OpSchema, variant: Variant) -> bool: + """Whether the schema at this version takes exactly the operands the variant declares.""" + variadic = ( + schema.inputs + and schema.inputs[-1].option == OpSchema.FormalParameterOption.Variadic + ) + if len(variant.shapes) > len(schema.inputs) and not variadic: + return False + provided = { + index for index, shape in enumerate(variant.shapes) if shape is not None + } + return all( + index in provided + for index, formal in enumerate(schema.inputs) + if formal.option == OpSchema.FormalParameterOption.Single + ) + + +# -------------------------------------------------------------------------------------- +# Generating a case: the model, and the values fed to it +# -------------------------------------------------------------------------------------- + + +def _model(case: Case) -> ModelProto: + """A single-node model importing exactly the opset the case's kernel revision claims. + + The output carries no declared type: shape inference derives it, so the generated model + states nothing about the op's result that could disagree with the oracle. + """ + names = [ + "" if shape is None else f"in{index}" + for index, shape in enumerate(case.variant.shapes) + ] + while names and not names[-1]: + names.pop() + results = [f"out{index}" for index in range(case.variant.outputs)] + node = helper.make_node( + case.op_type, + names, + results, + name="node", + domain=case.domain, + **dict(case.variant.attributes), + ) + constants = SWEEP[(case.domain, case.op_type)].constant_operands + graph = helper.make_graph( + [node], + "sweep", + [ + helper.make_tensor_value_info(name, _operand_type(case, index), list(shape)) + for index, (name, shape) in enumerate(zip(names, case.variant.shapes)) + if shape is not None and index not in constants + ], + [helper.make_empty_tensor_value_info(name) for name in results], + initializer=[ + numpy_helper.from_array(_operand(case, index, shape), name) + for index, (name, shape) in enumerate(zip(names, case.variant.shapes)) + if shape is not None and index in constants + ], + ) + return helper.make_model( + graph, opset_imports=[helper.make_opsetid(case.domain, case.version)] + ) + + +def _operand_type(case: Case, index: int) -> int: + """The element type of operand `index`: the case's, unless the variant or sweep pins it.""" + pinned = SWEEP[(case.domain, case.op_type)].operand_types.get(index, case.elem_type) + return case.variant.operand_types.get(index, pinned) + + +def _feeds(case: Case) -> dict[str, Any]: + constants = SWEEP[(case.domain, case.op_type)].constant_operands + return { + f"in{index}": _operand(case, index, shape) + for index, shape in enumerate(case.variant.shapes) + if shape is not None and index not in constants + } + + +def _operand(case: Case, index: int, shape: tuple[int, ...]) -> Any: + """What operand `index` holds: the value the variant pins, or a seeded draw.""" + sweep = SWEEP[(case.domain, case.op_type)] + elem_type = _operand_type(case, index) + pinned = case.variant.values.get(index) + if pinned is None: + domains = {**sweep.operand_domains, **case.variant.domains} + return _values(shape, elem_type, case.kind, index, domains.get(index)) + dtype = numpy_dtype_name(elem_type) + if isinstance(pinned, (int, float)): + return np.full(shape, pinned, dtype) + return np.array(pinned, dtype).reshape(shape) + + +def _values( + shape: tuple[int, ...], + elem_type: int, + kind: Kind, + operand: int, + domain: Domain | None = None, +) -> Any: + """The operand's values: its dtype's special values first, then seeded random draws. + + The specials are permuted per operand, so a binary op sees them paired up differently + (NaN against a finite value, +Inf against -Inf) rather than always against themselves. + """ + generator = np.random.default_rng([SEED, operand]) + dtype = np.dtype(numpy_dtype_name(elem_type)) + size = math.prod(shape) + specials = generator.permutation(_special_values(dtype, kind))[:size] + filler = _random_values(generator, size - len(specials), dtype, kind) + values = np.concatenate([specials, filler]).astype(dtype) + return _restrict(values, domain).reshape(shape) + + +def _restrict(values: Any, domain: Domain | None) -> Any: + """Move an operand into the range the op is defined over.""" + if domain is None: + return values + if domain in (Domain.CONVERTIBLE, Domain.CONVERTIBLE_UNSIGNED): + if values.dtype.kind != "f": + return values + # Every integer type of the target's signedness holds this range, so the clipped + # values convert to any of them; the fractions, the signed zeros and the subnormals + # the draw carries survive it. + smallest = -8 if domain is Domain.CONVERTIBLE else 0 + return np.clip(np.nan_to_num(values), smallest, 8) + if domain is Domain.NONNEGATIVE: + return np.abs(values) + if values.dtype.kind not in "iu": + return values + if domain is Domain.NONZERO: + return np.where(values == 0, values.dtype.type(1), values) + if domain is Domain.SMALL_FACTOR: + # `% 3 - 1` on an unsigned dtype would wrap a zero into the dtype's maximum. + return values % 2 if values.dtype.kind == "u" else values % 3 - 1 + return np.abs(values) % 4 + + +def _special_values(dtype: Any, kind: Kind) -> Any: + if kind is Kind.ACCUMULATING: + return np.empty(0, dtype=dtype) + if dtype.kind == "f": + info = np.finfo(dtype) + return np.array( + [ + 0.0, + -0.0, + 1.0, + -1.0, + np.nan, + np.inf, + -np.inf, + info.max, + -info.max, + info.tiny, + -info.tiny, + info.smallest_subnormal, + -info.smallest_subnormal, + ], + dtype=dtype, + ) + if dtype == np.bool_: + return np.array([False, True]) + info = np.iinfo(dtype) + values = [0, 1, 2] if info.min == 0 else [0, 1, -1] + if kind is Kind.POINTWISE: + values += [info.min, info.max, info.min + 1, info.max - 1] + return np.array(values, dtype=dtype) + + +def _random_values(generator: Any, count: int, dtype: Any, kind: Kind) -> Any: + if count <= 0: + return np.empty(0, dtype=dtype) + if dtype.kind == "f": + return generator.normal(size=count).astype(dtype) + if dtype == np.bool_: + return generator.integers(0, 2, size=count).astype(dtype) + info = np.iinfo(dtype) + low, high = info.min, info.max + if kind is not Kind.POINTWISE: + # Bounded so that neither the sum nor the product of two draws can leave this + # dtype's range. + limit = min(100, math.isqrt(info.max)) + low, high = max(low, -limit), min(high, limit) + return generator.integers(low, high, size=count, endpoint=True, dtype=dtype) + + +# -------------------------------------------------------------------------------------- +# Running a case and comparing it against the evaluator +# -------------------------------------------------------------------------------------- + + +def _execute(case: Case, directory: Path) -> tuple[list[Any], list[Any]]: + """Compile, build and run the case, and run the same model through the evaluator. + + The same model, unless the op's sweep names an equivalent one: an op the evaluator cannot + be trusted on is run as the op ONNX defines it to be equal to, on the same operands. An + op the evaluator does not implement at all takes the oracle its sweep names instead. + """ + model = _model(case) + feeds = _feeds(case) + sweep = SWEEP[(case.domain, case.op_type)] + oracle = model if sweep.equivalent_model is None else sweep.equivalent_model(case) + with np.errstate(all="ignore"): + expected = ( + list(ReferenceEvaluator(oracle).run(None, feeds)) + if sweep.oracle is None + else sweep.oracle(oracle, feeds) + ) + outputs = compile_onnx(model, directory).load().run(feeds) + return [outputs[entry.name] for entry in model.graph.output], expected + + +def _assert_matches( + case: Case, outputs: Sequence[Any], expected: Sequence[Any] +) -> None: + try: + Runner.assert_similar_outputs(expected, outputs, rtol=RTOL, atol=ATOL) + for got, want in zip(outputs, expected): + _assert_zero_signs_match(case, got, want) + except AssertionError as error: + raise AssertionError( + f"{case} diverges from the ONNX reference evaluator.\n{error}" + ) from None + + +def _assert_zero_signs_match(case: Case, got: Any, want: Any) -> None: + """-0.0 and 0.0 are `allclose`, so the sign of every zero is compared separately. + + A signed zero is one of the values the sweep feeds in, and a pointwise kernel applies + the same IEEE operation to the same operand as the reference, so it must come out with + the same sign. Accumulating kernels are exempt: 0 + (-0) is +0, which makes a sum's zero + sign a function of the summation order the spec does not fix. + """ + if case.kind is Kind.ACCUMULATING or want.dtype.kind != "f": + return + zeros = want == 0 + np.testing.assert_array_equal( + np.signbit(got[zeros]), + np.signbit(want[zeros]), + err_msg="the sign of a zero differs from the reference", + ) + + +@pytest.mark.parametrize("case", _cases(), ids=lambda case: case.id) +def test_the_kernel_matches_the_reference_evaluator(case, tmp_path): + _assert_matches(case, *_execute(case, tmp_path)) + + +# -------------------------------------------------------------------------------------- +# The suite's own teeth +# -------------------------------------------------------------------------------------- + + +def _relu_that_drops_nan(context: NodeContext) -> NodeEmission: + """A Relu that is right everywhere except the edge: `x > 0 ? x : 0` sends NaN to zero.""" + source = context.require_input(0) + result = context.require_output(0) + element = c_type(result.elem_type) + name = f"{context.prefix}_diverging_relu_{element}" + definition = "\n".join( + [ + f"static void {name}({element}* out, const {element}* in, size_t count)", + "{", + " size_t index;", + " for (index = 0; index < count; ++index) {", + f" out[index] = in[index] > 0 ? in[index] : ({element})0;", + " }", + "}", + ] + ) + return NodeEmission( + functions=(CFunction(name, definition),), + statements=(f"{name}({result.expr}, {source.expr}, {result.elem_count}u);",), + ) + + +def test_a_kernel_that_diverges_on_an_edge_input_is_reported(tmp_path, monkeypatch): + """Divergence on NaN alone still fails, and the report names what to reproduce it with.""" + case = Case( + domain="", + op_type="Relu", + version=14, + elem_type=TensorProto.FLOAT, + kind=Kind.POINTWISE, + variant=Variant("wide", ((4, 8),)), + ) + select = KERNELS.select + monkeypatch.setattr( + KERNELS, + "select", + lambda domain, op_type, version: ( + KernelSpec(domain, op_type, version, _relu_that_drops_nan) + if op_type == "Relu" + else select(domain, op_type, version) + ), + ) + + with pytest.raises(AssertionError) as error: + _assert_matches(case, *_execute(case, tmp_path)) + + message = str(error.value) + assert "`Relu`" in message + assert "opset version 14" in message + assert "float32" in message + assert str(SEED) in message + + +def test_a_kernel_that_loses_a_zero_sign_is_reported(tmp_path): + """The comparison sees -0.0 against 0.0, which `assert_allclose` alone would not.""" + case = Case( + domain="", + op_type="Relu", + version=14, + elem_type=TensorProto.FLOAT, + kind=Kind.POINTWISE, + variant=Variant("wide", ((4, 8),)), + ) + outputs, expected = _execute(case, tmp_path) + flipped = [np.where(value == 0, np.float32(-0.0), value) for value in expected] + + with pytest.raises(AssertionError, match="sign of a zero"): + _assert_matches(case, outputs, flipped) + + +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +def test_the_generator_feeds_every_float_edge(dtype): + values = _values((4, 8), _elem_type(dtype), Kind.POINTWISE, 0) + info = np.finfo(dtype) + + assert np.isnan(values).any() + assert (values == np.inf).any() and (values == -np.inf).any() + assert (np.signbit(values) & (values == 0)).any() + assert (~np.signbit(values) & (values == 0)).any() + assert (values == info.max).any() and (values == -info.max).any() + assert (np.abs(values) == info.smallest_subnormal).any() + assert (np.abs(values) == info.tiny).any() + + +@pytest.mark.parametrize("dtype", ["int8", "int64", "uint8", "uint64"]) +def test_the_generator_feeds_the_integer_extremes(dtype): + values = _values((4, 8), _elem_type(dtype), Kind.POINTWISE, 0) + info = np.iinfo(dtype) + + assert (values == info.min).any() + assert (values == info.max).any() + + +@pytest.mark.parametrize("dtype", ["int8", "uint8", "int64"]) +def test_arithmetic_operands_cannot_overflow_the_dtype(dtype): + """Integer overflow is undefined in C and undefined by ONNX, so it is never generated.""" + values = _values((8, 8), _elem_type(dtype), Kind.ARITHMETIC, 0).astype(np.int64) + info = np.iinfo(dtype) + + assert int(np.abs(values).max()) ** 2 <= info.max + assert int(np.abs(values).max()) * 2 <= info.max + + +def test_accumulating_operands_are_finite(): + values = _values((4, 8), TensorProto.FLOAT, Kind.ACCUMULATING, 0) + + assert np.isfinite(values).all() + + +def test_the_generator_reproduces_from_the_seed_and_varies_by_operand(): + """A reported case has to reproduce, and a binary op must not see two identical sides.""" + first = _values((4, 8), TensorProto.FLOAT, Kind.POINTWISE, 0) + again = _values((4, 8), TensorProto.FLOAT, Kind.POINTWISE, 0) + second = _values((4, 8), TensorProto.FLOAT, Kind.POINTWISE, 1) + + assert np.array_equal(first, again, equal_nan=True) + assert not np.array_equal(first, second, equal_nan=True) + + +def test_zero_element_shapes_generate_empty_operands(): + values = _values((0, 3), TensorProto.FLOAT, Kind.POINTWISE, 0) + + assert values.shape == (0, 3) + assert values.size == 0 + + +def test_every_operand_of_every_swept_op_is_fed_the_whole_special_set(): + """Specials are sliced to the operand's element count, so a family of shapes that are + all smaller than the list feeds only part of it -- and which part is an accident of the + seed. Each operand needs one variant at least as wide as the list, or the dtype edges + the sweep exists to cover reach one side of a binary op and not the other.""" + problems = [] + for (domain, op_type), sweep in sorted(SWEEP.items()): + # float carries the longest list, so a variant wide enough for it is wide enough + # for every dtype the op takes. + required = len(_special_values(np.dtype("float64"), sweep.kind)) + arity = max(len(variant.shapes) for variant in sweep.variants) + for operand in range(arity): + drawn = [ + variant.shapes[operand] + for variant in sweep.variants + if operand < len(variant.shapes) + and variant.shapes[operand] is not None + and operand not in variant.values + ] + # An operand no variant both draws and shapes — Clip's bounds, Dropout's ratio — + # has no room for the list to begin with: it is a parameter, and the values that + # matter for it are the ones its variants pin. + if not any(drawn): + continue + widest = max(math.prod(shape) for shape in drawn) + if widest < required: + problems.append( + f"`{op_type}` (domain `{display_domain(domain)}`) operand {operand}: " + f"its widest variant holds {widest} elements, short of the " + f"{required} special values, so the rest are never fed to it." + ) + + assert not problems, "\n".join(problems) + + +def _elem_type(dtype: str) -> int: + return helper.np_dtype_to_tensor_dtype(np.dtype(dtype)) + + +# -------------------------------------------------------------------------------------- +# The acceptance rule +# -------------------------------------------------------------------------------------- + + +def test_every_registered_op_is_swept_here(): + """Half the acceptance rule: a kernel with no differential coverage is not implemented.""" + covered = {(case.domain, case.op_type) for case in _cases()} + missing = [ + f"`{op_type}` (domain `{display_domain(domain)}`)" + for domain, op_type in KERNELS.registered_ops() + if (domain, op_type) not in covered + ] + + assert not missing, ( + f"The kernel registry serves {', '.join(missing)}, which this sweep executes no " + "case for; add the op to SWEEP with the attribute combinations its kernel reads." + ) + + +def test_the_sweep_claims_exactly_the_revisions_the_evaluator_can_vouch_for(): + """The oracle-validity restriction, in both directions. + + Sweeping a revision the evaluator is not faithful for would compare a kernel against the + wrong semantics; dropping one it is faithful for would quietly lose coverage. + """ + swept = {(case.domain, case.op_type, case.version) for case in _cases()} + provable = { + (domain, op_type, version) + for domain, op_type in KERNELS.registered_ops() + if (domain, op_type) in SWEEP + for version in _faithful_revisions(domain, op_type) + } + + assert swept == provable + + +def _signature(schema: OpSchema) -> Any: + """Everything about an op's interface a revision could have changed.""" + return ( + schema.doc, + [(formal.name, formal.type_str, formal.option) for formal in schema.inputs], + [(formal.name, formal.type_str, formal.option) for formal in schema.outputs], + { + name: (attribute.type, attribute.required, str(attribute.default_value)) + for name, attribute in schema.attributes.items() + }, + ) + + +def test_the_maxroipool_oracle_runs_the_same_op(): + """The one sweep whose oracle is handed a revision other than the kernel's own. + + onnxruntime implements `MaxRoiPool` up to opset 21 and the kernel is registered at the + revision after it, so running the oracle at 21 only proves anything if ONNX changed + nothing between the two but the element types it accepts. That is read off the schemas + rather than taken on trust — and the one type the newer revision adds is one this + compiler supports at neither. + """ + (claimed,) = KERNELS.registered_versions("", "MaxRoiPool") + newer = get_schema("MaxRoiPool", claimed, "") + older = get_schema("MaxRoiPool", _MAX_ROI_POOL_ORACLE_VERSION, "") + + assert older.since_version < newer.since_version + assert _signature(newer) == _signature(older) + added = _allowed_type_strings(newer, 0) - _allowed_type_strings(older, 0) + assert added == {"tensor(bfloat16)"} + assert not _allowed_type_strings(older, 0) - _allowed_type_strings(newer, 0) + assert not added & {_type_string(elem_type) for elem_type in C_TYPES} + + +def test_the_two_scatter_revisions_are_one_op(): + """The other op claimed at a revision this sweep does not run, and why that is sound. + + Scatter is registered at 9 and at 11, the revision that deprecated it. Only 11 is swept — + the evaluator is version-faithful for it — while the corpus's own Scatter tests import + opset 10, which selects 9. Claiming both from one generator is only sound if ONNX changed + nothing but the deprecation between them, which is read off the schemas rather than taken + on trust; and the deprecation notice is what points at the op this sweep's oracle runs. + """ + assert KERNELS.registered_versions("", "Scatter") == [9, 11] + older = get_schema("Scatter", 9, "") + newer = get_schema("Scatter", 11, "") + + assert not older.deprecated and newer.deprecated + # Everything but the document, which gained the notice quoted below. + assert _signature(older)[1:] == _signature(newer)[1:] + assert _allowed_type_strings(older, 0) == _allowed_type_strings(newer, 0) + assert "Please use ScatterElements" in newer.doc + + +@pytest.mark.parametrize("op_type", ["TreeEnsembleRegressor", "TreeEnsembleClassifier"]) +def test_the_legacy_ensemble_revisions_are_one_op(op_type): + """The remaining ops claimed at revisions this sweep does not run, and why that is sound. + + Both are registered at 1, 3 and 5 while only 5 -- the revision that deprecated them -- + is swept, the evaluator being version-faithful for no earlier one. Opset 1 is what + scikit-learn's own converter emits, which is where the compiler meets these ops in + practice and what the parity tests in `test_extra_compiler_trees.py` run at, with + onnxruntime as their oracle. + + Claiming all three from one generator is only sound if the revisions differ in nothing + the emitted code reads. That is read off the schemas rather than taken on trust: 3 added + the `*_as_tensor` attribute families, which the compiler refuses outright, and 5 added a + deprecation notice. + """ + assert KERNELS.registered_versions(ML_DOMAIN, op_type) == [1, 3, 5] + schemas = { + version: get_schema(op_type, version, ML_DOMAIN) for version in (1, 3, 5) + } + + assert not schemas[1].deprecated and not schemas[3].deprecated + assert schemas[5].deprecated + # The documents differ by those two notices alone, and the interfaces not at all. + assert _signature(schemas[3])[1:3] == _signature(schemas[1])[1:3] + assert _signature(schemas[5])[1:] == _signature(schemas[3])[1:] + assert _allowed_type_strings(schemas[1], 0) == _allowed_type_strings(schemas[5], 0) + added = set(schemas[3].attributes) - set(schemas[1].attributes) + assert added and all(name.endswith("_as_tensor") for name in added) + assert { + name: attribute + for name, attribute in _signature(schemas[3])[3].items() + if name not in added + } == _signature(schemas[1])[3] + + +@pytest.mark.skipif( + not onnx.__version__.startswith(f"{PINNED_ONNX}."), + reason=( + f"the conformance corpus and pass list are pinned to onnx {PINNED_ONNX}.*, " + f"but onnx {onnx.__version__} is installed" + ), +) +def test_every_registered_op_passes_the_backend_suite_too(): + """The other half: an op no ratcheted corpus test exercises is not implemented either. + + Unless the corpus holds no test of the op that could ever run: every node test for + `Expand`, `Tile`, `Pad`, `OneHot` and `TopK` hands the op the repeats, shape, pads, + depth or `k` that decide the shape of its result as a *run-time input*, which makes + those models uncompilable whatever the kernel does, and every one of them is ledgered + for exactly that. Or no test of the op at all, as for `GlobalLpPool`. An op in either + position has no backend evidence to offer either way and rests on the sweep above. Both + exemptions are derived from the corpus and the ledger rather than listed, so an op leaves + them the moment a test of it exists and compiles -- at which point the pass list has to + cover it. + """ + exercised, ledgered_only, mentioned = _corpus_op_types() + missing = [ + f"`{op_type}` (domain `{display_domain(domain)}`)" + for domain, op_type in KERNELS.registered_ops() + if (domain, op_type) in mentioned + and (domain, op_type) not in exercised | ledgered_only + ] + + assert not missing, ( + f"The kernel registry serves {', '.join(missing)}, which no test in " + f"`{RATCHET_PATH.name}` exercises; the backend conformance suite has to pass for " + "an op before it counts as implemented." + ) + + +@pytest.mark.skipif( + not onnx.__version__.startswith(f"{PINNED_ONNX}."), + reason=( + f"the conformance corpus and pass list are pinned to onnx {PINNED_ONNX}.*, " + f"but onnx {onnx.__version__} is installed" + ), +) +def test_the_ops_resting_on_the_sweep_alone_are_spelled_out(): + """Which registered ops the backend suite has no passable test for, written down. + + The exemptions above are derived, so nothing would otherwise show when they grew. These + lists are the record: an op joins the first only because every corpus test of it is + ledgered and the second only because the corpus has no test of it at all, and leaves as + soon as a test of it compiles -- at which point the pass list has to cover the op and + this expectation shrinks in the same change. + + `AffineGrid`, `Col2Im` and `STFT` join the first for the same reason the rest of it is + there: every node test of them hands the op the extents that decide the shape of its + result as a run-time input — for `STFT`, the frame step every frame of the result is one + of, and the frame length each of those transforms. `LabelEncoder` joins it for a reason + of its own: every node test of it maps to or from a string tensor, which is a run-time + string whatever the kernel does, and all four are ledgered as exactly that. + + The second list is where the ONNX-ML preprocessing ops, the two legacy tree ensembles and + the four support-vector and linear predictors sit, alongside `GlobalLpPool` and + `MaxRoiPool`: the corpus carries a node test for three of the fifteen ONNX-ML ops this + compiler serves and none at all for the rest, which is the thinness their own targeted + tests make up for. `GlobalLpPool` and `MaxRoiPool` are the two ops ONNX ships neither a + node test nor a reference implementation for, and their sweeps above are the whole of what + covers them: `GlobalLpPool` against the LpPool its own schema defines it to equal, and + `MaxRoiPool` against onnxruntime. + """ + exercised, ledgered_only, mentioned = _corpus_op_types() + + assert sorted(set(KERNELS.registered_ops()) & ledgered_only) == [ + ("", "AffineGrid"), + ("", "Col2Im"), + ("", "Expand"), + ("", "OneHot"), + ("", "Pad"), + ("", "STFT"), + ("", "Tile"), + ("", "TopK"), + (ML_DOMAIN, "LabelEncoder"), + ] + assert sorted(set(KERNELS.registered_ops()) - mentioned) == [ + ("", "GlobalLpPool"), + ("", "MaxRoiPool"), + (ML_DOMAIN, "FeatureVectorizer"), + (ML_DOMAIN, "Imputer"), + (ML_DOMAIN, "LinearClassifier"), + (ML_DOMAIN, "LinearRegressor"), + (ML_DOMAIN, "Normalizer"), + (ML_DOMAIN, "OneHotEncoder"), + (ML_DOMAIN, "SVMClassifier"), + (ML_DOMAIN, "SVMRegressor"), + (ML_DOMAIN, "Scaler"), + (ML_DOMAIN, "TreeEnsembleClassifier"), + (ML_DOMAIN, "TreeEnsembleRegressor"), + ] + assert ("", "Add") in exercised + assert (ML_DOMAIN, "Binarizer") in exercised + assert (ML_DOMAIN, "TreeEnsemble") in exercised + + +@cache +def _corpus_op_types() -> tuple[ + frozenset[tuple[str, str]], frozenset[tuple[str, str]], frozenset[tuple[str, str]] +]: + """The ops the ratcheted corpus tests run, the ops only ledgered tests mention, and + every op the corpus names at all.""" + ledger = set(json.loads(LEDGER_PATH.read_text(encoding="utf-8"))) + ratchet = { + line.strip() + for line in RATCHET_PATH.read_text(encoding="utf-8").splitlines() + if line.strip() and not line.lstrip().startswith("#") + } + corpus = {case.name: case for case in load_model_tests(kind="node")} + assert not ratchet - set(corpus), ( + "the pass list names tests the corpus does not have" + ) + + exercised: set[tuple[str, str]] = set() + ledgered: set[tuple[str, str]] = set() + for name, case in corpus.items(): + if name not in ratchet and name not in ledger: + continue + assert case.model_dir is not None, f"the corpus test `{name}` ships no model" + model = onnx.load(Path(case.model_dir) / "model.onnx") + op_types = { + (normalize_domain(node.domain), node.op_type) for node in model.graph.node + } + (exercised if name in ratchet else ledgered).update(op_types) + return ( + frozenset(exercised), + frozenset(ledgered - exercised), + frozenset(exercised | ledgered), + ) diff --git a/src/python/tests/test_extra_compiler_dispatch.py b/src/python/tests/test_extra_compiler_dispatch.py new file mode 100644 index 0000000..2147799 --- /dev/null +++ b/src/python/tests/test_extra_compiler_dispatch.py @@ -0,0 +1,195 @@ +"""Kernel-registry dispatch, the semantic-revision guard, and the unsupported-op error. + +Expected dispatch behaviour is derived from the ONNX schema registry itself, so the +tests stay valid as the installed `onnx` package (and its op revisions) change. +""" + +from __future__ import annotations + +import unittest + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError + +onnx = pytest.importorskip("onnx") +registry_module = pytest.importorskip("fnnx.extras.compilers.c.onnx.registry") + +KernelRegistry = registry_module.KernelRegistry +latest_semantic_revision = registry_module.latest_semantic_revision + +ML_DOMAIN = "ai.onnx.ml" + + +def _revisions(op_type: str, domain: str = "") -> list[int]: + """Opset versions at which ONNX revised `op_type`, oldest first.""" + return sorted( + schema.since_version + for schema in onnx.defs.get_all_schemas_with_history() + if schema.name == op_type and schema.domain == domain + ) + + +ADD_REVISIONS = _revisions("Add") +PREVIOUS_ADD_REVISION, LATEST_ADD_REVISION = ADD_REVISIONS[-2:] + + +class SelectTest(unittest.TestCase): + def setUp(self): + self.registry: KernelRegistry[str] = KernelRegistry() + + def test_selects_highest_version_at_or_below_the_requested_opset(self): + self.registry.register("", "Add", PREVIOUS_ADD_REVISION, "old") + self.registry.register("", "Add", LATEST_ADD_REVISION, "new") + + selected = self.registry.select("", "Add", LATEST_ADD_REVISION) + + self.assertIsNotNone(selected) + self.assertEqual(selected.since_version, LATEST_ADD_REVISION) + self.assertEqual(selected.generator, "new") + + def test_older_kernel_serves_opsets_before_the_next_revision(self): + self.registry.register("", "Add", PREVIOUS_ADD_REVISION, "old") + self.registry.register("", "Add", LATEST_ADD_REVISION, "new") + + selected = self.registry.select("", "Add", LATEST_ADD_REVISION - 1) + + self.assertIsNotNone(selected) + self.assertEqual(selected.since_version, PREVIOUS_ADD_REVISION) + + def test_semantic_revision_guard_rejects_stale_kernel(self): + self.registry.register("", "Add", PREVIOUS_ADD_REVISION, "old") + + self.assertIsNone(self.registry.select("", "Add", LATEST_ADD_REVISION)) + + def test_unregistered_op_selects_nothing(self): + self.assertIsNone(self.registry.select("", "Add", LATEST_ADD_REVISION)) + + def test_kernel_newer_than_the_requested_opset_selects_nothing(self): + self.registry.register("", "Add", LATEST_ADD_REVISION, "new") + + self.assertIsNone(self.registry.select("", "Add", LATEST_ADD_REVISION - 1)) + + def test_domain_alias_is_normalized_on_both_sides(self): + self.registry.register("ai.onnx", "Add", LATEST_ADD_REVISION, "new") + + by_empty = self.registry.select("", "Add", LATEST_ADD_REVISION) + by_alias = self.registry.select("ai.onnx", "Add", LATEST_ADD_REVISION) + + self.assertEqual(by_empty, by_alias) + self.assertIsNotNone(by_empty) + self.assertEqual( + self.registry.registered_versions("", "Add"), [LATEST_ADD_REVISION] + ) + + def test_ml_domain_dispatch(self): + since = _revisions("Scaler", ML_DOMAIN)[-1] + self.registry.register(ML_DOMAIN, "Scaler", since, "scaler") + + selected = self.registry.select(ML_DOMAIN, "Scaler", since) + + self.assertIsNotNone(selected) + self.assertEqual(selected.generator, "scaler") + self.assertIsNone(self.registry.select("", "Scaler", since)) + + def test_registered_versions_are_sorted(self): + for version in reversed(ADD_REVISIONS): + self.registry.register("", "Add", version, f"add{version}") + + self.assertEqual(self.registry.registered_versions("", "Add"), ADD_REVISIONS) + + +class RegisterTest(unittest.TestCase): + def setUp(self): + self.registry: KernelRegistry[str] = KernelRegistry() + + def test_unknown_op_is_rejected(self): + with self.assertRaises(ValueError): + self.registry.register("", "NotAnOnnxOp", 1, "kernel") + + def test_version_before_the_op_existed_is_rejected(self): + introduced = _revisions("Add")[0] + with self.assertRaises(ValueError): + self.registry.register("", "Add", introduced - 1, "kernel") + + def test_duplicate_registration_is_rejected(self): + self.registry.register("", "Add", LATEST_ADD_REVISION, "first") + with self.assertRaises(ValueError): + self.registry.register("ai.onnx", "Add", LATEST_ADD_REVISION, "second") + + +class LatestSemanticRevisionTest(unittest.TestCase): + def test_matches_the_schema_history(self): + for version in range(ADD_REVISIONS[0], LATEST_ADD_REVISION + 2): + expected = max((r for r in ADD_REVISIONS if r <= version), default=None) + self.assertEqual(latest_semantic_revision("", "Add", version), expected) + + def test_unknown_op_has_no_revision(self): + self.assertIsNone(latest_semantic_revision("", "NotAnOnnxOp", 1)) + + +class UnsupportedOpErrorTest(unittest.TestCase): + def setUp(self): + self.registry: KernelRegistry[str] = KernelRegistry() + + def test_names_op_domain_version_and_nearest_supported_version(self): + self.registry.register("", "Add", PREVIOUS_ADD_REVISION, "old") + + error = self.registry.unsupported_op_error("", "Add", LATEST_ADD_REVISION) + + self.assertIsInstance(error, CompileError) + message = str(error) + self.assertIn("Add", message) + self.assertIn("ai.onnx", message) + self.assertIn(str(LATEST_ADD_REVISION), message) + self.assertIn(f"Nearest supported version: {PREVIOUS_ADD_REVISION}", message) + + def test_nearest_supported_version_can_be_above_the_requested_opset(self): + self.registry.register("", "Add", ADD_REVISIONS[0], "oldest") + self.registry.register("", "Add", LATEST_ADD_REVISION, "newest") + + message = str( + self.registry.unsupported_op_error("", "Add", PREVIOUS_ADD_REVISION) + ) + + self.assertIn(f"Nearest supported version: {LATEST_ADD_REVISION}", message) + + def test_nearest_supported_version_can_be_below_the_requested_opset(self): + older_revision, requested = ADD_REVISIONS[1], ADD_REVISIONS[2] + self.registry.register("", "Add", older_revision, "older") + self.registry.register("", "Add", LATEST_ADD_REVISION, "newest") + + message = str(self.registry.unsupported_op_error("", "Add", requested)) + + self.assertIn(f"Nearest supported version: {older_revision}", message) + + def test_reports_when_no_kernel_is_registered(self): + message = str( + self.registry.unsupported_op_error("", "Add", LATEST_ADD_REVISION) + ) + + self.assertIn("Add", message) + self.assertIn("ai.onnx", message) + self.assertIn(str(LATEST_ADD_REVISION), message) + self.assertIn("no kernel is registered", message) + + def test_reports_kernels_that_are_all_newer(self): + self.registry.register("", "Add", LATEST_ADD_REVISION, "new") + + message = str(self.registry.unsupported_op_error("", "Add", ADD_REVISIONS[0])) + + self.assertIn("newer opset version", message) + self.assertIn(f"Nearest supported version: {LATEST_ADD_REVISION}", message) + + def test_names_the_node(self): + message = str( + self.registry.unsupported_op_error( + "", "Add", LATEST_ADD_REVISION, node_name="adder" + ) + ) + + self.assertIn("adder", message) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/python/tests/test_extra_compiler_dispositions.py b/src/python/tests/test_extra_compiler_dispositions.py new file mode 100644 index 0000000..ddc7c91 --- /dev/null +++ b/src/python/tests/test_extra_compiler_dispositions.py @@ -0,0 +1,792 @@ +"""Op-coverage closure: every op of the supported domains is dispositioned, and provably so. + +`conformance/dispositions.json` classifies every op schema the pinned `onnx` package defines +for `ai.onnx` and `ai.onnx.ml` as exactly one of + +* **native-kernel** — a generator in the kernel registry serves it, +* **function-expansion** — no kernel does, and the compiler inlines the function body ONNX + defines for it, +* **folding-or-graph-pass** — no kernel does, and a compiler pass resolves the node away + before dispatch: constant folding, or the ZipMap removal pass, +* **unsupported** — with a reason drawn from the structural part of the conformance + ledger's own closed set of categories. + +Coverage is then a closed property rather than an aspiration: an op the table does not name +— one an `onnx` upgrade adds, say — fails this suite, and so does a table entry the compiler +itself contradicts. Nothing here is taken on the table's word. A kernel claim is checked +against the registry, an expansion claim against the schema's function body, a reason against +the very set the compiler rejects that family of ops from, and every claim to serve an op +without a kernel against evidence that it compiles: a corpus test in the conformance pass +list, or a model this module compiles and runs against the ONNX reference evaluator. + +The table and the ledger have to agree, in both directions: an op the table calls unsupported +cannot appear in a corpus test that passes, and a corpus test excluded as `op-not-implemented` +has to hold an op the table does not claim a kernel or a function body for. +""" + +from __future__ import annotations + +import json +import shutil +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from functools import cache +from pathlib import Path +from typing import Any + +import pytest + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +# The harness refuses to import without numpy, so this covers both dependencies. +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from onnx import ModelProto, TensorProto, helper, numpy_helper # noqa: E402 +from onnx.backend.test.loader import load_model_tests # noqa: E402 +from onnx.backend.test.runner import Runner # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +from fnnx.extras.compilers.c import compile_onnx # noqa: E402 +from fnnx.extras.compilers.c.errors import CompileError # noqa: E402 +from fnnx.extras.compilers.c.onnx.folding import NONDETERMINISTIC_OPS # noqa: E402 +from fnnx.extras.compilers.c.onnx.kernels import KERNELS # noqa: E402 +from fnnx.extras.compilers.c.onnx.loader import ( # noqa: E402 + ML_DOMAIN, + SUPPORTED_DOMAINS, + display_domain, + max_supported_opset, + normalize_domain, +) +from fnnx.extras.compilers.c.onnx.verify import ( # noqa: E402 + CONTROL_FLOW_OPS, + DATA_DEPENDENT_SHAPE_OPS, +) +from fnnx.extras.compilers.c.onnx.zipmap import ZIP_MAP # noqa: E402 + +# The ledger's categories, the corpus and the schema set are all one release's; this module +# reads the conformance suite's own definitions rather than restating them, so the two files +# cannot drift into disagreeing about what a category is. +from test_extra_compiler_conformance import ( # noqa: E402 + CODEC_OPS, + LEDGER_CATEGORIES, + LEDGER_PATH, + PINNED_ONNX, + RATCHET_PATH, +) + +TABLE_PATH = Path(__file__).parent / "conformance" / "dispositions.json" + +DISPOSITIONS = ( + "native-kernel", + "function-expansion", + "folding-or-graph-pass", + "unsupported", +) + +# The ledger categories that cannot be a reason an *op* is unsupported. Two describe a test +# rather than an op: a corpus test is out of scope because of the domain it imports and +# dtype-limited because of the tensors it carries, neither of which is a property of an op the +# supported domains define. The third, `op-not-implemented`, is the milestone category the +# kernel tasks drive down, and admitting it here would reopen as a *disposition* the very +# "not implemented yet" excuse the closed category set exists to refuse: a straggler is to be +# implemented or refused for a structural reason, never parked in the table. What remains is +# exactly the reasons this module re-derives from the compiler or from the op's own schema. +_NOT_OP_REASONS = ("out-of-scope-domain", "unsupported-dtype", "op-not-implemented") +UNSUPPORTED_REASONS = tuple( + category for category in LEDGER_CATEGORIES if category not in _NOT_OP_REASONS +) + +# Where a reason names a family the compiler itself decides membership of, the set it decides +# it from. A table cannot declare that an op is control flow, or a draw: it can only record +# what the compiler already rejects it as. +_REASON_FAMILIES: Mapping[str, frozenset[str]] = { + "control-flow": CONTROL_FLOW_OPS, + "data-dependent-shape": DATA_DEPENDENT_SHAPE_OPS, + "random-op": NONDETERMINISTIC_OPS, + "external-codec": CODEC_OPS, +} + +# The dispositions that claim the compiler serves an op whatever a model does with it, as +# against `folding-or-graph-pass`, which serves the forms a pass can resolve and refuses the +# rest. +_UNCONDITIONAL = ("native-kernel", "function-expansion") + +pytestmark = [ + pytest.mark.skipif( + not onnx.__version__.startswith(f"{PINNED_ONNX}."), + reason=( + f"the disposition table is pinned to onnx {PINNED_ONNX}.*, " + f"but onnx {onnx.__version__} is installed" + ), + ), + pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", + ), +] + + +# -------------------------------------------------------------------------------------- +# The table, the schemas it has to cover, and the corpus it has to agree with +# -------------------------------------------------------------------------------------- + + +def _table() -> dict[tuple[str, str], dict[str, Any]]: + """The checked-in table, keyed by the normalized `(domain, op_type)` the compiler uses.""" + raw = json.loads(TABLE_PATH.read_text(encoding="utf-8")) + return { + (normalize_domain(domain), op_type): entry + for domain, entries in raw.items() + for op_type, entry in entries.items() + } + + +@cache +def _schemas() -> dict[tuple[str, str], Any]: + """Every op the pinned package defines for the supported domains, at its newest schema. + + Enumerated from the package rather than listed, so an op an upgrade adds is one this + module already asks the table about. + """ + names = { + (schema.domain, schema.name) + for schema in onnx.defs.get_all_schemas_with_history() + if schema.domain in SUPPORTED_DOMAINS + } + return { + (domain, op_type): onnx.defs.get_schema( + op_type, max_supported_opset(domain), domain + ) + for domain, op_type in names + } + + +def _defines_a_function(domain: str, op_type: str) -> bool: + """Whether ONNX defines the op as a function body the compiler could inline.""" + schema = _schemas()[(domain, op_type)] + return bool( + schema.has_function # type: ignore[attr-defined] + or schema.has_context_dependent_function # type: ignore[attr-defined] + ) + + +@cache +def _corpus_ops() -> tuple[ + frozenset[tuple[str, str]], dict[str, frozenset[tuple[str, str]]] +]: + """The ops every ratcheted test runs, and the ops of each `op-not-implemented` exclusion.""" + ledger: dict[str, str] = json.loads(LEDGER_PATH.read_text(encoding="utf-8")) + ratchet = { + line.strip() + for line in RATCHET_PATH.read_text(encoding="utf-8").splitlines() + if line.strip() and not line.lstrip().startswith("#") + } + exercised: set[tuple[str, str]] = set() + unimplemented: dict[str, frozenset[tuple[str, str]]] = {} + for case in load_model_tests(kind="node"): + wanted = case.name in ratchet or ledger.get(case.name) == "op-not-implemented" + if not wanted: + continue + assert case.model_dir is not None, ( + f"the corpus test `{case.name}` ships no model" + ) + model = onnx.load(Path(case.model_dir) / "model.onnx") + ops = frozenset( + (normalize_domain(node.domain), node.op_type) for node in _nodes(model) + ) + if case.name in ratchet: + exercised |= ops + else: + unimplemented[case.name] = ops + return frozenset(exercised), unimplemented + + +def _nodes(model: ModelProto) -> Iterator[Any]: + def walk(graph: Any) -> Iterator[Any]: + for node in graph.node: + yield node + for attribute in node.attribute: + if attribute.HasField("g"): + yield from walk(attribute.g) + for subgraph in attribute.graphs: + yield from walk(subgraph) + + yield from walk(model.graph) + + +# -------------------------------------------------------------------------------------- +# What is wrong with a table +# -------------------------------------------------------------------------------------- + + +def _problems(table: Mapping[tuple[str, str], Mapping[str, Any]]) -> list[str]: + """Everything the compiler, the corpus and the ledger contradict in `table`.""" + return [ + *_coverage_problems(table), + *_claim_problems(table), + *_ledger_problems(table), + ] + + +def _coverage_problems(table: Mapping[tuple[str, str], Mapping[str, Any]]) -> list[str]: + problems = [] + for domain, op_type in sorted(set(_schemas()) - set(table)): + problems.append( + f"`{op_type}` (domain `{display_domain(domain)}`) is defined by the installed " + "`onnx` package but has no disposition." + ) + for domain, op_type in sorted(set(table) - set(_schemas())): + problems.append( + f"`{op_type}` (domain `{display_domain(domain)}`) is dispositioned but the " + "installed `onnx` package defines no such op in a supported domain." + ) + return problems + + +def _claim_problems(table: Mapping[tuple[str, str], Mapping[str, Any]]) -> list[str]: + """Every entry against what the compiler itself does with the op.""" + problems = [] + for (domain, op_type), entry in sorted(table.items()): + if (domain, op_type) not in _schemas(): + continue + label = f"`{op_type}` (domain `{display_domain(domain)}`)" + disposition = entry.get("disposition") + served_by_kernel = bool(KERNELS.registered_versions(domain, op_type)) + if disposition not in DISPOSITIONS: + problems.append( + f"{label} is dispositioned `{disposition}`, which is not one of " + f"{', '.join(DISPOSITIONS)}." + ) + continue + if served_by_kernel != (disposition == "native-kernel"): + problems.append( + f"{label} is dispositioned `{disposition}` while the kernel registry " + + ("serves it." if served_by_kernel else "does not serve it.") + ) + if disposition == "function-expansion" and not _defines_a_function( + domain, op_type + ): + problems.append( + f"{label} is dispositioned `function-expansion`, but ONNX defines no " + "function body for it to be expanded into." + ) + if ( + disposition in ("folding-or-graph-pass", "unsupported") + and not str(entry.get("note", "")).strip() + ): + problems.append(f"{label} is dispositioned `{disposition}` with no note.") + if disposition == "unsupported": + problems.extend(_reason_problems(label, domain, op_type, entry)) + return problems + + +def _reason_problems( + label: str, domain: str, op_type: str, entry: Mapping[str, Any] +) -> list[str]: + """An unsupported entry against the closed reason set and the schema's own types.""" + reason = entry.get("reason") + if reason not in UNSUPPORTED_REASONS: + return [ + f"{label} is unsupported for reason `{reason}`, which is not one of " + f"{', '.join(UNSUPPORTED_REASONS)}." + ] + family = _REASON_FAMILIES.get(reason) + if family is not None and op_type not in family: + return [ + f"{label} is unsupported as `{reason}`, but the compiler does not count it " + f"as one: its `{reason}` family is {', '.join(sorted(family))}." + ] + if reason == "non-tensor-io" and not _admits( + domain, op_type, ("seq(", "map(", "optional(") + ): + return [ + f"{label} is unsupported as `non-tensor-io`, but no operand of its schema " + "takes or produces a sequence, a map or an optional." + ] + if reason == "runtime-strings" and not _admits( + domain, op_type, ("tensor(string)",) + ): + return [ + f"{label} is unsupported as `runtime-strings`, but no operand of its schema " + "takes or produces a string tensor." + ] + return [] + + +def _admits(domain: str, op_type: str, kinds: Sequence[str]) -> bool: + """Whether any operand of the op's newest schema admits one of these type forms.""" + schema = _schemas()[(domain, op_type)] + constraints = { + constraint.type_param_str: set(constraint.allowed_type_strs) + for constraint in schema.type_constraints + } + return any( + allowed.startswith(tuple(kinds)) + for formal in (*schema.inputs, *schema.outputs) + for allowed in constraints.get(formal.type_str, {formal.type_str}) + ) + + +def _ledger_problems(table: Mapping[tuple[str, str], Mapping[str, Any]]) -> list[str]: + """The two directions the table and the conformance ledger have to agree in.""" + exercised, unimplemented = _corpus_ops() + problems = [] + for domain, op_type in sorted(exercised): + entry = table.get((domain, op_type), {}) + if entry.get("disposition") == "unsupported": + problems.append( + f"`{op_type}` (domain `{display_domain(domain)}`) is dispositioned " + f"unsupported as `{entry.get('reason')}`, but a corpus test in the " + "conformance pass list compiles and runs it." + ) + for name, ops in sorted(unimplemented.items()): + served = [ + op_type + for domain, op_type in sorted(ops) + if table.get((domain, op_type), {}).get("disposition") in _UNCONDITIONAL + ] + if len(served) == len(ops): + problems.append( + f"`{name}` is ledgered as `op-not-implemented`, but every op it runs " + f"({', '.join(served)}) is dispositioned as served by a kernel or a " + "function body." + ) + return problems + + +# -------------------------------------------------------------------------------------- +# Evidence that an op served without a kernel really compiles +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Proof: + """A model an op's disposition is proven on, and what the oracle is run on. + + `oracle` is the model the reference evaluator computes the expected outputs from, which is + the compiled model itself except where ONNX ships no reference implementation for the op + at all -- `ZipMap`, whose entry says what stands in for it. + """ + + model: ModelProto + feeds: dict[str, Any] + oracle: ModelProto | None = None + + +def _windowed(op_type: str) -> Proof: + """A window of a fixed length: the one operand the graph has to fix for it to compile.""" + return Proof( + _proof_model( + [helper.make_node(op_type, ["size"], ["Y"])], + initializer=[_constant("size", np.array(8, dtype=np.int64))], + ), + {}, + ) + + +def _mel_weight_matrix() -> Proof: + return Proof( + _proof_model( + [ + helper.make_node( + "MelWeightMatrix", + ["bins", "dft_length", "sample_rate", "lower", "upper"], + ["Y"], + ) + ], + initializer=[ + _constant("bins", np.array(4, dtype=np.int64)), + _constant("dft_length", np.array(16, dtype=np.int64)), + _constant("sample_rate", np.array(16000, dtype=np.int64)), + _constant("lower", np.array(0.0, dtype=np.float32)), + _constant("upper", np.array(8000.0, dtype=np.float32)), + ], + ), + {}, + ) + + +def _range() -> Proof: + return Proof( + _proof_model( + [helper.make_node("Range", ["start", "limit", "delta"], ["Y"])], + initializer=[ + _constant("start", np.array(1.0, dtype=np.float32)), + _constant("limit", np.array(6.0, dtype=np.float32)), + _constant("delta", np.array(0.5, dtype=np.float32)), + ], + ), + {}, + ) + + +def _center_crop_pad() -> Proof: + """A crop on one axis and a pad on the other, through the function body ONNX defines.""" + return Proof( + _proof_model( + [helper.make_node("CenterCropPad", ["X", "shape"], ["Y"])], + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 6])], + initializer=[_constant("shape", np.array([2, 8], dtype=np.int64))], + ), + {"X": np.arange(24, dtype=np.float32).reshape(4, 6)}, + ) + + +def _zip_map() -> Proof: + """A scaled tensor keyed into a map output, which the pass removes to leave the tensor. + + ONNX ships no reference implementation for `ZipMap` at all, so the oracle runs the same + graph with the node already gone: what the pass promises is that the tensor `ZipMap` read + reaches the caller in its place, which is exactly what that graph computes. The pairing of + label to column — the part this cannot show — is covered in `test_extra_compiler_ml.py`, + against onnxruntime. + """ + scaler = helper.make_node( + "Scaler", + ["X"], + ["scores"], + domain=ML_DOMAIN, + offset=[1.0, 0.0, -1.0], + scale=[0.5, 0.25, 2.0], + ) + keyed = helper.make_node( + ZIP_MAP, ["scores"], ["Z"], domain=ML_DOMAIN, classlabels_int64s=[7, 9, 11] + ) + maps = helper.make_value_info( + "Z", + helper.make_sequence_type_proto( + helper.make_map_type_proto( + TensorProto.INT64, helper.make_tensor_type_proto(TensorProto.FLOAT, []) + ) + ), + ) + fed = [helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3])] + return Proof( + _proof_model([scaler, keyed], inputs=fed, outputs=[maps], ml=True), + {"X": np.array([[0.2, 0.3, 0.5], [0.1, 0.6, 0.3]], dtype=np.float32)}, + oracle=_proof_model( + [scaler], + inputs=fed, + outputs=[helper.make_empty_tensor_value_info("scores")], + ml=True, + ), + ) + + +def _proof_model( + nodes: Sequence[Any], + *, + inputs: Sequence[Any] = (), + outputs: Sequence[Any] | None = None, + initializer: Sequence[Any] = (), + ml: bool = False, +) -> ModelProto: + """A single-node model at the pinned package's newest opset, typed by inference alone.""" + graph = helper.make_graph( + list(nodes), + "proof", + list(inputs), + list(outputs) + if outputs is not None + else [helper.make_empty_tensor_value_info("Y")], + initializer=list(initializer), + ) + imports = [helper.make_opsetid("", max_supported_opset(""))] + if ml: + imports.append(helper.make_opsetid(ML_DOMAIN, max_supported_opset(ML_DOMAIN))) + return helper.make_model(graph, opset_imports=imports) + + +def _constant(name: str, value: Any) -> Any: + return numpy_helper.from_array(value, name) + + +PROOFS = { + ("", "BlackmanWindow"): lambda: _windowed("BlackmanWindow"), + ("", "CenterCropPad"): _center_crop_pad, + ("", "HammingWindow"): lambda: _windowed("HammingWindow"), + ("", "HannWindow"): lambda: _windowed("HannWindow"), + ("", "MelWeightMatrix"): _mel_weight_matrix, + ("", "Range"): _range, + (ML_DOMAIN, ZIP_MAP): _zip_map, +} + + +def _served_without_a_kernel( + table: Mapping[tuple[str, str], Mapping[str, Any]], +) -> list[tuple[str, str]]: + return sorted( + key + for key, entry in table.items() + if entry.get("disposition") in ("function-expansion", "folding-or-graph-pass") + ) + + +def _unevidenced(table: Mapping[tuple[str, str], Mapping[str, Any]]) -> list[str]: + """Ops claimed to be served without a kernel that neither the corpus nor a proof covers.""" + exercised, _ = _corpus_ops() + return [ + f"`{op_type}` (domain `{display_domain(domain)}`)" + for domain, op_type in _served_without_a_kernel(table) + if (domain, op_type) not in exercised and (domain, op_type) not in PROOFS + ] + + +# -------------------------------------------------------------------------------------- +# The suite +# -------------------------------------------------------------------------------------- + + +def test_every_op_in_the_supported_domains_is_dispositioned(): + problems = _problems(_table()) + + assert not problems, "\n".join(problems) + + +def test_an_op_the_table_does_not_name_fails_the_suite(): + """What an `onnx` upgrade looks like here: a new op is a missing disposition.""" + table = {key: entry for key, entry in _table().items() if key != ("", "Add")} + + problems = _problems(table) + + assert [ + problem + for problem in problems + if "`Add`" in problem and "has no disposition" in problem + ] + + +def test_a_disposition_the_registry_contradicts_fails_the_suite(): + """Both ways round: a kernel claimed for an op with none, and none for an op with one.""" + table = { + **_table(), + ("", "NonZero"): {"disposition": "native-kernel"}, + ("", "Add"): { + "disposition": "unsupported", + "reason": "data-dependent-shape", + "note": "invented", + }, + } + + problems = _problems(table) + + assert [ + problem + for problem in problems + if "`NonZero`" in problem and "does not serve it" in problem + ] + assert [ + problem for problem in problems if "`Add`" in problem and "serves it" in problem + ] + + +def test_an_expansion_claim_without_a_function_body_fails_the_suite(): + table = {**_table(), ("", "NonZero"): {"disposition": "function-expansion"}} + + problems = _problems(table) + + assert [ + problem + for problem in problems + if "`NonZero`" in problem and "no function body" in problem + ] + + +def test_a_reason_outside_the_ledger_categories_fails_the_suite(): + table = { + **_table(), + ("", "NonZero"): { + "disposition": "unsupported", + "reason": "too difficult", + "note": "invented", + }, + } + + problems = _problems(table) + + assert [problem for problem in problems if "which is not one of" in problem] + + +def test_an_op_parked_as_not_yet_implemented_fails_the_suite(): + """The ledger's milestone category is not a disposition: closure admits no backlog.""" + table = { + **_table(), + ("", "NonZero"): { + "disposition": "unsupported", + "reason": "op-not-implemented", + "note": "invented", + }, + } + + problems = _problems(table) + + assert [ + problem + for problem in problems + if "`NonZero`" in problem and "which is not one of" in problem + ] + + +def test_a_reason_the_compiler_does_not_apply_to_the_op_fails_the_suite(): + """A table cannot declare that an op is control flow, or a string op, or a draw.""" + table = { + **_table(), + ("", "NonZero"): { + "disposition": "unsupported", + "reason": "control-flow", + "note": "invented", + }, + (ML_DOMAIN, "CastMap"): { + "disposition": "unsupported", + "reason": "runtime-strings", + "note": "invented", + }, + ("", "Compress"): { + "disposition": "unsupported", + "reason": "non-tensor-io", + "note": "invented", + }, + } + + problems = _problems(table) + + assert [ + problem + for problem in problems + if "`NonZero`" in problem and "does not count it as one" in problem + ] + assert [ + problem + for problem in problems + if "`Compress`" in problem and "no operand of its schema" in problem + ] + # CastMap really does produce a string tensor for one of its output types, so the claim + # that strings reach it is the one the schema cannot refute; the note has to carry the + # rest of the story. + assert not [ + problem + for problem in problems + if "`CastMap`" in problem and "no operand of its schema" in problem + ] + + +def test_an_undocumented_unsupported_op_fails_the_suite(): + table = { + **_table(), + ("", "NonZero"): { + "disposition": "unsupported", + "reason": "data-dependent-shape", + "note": " ", + }, + } + + problems = _problems(table) + + assert [problem for problem in problems if "with no note" in problem] + + +def test_calling_an_op_the_corpus_runs_unsupported_fails_the_suite(): + """The ratchet is the table's second opinion: a passing test disproves `unsupported`.""" + table = { + **_table(), + ("", "Relu"): { + "disposition": "unsupported", + "reason": "data-dependent-shape", + "note": "invented", + }, + } + + problems = _problems(table) + + assert [ + problem + for problem in problems + if "`Relu`" in problem + and "conformance pass list compiles and runs it" in problem + ] + + +def test_an_exclusion_the_table_says_is_served_fails_the_suite(): + """`op-not-implemented` cannot be claimed for a graph of ops the table says are served.""" + _, unimplemented = _corpus_ops() + assert unimplemented, "the ledger excludes nothing as `op-not-implemented`" + name, ops = sorted(unimplemented.items())[0] + table = { + **_table(), + **{key: {"disposition": "native-kernel"} for key in ops}, + } + + problems = _problems(table) + + assert [ + problem + for problem in problems + if f"`{name}`" in problem + and "is dispositioned as served by a kernel" in problem + ] + + +def test_every_op_served_without_a_kernel_has_evidence(): + """No claim to serve an op through a pass or a function body rests on the table alone.""" + unevidenced = _unevidenced(_table()) + + assert not unevidenced, ( + f"{', '.join(unevidenced)} are dispositioned as served without a kernel, which no " + "test in the conformance pass list exercises and no proof model here compiles; add " + "one to PROOFS." + ) + + +def test_an_unproven_claim_to_serve_an_op_fails_the_suite(): + table = { + **_table(), + ("", "NonZero"): { + "disposition": "folding-or-graph-pass", + "note": "invented", + }, + } + + unevidenced = _unevidenced(table) + + assert [entry for entry in unevidenced if "`NonZero`" in entry] + + +@pytest.mark.parametrize( + "domain,op_type", sorted(PROOFS), ids=lambda value: value or "ai.onnx" +) +def test_an_op_served_by_a_pass_or_a_function_body_compiles(domain, op_type, tmp_path): + """The evidence itself: the compiler serves the op, and computes what ONNX says it does. + + Every expected value comes from the ONNX reference evaluator, never from this module. + """ + proof = PROOFS[(domain, op_type)]() + + compiled = compile_onnx(proof.model, tmp_path).load() + outputs = compiled.run(proof.feeds) + + expected = ReferenceEvaluator(proof.oracle or proof.model).run(None, proof.feeds) + Runner.assert_similar_outputs( + list(expected), + [outputs[spec.name] for spec in compiled.outputs], + rtol=1e-3, + atol=1e-7, + ) + + +def test_a_range_below_the_revision_the_evaluator_can_be_vouched_for_is_refused( + tmp_path, +): + """The other half of `Range`'s entry: what the folding pass declines is refused outright. + + ONNX revised the op at opset 27, so at any older one the evaluator folding runs the node + through implements a revision the compiler cannot vouch for. The rule the whole compiler + is built on — never serve semantics nothing can confirm — makes that a compile error + naming the op, not a folded value, and that is what the ledger's remaining + `op-not-implemented` exclusion is. + """ + model = _range().model + del model.opset_import[:] + model.opset_import.append(helper.make_opsetid("", 24)) + + with pytest.raises(CompileError, match="`Range`"): + compile_onnx(model, tmp_path) diff --git a/src/python/tests/test_extra_compiler_emit.py b/src/python/tests/test_extra_compiler_emit.py new file mode 100644 index 0000000..b145343 --- /dev/null +++ b/src/python/tests/test_extra_compiler_emit.py @@ -0,0 +1,823 @@ +"""C emission: the single-header artifact, its metadata, and its determinism. + +What is asserted here is structural — symbols, macros, buffers, byte-identity — and the +values that reach the header are whatever the ONNX reference evaluator folded, never +hand-written. Op semantics belong to the conformance and differential suites. +""" + +from __future__ import annotations + +import json +import math +import re +import shutil +import subprocess +from pathlib import Path +from typing import TYPE_CHECKING + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError + +if TYPE_CHECKING: + import numpy + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +codegen = pytest.importorskip("fnnx.extras.compilers.c.onnx.codegen") +dtypes = pytest.importorskip("fnnx.extras.compilers.c.onnx.dtypes") +emit = pytest.importorskip("fnnx.extras.compilers.c.onnx.emit") +frontend = pytest.importorskip("fnnx.extras.compilers.c.onnx.frontend") +kernels = pytest.importorskip("fnnx.extras.compilers.c.onnx.kernels") +registry = pytest.importorskip("fnnx.extras.compilers.c.onnx.registry") + +from fnnx import __version__ # noqa: E402 +from fnnx.extras.compilers.c import compile_onnx # noqa: E402 +from onnx import TensorProto, helper # noqa: E402 + +OPSET = 21 +STRICT_FLAGS = ("-std=c99", "-Wall", "-Wextra", "-Werror", "-Werror=vla") +C_COMPILERS = [name for name in ("gcc", "clang") if shutil.which(name)] +ALLOCATION_TOKENS = ("malloc", "calloc", "realloc", "free", "alloca") +SEED = 20260725 + + +def _model(nodes, inputs, outputs, *, initializer=(), name="graph", opset=OPSET): + graph = helper.make_graph( + nodes, name, list(inputs), list(outputs), initializer=list(initializer) + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + + +def _tensor(name, elem_type, shape): + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _constant_node(name, array): + return helper.make_node( + "Constant", + [], + [name], + name=f"const_{name}", + value=onnx.numpy_helper.from_array(array, f"{name}_value"), + ) + + +def _alias_model(name="demo"): + """A folded constant output, an aliased input, and an input nothing reads.""" + left = onnx.numpy_helper.from_array(np.array([1.5, 2.5], dtype=np.float32), "left") + right = onnx.numpy_helper.from_array( + np.array([1.0, 4.0], dtype=np.float32), "right" + ) + return _model( + [helper.make_node("Add", ["left", "right"], ["y"], name="add")], + [ + _tensor("x", TensorProto.FLOAT, ["batch", 3]), + _tensor("unread", TensorProto.FLOAT, [2]), + ], + [ + _tensor("y", TensorProto.FLOAT, [2]), + _tensor("x", TensorProto.FLOAT, ["batch", 3]), + ], + initializer=[left, right], + name=name, + ) + + +def _sample_values(elem_type: int) -> numpy.ndarray: + """Special values plus seeded random ones, at `elem_type`.""" + dtype = onnx.helper.tensor_dtype_to_np_dtype(elem_type) + generator = np.random.default_rng(SEED) + if elem_type == TensorProto.BOOL: + return np.array([True, False, True], dtype=dtype) + if elem_type in (TensorProto.FLOAT, TensorProto.DOUBLE): + info = np.finfo(dtype) + special = [ + 0.0, + -0.0, + np.nan, + np.inf, + -np.inf, + info.max, + -info.max, + info.tiny, + info.smallest_subnormal, + info.eps, + ] + random = generator.uniform(-1e6, 1e6, size=8) + return np.array([*special, *random], dtype=dtype) + info = np.iinfo(dtype) + random = generator.integers(info.min, info.max, size=8, endpoint=True, dtype=dtype) + return np.array([info.min, info.max, 0, 1, *random], dtype=dtype) + + +def _parse_c_literal(text: str) -> float | int: + """Read a literal `scalar_literal` produced back into Python.""" + expression = re.sub(r"\b(?:INT64_C|UINT64_C)\((-?\d+)\)", r"\1", text.strip()) + if expression == "NAN": + return math.nan + if expression in ("INFINITY", "-INFINITY"): + return math.inf if expression == "INFINITY" else -math.inf + compound = re.fullmatch(r"\((-\d+) - (\d+)\)", expression) + if compound: + return int(compound.group(1)) - int(compound.group(2)) + expression = re.sub(r"[fu]$", "", expression) + if re.search(r"[.eE]", expression): + return float(expression) + return int(expression) + + +def _same_value(parsed: float | int, expected: float | int, dtype: numpy.dtype) -> bool: + """Whether a literal, read back at its own precision, reproduces the source value.""" + if not np.issubdtype(dtype, np.floating): + return int(parsed) == int(expected) + read = dtype.type(parsed) + if math.isnan(expected): + return bool(np.isnan(read)) + if expected == 0.0: + return read == 0.0 and math.copysign(1.0, float(read)) == math.copysign( + 1.0, expected + ) + return bool(read == expected) + + +def _weight_arrays(header: str) -> dict[str, list[float | int]]: + """Every embedded weight in the header, parsed back into Python values.""" + blocks = re.findall( + r"static const \w+ (\w+)\[\d+\] = \{(.*?)\};", header, flags=re.DOTALL + ) + return { + symbol: [_parse_c_literal(piece) for piece in body.split(",") if piece.strip()] + for symbol, body in blocks + } + + +def _driver(report: dict) -> str: + """A declarations-only unit that calls the entrypoint with correctly sized buffers.""" + entrypoint = report["entrypoint"] + lines = [f'#include "{report["header"]}"', "", "int main(void)", "{"] + arguments = [] + for index, tensor in enumerate([*entrypoint["inputs"], *entrypoint["outputs"]]): + buffer = f"buffer_{index}" + lines.append( + f" static {tensor['c_type']} {buffer}[{max(1, tensor['elem_count'])}];" + ) + arguments.append(buffer) + lines += [f" return {entrypoint['symbol']}({', '.join(arguments)});", "}", ""] + return "\n".join(lines) + + +def _build(compiler: str, directory: Path, sources: list[Path]) -> Path: + binary = directory / f"{compiler}_artifact" + result = subprocess.run( + [ + compiler, + *STRICT_FLAGS, + f"-I{directory}", + *[str(source) for source in sources], + "-o", + str(binary), + "-lm", + ], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + return binary + + +def _build_and_run(compiler: str, result) -> None: + """Build the artifact from two units — one asking for the implementation — and run it.""" + directory = result.header_path.parent + main = directory / "main.c" + main.write_text(_driver(result.report), encoding="utf-8") + unit = directory / "implementation.c" + unit.write_text( + f"#define {result.report['prefix'].upper()}_IMPLEMENTATION\n" + f'#include "{result.report["header"]}"\n', + encoding="utf-8", + ) + binary = _build(compiler, directory, [main, unit]) + run = subprocess.run([str(binary)], capture_output=True, text=True) + assert run.returncode == 0, run.stderr + + +needs_c_compiler = pytest.mark.skipif( + not C_COMPILERS, reason="no system C compiler available" +) + + +def test_compile_writes_header_and_report(tmp_path): + result = compile_onnx(_alias_model(), tmp_path) + + assert result.header_path == tmp_path / "demo.h" + assert result.report_path == tmp_path / "demo_report.json" + assert json.loads(result.report_path.read_text()) == result.report + assert "int demo_run(" in result.header_path.read_text() + + +def test_declarations_and_implementation_are_separated(tmp_path): + header = compile_onnx(_alias_model(), tmp_path).header_path.read_text() + declarations, _, implementation = header.partition("#ifdef DEMO_IMPLEMENTATION") + + assert ( + "int demo_run(const float* x, const float* unread, float* y, float* x_2);" + in declarations + ) + assert "static const float demo_w_y[2]" not in declarations + assert "static const float demo_w_y[2]" in implementation + assert "#ifndef DEMO_H_INCLUDED" in declarations + + +@needs_c_compiler +@pytest.mark.parametrize("compiler", C_COMPILERS) +def test_header_builds_and_runs_under_strict_flags(tmp_path, compiler): + _build_and_run(compiler, compile_onnx(_alias_model(), tmp_path)) + + +@needs_c_compiler +@pytest.mark.parametrize("compiler", C_COMPILERS) +def test_header_is_include_guarded(tmp_path, compiler): + compile_onnx(_alias_model(), tmp_path) + source = tmp_path / "twice.c" + source.write_text( + "#define DEMO_IMPLEMENTATION\n" + '#include "demo.h"\n#include "demo.h"\n' + "int main(void) { return 0; }\n", + encoding="utf-8", + ) + _build(compiler, tmp_path, [source]) + + +def _comment_hostile_model(): + """Names carrying both block-comment delimiters, which reach the header's comments.""" + return _model( + [], + [_tensor("in/*x", TensorProto.FLOAT, [2])], + [ + _tensor("in/*x", TensorProto.FLOAT, [2]), + _tensor("w*/z", TensorProto.FLOAT, [2]), + ], + initializer=[ + onnx.numpy_helper.from_array(np.ones(2, dtype=np.float32), "w*/z") + ], + name="hostile/*name*/", + ) + + +def test_names_never_nest_a_block_comment(tmp_path): + header = compile_onnx(_comment_hostile_model(), tmp_path).header_path.read_text() + + for comment in re.findall(r"/\*.*?\*/", header, flags=re.DOTALL): + assert "/*" not in comment[2:] and "*/" not in comment[:-2], comment + + +@needs_c_compiler +@pytest.mark.parametrize("compiler", C_COMPILERS) +def test_comment_hostile_names_build_under_strict_flags(tmp_path, compiler): + _build_and_run(compiler, compile_onnx(_comment_hostile_model(), tmp_path)) + + +def test_generated_source_has_no_allocation_calls(tmp_path): + header = compile_onnx(_alias_model(), tmp_path).header_path.read_text() + + for token in ALLOCATION_TOKENS: + assert not re.search(rf"\b{token}\b", header), token + assert "#include" in header + assert set(re.findall(r"#include <(\w+\.h)>", header)) <= { + "stdint.h", + "stddef.h", + "string.h", + "math.h", + } + + +def test_repeated_compiles_are_byte_identical(tmp_path): + first = compile_onnx(_alias_model(), tmp_path / "first", dim_bindings={"batch": 3}) + second = compile_onnx( + _alias_model(), tmp_path / "second", dim_bindings={"batch": 3} + ) + + assert first.header_path.read_bytes() == second.header_path.read_bytes() + assert first.report_path.read_bytes() == second.report_path.read_bytes() + + +def test_report_records_options_bindings_opsets_and_footprint(tmp_path): + result = compile_onnx(_alias_model(), tmp_path, dim_bindings={"batch": 4}) + report = result.report + + assert report["fnnx_version"] == __version__ + assert report["options"] == { + "prefix": None, + "dim_bindings": {"batch": 4}, + "runtime_dims": {}, + } + assert report["dim_bindings"] == {"batch": 4} + assert report["runtime_dims"] == [] + assert report["opsets"] == {"ai.onnx": OPSET} + assert report["kernels"] == [] + assert report["memory"] == { + "weights_bytes": 8, + "arena_bytes": 0, + "static_bytes": 8, + } + assert [tensor["name"] for tensor in report["entrypoint"]["inputs"]] == [ + "x", + "unread", + ] + assert report["entrypoint"]["inputs"][0] == { + "name": "x", + "c_name": "x", + "macro": "DEMO_INPUT_X", + "dtype": "float32", + "c_type": "float", + "shape": [4, 3], + "elem_count": 12, + "bytes": 48, + } + + +def test_unbound_dimensions_default_to_one(tmp_path): + report = compile_onnx(_alias_model(), tmp_path).report + + assert report["dim_bindings"] == {"batch": 1} + assert report["entrypoint"]["inputs"][0]["shape"] == [1, 3] + assert "Dimension bindings: batch=1" in (tmp_path / "demo.h").read_text() + + +def test_the_preamble_documents_usage_reentrancy_bindings_and_footprint(tmp_path): + result = compile_onnx(_alias_model(), tmp_path, dim_bindings={"batch": 4}) + preamble = result.header_path.read_text().split("*/", 1)[0] + memory = result.report["memory"] + + assert f"#define {result.report['prefix'].upper()}_IMPLEMENTATION" in preamble + assert f'#include "{result.report["header"]}"' in preamble + assert f"{result.report['entrypoint']['symbol']}` runs the whole model" in preamble + assert "DEMO_OK on success" in preamble + assert " ".join(STRICT_FLAGS) in preamble + assert "Not reentrant" in preamble + assert ( + f"Static memory: {memory['static_bytes']} bytes " + f"({memory['weights_bytes']} of weights, {memory['arena_bytes']} of scratch)" + in preamble + ) + assert f"Opset imports: ai.onnx={OPSET}" in preamble + assert "Dimension bindings: batch=4" in preamble + + +def test_metadata_macros_describe_every_tensor(tmp_path): + result = compile_onnx(_alias_model(), tmp_path, dim_bindings={"batch": 4}) + header = result.header_path.read_text() + defined = dict(re.findall(r"#define (\w+) (\S+)", header)) + + for tensor in ( + *result.report["entrypoint"]["inputs"], + *result.report["entrypoint"]["outputs"], + ): + macro = tensor["macro"] + assert defined[f"{macro}_RANK"] == str(len(tensor["shape"])) + assert defined[f"{macro}_COUNT"] == str(tensor["elem_count"]) + for axis, size in enumerate(tensor["shape"]): + assert defined[f"{macro}_DIM_{axis}"] == str(size) + assert defined["DEMO_ARENA_BYTES"] == "0" + assert defined["DEMO_WEIGHTS_BYTES"] == "8" + assert defined["DEMO_STATIC_BYTES"] == "8" + + +def test_every_public_name_carries_the_prefix(tmp_path): + header = compile_onnx( + _alias_model(), tmp_path, prefix="my model.v2" + ).header_path.read_text() + + assert "int my_model_v2_run(" in header + for macro in re.findall(r"#define (\w+)", header): + assert macro.startswith("MY_MODEL_V2_"), macro + for line in header.splitlines(): + declaration = re.match(r"(?:static )?[A-Za-z_][\w ]*?[ *](\w+)\(", line) + if declaration: + assert declaration.group(1).startswith("my_model_v2_"), line + + +def test_prefix_falls_back_when_the_graph_name_is_unusable(tmp_path): + result = compile_onnx(_alias_model(name="***"), tmp_path) + + assert result.report["prefix"] == codegen.DEFAULT_PREFIX + assert result.header_path.name == f"{codegen.DEFAULT_PREFIX}.h" + + +def test_colliding_tensor_names_get_distinct_identifiers(tmp_path): + model = _model( + [], + [ + _tensor("x.1", TensorProto.FLOAT, [2]), + _tensor("x-1", TensorProto.FLOAT, [2]), + ], + [ + _tensor("x.1", TensorProto.FLOAT, [2]), + _tensor("x-1", TensorProto.FLOAT, [2]), + ], + ) + report = compile_onnx(model, tmp_path).report + names = [tensor["c_name"] for tensor in report["entrypoint"]["inputs"]] + macros = [tensor["macro"] for tensor in report["entrypoint"]["outputs"]] + + assert names == ["x_1", "x_1_2"] + assert len(set(macros)) == 2 + + +def test_a_parameter_never_shadows_a_static_buffer(tmp_path): + """An input named like a weight's symbol would hide it inside the entrypoint.""" + left = onnx.numpy_helper.from_array(np.array([1.5, 2.5], dtype=np.float32), "left") + right = onnx.numpy_helper.from_array( + np.array([1.0, 4.0], dtype=np.float32), "right" + ) + model = _model( + [helper.make_node("Add", ["left", "right"], ["y"], name="add")], + [_tensor("demo_w_y", TensorProto.FLOAT, [2])], + [ + _tensor("y", TensorProto.FLOAT, [2]), + _tensor("demo_w_y", TensorProto.FLOAT, [2]), + ], + initializer=[left, right], + name="demo", + ) + header = compile_onnx(model, tmp_path).header_path.read_text() + weights = re.findall(r"static const float (\w+)\[2\]", header) + signature = re.findall(r"int demo_run\(([^)]*)\)", header) + parameters = re.findall(r"\*\s*(\w+)", signature[0]) + + assert len(weights) == 1 + assert weights[0] not in parameters + assert f"memcpy(y, {weights[0]}, 2u * sizeof(*y));" in header + + +@pytest.mark.parametrize("elem_type", sorted(dtypes.C_TYPES)) +def test_scalar_literals_round_trip(elem_type): + values = _sample_values(elem_type) + for value in values.tolist(): + literal = emit.scalar_literal(value, elem_type) + parsed = _parse_c_literal(literal) + expected = int(value) if isinstance(value, bool) else value + assert _same_value(parsed, expected, values.dtype), literal + + +def _external_data_model(values, location: str): + """A weight whose bytes live in a side file next to the model.""" + tensor = onnx.numpy_helper.from_array(values, "w") + onnx.external_data_helper.set_external_data(tensor, location=location) + tensor.ClearField("raw_data") + return _model( + [], + [], + [_tensor("w", TensorProto.FLOAT, list(values.shape))], + initializer=[tensor], + name="external", + ) + + +def test_external_data_weights_are_embedded_in_the_header(tmp_path): + values = np.arange(6, dtype=np.float32).reshape(2, 3) + (tmp_path / "w.bin").write_bytes(values.tobytes()) + model_path = tmp_path / "model.onnx" + onnx.save_model(_external_data_model(values, "w.bin"), str(model_path)) + + result = compile_onnx(model_path, tmp_path / "out") + emitted = _weight_arrays(result.header_path.read_text()) + + assert emitted["external_w_w"] == values.reshape(-1).tolist() + + +def test_embedded_weights_preserve_every_supported_dtype(tmp_path): + arrays = { + dtypes.numpy_dtype_name(elem_type): _sample_values(elem_type) + for elem_type in sorted(dtypes.C_TYPES) + } + model = _model( + [_constant_node(name, array) for name, array in arrays.items()], + [], + [ + _tensor( + name, onnx.helper.np_dtype_to_tensor_dtype(array.dtype), array.shape + ) + for name, array in arrays.items() + ], + name="weights", + ) + result = compile_onnx(model, tmp_path) + emitted = _weight_arrays(result.header_path.read_text()) + + for name, array in arrays.items(): + parsed = emitted[f"weights_w_{name}"] + expected = [ + int(value) if isinstance(value, bool) else value for value in array.tolist() + ] + assert len(parsed) == len(expected) + assert all( + _same_value(read, source, array.dtype) + for read, source in zip(parsed, expected) + ), name + + +@needs_c_compiler +@pytest.mark.parametrize("compiler", C_COMPILERS) +def test_every_supported_dtype_builds(tmp_path, compiler): + arrays = { + dtypes.numpy_dtype_name(elem_type): _sample_values(elem_type) + for elem_type in sorted(dtypes.C_TYPES) + } + model = _model( + [_constant_node(name, array) for name, array in arrays.items()], + [], + [ + _tensor( + name, onnx.helper.np_dtype_to_tensor_dtype(array.dtype), array.shape + ) + for name, array in arrays.items() + ], + name="weights", + ) + _build_and_run(compiler, compile_onnx(model, tmp_path)) + + +def _zero_element_model(): + empty = np.zeros((0,), dtype=np.float32) + return _model( + [_constant_node("empty", empty)], + [_tensor("x", TensorProto.FLOAT, [0, 3])], + [ + _tensor("empty", TensorProto.FLOAT, [0]), + _tensor("x", TensorProto.FLOAT, [0, 3]), + ], + name="zero", + ) + + +def test_zero_element_tensors_declare_no_empty_arrays(tmp_path): + result = compile_onnx(_zero_element_model(), tmp_path) + header = result.header_path.read_text() + + assert "[0]" not in header.split("#ifdef ZERO_IMPLEMENTATION")[1] + assert "memcpy" not in header + assert result.report["entrypoint"]["outputs"][0]["elem_count"] == 0 + assert result.report["memory"]["static_bytes"] == 0 + + +@needs_c_compiler +@pytest.mark.parametrize("compiler", C_COMPILERS) +def test_zero_element_artifact_builds_and_runs(tmp_path, compiler): + _build_and_run(compiler, compile_onnx(_zero_element_model(), tmp_path)) + + +def _relu_generator(context): + """A stand-in kernel: the emitter's contract, not an ONNX-conformant Relu.""" + source, target = context.inputs[0], context.outputs[0] + c_type = dtypes.c_type(source.elem_type) + name = f"{context.prefix}_relu_{c_type}" + definition = "\n".join( + [ + f"static void {name}({c_type}* out, const {c_type}* in, size_t count)", + "{", + " size_t index;", + " for (index = 0; index < count; ++index) {", + " out[index] = in[index] > 0 ? in[index] : 0;", + " }", + "}", + ] + ) + return kernels.NodeEmission( + functions=(kernels.CFunction(name, definition),), + statements=(f"{name}({target.expr}, {source.expr}, {target.elem_count}u);",), + ) + + +@pytest.fixture +def relu_registry(monkeypatch): + stub = registry.KernelRegistry() + stub.register("", "Relu", 14, _relu_generator) + monkeypatch.setattr(codegen, "KERNELS", stub) + return stub + + +def _chain_model(): + """Two chained Relus on float and one on double: a shared kernel and a private one.""" + return _model( + [ + helper.make_node("Relu", ["x"], ["hidden"], name="first"), + helper.make_node("Relu", ["hidden"], ["y"], name="second"), + helper.make_node("Relu", ["xd"], ["yd"], name="third"), + ], + [ + _tensor("x", TensorProto.FLOAT, [2, 3]), + _tensor("xd", TensorProto.DOUBLE, [4]), + ], + [ + _tensor("y", TensorProto.FLOAT, [2, 3]), + _tensor("yd", TensorProto.DOUBLE, [4]), + ], + name="chain", + ) + + +def test_kernels_are_shared_and_intermediates_are_static(tmp_path, relu_registry): + result = compile_onnx(_chain_model(), tmp_path) + header = result.header_path.read_text() + + assert result.report["kernels"] == ["chain_relu_float", "chain_relu_double"] + assert header.count("static void chain_relu_float(") == 1 + assert header.count("chain_relu_float(") == 3 + assert "static float chain_t_hidden[6];" in header + assert result.report["memory"] == { + "weights_bytes": 0, + "arena_bytes": 24, + "static_bytes": 24, + } + calls = re.findall(r"chain_relu_\w+\([^;]+\);", header) + assert calls == [ + "chain_relu_float(chain_t_hidden, x, 6u);", + "chain_relu_float(y, chain_t_hidden, 6u);", + "chain_relu_double(yd, xd, 4u);", + ] + + +@needs_c_compiler +@pytest.mark.parametrize("compiler", C_COMPILERS) +def test_kernel_artifact_builds_and_runs(tmp_path, compiler, relu_registry): + _build_and_run(compiler, compile_onnx(_chain_model(), tmp_path)) + + +def test_fan_out_output_is_read_from_the_caller_buffer(tmp_path, relu_registry): + """A node output that is both a graph output and a downstream input.""" + model = _model( + [ + helper.make_node("Relu", ["x"], ["shared"], name="first"), + helper.make_node("Relu", ["shared"], ["y"], name="second"), + ], + [_tensor("x", TensorProto.FLOAT, [2])], + [ + _tensor("shared", TensorProto.FLOAT, [2]), + _tensor("y", TensorProto.FLOAT, [2]), + ], + name="fanout", + ) + header = compile_onnx(model, tmp_path).header_path.read_text() + calls = re.findall(r"fanout_relu_\w+\([^;]+\);", header) + + assert calls == [ + "fanout_relu_float(shared, x, 2u);", + "fanout_relu_float(y, shared, 2u);", + ] + assert "fanout_t_shared" not in header + + +def test_a_kernel_colliding_with_a_tensor_name_is_rejected(tmp_path, relu_registry): + """Such a tensor would shadow the kernel inside the entrypoint, breaking the build.""" + model = _model( + [helper.make_node("Relu", ["demo_relu_float"], ["y"], name="relu")], + [_tensor("demo_relu_float", TensorProto.FLOAT, [2])], + [_tensor("y", TensorProto.FLOAT, [2])], + name="demo", + ) + + with pytest.raises(CompileError, match="`demo_relu_float` collides with"): + compile_onnx(model, tmp_path) + + assert not tmp_path.exists() or not list(tmp_path.iterdir()) + + +def test_kernel_scratch_colliding_with_a_tensor_name_is_rejected(tmp_path): + """A kernel's working buffer is named like the kernel, and shadowed the same way.""" + model = _model( + [helper.make_node("Det", ["demo_det_float_work"], ["y"], name="det")], + [_tensor("demo_det_float_work", TensorProto.FLOAT, [2, 2])], + [_tensor("y", TensorProto.FLOAT, [])], + name="demo", + opset=22, + ) + + with pytest.raises(CompileError, match="`demo_det_float_work` collides with"): + compile_onnx(model, tmp_path) + + assert not tmp_path.exists() or not list(tmp_path.iterdir()) + + +def test_kernels_sharing_a_name_must_share_a_definition(tmp_path, monkeypatch): + def clashing(context): + target = context.outputs[0] + definition = f"static void clash(void) {{ /* {context.node.name} */ }}" + return kernels.NodeEmission( + functions=(kernels.CFunction("clash", definition),), + statements=(f"(void){target.expr};",), + ) + + stub = registry.KernelRegistry() + stub.register("", "Relu", 14, clashing) + monkeypatch.setattr(codegen, "KERNELS", stub) + model = _model( + [ + helper.make_node("Relu", ["x"], ["hidden"], name="first"), + helper.make_node("Relu", ["hidden"], ["y"], name="second"), + ], + [_tensor("x", TensorProto.FLOAT, [2])], + [_tensor("y", TensorProto.FLOAT, [2])], + ) + with pytest.raises(CompileError, match="emitted twice with different definitions"): + compile_onnx(model, tmp_path) + + +def test_unsupported_op_writes_no_files(tmp_path, monkeypatch): + """An empty registry stands in for any op no kernel covers, whatever is implemented.""" + monkeypatch.setattr(codegen, "KERNELS", registry.KernelRegistry()) + model = _model( + [helper.make_node("Relu", ["x"], ["y"], name="relu")], + [_tensor("x", TensorProto.FLOAT, [2])], + [_tensor("y", TensorProto.FLOAT, [2])], + ) + output_dir = tmp_path / "out" + with pytest.raises(CompileError) as error: + compile_onnx(model, output_dir) + + message = str(error.value) + assert "`relu`" in message and "`Relu`" in message + assert "ai.onnx" in message and str(OPSET) in message + assert not output_dir.exists() + + +def test_data_dependent_shape_op_writes_no_files(tmp_path): + model = _model( + [helper.make_node("NonZero", ["x"], ["y"], name="nonzero")], + [_tensor("x", TensorProto.FLOAT, [2, 3])], + [_tensor("y", TensorProto.INT64, [2, "n"])], + ) + output_dir = tmp_path / "out" + with pytest.raises(CompileError) as error: + compile_onnx(model, output_dir) + + message = str(error.value) + assert "`nonzero`" in message and "`NonZero`" in message + assert "depends on input data" in message + assert not output_dir.exists() + + +def test_output_without_a_producer_is_rejected(tmp_path): + model = _model( + [], + [_tensor("x", TensorProto.FLOAT, [2])], + [_tensor("missing", TensorProto.FLOAT, [2])], + ) + with pytest.raises(CompileError, match="`missing` is not produced"): + compile_onnx(model, tmp_path) + + +def test_node_reading_an_undefined_tensor_is_rejected(relu_registry): + """Nodes out of topological order, or reading a tensor nothing defines. + + Shape inference rejects most such graphs first, so codegen is driven directly here: + the emitter must still refuse rather than reference an undeclared C symbol. + """ + model = _model( + [ + helper.make_node("Relu", ["hidden"], ["y"], name="second"), + helper.make_node("Relu", ["x"], ["hidden"], name="first"), + ], + [_tensor("x", TensorProto.FLOAT, [2])], + [_tensor("y", TensorProto.FLOAT, [2])], + ) + prepared = frontend.PreparedModel(model=model, opsets={"": OPSET}, dim_bindings={}) + + with pytest.raises(CompileError, match="reads tensor `hidden`"): + codegen.build_program(prepared) + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("plain", "plain"), + ("with.dots-and spaces", "with_dots_and_spaces"), + ("2fast", "v_2fast"), + ("int", "int_"), + ("", "fallback"), + ("***", "fallback"), + ], +) +def test_sanitize_identifier(name, expected): + assert emit.sanitize_identifier(name, fallback="fallback") == expected + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ("plain", "plain"), + ("ends*/here", "ends* /here"), + ("opens/*here", "opens/ *here"), + ("/**/", "/ ** /"), + ], +) +def test_comment_safe_neutralizes_both_delimiters(text, expected): + assert emit.comment_safe(text) == expected + + +def test_unique_names_disambiguate_case_insensitively(): + names = emit.UniqueNames() + assigned = [names.assign(name, fallback="v") for name in ("a", "A", "a.", "b")] + + assert assigned == ["a", "A_2", "a_", "b"] diff --git a/src/python/tests/test_extra_compiler_frontend.py b/src/python/tests/test_extra_compiler_frontend.py new file mode 100644 index 0000000..c40b62d --- /dev/null +++ b/src/python/tests/test_extra_compiler_frontend.py @@ -0,0 +1,1398 @@ +"""Dimension binding, constant folding, and static verification in the C compiler frontend. + +Every expected value comes from the ONNX reference evaluator — the executable form of the +spec — or from the ONNX schema registry; none is hand-written. +""" + +from __future__ import annotations + +import unittest +from pathlib import Path + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +loader = pytest.importorskip("fnnx.extras.compilers.c.onnx.loader") +dtypes = pytest.importorskip("fnnx.extras.compilers.c.onnx.dtypes") +folding = pytest.importorskip("fnnx.extras.compilers.c.onnx.folding") +frontend = pytest.importorskip("fnnx.extras.compilers.c.onnx.frontend") +shapes = pytest.importorskip("fnnx.extras.compilers.c.onnx.shapes") +verify = pytest.importorskip("fnnx.extras.compilers.c.onnx.verify") + +from onnx import ModelProto, TensorProto, helper # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +MODELS_DIR = Path(__file__).parent / "models" +OPS_ARTIFACTS = MODELS_DIR / "onnx_pipeline.fnnx" / "ops_artifacts" + +MAX_OPSET = onnx.defs.onnx_opset_version() +# The opset the bundle's own node models import; folding must work at a real model's +# version, not only at the newest one the installed `onnx` package defines. +BUNDLE_OPSET = 21 + + +def _model( + nodes, + inputs, + outputs, + *, + initializer=(), + sparse_initializer=(), + opset: int = BUNDLE_OPSET, + ir_version: int | None = None, +): + graph = helper.make_graph( + nodes, + "g", + inputs, + outputs, + initializer=list(initializer), + sparse_initializer=list(sparse_initializer), + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + if ir_version is not None: + model.ir_version = ir_version + return model + + +def _prepare(model, **kwargs): + return frontend.prepare_model(loader.load_model(model), **kwargs) + + +def _tensor(name, elem_type, shape): + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _unshaped(name, elem_type=TensorProto.FLOAT): + return helper.make_tensor_value_info(name, elem_type, None) + + +def _constant_node(name, values, dtype=TensorProto.INT64): + array = np.asarray(values) + return helper.make_node( + "Constant", + [], + [name], + value=helper.make_tensor( + f"{name}_value", dtype, list(array.shape), array.flatten().tolist() + ), + ) + + +def _constant_branch(name, values): + """A subgraph that produces a constant, for use as an `If` branch.""" + return helper.make_graph( + [_constant_node(name, values, dtype=TensorProto.FLOAT)], + name, + [], + [_tensor(name, TensorProto.FLOAT, [len(values)])], + ) + + +def _random_branch(name): + """A subgraph that samples a distribution, for use as an `If` branch.""" + return helper.make_graph( + [ + helper.make_node( + "RandomNormal", [], [name], dtype=TensorProto.FLOAT, shape=[2] + ) + ], + name, + [], + [_tensor(name, TensorProto.FLOAT, [2])], + ) + + +def _random_op_models() -> dict[str, ModelProto]: + """Models whose single node samples a distribution, with every input a constant. + + No `seed` attribute is set, so each evaluation of these nodes draws different values. + """ + probabilities = onnx.numpy_helper.from_array( + np.full((4, 4), 0.5, dtype=np.float32), "p" + ) + return { + "RandomNormal": _model( + [ + helper.make_node( + "RandomNormal", + [], + ["y"], + name="draw", + dtype=TensorProto.FLOAT, + shape=[2, 3], + ) + ], + [], + [_unshaped("y")], + opset=MAX_OPSET, + ), + "Bernoulli": _model( + [helper.make_node("Bernoulli", ["p"], ["y"], name="draw")], + [], + [_unshaped("y")], + initializer=[probabilities], + opset=MAX_OPSET, + ), + # Dropout samples a mask in training mode; only its inference-mode identity is in scope. + "Dropout": _model( + [ + helper.make_node( + "Dropout", ["p", "ratio", "training"], ["y"], name="drop" + ) + ], + [], + [_unshaped("y")], + initializer=[ + probabilities, + onnx.numpy_helper.from_array(np.float32(0.5), "ratio"), + onnx.numpy_helper.from_array(np.array(True), "training"), + ], + opset=MAX_OPSET, + ), + } + + +def _op_types(prepared) -> list[str]: + return [node.op_type for node in prepared.model.graph.node] + + +def _initializer(prepared, name): + for initializer in prepared.model.graph.initializer: + if initializer.name == name: + return onnx.numpy_helper.to_array(initializer) + raise AssertionError(f"no initializer named `{name}`") + + +def _output_shape(prepared, name) -> tuple[int, ...]: + for value_info in prepared.model.graph.output: + if value_info.name == name: + return shapes.static_shape(value_info.type) + raise AssertionError(f"no graph output named `{name}`") + + +def _schema_revisions(op_type: str, domain: str = "") -> list[int]: + return sorted( + schema.since_version + for schema in onnx.defs.get_all_schemas_with_history() + if schema.name == op_type and schema.domain == domain + ) + + +def _reference_outputs(model, feeds: dict) -> list: + """What ONNX's own evaluator — the executable spec — computes for `model`.""" + return list(ReferenceEvaluator(model).run(None, feeds)) + + +def _assert_matches_reference(test: unittest.TestCase, original, prepared, feeds: dict): + """The prepared graph computes what the original did, at the shapes it claims.""" + expected = _reference_outputs(original, feeds) + actual = _reference_outputs(prepared.model, feeds) + test.assertEqual(len(expected), len(actual)) + for value_info, want, got in zip(prepared.model.graph.output, expected, actual): + np.testing.assert_allclose(got, want, rtol=1e-6, atol=1e-6) + test.assertEqual( + _output_shape(prepared, value_info.name), + tuple(np.asarray(got).shape), + f"declared shape of `{value_info.name}` does not match the computed one", + ) + + +class BindDimsTest(unittest.TestCase): + def _identity_model(self, shape, elem_type=TensorProto.FLOAT): + return _model( + [helper.make_node("Identity", ["x"], ["y"])], + [_tensor("x", elem_type, shape)], + [_unshaped("y", elem_type)], + ) + + def test_unbound_symbolic_dim_defaults_to_one_and_is_recorded(self): + prepared = _prepare(self._identity_model(["batch", 3])) + + self.assertEqual(prepared.dim_bindings, {"batch": 1}) + self.assertEqual(_output_shape(prepared, "y"), (1, 3)) + + def test_explicit_binding_reaches_inputs_and_inferred_outputs(self): + original = self._identity_model(["batch", 3]) + + prepared = _prepare(original, dim_bindings={"batch": 4}) + + self.assertEqual(prepared.dim_bindings, {"batch": 4}) + self.assertEqual( + shapes.static_shape(prepared.model.graph.input[0].type), (4, 3) + ) + _assert_matches_reference( + self, original, prepared, {"x": np.zeros((4, 3), dtype=np.float32)} + ) + + def test_unnamed_unknown_dim_defaults_to_one(self): + model = self._identity_model([None, 3]) + + prepared = _prepare(model, dim_bindings={"batch": 4}) + + self.assertEqual(prepared.dim_bindings, {}) + self.assertEqual(_output_shape(prepared, "y"), (1, 3)) + + def test_binding_for_an_absent_dim_is_ignored(self): + prepared = _prepare(self._identity_model([2, 3]), dim_bindings={"unused": 8}) + + self.assertEqual(prepared.dim_bindings, {}) + self.assertEqual(_output_shape(prepared, "y"), (2, 3)) + + def test_one_name_binds_every_occurrence(self): + original = _model( + [helper.make_node("Add", ["x", "z"], ["y"])], + [ + _tensor("x", TensorProto.FLOAT, ["batch", 3]), + _tensor("z", TensorProto.FLOAT, ["batch", 3]), + ], + [_unshaped("y")], + ) + + prepared = _prepare(original, dim_bindings={"batch": 5}) + + self.assertEqual(prepared.dim_bindings, {"batch": 5}) + feeds = { + "x": np.arange(15, dtype=np.float32).reshape(5, 3), + "z": np.ones((5, 3), dtype=np.float32), + } + _assert_matches_reference(self, original, prepared, feeds) + + def test_zero_is_a_valid_binding(self): + original = self._identity_model(["batch", 3]) + + prepared = _prepare(original, dim_bindings={"batch": 0}) + + self.assertEqual(_output_shape(prepared, "y"), (0, 3)) + _assert_matches_reference( + self, original, prepared, {"x": np.zeros((0, 3), dtype=np.float32)} + ) + + def test_invalid_binding_values_are_rejected(self): + for value in (-1, "4", 2.0, True, None): + with self.subTest(value=value): + with self.assertRaises(CompileError) as ctx: + _prepare( + self._identity_model(["batch", 3]), + dim_bindings={"batch": value}, + ) + self.assertIn("batch", str(ctx.exception)) + + def test_initializer_shadowing_an_input_becomes_a_weight(self): + """Pre-IR-4 models declare initializers as inputs; they compile as static weights.""" + weights = np.array([[1.0], [2.0], [3.0]], dtype=np.float32) + original = _model( + [helper.make_node("MatMul", ["x", "w"], ["y"])], + [ + _tensor("x", TensorProto.FLOAT, ["batch", 3]), + _tensor("w", TensorProto.FLOAT, [3, 1]), + ], + [_unshaped("y")], + initializer=[onnx.numpy_helper.from_array(weights, "w")], + ir_version=3, + ) + + prepared = _prepare(original, dim_bindings={"batch": 2}) + + self.assertEqual([entry.name for entry in prepared.model.graph.input], ["x"]) + self.assertEqual( + [entry.name for entry in prepared.model.graph.initializer], ["w"] + ) + self.assertEqual(_output_shape(prepared, "y"), (2, 1)) + + def test_an_output_the_graph_does_not_compute_aliases_its_input(self): + """Echoing an input as an output must not cost it the shape the binding gave it.""" + original = _model( + [helper.make_node("Relu", ["x"], ["y"])], + [_tensor("x", TensorProto.FLOAT, ["batch", 2])], + [_tensor("x", TensorProto.FLOAT, ["batch", 2]), _unshaped("y")], + ) + + prepared = _prepare(original, dim_bindings={"batch": 3}) + + self.assertEqual(_output_shape(prepared, "x"), (3, 2)) + _assert_matches_reference( + self, + original, + prepared, + {"x": np.arange(6, dtype=np.float32).reshape(3, 2) - 2}, + ) + + def test_binding_reaches_a_declared_ml_output_shape(self): + """`ai.onnx.ml` inference drops the batch dim; the binding reaches it through the + output declaration, which is the only static description of it there is.""" + graph = helper.make_graph( + [ + helper.make_node( + "LinearRegressor", + ["x"], + ["y"], + domain="ai.onnx.ml", + coefficients=[1.0, 2.0, 3.0], + intercepts=[0.5], + targets=1, + ) + ], + "g", + [_tensor("x", TensorProto.FLOAT, ["batch", 3])], + [_tensor("y", TensorProto.FLOAT, ["batch", 1])], + ) + original = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", BUNDLE_OPSET), + helper.make_opsetid("ai.onnx.ml", 1), + ], + ) + + prepared = _prepare(original, dim_bindings={"batch": 4}) + + self.assertEqual(prepared.dim_bindings, {"batch": 4}) + self.assertEqual(_output_shape(prepared, "y"), (4, 1)) + _assert_matches_reference( + self, + original, + prepared, + {"x": np.arange(12, dtype=np.float32).reshape(4, 3)}, + ) + + +class ConstantFoldingTest(unittest.TestCase): + def _shape_chain_model(self, opset): + """Flatten x by computing its element count from its own shape.""" + return _model( + [ + helper.make_node("Shape", ["x"], ["s"]), + _constant_node("i0", [0]), + _constant_node("i1", [1]), + helper.make_node("Gather", ["s", "i0"], ["rows"]), + helper.make_node("Gather", ["s", "i1"], ["cols"]), + helper.make_node("Mul", ["rows", "cols"], ["count"]), + helper.make_node("Reshape", ["x", "count"], ["y"]), + ], + [_tensor("x", TensorProto.FLOAT, ["batch", 3])], + [_unshaped("y")], + opset=opset, + ) + + def test_shape_computation_chain_folds_to_a_constant(self): + for opset in (BUNDLE_OPSET, MAX_OPSET): + with self.subTest(opset=opset): + original = self._shape_chain_model(opset) + + prepared = _prepare(original, dim_bindings={"batch": 4}) + + self.assertEqual(_op_types(prepared), ["Reshape"]) + np.testing.assert_array_equal(_initializer(prepared, "count"), [12]) + self.assertEqual(_output_shape(prepared, "y"), (12,)) + _assert_matches_reference( + self, + original, + prepared, + {"x": np.arange(12, dtype=np.float32).reshape(4, 3)}, + ) + + def test_folding_reaches_a_fixpoint_through_chained_constants(self): + original = _model( + [ + _constant_node("a", [1, 2, 3]), + _constant_node("b", [10, 20, 30]), + helper.make_node("Add", ["a", "b"], ["c"]), + helper.make_node("Mul", ["c", "b"], ["y"]), + ], + [], + [_unshaped("y", TensorProto.INT64)], + ) + + prepared = _prepare(original) + + self.assertEqual(_op_types(prepared), []) + np.testing.assert_array_equal( + _initializer(prepared, "y"), _reference_outputs(original, {})[0] + ) + + def test_a_node_reading_one_constant_twice_folds(self): + original = _model( + [helper.make_node("Mul", ["c", "c"], ["y"])], + [], + [_unshaped("y", TensorProto.INT64)], + initializer=[ + onnx.numpy_helper.from_array(np.array([2, 3], dtype=np.int64), "c") + ], + ) + + prepared = _prepare(original) + + self.assertEqual(_op_types(prepared), []) + np.testing.assert_array_equal( + _initializer(prepared, "y"), _reference_outputs(original, {})[0] + ) + + def _generated_tensor_model(self, opset): + """A tensor built out of nothing but the operands describing it.""" + return _model( + [ + _constant_node("shape", [2, 3]), + helper.make_node( + "ConstantOfShape", + ["shape"], + ["filled"], + value=helper.make_tensor("value", TensorProto.FLOAT, [1], [1.5]), + ), + _constant_node("start", 0, dtype=TensorProto.FLOAT), + _constant_node("limit", 3, dtype=TensorProto.FLOAT), + _constant_node("step", 1, dtype=TensorProto.FLOAT), + helper.make_node("Range", ["start", "limit", "step"], ["counted"]), + helper.make_node("Add", ["filled", "counted"], ["y"]), + ], + [], + [_unshaped("y")], + opset=opset, + ) + + def test_a_tensor_generated_out_of_fixed_operands_is_resolved_by_folding(self): + """`ConstantOfShape` and `Range` state their whole result through their operands. + + Nothing of either is left once the graph fixes those operands — which is the only + form the compiler accepts them in at all, since the result's shape is a function of + their values — so neither op carries a kernel of its own. + """ + original = self._generated_tensor_model(MAX_OPSET) + + prepared = _prepare(original) + + self.assertEqual(_op_types(prepared), []) + self.assertEqual(_output_shape(prepared, "y"), (2, 3)) + np.testing.assert_array_equal( + _initializer(prepared, "y"), _reference_outputs(original, {})[0] + ) + + def test_a_window_and_a_mel_matrix_are_resolved_by_folding(self): + """A window and a mel matrix are functions of their operands alone. + + Nothing of either is left once the graph fixes those operands — which is the only + form the compiler accepts them in, their operands being their result's own shape — + so neither carries a kernel, and the values come from the evaluator rather than from + a reimplementation of the formulas. + """ + original = _model( + [ + _constant_node("size", 8, dtype=TensorProto.INT32), + _constant_node("bins", 4, dtype=TensorProto.INT32), + _constant_node("rate", 16000, dtype=TensorProto.INT32), + _constant_node("low", 0.0, dtype=TensorProto.FLOAT), + _constant_node("high", 8000.0, dtype=TensorProto.FLOAT), + helper.make_node("HannWindow", ["size"], ["hann"]), + helper.make_node("HammingWindow", ["size"], ["hamming"]), + helper.make_node("BlackmanWindow", ["size"], ["blackman"], periodic=0), + helper.make_node("Add", ["hann", "hamming"], ["summed"]), + helper.make_node("Add", ["summed", "blackman"], ["windows"]), + helper.make_node( + "MelWeightMatrix", + ["bins", "size", "rate", "low", "high"], + ["mel"], + ), + ], + [], + [_unshaped("windows"), _unshaped("mel")], + opset=MAX_OPSET, + ) + + prepared = _prepare(original) + + self.assertEqual(_op_types(prepared), []) + self.assertEqual(_output_shape(prepared, "windows"), (8,)) + self.assertEqual(_output_shape(prepared, "mel"), (5, 4)) + expected = _reference_outputs(original, {}) + for name, want in zip(("windows", "mel"), expected): + np.testing.assert_allclose(_initializer(prepared, name), want, rtol=1e-6) + + def test_a_generated_tensor_at_an_unvouchable_revision_is_left_for_dispatch(self): + """Resting on folding means resting on where the evaluator is a valid oracle. + + Below the revision it implements, folding declines rather than applying semantics + nothing can vouch for — and the node reaches dispatch, which has no kernel for it. + """ + prepared = _prepare(self._generated_tensor_model(BUNDLE_OPSET)) + + self.assertEqual(_op_types(prepared), ["ConstantOfShape", "Range", "Add"]) + + def test_nodes_reading_runtime_values_are_left_alone(self): + original = _model( + [ + _constant_node("w", [1.0, 2.0, 3.0], dtype=TensorProto.FLOAT), + helper.make_node("Add", ["x", "w"], ["y"]), + ], + [_tensor("x", TensorProto.FLOAT, ["batch", 3])], + [_unshaped("y")], + ) + + prepared = _prepare(original, dim_bindings={"batch": 2}) + + self.assertEqual(_op_types(prepared), ["Add"]) + _assert_matches_reference( + self, original, prepared, {"x": np.ones((2, 3), dtype=np.float32)} + ) + + def test_op_the_evaluator_cannot_vouch_for_is_not_folded(self): + """Opset-7 `Add` carries `broadcast`/`axis`; only its modern semantics are implemented.""" + constants = [ + onnx.numpy_helper.from_array(np.array([1, 2, 3], dtype=np.int64), "a"), + onnx.numpy_helper.from_array(np.array([10, 20, 30], dtype=np.int64), "b"), + ] + add = [helper.make_node("Add", ["a", "b"], ["y"])] + output = [_unshaped("y", TensorProto.INT64)] + + stale = _prepare(_model(add, [], output, initializer=constants, opset=7)) + current = _prepare( + _model( + add, + [], + output, + initializer=constants, + opset=max(_schema_revisions("Add")), + ) + ) + + self.assertEqual(_op_types(stale), ["Add"]) + self.assertEqual(_op_types(current), []) + + def test_constant_condition_selects_a_branch(self): + original = _model( + [ + helper.make_node( + "If", + ["cond"], + ["y"], + name="branch", + then_branch=_constant_branch("then", [1.0, 2.0]), + else_branch=_constant_branch("else", [3.0, 4.0]), + ) + ], + [], + [_unshaped("y")], + initializer=[onnx.numpy_helper.from_array(np.array(True), "cond")], + opset=MAX_OPSET, + ) + + prepared = _prepare(original) + + self.assertEqual(_op_types(prepared), []) + np.testing.assert_array_equal( + _initializer(prepared, "y"), _reference_outputs(original, {})[0] + ) + + def test_subgraph_reading_a_runtime_value_keeps_the_node(self): + """An `If` whose branch reads a runtime tensor cannot be folded, so it is rejected.""" + passthrough = helper.make_graph( + [helper.make_node("Identity", ["x"], ["t"])], "then", [], [_unshaped("t")] + ) + original = _model( + [ + helper.make_node( + "If", + ["cond"], + ["y"], + name="branch", + then_branch=passthrough, + else_branch=_constant_branch("else", [0.0, 0.0]), + ) + ], + [_tensor("x", TensorProto.FLOAT, [2])], + [_unshaped("y")], + initializer=[onnx.numpy_helper.from_array(np.array(True), "cond")], + opset=MAX_OPSET, + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(original) + self.assertIn("branch", str(ctx.exception)) + + def test_initializers_left_unused_by_folding_are_dropped(self): + original = _model( + [ + helper.make_node("Shape", ["x"], ["s"]), + helper.make_node("Reshape", ["x", "s"], ["y"]), + ], + [_tensor("x", TensorProto.FLOAT, ["batch", 3])], + [_unshaped("y")], + initializer=[ + onnx.numpy_helper.from_array(np.zeros(4, dtype=np.float32), "unused") + ], + ) + + prepared = _prepare(original, dim_bindings={"batch": 2}) + + self.assertEqual( + [entry.name for entry in prepared.model.graph.initializer], ["s"] + ) + self.assertEqual(_op_types(prepared), ["Reshape"]) + + def test_shapes_of_folded_away_tensors_are_dropped(self): + """Stale `value_info` would otherwise offer the emitter buffers for absent tensors.""" + prepared = _prepare( + self._shape_chain_model(BUNDLE_OPSET), dim_bindings={"batch": 4} + ) + + self.assertEqual([entry.name for entry in prepared.model.graph.value_info], []) + + def test_random_draws_are_never_folded(self): + """A draw is not a constant; baking one in would compile an unsupported op wrongly.""" + for op_type, model in _random_op_models().items(): + with self.subTest(op_type=op_type): + prepared = _prepare(model) + + self.assertEqual(_op_types(prepared), [op_type]) + + def test_a_random_draw_inside_a_branch_blocks_folding(self): + """Folding the enclosing `If` would bake the draw in; it is refused instead.""" + original = _model( + [ + helper.make_node( + "If", + ["cond"], + ["y"], + name="branch", + then_branch=_random_branch("then"), + else_branch=_constant_branch("else", [0.0, 0.0]), + ) + ], + [], + [_unshaped("y")], + initializer=[onnx.numpy_helper.from_array(np.array(True), "cond")], + opset=MAX_OPSET, + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(original) + + message = str(ctx.exception) + self.assertIn("branch", message) + self.assertIn("control flow", message) + + def test_a_random_op_does_not_make_preparing_non_deterministic(self): + loaded = loader.load_model(_random_op_models()["RandomNormal"]) + + first = frontend.prepare_model(loaded) + second = frontend.prepare_model(loaded) + + self.assertEqual( + first.model.SerializeToString(), second.model.SerializeToString() + ) + + +class EvaluatorFaithfulnessTest(unittest.TestCase): + def test_latest_revision_is_faithful(self): + for op_type in ("Add", "Cast", "Concat", "Shape", "Reshape"): + with self.subTest(op_type=op_type): + self.assertTrue( + folding.evaluator_is_version_faithful("", op_type, MAX_OPSET) + ) + + def test_superseded_semantics_are_refused(self): + """`Add` gained numpy broadcasting at opset 7 and dropped `axis`; only one is implemented.""" + self.assertFalse(folding.evaluator_is_version_faithful("", "Add", 6)) + self.assertFalse(folding.evaluator_is_version_faithful("", "Add", 7)) + self.assertTrue( + folding.evaluator_is_version_faithful( + "", "Add", max(_schema_revisions("Add")) + ) + ) + + def test_a_versioned_implementation_covers_its_own_revision(self): + """`Cast` has implementations at revisions 1 and 19; opset 13 lies between them.""" + self.assertTrue(folding.evaluator_is_version_faithful("", "Cast", 19)) + self.assertFalse(folding.evaluator_is_version_faithful("", "Cast", 13)) + + def test_domain_alias_is_accepted(self): + self.assertEqual( + folding.evaluator_is_version_faithful("ai.onnx", "Concat", MAX_OPSET), + folding.evaluator_is_version_faithful("", "Concat", MAX_OPSET), + ) + + def test_unknown_op_or_domain_is_not_faithful(self): + self.assertFalse( + folding.evaluator_is_version_faithful("", "NoSuchOp", MAX_OPSET) + ) + self.assertFalse( + folding.evaluator_is_version_faithful("com.example", "Add", MAX_OPSET) + ) + + +class StaticVerificationTest(unittest.TestCase): + def test_data_dependent_output_shape_names_the_node_and_op(self): + model = _model( + [helper.make_node("NonZero", ["x"], ["y"], name="find")], + [_tensor("x", TensorProto.FLOAT, ["batch", 3])], + [_unshaped("y", TensorProto.INT64)], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("find", message) + self.assertIn("NonZero", message) + self.assertIn("depends on input data", message) + + def test_shape_that_only_runtime_data_determines_is_rejected(self): + """A declared symbolic output shape never substitutes for one the graph must compute. + + Binding the names the model declares its result under is not enough: the shape the + graph actually computes comes from an operand nothing fixes, and the error names it. + """ + model = _model( + [helper.make_node("Reshape", ["x", "s"], ["y"], name="reshape")], + [ + _tensor("x", TensorProto.FLOAT, ["batch", 3]), + _tensor("s", TensorProto.INT64, [2]), + ], + [_tensor("y", TensorProto.FLOAT, ["rows", "cols"])], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model, dim_bindings={"batch": 2, "rows": 2, "cols": 3}) + + message = str(ctx.exception) + self.assertIn("reshape", message) + self.assertIn("`s`", message) + self.assertIn("depends on input data", message) + + def test_every_operand_a_view_takes_its_shape_from_is_checked(self): + """Each of these ops reads a shape, an axis list or a repeat count as data. + + Whichever operand carries it, the result's shape is a function of its *values*, so a + model that computes one at run time is rejected by name rather than compiled for + whatever the operand happened to hold. + """ + cases = ( + ("Reshape", ["x", "p"], {}), + ("Expand", ["x", "p"], {}), + ("Tile", ["x", "p"], {}), + ("Squeeze", ["x", "p"], {}), + ("Unsqueeze", ["x", "p"], {}), + ("Split", ["x", "p"], {"axis": 0}), + ("Slice", ["x", "p", "p"], {}), + ) + for op_type, inputs, attributes in cases: + with self.subTest(op_type=op_type): + model = _model( + [ + helper.make_node( + op_type, inputs, ["y"], name="view", **attributes + ) + ], + [ + _tensor("x", TensorProto.FLOAT, [2, 3]), + _tensor("p", TensorProto.INT64, [2]), + ], + [_unshaped("y")], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("view", message) + self.assertIn("`p`", message) + self.assertIn("depends on input data", message) + + def test_every_operand_that_sizes_a_result_is_checked(self): + """These ops read their result's own extents out of an operand's values. + + A padding, a depth, a `k`, the shape a tensor is generated at: none of them is a + shape the graph states, so an operand nothing fixes leaves the result's size a + function of input data — and a buffer the compiler cannot size. + """ + float_result = {"y": TensorProto.FLOAT} + cases = ( + ("ConstantOfShape", ["p"], float_result, {}), + ("Range", ["s", "s", "s"], {"y": TensorProto.INT64}, {}), + ("OneHot", ["i", "s", "v"], float_result, {}), + ("Pad", ["x", "p"], float_result, {}), + ("CenterCropPad", ["x", "p"], float_result, {}), + ( + "TopK", + ["x", "p"], + {"y": TensorProto.FLOAT, "z": TensorProto.INT64}, + {"axis": 0}, + ), + ) + for op_type, inputs, outputs, attributes in cases: + with self.subTest(op_type=op_type): + model = _model( + [ + helper.make_node( + op_type, inputs, list(outputs), name="size", **attributes + ) + ], + [ + _tensor("x", TensorProto.FLOAT, [2, 3]), + _tensor("i", TensorProto.INT64, [2]), + _tensor("p", TensorProto.INT64, [2]), + _tensor("s", TensorProto.INT64, []), + _tensor("v", TensorProto.FLOAT, [2]), + ], + [_unshaped(name, elem_type) for name, elem_type in outputs.items()], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("size", message) + self.assertIn("depends on input data", message) + + def test_every_operand_a_signal_op_takes_its_extents_from_is_checked(self): + """The transforms and the windows state their result's shape through operands too. + + A window is `size` samples and nothing else, a mel matrix is its bin counts, a + short-time transform is one frame per step of `frame_step` each `frame_length` + long, and a `dft_length` is the extent of the axis it transforms — so an operand + nothing fixes leaves a buffer the compiler cannot size. Each is declared with the + result shape it would have had, since ONNX's own inference stops at the same + operands and would otherwise report an unknown shape instead. + """ + cases = ( + ("HannWindow", ["n"], [_tensor("y", TensorProto.FLOAT, [8])], {}), + ("HammingWindow", ["n"], [_tensor("y", TensorProto.FLOAT, [8])], {}), + ("BlackmanWindow", ["n"], [_tensor("y", TensorProto.FLOAT, [8])], {}), + ( + "MelWeightMatrix", + ["n", "n", "n", "f", "f"], + [_tensor("y", TensorProto.FLOAT, [5, 8])], + {}, + ), + ( + "STFT", + ["signal", "s", "", "s"], + [_tensor("y", TensorProto.FLOAT, [1, 3, 5, 2])], + {}, + ), + ("DFT", ["signal", "s"], [_tensor("y", TensorProto.FLOAT, [1, 16, 2])], {}), + ) + for op_type, inputs, outputs, attributes in cases: + with self.subTest(op_type=op_type): + model = _model( + [ + helper.make_node( + op_type, + inputs, + [entry.name for entry in outputs], + name="sized", + **attributes, + ) + ], + [ + _tensor("signal", TensorProto.FLOAT, [1, 16, 1]), + _tensor("n", TensorProto.INT32, []), + _tensor("s", TensorProto.INT64, []), + _tensor("f", TensorProto.FLOAT, []), + ], + outputs, + opset=17, + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("sized", message) + self.assertIn("depends on input data", message) + + def test_a_transform_axis_is_a_shape_operand_only_where_it_resizes_an_axis(self): + """`onesided` halves the axis it lands on, which makes which axis that is part of + the result's shape; without it the result carries the operand's own extents + whichever axis is transformed, and a run-time axis is a value, not a shape.""" + + def transform(**attributes): + return _model( + [ + helper.make_node( + "DFT", ["signal", "", "axis"], ["y"], name="dft", **attributes + ) + ], + [ + _tensor("signal", TensorProto.FLOAT, [1, 8, 1]), + _tensor("axis", TensorProto.INT64, []), + ], + [_unshaped("y")], + opset=MAX_OPSET, + ) + + prepared = _prepare(transform()) + + self.assertEqual(_op_types(prepared), ["DFT"]) + self.assertEqual(_output_shape(prepared, "y"), (1, 8, 2)) + with self.assertRaises(CompileError) as ctx: + _prepare(transform(onesided=1)) + self.assertIn("depends on input data", str(ctx.exception)) + + def test_a_view_shape_an_initializer_fixes_is_accepted(self): + model = _model( + [helper.make_node("Tile", ["x", "repeats"], ["y"], name="tile")], + [_tensor("x", TensorProto.FLOAT, [2, 3])], + [_unshaped("y")], + initializer=[onnx.numpy_helper.from_array(np.array([2, 3]), "repeats")], + ) + + prepared = _prepare(model) + + self.assertEqual( + shapes.static_shape(shapes.tensor_types(prepared.model.graph)["y"]), (4, 9) + ) + + def test_a_squeeze_axes_operand_with_no_elements_is_kept(self): + """An empty axes list squeezes nothing, where an absent one squeezes everything. + + The reductions are the one family ONNX defines the two the same way for, so they are + the only one whose empty operand the frontend drops; dropping this one would turn a + `[1, 3]` result into a `[3]` one. + """ + model = _model( + [helper.make_node("Squeeze", ["x", "axes"], ["y"], name="squeeze")], + [_tensor("x", TensorProto.FLOAT, [1, 3])], + [_unshaped("y")], + initializer=[ + onnx.numpy_helper.from_array(np.array([], dtype=np.int64), "axes") + ], + ) + + prepared = _prepare(model) + + self.assertEqual(list(prepared.model.graph.node[0].input), ["x", "axes"]) + self.assertEqual( + shapes.static_shape(shapes.tensor_types(prepared.model.graph)["y"]), (1, 3) + ) + + def test_reduction_axes_that_only_runtime_data_names_is_rejected(self): + """Which axes a reduction removes decides its output shape, so they must be static.""" + model = _model( + [helper.make_node("ReduceSum", ["x", "axes"], ["y"], name="total")], + [ + _tensor("x", TensorProto.FLOAT, [2, 3]), + _tensor("axes", TensorProto.INT64, [1]), + ], + [_unshaped("y")], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("total", message) + self.assertIn("`axes`", message) + self.assertIn("depends on input data", message) + + def test_reduction_axes_an_initializer_fixes_are_accepted(self): + model = _model( + [helper.make_node("ReduceSum", ["x", "axes"], ["y"], name="total")], + [_tensor("x", TensorProto.FLOAT, [2, 3])], + [_unshaped("y")], + initializer=[onnx.numpy_helper.from_array(np.array([1]), "axes")], + ) + + prepared = _prepare(model) + + self.assertEqual( + shapes.static_shape(shapes.tensor_types(prepared.model.graph)["y"]), (2, 1) + ) + + def test_an_axes_operand_with_no_elements_is_left_out(self): + """An empty axes tensor names no axes, which is what passing none at all means. + + ONNX's shape inference does not reason about the values of an operand it cannot see, + so it leaves the result of such a reduction untyped; dropping the operand is what + makes the shape follow from the graph. + """ + model = _model( + [helper.make_node("ReduceSum", ["x", "axes"], ["y"], name="total")], + [ + _tensor("x", TensorProto.FLOAT, [2, 3]), + _tensor("axes", TensorProto.INT64, [0]), + ], + [_unshaped("y")], + ) + + prepared = _prepare(model) + + self.assertEqual(list(prepared.model.graph.node[0].input), ["x", ""]) + _assert_matches_reference( + self, + model, + prepared, + { + "x": np.arange(6, dtype=np.float32).reshape(2, 3), + "axes": np.array([], dtype=np.int64), + }, + ) + + def test_a_short_time_transform_states_the_onesided_default_it_is_read_under(self): + """ONNX's shape inference reads a default for `onesided` its schema does not declare. + + The schema's default is 1 and the reference evaluator applies it, so the op returns + the non-redundant half of each frame's spectrum; inference falls back to 0 and would + size the result at the whole frame length. Stating the schema's own default leaves + what the node computes alone and makes the two agree. + """ + model = _model( + [helper.make_node("STFT", ["x", "step", "", "length"], ["y"], name="stft")], + [_tensor("x", TensorProto.FLOAT, [1, 16, 1])], + [_unshaped("y")], + initializer=[ + onnx.numpy_helper.from_array(np.array(value, dtype=np.int64), name) + for name, value in (("step", 4), ("length", 8)) + ], + opset=17, + ) + declared = onnx.defs.get_schema("STFT").attributes["onesided"].default_value.i + + prepared = _prepare(model) + + (stated,) = prepared.model.graph.node[0].attribute + self.assertEqual((stated.name, stated.i), ("onesided", declared)) + self.assertEqual(_output_shape(prepared, "y"), (1, 3, 5, 2)) + _assert_matches_reference( + self, + model, + prepared, + {"x": np.arange(16, dtype=np.float32).reshape(1, 16, 1)}, + ) + + def test_shape_inference_failure_names_the_graph_and_node(self): + model = _model( + [helper.make_node("Add", ["x", "z"], ["y"], name="mismatch")], + [ + _tensor("x", TensorProto.FLOAT, [2, 3]), + _tensor("z", TensorProto.FLOAT, [4, 5]), + ], + [_unshaped("y")], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("shape inference failed", message) + self.assertIn("mismatch", message) + + def test_tensor_without_an_inferred_type_is_reported(self): + """`verify_static` refuses a tensor nothing described, rather than assuming one.""" + model = _model( + [helper.make_node("Identity", ["x"], ["y"], name="copy")], + [_tensor("x", TensorProto.FLOAT, [2, 2])], + [_unshaped("y")], + ) + model.graph.output[0].ClearField("type") + + with self.assertRaises(CompileError) as ctx: + verify.verify_static(model) + + message = str(ctx.exception) + self.assertIn("copy", message) + self.assertIn("`y`", message) + + def test_unsupported_element_types_name_the_tensor_and_type(self): + for elem_type in ( + TensorProto.FLOAT16, + TensorProto.BFLOAT16, + TensorProto.STRING, + TensorProto.COMPLEX64, + ): + with self.subTest(elem_type=elem_type): + model = _model( + [helper.make_node("Identity", ["x"], ["y"])], + [_tensor("x", elem_type, [2, 2])], + [_unshaped("y", elem_type)], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("`x`", message) + self.assertIn(onnx.TensorProto.DataType.Name(elem_type), message) + + def test_unsupported_weight_type_names_the_initializer(self): + model = _model( + [ + helper.make_node("CastLike", ["w", "x"], ["wf"], name="cast"), + helper.make_node("Add", ["x", "wf"], ["y"], name="add"), + ], + [_tensor("x", TensorProto.FLOAT, [2, 2])], + [_unshaped("y")], + initializer=[ + onnx.numpy_helper.from_array(np.ones((2, 2), dtype=np.float16), "w") + ], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("initializer `w`", message) + self.assertIn("FLOAT16", message) + + def test_non_tensor_io_is_rejected(self): + sequence = helper.make_value_info( + "s", + helper.make_sequence_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [2]) + ), + ) + model = _model( + [helper.make_node("ConcatFromSequence", ["s"], ["y"], axis=0)], + [sequence], + [_unshaped("y")], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("`s`", message) + self.assertIn("sequence", message) + + def test_sparse_initializers_are_rejected(self): + sparse = helper.make_sparse_tensor( + helper.make_tensor("sv", TensorProto.FLOAT, [1], [1.0]), + helper.make_tensor("si", TensorProto.INT64, [1, 2], [0, 0]), + [2, 2], + ) + model = _model( + [helper.make_node("Identity", ["x"], ["y"])], + [_tensor("x", TensorProto.FLOAT, [2, 2])], + [_unshaped("y")], + sparse_initializer=[sparse], + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + self.assertIn("sparse", str(ctx.exception)) + + def test_control_flow_that_survives_folding_is_rejected(self): + model = _model( + [ + helper.make_node( + "If", + ["cond"], + ["y"], + name="branch", + then_branch=_constant_branch("then", [1.0]), + else_branch=_constant_branch("else", [2.0]), + ) + ], + [_tensor("cond", TensorProto.BOOL, [])], + [_unshaped("y")], + opset=MAX_OPSET, + ) + + with self.assertRaises(CompileError) as ctx: + _prepare(model) + + message = str(ctx.exception) + self.assertIn("branch", message) + self.assertIn("control flow", message) + + def test_zero_element_tensors_are_static(self): + original = _model( + [helper.make_node("Concat", ["a", "b"], ["y"], axis=0)], + [ + _tensor("a", TensorProto.FLOAT, [0, 3]), + _tensor("b", TensorProto.FLOAT, ["batch", 3]), + ], + [_unshaped("y")], + ) + + prepared = _prepare(original, dim_bindings={"batch": 2}) + + self.assertEqual(_output_shape(prepared, "y"), (2, 3)) + _assert_matches_reference( + self, + original, + prepared, + { + "a": np.zeros((0, 3), dtype=np.float32), + "b": np.ones((2, 3), dtype=np.float32), + }, + ) + + +class ShapeInferenceGapTest(unittest.TestCase): + """Where ONNX's own inference stops short of a shape it nonetheless guarantees.""" + + def test_a_body_it_cannot_infer_through_is_retried_without_strict_mode(self): + """Strict inference recurses into the function body ONNX defines for the op. + + MeanVarianceNormalization's builds its `axes` from a Constant that carries nothing + at all unless the node sets the attribute, and inference raises on it — while the + shapes it derives without strict mode are complete, so the model still compiles. + """ + model = _model( + [helper.make_node("MeanVarianceNormalization", ["x"], ["y"], name="mvn")], + [_tensor("x", TensorProto.FLOAT, [2, 3, 2, 2])], + [_unshaped("y")], + opset=13, + ) + + with self.assertRaises(onnx.shape_inference.InferenceError): + onnx.shape_inference.infer_shapes(model, strict_mode=True) + prepared = _prepare(model) + + self.assertEqual(_output_shape(prepared, "y"), (2, 3, 2, 2)) + + def test_group_normalization_gives_its_result_the_shape_of_its_operand(self): + """Inference stops inside GroupNormalization's body, which reshapes through shapes + it computes, leaving its result a rank it never states — and every tensor after it + none either, though the schema says `Y` has the shape of `X`.""" + model = _model( + [ + helper.make_node( + "GroupNormalization", + ["x", "scale", "bias"], + ["h"], + name="norm", + num_groups=2, + ), + helper.make_node("Relu", ["h"], ["y"], name="relu"), + ], + [ + _tensor("x", TensorProto.FLOAT, [2, 4, 3]), + _tensor("scale", TensorProto.FLOAT, [4]), + _tensor("bias", TensorProto.FLOAT, [4]), + ], + [_unshaped("y")], + ) + + prepared = _prepare(model) + + self.assertEqual(_output_shape(prepared, "y"), (2, 4, 3)) + _assert_matches_reference( + self, + model, + prepared, + { + "x": np.arange(24, dtype=np.float32).reshape(2, 4, 3), + "scale": np.arange(4, dtype=np.float32), + "bias": np.ones(4, dtype=np.float32), + }, + ) + + def test_a_shape_only_folding_makes_derivable_is_not_shadowed(self): + """A tensor inference could not size on the first pass takes the shape of the last. + + Folding is what gives inference the operand values it was missing, so the two run in + turn — and ONNX leaves behind an entry naming only the element type for every tensor + it stopped at. Those entries outlive the round that produced them and are read before + the graph's own outputs, so a stale one would hide the shape the next round derives. + """ + model = _model( + [ + helper.make_node("Shape", ["x"], ["s"], name="shape"), + helper.make_node("Sub", ["s", "k"], ["ends"], name="sub"), + helper.make_node("Slice", ["x", "zeros", "ends"], ["y"], name="slice"), + ], + [_tensor("x", TensorProto.FLOAT, [2, 4, 8])], + [_unshaped("y")], + initializer=[ + onnx.numpy_helper.from_array(np.array([0, 1, 3], np.int64), "k"), + onnx.numpy_helper.from_array(np.zeros(3, np.int64), "zeros"), + ], + ) + + prepared = _prepare(model) + + self.assertEqual(_output_shape(prepared, "y"), (2, 3, 5)) + _assert_matches_reference( + self, + model, + prepared, + {"x": np.arange(64, dtype=np.float32).reshape(2, 4, 8)}, + ) + + +class PrepareBundleModelTest(unittest.TestCase): + def test_ml_node_model_takes_its_shape_from_the_declared_output(self): + """`ai.onnx.ml` inference stops at the batch dimension; the declaration carries it.""" + original = onnx.load(str(OPS_ARTIFACTS / "linreg" / "model.onnx")) + + prepared = _prepare(original) + + self.assertEqual( + shapes.static_shape(prepared.model.graph.input[0].type), (1, 3) + ) + self.assertEqual(_output_shape(prepared, "variable"), (1, 1)) + _assert_matches_reference( + self, + original, + prepared, + {"float_input": np.ones((1, 3), dtype=np.float32)}, + ) + + def test_standard_node_model_folds_its_constant_node(self): + original = onnx.load(str(OPS_ARTIFACTS / "concat_reduce" / "model.onnx")) + + prepared = _prepare(original) + + self.assertEqual(_op_types(prepared), ["Concat", "ReduceSum"]) + _assert_matches_reference( + self, + original, + prepared, + { + name: np.full((1, 1), index + 1, dtype=np.float32) + for index, name in enumerate(["input1", "input2", "input3"]) + }, + ) + + def test_preparing_does_not_mutate_the_loaded_model(self): + loaded = loader.load_model(OPS_ARTIFACTS / "concat_reduce" / "model.onnx") + before = loaded.model.SerializeToString() + + frontend.prepare_model(loaded, dim_bindings={"batch": 3}) + + self.assertEqual(loaded.model.SerializeToString(), before) + + def test_preparing_twice_produces_identical_models(self): + loaded = loader.load_model(OPS_ARTIFACTS / "linreg" / "model.onnx") + + first = frontend.prepare_model(loaded, dim_bindings={"batch": 2}) + second = frontend.prepare_model(loaded, dim_bindings={"batch": 2}) + + self.assertEqual( + first.model.SerializeToString(), second.model.SerializeToString() + ) + + +class ElementTypeTest(unittest.TestCase): + def test_every_supported_type_has_a_fixed_width_c_type(self): + expected = { + TensorProto.FLOAT: "float", + TensorProto.DOUBLE: "double", + TensorProto.INT8: "int8_t", + TensorProto.INT16: "int16_t", + TensorProto.INT32: "int32_t", + TensorProto.INT64: "int64_t", + TensorProto.UINT8: "uint8_t", + TensorProto.UINT16: "uint16_t", + TensorProto.UINT32: "uint32_t", + TensorProto.UINT64: "uint64_t", + TensorProto.BOOL: "uint8_t", + } + self.assertEqual(dtypes.C_TYPES, expected) + for elem_type, c_type in expected.items(): + self.assertTrue(dtypes.is_supported(elem_type)) + self.assertEqual(dtypes.c_type(elem_type), c_type) + + def test_unsupported_type_raises_and_names_the_type(self): + self.assertFalse(dtypes.is_supported(TensorProto.FLOAT16)) + with self.assertRaises(CompileError) as ctx: + dtypes.c_type(TensorProto.FLOAT16) + self.assertIn("FLOAT16", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/python/tests/test_extra_compiler_functions.py b/src/python/tests/test_extra_compiler_functions.py new file mode 100644 index 0000000..28ecf19 --- /dev/null +++ b/src/python/tests/test_extra_compiler_functions.py @@ -0,0 +1,454 @@ +"""Function expansion: compiling an op through the body ONNX defines for it. + +Dispatch is native kernel → function expansion → compile error, and this module covers the +middle step. Nothing here decides what an op computes: the bodies come from the `onnx` +schema registry, the values a body is specialized with come from the model's own node, and +what an expanded artifact must produce comes from `onnx.reference.ReferenceEvaluator`. +""" + +from __future__ import annotations + +import shutil +from typing import Any + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +# The harness refuses to import without numpy, so this covers both dependencies. +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from onnx import TensorProto, helper # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +from fnnx.extras.compilers.c import compile_onnx # noqa: E402 +from fnnx.extras.compilers.c.onnx import codegen # noqa: E402 +from fnnx.extras.compilers.c.onnx.frontend import prepare_model # noqa: E402 +from fnnx.extras.compilers.c.onnx.functions import ( # noqa: E402 + Expansion, + expand_function, + function_body, +) +from fnnx.extras.compilers.c.onnx.kernels import KERNELS # noqa: E402 +from fnnx.extras.compilers.c.onnx.loader import LoadedModel, resolve_opsets # noqa: E402 + +# `HardSwish` at 22 is the fallback's smallest complete case: no native kernel serves it, +# ONNX defines it by a function body, and that body — one `HardSigmoid` and one `Mul` — is +# made of ops this compiler does serve. +HARD_SWISH_OPSET = 22 +# `Bernoulli` at 22 is the other side of it: a body built on `RandomUniformLike`, which +# draws at random and can never be compiled into a static artifact — so it stays the +# fallback's failing case however far kernel coverage grows. +BERNOULLI_OPSET = 22 +# `Clip`'s body is context-dependent, which is what the specialization tests need; dispatch +# never reaches it, since a native kernel serves the op. +CLIP_OPSET = 13 +LAYER_NORM_OPSET = 17 +# `CenterCropPad` at 18 is the case for an operand's *value* reaching a body: its own +# reads the extents it crops to out of a tensor, and no kernel serves the op. +CENTER_CROP_PAD_OPSET = 18 +# Older than every registered `Add` kernel, and an opset at which ONNX defines no body for +# the op — so dispatch runs out of options with kernels registered for the op all the same. +LEGACY_ADD_OPSET = 6 + +requires_c_compiler = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + + +def _model(node, inputs, outputs, *, opset, name="expansion"): + graph = helper.make_graph([node], name, list(inputs), list(outputs)) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + + +def _tensor(name, elem_type, shape): + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _type(elem_type, shape) -> Any: + return helper.make_tensor_type_proto(elem_type, list(shape)) + + +def _hard_swish(elem_type=TensorProto.FLOAT, shape=(2, 3)): + return _model( + helper.make_node("HardSwish", ["x"], ["y"], name="swish"), + [_tensor("x", elem_type, list(shape))], + [helper.make_empty_tensor_value_info("y")], + opset=HARD_SWISH_OPSET, + ) + + +def _layer_norm_node(): + """`LayerNormalization` asking only for `Y`, leaving `Mean` and `InvStdDev` out.""" + return helper.make_node("LayerNormalization", ["x", "scale"], ["y"], axis=-1) + + +_LAYER_NORM_TYPES = (_type(TensorProto.FLOAT, (2, 3)), _type(TensorProto.FLOAT, (3,))) + + +# -------------------------------------------------------------------------------------- +# Dispatch: kernel, then body, then error +# -------------------------------------------------------------------------------------- + + +@requires_c_compiler +@pytest.mark.parametrize( + ("elem_type", "dtype"), + [(TensorProto.FLOAT, "float32"), (TensorProto.DOUBLE, "float64")], +) +def test_a_function_defined_op_compiles_and_matches_the_reference( + tmp_path, elem_type, dtype +): + """The scenario itself: no kernel, a body, and output the reference agrees with.""" + model = _hard_swish(elem_type) + values = np.arange(-3, 3, dtype=dtype).reshape(2, 3) + + compiled = compile_onnx(model, tmp_path).load() + outputs = compiled.run({"x": values}) + + expected = ReferenceEvaluator(model).run(None, {"x": values}) + np.testing.assert_array_equal(outputs["y"], expected[0]) + + +@requires_c_compiler +def test_a_context_dependent_body_carries_the_type_castlike_reads_off_its_operand( + tmp_path, +): + """`CastLike` has no kernel of its own, and could not have a fixed one: the type it + converts to is the second operand's, which ONNX writes into the body it builds per call + site. Compiling that body is what serves the op.""" + assert not KERNELS.registered_versions("", "CastLike") + model = _model( + helper.make_node("CastLike", ["x", "like"], ["y"], name="convert"), + [ + _tensor("x", TensorProto.FLOAT, [2, 3]), + _tensor("like", TensorProto.INT16, [1]), + ], + [helper.make_empty_tensor_value_info("y")], + opset=21, + ) + feeds = { + "x": np.array([[1.5, -2.5, 0.0], [3.9, -4.9, 6.0]], dtype="float32"), + "like": np.zeros(1, dtype="int16"), + } + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + expected = ReferenceEvaluator(model).run(None, feeds) + np.testing.assert_array_equal(outputs["y"], expected[0]) + + +def test_a_native_kernel_takes_precedence_over_the_function_body(tmp_path): + """`Relu` has both; the body needs ops no kernel serves, so using it would fail.""" + model = _model( + helper.make_node("Relu", ["x"], ["y"], name="relu"), + [_tensor("x", TensorProto.FLOAT, [2, 3])], + [helper.make_empty_tensor_value_info("y")], + opset=14, + ) + assert onnx.defs.get_schema("Relu", 14, "").has_function + + report = compile_onnx(model, tmp_path).report + + assert [ + kernel + for kernel in report["kernels"] + if kernel.startswith(f"{report['prefix']}_relu_") + ] == report["kernels"] + + +def test_a_kernel_the_registry_cannot_vouch_for_falls_through_to_the_body( + tmp_path, monkeypatch +): + """A selection the semantic-revision guard rejects is not the end of dispatch. + + `Relu` has both a kernel and a body, so a registry that selects nothing for it is what + a guard rejection looks like from the emitter's side. + """ + monkeypatch.setattr(KERNELS, "select", lambda domain, op_type, version: None) + model = _model( + helper.make_node("Relu", ["x"], ["y"], name="relu"), + [_tensor("x", TensorProto.FLOAT, [2, 3])], + [helper.make_empty_tensor_value_info("y")], + opset=14, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "function body of `Relu`" in str(error.value) + + +def test_an_op_with_neither_a_kernel_nor_a_body_is_rejected(tmp_path): + """`RandomUniform` draws rather than computes, which is off the v1 supported surface.""" + model = _model( + helper.make_node("RandomUniform", [], ["y"], name="draw", shape=[2, 3]), + [], + [helper.make_empty_tensor_value_info("y")], + opset=14, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`draw`" in message + assert "`RandomUniform`" in message + assert "ai.onnx" in message + assert "14" in message + + +def test_a_version_no_kernel_and_no_body_covers_keeps_the_dispatch_guidance(tmp_path): + """Falling through to expansion must not cost the error its nearest-version guidance. + + Dispatch declines twice here — no kernel is valid at this opset and ONNX defines no + body — and what surfaces has to be the registry's own error, not the fallback's. + """ + nearest = min(KERNELS.registered_versions("", "Add")) + assert nearest > LEGACY_ADD_OPSET + model = _model( + helper.make_node("Add", ["a", "b"], ["y"], name="adder"), + [_tensor("a", TensorProto.FLOAT, [2]), _tensor("b", TensorProto.FLOAT, [2])], + [helper.make_empty_tensor_value_info("y")], + opset=LEGACY_ADD_OPSET, + ) + output_dir = tmp_path / "out" + + with pytest.raises(CompileError) as error: + compile_onnx(model, output_dir) + + message = str(error.value) + assert "`adder`" in message + assert "`Add`" in message + assert "ai.onnx" in message + assert f"opset version {LEGACY_ADD_OPSET}" in message + assert f"Nearest supported version: {nearest}" in message + assert not output_dir.exists() + + +def test_an_op_whose_body_needs_a_missing_kernel_is_rejected(tmp_path): + """The error names the model's own node and the primitive the body ran aground on.""" + model = _model( + helper.make_node("Bernoulli", ["x"], ["y"], name="draw"), + [_tensor("x", TensorProto.FLOAT, [2, 3])], + [helper.make_empty_tensor_value_info("y")], + opset=BERNOULLI_OPSET, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`draw`" in message + assert "`Bernoulli`" in message + assert "`RandomUniformLike`" in message + + +@requires_c_compiler +def test_a_body_made_of_function_defined_ops_expands_recursively(tmp_path, monkeypatch): + """`HardSwish`'s body calls `HardSigmoid`, which is function-defined in turn. + + With no kernel serving `HardSigmoid` — what a semantic-revision guard rejection looks + like — the inner node has to expand through its own body, so the artifact is built out + of that body's primitives rather than out of a `HardSigmoid` kernel, and still has to + compute what the reference computes. + """ + select = KERNELS.select + monkeypatch.setattr( + KERNELS, + "select", + lambda domain, op_type, version: ( + None if op_type == "HardSigmoid" else select(domain, op_type, version) + ), + ) + model = _hard_swish() + values = np.arange(-3, 3, dtype="float32").reshape(2, 3) + + result = compile_onnx(model, tmp_path) + outputs = result.load().run({"x": values}) + + assert not [ + kernel for kernel in result.report["kernels"] if "hardsigmoid" in kernel + ] + expected = ReferenceEvaluator(model).run(None, {"x": values}) + np.testing.assert_array_equal(outputs["y"], expected[0]) + + +def test_expansion_stops_at_a_bounded_nesting_depth(tmp_path, monkeypatch): + """A body that never bottoms out is a compile error, not a stack overflow.""" + monkeypatch.setattr(codegen, "MAX_EXPANSION_DEPTH", 0) + + with pytest.raises(CompileError, match="nested more than"): + compile_onnx(_hard_swish(), tmp_path) + + +# -------------------------------------------------------------------------------------- +# The body a node expands to +# -------------------------------------------------------------------------------------- + + +def test_the_body_is_specialized_to_the_operands_the_node_passes(): + """A context-dependent body sees which optional operands the node actually has.""" + float_type = _type(TensorProto.FLOAT, (2, 3)) + scalar = _type(TensorProto.FLOAT, ()) + unbounded = helper.make_node("Clip", ["x", "", ""], ["y"]) + bounded = helper.make_node("Clip", ["x", "lo", "hi"], ["y"]) + + without = expand_function(unbounded, "", CLIP_OPSET, [float_type, None, None]) + with_bounds = expand_function(bounded, "", CLIP_OPSET, [float_type, scalar, scalar]) + + assert [node.op_type for node in without.prepared.model.graph.node] == ["Identity"] + assert without.inputs == (("input", 0),) + assert with_bounds.inputs == (("input", 0), ("min", 1), ("max", 2)) + assert len(with_bounds.prepared.model.graph.node) > 1 + + +def test_an_attribute_the_node_sets_reaches_the_body(): + """`LeakyRelu`'s body reads `alpha` through the caller's node.""" + node = helper.make_node("LeakyRelu", ["x"], ["y"], alpha=0.25) + + expansion = expand_function(node, "", 16, [_type(TensorProto.FLOAT, (2, 3))]) + + assert 0.25 in _folded_values(expansion) + + +def test_an_attribute_the_node_omits_falls_back_on_the_schema_default(): + """Without the default filled in, the body's `Constant` would have no value at all.""" + node = helper.make_node("LeakyRelu", ["x"], ["y"]) + default = ( + onnx.defs.get_schema("LeakyRelu", 16, "").attributes["alpha"].default_value + ) + + expansion = expand_function(node, "", 16, [_type(TensorProto.FLOAT, (2, 3))]) + + assert np.float32(default.f) in _folded_values(expansion) + + +def test_the_body_of_an_omitted_optional_output_is_not_compiled(): + """A body computes every output its op declares; the caller pays only for the ones it + asked for. Asserted as the property pruning establishes: nothing left in the compiled + body fails to reach an output the node wants.""" + graph = expand_function( + _layer_norm_node(), "", LAYER_NORM_OPSET, list(_LAYER_NORM_TYPES) + ).prepared.model.graph + + assert [entry.name for entry in graph.output] == ["Y"] + live = {entry.name for entry in graph.output} + for node in reversed(graph.node): + assert any(name in live for name in node.output), ( + f"`{node.op_type}` computes nothing the expanded node asked for" + ) + live |= {name for name in node.input if name} + + +def test_the_unpruned_body_does_compute_the_omitted_outputs(): + """Teeth for the test above: ONNX's body really is bigger than what gets compiled.""" + node = _layer_norm_node() + + body = function_body(node, "", LAYER_NORM_OPSET, list(_LAYER_NORM_TYPES)) + compiled = expand_function( + node, "", LAYER_NORM_OPSET, list(_LAYER_NORM_TYPES) + ).prepared.model.graph + + assert list(body.output) == ["Y", "Mean", "InvStdDev"] + assert len(compiled.node) < len( + [entry for entry in body.node if entry.op_type != "Constant"] + ) + + +def test_an_op_onnx_defines_no_body_for_has_no_expansion(): + node = helper.make_node("Sub", ["a", "b"], ["y"]) + + assert function_body(node, "", 14, [None, None]) is None + assert expand_function(node, "", 14, [None, None]) is None + + +def test_an_op_the_installed_onnx_does_not_define_has_no_expansion(): + node = helper.make_node("NotAnOnnxOp", ["a"], ["y"]) + + assert expand_function(node, "", 14, [None]) is None + + +def _folded_values(expansion: Expansion) -> list[Any]: + """Every scalar the prepared body holds as constant data.""" + return [ + onnx.numpy_helper.to_array(initializer).reshape(-1)[0] + for initializer in expansion.prepared.model.graph.initializer + if onnx.numpy_helper.to_array(initializer).size == 1 + ] + + +# -------------------------------------------------------------------------------------- +# Splicing a body into the caller's buffers +# -------------------------------------------------------------------------------------- + + +@requires_c_compiler +def test_a_body_output_no_node_writes_still_reaches_the_callers_buffer( + tmp_path, monkeypatch +): + """A body output that folding resolves to a constant is written by nothing; the value + still has to land in the buffer the caller passed for it.""" + values = np.array([[1.5, 2.5, 3.5], [4.5, 5.5, 6.5]], dtype=np.float32) + body = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Constant", + [], + ["output"], + value=onnx.numpy_helper.from_array(values, "value"), + ) + ], + "constant_body", + [_tensor("input", TensorProto.FLOAT, [2, 3])], + [helper.make_empty_tensor_value_info("output")], + ), + opset_imports=[helper.make_opsetid("", CLIP_OPSET)], + ) + expansion = Expansion( + prepared=prepare_model(LoadedModel(model=body, opsets=resolve_opsets(body))), + inputs=(("input", 0),), + outputs=(("output", 0),), + ) + monkeypatch.setattr(codegen, "expand_function", lambda *_: expansion) + + compiled = compile_onnx(_hard_swish(), tmp_path).load() + outputs = compiled.run({"x": np.zeros((2, 3), dtype=np.float32)}) + + np.testing.assert_array_equal(outputs["y"], values) + + +@requires_c_compiler +def test_an_operand_the_graph_fixes_reaches_the_body(tmp_path): + """A body that computes its own result shape from an operand needs the operand itself. + + `CenterCropPad` is defined by a body that pads and then slices by extents it derives from + the `shape` tensor; with only that tensor's type, the `Pad` inside takes a shape no + folding can settle and the body is refused. The values the caller's graph fixes therefore + travel with the types, and here they are what makes the op compilable at all. + """ + extents = np.array([9, 5], dtype=np.int64) + node = helper.make_node( + "CenterCropPad", ["x", "shape"], ["y"], name="crop", axes=[0, 1] + ) + graph = helper.make_graph( + [node], + "center_crop_pad", + [_tensor("x", TensorProto.FLOAT, [10, 7, 3])], + [helper.make_empty_tensor_value_info("y")], + initializer=[onnx.numpy_helper.from_array(extents, "shape")], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", CENTER_CROP_PAD_OPSET)] + ) + feeds = {"x": np.arange(210, dtype=np.float32).reshape(10, 7, 3)} + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + assert not KERNELS.registered_versions("", "CenterCropPad") + expected = ReferenceEvaluator(model).run(None, feeds) + np.testing.assert_array_equal(outputs["y"], expected[0]) diff --git a/src/python/tests/test_extra_compiler_harness.py b/src/python/tests/test_extra_compiler_harness.py new file mode 100644 index 0000000..c443cad --- /dev/null +++ b/src/python/tests/test_extra_compiler_harness.py @@ -0,0 +1,797 @@ +"""The load-and-run harness, and the starter kernels driven end to end through it. + +Every expected value comes from the ONNX reference evaluator — the executable form of the +spec — never from a hand-written expectation. The only hand-written C here is the stub +artifact used to cover harness behaviour (status codes, per-node entrypoints, the strict +build flags) that the emitted artifacts cannot yet exercise. +""" + +from __future__ import annotations + +import importlib +import json +import re +import shutil +import sys +from pathlib import Path +from typing import Any +from unittest import mock + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError, HarnessError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +codegen = pytest.importorskip("fnnx.extras.compilers.c.onnx.codegen") +frontend = pytest.importorskip("fnnx.extras.compilers.c.onnx.frontend") +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from fnnx.extras.compilers.c import compile_onnx, load_compiled # noqa: E402 +from onnx import TensorProto, helper # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +OPSET = 21 +SEED = 20260725 + +# Building the artifact is the point of this module, so the whole file needs a compiler. +pytestmark = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + + +def _tensor(name: str, elem_type: int, shape) -> Any: + return helper.make_tensor_value_info(name, elem_type, list(shape)) + + +def _model(nodes, inputs, outputs, *, initializer=(), name="graph", opset=OPSET): + graph = helper.make_graph( + nodes, name, list(inputs), list(outputs), initializer=list(initializer) + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + + +def _elem_type(dtype) -> int: + return helper.np_dtype_to_tensor_dtype(np.dtype(dtype)) + + +def _values(shape, dtype, *, seed: int = SEED): + """Seeded inputs; integers stay small so C's undefined signed overflow is never hit.""" + generator = np.random.default_rng(seed) + if np.issubdtype(np.dtype(dtype), np.floating): + return generator.normal(size=shape).astype(dtype) + if np.dtype(dtype) == np.bool_: + return generator.integers(0, 2, size=shape).astype(dtype) + info = np.iinfo(dtype) + low, high = max(info.min, -100), min(info.max, 100) + return generator.integers(low, high, size=shape, endpoint=True).astype(dtype) + + +def _reference(model, feeds: dict[str, Any]) -> dict[str, Any]: + outputs = ReferenceEvaluator(model).run(None, dict(feeds)) + return dict(zip([entry.name for entry in model.graph.output], outputs)) + + +def _assert_matches(actual: dict[str, Any], expected: dict[str, Any]) -> None: + assert sorted(actual) == sorted(expected) + for name, want in expected.items(): + got = actual[name] + assert got.dtype == want.dtype, name + assert got.shape == want.shape, name + if np.issubdtype(want.dtype, np.floating): + np.testing.assert_allclose(got, want, rtol=1e-6, atol=1e-6, err_msg=name) + else: + np.testing.assert_array_equal(got, want, err_msg=name) + + +def _run_against_reference(model, feeds, tmp_path): + compiled = compile_onnx(model, tmp_path).load() + _assert_matches(compiled.run(feeds), _reference(model, feeds)) + return compiled + + +def _pipeline_model(name="demo"): + """Gemm -> Add -> Relu -> Identity: every starter kernel in one graph.""" + weights = (np.arange(12, dtype=np.float32).reshape(4, 3) - 5.0) / 7.0 + intercept = np.array([0.25, -0.5, 1.0, 2.0], dtype=np.float32) + offset = np.array([[0.1], [-0.2]], dtype=np.float32) + return _model( + [ + helper.make_node( + "Gemm", ["x", "w", "b"], ["h"], name="gemm", alpha=0.5, transB=1 + ), + helper.make_node("Add", ["h", "offset"], ["s"], name="add"), + helper.make_node("Relu", ["s"], ["r"], name="relu"), + helper.make_node("Identity", ["r"], ["y"], name="identity"), + ], + [_tensor("x", TensorProto.FLOAT, [2, 3])], + [_tensor("y", TensorProto.FLOAT, [2, 4])], + initializer=[ + onnx.numpy_helper.from_array(weights, "w"), + onnx.numpy_helper.from_array(intercept, "b"), + onnx.numpy_helper.from_array(offset, "offset"), + ], + name=name, + ) + + +def _symbolic_batch_model(): + """The same layer over a symbolic batch, so a binding decides every buffer's extent.""" + weights = (np.arange(12, dtype=np.float32).reshape(4, 3) - 5.0) / 7.0 + return _model( + [ + helper.make_node("Gemm", ["x", "w"], ["h"], name="gemm", transB=1), + helper.make_node("Relu", ["h"], ["y"], name="relu"), + ], + [_tensor("x", TensorProto.FLOAT, ["batch", 3])], + [_tensor("y", TensorProto.FLOAT, ["batch", 4])], + initializer=[onnx.numpy_helper.from_array(weights, "w")], + ) + + +def _external_weight_model(weights, location: str): + """A layer whose weight lives in a side file next to the model rather than inside it.""" + tensor = onnx.numpy_helper.from_array(weights, "w") + onnx.external_data_helper.set_external_data(tensor, location=location) + tensor.ClearField("raw_data") + return _model( + [helper.make_node("Gemm", ["x", "w"], ["y"], name="gemm", transB=1)], + [_tensor("x", TensorProto.FLOAT, [2, 3])], + [_tensor("y", TensorProto.FLOAT, [2, 4])], + initializer=[tensor], + name="external", + ) + + +# -------------------------------------------------------------------------------------- +# Loading and running an emitted artifact +# -------------------------------------------------------------------------------------- + + +def test_repeated_runs_match_the_reference_and_carry_no_state(tmp_path): + model = _pipeline_model() + compiled = compile_onnx(model, tmp_path).load() + first = {"x": _values((2, 3), np.float32, seed=1)} + second = {"x": _values((2, 3), np.float32, seed=2)} + + first_outputs = compiled.run(first) + second_outputs = compiled.run(second) + repeat_outputs = compiled.run(first) + + _assert_matches(first_outputs, _reference(model, first)) + _assert_matches(second_outputs, _reference(model, second)) + _assert_matches(repeat_outputs, _reference(model, first)) + + +def test_a_bound_dimension_sizes_the_code_the_artifact_runs(tmp_path): + """A binding is not only report metadata: the emitted code has to compute at that size. + + The report and the macros are checked where they are emitted; what is checked here is + the half a buffer sized from the pre-binding shape would pass — that the artifact run at + the bound size produces what the spec says it should. + """ + model = _symbolic_batch_model() + result = compile_onnx(model, tmp_path, dim_bindings={"batch": 4}) + feeds = {"x": _values((4, 3), np.float32)} + + outputs = result.load().run(feeds) + + assert result.report["entrypoint"]["inputs"][0]["shape"] == [4, 3] + _assert_matches(outputs, _reference(model, feeds)) + + +def test_an_external_data_model_runs_without_its_side_file(tmp_path): + """External tensors are resolved at compile time, so the artifact never reads them again. + + The side file is deleted after the build and before the run, which is what makes "no + runtime file access" an assertion rather than a claim. + """ + weights = _values((4, 3), np.float32) + side_file = tmp_path / "w.bin" + side_file.write_bytes(weights.tobytes()) + model_path = tmp_path / "model.onnx" + onnx.save_model(_external_weight_model(weights, side_file.name), str(model_path)) + feeds = {"x": _values((2, 3), np.float32, seed=SEED + 1)} + expected = _reference(onnx.load(str(model_path)), feeds) + + compiled = compile_onnx(model_path, tmp_path / "out").load() + side_file.unlink() + + _assert_matches(compiled.run(feeds), expected) + + +def test_load_compiled_accepts_the_header_or_the_report(tmp_path): + result = compile_onnx(_pipeline_model(), tmp_path) + feeds = {"x": _values((2, 3), np.float32)} + + from_header = load_compiled(result.header_path).run(feeds) + from_report = load_compiled(result.report_path).run(feeds) + + np.testing.assert_array_equal(from_header["y"], from_report["y"]) + + +def test_the_artifact_directory_keeps_only_the_emitted_files(tmp_path): + result = compile_onnx(_pipeline_model(), tmp_path) + result.load() + + assert sorted(path.name for path in tmp_path.iterdir()) == [ + "demo.h", + "demo_report.json", + ] + + +def test_compiling_a_kernel_graph_twice_is_byte_identical(tmp_path): + """Kernel dedup and registry lookup are the new orderings determinism rests on.""" + first = compile_onnx(_pipeline_model(), tmp_path / "first") + second = compile_onnx(_pipeline_model(), tmp_path / "second") + + assert first.header_path.read_bytes() == second.header_path.read_bytes() + + +def test_the_starter_kernels_allocate_nothing(tmp_path): + """Every other test builds under `-Werror=vla`; the tokens are checked here.""" + header = compile_onnx(_pipeline_model(), tmp_path).header_path.read_text() + + for token in ("malloc", "calloc", "realloc", "free", "alloca"): + assert not re.search(rf"\b{token}\b", header), token + + +def test_metadata_comes_from_the_compile_report(tmp_path): + compiled = compile_onnx(_pipeline_model(), tmp_path).load() + + assert compiled.inputs == (harness.TensorSpec("x", np.dtype("float32"), (2, 3)),) + assert compiled.outputs == (harness.TensorSpec("y", np.dtype("float32"), (2, 4)),) + assert compiled.node_ids == () + + +def test_inputs_may_be_named_or_passed_as_a_mapping(tmp_path): + """Tensor names are not always Python identifiers, so both forms have to work.""" + model = _model( + [helper.make_node("Relu", ["in.1"], ["out.1"], name="relu")], + [_tensor("in.1", TensorProto.FLOAT, [3])], + [_tensor("out.1", TensorProto.FLOAT, [3])], + ) + compiled = compile_onnx(model, tmp_path).load() + values = _values((3,), np.float32) + + _assert_matches(compiled.run({"in.1": values}), _reference(model, {"in.1": values})) + with pytest.raises(HarnessError, match="unexpected `x`"): + compiled.run(x=values) + + +def test_a_name_given_twice_is_rejected(tmp_path): + compiled = compile_onnx(_pipeline_model(), tmp_path).load() + values = _values((2, 3), np.float32) + + with pytest.raises(HarnessError, match="both in the mapping and as keyword"): + compiled.run({"x": values}, x=values) + + +def test_non_contiguous_inputs_are_accepted(tmp_path): + model = _model( + [helper.make_node("Relu", ["x"], ["y"], name="relu")], + [_tensor("x", TensorProto.FLOAT, [2, 3])], + [_tensor("y", TensorProto.FLOAT, [2, 3])], + ) + compiled = compile_onnx(model, tmp_path).load() + view = _values((3, 2), np.float32).T + + assert not view.flags["C_CONTIGUOUS"] + _assert_matches(compiled.run(x=view), _reference(model, {"x": view})) + + +def _forbid_the_c_call(compiled, monkeypatch) -> None: + def explode(*arguments): + raise AssertionError("the C entrypoint must not be called") + + monkeypatch.setattr(compiled._entry, "call", explode) + + +@pytest.mark.parametrize( + ("value", "expected_message"), + [ + ( + np.zeros((2, 3), dtype=np.float64), + r"input `x` has dtype `float64`.*`float32`", + ), + (np.zeros((3, 3), dtype=np.float32), r"input `x` has shape \(3, 3\).*\(2, 3\)"), + (np.zeros((2, 3), dtype=np.int32), r"input `x` has dtype `int32`.*`float32`"), + ], +) +def test_mismatched_inputs_are_rejected_before_the_c_call( + tmp_path, monkeypatch, value, expected_message +): + compiled = compile_onnx(_pipeline_model(), tmp_path).load() + _forbid_the_c_call(compiled, monkeypatch) + + with pytest.raises(HarnessError, match=expected_message): + compiled.run(x=value) + + +def test_missing_and_unexpected_inputs_are_named(tmp_path, monkeypatch): + compiled = compile_onnx(_pipeline_model(), tmp_path).load() + _forbid_the_c_call(compiled, monkeypatch) + + with pytest.raises(HarnessError, match="missing `x` and unexpected `wrong`"): + compiled.run(wrong=_values((2, 3), np.float32)) + + +# -------------------------------------------------------------------------------------- +# Harness behaviour the emitted artifacts cannot exercise yet +# -------------------------------------------------------------------------------------- + +_STUB_HEADER = """\ +#ifndef STUB_H_INCLUDED +#define STUB_H_INCLUDED + +int stub_run(const float* x, float* y); +int stub_node_double_run(const float* x, float* y); + +#endif /* STUB_H_INCLUDED */ + +#ifdef STUB_IMPLEMENTATION + +int stub_run(const float* x, float* y) +{ + (void)x; + (void)y; + return 7; +} + +int stub_node_double_run(const float* x, float* y) +{ + y[0] = x[0] * 2.0f; + return 0; +} + +#endif /* STUB_IMPLEMENTATION */ +""" + + +def _stub_artifact(tmp_path: Path, *, header: str = _STUB_HEADER, **overrides) -> Path: + """A hand-written artifact standing in for one the bundle layer will emit later.""" + scalar = [{"name": "x", "dtype": "float32", "shape": [1]}] + result = [{"name": "y", "dtype": "float32", "shape": [1]}] + report = { + "prefix": "stub", + "header": "stub.h", + "entrypoint": {"symbol": "stub_run", "inputs": scalar, "outputs": result}, + "nodes": [ + { + "id": "double", + "symbol": "stub_node_double_run", + "inputs": scalar, + "outputs": result, + } + ], + } + report.update(overrides) + (tmp_path / "stub.h").write_text(header, encoding="utf-8") + report_path = tmp_path / "stub_report.json" + report_path.write_text(json.dumps(report), encoding="utf-8") + return report_path + + +def test_a_nonzero_status_becomes_an_exception(tmp_path): + compiled = load_compiled(_stub_artifact(tmp_path)) + + with pytest.raises(HarnessError, match="`stub_run` returned status 7"): + compiled.run(x=np.ones(1, dtype=np.float32)) + + +def test_node_entrypoints_are_bound_from_the_report(tmp_path): + compiled = load_compiled(_stub_artifact(tmp_path)) + + outputs = compiled.run_node("double", {"x": np.array([1.5], dtype=np.float32)}) + + assert compiled.node_ids == ("double",) + np.testing.assert_array_equal(outputs["y"], np.array([3.0], dtype=np.float32)) + + +def test_node_inputs_are_validated_like_the_model_inputs(tmp_path): + compiled = load_compiled(_stub_artifact(tmp_path)) + + with pytest.raises( + HarnessError, match=r"Node `double`: input `x` has shape \(2,\)" + ): + compiled.run_node("double", {"x": np.ones(2, dtype=np.float32)}) + + +def test_an_unknown_node_id_lists_the_available_ones(tmp_path): + compiled = load_compiled(_stub_artifact(tmp_path)) + + with pytest.raises( + HarnessError, match="no entrypoint for node `missing`.*`double`" + ): + compiled.run_node("missing", {}) + + +def test_an_artifact_without_node_entrypoints_says_so(tmp_path): + compiled = compile_onnx(_pipeline_model(), tmp_path).load() + + with pytest.raises(HarnessError, match="no entrypoint for node `gemm`.*none"): + compiled.run_node("gemm", {}) + + +def test_a_symbol_the_library_lacks_is_reported(tmp_path): + report_path = _stub_artifact( + tmp_path, + entrypoint={ + "symbol": "stub_absent_run", + "inputs": [], + "outputs": [], + }, + ) + + with pytest.raises(HarnessError, match="exports no symbol `stub_absent_run`"): + load_compiled(report_path) + + +def test_the_artifact_is_built_under_the_strict_flags(tmp_path): + """An unused parameter only fails the build because of `-Wextra -Werror`.""" + lax_header = _STUB_HEADER.replace(" (void)x;\n", "") + + with pytest.raises(HarnessError, match="unused parameter"): + load_compiled(_stub_artifact(tmp_path, header=lax_header)) + + +def test_a_missing_header_is_reported(tmp_path): + report_path = _stub_artifact(tmp_path) + (tmp_path / "stub.h").unlink() + + with pytest.raises( + HarnessError, match="stub.h` the compile report names is missing" + ): + load_compiled(report_path) + + +@pytest.mark.parametrize( + ("name", "expected_message"), + [ + ("model.onnx", "neither a generated header"), + ("absent_report.json", "Compile report not found"), + ("absent.h", "Compile report not found"), + ], +) +def test_paths_that_are_not_an_artifact_are_rejected(tmp_path, name, expected_message): + with pytest.raises(HarnessError, match=expected_message): + load_compiled(tmp_path / name) + + +def test_a_report_that_is_not_ours_is_rejected(tmp_path): + report_path = tmp_path / "other_report.json" + report_path.write_text(json.dumps({"prefix": "other"}), encoding="utf-8") + + with pytest.raises(HarnessError, match="missing the `header`, `entrypoint` field"): + load_compiled(report_path) + + +def test_an_unknown_compiler_is_reported(tmp_path): + result = compile_onnx(_pipeline_model(), tmp_path) + + with pytest.raises(HarnessError, match="`not-a-real-compiler` was not found"): + result.load(compiler="not-a-real-compiler") + + +def test_a_missing_compiler_is_reported(tmp_path, monkeypatch): + result = compile_onnx(_pipeline_model(), tmp_path) + monkeypatch.setattr(harness.shutil, "which", lambda name: None) + + with pytest.raises(HarnessError, match="No system C compiler was found"): + result.load() + + +def test_a_missing_numpy_raises_an_actionable_error(monkeypatch): + """A `ModuleNotFoundError` keeps `pytest.importorskip` skipping rather than erroring.""" + module = "fnnx.extras.compilers.c.harness" + monkeypatch.delitem(sys.modules, module) + + with mock.patch("importlib.util.find_spec", return_value=None): + with pytest.raises(ModuleNotFoundError, match=r"numpy.*fnnx\[core\]"): + importlib.import_module(module) + + +def test_the_cc_environment_variable_is_preferred(tmp_path, monkeypatch): + result = compile_onnx(_pipeline_model(), tmp_path) + monkeypatch.setenv("CC", "cc-from-the-environment") + monkeypatch.setattr( + harness.shutil, + "which", + lambda name: f"/usr/bin/{name}" if "env" in name else None, + ) + + with pytest.raises(HarnessError, match="cc-from-the-environment"): + result.load() + + +# -------------------------------------------------------------------------------------- +# Starter kernels +# -------------------------------------------------------------------------------------- + +_BINARY_SHAPES = [ + ((2, 3), (2, 3)), + ((2, 3), (3,)), + ((1, 3), (2, 1)), + ((2, 3), ()), + ((2, 1, 3), (4, 3)), + ((), ()), + ((0, 3), (3,)), +] + + +@pytest.mark.parametrize("op_type", ["Add", "Mul"]) +@pytest.mark.parametrize(("left_shape", "right_shape"), _BINARY_SHAPES) +def test_binary_ops_broadcast_like_the_reference( + tmp_path, op_type, left_shape, right_shape +): + output_shape = np.broadcast_shapes(left_shape, right_shape) + model = _model( + [helper.make_node(op_type, ["a", "b"], ["y"], name="op")], + [ + _tensor("a", TensorProto.FLOAT, left_shape), + _tensor("b", TensorProto.FLOAT, right_shape), + ], + [_tensor("y", TensorProto.FLOAT, output_shape)], + ) + feeds = { + "a": _values(left_shape, np.float32, seed=1), + "b": _values(right_shape, np.float32, seed=2), + } + + _run_against_reference(model, feeds, tmp_path) + + +@pytest.mark.parametrize( + "dtype", [np.float32, np.float64, np.int32, np.int64, np.uint8, np.int8] +) +def test_binary_ops_cover_every_supported_family(tmp_path, dtype): + elem_type = _elem_type(dtype) + model = _model( + [helper.make_node("Add", ["a", "b"], ["y"], name="op")], + [_tensor("a", elem_type, (2, 3)), _tensor("b", elem_type, (3,))], + [_tensor("y", elem_type, (2, 3))], + ) + feeds = { + "a": _values((2, 3), dtype, seed=1), + "b": _values((3,), dtype, seed=2), + } + + _run_against_reference(model, feeds, tmp_path) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_relu_handles_nan_infinity_and_signed_zero(tmp_path, dtype): + info = np.finfo(dtype) + values = np.array( + [ + 0.0, + -0.0, + np.nan, + np.inf, + -np.inf, + info.max, + -info.max, + info.tiny, + -info.tiny, + info.smallest_subnormal, + -1.5, + 1.5, + ], + dtype=dtype, + ) + elem_type = _elem_type(dtype) + model = _model( + [helper.make_node("Relu", ["x"], ["y"], name="relu")], + [_tensor("x", elem_type, values.shape)], + [_tensor("y", elem_type, values.shape)], + ) + + compiled = _run_against_reference(model, {"x": values}, tmp_path) + outputs = compiled.run(x=values) + + # `allclose` reads -0.0 and 0.0 as equal, so the sign of every zero is checked too. + expected = _reference(model, {"x": values})["y"] + assert np.array_equal(np.signbit(outputs["y"]), np.signbit(expected)) + + +@pytest.mark.parametrize("dtype", [np.int32, np.int64, np.int8]) +def test_relu_clamps_integers(tmp_path, dtype): + elem_type = _elem_type(dtype) + values = _values((3, 4), dtype) + model = _model( + [helper.make_node("Relu", ["x"], ["y"], name="relu")], + [_tensor("x", elem_type, values.shape)], + [_tensor("y", elem_type, values.shape)], + ) + + _run_against_reference(model, {"x": values}, tmp_path) + + +@pytest.mark.parametrize("transpose_a", [0, 1]) +@pytest.mark.parametrize("transpose_b", [0, 1]) +@pytest.mark.parametrize( + ("alpha", "beta", "bias_shape"), + [ + (1.0, 1.0, (4,)), + (0.5, 2.0, (2, 4)), + (1.5, 0.0, (2, 4)), + (0.25, 3.0, ()), + (2.0, 1.0, None), + ], +) +def test_gemm_covers_transposes_scaling_and_bias( + tmp_path, transpose_a, transpose_b, alpha, beta, bias_shape +): + left_shape = (3, 2) if transpose_a else (2, 3) + right_shape = (4, 3) if transpose_b else (3, 4) + inputs = ["a", "b"] if bias_shape is None else ["a", "b", "c"] + model = _model( + [ + helper.make_node( + "Gemm", + inputs, + ["y"], + name="gemm", + alpha=alpha, + beta=beta, + transA=transpose_a, + transB=transpose_b, + ) + ], + [ + _tensor("a", TensorProto.FLOAT, left_shape), + _tensor("b", TensorProto.FLOAT, right_shape), + *( + [] + if bias_shape is None + else [_tensor("c", TensorProto.FLOAT, bias_shape)] + ), + ], + [_tensor("y", TensorProto.FLOAT, (2, 4))], + ) + feeds = { + "a": _values(left_shape, np.float32, seed=1), + "b": _values(right_shape, np.float32, seed=2), + } + if bias_shape is not None: + feeds["c"] = _values(bias_shape, np.float32, seed=3) + + _run_against_reference(model, feeds, tmp_path) + + +def test_gemm_scales_integer_operands_like_the_reference(tmp_path): + """Integer Gemm scales through float64 and truncates, as the reference does.""" + model = _model( + [ + helper.make_node( + "Gemm", ["a", "b", "c"], ["y"], name="gemm", alpha=0.5, beta=1.5 + ) + ], + [ + _tensor("a", TensorProto.INT32, (2, 3)), + _tensor("b", TensorProto.INT32, (3, 4)), + _tensor("c", TensorProto.INT32, (4,)), + ], + [_tensor("y", TensorProto.INT32, (2, 4))], + ) + feeds = { + "a": _values((2, 3), np.int32, seed=1), + "b": _values((3, 4), np.int32, seed=2), + "c": _values((4,), np.int32, seed=3), + } + + _run_against_reference(model, feeds, tmp_path) + + +def test_gemm_rejects_operands_that_are_not_matrices(): + """ONNX shape inference rejects such a graph first, so codegen is driven directly. + + The kernel must still refuse rather than fail on an unpacking error. + """ + model = _model( + [helper.make_node("Gemm", ["a", "b"], ["y"], name="gemm")], + [ + _tensor("a", TensorProto.FLOAT, (2, 3, 4)), + _tensor("b", TensorProto.FLOAT, (3, 4)), + ], + [_tensor("y", TensorProto.FLOAT, (2, 4))], + ) + prepared = frontend.PreparedModel(model=model, opsets={"": OPSET}, dim_bindings={}) + + with pytest.raises(CompileError, match="Gemm takes 2-D operands"): + codegen.build_program(prepared) + + +@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.uint8, np.bool_]) +def test_identity_round_trips_every_element_family(tmp_path, dtype): + elem_type = _elem_type(dtype) + values = _values((2, 3), dtype) + model = _model( + [helper.make_node("Identity", ["x"], ["y"], name="identity")], + [_tensor("x", elem_type, values.shape)], + [_tensor("y", elem_type, values.shape)], + ) + + _run_against_reference(model, {"x": values}, tmp_path) + + +def test_zero_element_tensors_flow_through_the_kernels(tmp_path): + model = _model( + [ + helper.make_node("Add", ["x", "b"], ["s"], name="add"), + helper.make_node("Relu", ["s"], ["r"], name="relu"), + helper.make_node("Identity", ["r"], ["y"], name="identity"), + ], + [ + _tensor("x", TensorProto.FLOAT, (0, 3)), + _tensor("b", TensorProto.FLOAT, (3,)), + ], + [_tensor("y", TensorProto.FLOAT, (0, 3))], + ) + feeds = { + "x": np.zeros((0, 3), dtype=np.float32), + "b": _values((3,), np.float32), + } + + _run_against_reference(model, feeds, tmp_path) + + +def test_a_zero_element_intermediate_no_statement_names_still_builds(tmp_path): + """The chained copies emit nothing at all, leaving `t` named by no statement. + + An intermediate the implementation declares but never mentions is a `static` the C + compiler rejects under the strict flags, so it must not be declared either. + """ + model = _model( + [ + helper.make_node("Identity", ["x"], ["t"], name="first"), + helper.make_node("Identity", ["t"], ["y"], name="second"), + ], + [_tensor("x", TensorProto.FLOAT, (0, 3))], + [_tensor("y", TensorProto.FLOAT, (0, 3))], + ) + feeds = {"x": np.zeros((0, 3), dtype=np.float32)} + + compiled = _run_against_reference(model, feeds, tmp_path) + + assert compiled.report["memory"]["arena_bytes"] == 0 + + +@pytest.mark.parametrize("opset", [7, 13, 14, OPSET]) +def test_kernels_dispatch_at_every_registered_revision(tmp_path, opset): + model = _model( + [helper.make_node("Add", ["a", "b"], ["y"], name="op")], + [ + _tensor("a", TensorProto.FLOAT, (2, 3)), + _tensor("b", TensorProto.FLOAT, (3,)), + ], + [_tensor("y", TensorProto.FLOAT, (2, 3))], + opset=opset, + ) + feeds = { + "a": _values((2, 3), np.float32, seed=1), + "b": _values((3,), np.float32, seed=2), + } + + _run_against_reference(model, feeds, tmp_path) + + +def test_an_opset_below_the_registered_semantics_is_refused(tmp_path): + """Add-6 broadcast through attributes; its semantics are not what the kernel implements.""" + model = _model( + [helper.make_node("Add", ["a", "b"], ["y"], name="op")], + [ + _tensor("a", TensorProto.FLOAT, (2, 3)), + _tensor("b", TensorProto.FLOAT, (2, 3)), + ], + [_tensor("y", TensorProto.FLOAT, (2, 3))], + opset=6, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path / "out") + + assert "`Add`" in str(error.value) + assert "Nearest supported version: 7." in str(error.value) + assert not (tmp_path / "out").exists() diff --git a/src/python/tests/test_extra_compiler_kernels.py b/src/python/tests/test_extra_compiler_kernels.py new file mode 100644 index 0000000..f75d3d5 --- /dev/null +++ b/src/python/tests/test_extra_compiler_kernels.py @@ -0,0 +1,7199 @@ +"""The elementwise kernel family: the loop it is emitted from, and what it refuses. + +What an op computes is settled by the conformance and differential suites, against ONNX's +own corpus and reference evaluator. What is asserted here is the emission contract — one +shared kernel per op, operand types and loop form; the flat loop when nothing broadcasts — +and the errors for the combinations the compiler will not compile at all. +""" + +from __future__ import annotations + +import re +import shutil + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError, HarnessError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +# The harness refuses to import without numpy, so this covers both dependencies. +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from onnx import TensorProto, helper # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +from fnnx.extras.compilers.c import compile_onnx # noqa: E402 +from fnnx.extras.compilers.c.onnx.kernels import ( # noqa: E402 + KERNELS, + NodeContext, + TensorRef, +) + +OPSET = 22 + +requires_c_compiler = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + + +def _tensor(name, elem_type, shape): + return helper.make_tensor_value_info(name, elem_type, list(shape)) + + +def _model(nodes, inputs, outputs, *, initializer=(), opset=OPSET): + graph = helper.make_graph( + nodes, "kernels", list(inputs), list(outputs), initializer=list(initializer) + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + + +def _compile(model, tmp_path): + result = compile_onnx(model, tmp_path) + return result.report, result.header_path.read_text(encoding="utf-8") + + +def _add_model(left_shape, right_shape, *, elem_type=TensorProto.FLOAT): + return _model( + [helper.make_node("Add", ["a", "b"], ["y"], name="add")], + [_tensor("a", elem_type, left_shape), _tensor("b", elem_type, right_shape)], + [helper.make_empty_tensor_value_info("y")], + ) + + +def _kernels(report, op_type: str) -> list[str]: + prefix = f"{report['prefix']}_{op_type}" + return [name for name in report["kernels"] if name.startswith(prefix)] + + +# -------------------------------------------------------------------------------------- +# The shared loop +# -------------------------------------------------------------------------------------- + + +def test_operands_of_the_results_shape_take_the_flat_loop(tmp_path): + """Nothing broadcasts, so the kernel indexes straight into its operands.""" + report, header = _compile(_add_model((2, 3), (2, 3)), tmp_path) + + (kernel,) = _kernels(report, "add") + assert "strides" not in header + assert f"{kernel}(\n y,\n a,\n b,\n 6u);" in header + + +def test_a_broadcasting_operand_takes_the_strided_loop(tmp_path): + """The stretched axis is a zero stride, and the shape and strides are call-site literals.""" + report, header = _compile(_add_model((2, 3), (3,)), tmp_path) + + (kernel,) = _kernels(report, "add") + assert "const size_t* strides0" in header + assert "offset1 += coordinate * strides1[axis];" in header + assert ( + f"{kernel}(\n y,\n a,\n b,\n 6u,\n 2,\n" + " (const size_t[]){2u, 3u},\n" + " (const size_t[]){3u, 1u},\n" + " (const size_t[]){0u, 1u});" in header + ) + + +def test_nodes_running_one_op_at_one_type_share_a_kernel(tmp_path): + """Kernels are shared statics, not code inlined per node.""" + model = _model( + [ + helper.make_node("Add", ["a", "b"], ["h"], name="first"), + helper.make_node("Add", ["h", "b"], ["y"], name="second"), + ], + [ + _tensor("a", TensorProto.FLOAT, (2, 3)), + _tensor("b", TensorProto.FLOAT, (2, 3)), + ], + [helper.make_empty_tensor_value_info("y")], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "add") + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +def test_a_kernel_is_emitted_per_element_type_and_per_loop_form(tmp_path): + """A kernel name has to encode everything the emitted code depends on.""" + model = _model( + [ + helper.make_node("Add", ["a", "b"], ["y"], name="aligned"), + helper.make_node("Add", ["a", "c"], ["z"], name="broadcasting"), + helper.make_node("Add", ["i", "i"], ["w"], name="integer"), + ], + [ + _tensor("a", TensorProto.FLOAT, (2, 3)), + _tensor("b", TensorProto.FLOAT, (2, 3)), + _tensor("c", TensorProto.FLOAT, (3,)), + _tensor("i", TensorProto.INT32, (2, 3)), + ], + [ + helper.make_empty_tensor_value_info("y"), + helper.make_empty_tensor_value_info("z"), + helper.make_empty_tensor_value_info("w"), + ], + ) + + report, _ = _compile(model, tmp_path) + + assert len(set(_kernels(report, "add"))) == 3 + + +def test_an_attribute_value_is_a_kernel_argument_rather_than_a_kernel(tmp_path): + """Two nodes differing only in an attribute value are one kernel called twice.""" + model = _model( + [ + helper.make_node("Elu", ["x"], ["h"], name="gentle", alpha=0.5), + helper.make_node("Elu", ["h"], ["y"], name="steep", alpha=2.0), + ], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "elu") + assert header.count(f"static void {kernel}(") == 1 + assert "0.5f" in header and "2.0f" in header + + +# -------------------------------------------------------------------------------------- +# Moving bytes rather than computing them +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("op_type", "source", "target"), + [ + ("Cast", TensorProto.FLOAT, TensorProto.FLOAT), + ("BitCast", TensorProto.FLOAT, TensorProto.INT32), + ], +) +def test_a_conversion_that_changes_no_bits_is_a_copy(tmp_path, op_type, source, target): + """Neither op has anything to compute per element, so neither emits a kernel at all.""" + model = _model( + [helper.make_node(op_type, ["x"], ["y"], name="convert", to=target)], + [_tensor("x", source, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + opset=26, + ) + + report, header = _compile(model, tmp_path) + + assert not _kernels(report, op_type.lower()) + assert "memcpy(y, x, 6u * sizeof(*y));" in header + + +def test_a_cast_kernel_is_emitted_per_target_type(tmp_path): + """The target type is what a cast's code depends on, so it has to name the kernel.""" + model = _model( + [ + helper.make_node("Cast", ["x"], ["a"], name="narrow", to=TensorProto.INT16), + helper.make_node("Cast", ["x"], ["b"], name="wide", to=TensorProto.INT32), + helper.make_node("Cast", ["x"], ["c"], name="again", to=TensorProto.INT32), + ], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info(name) for name in ("a", "b", "c")], + ) + + report, header = _compile(model, tmp_path) + + assert len(set(_kernels(report, "cast"))) == 2 + assert "(int16_t)x0" in header and "(int32_t)x0" in header + + +@requires_c_compiler +def test_casting_to_and_from_bool_emits_a_kernel_for_each_direction(tmp_path): + """`bool` and `uint8` are one C type, and the two directions are not one formula. + + A byte casts to true when it is nonzero, while a boolean carries the value it already + holds, so a kernel named after its C types alone would have the two collide. + """ + model = _model( + [ + helper.make_node("Cast", ["x"], ["b"], name="to_bool", to=TensorProto.BOOL), + helper.make_node( + "Cast", ["b"], ["y"], name="to_byte", to=TensorProto.UINT8 + ), + ], + [_tensor("x", TensorProto.UINT8, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + ) + values = np.array([[0, 1, 2], [3, 200, 255]], dtype=np.uint8) + + result = compile_onnx(model, tmp_path) + outputs = result.load().run({"x": values}) + + assert len(set(_kernels(result.report, "cast"))) == 2 + expected = ReferenceEvaluator(model).run(None, {"x": values}) + np.testing.assert_array_equal(outputs["y"], expected[0]) + + +@requires_c_compiler +def test_isinf_detecting_neither_infinity_never_reads_its_operand(tmp_path): + """Nothing is detected, so the result is false everywhere. + + A kernel that still read the operand would leave the local unused, which the artifact's + `-Werror` build contract turns into a failure rather than a warning. + """ + model = _model( + [ + helper.make_node( + "IsInf", + ["x"], + ["y"], + name="never", + detect_positive=0, + detect_negative=0, + ) + ], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + ) + + result = compile_onnx(model, tmp_path) + outputs = result.load().run({"x": np.full((2, 3), np.inf, dtype=np.float32)}) + + (kernel,) = _kernels(result.report, "isinf") + assert f"{kernel}(\n y,\n 6u);" in result.header_path.read_text() + assert not outputs["y"].any() + + +# -------------------------------------------------------------------------------------- +# Rearranging elements rather than computing them +# -------------------------------------------------------------------------------------- + + +def _int64(name, values): + return onnx.numpy_helper.from_array(np.array(values, dtype=np.int64), name) + + +@pytest.mark.parametrize( + ("op_type", "shape", "initializer"), + [ + ("Reshape", (2, 3), [_int64("p", [6])]), + ("Flatten", (2, 3), []), + ("Squeeze", (1, 2, 3), [_int64("p", [0])]), + ("Unsqueeze", (2, 3), [_int64("p", [0])]), + ], +) +def test_an_op_that_only_relabels_axes_is_a_copy(tmp_path, op_type, shape, initializer): + """None of them moves an element: the row-major buffer is the same, under other axes.""" + inputs = ["x", "p"] if initializer else ["x"] + model = _model( + [helper.make_node(op_type, inputs, ["y"], name="view")], + [_tensor("x", TensorProto.FLOAT, shape)], + [helper.make_empty_tensor_value_info("y")], + initializer=initializer, + ) + + report, header = _compile(model, tmp_path) + + assert not _kernels(report, "copy") + assert "memcpy(y, x, 6u * sizeof(*y));" in header + + +def test_a_move_that_stays_contiguous_is_a_memcpy(tmp_path): + """Concatenating along the outermost axis lays each operand down in one unbroken run.""" + model = _model( + [helper.make_node("Concat", ["a", "b"], ["y"], name="join", axis=0)], + [ + _tensor("a", TensorProto.FLOAT, (2, 3)), + _tensor("b", TensorProto.FLOAT, (4, 3)), + ], + [helper.make_empty_tensor_value_info("y")], + ) + + report, header = _compile(model, tmp_path) + + assert not _kernels(report, "copy") + assert "memcpy(y, a, 6u * sizeof(*y));" in header + assert "memcpy(y + 6, b, 12u * sizeof(*y));" in header + + +def test_a_reordering_move_passes_its_strides_as_call_site_literals(tmp_path): + """A transpose is the operand read along permuted strides, written out in order.""" + model = _model( + [helper.make_node("Transpose", ["x"], ["y"], name="swap", perm=[1, 0])], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "copy") + assert ( + f"{kernel}(\n y,\n x,\n 6u,\n 2,\n" + " (const size_t[]){3u, 2u},\n" + " (const ptrdiff_t[]){2, 1},\n" + " (const ptrdiff_t[]){1, 3},\n" + " 0,\n" + " 0);" in header + ) + + +def test_the_views_share_one_move_kernel_per_element_type(tmp_path): + """Every one of them walks the same addressing, so they are one shared static.""" + model = _model( + [ + helper.make_node("Transpose", ["x"], ["t"], name="swap", perm=[1, 0]), + helper.make_node("Tile", ["t", "r"], ["u"], name="repeat"), + helper.make_node("Slice", ["u", "s", "e"], ["y"], name="cut"), + helper.make_node("Transpose", ["i"], ["z"], name="swap_ints", perm=[1, 0]), + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3)), + _tensor("i", TensorProto.INT32, (2, 3)), + ], + [ + helper.make_empty_tensor_value_info("y"), + helper.make_empty_tensor_value_info("z"), + ], + initializer=[_int64("r", [2, 2]), _int64("s", [1, 1]), _int64("e", [5, 3])], + ) + + report, header = _compile(model, tmp_path) + + kernels = set(_kernels(report, "copy")) + assert len(kernels) == 2 + for kernel in kernels: + assert header.count(f"static void {kernel}(") == 1 + assert f"{report['prefix']}_copy_float" in kernels + assert f"{report['prefix']}_copy_int32_t" in kernels + + +@requires_c_compiler +def test_a_view_of_a_zero_element_tensor_emits_no_move(tmp_path): + """There is nothing to move, and an empty `memcpy` would still need a valid pointer.""" + model = _model( + [helper.make_node("Tile", ["x", "r"], ["y"], name="repeat")], + [_tensor("x", TensorProto.FLOAT, (2, 0))], + [helper.make_empty_tensor_value_info("y")], + initializer=[_int64("r", [2, 3])], + ) + + result = compile_onnx(model, tmp_path) + outputs = result.load().run({"x": np.zeros((2, 0), dtype=np.float32)}) + + assert "memcpy" not in result.header_path.read_text() + assert outputs["y"].shape == (4, 0) + + +# -------------------------------------------------------------------------------------- +# Walking a tensor by axes +# -------------------------------------------------------------------------------------- + + +def test_reductions_of_one_op_and_type_share_a_kernel(tmp_path): + """Extents, strides and group counts are arguments, so the axes do not name a kernel.""" + model = _model( + [ + helper.make_node("ReduceSum", ["x"], ["a"], name="rows", axes=[0]), + helper.make_node( + "ReduceSum", ["x"], ["b"], name="columns", axes=[1], keepdims=0 + ), + ], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info(name) for name in ("a", "b")], + opset=11, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "reducesum") + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +def test_a_scan_axis_the_graph_fixes_compiles_to_a_single_call(tmp_path): + model = _model( + [helper.make_node("CumSum", ["x", "axis"], ["y"], name="scan")], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + initializer=[onnx.numpy_helper.from_array(np.array(1, dtype=np.int64), "axis")], + opset=14, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "cumsum") + assert "switch" not in header + assert header.count(f"{kernel}(\n") == 2 + + +@requires_c_compiler +def test_a_runtime_scan_axis_compiles_to_a_call_per_axis(tmp_path): + """The axis decides which elements a scan visits and no shape at all, so a graph that + computes it still compiles: every value it can take is a call site of its own.""" + model = _model( + [helper.make_node("CumSum", ["x", "axis"], ["y"], name="scan")], + [ + _tensor("x", TensorProto.FLOAT, (2, 3)), + _tensor("axis", TensorProto.INT32, ()), + ], + [helper.make_empty_tensor_value_info("y")], + opset=14, + ) + values = np.arange(6, dtype=np.float32).reshape(2, 3) + + result = compile_onnx(model, tmp_path) + compiled = result.load() + + header = result.header_path.read_text(encoding="utf-8") + assert header.count("case 0:") == 1 and header.count("case 1:") == 1 + for axis in (0, 1, -1, -2): + feeds = {"x": values, "axis": np.array(axis, dtype=np.int32)} + expected = ReferenceEvaluator(model).run(None, feeds) + np.testing.assert_array_equal(compiled.run(feeds)["y"], expected[0]) + + +@requires_c_compiler +@pytest.mark.parametrize("axis", [2, -3]) +def test_a_scan_axis_outside_the_rank_returns_a_nonzero_status(tmp_path, axis): + """The status enum is what a run-time-checked operand reports through.""" + model = _model( + [helper.make_node("CumSum", ["x", "axis"], ["y"], name="scan")], + [ + _tensor("x", TensorProto.FLOAT, (2, 3)), + _tensor("axis", TensorProto.INT32, ()), + ], + [helper.make_empty_tensor_value_info("y")], + opset=14, + ) + + compiled = compile_onnx(model, tmp_path).load() + + with pytest.raises(HarnessError, match="status 1"): + compiled.run( + { + "x": np.zeros((2, 3), dtype=np.float32), + "axis": np.array(axis, dtype=np.int32), + } + ) + + +@pytest.mark.parametrize("axis", [2, -3]) +def test_a_scan_axis_the_graph_fixes_outside_the_rank_is_rejected(tmp_path, axis): + """An axis the graph pins is checked where it is known: at compile time, not at run.""" + model = _model( + [helper.make_node("CumSum", ["x", "axis"], ["y"], name="scan")], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + initializer=[ + onnx.numpy_helper.from_array(np.array(axis, dtype=np.int64), "axis") + ], + opset=14, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`scan`" in message + assert f"axis {axis}" in message + assert "rank-2" in message + + +def test_a_scan_of_a_scalar_is_rejected(tmp_path): + model = _model( + [helper.make_node("CumSum", ["x", "axis"], ["y"], name="scan")], + [ + _tensor("x", TensorProto.FLOAT, ()), + _tensor("axis", TensorProto.INT32, ()), + ], + [helper.make_empty_tensor_value_info("y")], + opset=14, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`scan`" in message + assert "`x`" in message + + +def test_axes_naming_one_dimension_twice_are_rejected(tmp_path): + """Reducing an axis twice has no meaning; numpy refuses it and so does the compiler.""" + model = _model( + [helper.make_node("ReduceSum", ["x"], ["y"], name="total", axes=[1, 1])], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + opset=11, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`total`" in message + assert "[1, 1]" in message + + +# -------------------------------------------------------------------------------------- +# Normalizing by a group's own statistics +# -------------------------------------------------------------------------------------- + + +def _instance_norm_node(name, data, **attributes): + return helper.make_node( + "InstanceNormalization", + [data, "s", "b"], + [f"y_{name}"], + name=name, + **attributes, + ) + + +def test_normalizations_of_one_op_and_type_share_a_kernel(tmp_path): + """Extents, strides and the epsilon are arguments, so neither names a kernel.""" + model = _model( + [ + _instance_norm_node("wide", "x"), + _instance_norm_node("narrow", "z", epsilon=0.01), + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3, 4, 5)), + _tensor("z", TensorProto.FLOAT, (2, 3, 4)), + _tensor("s", TensorProto.FLOAT, (3,)), + _tensor("b", TensorProto.FLOAT, (3,)), + ], + [ + helper.make_empty_tensor_value_info(f"y_{name}") + for name in ("wide", "narrow") + ], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "instancenormalization") + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +def test_a_kernel_is_emitted_per_set_of_statistics_asked_for(tmp_path): + """The mean and the inverse deviation are buffers the kernel writes into. + + Two nodes that report different ones do not run the same code, so a name that did not + say which they are would have the two collide. + """ + model = _model( + [ + helper.make_node("LayerNormalization", ["x", "s"], ["y"], name="plain"), + helper.make_node( + "LayerNormalization", ["x", "s"], ["z", "m", "d"], name="reporting" + ), + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3)), + _tensor("s", TensorProto.FLOAT, (3,)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("y", "z", "m", "d")], + opset=17, + ) + + report, header = _compile(model, tmp_path) + + kernels = set(_kernels(report, "layernormalization")) + assert len(kernels) == 2 + for kernel in kernels: + assert header.count(f"static void {kernel}(") == 1 + + +@requires_c_compiler +def test_a_running_statistic_the_node_skips_is_not_computed(tmp_path): + """ONNX lets a node drop an optional output by naming it the empty string.""" + model = _model( + [ + helper.make_node( + "BatchNormalization", + ["x", "s", "b", "m", "v"], + ["y", "", "running_var"], + name="norm", + training_mode=1, + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3, 2)), + _tensor("s", TensorProto.FLOAT, (3,)), + _tensor("b", TensorProto.FLOAT, (3,)), + _tensor("m", TensorProto.FLOAT, (3,)), + _tensor("v", TensorProto.FLOAT, (3,)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("y", "running_var")], + opset=15, + ) + feeds = { + "x": np.arange(12, dtype=np.float32).reshape(2, 3, 2), + "s": np.array([1.0, 2.0, 3.0], dtype=np.float32), + "b": np.array([0.0, 1.0, 2.0], dtype=np.float32), + "m": np.array([0.5, 0.5, 0.5], dtype=np.float32), + "v": np.array([1.0, 2.0, 3.0], dtype=np.float32), + } + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + expected = ReferenceEvaluator(model).run(None, feeds) + for name, want in zip(("y", "running_var"), expected): + np.testing.assert_allclose(outputs[name], want, rtol=1e-6, atol=1e-6) + + +@requires_c_compiler +def test_a_training_node_reporting_no_statistic_still_builds(tmp_path): + """`momentum` blends the running statistics and is read nowhere else. + + A node that reports neither is one ONNX's own shape inference refuses, but the + reference evaluator computes its `Y` all the same, so the kernel is emitted — and has + to leave the argument out rather than take one it never reads, which the artifact's + own `-Werror=unused-parameter` build would refuse. The result shape is declared + because inference derives none for a node it will not vouch for. + """ + model = _model( + [ + helper.make_node( + "BatchNormalization", + ["x", "s", "b", "m", "v"], + ["y"], + name="norm", + training_mode=1, + momentum=0.7, + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3, 2)), + *(_tensor(name, TensorProto.FLOAT, (3,)) for name in ("s", "b", "m", "v")), + ], + [_tensor("y", TensorProto.FLOAT, (2, 3, 2))], + opset=15, + ) + feeds = { + "x": np.arange(12, dtype=np.float32).reshape(2, 3, 2), + "s": np.array([1.0, 2.0, 3.0], dtype=np.float32), + "b": np.array([0.0, 1.0, 2.0], dtype=np.float32), + "m": np.array([0.5, 0.5, 0.5], dtype=np.float32), + "v": np.array([1.0, 2.0, 3.0], dtype=np.float32), + } + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + (expected,) = ReferenceEvaluator(model).run(None, feeds) + np.testing.assert_allclose(outputs["y"], expected, rtol=1e-6, atol=1e-6) + + +def test_a_stash_type_that_takes_no_statistics_is_rejected(tmp_path): + """`stash_type` names the precision stage one runs in, which has to be a float one.""" + model = _model( + [ + helper.make_node( + "LayerNormalization", + ["x", "s"], + ["y"], + name="norm", + stash_type=TensorProto.INT32, + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3)), + _tensor("s", TensorProto.FLOAT, (3,)), + ], + [helper.make_empty_tensor_value_info("y")], + opset=17, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`norm`" in message + assert "INT32" in message + + +def test_groups_that_do_not_divide_the_channels_are_rejected(tmp_path): + model = _model( + [ + helper.make_node( + "GroupNormalization", ["x", "s", "b"], ["y"], name="norm", num_groups=2 + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 5, 3)), + _tensor("s", TensorProto.FLOAT, (5,)), + _tensor("b", TensorProto.FLOAT, (5,)), + ], + [helper.make_empty_tensor_value_info("y")], + opset=21, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`norm`" in message + assert "5 channel(s)" in message + + +def test_an_lp_order_onnx_does_not_define_is_rejected(tmp_path): + """`p` selects the norm, and ONNX defines the op for the first two only.""" + model = _model( + [helper.make_node("LpNormalization", ["x"], ["y"], name="norm", p=3)], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`norm`" in message + assert "[1, 2]" in message + + +@requires_c_compiler +def test_the_l1_norm_sums_absolute_values(tmp_path): + """The one place onnxruntime is the oracle, because ONNX's own two are blind here. + + The reference evaluator raises the elements to the power `p` without taking their + absolute value, so at `p` = 1 it computes a signed sum rather than a norm, and the + corpus's `l1normalization` tests all carry non-negative data — nothing else in this + suite can tell the two apart. onnxruntime, the second oracle the compiler's parity + testing rests on, computes the norm ONNX defines. + """ + ort = pytest.importorskip("onnxruntime") + model = _model( + [helper.make_node("LpNormalization", ["x"], ["y"], name="norm", p=1, axis=1)], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + # onnxruntime implements the op's first revision only, which is the one every + # model below opset 22 -- and so this one -- dispatches to. + opset=21, + ) + feeds = {"x": np.array([[1.0, -2.0, 3.0], [-1.0, -1.0, 2.0]], dtype=np.float32)} + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + path = tmp_path / "lpnormalization.onnx" + onnx.save(model, str(path)) + session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + expected = session.run(None, feeds)[0] + np.testing.assert_allclose(outputs["y"], expected, rtol=1e-6, atol=1e-7) + + +@pytest.mark.parametrize( + ("op_type", "inputs", "shapes"), + [ + ("InstanceNormalization", ["x", "s", "b"], [(4,), (1,), (1,)]), + ("LRN", ["x"], [(4,)]), + ], +) +def test_a_normalization_without_a_channel_axis_is_rejected( + tmp_path, op_type, inputs, shapes +): + """Both read their operand as instances by channels; a vector has no channels.""" + attributes = {"size": 3} if op_type == "LRN" else {} + model = _model( + [helper.make_node(op_type, inputs, ["y"], name="norm", **attributes)], + [ + _tensor(name, TensorProto.FLOAT, shape) + for name, shape in zip(inputs, shapes) + ], + [helper.make_empty_tensor_value_info("y")], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`norm`" in message + assert "rank of at least 2" in message + + +def test_batch_normalization_at_inference_refuses_the_training_outputs(tmp_path): + """ONNX leaves the running statistics undefined outside training mode. + + Its own shape inference refuses the node too, but only where it can derive the extra + outputs' types; a model that declares them itself gets this far, and is stopped here + rather than handed a buffer nothing writes. + """ + model = _model( + [ + helper.make_node( + "BatchNormalization", + ["x", "s", "b", "m", "v"], + ["y", "rm", "rv"], + name="norm", + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3)), + *(_tensor(name, TensorProto.FLOAT, (3,)) for name in ("s", "b", "m", "v")), + ], + [ + _tensor("y", TensorProto.FLOAT, (2, 3)), + _tensor("rm", TensorProto.FLOAT, (3,)), + _tensor("rv", TensorProto.FLOAT, (3,)), + ], + opset=15, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`norm`" in message + assert "training mode only" in message + + +# -------------------------------------------------------------------------------------- +# Reading through an index +# -------------------------------------------------------------------------------------- + + +def _gather_model(op_type, data_shape, index_shape, *, index_type=TensorProto.INT64): + return _model( + [helper.make_node(op_type, ["x", "i"], ["y"], name="pick", axis=0)], + [ + _tensor("x", TensorProto.FLOAT, data_shape), + _tensor("i", index_type, index_shape), + ], + [helper.make_empty_tensor_value_info("y")], + ) + + +def test_gathering_at_one_element_and_index_type_shares_a_kernel(tmp_path): + """Two nodes reading the same types read them through the same emitted loop.""" + model = _model( + [ + helper.make_node("Gather", ["x", "i"], ["a"], name="first", axis=0), + helper.make_node("Gather", ["x", "j"], ["b"], name="second", axis=1), + ], + [ + _tensor("x", TensorProto.FLOAT, (3, 4)), + _tensor("i", TensorProto.INT64, (2,)), + _tensor("j", TensorProto.INT64, (3,)), + ], + [ + helper.make_empty_tensor_value_info("a"), + helper.make_empty_tensor_value_info("b"), + ], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "gather") + assert header.count(f"static int {kernel}(") == 1 + assert header.count(f"if ({kernel}(") == 2 + + +def test_a_gather_kernel_is_emitted_per_index_type(tmp_path): + """The index type is part of the loop's signature, so int32 and int64 are two kernels.""" + model = _model( + [ + helper.make_node("Gather", ["x", "i"], ["a"], name="wide", axis=0), + helper.make_node("Gather", ["x", "j"], ["b"], name="narrow", axis=0), + ], + [ + _tensor("x", TensorProto.FLOAT, (3, 4)), + _tensor("i", TensorProto.INT64, (2,)), + _tensor("j", TensorProto.INT32, (2,)), + ], + [ + helper.make_empty_tensor_value_info("a"), + helper.make_empty_tensor_value_info("b"), + ], + ) + + report, _ = _compile(model, tmp_path) + + assert sorted(_kernels(report, "gather")) == [ + f"{report['prefix']}_gather_float_int32_t", + f"{report['prefix']}_gather_float_int64_t", + ] + + +@requires_c_compiler +@pytest.mark.parametrize( + ("op_type", "index"), + [("Gather", 3), ("Gather", -4), ("GatherElements", 3), ("GatherND", 3)], +) +def test_an_index_outside_its_axis_returns_a_nonzero_status(tmp_path, op_type, index): + """An index operand comes from the caller, so the artifact refuses one it cannot serve. + + ONNX defines an index only within its axis, counted from either end; anything else would + read past the buffer, which is what the status enum exists to report instead. + """ + index_shape = {"GatherND": (2, 1), "GatherElements": (2, 4)}.get(op_type, (2,)) + compiled = compile_onnx( + _gather_model(op_type, (3, 4), index_shape), tmp_path + ).load() + + with pytest.raises(HarnessError, match="status 1"): + compiled.run( + { + "x": np.zeros((3, 4), dtype=np.float32), + "i": np.full(index_shape, index, dtype=np.int64), + } + ) + + +@requires_c_compiler +def test_a_sequence_length_outside_the_time_axis_returns_a_nonzero_status(tmp_path): + """ONNX defines a length as being in `[1, s]`; the reversal of anything else is nothing.""" + model = _model( + [ + helper.make_node( + "ReverseSequence", + ["x", "lens"], + ["y"], + name="reverse", + batch_axis=0, + time_axis=1, + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3)), + _tensor("lens", TensorProto.INT64, (2,)), + ], + [helper.make_empty_tensor_value_info("y")], + ) + + compiled = compile_onnx(model, tmp_path).load() + + for lengths in ([0, 2], [2, 4]): + with pytest.raises(HarnessError, match="status 1"): + compiled.run( + { + "x": np.zeros((2, 3), dtype=np.float32), + "lens": np.array(lengths, dtype=np.int64), + } + ) + + +def test_an_eye_reads_nothing_but_the_shape_of_its_operand(tmp_path): + """EyeLike is a function of coordinates, so the operand it is shaped like goes unread.""" + model = _model( + [helper.make_node("EyeLike", ["x"], ["y"], name="eye")], + [_tensor("x", TensorProto.FLOAT, (3, 3))], + [helper.make_empty_tensor_value_info("y")], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "eyelike") + entrypoint = header.split(f"int {report['prefix']}_run(")[-1] + assert "(void)x;" in entrypoint + assert "x" not in entrypoint.split(f"{kernel}(")[1].split(");")[0] + + +def test_a_pad_kernel_is_emitted_per_mode(tmp_path): + """What a mode does with a coordinate outside the operand is the kernel's whole body.""" + model = _model( + [ + helper.make_node("Pad", ["x", "p"], ["a"], name="fill", mode="constant"), + helper.make_node("Pad", ["x", "p"], ["b"], name="mirror", mode="reflect"), + helper.make_node("Pad", ["x", "p"], ["c"], name="repeat", mode="edge"), + ], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info(name) for name in ("a", "b", "c")], + initializer=[_int64("p", [1, 0, 1, 0])], + ) + + report, _ = _compile(model, tmp_path) + + assert sorted(_kernels(report, "pad")) == [ + f"{report['prefix']}_pad_{mode}_float" + for mode in ("constant", "edge", "reflect") + ] + + +# -------------------------------------------------------------------------------------- +# The matrix products and the determinant +# -------------------------------------------------------------------------------------- + + +def test_matmuls_of_one_element_type_share_a_kernel(tmp_path): + """Shapes reach the kernel as call-site literals, so batching is not a kernel of its own.""" + model = _model( + [ + helper.make_node("MatMul", ["a", "b"], ["p"], name="plain"), + helper.make_node("MatMul", ["c", "d"], ["q"], name="batched"), + ], + [ + _tensor("a", TensorProto.FLOAT, (2, 3)), + _tensor("b", TensorProto.FLOAT, (3, 4)), + _tensor("c", TensorProto.FLOAT, (5, 2, 3)), + _tensor("d", TensorProto.FLOAT, (3,)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q")], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "matmul") + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +def test_determinants_share_one_working_matrix_sized_for_the_largest(tmp_path): + """Elimination needs a copy of its matrix, which the artifact reserves at compile time. + + Both nodes run the same kernel, so both eliminate in the same static buffer — they run one + after the other — and it is large enough for the bigger of the two. + """ + model = _model( + [ + helper.make_node("Det", ["x"], ["a"], name="small"), + helper.make_node("Det", ["z"], ["b"], name="large"), + ], + [ + _tensor("x", TensorProto.FLOAT, (3, 2, 2)), + _tensor("z", TensorProto.FLOAT, (4, 4)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("a", "b")], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "det") + assert f"static float {kernel}_work[16];" in header + assert header.count(f"{kernel}_work,") == 2 + assert report["memory"]["arena_bytes"] == 16 * 4 + + +def test_a_determinant_of_another_element_type_gets_its_own_working_matrix(tmp_path): + """The buffer a kernel eliminates in is typed like the kernel, so the two do not share.""" + model = _model( + [ + helper.make_node("Det", ["x"], ["a"], name="single"), + helper.make_node("Det", ["z"], ["b"], name="wide"), + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 2)), + _tensor("z", TensorProto.DOUBLE, (2, 2)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("a", "b")], + ) + + report, header = _compile(model, tmp_path) + + prefix = report["prefix"] + assert f"static float {prefix}_det_float_work[4];" in header + assert f"static double {prefix}_det_double_work[4];" in header + assert report["memory"]["arena_bytes"] == 4 * 4 + 4 * 8 + + +# -------------------------------------------------------------------------------------- +# The convolutions +# -------------------------------------------------------------------------------------- + + +def _conv_model( + x_shape, + w_shape, + *, + op_type="Conv", + bias=None, + output=None, + opset=OPSET, + **attributes, +): + names = ["x", "w"] + (["b"] if bias else []) + inputs = [ + _tensor("x", TensorProto.FLOAT, x_shape), + _tensor("w", TensorProto.FLOAT, w_shape), + ] + if bias: + inputs.append(_tensor("b", TensorProto.FLOAT, bias)) + return _model( + [helper.make_node(op_type, names, ["y"], name="conv", **attributes)], + inputs, + [ + helper.make_empty_tensor_value_info("y") + if output is None + else _tensor("y", TensorProto.FLOAT, output) + ], + opset=opset, + ) + + +def test_convolutions_of_one_element_type_share_a_kernel(tmp_path): + """The geometry reaches the kernel as call-site literals, so rank is not a kernel of its + own: a 1-D and a grouped 2-D convolution run the same code at different arguments.""" + model = _model( + [ + helper.make_node("Conv", ["x", "w"], ["p"], name="signal"), + helper.make_node("Conv", ["i", "k"], ["q"], name="image", group=2), + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 2, 6)), + _tensor("w", TensorProto.FLOAT, (3, 2, 3)), + _tensor("i", TensorProto.FLOAT, (2, 4, 5, 5)), + _tensor("k", TensorProto.FLOAT, (6, 2, 3, 3)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q")], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "conv") + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +def test_the_resolved_padding_reaches_the_kernel_as_literals(tmp_path): + """`auto_pad` is resolved at compile time: the kernel only ever sees concrete pads. + + A 4-wide window at stride 2 over a 5-wide axis needs three pads, and SAME_LOWER puts the + odd one at the beginning — two here against SAME_UPPER's one. The second axis is padded + for the reach of a 2-dilated window instead, which splits evenly. + """ + model = _conv_model( + (1, 1, 5, 5), + (1, 1, 4, 2), + auto_pad="SAME_LOWER", + strides=[2, 1], + dilations=[1, 2], + ) + + _, header = _compile(model, tmp_path) + + assert "(const ptrdiff_t[]){2, 1}" in header + + +def test_a_valid_convolution_pads_nothing(tmp_path): + """`VALID` is not `SAME`: it drops the positions a full window does not fit in. + + ONNX's own shape inference arbitrates this. It derives the result's shape from the pads + the mode implies, and the kernel refuses to be emitted against a buffer its addressing + disagrees with — so a compiler reading `VALID` as `SAME` would fail to compile this + model rather than quietly pad it out to the 5x5 SAME shape. + """ + report, header = _compile( + _conv_model((1, 1, 5, 5), (1, 1, 3, 3), auto_pad="VALID"), tmp_path + ) + + assert report["entrypoint"]["outputs"][0]["shape"] == [1, 1, 3, 3] + assert "(const ptrdiff_t[]){0, 0}" in header + + +def test_a_convolution_writing_no_elements_emits_no_call(tmp_path): + """An empty batch leaves nothing to convolve, and no loop that could read past a buffer.""" + report, header = _compile(_conv_model((0, 2, 5, 5), (3, 2, 3, 3)), tmp_path) + + assert not _kernels(report, "conv") + assert "conv" not in header.split(f"int {report['prefix']}_run(")[-1] + + +@requires_c_compiler +def test_a_grouped_convolution_reads_only_its_own_channels(tmp_path): + """Each group convolves its own slice of the channels, which zeroed weights expose. + + Zeroing the second group's filters leaves the first group's output untouched, so a + kernel addressing the whole channel stack per filter would change the answer. What the + surviving group *computes* is settled by the conformance and differential suites. + """ + model = _conv_model((1, 4, 5, 5), (4, 2, 3, 3), group=2) + compiled = compile_onnx(model, tmp_path).load() + x = np.arange(100, dtype=np.float32).reshape(1, 4, 5, 5) / 100 + w = np.ones((4, 2, 3, 3), dtype=np.float32) + masked = w.copy() + masked[2:] = 0.0 + + full = compiled.run({"x": x, "w": w})["y"] + partial = compiled.run({"x": x, "w": masked})["y"] + + np.testing.assert_array_equal(full[:, :2], partial[:, :2]) + np.testing.assert_array_equal(partial[:, 2:], np.zeros_like(partial[:, 2:])) + + +@pytest.mark.parametrize( + ("attributes", "message"), + [ + ({"auto_pad": "SAME_UPPER", "pads": [1, 1, 1, 1]}, "mutually exclusive"), + ({"kernel_shape": [2, 2]}, "the filter it is handed measures [3, 3]"), + ({"auto_pad": "SAME"}, "is not one of the modes ONNX defines"), + ({"group": 3}, "3 group(s) takes"), + ({"strides": [0, 1]}, "ONNX defines them as positive"), + ({"dilations": [1, 0]}, "ONNX defines them as positive"), + ({"group": 0}, "positive count"), + ], +) +def test_a_convolution_the_compiler_cannot_place_is_rejected( + tmp_path, attributes, message +): + model = _conv_model((1, 2, 5, 5), (2, 2, 3, 3), output=(1, 2, 3, 3), **attributes) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("x_shape", "w_shape", "output", "message"), + [ + ((1, 3), (2, 3), (1, 2), "rank 3 or more"), + ((1, 1, 5, 5), (1, 1, 3), (1, 1, 3, 3), "ONNX defines both as rank 4"), + ], +) +def test_a_convolution_of_mismatched_ranks_is_rejected( + tmp_path, x_shape, w_shape, output, message +): + """Declared shapes get past ONNX's own inference, so the kernel checks them itself.""" + model = _conv_model(x_shape, w_shape, output=output) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +def test_a_convolution_bias_is_one_value_per_output_channel(tmp_path): + model = _conv_model((1, 2, 5, 5), (2, 2, 3, 3), bias=(3,)) + + with pytest.raises(CompileError, match="one bias per output channel"): + compile_onnx(model, tmp_path) + + +# -------------------------------------------------------------------------------------- +# The transposed convolution +# -------------------------------------------------------------------------------------- + + +def _transpose_model(x_shape, w_shape, **kwargs): + return _conv_model(x_shape, w_shape, op_type="ConvTranspose", **kwargs) + + +def test_transposed_convolutions_of_one_element_type_share_a_kernel(tmp_path): + """One kernel for the op, as with Conv: the geometry is what the call sites differ in.""" + model = _model( + [ + helper.make_node("ConvTranspose", ["x", "w"], ["p"], name="signal"), + helper.make_node("ConvTranspose", ["i", "k"], ["q"], name="image", group=2), + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 2, 6)), + _tensor("w", TensorProto.FLOAT, (2, 3, 3)), + _tensor("i", TensorProto.FLOAT, (2, 4, 5, 5)), + _tensor("k", TensorProto.FLOAT, (4, 3, 3, 3)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q")], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "convtranspose") + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +def test_a_transposed_convolution_is_not_a_convolution(tmp_path): + """The two walk the same geometry in opposite directions, so neither may serve for the + other: one kernel each, however alike their arguments look.""" + model = _model( + [ + helper.make_node("Conv", ["x", "w"], ["p"], name="forward"), + helper.make_node("ConvTranspose", ["x", "k"], ["q"], name="backward"), + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 2, 5, 5)), + _tensor("w", TensorProto.FLOAT, (2, 2, 3, 3)), + _tensor("k", TensorProto.FLOAT, (2, 2, 3, 3)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q")], + ) + + report, _ = _compile(model, tmp_path) + + assert len(_kernels(report, "conv")) == 2 + + +def test_the_padding_a_same_mode_implies_reaches_the_kernel_as_literals(tmp_path): + """`auto_pad` is resolved at compile time: the kernel only ever sees concrete pads. + + A transposed convolution pads so that the result measures `extent * stride`, cropping + the reach the window has beyond it. A 3-tap window reaches two past a stride of 2 and + nothing past a stride of 3, and SAME_LOWER puts the odd pad of the first at the front. + """ + model = _transpose_model( + (1, 1, 3, 3), (1, 1, 3, 3), auto_pad="SAME_LOWER", strides=[2, 3] + ) + + report, header = _compile(model, tmp_path) + + assert report["entrypoint"]["outputs"][0]["shape"] == [1, 1, 6, 9] + assert "(const ptrdiff_t[]){1, 0}" in header + + +def test_a_result_wider_than_the_window_reaches_is_not_padded(tmp_path): + """`output_shape` may name a result the taps do not cover; the rest of it is bias alone. + + The pads are what crops the reach down to the result, so a result the reach falls short + of needs none — the positions past it are simply ones no tap contributes to. + """ + model = _transpose_model( + (1, 1, 3, 3), (1, 2, 3, 3), strides=[3, 2], output_shape=[10, 8] + ) + + report, header = _compile(model, tmp_path) + + assert report["entrypoint"]["outputs"][0]["shape"] == [1, 2, 10, 8] + assert "(const ptrdiff_t[]){0, 0}" in header + + +def test_a_transposed_convolution_writing_no_elements_emits_no_call(tmp_path): + report, header = _compile(_transpose_model((0, 2, 5, 5), (2, 3, 3, 3)), tmp_path) + + assert not _kernels(report, "convtranspose") + assert "convtranspose" not in header.split(f"int {report['prefix']}_run(")[-1] + + +@requires_c_compiler +def test_a_grouped_transposed_convolution_is_its_groups_run_separately(tmp_path): + """`group` splits both channel stacks into independent transposed convolutions. + + That is ONNX's definition of the attribute, and here it is also the only oracle for the + general case: the reference evaluator's own grouped path slices `W` by output rather + than input channels and hands every group the whole bias, so it can evaluate a grouped + node only where each group holds exactly one channel of each — which is what the + conformance corpus and the differential sweep are left covering. Splitting the operands + here and running each group through the evaluator ungrouped puts the general case back + within reach of the same oracle. + """ + x_shape, w_shape, groups = (2, 4, 4, 3), (4, 3, 3, 2), 2 + attributes = {"strides": [2, 1], "pads": [1, 0, 0, 1], "dilations": [1, 2]} + generator = np.random.default_rng(20260726) + x = generator.normal(size=x_shape).astype(np.float32) + w = generator.normal(size=w_shape).astype(np.float32) + compiled = compile_onnx( + _transpose_model(x_shape, w_shape, group=groups, **attributes), tmp_path + ).load() + + got = compiled.run({"x": x, "w": w})["y"] + + channels, filters = x_shape[1] // groups, w_shape[1] + expected = np.concatenate( + [ + ReferenceEvaluator( + _transpose_model( + (x_shape[0], channels, *x_shape[2:]), + (channels, filters, *w_shape[2:]), + **attributes, + ) + ).run( + None, + { + "x": x[:, group * channels : (group + 1) * channels], + "w": w[group * channels : (group + 1) * channels], + }, + )[0] + for group in range(groups) + ], + axis=1, + ) + np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize( + ("attributes", "message"), + [ + ({"auto_pad": "SAME_UPPER", "pads": [1, 1, 1, 1]}, "mutually exclusive"), + ({"kernel_shape": [2, 2]}, "the filter it is handed measures [3, 3]"), + ({"auto_pad": "SAME"}, "is not one of the modes ONNX defines"), + ({"group": 3}, "3 group(s) takes a filter"), + ({"strides": [0, 1]}, "ONNX defines them as positive"), + ({"output_padding": [-1, 0]}, "ONNX defines them as nonnegative"), + ({"output_shape": [4]}, "was given 1 `output_shape` for 2 spatial axis/axes"), + ({"pads": [1, 1]}, "was given 2 pad(s) for 2 spatial axis/axes"), + ], +) +def test_a_transposed_convolution_the_compiler_cannot_place_is_rejected( + tmp_path, attributes, message +): + model = _transpose_model( + (1, 2, 5, 5), (2, 2, 3, 3), output=(1, 2, 7, 7), **attributes + ) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +def test_a_transposed_filter_is_indexed_by_input_channel(tmp_path): + """`W` is (C, M/group, ...) here, not Conv's (M, C/group, ...); a node that hands over + the other layout is convolving something other than what it declares.""" + model = _transpose_model((1, 2, 5, 5), (3, 2, 3, 3), output=(1, 2, 7, 7)) + + with pytest.raises(CompileError, match="one stack per input channel"): + compile_onnx(model, tmp_path) + + +# -------------------------------------------------------------------------------------- +# The deformable convolution +# -------------------------------------------------------------------------------------- + + +def _deform_model( + x_shape, + w_shape, + offset_shape, + *, + bias=None, + mask=None, + output=None, + **attributes, +): + operands = [("x", x_shape), ("w", w_shape), ("offset", offset_shape)] + operands.append(("b" if bias else "", bias)) + if mask: + operands.append(("mask", mask)) + names = [name for name, _ in operands] + while names and not names[-1]: + names.pop() + return _model( + [helper.make_node("DeformConv", names, ["y"], name="deform", **attributes)], + [ + _tensor(name, TensorProto.FLOAT, shape) + for name, shape in operands + if name and shape + ], + [ + helper.make_empty_tensor_value_info("y") + if output is None + else _tensor("y", TensorProto.FLOAT, output) + ], + ) + + +def test_a_deformable_convolution_emits_one_sampler_per_element_type(tmp_path): + """The interpolation is a shared static of its own, so nodes reading the same element + type share it however their geometries differ.""" + model = _model( + [ + helper.make_node("DeformConv", ["x", "w", "o"], ["p"], name="small"), + helper.make_node("DeformConv", ["x", "k", "n"], ["q"], name="wide"), + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 1, 4, 4)), + _tensor("w", TensorProto.FLOAT, (1, 1, 2, 2)), + _tensor("o", TensorProto.FLOAT, (1, 8, 3, 3)), + _tensor("k", TensorProto.FLOAT, (2, 1, 3, 3)), + _tensor("n", TensorProto.FLOAT, (1, 18, 2, 2)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q")], + ) + + report, header = _compile(model, tmp_path) + + (sampler,) = _kernels(report, "bilinear") + assert header.count(f"static float {sampler}(") == 1 + (kernel,) = _kernels(report, "deformconv") + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +def test_a_deformable_convolution_writing_no_elements_emits_no_call(tmp_path): + """An empty batch leaves nothing to sample, and no sampler either.""" + report, header = _compile( + _deform_model((0, 1, 4, 4), (1, 1, 2, 2), (0, 8, 3, 3)), tmp_path + ) + + assert not _kernels(report, "deformconv") + assert not _kernels(report, "bilinear") + assert "deformconv" not in header.split(f"int {report['prefix']}_run(")[-1] + + +@requires_c_compiler +def test_a_deformation_that_leaves_the_operand_samples_nothing(tmp_path): + """A sampling point outside the operand contributes nothing — and one that is not a + number is not a point at all, so it cannot be floored into an index. + + Neither is reachable through the suites that settle what this op computes: the corpus + exercises ordinary offsets and the reference evaluator raises outright on a coordinate + it cannot floor. They are asserted here because the artifact has to stay defined on + whatever the caller passes, and what is left when every tap misses is the bias alone. + """ + compiled = compile_onnx( + _deform_model((1, 1, 4, 4), (1, 1, 2, 2), (1, 8, 3, 3), bias=(1,)), tmp_path + ).load() + operands = { + "x": np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4), + "w": np.ones((1, 1, 2, 2), dtype=np.float32), + "b": np.array([0.5], dtype=np.float32), + } + bias = np.full((1, 1, 3, 3), 0.5, dtype=np.float32) + + for offset in (1e30, -1e30, np.nan): + outside = compiled.run( + {**operands, "offset": np.full((1, 8, 3, 3), offset, dtype=np.float32)} + )["y"] + + np.testing.assert_array_equal(outside, bias) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"offset_shape": (1, 8, 4, 4)}, "addresses `offset` as [1, 8, 3, 3]"), + ({"offset_shape": (1, 6, 3, 3)}, "addresses `offset` as [1, 8, 3, 3]"), + ( + {"offset_shape": (1, 8, 3, 3), "mask": (1, 3, 3, 3)}, + "addresses `mask` as [1, 4, 3, 3]", + ), + ({"offset_shape": (1, 8, 3, 3), "bias": (2,)}, "one bias per output channel"), + ( + {"offset_shape": (1, 16, 3, 3), "offset_group": 3}, + "ONNX defines `offset_group` as a positive count that divides them", + ), + ], +) +def test_a_deformation_shaped_for_another_geometry_is_rejected( + tmp_path, kwargs, message +): + """ONNX's own shape inference reads none of these operands, so the kernel checks them + itself rather than addressing past the end of one.""" + model = _deform_model((1, 2, 4, 4), (1, 2, 2, 2), **kwargs) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +def test_a_deformable_convolution_of_another_rank_is_rejected(tmp_path): + """ONNX defines the op for any rank; nothing can vouch for a sampler of another one.""" + model = _deform_model((1, 1, 4, 4, 4), (1, 1, 2, 2, 2), (1, 24, 3, 3, 3)) + + with pytest.raises(CompileError, match="compiled for 2 spatial axes"): + compile_onnx(model, tmp_path) + + +# -------------------------------------------------------------------------------------- +# The poolings +# -------------------------------------------------------------------------------------- + + +def _pool_model( + x_shape, + *, + op_type="AveragePool", + elem_type=TensorProto.FLOAT, + outputs=1, + output=None, + **attributes, +): + names = ["y"] + [f"y{index}" for index in range(1, outputs)] + declared = [ + helper.make_empty_tensor_value_info(name) if output is None else output + for name in names + ] + return _model( + [helper.make_node(op_type, ["x"], names, name="pool", **attributes)], + [_tensor("x", elem_type, x_shape)], + declared, + ) + + +def test_poolings_of_one_fold_and_element_type_share_a_kernel(tmp_path): + """The geometry reaches the kernel as call-site literals, so neither rank nor a window + the size of the whole operand is a kernel of its own: a 1-D average pooling, a 2-D one + and a GlobalAveragePool all run the same code at different arguments.""" + model = _model( + [ + helper.make_node( + "AveragePool", ["x"], ["p"], name="signal", kernel_shape=[3] + ), + helper.make_node( + "AveragePool", ["i"], ["q"], name="image", kernel_shape=[3, 3] + ), + helper.make_node("GlobalAveragePool", ["i"], ["r"], name="whole"), + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 2, 6)), + _tensor("i", TensorProto.FLOAT, (2, 4, 5, 5)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q", "r")], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "pool") + assert kernel.endswith("_pool_average_float") + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 4 + + +def test_a_global_pooling_takes_the_whole_spatial_extent_as_one_window(tmp_path): + """One window per channel, covering every position, which is what ONNX defines it as. + + Six planes of twenty positions each, folded into one position by a window of twenty. + """ + report, header = _compile( + _pool_model((2, 3, 5, 4), op_type="GlobalAveragePool"), tmp_path + ) + + (kernel,) = _kernels(report, "pool") + call = header.split(f"{kernel}(\n")[-1].splitlines()[:6] + assert [line.strip(" ,);") for line in call] == ["y", "x", "6u", "20u", "1u", "20u"] + assert report["entrypoint"]["outputs"][0]["shape"] == [2, 3, 1, 1] + + +def test_count_include_pad_is_a_kernel_argument_rather_than_a_kernel(tmp_path): + """Two poolings differing only in what they divide by share one kernel, at 0u and 1u.""" + model = _model( + [ + helper.make_node( + "AveragePool", ["x"], ["p"], name="bare", kernel_shape=[3], pads=[1, 1] + ), + helper.make_node( + "AveragePool", + ["x"], + ["q"], + name="padded", + kernel_shape=[3], + pads=[1, 1], + count_include_pad=1, + ), + ], + [_tensor("x", TensorProto.FLOAT, (1, 2, 6))], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q")], + ) + + report, header = _compile(model, tmp_path) + + assert len(_kernels(report, "pool")) == 1 + assert header.count(" 0u);") == 1 + assert header.count(" 1u);") == 1 + + +def test_the_padding_a_dilated_same_mode_implies_reaches_the_kernel_as_literals( + tmp_path, +): + """`auto_pad` is resolved at compile time, for the window's *dilated* reach. + + A 2-tap window dilated to a reach of 4 over a 5-wide axis at stride 2 needs three pads, + and SAME_LOWER puts the odd one at the beginning — two here against SAME_UPPER's one. The + second axis takes an undilated 2-tap window at unit stride, which needs one. + """ + model = _pool_model( + (1, 1, 5, 5), + kernel_shape=[2, 2], + auto_pad="SAME_LOWER", + strides=[2, 1], + dilations=[3, 1], + ) + + _, header = _compile(model, tmp_path) + + assert "(const ptrdiff_t[]){2, 1}" in header + + +def test_a_valid_pooling_pads_nothing(tmp_path): + """`VALID` is not `SAME`: it drops the positions a full window does not fit in. + + ONNX's own shape inference arbitrates this, and the dilated reach with it. It derives the + result's shape from the pads the mode implies, and the kernel refuses to be emitted + against a buffer its addressing disagrees with — so a compiler reading `VALID` as `SAME`, + or measuring the window by its tap count rather than its reach, would fail to compile + this model rather than quietly pool the wrong positions. + """ + report, header = _compile( + _pool_model( + (1, 1, 7, 7), kernel_shape=[3, 3], auto_pad="VALID", dilations=[2, 1] + ), + tmp_path, + ) + + assert report["entrypoint"]["outputs"][0]["shape"] == [1, 1, 3, 5] + assert "(const ptrdiff_t[]){0, 0}" in header + + +@pytest.mark.parametrize( + ("extent", "attributes", "shape"), + [ + # A 3-tap window at stride 2 fits a 4-wide axis one and a half times: rounding down + # drops the half window, rounding up keeps it and reads the two taps of it that land + # on the operand. + (4, {"kernel_shape": [3, 3], "strides": [2, 2]}, [1, 3, 1, 1]), + (4, {"kernel_shape": [3, 3], "strides": [2, 2], "ceil_mode": 1}, [1, 3, 2, 2]), + # Rounding up here would put a second window's own start at 3, level with the end of + # the padded operand, where it would cover nothing but pad: ONNX drops it again. + ( + 2, + { + "kernel_shape": [3, 3], + "strides": [3, 3], + "pads": [1, 1, 1, 1], + "ceil_mode": 1, + }, + [1, 3, 1, 1], + ), + ], +) +def test_ceil_mode_decides_how_many_windows_fit(tmp_path, extent, attributes, shape): + """What the compiler counts is checked against what ONNX's shape inference counted.""" + report, _ = _compile(_pool_model((1, 3, extent, extent), **attributes), tmp_path) + + assert report["entrypoint"]["outputs"][0]["shape"] == shape + + +def test_a_pooling_writing_no_elements_emits_no_call(tmp_path): + """An empty batch leaves nothing to pool, and no loop that could read past a buffer.""" + report, header = _compile(_pool_model((0, 2, 5, 5), kernel_shape=[3, 3]), tmp_path) + + assert not _kernels(report, "pool") + assert "pool" not in header.split(f"int {report['prefix']}_run(")[-1] + + +def test_the_indexed_kernel_is_emitted_only_for_a_node_that_asks_for_indices(tmp_path): + """Reporting where each maximum came from is a kernel of its own; the plain fold is not + burdened with it, and the two do not share a name.""" + plain, _ = _compile( + _pool_model((1, 2, 5, 5), op_type="MaxPool", kernel_shape=[2, 2]), tmp_path + ) + indexed, _ = _compile( + _pool_model((1, 2, 5, 5), op_type="MaxPool", kernel_shape=[2, 2], outputs=2), + tmp_path / "indexed", + ) + + assert [name.split("_pool_")[-1] for name in _kernels(plain, "pool")] == [ + "max_float" + ] + assert [name.split("_pool_")[-1] for name in _kernels(indexed, "pool")] == [ + "max_indexed_float" + ] + + +@pytest.mark.parametrize( + ("storage_order", "strides", "rejected"), + [ + (0, "(const size_t[]){5u, 1u}", "(const size_t[]){1u, 5u}"), + (1, "(const size_t[]){1u, 5u}", "(const size_t[]){5u, 1u}"), + ], +) +def test_storage_order_chooses_the_strides_the_indices_are_reported_in( + tmp_path, storage_order, strides, rejected +): + """ONNX defines the second output as one flat index per maximum, laid out row-major or + column-major; both are the same walk over the operand at different strides.""" + model = _pool_model( + (1, 2, 5, 5), + op_type="MaxPool", + kernel_shape=[2, 2], + outputs=2, + storage_order=storage_order, + ) + + _, header = _compile(model, tmp_path) + + assert strides in header + assert rejected not in header + + +def _positions(*values): + return np.array(values, np.int64).reshape(1, 1, 2, 2) + + +@requires_c_compiler +def test_an_unpooled_position_outside_the_result_returns_a_nonzero_status(tmp_path): + """ONNX leaves an index past the end of the result undefined; the artifact reports it.""" + model = _model( + [ + helper.make_node( + "MaxUnpool", + ["x", "i"], + ["y"], + name="unpool", + kernel_shape=[2, 2], + strides=[2, 2], + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 1, 2, 2)), + _tensor("i", TensorProto.INT64, (1, 1, 2, 2)), + ], + [helper.make_empty_tensor_value_info("y")], + ) + compiled = compile_onnx(model, tmp_path).load() + x = np.arange(4, dtype=np.float32).reshape(1, 1, 2, 2) + + inside = compiled.run({"x": x, "i": _positions(0, 5, 10, 15)})["y"] + with pytest.raises(HarnessError, match="status"): + compiled.run({"x": x, "i": _positions(0, 5, 10, 16)}) + + assert inside.reshape(-1).tolist() == [ + 0, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 2, + 0, + 0, + 0, + 0, + 3, + ] + + +@pytest.mark.parametrize( + ("op_type", "attributes", "message"), + [ + ("AveragePool", {}, "states no `kernel_shape`"), + ("MaxPool", {}, "states no `kernel_shape`"), + ("LpPool", {}, "states no `kernel_shape`"), + ( + "AveragePool", + {"kernel_shape": [2]}, + "was given 1 `kernel_shape` for 2 spatial axis/axes", + ), + ( + "AveragePool", + {"kernel_shape": [2, 2], "auto_pad": "SAME_UPPER", "pads": [1, 1, 1, 1]}, + "mutually exclusive", + ), + ( + "AveragePool", + {"kernel_shape": [2, 2], "auto_pad": "SAME"}, + "is not one of the modes ONNX defines", + ), + ( + "AveragePool", + {"kernel_shape": [2, 2], "strides": [0, 1]}, + "ONNX defines them as positive", + ), + ( + "MaxPool", + {"kernel_shape": [2, 2], "dilations": [1, 0]}, + "ONNX defines them as positive", + ), + ( + "AveragePool", + {"kernel_shape": [2, 2], "pads": [1, 1]}, + "was given 2 pad(s) for 2 spatial axis/axes", + ), + ("LpPool", {"kernel_shape": [2, 2], "p": 0}, "which is positive"), + ], +) +def test_a_pooling_the_compiler_cannot_place_is_rejected( + tmp_path, op_type, attributes, message +): + model = _pool_model( + (1, 2, 5, 5), + op_type=op_type, + output=_tensor("y", TensorProto.FLOAT, (1, 2, 4, 4)), + **attributes, + ) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +def test_a_pooling_of_a_tensor_with_no_channel_axis_is_rejected(tmp_path): + model = _pool_model( + (2, 5), + kernel_shape=[2], + output=_tensor("y", TensorProto.FLOAT, (2, 4)), + ) + + with pytest.raises(CompileError, match="rank 3 or more"): + compile_onnx(model, tmp_path) + + +def test_unpooling_one_index_per_value_is_required(tmp_path): + """ONNX's shape inference reads neither operand's extent against the other's.""" + model = _model( + [ + helper.make_node( + "MaxUnpool", ["x", "i"], ["y"], name="unpool", kernel_shape=[2, 2] + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 1, 2, 2)), + _tensor("i", TensorProto.INT64, (1, 1, 3, 3)), + ], + [_tensor("y", TensorProto.FLOAT, (1, 1, 3, 3))], + ) + + with pytest.raises(CompileError, match="one index per value"): + compile_onnx(model, tmp_path) + + +def test_an_unpooling_shaped_at_run_time_is_rejected(tmp_path): + """`output_shape` is an operand, so a graph may compute it; one it does not fix makes the + result's shape a function of input data, which no binding can make static.""" + model = _model( + [ + helper.make_node( + "MaxUnpool", + ["x", "i", "s"], + ["y"], + name="unpool", + kernel_shape=[2, 2], + strides=[2, 2], + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 1, 2, 2)), + _tensor("i", TensorProto.INT64, (1, 1, 2, 2)), + _tensor("s", TensorProto.INT64, (4,)), + ], + [_tensor("y", TensorProto.FLOAT, (1, 1, 4, 4))], + ) + + with pytest.raises(CompileError, match="takes the shape of its `MaxUnpool` output"): + compile_onnx(model, tmp_path) + + +# -------------------------------------------------------------------------------------- +# The recurrent layers +# -------------------------------------------------------------------------------------- + +# What separates the three layers: how many gate rows the weights carry, and the two operands +# and third result only the LSTM's cell state brings. Everything else — the batch loop, the +# per-item sequence length, the padding past a sequence's end, the state outputs — is one +# shared frame, which is why most of what follows is parametrized over the family. +_LAYERS = ("LSTM", "GRU", "RNN") +_GATES = {"LSTM": 4, "GRU": 3, "RNN": 1} +_ALL_OPERANDS = ("X", "W", "R", "B", "sequence_lens", "initial_h", "initial_c", "P") +_ALL_RESULTS = ("Y", "Y_h", "Y_c") + + +def _layer_operands(op): + return _ALL_OPERANDS if op == "LSTM" else _ALL_OPERANDS[:6] + + +def _layer_results(op): + return _ALL_RESULTS if op == "LSTM" else _ALL_RESULTS[:2] + + +def _recurrent_feeds( + op, + *, + hidden=3, + seq=4, + batch=2, + inputs=3, + direction="forward", + optional=(), + lengths=None, + layout=0, + seed=11, +): + """Seeded operands for one recurrent node; `optional` names those past `X`, `W`, `R`.""" + generator = np.random.default_rng(seed) + rows = _GATES[op] * hidden + directions = 2 if direction == "bidirectional" else 1 + sequences = (batch, seq, inputs) if layout else (seq, batch, inputs) + state = (batch, directions, hidden) if layout else (directions, batch, hidden) + + def draw(shape, scale=1.0): + return (generator.normal(size=shape) * scale).astype(np.float32) + + feeds = { + "X": draw(sequences), + "W": draw((directions, rows, inputs), 0.3), + "R": draw((directions, rows, hidden), 0.3), + } + available = { + "B": lambda: draw((directions, 2 * rows), 0.2), + "sequence_lens": lambda: np.array(lengths, dtype=np.int32), + "initial_h": lambda: draw(state), + "initial_c": lambda: draw(state), + "P": lambda: draw((directions, 3 * hidden), 0.4), + } + for name in _layer_operands(op): + if name in optional: + feeds[name] = available[name]() + return feeds + + +def _recurrent_model(op, feeds, *, outputs=None, hidden=3, **attributes): + outputs = _layer_results(op) if outputs is None else outputs + names = [name if name in feeds else "" for name in _layer_operands(op)] + while names and not names[-1]: + names.pop() + node = helper.make_node( + op, + names, + list(outputs), + name=op.lower(), + **{"hidden_size": hidden, **attributes}, + ) + return _model( + [node], + [ + _tensor( + name, + TensorProto.INT32 if name == "sequence_lens" else TensorProto.FLOAT, + value.shape, + ) + for name, value in feeds.items() + ], + [helper.make_empty_tensor_value_info(name) for name in outputs if name], + ) + + +def _recurrent_error_model(op, *, outputs=("Y",), shapes=None, **attributes): + """A model whose every tensor is declared, so ONNX's own inference derives none of them. + + Which is what leaves the kernel's own checks reachable: a shape ONNX would have inferred, + and refused to, is one the compiler has to refuse itself. + """ + rows = 3 * _GATES[op] + declared = { + "X": (4, 2, 3), + "W": (1, rows, 3), + "R": (1, rows, 3), + "B": (1, 2 * rows), + "sequence_lens": (2,), + "initial_h": (1, 2, 3), + "initial_c": (1, 2, 3), + "P": (1, 9), + } + operands = {name: declared[name] for name in _layer_operands(op)} + operands.update(shapes or {}) + results = {"Y": (4, 1, 2, 3), "Y_h": (1, 2, 3), "Y_c": (1, 2, 3)} + return _model( + [ + helper.make_node( + op, + list(operands), + list(outputs), + name=op.lower(), + **{"hidden_size": 3, **attributes}, + ) + ], + [ + _tensor( + name, + TensorProto.INT32 if name == "sequence_lens" else TensorProto.FLOAT, + shape, + ) + for name, shape in operands.items() + ], + [_tensor(name, TensorProto.FLOAT, results[name]) for name in outputs], + ) + + +# Where a call site's argument list carries what, in the kernel's own parameter order: the +# operand pointers and scratch buffers, the four extents, the five strides that place them, +# and the flags the attributes become. The LSTM's pointers run five longer — a cell state to +# read, one to report and peephole weights, plus the buffer it carries the cell in — so +# everything after them sits further along its list. +_STRIDES = {"LSTM": slice(18, 23), "GRU": slice(14, 19), "RNN": slice(14, 19)} +_FLAGS = {"LSTM": slice(23, 27), "GRU": slice(19, 23), "RNN": slice(19, 22)} + + +def _recurrent_call(header, kernel, index=0): + """One emitted call site's arguments, in order.""" + body = header.split(f"if ({kernel}(\n")[index + 1].split(") != 0)")[0] + return [line.strip().rstrip(",") for line in body.splitlines()] + + +def _recurrent_kernel(report, op): + """The op's own kernel, which the `rnnact_`/`rnnclip_` helpers must not be mistaken for.""" + (kernel,) = _kernels(report, f"{op.lower()}_float") + return kernel + + +def _assert_matches_onnxruntime(tmp_path, model, feeds): + """Run the compiled artifact and onnxruntime on the same node, and compare. + + The reference evaluator implements a slice of each recurrent op only -- one forward + direction, with `clip` and `sequence_lens` ignored outright, and for the LSTM the + activations too -- so the differential sweep covers that slice and everything else rests + on the second oracle the compiler's parity test already stands on: onnxruntime, which is + neither this compiler nor the evaluator. Nothing here states an expected value of its own. + """ + runtime = pytest.importorskip("onnxruntime") + runtime.set_default_logger_severity(3) + session = runtime.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + expected = session.run(None, feeds) + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + for entry, want in zip(model.graph.output, expected): + np.testing.assert_allclose( + outputs[entry.name], want, rtol=1e-5, atol=1e-6, err_msg=entry.name + ) + + +@pytest.mark.parametrize("op", _LAYERS) +def test_recurrent_nodes_of_one_element_type_share_a_kernel(tmp_path, op): + """Sequence length, batch, width and layout are call-site literals, not kernels.""" + rows = _GATES[op] + model = _model( + [ + helper.make_node(op, ["x", "w", "r"], ["y"], name="wide", hidden_size=3), + helper.make_node( + op, ["s", "v", "q"], ["z"], name="narrow", hidden_size=2, layout=1 + ), + ], + [ + _tensor("x", TensorProto.FLOAT, (4, 2, 3)), + _tensor("w", TensorProto.FLOAT, (1, 3 * rows, 3)), + _tensor("r", TensorProto.FLOAT, (1, 3 * rows, 3)), + _tensor("s", TensorProto.FLOAT, (2, 5, 1)), + _tensor("v", TensorProto.FLOAT, (1, 2 * rows, 1)), + _tensor("q", TensorProto.FLOAT, (1, 2 * rows, 2)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("y", "z")], + ) + + report, header = _compile(model, tmp_path) + + kernel = _recurrent_kernel(report, op) + assert header.count(f"static int {kernel}(") == 1 + assert header.count(f"if ({kernel}(") == 2 + # The wider node's state decides the shared scratch; the narrower one reuses it. + symbol = f"{report['prefix']}_{op.lower()}" + assert f"static float {symbol}_hidden_float[3];" in header + assert f"static float {symbol}_gates_float[{3 * rows}];" in header + + +@pytest.mark.parametrize( + ("op", "weights"), + [ + ("LSTM", ("W + 48", "R + 36")), + ("GRU", ("W + 36", "R + 27")), + ("RNN", ("W + 12", "R + 9")), + ], +) +def test_a_bidirectional_layer_is_its_two_passes_over_one_kernel(tmp_path, op, weights): + """One direction per call, differing only in what the second direction reads and writes. + + Its own half of the weights, its own slice of the results — one direction's worth of + hidden units into each — and the flag that walks time backwards; everything else is the + same call. + """ + feeds = _recurrent_feeds( + op, direction="bidirectional", hidden=3, seq=4, batch=2, inputs=4 + ) + model = _recurrent_model(op, feeds, outputs=("Y", "Y_h"), direction="bidirectional") + + report, header = _compile(model, tmp_path) + + kernel = _recurrent_kernel(report, op) + assert header.count(f"if ({kernel}(") == 2 + forward = _recurrent_call(header, kernel) + backward = _recurrent_call(header, kernel, 1) + assert len(forward) == len(backward) + assert [pair for pair in zip(forward, backward) if pair[0] != pair[1]] == [ + ("Y", "Y + 6"), + ("Y_h", "Y_h + 6"), + ("W", weights[0]), + ("R", weights[1]), + ("0", "1"), + ] + + +@pytest.mark.parametrize("op", _LAYERS) +def test_the_layout_reaches_the_kernel_as_strides(tmp_path, op): + """Layout 1 packs the batch outermost, which changes only how far apart two steps are.""" + emissions = {} + for layout in (0, 1): + feeds = _recurrent_feeds(op, hidden=3, seq=4, batch=2, inputs=3, layout=layout) + emissions[layout] = _compile( + _recurrent_model(op, feeds, layout=layout), tmp_path / str(layout) + ) + + # Time first: one step of a sequence is a whole batch of input rows away (2 * 3) and one + # batch item is a single row (3). Batch first: a step is one row away and an item a whole + # sequence of them (4 * 3). The results follow the same reordering. + kernel = _recurrent_kernel(emissions[0][0], op) + assert _recurrent_kernel(emissions[1][0], op) == kernel + strides = { + layout: _recurrent_call(header, kernel)[_STRIDES[op]] + for layout, (_, header) in emissions.items() + } + assert strides[0] == ["6u", "3u", "6u", "3u", "3u"] + assert strides[1] == ["3u", "12u", "3u", "12u", "3u"] + + +@pytest.mark.parametrize( + ("op", "attributes", "flags"), + [ + ("LSTM", {}, ["0", "0", "0", "0.0f"]), + ("LSTM", {"direction": "reverse"}, ["1", "0", "0", "0.0f"]), + ("LSTM", {"input_forget": 1}, ["0", "1", "0", "0.0f"]), + ("LSTM", {"clip": 0.5}, ["0", "0", "1", "0.5f"]), + ("GRU", {}, ["0", "0", "0", "0.0f"]), + ("GRU", {"direction": "reverse"}, ["1", "0", "0", "0.0f"]), + ("GRU", {"linear_before_reset": 1}, ["0", "1", "0", "0.0f"]), + ("GRU", {"clip": 0.5}, ["0", "0", "1", "0.5f"]), + # An RNN has no second mode to switch, so its flags are the direction and the clip. + ("RNN", {}, ["0", "0", "0.0f"]), + ("RNN", {"direction": "reverse"}, ["1", "0", "0.0f"]), + ("RNN", {"clip": 0.5}, ["0", "1", "0.5f"]), + ], +) +def test_the_attributes_that_only_switch_a_branch_are_kernel_arguments( + tmp_path, op, attributes, flags +): + """Direction, the op's own mode and the cell clip pick a branch: one kernel serves all.""" + feeds = _recurrent_feeds(op) + report, header = _compile(_recurrent_model(op, feeds, **attributes), tmp_path) + + kernel = _recurrent_kernel(report, op) + assert _recurrent_call(header, kernel)[_FLAGS[op]] == flags + + +@pytest.mark.parametrize("op", _LAYERS) +def test_an_output_the_node_drops_is_never_written(tmp_path, op): + """The optional results are pointers the kernel checks, so a dropped one costs no buffer.""" + feeds = _recurrent_feeds(op) + report, header = _compile( + _recurrent_model(op, feeds, outputs=("", "Y_h")), tmp_path + ) + + kernel = _recurrent_kernel(report, op) + dropped = ["NULL", "Y_h"] + ["NULL"] * (len(_layer_results(op)) - 2) + assert _recurrent_call(header, kernel)[: len(dropped)] == dropped + assert [entry["name"] for entry in report["entrypoint"]["outputs"]] == ["Y_h"] + + +@pytest.mark.parametrize( + ("op", "chosen", "dropped"), + [ + ("LSTM", ["Relu", "Softsign", "Tanh"], "sigmoid"), + ("GRU", ["Relu", "Softsign"], "sigmoid"), + # An RNN runs `Tanh` by default and nothing else, so naming another drops it too. + ("RNN", ["Relu"], "tanh"), + ], +) +def test_a_recurrent_layer_emits_only_the_activations_it_names( + tmp_path, op, chosen, dropped +): + """One function per activation, shared by the gates that run it and by other nodes.""" + feeds = _recurrent_feeds(op) + report, default = _compile(_recurrent_model(op, feeds), tmp_path / "default") + prefix = report["prefix"] + picked = _compile( + _recurrent_model(op, feeds, activations=chosen), tmp_path / "chosen" + )[1] + + assert f"static float {prefix}_rnnact_{dropped}_float(" in default + assert f"{prefix}_rnnact_relu_float" not in default + assert f"static float {prefix}_rnnact_relu_float(" in picked + assert f"{prefix}_rnnact_{dropped}_float" not in picked + + +@pytest.mark.parametrize( + ("op", "case", "feeds", "attributes"), + [ + ( + "LSTM", + "reverse", + {"optional": ("B", "initial_h", "initial_c", "P")}, + {"direction": "reverse"}, + ), + ( + "LSTM", + "bidirectional", + { + "direction": "bidirectional", + "optional": ("B", "initial_h", "initial_c", "P"), + }, + {"direction": "bidirectional"}, + ), + ("LSTM", "clip", {"optional": ("B",)}, {"clip": 0.4}), + # The clip bounds a gate's whole pre-activation, its peephole term included, and the + # output gate's after the cell that term reads has been updated. + ("LSTM", "clip_with_peepholes", {"optional": ("B", "P")}, {"clip": 0.4}), + ("LSTM", "input_forget", {"optional": ("B", "P")}, {"input_forget": 1}), + ( + "LSTM", + "activations", + {"optional": ("B",)}, + {"activations": ["Relu", "Softsign", "HardSigmoid"]}, + ), + ( + "LSTM", + "parameterized_activations", + {"optional": ("B",)}, + { + "activations": ["LeakyRelu", "LeakyRelu", "LeakyRelu"], + "activation_alpha": [0.3, 0.4, 0.5], + }, + ), + # `activation_alpha` and `activation_beta` carry a value for the activations that + # take one and for no others, so `Sigmoid` here consumes neither: the `LeakyRelu` + # reads the first alpha and the `HardSigmoid` the second, along with the only beta. + ( + "LSTM", + "activation_parameters_consumed_where_they_are_taken", + {"optional": ("B",)}, + { + "activations": ["Sigmoid", "LeakyRelu", "HardSigmoid"], + "activation_alpha": [0.3, 0.9], + "activation_beta": [0.7], + }, + ), + # And they are consumed over both directions together, not restarted per direction. + ( + "LSTM", + "activation_parameters_across_directions", + {"direction": "bidirectional", "optional": ("B",)}, + { + "direction": "bidirectional", + "activations": ["LeakyRelu", "Tanh", "Tanh"] * 2, + "activation_alpha": [0.3, 0.8], + }, + ), + ( + "LSTM", + "short_sequences", + { + "optional": ("B", "sequence_lens", "initial_h", "initial_c"), + "lengths": [4, 2], + }, + {}, + ), + ( + "LSTM", + "empty_sequence", + { + "optional": ("B", "sequence_lens", "initial_h", "initial_c"), + "lengths": [3, 0], + }, + {}, + ), + ( + "LSTM", + "backwards_over_short_sequences", + {"optional": ("B", "sequence_lens"), "lengths": [1, 3]}, + {"direction": "reverse"}, + ), + ( + "GRU", + "reverse", + {"optional": ("B", "initial_h")}, + {"direction": "reverse"}, + ), + ( + "GRU", + "bidirectional", + {"direction": "bidirectional", "optional": ("B", "initial_h")}, + {"direction": "bidirectional"}, + ), + ("GRU", "clip", {"optional": ("B",)}, {"clip": 0.4}), + # `linear_before_reset` moves the candidate's recurrent bias inside the term the + # reset gate scales, so the two branches differ only where a bias is present — and + # differ in what the reset gate multiplies whether one is or not. + ( + "GRU", + "linear_before_reset", + {"optional": ("B",)}, + {"linear_before_reset": 1}, + ), + ("GRU", "linear_before_reset_unbiased", {}, {"linear_before_reset": 1}), + ( + "GRU", + "linear_before_reset_clipped", + {"optional": ("B",)}, + {"linear_before_reset": 1, "clip": 0.3}, + ), + ( + "GRU", + "activations", + {"optional": ("B",)}, + {"activations": ["Relu", "Softsign"]}, + ), + ( + "GRU", + "parameterized_activations", + {"optional": ("B",)}, + { + "activations": ["LeakyRelu", "HardSigmoid"], + "activation_alpha": [0.3, 0.9], + "activation_beta": [0.7], + }, + ), + ( + "GRU", + "activation_parameters_across_directions", + {"direction": "bidirectional", "optional": ("B",)}, + { + "direction": "bidirectional", + "activations": ["LeakyRelu", "Tanh"] * 2, + "activation_alpha": [0.3, 0.8], + }, + ), + ( + "GRU", + "short_sequences", + {"optional": ("B", "sequence_lens", "initial_h"), "lengths": [4, 2]}, + {}, + ), + ( + "GRU", + "empty_sequence", + {"optional": ("B", "sequence_lens", "initial_h"), "lengths": [3, 0]}, + {}, + ), + ( + "GRU", + "backwards_over_short_sequences", + {"optional": ("B", "sequence_lens"), "lengths": [1, 3]}, + {"direction": "reverse"}, + ), + ( + "RNN", + "reverse", + {"optional": ("B", "initial_h")}, + {"direction": "reverse"}, + ), + ( + "RNN", + "bidirectional", + {"direction": "bidirectional", "optional": ("B", "initial_h")}, + {"direction": "bidirectional"}, + ), + ("RNN", "clip", {"optional": ("B",)}, {"clip": 0.4}), + ("RNN", "activations", {"optional": ("B",)}, {"activations": ["Relu"]}), + ( + "RNN", + "parameterized_activations", + {"optional": ("B",)}, + { + "activations": ["HardSigmoid"], + "activation_alpha": [0.3], + "activation_beta": [0.7], + }, + ), + # onnxruntime reads this attribute per direction rather than per parameterized + # activation, and walks off the end of its own vector when given fewer values than + # it has directions -- so the case that separates the two readings is the LSTM's. + ( + "RNN", + "activation_parameters_across_directions", + {"direction": "bidirectional", "optional": ("B",)}, + { + "direction": "bidirectional", + "activations": ["LeakyRelu", "LeakyRelu"], + "activation_alpha": [0.3, 0.8], + }, + ), + ( + "RNN", + "short_sequences", + {"optional": ("B", "sequence_lens", "initial_h"), "lengths": [4, 2]}, + {}, + ), + ( + "RNN", + "backwards_over_short_sequences", + {"optional": ("B", "sequence_lens"), "lengths": [1, 3]}, + {"direction": "reverse"}, + ), + ], +) +@requires_c_compiler +def test_the_recurrent_surface_the_evaluator_cannot_vouch_for( + tmp_path, op, case, feeds, attributes +): + """Every attribute ONNX's reference evaluator drops, against onnxruntime instead. + + An empty sequence is asked of the LSTM and the GRU only: onnxruntime leaves an RNN's + state for a zero-length item uninitialized, so there is nothing there to compare against. + """ + operands = _recurrent_feeds(op, **feeds) + + _assert_matches_onnxruntime( + tmp_path, _recurrent_model(op, operands, **attributes), operands + ) + + +@requires_c_compiler +def test_the_lstm_cell_state_output_matches_onnxruntime(tmp_path): + """`Y_c` is the third output, which the reference evaluator never returns at all.""" + feeds = _recurrent_feeds("LSTM", optional=("B", "initial_h", "initial_c", "P")) + + _assert_matches_onnxruntime( + tmp_path, _recurrent_model("LSTM", feeds, outputs=("", "", "Y_c")), feeds + ) + + +@pytest.mark.parametrize("op", _LAYERS) +@pytest.mark.parametrize("direction", ["forward", "reverse", "bidirectional"]) +@requires_c_compiler +def test_a_batchwise_layer_runs_what_the_time_first_one_does(tmp_path, op, direction): + """Layout 1 over several steps, against the layout its two oracles cover. + + onnxruntime refuses `layout` 1 outright and the reference evaluator's own layout path + holds for a single step only, which leaves the corpus's one batchwise test per op -- also + a single step -- as the whole of the direct evidence. But layout 1 means only that the + operands are packed batch-outermost, so transposing them onto a time-first node of the + same layer has to reproduce the run down to the bit: same arithmetic in the same order, + reaching the same elements from elsewhere in memory. + """ + packed = ("X", "initial_h", "initial_c") + optional = ("B", "sequence_lens", "initial_h", "initial_c", "P") + flat = _recurrent_feeds(op, direction=direction, optional=optional, lengths=[4, 2]) + batchwise = { + name: value.transpose(1, 0, 2).copy() if name in packed else value + for name, value in flat.items() + } + + time_first = ( + compile_onnx(_recurrent_model(op, flat, direction=direction), tmp_path / "time") + .load() + .run(flat) + ) + batch_first = ( + compile_onnx( + _recurrent_model(op, batchwise, direction=direction, layout=1), + tmp_path / "batch", + ) + .load() + .run(batchwise) + ) + + assert np.array_equal(batch_first["Y"], time_first["Y"].transpose(2, 0, 1, 3)) + for state in _layer_results(op)[1:]: + assert np.array_equal(batch_first[state], time_first[state].transpose(1, 0, 2)) + + +@pytest.mark.parametrize("op", _LAYERS) +@requires_c_compiler +def test_a_sequence_length_past_the_padded_end_returns_a_nonzero_status(tmp_path, op): + """A length names a step of the padded sequence; anything else names one that is absent.""" + feeds = _recurrent_feeds(op, optional=("sequence_lens",), lengths=[4, 4]) + compiled = compile_onnx(_recurrent_model(op, feeds), tmp_path).load() + + for lengths in ([4, 5], [-1, 2]): + with pytest.raises(HarnessError, match="status 1"): + compiled.run({**feeds, "sequence_lens": np.array(lengths, dtype=np.int32)}) + + +@pytest.mark.parametrize("op", _LAYERS) +@pytest.mark.parametrize( + ("attributes", "message"), + [ + ({"direction": "backward"}, "not one of the directions ONNX defines"), + ({"layout": 2}, "ONNX defines only 0 (time first) and 1 (batch first)"), + ({"clip": -1.0}, "which is not negative"), + ({"hidden_size": 4}, "states a `hidden_size` of 4"), + ], +) +def test_a_recurrent_node_the_compiler_cannot_place_is_rejected( + tmp_path, op, attributes, message +): + model = _recurrent_error_model(op, **attributes) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("op", "attributes", "message"), + [ + ( + "LSTM", + {"activations": ["Sigmoid", "Tanh"]}, + "runs 3 activations — 3 per direction — but this node names 2", + ), + ( + "GRU", + {"activations": ["Sigmoid", "Tanh", "Tanh"]}, + "runs 2 activations — 2 per direction — but this node names 3", + ), + ( + "RNN", + {"activations": ["Tanh", "Tanh"]}, + "runs 1 activation — 1 per direction — but this node names 2", + ), + ( + "RNN", + {"direction": "bidirectional", "activations": ["Tanh"]}, + "runs 2 activations — 1 per direction — but this node names 1", + ), + ("RNN", {"activations": ["Swish"]}, "names the activation `Swish`"), + ( + "GRU", + {"activations": ["Affine", "Tanh"]}, + "there is no default `activation_alpha`", + ), + ( + "RNN", + {"activations": ["ScaledTanh"], "activation_alpha": [2.0]}, + "there is no default `activation_beta`", + ), + ], +) +def test_an_activation_a_layer_cannot_run_is_rejected( + tmp_path, op, attributes, message +): + """How many activations an op runs is the op's own, and so is what it may name.""" + model = _recurrent_error_model(op, **attributes) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("op", "name", "shape", "message"), + [ + ("LSTM", "X", (4, 2), "a tensor of rank 3"), + ("LSTM", "R", (12, 3), "recurrence weights as rank 3"), + ("LSTM", "W", (1, 12, 2), "reads `W` as its input weights of shape [1, 12, 3]"), + ("LSTM", "B", (1, 12), "reads `B` as its biases of shape [1, 24]"), + ("LSTM", "P", (1, 12), "reads `P` as its peephole weights of shape [1, 9]"), + ( + "LSTM", + "initial_h", + (1, 3, 3), + "as an initial hidden state of shape [1, 2, 3]", + ), + ("LSTM", "sequence_lens", (3,), "as its sequence lengths of shape [2]"), + ( + "GRU", + "W", + (1, 12, 3), + "`GRU` reads `W` as its input weights of shape [1, 9, 3]", + ), + ("GRU", "B", (1, 24), "`GRU` reads `B` as its biases of shape [1, 18]"), + ( + "GRU", + "initial_h", + (1, 3, 3), + "as an initial hidden state of shape [1, 2, 3]", + ), + ( + "RNN", + "W", + (1, 12, 3), + "`RNN` reads `W` as its input weights of shape [1, 3, 3]", + ), + ("RNN", "B", (1, 24), "`RNN` reads `B` as its biases of shape [1, 6]"), + ("RNN", "X", (4, 2), "a tensor of rank 3"), + ], +) +def test_a_recurrent_operand_shaped_for_another_layer_is_rejected( + tmp_path, op, name, shape, message +): + """Declared shapes get past ONNX's own inference, so the kernel checks them itself.""" + model = _recurrent_error_model(op, shapes={name: shape}) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize("op", _LAYERS) +@pytest.mark.parametrize("output", ["Y", "Y_h"]) +def test_a_recurrent_result_shaped_for_another_layer_is_rejected(tmp_path, op, output): + """The same, for the buffers it writes: a declared result is checked, never trusted.""" + model = _recurrent_error_model(op, outputs=(output,)) + entry = model.graph.output[0] + entry.type.tensor_type.shape.dim[0].dim_value = 7 + + with pytest.raises(CompileError, match="addresses a result of shape"): + compile_onnx(model, tmp_path) + + +# -------------------------------------------------------------------------------------- +# The resizes +# -------------------------------------------------------------------------------------- + +# What a resize reads its geometry from — the scales, the sizes and the region of interest — +# is passed as operands rather than attributes, so every test below has to say where each of +# them comes from. Carried in the model, they let ONNX's own shape inference derive the +# result; fed at run time, they leave the declared result shape as the only static one, which +# is the form the backend corpus's own Resize tests take. +_RESIZE_OPSET = 19 + +_RESIZE_OPERANDS = ( + ("roi", TensorProto.FLOAT), + ("scales", TensorProto.FLOAT), + ("sizes", TensorProto.INT64), +) + + +def _resize_model( + x_shape, + *, + op_type="Resize", + elem_type=TensorProto.FLOAT, + roi=None, + scales=None, + sizes=None, + runtime=(), + output=None, + opset=_RESIZE_OPSET, + name="resize", + **attributes, +): + """One resize node, its operands carried in the model unless `runtime` names them.""" + values = {"roi": roi, "scales": scales, "sizes": sizes} + inputs = [_tensor("x", elem_type, x_shape)] + initializer = [] + names = ["x"] + for operand, operand_type in _RESIZE_OPERANDS: + given = values[operand] + names.append("" if given is None else operand) + if given is None: + continue + array = np.array(given, dtype=helper.tensor_dtype_to_np_dtype(operand_type)) + if operand in runtime: + inputs.append(_tensor(operand, operand_type, array.shape)) + else: + initializer.append(onnx.numpy_helper.from_array(array, operand)) + if op_type == "Upsample": + names = ["x", "scales"] + while names and not names[-1]: + names.pop() + declared = ( + helper.make_empty_tensor_value_info("y") + if output is None + else _tensor("y", elem_type, output) + ) + return _model( + [helper.make_node(op_type, names, ["y"], name=name, **attributes)], + inputs, + [declared], + initializer=initializer, + opset=opset, + ) + + +def test_resizes_of_one_element_type_share_a_kernel(tmp_path): + """Every setting a resize reads is a kernel argument, not a kernel of its own. + + Three nodes interpolating differently, over different ranks and in both directions, run + the same code at different literals — the mode, the mapping and the rounding rule among + them. + """ + model = _model( + [ + helper.make_node( + "Resize", ["x", "", "up"], ["p"], name="grown", mode="linear" + ), + helper.make_node( + "Resize", + ["x", "", "down"], + ["q"], + name="shrunk", + mode="cubic", + antialias=1, + ), + helper.make_node( + "Resize", + ["s", "", "signal"], + ["r"], + name="signal", + mode="nearest", + nearest_mode="ceil", + ), + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 2, 4, 5)), + _tensor("s", TensorProto.FLOAT, (2, 3, 7)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q", "r")], + initializer=[ + onnx.numpy_helper.from_array(np.array(values, dtype=np.float32), name) + for name, values in ( + ("up", [1.0, 1.0, 2.0, 3.0]), + ("down", [1.0, 1.0, 0.5, 0.5]), + ("signal", [1.0, 1.0, 1.5]), + ) + ], + opset=_RESIZE_OPSET, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "resize_float") + assert kernel.endswith("_resize_float_float") + assert header.count(f"static int {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 4 + + +def test_the_geometry_reaches_the_kernel_as_call_site_literals(tmp_path): + """The shapes, the axes and the settings are compile-time constants at the call site.""" + model = _resize_model( + (1, 2, 4, 5), + scales=[0.5, 2.0], + axes=[3, 2], + mode="cubic", + coordinate_transformation_mode="align_corners", + exclude_outside=1, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "resize_float") + call = header.split(f"{kernel}(\n")[-1].splitlines() + assert [line.strip(" ,);") for line in call[:14]] == [ + "y", + "x", + "NULL", + f"{report['prefix']}_w_scales", + "NULL", + f"{report['prefix']}_resize_work", + f"{report['prefix']}_resize_spare", + "40u", + "32u", + "4", + "(const size_t[]){1u, 2u, 4u, 5u}", + "(const size_t[]){1u, 2u, 8u, 2u}", + "2", + "(const size_t[]){3u, 2u}", + ] + # mode `cubic`, the default rounding, `align_corners`, no antialias, exclude outside. + assert [line.strip(" ,);") for line in call[14:20]] == [ + "2", + "0", + "3", + "0", + "1", + "0", + ] + + +def test_an_upsample_is_the_resize_its_successor_defines_it_to_be(tmp_path): + """ONNX deprecated Upsample in favour of Resize's asymmetric mapping at the floor. + + Compiling the two ops on the same operands emits the same call to the same kernel — the + deprecated op is not a walk of its own, it is that one at the settings its successor's + specification spells out for it. + """ + emissions = [ + _compile( + _resize_model( + (1, 1, 4, 4), + op_type=op_type, + scales=[1.0, 1.0, 2.0, 1.5], + opset=opset, + **attributes, + ), + tmp_path / op_type, + ) + for op_type, opset, attributes in ( + ("Upsample", 9, {}), + ( + "Resize", + _RESIZE_OPSET, + { + "mode": "nearest", + "coordinate_transformation_mode": "asymmetric", + "nearest_mode": "floor", + }, + ), + ) + ] + + (kernel,) = _kernels(emissions[0][0], "resize_float") + assert _kernels(emissions[1][0], "resize_float") == [kernel] + settings = [ + [ + line.strip(" ,);") + for line in header.split(f"{kernel}(\n")[-1].splitlines()[14:20] + ] + for _, header in emissions + ] + # `nearest` at the floor of an asymmetrically mapped coordinate, no antialias, no + # exclusion and no aspect-ratio policy. + assert settings[0] == ["0", "2", "4", "0", "0", "0"] + assert settings[1] == settings[0] + + +def test_the_working_buffers_are_static_and_sized_for_the_widest_pass(tmp_path): + """A pass reads the whole result of the one before it, so two buffers are reserved. + + They are shared by every resize in the model and sized for the largest — each axis at + the larger of the two extents it carries — which is what the reported footprint holds. + """ + model = _model( + [ + helper.make_node("Resize", ["x", "", "up"], ["p"], name="grown"), + helper.make_node("Resize", ["x", "", "down"], ["q"], name="shrunk"), + ], + [_tensor("x", TensorProto.FLOAT, (1, 2, 4, 5))], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q")], + initializer=[ + onnx.numpy_helper.from_array(np.array(values, dtype=np.float32), name) + for name, values in ( + ("up", [1.0, 1.0, 3.0, 2.0]), + ("down", [1.0, 1.0, 0.5, 0.5]), + ) + ], + opset=_RESIZE_OPSET, + ) + + report, header = _compile(model, tmp_path) + + widest = 1 * 2 * 12 * 10 + for role in ("work", "spare"): + assert f"static double {report['prefix']}_resize_{role}[{widest}];" in header + # The two buffers, plus the results of the two nodes and the embedded scales. + assert report["memory"]["arena_bytes"] >= 2 * widest * 8 + + +@requires_c_compiler +def test_a_run_time_scale_the_artifact_was_not_compiled_for_is_reported(tmp_path): + """The result's extent follows from operand *values*, which only the caller has. + + Compiled against the shape the model declares, the artifact still derives the extents + the operands ask for and refuses, through the status enum, to write a result of any + other shape — so a scale it was not compiled for is an error rather than a buffer + written past. + """ + model = _resize_model( + (1, 1, 2, 2), + scales=[1.0, 1.0, 2.0, 3.0], + runtime=("scales",), + output=(1, 1, 4, 6), + mode="nearest", + ) + values = np.arange(4, dtype=np.float32).reshape(1, 1, 2, 2) + + compiled = compile_onnx(model, tmp_path).load() + + compiled_output = compiled.run( + {"x": values, "scales": np.array([1.0, 1.0, 2.0, 3.0], dtype=np.float32)} + )["y"] + expected = ReferenceEvaluator(model).run( + None, {"x": values, "scales": np.array([1.0, 1.0, 2.0, 3.0], dtype=np.float32)} + )[0] + np.testing.assert_array_equal(compiled_output, expected) + + with pytest.raises(HarnessError, match="status"): + compiled.run( + {"x": values, "scales": np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32)} + ) + + +@requires_c_compiler +def test_a_resize_writing_nothing_emits_no_kernel_at_all(tmp_path): + """A scale that shrinks an axis to nothing leaves a result with no elements in it.""" + model = _resize_model((1, 2, 4, 5), scales=[1.0, 1.0, 0.2, 1.0], mode="linear") + + report, header = _compile(model, tmp_path) + + assert report["entrypoint"]["outputs"][0]["shape"] == [1, 2, 0, 5] + assert not _kernels(report, "resize") + assert f"{report['prefix']}_resize" not in header + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({}, "exactly one of `scales` and `sizes`"), + # An operand holding nothing at all is one left out, which is how an exporter + # passes `sizes` without leaving the `scales` position blank. + ({"scales": []}, "exactly one of `scales` and `sizes`"), + ( + {"scales": [1.0, 1.0, 2.0, 2.0], "sizes": [1, 2, 8, 10]}, + "exactly one of `scales` and `sizes`", + ), + ({"scales": [1.0, 2.0]}, "takes `scales` as 4 value(s)"), + ( + { + "scales": [1.0, 1.0, 2.0, 2.0], + "roi": [0.0, 1.0], + "coordinate_transformation_mode": "tf_crop_and_resize", + }, + "takes `roi` as 8 value(s)", + ), + ( + {"scales": [1.0, 1.0, 2.0, 2.0], "mode": "bilinear"}, + "asks for `mode` `bilinear`", + ), + ( + { + "scales": [1.0, 1.0, 2.0, 2.0], + "coordinate_transformation_mode": "corners", + }, + "asks for `coordinate_transformation_mode` `corners`", + ), + ( + {"scales": [1.0, 1.0, 2.0, 2.0], "nearest_mode": "nearest"}, + "asks for `nearest_mode` `nearest`", + ), + ( + {"sizes": [1, 2, 8, 10], "keep_aspect_ratio_policy": "squash"}, + "asks for `keep_aspect_ratio_policy` `squash`", + ), + ( + {"scales": [1.0, 1.0, 2.0, 2.0], "antialias": 1}, + "`antialias` in `nearest` mode", + ), + ( + {"scales": [2.0, 2.0], "axes": [2, 2]}, + "names the same dimension more than once", + ), + ({"scales": [2.0, 2.0], "axes": [2, 4]}, "axis 4 is out of range"), + ], +) +def test_a_resize_the_compiler_cannot_serve_is_rejected(tmp_path, kwargs, message): + """Everything about a resize that its operands' shapes and its attributes settle. + + The operands are fed at run time here, which is what leaves these to the compiler at + all: given their values, ONNX's own inference rejects most of these models first. + """ + model = _resize_model( + (1, 2, 4, 5), runtime=("roi", "scales", "sizes"), output=(1, 2, 8, 10), **kwargs + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "`resize`" in str(error.value) + assert message in str(error.value) + + +def test_a_resize_of_a_boolean_tensor_is_rejected(tmp_path): + """ONNX allows one; there is no value between two truth values to interpolate to.""" + model = _resize_model( + (1, 2, 4, 5), + elem_type=TensorProto.BOOL, + scales=[1.0, 1.0, 2.0, 2.0], + mode="nearest", + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "`resize`" in str(error.value) + assert "BOOL" in str(error.value) + + +@pytest.mark.parametrize( + ("x_shape", "output", "message"), + [ + ((1, 2, 4, 5), (2, 2, 8, 10), "does not resize axis 0"), + ((1, 2, 0, 5), (1, 2, 3, 10), "which holds no elements"), + ], +) +def test_a_declared_result_the_axes_do_not_allow_is_rejected( + tmp_path, x_shape, output, message +): + """A result shape the model declares rather than computes is checked, never trusted. + + With the scales fed at run time, nothing but the declaration says what shape the result + has — so the axes the node does not resize have to carry the operand's own extents, and + an axis holding nothing has nothing to interpolate into one that does. + """ + model = _resize_model( + x_shape, + scales=[2.0, 2.0], + axes=[2, 3], + runtime=("scales",), + output=output, + mode="nearest", + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "`resize`" in str(error.value) + assert message in str(error.value) + + +@pytest.mark.parametrize( + ("op_type", "opset", "nearest"), + [("Resize", 18, 19), ("Upsample", 7, 9)], +) +def test_a_resize_below_its_supported_revision_names_the_nearest_version( + tmp_path, op_type, opset, nearest +): + """Resize was revised at 19 and Upsample took its scales as an attribute up to 7. + + Neither older revision has an oracle — the reference evaluator applies today's + implementation to them and the backend corpus has no test at one — so no kernel claims + them, and dispatch says so rather than serving the current walk under the old op's name. + """ + scales = [1.0, 1.0, 2.0, 2.0] + if op_type == "Upsample": + # The revision that took its scales as an attribute rather than as an operand. + model = _model( + [helper.make_node(op_type, ["x"], ["y"], name="node", scales=scales)], + [_tensor("x", TensorProto.FLOAT, (1, 1, 4, 4))], + [helper.make_empty_tensor_value_info("y")], + opset=opset, + ) + else: + model = _resize_model( + (1, 1, 4, 4), op_type=op_type, scales=scales, opset=opset, name="node" + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`node`" in message + assert f"opset version {opset}" in message + assert f"Nearest supported version: {nearest}" in message + + +# -------------------------------------------------------------------------------------- +# The samplers +# -------------------------------------------------------------------------------------- + +# What a sampler reads at run time is a coordinate, not an index, so what is asserted here is +# what the sweeps cannot reach: that the mode, the padding and the geometry are call-site +# literals rather than kernels of their own; that a coordinate no conversion is defined for +# reads nothing instead of reading past a buffer; and the errors for the combinations the +# compiler refuses outright. The two block shuffles are neither sampled nor computed — they +# are the strided move every view op runs — which is what their tests say. + +_AFFINE_GRID_OPSET = 20 +_COL2IM_OPSET = 18 +_BLOCK_OPSET = 13 + + +def _grid_sample_model( + x_shape, + grid_shape, + *, + elem_type=TensorProto.FLOAT, + grid_type=TensorProto.FLOAT, + name="sample", + **attributes, +): + return _model( + [helper.make_node("GridSample", ["x", "g"], ["y"], name=name, **attributes)], + [_tensor("x", elem_type, x_shape), _tensor("g", grid_type, grid_shape)], + [helper.make_empty_tensor_value_info("y")], + ) + + +def _roi_model( + op_type, + x_shape, + roi_count, + *, + elem_type=TensorProto.FLOAT, + columns=None, + indices=None, + name="roi", + **attributes, +): + """One region-of-interest pooling; MaxRoiPool carries its batch in the region itself.""" + columns = (4 if op_type == "RoiAlign" else 5) if columns is None else columns + names = ["x", "rois"] + inputs = [ + _tensor("x", elem_type, x_shape), + _tensor("rois", elem_type, (roi_count, columns)), + ] + if op_type == "RoiAlign": + names.append("batch_indices") + inputs.append( + _tensor( + "batch_indices", + TensorProto.INT64, + (roi_count,) if indices is None else indices, + ) + ) + return _model( + [helper.make_node(op_type, names, ["y"], name=name, **attributes)], + inputs, + [helper.make_empty_tensor_value_info("y")], + ) + + +def _affine_grid_model( + size, + *, + theta_shape=None, + elem_type=TensorProto.FLOAT, + runtime=False, + name="grid", + **attributes, +): + """One AffineGrid, its size carried in the model unless `runtime` feeds it instead.""" + rank = len(size) - 2 + theta_shape = (size[0], rank, rank + 1) if theta_shape is None else theta_shape + values = np.array(size, dtype=np.int64) + inputs = [_tensor("theta", elem_type, theta_shape)] + initializer = [] + if runtime: + inputs.append(_tensor("size", TensorProto.INT64, values.shape)) + else: + initializer.append(onnx.numpy_helper.from_array(values, "size")) + return _model( + [ + helper.make_node( + "AffineGrid", ["theta", "size"], ["y"], name=name, **attributes + ) + ], + inputs, + [helper.make_empty_tensor_value_info("y")], + initializer=initializer, + opset=_AFFINE_GRID_OPSET, + ) + + +def _col2im_model( + x_shape, + image, + block, + *, + elem_type=TensorProto.FLOAT, + runtime=(), + output=None, + name="fold", + **attributes, +): + """One Col2Im, its extents carried in the model unless `runtime` names them.""" + inputs = [_tensor("x", elem_type, x_shape)] + initializer = [] + for operand, extents in (("image_shape", image), ("block_shape", block)): + values = np.array(extents, dtype=np.int64) + if operand in runtime: + inputs.append(_tensor(operand, TensorProto.INT64, values.shape)) + else: + initializer.append(onnx.numpy_helper.from_array(values, operand)) + declared = ( + helper.make_empty_tensor_value_info("y") + if output is None + else _tensor("y", elem_type, output) + ) + return _model( + [ + helper.make_node( + "Col2Im", + ["x", "image_shape", "block_shape"], + ["y"], + name=name, + **attributes, + ) + ], + inputs, + [declared], + initializer=initializer, + opset=_COL2IM_OPSET, + ) + + +def _block_model( + op_type, x_shape, *, elem_type=TensorProto.FLOAT, name="block", **attributes +): + return _model( + [helper.make_node(op_type, ["x"], ["y"], name=name, **attributes)], + [_tensor("x", elem_type, x_shape)], + [helper.make_empty_tensor_value_info("y")], + opset=_BLOCK_OPSET, + ) + + +def _samplers_model(name="sample"): + """Four samplers over one operand: what they share and what they do not.""" + return _model( + [ + helper.make_node("GridSample", ["x", "g"], ["p"], name=name, mode="linear"), + helper.make_node( + "GridSample", + ["x", "g"], + ["q"], + name="cubic", + mode="cubic", + padding_mode="border", + align_corners=1, + ), + helper.make_node( + "GridSample", ["x", "g"], ["r"], name="near", mode="nearest" + ), + helper.make_node( + "RoiAlign", + ["x", "rois", "batch_indices"], + ["s"], + name="align", + output_height=2, + output_width=2, + ), + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 2, 4, 5)), + _tensor("g", TensorProto.FLOAT, (1, 3, 3, 2)), + _tensor("rois", TensorProto.FLOAT, (2, 4)), + _tensor("batch_indices", TensorProto.INT64, (2,)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q", "r", "s")], + ) + + +def test_interpolating_grid_samples_of_one_element_type_share_a_kernel(tmp_path): + """The mode, the padding and the corner convention are arguments, not kernels. + + Two nodes interpolating differently — one cubic against a clamped border, one linear + against zeros — run the same code at different literals. `nearest` is the one that is + not: it reads a single element and computes nothing, which is what lets it serve the + element types no interpolation is defined for. + """ + report, header = _compile(_samplers_model(), tmp_path) + + interpolating, nearest = sorted(_kernels(report, "gridsample")) + assert interpolating.endswith("_gridsample_float_float") + assert nearest.endswith("_gridsample_nearest_float_float") + assert header.count(f"static void {interpolating}(") == 1 + assert header.count(f"{interpolating}(\n") == 3 + + +def test_every_sampler_shares_one_padding_resolution(tmp_path): + """What a coordinate outside the operand reads is one decision, made in one place.""" + report, header = _compile(_samplers_model(), tmp_path) + + for helper_name, returns in ( + ("sample_reflect", "double"), + ("sample_index", "ptrdiff_t"), + ("sample_locate_float", "float"), + ("sample_coefficient_float", "float"), + ): + (shared,) = _kernels(report, helper_name) + assert header.count(f"static {returns} {shared}(") == 1 + + +def test_the_grid_sample_geometry_reaches_the_kernel_as_call_site_literals(tmp_path): + """The extents, the mode and the padding are compile-time constants at the call site.""" + model = _grid_sample_model( + (2, 3, 4, 5, 6), + (2, 2, 3, 4, 3), + mode="cubic", + padding_mode="reflection", + align_corners=1, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "gridsample") + call = header.split(f"{kernel}(\n")[-1].splitlines() + assert [line.strip(" ,);") for line in call[:12]] == [ + "y", + "x", + "g", + "2u", + "3u", + "120u", + "24u", + "3", + "(const size_t[]){4u, 5u, 6u}", + # `reflection` padding, aligned corners, and the cubic mode. + "2", + "1", + "2", + ] + + +@requires_c_compiler +@pytest.mark.parametrize("padding_mode", ["zeros", "border", "reflection"]) +def test_a_grid_coordinate_no_index_could_hold_reads_nothing(tmp_path, padding_mode): + """A coordinate arrives at run time, and not every one of them names an element. + + ONNX says nothing about a grid holding infinities or values that are not numbers at all — + its own reference raises on them — so what is asserted is only that the artifact answers + with a value from the operand or with zero, rather than reading wherever an out-of-range + conversion would point. + """ + model = _grid_sample_model( + (1, 1, 4, 4), (1, 2, 3, 2), mode="nearest", padding_mode=padding_mode + ) + grid = np.array( + [ + [np.nan, 0.0], + [np.inf, 0.5], + [-np.inf, -0.5], + [1e30, 0.0], + [-1e30, 0.0], + [np.nan, np.nan], + ], + dtype=np.float32, + ).reshape(1, 2, 3, 2) + values = np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4) + + outputs = compile_onnx(model, tmp_path).load().run({"x": values, "g": grid}) + + assert np.isin(outputs["y"], np.append(values, 0.0)).all() + + +@requires_c_compiler +def test_a_grid_sample_writing_nothing_emits_no_kernel_at_all(tmp_path): + """An empty batch leaves no position to sample, and no loop that could read a buffer.""" + model = _grid_sample_model((0, 2, 4, 5), (0, 3, 3, 2)) + + report, header = _compile(model, tmp_path) + + assert report["entrypoint"]["outputs"][0]["shape"] == [0, 2, 3, 3] + assert not _kernels(report, "gridsample") + assert f"{report['prefix']}_gridsample" not in header + + +@requires_c_compiler +def test_a_nearest_grid_sample_carries_every_value_through_unchanged(tmp_path): + """It selects rather than interpolates, which is why it serves the integer types too.""" + model = _grid_sample_model( + (1, 1, 2, 2), + (1, 2, 2, 2), + elem_type=TensorProto.INT64, + mode="nearest", + align_corners=1, + ) + values = np.array([[[[-(2**62), 2**62 - 1], [7, -7]]]], dtype=np.int64) + grid = np.array( + [[[[-1.0, -1.0], [1.0, -1.0]], [[-1.0, 1.0], [1.0, 1.0]]]], np.float32 + ) + + outputs = compile_onnx(model, tmp_path).load().run({"x": values, "g": grid}) + + np.testing.assert_array_equal(outputs["y"], values) + + +@requires_c_compiler +def test_the_regions_a_roi_align_reads_are_checked_against_the_batch(tmp_path): + """A batch index outside the operand is an argument the artifact reports, not a read.""" + model = _roi_model("RoiAlign", (2, 1, 4, 4), 1, output_height=2, output_width=2) + compiled = compile_onnx(model, tmp_path).load() + feeds = { + "x": np.arange(32, dtype=np.float32).reshape(2, 1, 4, 4), + "rois": np.array([[0.0, 0.0, 3.0, 3.0]], dtype=np.float32), + } + + inside = compiled.run({**feeds, "batch_indices": np.array([1], dtype=np.int64)}) + + assert inside["y"].shape == (1, 1, 2, 2) + for outside in (2, -1): + with pytest.raises(HarnessError, match="status 1"): + compiled.run( + {**feeds, "batch_indices": np.array([outside], dtype=np.int64)} + ) + + +@requires_c_compiler +def test_the_batch_a_max_roi_pool_reads_is_checked_against_the_operand(tmp_path): + """MaxRoiPool carries the batch in the region's first column, and checks it there.""" + model = _roi_model("MaxRoiPool", (2, 1, 4, 4), 1, pooled_shape=[2, 2]) + compiled = compile_onnx(model, tmp_path).load() + values = np.arange(32, dtype=np.float32).reshape(2, 1, 4, 4) + + inside = compiled.run( + {"x": values, "rois": np.array([[1.0, 0.0, 0.0, 3.0, 3.0]], np.float32)} + ) + + assert inside["y"].shape == (1, 1, 2, 2) + for region in ( + [2.0, 0.0, 0.0, 3.0, 3.0], + [-1.0, 0.0, 0.0, 3.0, 3.0], + [0.0, 0.0, 0.0, 1e30, 3.0], + ): + with pytest.raises(HarnessError, match="status 1"): + compiled.run({"x": values, "rois": np.array([region], np.float32)}) + + +@requires_c_compiler +def test_a_max_roi_pool_pools_the_same_regions_at_either_precision(tmp_path): + """The double kernel is the float one at another type, which is what grounds it. + + onnxruntime — the only implementation ONNX has for this op, and so the differential + sweep's oracle for it — is registered for float alone. On data every value of which both + types hold exactly, the two kernels select from the same elements, so the wider one is + pinned to the narrower one the sweep covers. + """ + values = np.arange(-24, 24, dtype=np.float64).reshape(2, 2, 4, 3) + regions = np.array( + [[0.0, 0.0, 0.0, 2.0, 3.0], [1.0, 1.0, 0.0, 3.0, 2.0]], dtype=np.float64 + ) + outputs = {} + for elem_type, dtype in ( + (TensorProto.FLOAT, np.float32), + (TensorProto.DOUBLE, np.float64), + ): + model = _roi_model( + "MaxRoiPool", + values.shape, + len(regions), + elem_type=elem_type, + pooled_shape=[2, 2], + ) + compiled = compile_onnx(model, tmp_path / dtype.__name__).load() + outputs[dtype] = compiled.run( + {"x": values.astype(dtype), "rois": regions.astype(dtype)} + )["y"] + + np.testing.assert_array_equal( + outputs[np.float64], outputs[np.float32].astype(np.float64) + ) + + +@requires_c_compiler +def test_an_affine_grid_maps_the_same_coordinates_at_either_precision(tmp_path): + """The one place onnxruntime is the oracle, because ONNX's own is blind to the type. + + The reference evaluator casts its result to float32 whatever type the transform arrives + in, so it cannot tell a double AffineGrid from a float one — the differential sweep runs + it at float32 for that reason. onnxruntime, the second oracle the compiler's parity + testing rests on, computes the grid in the type ONNX defines for it — save for the + spacing between two positions, which it rounds to float32 whatever the type. The extents + here are the ones whose spacing that rounding leaves exact, so the two are compared + element for element. + """ + runtime = pytest.importorskip("onnxruntime") + runtime.set_default_logger_severity(3) + model = _affine_grid_model( + (2, 3, 5, 5), elem_type=TensorProto.DOUBLE, align_corners=1 + ) + theta = np.array( + [[[0.7, -0.3, 0.1], [0.25, 0.9, -0.2]], [[1.0, 0.0, 0.0], [0.0, 1.0, 0.5]]], + dtype=np.float64, + ) + + outputs = compile_onnx(model, tmp_path).load().run({"theta": theta}) + + session = runtime.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + expected = session.run(None, {"theta": theta})[0] + assert outputs["y"].dtype == np.float64 + np.testing.assert_array_equal(outputs["y"], expected) + + +@requires_c_compiler +def test_a_col2im_folding_nothing_emits_no_call(tmp_path): + """An empty batch leaves no image to fold into, and no loop over one.""" + report, header = _compile(_col2im_model((0, 5, 5), (5, 5), (1, 5)), tmp_path) + + assert report["entrypoint"]["outputs"][0]["shape"] == [0, 1, 5, 5] + assert not _kernels(report, "col2im") + assert f"{report['prefix']}_col2im" not in header + + +def test_the_col2im_geometry_reaches_the_kernel_as_call_site_literals(tmp_path): + """The block, the image and the positions the blocks sat at are all literals.""" + model = _col2im_model( + (2, 12, 9), (4, 4), (2, 2), dilations=[1, 1], strides=[1, 1], pads=[0, 0, 0, 0] + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "col2im") + call = header.split(f"{kernel}(\n")[-1].splitlines() + assert [line.strip(" ,);") for line in call[:13]] == [ + "y", + "x", + "6u", + "9u", + "16u", + "4u", + "2", + "(const size_t[]){4u, 4u}", + "(const size_t[]){2u, 2u}", + "(const size_t[]){3u, 3u}", + "(const size_t[]){1u, 1u}", + "(const size_t[]){1u, 1u}", + "(const ptrdiff_t[]){0, 0}", + ] + + +@pytest.mark.parametrize( + ("op_type", "attributes", "shape", "strides"), + [ + # A block shuffle is a transpose of the operand read as blocks, so what it emits is + # the shared strided move at the strides that transpose says. + ( + "DepthToSpace", + {"blocksize": 2, "mode": "DCR"}, + "(const size_t[]){1u, 2u, 2u, 2u, 3u, 2u}", + "(const ptrdiff_t[]){48, 6, 3, 24, 1, 12}", + ), + ( + "DepthToSpace", + {"blocksize": 2, "mode": "CRD"}, + "(const size_t[]){1u, 2u, 2u, 2u, 3u, 2u}", + "(const ptrdiff_t[]){48, 24, 3, 12, 1, 6}", + ), + ( + "SpaceToDepth", + {"blocksize": 2}, + "(const size_t[]){1u, 2u, 2u, 2u, 2u, 3u}", + "(const ptrdiff_t[]){48, 6, 1, 24, 12, 2}", + ), + ], +) +def test_a_block_shuffle_is_emitted_as_the_shared_strided_move( + tmp_path, op_type, attributes, shape, strides +): + x_shape = (1, 8, 2, 3) if op_type == "DepthToSpace" else (1, 2, 4, 6) + result_strides = ( + "(const ptrdiff_t[]){48, 24, 12, 6, 2, 1}" + if op_type == "DepthToSpace" + else "(const ptrdiff_t[]){48, 24, 12, 6, 3, 1}" + ) + model = _block_model(op_type, x_shape, **attributes) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "copy") + assert not _kernels(report, op_type.lower()) + call = header.split(f"{kernel}(\n")[-1].splitlines() + assert [line.strip(" ,);") for line in call[:7]] == [ + "y", + "x", + "48u", + "6", + shape, + result_strides, + strides, + ] + + +@requires_c_compiler +@pytest.mark.parametrize("op_type", ["DepthToSpace", "SpaceToDepth"]) +def test_a_block_shuffle_of_one_element_moves_nothing(tmp_path, op_type): + """A block of one leaves every element where it is, which is one `memcpy`.""" + model = _block_model(op_type, (2, 3, 4, 5), blocksize=1) + + report, header = _compile(model, tmp_path) + + assert not _kernels(report, "copy") + assert "memcpy(y, x, 120u * sizeof(*y));" in header + + +@requires_c_compiler +def test_a_block_shuffle_writing_nothing_emits_no_move(tmp_path): + report, header = _compile( + _block_model("SpaceToDepth", (0, 2, 6, 4), blocksize=2), tmp_path + ) + + assert report["entrypoint"]["outputs"][0]["shape"] == [0, 8, 3, 2] + assert not _kernels(report, "copy") + assert "memcpy" not in header.split(f"int {report['prefix']}_run(")[-1] + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + # Everything a grid sample settles that ONNX's own inference does not: given the + # shapes, inference rejects a grid of the wrong rank or width before this does. + ( + {"elem_type": TensorProto.INT32, "mode": "linear"}, + "weights the elements around each coordinate", + ), + ( + {"elem_type": TensorProto.BOOL, "mode": "cubic"}, + "weights the elements around each coordinate", + ), + ({"mode": "bilinear"}, "is not one of the values ONNX defines for it"), + ({"padding_mode": "wrap"}, "is not one of the values ONNX defines for it"), + ], +) +def test_a_grid_sample_the_compiler_cannot_serve_is_rejected(tmp_path, kwargs, message): + x_shape = kwargs.pop("x_shape", (1, 2, 4, 5)) + grid_shape = kwargs.pop("grid_shape", (1, 3, 3, 2)) + + with pytest.raises(CompileError) as error: + compile_onnx(_grid_sample_model(x_shape, grid_shape, **kwargs), tmp_path) + + assert "`sample`" in str(error.value) + assert message in str(error.value) + + +@pytest.mark.parametrize( + ("op_type", "kwargs", "message"), + [ + ("RoiAlign", {"columns": 5}, "one row of 4 value(s) per region"), + ("MaxRoiPool", {"columns": 4}, "one row of 5 value(s) per region"), + ("RoiAlign", {"x_shape": (1, 2, 0, 4)}, "holds no elements to sample"), + ("MaxRoiPool", {"x_shape": (1, 2, 4, 0)}, "holds no elements to sample"), + ( + "RoiAlign", + {"mode": "median"}, + "is not one of the values ONNX defines for it", + ), + ( + "RoiAlign", + {"coordinate_transformation_mode": "asymmetric"}, + "is not one of the values ONNX defines for it", + ), + ], +) +def test_a_region_pooling_the_compiler_cannot_serve_is_rejected( + tmp_path, op_type, kwargs, message +): + defaults = ( + {"output_height": 2, "output_width": 2} + if op_type == "RoiAlign" + else {"pooled_shape": [2, 2]} + ) + x_shape = kwargs.pop("x_shape", (2, 3, 4, 5)) + + with pytest.raises(CompileError) as error: + compile_onnx( + _roi_model(op_type, x_shape, 2, **{**defaults, **kwargs}), tmp_path + ) + + assert "`roi`" in str(error.value) + assert message in str(error.value) + + +def test_a_max_roi_pool_without_a_pooled_shape_is_rejected(tmp_path): + """ONNX defines the attribute as required, so its absence is not a default.""" + model = _model( + [helper.make_node("MaxRoiPool", ["x", "rois"], ["y"], name="roi")], + [ + _tensor("x", TensorProto.FLOAT, (1, 2, 4, 5)), + _tensor("rois", TensorProto.FLOAT, (2, 5)), + ], + [_tensor("y", TensorProto.FLOAT, (2, 2, 2, 2))], + ) + + with pytest.raises(CompileError, match=re.escape("states no `pooled_shape`")): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"theta_shape": (2, 2, 2)}, "ONNX defines as a transform of shape [2, 2, 3]"), + ({"theta_shape": (1, 2, 3)}, "ONNX defines as a transform of shape [2, 2, 3]"), + ], +) +def test_an_affine_grid_the_compiler_cannot_serve_is_rejected( + tmp_path, kwargs, message +): + size = kwargs.pop("size", (2, 3, 4, 5)) + # A theta the size disagrees with leaves ONNX's own inference nothing to derive, so the + # result's shape is declared rather than inferred. + model = _affine_grid_model(size, **kwargs) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "`grid`" in str(error.value) + assert message in str(error.value) + + +@pytest.mark.parametrize( + ("op_type", "runtime", "operand"), + [ + ("AffineGrid", ("size",), "size"), + ("Col2Im", ("image_shape",), "image_shape"), + ("Col2Im", ("block_shape",), "block_shape"), + ], +) +def test_an_extent_read_at_run_time_is_rejected_as_a_dynamic_shape( + tmp_path, op_type, runtime, operand +): + """The operands that decide these results' shapes have to be fixed by the graph. + + Which is what the corpus's own tests of both ops do not do — every one of them feeds the + extents — so the frontend names the operand rather than letting a kernel reach for values + that are not there. + """ + model = ( + _affine_grid_model((2, 3, 4, 5), runtime=True) + if op_type == "AffineGrid" + else _col2im_model( + (1, 5, 5), (5, 5), (1, 5), runtime=runtime, output=(1, 1, 5, 5) + ) + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert f"`{operand}`" in str(error.value) + assert "depends on input data" in str(error.value) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"block": (2, 2)}, "which does not divide the 5 row(s)"), + ({"x_shape": (1, 5, 7)}, "block(s) over an image of [5, 5]"), + ({"attributes": {"strides": [2, 2]}}, "block(s) over an image of [5, 5]"), + ({"image": (5, 0)}, "ONNX defines them as extents, which are positive"), + ({"block": (1, 5, 5)}, "for a 2-dimensional image"), + ({"x_shape": (1, 1, 5, 5)}, "a tensor of rank 3"), + ( + {"elem_type": TensorProto.BOOL}, + "summing truth values has no defined result", + ), + ], +) +def test_a_col2im_the_compiler_cannot_serve_is_rejected(tmp_path, kwargs, message): + shape = kwargs.pop("x_shape", (1, 5, 5)) + image = kwargs.pop("image", (5, 5)) + block = kwargs.pop("block", (1, 5)) + elem_type = kwargs.pop("elem_type", TensorProto.FLOAT) + # ONNX's own inference rejects most of these first, given the extents; the declared + # result shape is what leaves them to the compiler at all. + model = _col2im_model( + shape, + image, + block, + elem_type=elem_type, + output=(shape[0], 1, 5, 5), + **kwargs.pop("attributes", {}), + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "`fold`" in str(error.value) + assert message in str(error.value) + + +@pytest.mark.parametrize( + ("op_type", "shape", "attributes", "message"), + [ + # What a block shuffle settles that ONNX's own inference does not: given the shape, + # inference rejects a non-positive blocksize and a rank other than 4 before this does. + ("DepthToSpace", (1, 7, 2, 3), {"blocksize": 2}, "does not divide them evenly"), + ("SpaceToDepth", (1, 2, 5, 4), {"blocksize": 2}, "do not tile it evenly"), + ( + "DepthToSpace", + (1, 8, 2, 3), + {"blocksize": 2, "mode": "RCD"}, + "not one of the modes ONNX defines", + ), + ], +) +def test_a_block_shuffle_the_compiler_cannot_serve_is_rejected( + tmp_path, op_type, shape, attributes, message +): + model = _block_model(op_type, shape, **attributes) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "`block`" in str(error.value) + assert message in str(error.value) + + +def test_a_block_shuffle_without_a_blocksize_is_rejected(tmp_path): + """ONNX defines the attribute as required, so its absence is not a default.""" + model = _model( + [helper.make_node("DepthToSpace", ["x"], ["y"], name="block")], + [_tensor("x", TensorProto.FLOAT, (1, 8, 2, 3))], + [_tensor("y", TensorProto.FLOAT, (1, 2, 4, 6))], + opset=_BLOCK_OPSET, + ) + + with pytest.raises(CompileError, match=re.escape("states no `blocksize`")): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("op_type", "opset", "nearest"), + [ + ("GridSample", 16, 22), + ("GridSample", 20, 22), + ("RoiAlign", 10, 22), + ("RoiAlign", 16, 22), + ("MaxRoiPool", 21, 22), + ("DepthToSpace", 11, 13), + ("SpaceToDepth", 1, 13), + ], +) +def test_a_sampler_below_its_supported_revision_names_the_nearest_version( + tmp_path, op_type, opset, nearest +): + """None of the older revisions has an oracle — the reference evaluator applies today's + implementation to them and the backend corpus tests none of them — so no kernel claims + them, and dispatch says so rather than serving the current walk under an older name.""" + models = { + "GridSample": lambda: _grid_sample_model( + (1, 2, 4, 5), (1, 3, 3, 2), name="node" + ), + "RoiAlign": lambda: _roi_model( + "RoiAlign", (1, 2, 4, 5), 2, name="node", output_height=2, output_width=2 + ), + "MaxRoiPool": lambda: _roi_model( + "MaxRoiPool", (1, 2, 4, 5), 2, name="node", pooled_shape=[2, 2] + ), + "DepthToSpace": lambda: _block_model( + "DepthToSpace", (1, 8, 2, 3), name="node", blocksize=2 + ), + "SpaceToDepth": lambda: _block_model( + "SpaceToDepth", (1, 2, 6, 4), name="node", blocksize=2 + ), + } + model = models[op_type]() + del model.opset_import[:] + model.opset_import.append(helper.make_opsetid("", opset)) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`node`" in message + assert f"opset version {opset}" in message + assert f"Nearest supported version: {nearest}" in message + + +# -------------------------------------------------------------------------------------- +# The scatters +# -------------------------------------------------------------------------------------- + +# What a scatter computes is settled by the conformance and differential suites. What is +# asserted here is what those cannot reach: that the result is the operand copied and then +# written into, that the geometry and the fold are compile-time literals rather than kernels +# of their own, that an index no element sits at is reported instead of written past, and the +# errors for the shapes and attributes the compiler refuses outright. + +_SCATTER_OPSET = 18 +_TENSOR_SCATTER_OPSET = 24 + + +def _scatter_model( + op_type, + data_shape, + indices_shape, + *, + updates_shape=None, + elem_type=TensorProto.FLOAT, + index_type=TensorProto.INT64, + name="scatter", + opset=_SCATTER_OPSET, + **attributes, +): + """One scatter; ScatterElements takes one update per index, ScatterND one per tuple.""" + return _model( + [ + helper.make_node( + op_type, + ["data", "indices", "updates"], + ["y"], + name=name, + **attributes, + ) + ], + [ + _tensor("data", elem_type, data_shape), + _tensor("indices", index_type, indices_shape), + _tensor( + "updates", + elem_type, + indices_shape if updates_shape is None else updates_shape, + ), + ], + [helper.make_empty_tensor_value_info("y")], + opset=opset, + ) + + +def _tensor_scatter_model( + cache_shape, + update_shape, + *, + elem_type=TensorProto.FLOAT, + indices_shape=(), + name="cache", + **attributes, +): + """One TensorScatter; `indices_shape` of None leaves the write indices out.""" + names = ["past", "update"] + inputs = [ + _tensor("past", elem_type, cache_shape), + _tensor("update", elem_type, update_shape), + ] + if indices_shape is not None: + names.append("written_at") + inputs.append( + _tensor( + "written_at", + TensorProto.INT64, + cache_shape[:1] if indices_shape == () else indices_shape, + ) + ) + return _model( + [helper.make_node("TensorScatter", names, ["y"], name=name, **attributes)], + inputs, + [helper.make_empty_tensor_value_info("y")], + opset=_TENSOR_SCATTER_OPSET, + ) + + +def _reference(model, feeds): + """What ONNX's own evaluator computes for the model, as the oracle for a kernel test.""" + return ReferenceEvaluator(model).run(None, feeds) + + +def test_scatters_of_one_fold_and_type_share_a_kernel(tmp_path): + """A kernel name encodes the fold and the types, and nothing else it does not depend on.""" + model = _model( + [ + helper.make_node("ScatterElements", ["data", "i", "u"], ["h"], name="one"), + helper.make_node("ScatterElements", ["h", "i", "u"], ["g"], name="two"), + helper.make_node( + "ScatterElements", + ["g", "i", "u"], + ["y"], + name="folded", + reduction="add", + ), + ], + [ + _tensor("data", TensorProto.FLOAT, (3, 4)), + _tensor("i", TensorProto.INT64, (2, 4)), + _tensor("u", TensorProto.FLOAT, (2, 4)), + ], + [helper.make_empty_tensor_value_info("y")], + opset=_SCATTER_OPSET, + ) + + report, header = _compile(model, tmp_path) + + plain, folded = sorted(_kernels(report, "scatterelements")) + assert plain.endswith("_add_float_int64_t") and folded.endswith( + "_none_float_int64_t" + ) + assert header.count(f"static int {folded}(") == 1 + assert header.count(f"{folded}(\n") == 3 + assert "out[offset] = updates[index];" in header + assert "out[offset] = out[offset] + updates[index];" in header + + +def test_the_scatter_geometry_reaches_the_kernel_as_call_site_literals(tmp_path): + """The result is the operand copied, and then written into through its own strides. + + The updates are walked by their own shape and addressed by the operand's strides, which + is what lets an index tensor cover only part of the axes it does not write along. + """ + model = _scatter_model("ScatterElements", (3, 4), (2, 3), axis=1) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "scatterelements") + assert "memcpy(y, data, 12u * sizeof(*y));" in header + assert ( + f"if ({kernel}(\n y,\n updates,\n indices,\n" + " 6u,\n" + " 2,\n" + " (const size_t[]){2u, 3u},\n" + " (const size_t[]){4u, 1u},\n" + " 1,\n" + " 4u) != 0) {" in header + ) + + +def test_a_scatter_nd_writes_a_slice_per_index_tuple(tmp_path): + """The tuple's depth decides how much of the operand one update replaces.""" + model = _scatter_model("ScatterND", (4, 2, 3), (2, 1), updates_shape=(2, 2, 3)) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "scatternd") + assert ( + f"if ({kernel}(\n y,\n updates,\n indices,\n" + " 2u,\n" + " 1u,\n" + " 6u,\n" + " (const size_t[]){4u},\n" + " (const size_t[]){6u}) != 0) {" in header + ) + + +def test_a_scatter_with_no_updates_is_the_copy_alone(tmp_path): + """Nothing is written, and no loop is emitted that could read an empty buffer.""" + model = _scatter_model("ScatterElements", (3, 4), (0, 4)) + + report, header = _compile(model, tmp_path) + + assert report["entrypoint"]["outputs"][0]["shape"] == [3, 4] + assert not _kernels(report, "scatterelements") + assert "memcpy(y, data, 12u * sizeof(*y));" in header + + +def test_the_two_families_fold_an_extremum_as_their_own_references_do(tmp_path): + """`max` means one thing in ScatterElements and another in ScatterND — on a NaN. + + ONNX documents both as folding with `max`, and its two reference implementations disagree + about the case neither document mentions: ScatterElements folds with Python's own `max`, + which keeps the element already in the result when a comparison against a NaN comes out + false, while ScatterND folds with `np.maximum`, which propagates a NaN from either side. + Each kernel follows the reference of its own op, which is what this pins. + """ + headers = { + op_type: _compile( + _scatter_model( + op_type, + (3, 4), + (2, 4) if op_type == "ScatterElements" else (2, 1), + updates_shape=None if op_type == "ScatterElements" else (2, 4), + reduction="max", + ), + tmp_path / op_type, + )[1] + for op_type in ("ScatterElements", "ScatterND") + } + + assert ( + "out[offset] = (updates[index] > out[offset]) ? updates[index] : out[offset];" + in headers["ScatterElements"] + ) + assert "isnan" not in headers["ScatterElements"] + assert "_maximum_float(out[offset + element]," in headers["ScatterND"] + assert ( + "return (left > right || isnan(left)) ? left : right;" in headers["ScatterND"] + ) + + +@requires_c_compiler +@pytest.mark.parametrize( + ("op_type", "indices_shape", "updates_shape", "outside"), + [ + ("ScatterElements", (2, 4), None, (4, -5)), + ("Scatter", (2, 4), None, (4, -5)), + ("ScatterND", (2, 1), (2, 4), (3, -4)), + ], +) +def test_an_index_no_element_sits_at_is_reported( + tmp_path, op_type, indices_shape, updates_shape, outside +): + """An index past either end of the axis is an argument the artifact reports.""" + model = _scatter_model( + op_type, (3, 4), indices_shape, updates_shape=updates_shape, axis=0 + ) + compiled = compile_onnx(model, tmp_path).load() + feeds = { + "data": np.zeros((3, 4), np.float32), + "updates": np.ones( + indices_shape if updates_shape is None else updates_shape, np.float32 + ), + } + + inside = compiled.run( + {**feeds, "indices": np.full(indices_shape, -1, dtype=np.int64)} + ) + + np.testing.assert_array_equal(inside["y"][-1], np.ones(4, np.float32)) + for index in outside: + with pytest.raises(HarnessError, match="status 1"): + compiled.run( + {**feeds, "indices": np.full(indices_shape, index, dtype=np.int64)} + ) + + +@requires_c_compiler +def test_a_deprecated_scatter_computes_what_its_successor_does(tmp_path): + """Scatter dispatches at the revisions before the deprecating one too. + + Only the deprecating revision has an oracle in the differential sweep — the corpus's own + Scatter tests import opset 10 — so what this adds is that every revision the kernel claims + reaches it, against the ScatterElements ONNX's own document says computes the same thing. + """ + feeds = { + "data": np.arange(12, dtype=np.float32).reshape(3, 4), + "indices": np.array([[0, 2], [1, 0]], np.int64), + "updates": np.array([[10.0, 20.0], [30.0, 40.0]], np.float32), + } + successor = _scatter_model("ScatterElements", (3, 4), (2, 2), axis=1) + (expected,) = _reference(successor, feeds) + + for opset in (9, 10, 11): + model = _scatter_model("Scatter", (3, 4), (2, 2), axis=1, opset=opset) + outputs = compile_onnx(model, tmp_path / str(opset)).load().run(feeds) + + np.testing.assert_array_equal(outputs["y"], expected) + + +def test_a_tensor_scatter_without_write_indices_reads_none(tmp_path): + """The operand is optional, and a kernel that took it anyway would read a buffer that + is not there; leaving it out is a kernel of its own, writing from the start of the axis.""" + report, header = _compile( + _tensor_scatter_model((2, 1, 4, 5), (2, 1, 2, 5), indices_shape=None), tmp_path + ) + + (kernel,) = _kernels(report, "tensorscatter") + assert kernel.endswith("_linear_appended_float") + assert "write_indices" not in header + assert "const ptrdiff_t written_at = 0;" in header + + +@requires_c_compiler +def test_a_linear_write_past_the_end_of_the_cache_is_reported(tmp_path): + """ONNX's own reference answers a write running off the axis with an exception; the + artifact answers it with the argument error the status enum exists for, having read + nothing outside the buffer.""" + model = _tensor_scatter_model((2, 1, 4, 5), (2, 1, 2, 5)) + compiled = compile_onnx(model, tmp_path).load() + feeds = { + "past": np.zeros((2, 1, 4, 5), np.float32), + "update": np.ones((2, 1, 2, 5), np.float32), + } + + inside = compiled.run({**feeds, "written_at": np.array([0, 2], np.int64)}) + + assert inside["y"][1, 0, 3].tolist() == [1.0] * 5 + for written_at in ((0, 3), (-1, 0), (4, 0)): + with pytest.raises(HarnessError, match="status 1"): + compiled.run({**feeds, "written_at": np.array(written_at, np.int64)}) + + +@requires_c_compiler +def test_a_circular_write_wraps_where_a_linear_one_is_refused(tmp_path): + """The mode is what decides whether running off the end is an error or a wrap.""" + written_at = np.array([3, 0], np.int64) + feeds = { + "past": np.zeros((2, 1, 4, 5), np.float32), + "update": np.arange(20, dtype=np.float32).reshape(2, 1, 2, 5), + "written_at": written_at, + } + model = _tensor_scatter_model((2, 1, 4, 5), (2, 1, 2, 5), mode="circular") + (expected,) = _reference(model, feeds) + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + np.testing.assert_array_equal(outputs["y"], expected) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ( + {"op_type": "ScatterElements", "indices_shape": (2, 4, 1)}, + "the same rank", + ), + ( + { + "op_type": "ScatterElements", + "indices_shape": (2, 4), + "updates_shape": (2, 3), + }, + "one update per index", + ), + ( + {"op_type": "ScatterElements", "indices_shape": (4, 4), "axis": 1}, + "reaches past it on axis 0", + ), + ( + { + "op_type": "ScatterElements", + "indices_shape": (2, 4), + "reduction": "mean", + }, + "not one of the reductions ONNX defines", + ), + ( + {"op_type": "Scatter", "indices_shape": (2, 4), "reduction": "add"}, + "not one of the reductions ONNX defines", + ), + ( + {"op_type": "ScatterND", "indices_shape": (2, 3), "updates_shape": (2,)}, + "addresses 3 dimension(s)", + ), + ( + {"op_type": "ScatterND", "indices_shape": (2, 1), "updates_shape": (2, 3)}, + "writes one slice per index tuple", + ), + ], +) +def test_a_scatter_the_compiler_cannot_serve_is_rejected(tmp_path, kwargs, message): + op_type = kwargs.pop("op_type") + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(_scatter_model(op_type, (3, 4), **kwargs), tmp_path) + + +def test_a_tensor_scatter_mode_onnx_does_not_define_is_rejected(tmp_path): + """The one thing about a TensorScatter that ONNX's own shape inference does not check.""" + model = _tensor_scatter_model((2, 1, 4, 5), (2, 1, 2, 5), mode="rolling") + + with pytest.raises(CompileError, match="not one of the modes ONNX defines"): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + "kwargs", + [ + # The sequence axis is never the batch the write indices are read by — ONNX's own + # inference calls axis 0 out of range for any rank, which a rank-2 cache reaches + # through the default axis alone. + {"axis": 0}, + {"cache_shape": (3, 4), "update_shape": (3, 2)}, + # An update longer than the cache along the sequence axis, one differing on another + # axis, one of another rank, and a write index per something other than the batch. + {"update_shape": (2, 1, 5, 5)}, + {"update_shape": (2, 1, 2, 4)}, + {"update_shape": (2, 1, 2)}, + {"indices_shape": (3,)}, + ], +) +def test_a_tensor_scatter_of_disagreeing_shapes_is_rejected(tmp_path, kwargs): + """Every one of these is refused before a kernel is asked for it. + + ONNX's own shape inference relates a TensorScatter's operands, so a model that puts them + at odds has no inferred result to compile against; the compiler reports that against the + node rather than emitting a kernel whose addressing would run off a buffer. The kernel + generator checks the same relations again, which is where a shape that reached it another + way would stop. + """ + model = _tensor_scatter_model( + kwargs.pop("cache_shape", (2, 1, 4, 5)), + kwargs.pop("update_shape", (2, 1, 2, 5)), + **kwargs, + ) + + with pytest.raises(CompileError, match="cache"): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("op_type", "opset"), + [("ScatterElements", 11), ("ScatterElements", 16), ("ScatterND", 13)], +) +def test_a_scatter_below_its_supported_revision_names_the_nearest_version( + tmp_path, op_type, opset +): + """`reduction` arrived at 16 and grew at 18, so the older revisions are a different op. + + The reference evaluator applies today's implementation to all of them and the backend + corpus tests none of them, so nothing can vouch for one; dispatch says so rather than + serving the current fold under an older revision's name. + """ + model = _scatter_model( + op_type, + (3, 4), + (2, 4) if op_type == "ScatterElements" else (2, 1), + updates_shape=None if op_type == "ScatterElements" else (2, 4), + name="node", + opset=opset, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`node`" in message + assert f"opset version {opset}" in message + assert "Nearest supported version: 18" in message + + +# -------------------------------------------------------------------------------------- +# The einsum contraction +# -------------------------------------------------------------------------------------- + +# What an equation computes is settled by the conformance and differential suites, against +# ONNX's corpus and reference evaluator. What is asserted here is what those cannot reach: +# that the equation is read at compile time into extents and strides — so that a diagonal is +# addressing rather than code, and every equation of one arity and element type is one shared +# kernel — and the errors for the equations the compiler refuses outright. + +_EINSUM_OPSET = 12 + + +def _einsum_model( + equation, + shapes, + *, + names=None, + elem_type=TensorProto.FLOAT, + result_shape=None, + name="node", +): + """One Einsum; `result_shape` declares the result rather than leaving it to inference.""" + names = list(names or [f"in{index}" for index in range(len(shapes))]) + result = ( + helper.make_empty_tensor_value_info("y") + if result_shape is None + else _tensor("y", elem_type, result_shape) + ) + return _model( + [helper.make_node("Einsum", names, ["y"], name=name, equation=equation)], + [_tensor(operand, elem_type, shape) for operand, shape in zip(names, shapes)], + [result], + opset=_EINSUM_OPSET, + ) + + +def test_equations_of_one_arity_and_element_type_share_a_kernel(tmp_path): + """The equation is addressing, so only the operand count and type reach the loop.""" + model = _model( + [ + helper.make_node( + "Einsum", ["a", "b"], ["p"], name="product", equation="ij,jk->ik" + ), + helper.make_node( + "Einsum", ["a", "a"], ["q"], name="hadamard", equation="ij,ij->ij" + ), + helper.make_node("Einsum", ["a"], ["r"], name="total", equation="ij->"), + ], + [ + _tensor("a", TensorProto.FLOAT, (2, 3)), + _tensor("b", TensorProto.FLOAT, (3, 4)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q", "r")], + opset=_EINSUM_OPSET, + ) + + report, header = _compile(model, tmp_path) + + unary, binary = sorted(_kernels(report, "einsum")) + assert unary.endswith("_1_float") and binary.endswith("_2_float") + assert header.count(f"static void {binary}(") == 1 + assert header.count(f"{binary}(\n") == 3 + + +def test_the_equation_reaches_the_kernel_as_call_site_literals(tmp_path): + """A label the result keeps is a stride per operand; one it drops is the summed loop.""" + model = _einsum_model("ij,jk->ik", [(2, 3), (3, 4)], names=["a", "b"]) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "einsum") + assert ( + f"{kernel}(\n y,\n a,\n b,\n" + " 8u,\n" + " 2,\n" + " (const size_t[]){2u, 4u},\n" + " (const size_t[]){3u, 0u},\n" + " (const size_t[]){0u, 1u},\n" + " 3u,\n" + " 1,\n" + " (const size_t[]){3u},\n" + " (const size_t[]){1u},\n" + " (const size_t[]){4u});" in header + ) + + +def test_a_label_repeated_in_a_term_is_a_stride_down_the_diagonal(tmp_path): + """The strides of the axes one label names add up: 3 + 1 walks a 3x3 diagonally.""" + model = _einsum_model("...ii->...i", [(2, 3, 3)], names=["a"]) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "einsum") + assert ( + f"{kernel}(\n y,\n a,\n" + " 6u,\n" + " 2,\n" + " (const size_t[]){2u, 3u},\n" + " (const size_t[]){9u, 4u},\n" + " 1u,\n" + " 0,\n" + " (const size_t[]){0u},\n" + " (const size_t[]){0u});" in header + ) + + +def test_an_operand_stretched_along_a_label_reads_it_at_a_zero_stride(tmp_path): + """numpy stretches an extent of 1 against another operand's, which is a stride of 0.""" + model = _einsum_model("ij,ij->j", [(1, 3), (2, 3)], names=["a", "b"]) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "einsum") + assert ( + f"{kernel}(\n y,\n a,\n b,\n" + " 3u,\n" + " 1,\n" + " (const size_t[]){3u},\n" + " (const size_t[]){1u},\n" + " (const size_t[]){1u},\n" + " 2u,\n" + " 1,\n" + " (const size_t[]){2u},\n" + " (const size_t[]){0u},\n" + " (const size_t[]){3u});" in header + ) + + +def test_an_equation_with_no_result_to_write_emits_no_loop(tmp_path): + """Nothing to write, so no kernel is emitted that could read an empty buffer.""" + report, _ = _compile(_einsum_model("ij->ji", [(0, 3)]), tmp_path) + + assert report["entrypoint"]["outputs"][0]["shape"] == [3, 0] + assert not _kernels(report, "einsum") + + +@requires_c_compiler +def test_a_contraction_over_an_empty_axis_sums_over_nothing(tmp_path): + """The result is written, from a sum of no terms at all, reading neither operand.""" + model = _einsum_model("ij,jk->ik", [(2, 0), (0, 3)]) + feeds = { + "in0": np.zeros((2, 0), dtype=np.float32), + "in1": np.zeros((0, 3), dtype=np.float32), + } + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + np.testing.assert_array_equal(outputs["y"], _reference(model, feeds)[0]) + + +@pytest.mark.parametrize( + ("label", "equation", "shapes", "result_shape", "expected"), + [ + ("empty", "", ((),), (), "`equation` is empty"), + ("not_a_label", "i1,j->ij", ((2,), (3,)), None, "is not a term"), + ("two_ellipses", "...i...->i", ((2, 3),), None, "is not a term"), + ( + "two_outputs", + "ij->i->j", + ((2, 3),), + None, + "states its output more than once", + ), + ("uneven_diagonal", "ii->i", ((2, 3),), None, "only over axes of equal extent"), + ( + "disagreeing_label", + "ij,jk->ik", + ((2, 3), (4, 5)), + None, + "measures 3 on one operand", + ), + ("repeated_output", "i->ii", ((3,),), None, "more than once; each axis"), + ("unknown_output", "i->ij", ((3,),), None, "which no operand's term carries"), + ( + "term_per_operand", + "ij,jk->ik", + ((2, 3),), + (2, 4), + "states 2 term(s) for 1 operand(s)", + ), + ( + "label_per_axis", + "ijk->ij", + ((2, 3),), + (2, 3), + "names 3 label(s) for `in0`", + ), + ], +) +def test_an_equation_onnx_does_not_define_is_rejected( + tmp_path, label, equation, shapes, result_shape, expected +): + """Every reading ONNX's Einsum leaves undefined is refused by name, never guessed at. + + The empty equation is one numpy would take for a scalar term and ONNX's own reference + implementation rejects outright; it is refused here for that reason. Declaring a result + is what reaches these readings at all: left underived, a node ONNX cannot infer a shape + for is stopped one step earlier, by shape inference rather than by the equation. + """ + model = _einsum_model(equation, shapes, result_shape=result_shape) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`node`" in message + assert expected in message + + +def test_an_equation_wider_than_the_shape_onnx_inferred_is_rejected(tmp_path): + """numpy stretches a labelled axis onto the result and ONNX's shape inference does not. + + The two disagree about this equation's result — numpy computes [2, 3] and the buffer ONNX + sized holds [1, 3] — so the node is refused rather than written past. + """ + model = _einsum_model("ij,ij->ij", [(1, 3), (2, 3)]) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`node`" in message + assert "[2, 3]" in message and "[1, 3]" in message + + +def test_an_einsum_below_its_supported_revision_names_the_nearest_version(tmp_path): + """Einsum arrived at opset 12; nothing defines it below that.""" + model = _model( + [helper.make_node("Einsum", ["x"], ["y"], name="node", equation="ij->ji")], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [_tensor("y", TensorProto.FLOAT, (3, 2))], + opset=11, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`Einsum`" in message + assert "opset version 11" in message + assert "Nearest supported version: 12" in message + + +# -------------------------------------------------------------------------------------- +# Opset-dependent dispatch +# -------------------------------------------------------------------------------------- + + +def test_softmax_below_its_supported_revision_names_the_nearest_version(tmp_path): + """Up to opset 12 Softmax flattened the axes from `axis` on, which nothing can vouch for. + + The reference evaluator applies the current semantics to those revisions and the backend + corpus has no test at one, so no kernel claims them — and dispatch says so rather than + serving the current formula under the old op's name. + """ + model = _model( + [helper.make_node("Softmax", ["x"], ["y"], name="soft", axis=1)], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + opset=12, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`soft`" in message + assert "opset version 12" in message + assert "Nearest supported version: 13" in message + + +@pytest.mark.parametrize( + ("op_type", "shapes", "opset", "nearest"), + [ + ("MatMul", ((2, 3), (3, 4)), 12, 13), + ("Det", ((2, 2),), 21, 22), + ("Conv", ((1, 1, 5, 5), (1, 1, 3, 3)), 21, 22), + ("ConvTranspose", ((1, 1, 5, 5), (1, 1, 3, 3)), 21, 22), + ], +) +def test_a_matrix_op_below_its_supported_revision_names_the_nearest_version( + tmp_path, op_type, shapes, opset, nearest +): + """None of these ops changed semantics at its claimed revision — but nothing vouches for + the older ones either: the reference evaluator applies today's implementation to them and + the backend corpus has no test at one, so no kernel claims them. + """ + names = ["x", "z"][: len(shapes)] + model = _model( + [helper.make_node(op_type, names, ["y"], name="node")], + [_tensor(name, TensorProto.FLOAT, shape) for name, shape in zip(names, shapes)], + [helper.make_empty_tensor_value_info("y")], + opset=opset, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`node`" in message + assert f"opset version {opset}" in message + assert f"Nearest supported version: {nearest}" in message + + +@requires_c_compiler +def test_each_opset_gets_the_semantics_of_its_own_revision(tmp_path): + """`Clip` reads its bounds from attributes up to opset 10 and from inputs after it. + + The same node is two different ops across that boundary: at 6 the attributes bound the + result, at 13 they are not attributes at all and an unbounded `Clip` is the identity. + Each model is compared against the reference evaluator reading that same model. + """ + values = np.arange(-3, 3, dtype=np.float32).reshape(2, 3) + results = {} + for opset, attributes in ((6, {"min": -1.0, "max": 1.0}), (13, {})): + model = _model( + [helper.make_node("Clip", ["x"], ["y"], name="clip", **attributes)], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + opset=opset, + ) + + results[opset] = ( + compile_onnx(model, tmp_path / str(opset)).load().run({"x": values})["y"] + ) + + expected = ReferenceEvaluator(model).run(None, {"x": values})[0] + np.testing.assert_array_equal(results[opset], expected) + assert not np.array_equal(results[6], results[13]) + + +# -------------------------------------------------------------------------------------- +# What the kernels refuse +# -------------------------------------------------------------------------------------- + + +def test_an_undefined_gelu_approximation_is_rejected(tmp_path): + """`approximate` selects a formula, so an unknown one has no code to emit.""" + model = _model( + [helper.make_node("Gelu", ["x"], ["y"], name="gelu", approximate="sigmoid")], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`gelu`" in message + assert "`sigmoid`" in message + assert "`tanh`" in message + + +def test_a_bitcast_to_bool_is_rejected(tmp_path): + """A boolean tensor is emitted as bytes holding 0 or 1; arbitrary bits are neither.""" + model = _model( + [helper.make_node("BitCast", ["x"], ["y"], name="bits", to=TensorProto.BOOL)], + [_tensor("x", TensorProto.INT8, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + opset=26, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`bits`" in message + assert "BOOL" in message + assert "INT8" in message + + +@pytest.mark.parametrize("direction", ["SIDEWAYS", None]) +def test_a_bit_shift_in_no_defined_direction_is_rejected(tmp_path, direction): + """`direction` selects the operator, so an unknown one has no code to emit.""" + attributes = {} if direction is None else {"direction": direction} + model = _model( + [helper.make_node("BitShift", ["a", "b"], ["y"], name="shift", **attributes)], + [ + _tensor("a", TensorProto.UINT8, (2, 3)), + _tensor("b", TensorProto.UINT8, (2, 3)), + ], + [helper.make_empty_tensor_value_info("y")], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`shift`" in message + assert "`LEFT`" in message and "`RIGHT`" in message + + +def test_mod_on_floats_without_fmod_is_rejected(tmp_path): + """ONNX requires `fmod=1` for the floating-point families; the compiler says so.""" + model = _model( + [helper.make_node("Mod", ["a", "b"], ["y"], name="remainder")], + [ + _tensor("a", TensorProto.FLOAT, (2, 3)), + _tensor("b", TensorProto.FLOAT, (2, 3)), + ], + [helper.make_empty_tensor_value_info("y")], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`remainder`" in message + assert "FLOAT" in message + assert "fmod" in message + + +@pytest.mark.parametrize("training_mode", [None, True]) +def test_dropout_that_is_not_provably_inference_is_rejected(tmp_path, training_mode): + """Training mode samples a mask, which no static artifact can reproduce.""" + initializer = ( + [] + if training_mode is None + else [onnx.numpy_helper.from_array(np.array(True), "mode")] + ) + inputs = [_tensor("x", TensorProto.FLOAT, (2, 3))] + if training_mode is None: + inputs.append(_tensor("mode", TensorProto.BOOL, ())) + model = _model( + [helper.make_node("Dropout", ["x", "", "mode"], ["y"], name="drop")], + inputs, + [helper.make_empty_tensor_value_info("y")], + initializer=initializer, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`drop`" in message + assert "`mode`" in message + assert "inference mode" in message + + +@requires_c_compiler +def test_dropout_pinned_to_inference_passes_its_input_and_mask_through(tmp_path): + """A `training_mode` the graph fixes to false is compiled as the identity it is.""" + model = _model( + [helper.make_node("Dropout", ["x", "", "mode"], ["y", "mask"], name="drop")], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [ + helper.make_empty_tensor_value_info("y"), + helper.make_empty_tensor_value_info("mask"), + ], + initializer=[onnx.numpy_helper.from_array(np.array(False), "mode")], + ) + values = np.arange(-3, 3, dtype=np.float32).reshape(2, 3) + + outputs = compile_onnx(model, tmp_path).load().run({"x": values}) + + expected = ReferenceEvaluator(model).run(None, {"x": values}) + np.testing.assert_array_equal(outputs["y"], expected[0]) + np.testing.assert_array_equal(outputs["mask"], expected[1]) + + +def test_a_perm_that_is_not_a_permutation_is_rejected(tmp_path): + """`perm` says where each axis comes from, so it has to name each of them once.""" + model = _model( + [helper.make_node("Transpose", ["x"], ["y"], name="swap", perm=[0, 0])], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [_tensor("y", TensorProto.FLOAT, (2, 2))], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`swap`" in message + assert "[0, 0]" in message + + +@pytest.mark.parametrize( + ("label", "starts", "ends", "axes", "steps", "expected"), + [ + ("zero_step", [0], [2], [0], [0], "steps by 0"), + ("repeated_axis", [0, 0], [2, 2], [1, 1], [1, 1], "more than once"), + ("mismatched_bounds", [0, 0], [2], [0], [1], "one of each"), + ], +) +def test_slice_bounds_onnx_does_not_define_are_rejected( + tmp_path, label, starts, ends, axes, steps, expected +): + model = _model( + [helper.make_node("Slice", ["x", "s", "e", "a", "t"], ["y"], name="cut")], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [_tensor("y", TensorProto.FLOAT, (2, 3))], + initializer=[ + _int64("s", starts), + _int64("e", ends), + _int64("a", axes), + _int64("t", steps), + ], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`cut`" in message + assert expected in message + + +def test_a_repeat_count_per_axis_is_required(tmp_path): + """ONNX defines `repeats` as one count per axis; anything else addresses nothing.""" + model = _model( + [helper.make_node("Tile", ["x", "r"], ["y"], name="repeat")], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [_tensor("y", TensorProto.FLOAT, (4, 3))], + initializer=[_int64("r", [2])], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`repeat`" in message + assert "1 repeat count(s)" in message + + +def test_a_pad_mode_onnx_does_not_define_is_rejected(tmp_path): + model = _model( + [helper.make_node("Pad", ["x", "p"], ["y"], name="fill", mode="mirror")], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [helper.make_empty_tensor_value_info("y")], + initializer=[_int64("p", [1, 0, 1, 0])], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`fill`" in message + assert "mirror" in message + assert "constant, edge, reflect, wrap" in message + + +@pytest.mark.parametrize("mode", ["edge", "reflect", "wrap"]) +def test_padding_an_empty_axis_from_the_operand_is_rejected(tmp_path, mode): + """Only a constant pad can widen an axis the operand has no values along.""" + model = _model( + [helper.make_node("Pad", ["x", "p"], ["y"], name="fill", mode=mode)], + [_tensor("x", TensorProto.FLOAT, (0, 3))], + [helper.make_empty_tensor_value_info("y")], + initializer=[_int64("p", [1, 0, 1, 0])], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`fill`" in message + assert f"`{mode}`" in message + assert "empty along it" in message + + +def test_gathering_elements_across_disagreeing_axes_is_rejected(tmp_path): + """Every axis but the gathered one addresses both operands, so both must measure alike.""" + model = _model( + [helper.make_node("GatherElements", ["x", "i"], ["y"], name="pick", axis=0)], + [ + _tensor("x", TensorProto.FLOAT, (3, 4)), + _tensor("i", TensorProto.INT64, (2, 5)), + ], + [helper.make_empty_tensor_value_info("y")], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`pick`" in message + assert "[2, 5]" in message and "[3, 4]" in message + + +@pytest.mark.parametrize( + ("label", "attributes", "lengths_shape", "expected"), + [ + ("axes", {"batch_axis": 0, "time_axis": 2}, (2,), "0 and 1 in either order"), + ("lengths", {"batch_axis": 0, "time_axis": 1}, (4,), "one per batch"), + ], +) +def test_reversing_what_onnx_does_not_define_is_rejected( + tmp_path, label, attributes, lengths_shape, expected +): + model = _model( + [ + helper.make_node( + "ReverseSequence", ["x", "l"], ["y"], name="reverse", **attributes + ) + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3, 4)), + _tensor("l", TensorProto.INT64, lengths_shape), + ], + [helper.make_empty_tensor_value_info("y")], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`reverse`" in message + assert expected in message + + +def test_an_op_only_folding_serves_is_refused_where_folding_declines(tmp_path): + """`ConstantOfShape` carries no kernel: a graph that fixes its shape folds it away. + + That rests on the reference evaluator being a valid oracle for the revision, which it is + only from 25 on. Below it the node survives folding and dispatch refuses it by name, + rather than a kernel serving semantics nothing can vouch for. + """ + model = _model( + [ + helper.make_node( + "ConstantOfShape", + ["s"], + ["y"], + name="fill", + value=helper.make_tensor("v", TensorProto.FLOAT, [1], [1.5]), + ) + ], + [], + [helper.make_empty_tensor_value_info("y")], + initializer=[_int64("s", [2, 3])], + opset=21, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`ConstantOfShape`" in message + assert "opset version 21" in message + assert "no kernel is registered" in message + + +def test_the_pad_revision_onnx_infers_no_shape_for_names_the_nearest_version(tmp_path): + """Pad-1 spells its pads `paddings`, and ONNX derives no shape for a node of it. + + A revision whose result the compiler can only take on the model's word is one it cannot + prove anything about, so dispatch refuses it by name instead of reading the attribute + under another spelling. + """ + model = _model( + [ + helper.make_node( + "Pad", ["x"], ["y"], name="fill", paddings=[1, 0, 1, 0], mode="constant" + ) + ], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [_tensor("y", TensorProto.FLOAT, (4, 3))], + opset=1, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`Pad`" in message + assert "opset version 1" in message + assert "Nearest supported version: 2" in message + + +@pytest.mark.parametrize( + ("op_type", "opset", "nearest"), + [("Tile", 5, 6), ("Concat", 1, 4)], +) +def test_a_view_revision_no_kernel_claims_names_the_nearest_version( + tmp_path, op_type, opset, nearest +): + """Tile-1 repeats along one named axis, and Concat-1's `axis` default is unreadable. + + Neither is the op the generator implements, and neither has an oracle to prove one + against, so dispatch says so rather than serving the current semantics under the old + revision's name. + """ + # Tile-1 takes a repeat count and the single axis to apply it to, both as operands of + # the tensor's own element type; Concat takes any number of operands and an axis. + tiling = op_type == "Tile" + model = _model( + [ + helper.make_node( + op_type, + ["x", "n", "a"] if tiling else ["x", "x"], + ["y"], + name="view", + **({} if tiling else {"axis": 0}), + ) + ], + [_tensor("x", TensorProto.FLOAT, (2, 3))], + [_tensor("y", TensorProto.FLOAT, (4, 3))], + initializer=[ + onnx.numpy_helper.from_array(np.array([value], dtype=np.float32), name) + for name, value in (("n", 2.0), ("a", 0.0)) + ] + if tiling + else [], + opset=opset, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert f"`{op_type}`" in message + assert f"opset version {opset}" in message + assert f"Nearest supported version: {nearest}" in message + + +# -------------------------------------------------------------------------------------- +# The Fourier transforms +# -------------------------------------------------------------------------------------- + +# What a transform computes is settled by the conformance and differential suites, against +# ONNX's corpus and reference evaluator. What is asserted here is what those cannot reach: +# that the axis, the length and the frame geometry are read at compile time into the +# addressing one shared kernel walks, that an axis named only at run time becomes a switch +# over those call sites where it cannot change the result's shape, and the errors for the +# models the compiler refuses outright. + +_DFT_OPSET = 20 +_DFT_ATTRIBUTE_OPSET = 19 +_STFT_OPSET = 17 + + +def _dft_model( + shape, + *, + result_shape=None, + inputs=("x",), + initializer=(), + opset=_DFT_OPSET, + elem_type=TensorProto.FLOAT, + **attributes, +): + result = ( + helper.make_empty_tensor_value_info("y") + if result_shape is None + else _tensor("y", elem_type, result_shape) + ) + declared = [_tensor("x", elem_type, shape)] + declared += [ + _tensor(name, TensorProto.INT64, ()) + for name in inputs[1:] + if name and name not in {entry.name for entry in initializer} + ] + return _model( + [helper.make_node("DFT", list(inputs), ["y"], name="dft", **attributes)], + declared, + [result], + initializer=initializer, + opset=opset, + ) + + +def _stft_model( + shape, + *, + inputs=("x", "step"), + initializer=(), + window_shape=None, + **attributes, +): + declared = [_tensor("x", TensorProto.FLOAT, shape)] + if window_shape is not None: + declared.append(_tensor("window", TensorProto.FLOAT, window_shape)) + return _model( + [helper.make_node("STFT", list(inputs), ["y"], name="stft", **attributes)], + declared, + [helper.make_empty_tensor_value_info("y")], + initializer=initializer, + opset=_STFT_OPSET, + ) + + +def test_every_transform_of_one_element_type_shares_a_kernel(tmp_path): + """Axis, length and mode are addressing, so only the element type reaches the code.""" + model = _model( + [ + helper.make_node("DFT", ["x", "", "first"], ["p"], name="first"), + helper.make_node("DFT", ["x", "long", "last"], ["q"], name="padded"), + helper.make_node("DFT", ["x"], ["r"], name="reversed", inverse=1), + ], + [_tensor("x", TensorProto.FLOAT, (2, 4, 2))], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q", "r")], + initializer=[_int64("first", 0), _int64("last", 1), _int64("long", 6)], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "dft") + assert header.count(f"static void {kernel}(") == 1 + # The call sites, told from the definition by the indent their arguments carry. + assert header.count(f"{kernel}(\n ") == 3 + + +def test_the_transform_reaches_the_kernel_as_call_site_literals(tmp_path): + """One block per leading coordinate, one stride per trailing one, and the bin count.""" + model = _dft_model( + (2, 6, 3, 1), + inputs=("x", "", "axis"), + initializer=[_int64("axis", 1)], + onesided=1, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "dft") + assert ( + f"{kernel}(\n y,\n x,\n" + " 2u,\n" + " 3u,\n" + " 6u,\n" + " 4u,\n" + " 6u,\n" + " 1u,\n" + " 2u,\n" + " 0,\n" + " 0);" in header + ) + + +def test_the_inverse_one_sided_transform_writes_a_real_result(tmp_path): + """It is the one transform whose result has no imaginary part, and the one that + mirrors its operand: the length it writes is twice its own extent, less two.""" + model = _dft_model((1, 5, 2), inverse=1, onesided=1) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "dft") + assert ( + f"{kernel}(\n y,\n x,\n" + " 1u,\n" + " 1u,\n" + " 5u,\n" + " 8u,\n" + " 8u,\n" + " 2u,\n" + " 1u,\n" + " 1,\n" + " 1);" in header + ) + + +def test_the_two_revisions_read_their_default_axis_differently(tmp_path): + """Revision 17 states the axis as an attribute defaulting to 1, revision 20 as an + operand defaulting to the last signal axis; a rank-4 operand tells the two apart.""" + older = _dft_model((2, 3, 4, 1), opset=_DFT_ATTRIBUTE_OPSET) + newer = _dft_model((2, 3, 4, 1)) + + (report, header), (newer_report, newer_header) = ( + _compile(older, tmp_path / "older"), + _compile(newer, tmp_path / "newer"), + ) + + (kernel,) = _kernels(report, "dft") + (newer_kernel,) = _kernels(newer_report, "dft") + # The transformed axis is the operand's, so the blocks before it and the strides after + # it are what the default moves: axis 1 against axis 2. + assert f"{kernel}(\n y,\n x,\n 2u,\n 4u,\n" in header + assert ( + f"{newer_kernel}(\n y,\n x,\n 6u,\n 1u,\n" + in newer_header + ) + + +@requires_c_compiler +def test_an_axis_named_at_run_time_switches_over_the_axes_it_could_name(tmp_path): + """The result keeps the operand's extents whichever axis is transformed, so every axis + is a call site of its own and the operand only chooses between them.""" + model = _dft_model((2, 4, 3, 1), inputs=("x", "", "axis")) + signal = np.arange(24, dtype=np.float32).reshape(2, 4, 3, 1) + + result = compile_onnx(model, tmp_path) + loaded = result.load() + + (kernel,) = _kernels(result.report, "dft") + header = result.header_path.read_text(encoding="utf-8") + prefix = result.report["prefix"].upper() + assert header.count(f"static void {kernel}(") == 1 + assert ( + f"switch ({result.report['prefix']}_normalized_axis((int64_t)axis[0], 4))" + in (header) + ) + assert [f"case {axis}:" in header for axis in range(4)] == [True, True, True, False] + assert f"default:\n return {prefix}_ERROR_INVALID_ARGUMENT;" in header + for axis in (0, 1, 2, -2, -3, -4): + outputs = loaded.run({"x": signal, "axis": np.array(axis, dtype=np.int64)}) + expected = ReferenceEvaluator(model).run( + None, {"x": signal, "axis": np.array(axis, dtype=np.int64)} + ) + np.testing.assert_allclose(outputs["y"], expected[0], rtol=1e-3, atol=1e-6) + + +@requires_c_compiler +def test_an_axis_outside_the_operands_rank_returns_the_argument_error(tmp_path): + """The last axis holds the real and imaginary parts, so it is one no transform names.""" + model = _dft_model((2, 4, 3, 1), inputs=("x", "", "axis")) + + loaded = compile_onnx(model, tmp_path).load() + + with pytest.raises(HarnessError, match="status 1"): + loaded.run( + { + "x": np.zeros((2, 4, 3, 1), dtype=np.float32), + "axis": np.array(3, dtype=np.int64), + } + ) + + +@pytest.mark.parametrize( + ("label", "attributes", "initializer"), + [ + ("onesided", {"onesided": 1}, []), + ("dft_length", {}, [_int64("length", 4)]), + ], +) +def test_an_axis_named_at_run_time_that_resizes_the_result_is_refused( + tmp_path, label, attributes, initializer +): + """Both resize the axis they land on, so which axis that is decides the result's shape, + and no buffer can be sized before the operand is read.""" + model = _dft_model( + (2, 4, 3, 1), + inputs=("x", "length" if initializer else "", "axis"), + initializer=initializer, + **attributes, + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`dft`" in message and "`axis`" in message + assert "depends on input data" in message + + +def test_a_transform_length_named_at_run_time_is_refused(tmp_path): + model = _dft_model((2, 4, 1), inputs=("x", "length")) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "`length`" in str(error.value) + assert "depends on input data" in str(error.value) + + +def test_a_transform_over_no_samples_at_all_is_refused(tmp_path): + """A one-sided transform of an empty axis still states a bin, which nothing defines: + numpy refuses the transform outright, so there is no such thing to compile.""" + model = _dft_model((2, 0, 1), result_shape=(2, 1, 2), onesided=1) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "transform length of 0" in str(error.value) + + +def test_a_transform_of_an_operand_that_is_not_a_signal_is_refused(tmp_path): + """The last axis is the real and imaginary parts, so it measures 1 or 2 and nothing + else; ONNX's own inference does not check it.""" + model = _dft_model((2, 4, 3), result_shape=(2, 4, 2)) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "measures 3" in str(error.value) + + +@requires_c_compiler +def test_a_window_is_read_at_run_time(tmp_path): + """Its values reach the kernel and nothing else about it does, so a window the graph + does not fix is a pointer the call site passes rather than a compile error.""" + model = _stft_model( + (1, 16, 1), + inputs=("x", "step", "window"), + initializer=[_int64("step", 4)], + window_shape=(8,), + ) + signal = np.arange(16, dtype=np.float32).reshape(1, 16, 1) + window = np.hanning(8).astype(np.float32) + + result = compile_onnx(model, tmp_path) + outputs = result.load().run({"x": signal, "window": window}) + + (kernel,) = _kernels(result.report, "stft") + assert f"{kernel}(\n y,\n x,\n window,\n" in ( + result.header_path.read_text(encoding="utf-8") + ) + expected = ReferenceEvaluator(model).run(None, {"x": signal, "window": window}) + np.testing.assert_allclose(outputs["y"], expected[0], rtol=1e-3, atol=1e-6) + + +def test_a_transform_with_no_window_passes_none_in_its_place(tmp_path): + """Every sample weighs the same, which the kernel reads as no window at all.""" + model = _stft_model( + (1, 16, 1), + inputs=("x", "step", "", "length"), + initializer=[_int64("step", 4), _int64("length", 8)], + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "stft") + assert ( + f"{kernel}(\n y,\n x,\n NULL,\n" + " 1u,\n" + " 16u,\n" + " 1u,\n" + " 3u,\n" + " 4u,\n" + " 8u,\n" + " 5u);" in header + ) + + +def test_the_frame_step_must_be_known_at_compile_time(tmp_path): + model = _stft_model( + (1, 16, 1), + inputs=("x", "step", "", "length"), + initializer=[_int64("length", 8)], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "`step`" in str(error.value) + assert "depends on input data" in str(error.value) + + +@pytest.mark.parametrize( + ("label", "initializer"), + [ + ("no_frame", [_int64("step", 4)]), + ("no_step", [_int64("step", 0), _int64("length", 8)]), + ], +) +def test_a_frame_layout_that_states_no_frames_is_refused(tmp_path, label, initializer): + """ONNX reads the frame from a window or a `frame_length`, and counts the frames by + dividing the signal by the step; without the first its own inference stops before a + shape and with a step of nothing it divides by zero, so neither reaches an artifact.""" + stated = ("x", "step") if label == "no_frame" else ("x", "step", "", "length") + model = _stft_model((1, 16, 1), inputs=stated, initializer=initializer) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "`stft`" in str(error.value) + + +def test_a_signal_shorter_than_one_frame_is_refused(tmp_path): + """Not one whole frame fits, and ONNX's own two readings of that disagree: its shape + inference truncates the frame count towards zero where the reference floors it.""" + model = _stft_model( + (1, 4, 1), + inputs=("x", "step", "", "length"), + initializer=[_int64("step", 3), _int64("length", 8)], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert "do not fit a signal of 4" in str(error.value) + + +# -------------------------------------------------------------------------------------- +# The quantization family +# -------------------------------------------------------------------------------------- + +# What these compute is settled by the conformance and differential suites, against ONNX's +# corpus and reference evaluator. What is asserted here is what those cannot reach: that the +# granularity of a scale — per tensor, per axis, per block — is resolved at compile time into +# the strides one shared kernel walks, that a filter parameter per output channel is addressed +# by channel at a rank the reference evaluator cannot evaluate at all, and the errors for the +# models the compiler refuses outright. + +_QUANTIZE_OPSET = 25 +_INTEGER_OPSET = 10 +_QLINEAR_MATMUL_OPSET = 21 + + +def _quantize_model( + shape, + scale_shape, + *, + zero_shape=(), + grid=TensorProto.UINT8, + opset=_QUANTIZE_OPSET, + output=None, + name="quantize", + **attributes, +): + declared = [ + _tensor("x", TensorProto.FLOAT, shape), + _tensor("y_scale", TensorProto.FLOAT, scale_shape), + ] + inputs = ["x", "y_scale"] + if zero_shape is not None: + declared.append(_tensor("y_zero_point", grid, zero_shape)) + inputs.append("y_zero_point") + return _model( + [helper.make_node("QuantizeLinear", inputs, ["y"], name=name, **attributes)], + declared, + [ + helper.make_empty_tensor_value_info("y") + if output is None + else _tensor("y", output, shape) + ], + opset=opset, + ) + + +def _dequantize_model( + shape, scale_shape, *, grid=TensorProto.UINT8, opset=_QUANTIZE_OPSET, **attributes +): + return _model( + [ + helper.make_node( + "DequantizeLinear", + ["x", "x_scale", "x_zero_point"], + ["y"], + name="dequantize", + **attributes, + ) + ], + [ + _tensor("x", grid, shape), + _tensor("x_scale", TensorProto.FLOAT, scale_shape), + _tensor("x_zero_point", grid, scale_shape), + ], + [helper.make_empty_tensor_value_info("y")], + opset=opset, + ) + + +def _conv_integer_model( + shape, + filter_shape, + *, + zero_shape=(), + filter_zero_shape=None, + grid=TensorProto.UINT8, + **attributes, +): + declared = [_tensor("x", grid, shape), _tensor("w", grid, filter_shape)] + inputs = ["x", "w"] + for name, operand_shape in (("x_zp", zero_shape), ("w_zp", filter_zero_shape)): + if operand_shape is None: + break + declared.append(_tensor(name, grid, operand_shape)) + inputs.append(name) + return _model( + [helper.make_node("ConvInteger", inputs, ["y"], name="conv", **attributes)], + declared, + [helper.make_empty_tensor_value_info("y")], + opset=_INTEGER_OPSET, + ) + + +def _qlinear_conv_model( + shape, + filter_shape, + *, + filter_scale_shape=(), + filter_zero_shape=(), + scale_shape=(), + grid=TensorProto.UINT8, + bias=False, + **attributes, +): + operands = ( + ("x", grid, shape), + ("x_scale", TensorProto.FLOAT, scale_shape), + ("x_zero_point", grid, scale_shape), + ("w", grid, filter_shape), + ("w_scale", TensorProto.FLOAT, filter_scale_shape), + ("w_zero_point", grid, filter_zero_shape), + ("y_scale", TensorProto.FLOAT, scale_shape), + ("y_zero_point", grid, scale_shape), + ) + declared = [_tensor(*operand) for operand in operands] + if bias: + declared.append(_tensor("b", TensorProto.INT32, (filter_shape[0],))) + return _model( + [ + helper.make_node( + "QLinearConv", + [entry.name for entry in declared], + ["y"], + name="qconv", + **attributes, + ) + ], + declared, + [helper.make_empty_tensor_value_info("y")], + opset=_INTEGER_OPSET, + ) + + +def _qlinear_matmul_model(left_shape, right_shape, *, parameter_shape=()): + operands = ( + ("a", TensorProto.UINT8, left_shape), + ("a_scale", TensorProto.FLOAT, parameter_shape), + ("a_zero_point", TensorProto.UINT8, parameter_shape), + ("b", TensorProto.UINT8, right_shape), + ("b_scale", TensorProto.FLOAT, parameter_shape), + ("b_zero_point", TensorProto.UINT8, parameter_shape), + ("y_scale", TensorProto.FLOAT, parameter_shape), + ("y_zero_point", TensorProto.UINT8, parameter_shape), + ) + declared = [_tensor(*operand) for operand in operands] + return _model( + [ + helper.make_node( + "QLinearMatMul", + [entry.name for entry in declared], + ["y"], + name="qmatmul", + ) + ], + declared, + [helper.make_empty_tensor_value_info("y")], + opset=_QLINEAR_MATMUL_OPSET, + ) + + +@pytest.mark.parametrize( + ("label", "kwargs", "block", "strides"), + [ + ("per_tensor", {"scale_shape": ()}, " -1,\n 1u,", "{0u, 0u}"), + ( + "single_element", + {"scale_shape": (1,)}, + " -1,\n 1u,", + "{0u, 0u}", + ), + ("per_axis", {"scale_shape": (8,)}, " -1,\n 1u,", "{0u, 1u}"), + ( + "per_axis_first", + {"scale_shape": (4,), "axis": 0}, + " -1,\n 1u,", + "{1u, 0u}", + ), + ( + "blocked", + {"scale_shape": (4, 2), "axis": 1, "block_size": 4}, + " 1,\n 4u,", + "{2u, 1u}", + ), + ( + "blocked_first_axis", + {"scale_shape": (2, 8), "axis": 0, "block_size": 2}, + " 0,\n 2u,", + "{8u, 1u}", + ), + ], +) +def test_the_granularity_of_a_scale_reaches_the_kernel_as_call_site_literals( + tmp_path, label, kwargs, block, strides +): + """All three granularities are one addressing: a stride per axis, and a divisor on the + axis a blocked scale repeats along. Which one a node states is settled at compile time.""" + scale_shape = kwargs.pop("scale_shape") + model = _quantize_model((4, 8), scale_shape, zero_shape=scale_shape, **kwargs) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "quantizelinear") + assert ( + f"{kernel}(\n y,\n x,\n y_scale,\n y_zero_point,\n" + " 32u,\n 2,\n (const size_t[]){4u, 8u},\n" + f"{block}\n" + f" (const size_t[]){strides},\n" + f" (const size_t[]){strides});" in header + ) + + +def test_maps_of_one_grid_and_precision_share_a_kernel(tmp_path): + """Granularity is addressing, so only the types the kernel reads and writes reach the + code — and whether a zero point shifts the grid at all, which is a formula of its own.""" + model = _model( + [ + helper.make_node("QuantizeLinear", ["x", "s", "z"], ["p"], name="tensor"), + helper.make_node( + "QuantizeLinear", ["x", "v", "w"], ["q"], name="axis", axis=1 + ), + helper.make_node("QuantizeLinear", ["x", "s"], ["r"], name="bare"), + ], + [ + _tensor("x", TensorProto.FLOAT, (4, 8)), + _tensor("s", TensorProto.FLOAT, ()), + _tensor("z", TensorProto.UINT8, ()), + _tensor("v", TensorProto.FLOAT, (8,)), + _tensor("w", TensorProto.UINT8, (8,)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("p", "q", "r")], + opset=_QUANTIZE_OPSET, + ) + + report, header = _compile(model, tmp_path) + + # The one that reads a zero point carries a suffix saying so, and so is the longer name. + bare, shifted = sorted(_kernels(report, "quantizelinear"), key=len) + assert header.count(f"static void {bare}(") == 1 + assert header.count(f"static void {shifted}(") == 1 + # The two that state a zero point share one definition and call it from both sites. + assert header.count(f"{shifted}(\n ") == 2 + assert f"{bare}(\n r,\n x,\n s,\n NULL," in header + + +def test_the_rounding_store_is_one_helper_shared_across_the_family(tmp_path): + """Every op that writes a grid saturates onto it the same way, so the store is one + function per grid rather than one per op.""" + model = _model( + [ + helper.make_node("QuantizeLinear", ["x", "s", "z"], ["a"], name="quantize"), + helper.make_node( + "QLinearMatMul", + ["a", "s", "z", "b", "s", "z", "s", "z"], + ["y"], + name="qmatmul", + ), + ], + [ + _tensor("x", TensorProto.FLOAT, (4, 8)), + _tensor("s", TensorProto.FLOAT, ()), + _tensor("z", TensorProto.UINT8, ()), + _tensor("b", TensorProto.UINT8, (8, 4)), + ], + [helper.make_empty_tensor_value_info("y")], + opset=_QLINEAR_MATMUL_OPSET, + ) + + report, header = _compile(model, tmp_path) + + (saturate,) = [name for name in report["kernels"] if "saturate" in name] + assert saturate.endswith("_uint8_t") + assert header.count(f"static uint8_t {saturate}(") == 1 + assert header.count(f"{saturate}(") == 3 + assert "return (uint8_t)rint(value);" in header + + +@pytest.mark.parametrize( + ("filter_scale_shape", "filter_zero_shape", "scale_stride", "zero_stride"), + [((), (), "0u", "0u"), ((2,), (2,), "1u", "1u"), ((2,), (), "1u", "0u")], +) +def test_a_filter_parameter_per_output_channel_reaches_the_kernel_as_a_stride( + tmp_path, filter_scale_shape, filter_zero_shape, scale_stride, zero_stride +): + """A filter's scale and zero point are one for the whole filter or one per output + channel, which is the difference between a stride of zero and a stride of one.""" + model = _qlinear_conv_model( + (1, 1, 5, 5), + (2, 1, 2, 2), + filter_scale_shape=filter_scale_shape, + filter_zero_shape=filter_zero_shape, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "qlinearconv") + assert ( + f"{kernel}(\n y,\n x,\n w,\n x_zero_point,\n" + f" w_zero_point,\n {zero_stride},\n x_scale,\n" + f" w_scale,\n {scale_stride},\n y_scale,\n" + " y_zero_point,\n NULL," in header + ) + + +@requires_c_compiler +def test_a_filter_zero_point_per_channel_is_read_by_channel_at_any_rank(tmp_path): + """ONNX defines it as one zero point per output channel, whatever the filter's rank. + + Its reference evaluator stretches the operand over four axes regardless, so it can + evaluate a per-channel zero point at two spatial axes and nowhere else — which is what + the differential sweep is left covering. Splitting the filter here and running each + output channel through the evaluator with a zero point of its own puts the one-spatial- + axis case back within reach of the same oracle. + """ + shape, filter_shape = (2, 2, 7), (3, 2, 3) + generator = np.random.default_rng(20260726) + x = generator.integers(0, 256, size=shape, dtype=np.uint8) + w = generator.integers(0, 256, size=filter_shape, dtype=np.uint8) + x_zero = np.uint8(37) + w_zero = np.array([3, 130, 255], np.uint8) + compiled = compile_onnx( + _conv_integer_model(shape, filter_shape, filter_zero_shape=(3,)), tmp_path + ).load() + + got = compiled.run({"x": x, "w": w, "x_zp": x_zero, "w_zp": w_zero})["y"] + + expected = np.concatenate( + [ + ReferenceEvaluator( + _conv_integer_model(shape, (1, *filter_shape[1:]), filter_zero_shape=()) + ).run( + None, + { + "x": x, + "w": w[channel : channel + 1], + "x_zp": x_zero, + "w_zp": w_zero[channel], + }, + )[0] + for channel in range(filter_shape[0]) + ], + axis=1, + ) + np.testing.assert_array_equal(got, expected) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"precision": TensorProto.DOUBLE}, "states `precision` `DOUBLE`"), + ( + {"scale_shape": (4, 2)}, + "with no `block_size`, ONNX defines a scale of more than one element", + ), + ( + {"scale_shape": (4, 3), "axis": 1, "block_size": 4}, + "takes a scale of [4, 2], but `y_scale` has shape [4, 3]", + ), + ( + {"scale_shape": (4, 2), "axis": 1, "block_size": 0 - 2}, + "states `block_size` -2", + ), + ({"scale_shape": (4,)}, "does not broadcast to [4, 8]"), + ({"scale_shape": (8,), "axis": 3}, "axis 3 is out of range"), + ], +) +def test_a_granularity_the_compiler_cannot_address_is_rejected( + tmp_path, kwargs, message +): + scale_shape = kwargs.pop("scale_shape", ()) + model = _quantize_model((4, 8), scale_shape, zero_shape=scale_shape, **kwargs) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +def test_a_zero_point_of_another_shape_than_its_scale_is_rejected(tmp_path): + """The two together fix the granularity, so ONNX defines them as one shape; a pair that + disagrees would be read at two granularities at once.""" + model = _quantize_model((4, 8), (8,), zero_shape=(4,)) + + with pytest.raises(CompileError, match="ONNX defines the two as one shape"): + compile_onnx(model, tmp_path) + + +def test_a_zero_point_off_the_grid_it_shifts_is_rejected(tmp_path): + """`y_zero_point` is what states the grid, so a result declared as another type is two + answers to what this node quantizes onto.""" + model = _quantize_model((4, 8), (), output=TensorProto.INT16) + + with pytest.raises(CompileError, match="ONNX defines the two as one type"): + compile_onnx(model, tmp_path) + + +def test_quantizing_onto_a_type_that_is_no_grid_is_rejected(tmp_path): + """The saturation range is the grid's own, so there is none to round onto here.""" + model = _quantize_model((4, 8), (), zero_shape=None, output=TensorProto.INT32) + + with pytest.raises(CompileError, match="quantizes onto the integer grids"): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("parameter_shape", "message"), + [((2, 1), "[2, 1]"), ((2,), "[2]")], +) +def test_a_matrix_product_quantized_per_row_is_rejected( + tmp_path, parameter_shape, message +): + """ONNX reads these per row of `A` and per column of `B`. In the form its own text + describes — an `M`-element vector against an `[M, K]` operand — the reference evaluator + stretches that vector along numpy's trailing axis instead, so nothing can vouch for what + a kernel should compute; the granularity goes unserved as a whole rather than be read one + way there and another in the `[M, 1]` form numpy does broadcast as written.""" + model = _qlinear_matmul_model((2, 4), (4, 3), parameter_shape=parameter_shape) + + with pytest.raises(CompileError, match=re.escape(message)) as error: + compile_onnx(model, tmp_path) + + assert "per-tensor granularity only" in str(error.value) + + +def test_a_filter_parameter_of_neither_granularity_is_rejected(tmp_path): + model = _qlinear_conv_model( + (1, 1, 5, 5), (2, 1, 2, 2), filter_scale_shape=(3,), filter_zero_shape=(2,) + ) + + with pytest.raises(CompileError, match=re.escape("a 1-D tensor of 2")): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("op_type", "nearest"), [("QuantizeLinear", 10), ("DequantizeLinear", 19)] +) +def test_a_map_below_its_supported_revision_names_the_nearest_version( + tmp_path, op_type, nearest +): + """Opset 13 revised both, and no oracle covers that revision: the reference evaluator + does not distinguish it and no corpus test imports it, so it is not claimed at all.""" + model = ( + _quantize_model((4, 8), (), opset=13) + if op_type == "QuantizeLinear" + else _dequantize_model((4, 8), (), opset=13) + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + assert f"`{op_type}`" in str(error.value) + assert f"Nearest supported version: {nearest}." in str(error.value) + + +@requires_c_compiler +def test_a_quantized_model_stays_on_the_grid_between_its_ops(tmp_path): + """The pipeline these ops exist for: quantize, convolve on the grid, read it back. + + Every tensor between the first op and the last is an integer one, which is the whole + point of compiling a quantized model — nothing widens back to float in between — and the + result is the reference evaluator's for the same chain. + """ + shape, filter_shape = (1, 1, 5, 5), (2, 1, 3, 3) + parameters = ( + ("x_scale", TensorProto.FLOAT, ()), + ("x_zero_point", TensorProto.UINT8, ()), + ("w", TensorProto.UINT8, filter_shape), + ("w_scale", TensorProto.FLOAT, (2,)), + ("w_zero_point", TensorProto.UINT8, (2,)), + ("y_scale", TensorProto.FLOAT, ()), + ("y_zero_point", TensorProto.UINT8, ()), + ) + model = _model( + [ + helper.make_node( + "QuantizeLinear", + ["x", "x_scale", "x_zero_point"], + ["q"], + name="quantize", + ), + helper.make_node( + "QLinearConv", + [ + "q", + "x_scale", + "x_zero_point", + *[name for name, _, _ in parameters[2:]], + ], + ["p"], + name="qconv", + ), + helper.make_node( + "DequantizeLinear", + ["p", "y_scale", "y_zero_point"], + ["y"], + name="dequantize", + ), + ], + [_tensor("x", TensorProto.FLOAT, shape), *(_tensor(*p) for p in parameters)], + [helper.make_empty_tensor_value_info("y")], + opset=_QUANTIZE_OPSET, + ) + generator = np.random.default_rng(20260726) + feeds = { + "x": generator.normal(size=shape).astype(np.float32), + "x_scale": np.float32(0.017), + "x_zero_point": np.uint8(128), + "w": generator.integers(0, 256, size=filter_shape, dtype=np.uint8), + "w_scale": np.array([0.011, 0.023], np.float32), + "w_zero_point": np.array([127, 130], np.uint8), + "y_scale": np.float32(0.09), + "y_zero_point": np.uint8(64), + } + result = compile_onnx(model, tmp_path) + + got = result.load().run(feeds)["y"] + + header = result.header_path.read_text(encoding="utf-8") + assert "static uint8_t" in header + assert not re.search(r"^static (float|double)", header, re.MULTILINE) + expected = ReferenceEvaluator(model).run(None, feeds)[0] + np.testing.assert_allclose(got, expected, rtol=1e-3, atol=1e-7) + + +# -------------------------------------------------------------------------------------- +# The normalization by a root mean square, and the cross-entropy loss +# -------------------------------------------------------------------------------------- + +_RMS_OPSET = 23 +_SCE_OPSET = 13 + + +def _rms_model(shape=(2, 3), *, opset=_RMS_OPSET, **attributes): + return _model( + [ + helper.make_node( + "RMSNormalization", ["x", "s"], ["y"], name="rms", **attributes + ) + ], + [ + _tensor("x", TensorProto.FLOAT, shape), + _tensor("s", TensorProto.FLOAT, shape[-1:]), + ], + [helper.make_empty_tensor_value_info("y")], + opset=opset, + ) + + +def _sce_model( + scores_shape=(3, 5), + *, + weighted=False, + log_prob=False, + labels_type=TensorProto.INT64, + **attributes, +): + inputs = ["scores", "labels"] + (["weights"] if weighted else []) + outputs = ["loss"] + (["log_prob"] if log_prob else []) + labels_shape = (scores_shape[0], *scores_shape[2:]) + return _model( + [ + helper.make_node( + "SoftmaxCrossEntropyLoss", inputs, outputs, name="sce", **attributes + ) + ], + [ + _tensor("scores", TensorProto.FLOAT, scores_shape), + _tensor("labels", labels_type, labels_shape), + *( + [_tensor("weights", TensorProto.FLOAT, scores_shape[1:2])] + if weighted + else [] + ), + ], + [helper.make_empty_tensor_value_info(name) for name in outputs], + opset=_SCE_OPSET, + ) + + +def test_a_stash_type_the_reference_refuses_is_a_compile_error(tmp_path): + """The reference evaluator computes RMSNormalization in the data's own type and raises + on any `stash_type` but its default, so nothing vouches for another one.""" + with pytest.raises(CompileError, match="stash_type"): + compile_onnx(_rms_model(stash_type=TensorProto.DOUBLE), tmp_path) + + +@requires_c_compiler +def test_nodes_normalizing_at_one_element_type_share_one_kernel(tmp_path): + """The extents and the epsilon are call-site literals, so the axis does not fork it.""" + model = _model( + [ + helper.make_node( + "RMSNormalization", ["x", "s"], ["h"], name="first", axis=1 + ), + helper.make_node( + "RMSNormalization", ["h", "s"], ["y"], name="second", axis=-1 + ), + ], + [ + _tensor("x", TensorProto.FLOAT, (2, 3)), + _tensor("s", TensorProto.FLOAT, (3,)), + ], + [helper.make_empty_tensor_value_info("y")], + opset=_RMS_OPSET, + ) + + report, header = _compile(model, tmp_path) + + assert len(_kernels(report, "rmsnormalization")) == 1 + assert header.count("static void kernels_rmsnormalization") == 1 + + +@requires_c_compiler +@pytest.mark.parametrize( + ("shape", "weighted", "log_prob", "attributes"), + [ + ((3, 5), False, False, {}), + ((3, 5), True, True, {"reduction": "sum"}), + ((3, 5, 2), True, False, {"reduction": "none", "ignore_index": -1}), + ((2, 3, 2, 2), False, True, {"ignore_index": 1}), + ], +) +def test_the_loss_matches_the_reference_across_its_operand_combinations( + tmp_path, shape, weighted, log_prob, attributes +): + model = _sce_model(shape, weighted=weighted, log_prob=log_prob, **attributes) + generator = np.random.default_rng(20260726) + labels_shape = (shape[0], *shape[2:]) + feeds = { + "scores": generator.normal(size=shape).astype(np.float32), + "labels": generator.integers(-1, shape[1], size=labels_shape).astype(np.int64), + } + if weighted: + feeds["weights"] = np.abs(generator.normal(size=shape[1])).astype(np.float32) + # A label ONNX does not ignore has to name a class; the draw above reaches -1, which + # only the variants naming it as `ignore_index` may see. + if attributes.get("ignore_index") != -1: + feeds["labels"] = np.abs(feeds["labels"]) % shape[1] + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + expected = ReferenceEvaluator(model).run(None, feeds) + for name, reference in zip(["loss", "log_prob"], expected): + np.testing.assert_allclose(outputs[name], reference, rtol=1e-3, atol=1e-7) + + +@requires_c_compiler +def test_a_label_outside_the_class_axis_is_an_argument_error(tmp_path): + """The reference raises on one; the artifact reads nothing and reports the status.""" + compiled = compile_onnx(_sce_model(), tmp_path).load() + + with pytest.raises(HarnessError, match="status 1"): + compiled.run( + { + "scores": np.zeros((3, 5), dtype=np.float32), + "labels": np.array([0, 5, 1], dtype=np.int64), + } + ) + + +@requires_c_compiler +def test_a_label_the_node_ignores_may_sit_outside_the_class_axis(tmp_path): + """`ignore_index` is checked first, so the entry it names is skipped rather than refused.""" + model = _sce_model(ignore_index=-1) + feeds = { + "scores": np.arange(15, dtype=np.float32).reshape(3, 5), + "labels": np.array([0, -1, 4], dtype=np.int64), + } + + outputs = compile_onnx(model, tmp_path).load().run(feeds) + + expected = ReferenceEvaluator(model).run(None, feeds) + np.testing.assert_allclose(outputs["loss"], expected[0], rtol=1e-3, atol=1e-7) + + +@requires_c_compiler +def test_the_unweighted_loss_takes_no_weight_operand(tmp_path): + """A parameter the kernel never reads is what the artifact's `-Werror` build refuses, + so the operand combination is part of the kernel's identity rather than a branch.""" + plain, _ = _compile(_sce_model(), tmp_path / "plain") + weighted, header = _compile(_sce_model(weighted=True), tmp_path / "weighted") + + (plain_kernel,) = _kernels(plain, "softmaxcrossentropyloss") + (weighted_kernel,) = _kernels(weighted, "softmaxcrossentropyloss") + assert plain_kernel != weighted_kernel + assert "const float* weights" in header + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"reduction": "median"}, "reduction"), + ({"scores_shape": (5,)}, "rank of at least 2"), + ], +) +def test_the_loss_refuses_what_it_cannot_group(tmp_path, kwargs, message): + with pytest.raises(CompileError, match=message): + compile_onnx(_sce_model(**kwargs), tmp_path) + + +# -------------------------------------------------------------------------------------- +# LinearAttention +# -------------------------------------------------------------------------------------- + +_LINEAR_ATTENTION_OPSET = 27 + +# The operands in schema order; the ones a node leaves out reach it as empty names. +_LINEAR_ATTENTION_INPUTS = ("query", "key", "value", "past_state", "decay", "beta") + +# Where the call site carries what, in the kernel's own parameter order: eight pointers, +# seven extents, then the strides that place each of the two optional gates and the scale. +_LINEAR_ATTENTION_DECAY_STRIDES = slice(15, 18) +_LINEAR_ATTENTION_BETA_STRIDES = slice(18, 20) + + +def _linear_attention_operands( + *, + batch=2, + steps=3, + q_heads=2, + kv_heads=2, + d_k=3, + d_v=2, + past=False, + decay=None, + beta=None, +): + """The shape of every operand a node passes, by name, in schema order. + + Every 3-D operand packs `H * D` into its last axis, so the shapes follow from the two + head counts and the two head widths. `decay` and `beta` name the granularity ONNX packs + each of them at -- one value per key dimension or one per head for the decay, one per + head or one the heads share for beta -- and leaving either out is how `update_rule` + reaches a node: the rule forbids the operand it does not read. + """ + operands = { + "query": (batch, steps, q_heads * d_k), + "key": (batch, steps, kv_heads * d_k), + "value": (batch, steps, kv_heads * d_v), + } + if past: + operands["past_state"] = (batch, kv_heads, d_k, d_v) + if decay is not None: + operands["decay"] = ( + batch, + steps, + kv_heads * (1 if decay == "per_head" else d_k), + ) + if beta is not None: + operands["beta"] = (batch, steps, kv_heads if beta == "per_head" else 1) + return operands + + +def _linear_attention_model( + *, q_heads=2, kv_heads=2, state_shape=None, attributes=None, **geometry +): + """One LinearAttention node, with `state_shape` declaring its second output by hand.""" + operands = _linear_attention_operands( + q_heads=q_heads, kv_heads=kv_heads, **geometry + ) + names = [name if name in operands else "" for name in _LINEAR_ATTENTION_INPUTS] + while names and not names[-1]: + names.pop() + node = helper.make_node( + "LinearAttention", + names, + ["output", "present_state"], + name="attn", + q_num_heads=q_heads, + kv_num_heads=kv_heads, + **(attributes or {}), + ) + state = ( + helper.make_empty_tensor_value_info("present_state") + if state_shape is None + else _tensor("present_state", TensorProto.FLOAT, state_shape) + ) + return _model( + [node], + [_tensor(name, TensorProto.FLOAT, shape) for name, shape in operands.items()], + [helper.make_empty_tensor_value_info("output"), state], + opset=_LINEAR_ATTENTION_OPSET, + ) + + +def _linear_attention_feeds(model, seed=20260726): + return { + entry.name: np.random.default_rng([seed, index]) + .normal(size=[dim.dim_value for dim in entry.type.tensor_type.shape.dim]) + .astype(np.float32) + for index, entry in enumerate(model.graph.input) + } + + +def _linear_attention_emission(operands, results, **attributes): + """The registered kernel generator, run on one node directly. + + ONNX's own shape inference vets this op thoroughly: it rejects every operand combination + the reference evaluator refuses, and does so before dispatch ever asks for a kernel, so a + model is no way to reach the generator's own refusals. They are the compiler's last word + for a node that arrives without inference having vetted it, and this is where they are + read. + """ + node = helper.make_node( + "LinearAttention", + [name if shape is not None else "" for name, shape in operands], + [name for name, _ in results], + name="attn", + **{"q_num_heads": 2, "kv_num_heads": 2, **attributes}, + ) + spec = KERNELS.select("", "LinearAttention", _LINEAR_ATTENTION_OPSET) + context = NodeContext( + node=node, + domain="", + opset_version=_LINEAR_ATTENTION_OPSET, + since_version=spec.since_version, + prefix="attn", + inputs=tuple( + None if shape is None else TensorRef(name, TensorProto.FLOAT, shape, name) + for name, shape in operands + ), + outputs=tuple( + TensorRef(name, TensorProto.FLOAT, shape, name) for name, shape in results + ), + ) + return spec.generator(context) + + +def _linear_attention_node(*, decay=None, beta=None): + """One node's operands and results as `_linear_attention_emission` takes them.""" + shapes = _linear_attention_operands(decay=decay, beta=beta) + return ( + [(name, shapes.get(name)) for name in _LINEAR_ATTENTION_INPUTS], + [("output", (2, 3, 4)), ("present_state", (2, 2, 3, 2))], + ) + + +def _linear_attention_call(header, kernel, index=0): + """One emitted call site's arguments, in order. + + The kernel's own definition opens the same way, so the split's first piece is its + parameter list and the call sites follow. + """ + body = header.split(f"{kernel}(\n")[index + 2].split(");")[0] + return [line.strip().rstrip(",") for line in body.splitlines()] + + +# Every combination of a rule and the two optional gates ONNX refuses, with the operand its +# refusal is about: `gated` reads the decay and forbids beta, `delta` the other way round, +# `gated_delta` reads both and `linear` neither. The reference raises for a stray operand as +# readily as for a missing one, so there is nothing for a kernel to compute for any of these. +_LINEAR_ATTENTION_MISMATCHES = ( + ("linear", "per_key_dim", None, "decay"), + ("linear", None, "per_head", "beta"), + ("gated", None, None, "decay"), + ("gated", "per_key_dim", "per_head", "beta"), + ("delta", "per_key_dim", "per_head", "decay"), + ("delta", None, None, "beta"), + ("gated_delta", None, "per_head", "decay"), + ("gated_delta", "per_key_dim", None, "beta"), +) + + +@pytest.mark.parametrize( + ("rule", "decay", "beta", "operand"), _LINEAR_ATTENTION_MISMATCHES +) +def test_a_rule_and_the_gates_it_reads_have_to_agree( + tmp_path, rule, decay, beta, operand +): + """The compiler refuses exactly the models the reference does, and for the same reason.""" + model = _linear_attention_model( + decay=decay, beta=beta, attributes={"update_rule": rule} + ) + + with pytest.raises(CompileError): + compile_onnx(model, tmp_path) + + with pytest.raises(ValueError, match=f"'{rule}' (requires|forbids) {operand}"): + ReferenceEvaluator(model).run(None, _linear_attention_feeds(model)) + + +@pytest.mark.parametrize( + ("rule", "decay", "beta", "operand"), _LINEAR_ATTENTION_MISMATCHES +) +def test_the_kernel_refuses_a_rule_its_operands_contradict(rule, decay, beta, operand): + """The same refusal as the kernel's own, naming the rule and the operand it is about.""" + operands, results = _linear_attention_node(decay=decay, beta=beta) + + with pytest.raises(CompileError, match=f"`{rule}`") as error: + _linear_attention_emission(operands, results, update_rule=rule) + + assert f"`{operand}`" in str(error.value) + + +def test_an_update_rule_onnx_does_not_define_is_refused(tmp_path): + """The four recurrences are the whole of what the attribute may name.""" + model = _linear_attention_model( + decay="per_key_dim", beta="per_head", attributes={"update_rule": "chunked"} + ) + operands, results = _linear_attention_node(decay="per_key_dim", beta="per_head") + + with pytest.raises(CompileError): + compile_onnx(model, tmp_path) + with pytest.raises(ValueError, match="chunked"): + ReferenceEvaluator(model).run(None, _linear_attention_feeds(model)) + + with pytest.raises(CompileError, match="`chunked`") as error: + _linear_attention_emission(operands, results, update_rule="chunked") + + assert "linear, gated, delta, gated_delta" in str(error.value) + + +def test_the_op_is_served_by_a_kernel_rather_than_by_its_function_body(tmp_path): + """Why this op has a kernel at all, written down. + + ONNX defines a function body for `LinearAttention` and the compiler prefers one over a + kernel wherever it can -- but this body drives the recurrence with a `Scan` over the + sequence, whose trip count is a run-time tensor rather than anything constant folding can + resolve, which puts it on the v1 unsupported surface. The registry claiming the op is + what keeps every model of it off that path. + """ + assert onnx.defs.get_schema( + "LinearAttention", _LINEAR_ATTENTION_OPSET, "" + ).has_context_dependent_function + assert KERNELS.registered_versions("", "LinearAttention") == [ + _LINEAR_ATTENTION_OPSET + ] + + report, _ = _compile( + _linear_attention_model(decay="per_key_dim", beta="per_head"), tmp_path + ) + + assert _kernels(report, "linearattention") + + +def test_linear_attention_nodes_of_one_element_type_share_a_kernel(tmp_path): + """Batch, sequence, head counts and head widths are call-site literals, not kernels.""" + model = _model( + [ + helper.make_node( + "LinearAttention", + ["q", "k", "v"], + ["y", "s"], + name="wide", + q_num_heads=2, + kv_num_heads=2, + update_rule="linear", + ), + helper.make_node( + "LinearAttention", + ["q2", "k2", "v2"], + ["y2", "s2"], + name="narrow", + q_num_heads=3, + kv_num_heads=1, + update_rule="linear", + ), + ], + [ + _tensor("q", TensorProto.FLOAT, (2, 3, 6)), + _tensor("k", TensorProto.FLOAT, (2, 3, 6)), + _tensor("v", TensorProto.FLOAT, (2, 3, 4)), + _tensor("q2", TensorProto.FLOAT, (1, 5, 6)), + _tensor("k2", TensorProto.FLOAT, (1, 5, 2)), + _tensor("v2", TensorProto.FLOAT, (1, 5, 7)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("y", "s", "y2", "s2")], + opset=_LINEAR_ATTENTION_OPSET, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "linearattention") + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +def test_the_running_state_is_the_present_state_buffer(tmp_path): + """`present_state` is a required output, so the recurrence needs no storage of its own: + it updates that buffer in place and the sequence ends with the answer already there.""" + model = _linear_attention_model(decay="per_key_dim", beta="per_head", past=True) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "linearattention") + assert _linear_attention_call(header, kernel)[:3] == [ + "output", + "present_state", + "query", + ] + assert f"static float {report['prefix']}_linearattention" not in header + + +@pytest.mark.parametrize( + ("decay", "strides"), + [("per_head", ["2u", "1u", "0u"]), ("per_key_dim", ["6u", "3u", "1u"])], +) +def test_the_decay_granularity_reaches_the_kernel_as_strides(tmp_path, decay, strides): + """Which granularity the gate is packed at is where its elements sit, not what the + kernel computes: the per-head packing reaches every key dimension through a zero + stride.""" + model = _linear_attention_model(decay=decay, attributes={"update_rule": "gated"}) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "linearattention") + arguments = _linear_attention_call(header, kernel) + assert arguments[_LINEAR_ATTENTION_DECAY_STRIDES] == strides + + +@pytest.mark.parametrize( + ("beta", "strides"), [("per_head", ["2u", "1u"]), ("shared", ["1u", "0u"])] +) +def test_the_beta_granularity_reaches_the_kernel_as_strides(tmp_path, beta, strides): + """A beta the heads share is one the head axis is addressed with a zero stride.""" + model = _linear_attention_model(beta=beta, attributes={"update_rule": "delta"}) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "linearattention") + arguments = _linear_attention_call(header, kernel) + assert arguments[_LINEAR_ATTENTION_BETA_STRIDES] == strides + + +def test_the_default_scale_is_the_one_the_schema_derives(tmp_path): + """A `scale` of 0 -- the attribute's own default -- asks for `1/sqrt(d_k)`, which no + model could mean literally: a zero factor answers every query with zero. The derived + factor and the same number stated outright have to reach the kernel alike.""" + geometry = {"d_k": 4, "decay": "per_key_dim", "beta": "per_head"} + report, derived = _compile( + _linear_attention_model(**geometry), tmp_path / "derived" + ) + _, stated = _compile( + _linear_attention_model(attributes={"scale": 0.5}, **geometry), + tmp_path / "stated", + ) + + (kernel,) = _kernels(report, "linearattention") + assert _linear_attention_call(derived, kernel)[-1] != "0.0" + assert _linear_attention_call(derived, kernel) == _linear_attention_call( + stated, kernel + ) + + +def test_a_state_output_shaped_otherwise_is_refused(tmp_path): + """The state's shape follows from the head counts and the two head widths; a graph that + declares another one describes a different op, and this is where that stops rather than + where the kernel writes past the buffer.""" + model = _linear_attention_model( + decay="per_key_dim", beta="per_head", state_shape=(2, 2, 2, 3) + ) + + with pytest.raises(CompileError, match=re.escape("[2, 2, 3, 2]")) as error: + compile_onnx(model, tmp_path) + + assert "present_state" in str(error.value) + + +@requires_c_compiler +def test_a_sequence_decoded_one_token_at_a_time_matches_the_prefill(tmp_path): + """What the two state operands exist for: a prefill over a whole sequence and one decode + step per token, each handed the state the last one reported, are the same recurrence. + + Neither suite reaches this -- the sweep and the corpus both run a single node once -- and + the expected values are the reference evaluator's, on the prefill the chain is supposed + to equal. + """ + steps = 4 + packing = {"decay": "per_key_dim", "beta": "per_head", "past": True} + prefill = _linear_attention_model(steps=steps, **packing) + feeds = _linear_attention_feeds(prefill) + expected = ReferenceEvaluator(prefill).run(None, feeds) + + compiled = compile_onnx( + _linear_attention_model(steps=1, **packing), tmp_path + ).load() + state = feeds["past_state"] + answers = [] + for step in range(steps): + result = compiled.run( + { + name: state if name == "past_state" else feeds[name][:, step : step + 1] + for name in feeds + } + ) + answers.append(result["output"]) + state = result["present_state"] + + np.testing.assert_allclose( + np.concatenate(answers, axis=1), expected[0], rtol=1e-3, atol=1e-6 + ) + np.testing.assert_allclose(state, expected[1], rtol=1e-3, atol=1e-6) + + +# -------------------------------------------------------------------------------------- +# Attention and RotaryEmbedding +# -------------------------------------------------------------------------------------- + +_TRANSFORMER_OPSET = 24 + +# The operands and the results in schema order; the ones a node leaves out reach it as +# empty names, and the ones it does not ask for are simply not there. +_ATTENTION_INPUTS = ( + "Q", + "K", + "V", + "attn_mask", + "past_key", + "past_value", + "nonpad_kv_seqlen", +) +_ATTENTION_OUTPUTS = ("Y", "present_key", "present_value", "qk_matmul_output") +_ROTARY_INPUTS = ("X", "cos_cache", "sin_cache", "position_ids") + +# Attention emits its scorer and two helpers under one `_attention_` family; these +# are the tokens that tell the helpers apart from the scorer itself. +_ATTENTION_HELPERS = ("bias", "present") + +_ATTENTION_OPERANDS = { + "Q": (TensorProto.FLOAT, (1, 2, 2, 4)), + "K": (TensorProto.FLOAT, (1, 2, 3, 4)), + "V": (TensorProto.FLOAT, (1, 2, 3, 4)), +} + +_ROTARY_OPERANDS = { + "X": (TensorProto.FLOAT, (1, 2, 3, 4)), + "cos_cache": (TensorProto.FLOAT, (5, 2)), + "sin_cache": (TensorProto.FLOAT, (5, 2)), + "position_ids": (TensorProto.INT64, (1, 3)), +} + + +def _transformer_model(op_type, order, operands, results, **attributes): + names = [entry if entry in operands else "" for entry in order] + while names and not names[-1]: + names.pop() + node = helper.make_node(op_type, names, list(results), name="node", **attributes) + return _model( + [node], + [_tensor(entry, *operands[entry]) for entry in order if entry in operands], + [helper.make_empty_tensor_value_info(name) for name in results], + opset=_TRANSFORMER_OPSET, + ) + + +def _attention_model(operands, *, outputs=1, **attributes): + return _transformer_model( + "Attention", + _ATTENTION_INPUTS, + operands, + _ATTENTION_OUTPUTS[:outputs], + **attributes, + ) + + +def _rotary_model(operands, **attributes): + return _transformer_model( + "RotaryEmbedding", _ROTARY_INPUTS, operands, ("Y",), **attributes + ) + + +def _attention_kernels(report, role=None): + """Attention's emitted kernels: the scorer itself, or one of its named helpers.""" + start = len(f"{report['prefix']}_attention_") + return [ + name + for name in _kernels(report, "attention") + if name[start:].split("_")[0] == role + or (role is None and name[start:].split("_")[0] not in _ATTENTION_HELPERS) + ] + + +def _feeds(model, seed=20260726): + """A seeded value for every input the model declares, at its own dtype and shape.""" + generator = np.random.default_rng(seed) + drawn = {} + for entry in model.graph.input: + tensor = entry.type.tensor_type + shape = tuple(dim.dim_value for dim in tensor.shape.dim) + dtype = np.dtype(helper.tensor_dtype_to_np_dtype(tensor.elem_type)) + drawn[entry.name] = ( + generator.normal(size=shape).astype(dtype) + if dtype.kind == "f" + else generator.integers(0, 2, size=shape).astype(dtype) + ) + return drawn + + +def test_both_attention_layouts_share_one_kernel(tmp_path): + """`(batch, head, sequence, size)` and `(batch, sequence, head * size)` are one tensor at + two sets of strides, so the loop nest is the same code and the layout arrives as + arguments.""" + model = _model( + [ + helper.make_node("Attention", ["q4", "k4", "v4"], ["y4"], name="packed"), + helper.make_node( + "Attention", + ["q3", "k3", "v3"], + ["y3"], + name="flat", + q_num_heads=2, + kv_num_heads=2, + ), + ], + [ + _tensor("q4", TensorProto.FLOAT, (1, 2, 2, 4)), + _tensor("k4", TensorProto.FLOAT, (1, 2, 3, 4)), + _tensor("v4", TensorProto.FLOAT, (1, 2, 3, 4)), + _tensor("q3", TensorProto.FLOAT, (1, 2, 8)), + _tensor("k3", TensorProto.FLOAT, (1, 3, 8)), + _tensor("v3", TensorProto.FLOAT, (1, 3, 8)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("y4", "y3")], + opset=_TRANSFORMER_OPSET, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _attention_kernels(report) + assert header.count(f"static void {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +def test_a_narrowed_softmax_precision_emits_a_scorer_of_its_own(tmp_path): + """The type the softmax runs in is part of the emitted code, so it is part of the name.""" + model = _model( + [ + helper.make_node("Attention", ["q", "k", "v"], ["y"], name="wide"), + helper.make_node( + "Attention", + ["q", "k", "v"], + ["z"], + name="narrow", + softmax_precision=TensorProto.FLOAT, + ), + ], + [ + _tensor("q", TensorProto.FLOAT, (1, 2, 2, 4)), + _tensor("k", TensorProto.FLOAT, (1, 2, 3, 4)), + _tensor("v", TensorProto.FLOAT, (1, 2, 3, 4)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("y", "z")], + opset=_TRANSFORMER_OPSET, + ) + + report, _ = _compile(model, tmp_path) + + kernels = sorted(_attention_kernels(report)) + assert len(kernels) == 2 + assert [name.split("_soft")[1] for name in kernels] == [ + "double_maskfloat", + "float_maskfloat", + ] + + +def test_a_boolean_mask_emits_a_bias_of_its_own(tmp_path): + """The two mask forms are different expressions over a differently typed operand.""" + model = _model( + [ + helper.make_node("Attention", ["q", "k", "v", "f"], ["y"], name="additive"), + helper.make_node("Attention", ["q", "k", "v", "b"], ["z"], name="boolean"), + ], + [ + _tensor("q", TensorProto.FLOAT, (1, 2, 2, 4)), + _tensor("k", TensorProto.FLOAT, (1, 2, 3, 4)), + _tensor("v", TensorProto.FLOAT, (1, 2, 3, 4)), + _tensor("f", TensorProto.FLOAT, (2, 3)), + _tensor("b", TensorProto.BOOL, (2, 3)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("y", "z")], + opset=_TRANSFORMER_OPSET, + ) + + report, header = _compile(model, tmp_path) + + assert len(_attention_kernels(report)) == 2 + assert len(_attention_kernels(report, "bias")) == 2 + # The NaN a boolean mask turns an entry that does take part into under `is_causal`; + # the additive form has no such value anywhere in it. + assert header.count("(causal ? (double)NAN : 0.0) : -INFINITY;") == 1 + + +def test_the_caches_are_written_only_where_the_node_asks_for_them(tmp_path): + """`present_key` and `present_value` are a concatenation nothing else needs.""" + past = (TensorProto.FLOAT, (1, 2, 2, 4)) + operands = {**_ATTENTION_OPERANDS, "past_key": past, "past_value": past} + + report, _ = _compile(_attention_model(operands), tmp_path) + assert not _attention_kernels(report, "present") + + report, header = _compile(_attention_model(operands, outputs=3), tmp_path) + (present,) = _attention_kernels(report, "present") + assert header.count(f"{present}(\n") == 3 + + +def test_the_row_of_scores_is_one_buffer_sized_for_the_longest_row(tmp_path): + """The softmax needs the whole row before it can normalize any of it, and the artifact + allocates nothing; nodes sharing the kernel share that row, sized for the largest.""" + model = _model( + [ + helper.make_node("Attention", ["q", "k", "v"], ["y"], name="short"), + helper.make_node("Attention", ["q", "k2", "v2"], ["z"], name="long"), + ], + [ + _tensor("q", TensorProto.FLOAT, (1, 2, 2, 4)), + _tensor("k", TensorProto.FLOAT, (1, 2, 3, 4)), + _tensor("v", TensorProto.FLOAT, (1, 2, 3, 4)), + _tensor("k2", TensorProto.FLOAT, (1, 2, 7, 4)), + _tensor("v2", TensorProto.FLOAT, (1, 2, 7, 4)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("y", "z")], + opset=_TRANSFORMER_OPSET, + ) + + report, header = _compile(model, tmp_path) + + assert ( + header.count(f"static double {report['prefix']}_attention_scores_double[") == 1 + ) + assert f"{report['prefix']}_attention_scores_double[7];" in header + + +@requires_c_compiler +def test_a_mask_of_one_axis_reaches_every_query_row(tmp_path): + """ONNX defines `attn_mask` as broadcastable to the whole score tensor, which a single + key axis is; the corpus only ever ships the 2-D and 4-D forms.""" + model = _attention_model( + {**_ATTENTION_OPERANDS, "attn_mask": (TensorProto.FLOAT, (3,))} + ) + feeds = _feeds(model) + + got = compile_onnx(model, tmp_path).load().run(feeds)["Y"] + + expected = ReferenceEvaluator(model).run(None, feeds)[0] + np.testing.assert_allclose(got, expected, rtol=1e-3, atol=1e-7) + + +def test_both_rotation_patterns_share_one_kernel(tmp_path): + """`interleaved` picks which two lanes a pair is made of and nothing else, so it reaches + the loop nest as an argument rather than forking it.""" + model = _model( + [ + helper.make_node( + "RotaryEmbedding", ["x", "cos", "sin"], ["y"], name="halves" + ), + helper.make_node( + "RotaryEmbedding", + ["x", "cos", "sin"], + ["z"], + name="interleaved", + interleaved=1, + ), + ], + [ + _tensor("x", TensorProto.FLOAT, (1, 2, 3, 4)), + _tensor("cos", TensorProto.FLOAT, (1, 3, 2)), + _tensor("sin", TensorProto.FLOAT, (1, 3, 2)), + ], + [helper.make_empty_tensor_value_info(name) for name in ("y", "z")], + opset=_TRANSFORMER_OPSET, + ) + + report, header = _compile(model, tmp_path) + + (kernel,) = _kernels(report, "rotaryembedding") + assert header.count(f"static int {kernel}(") == 1 + assert header.count(f"{kernel}(\n") == 3 + + +@requires_c_compiler +def test_a_position_the_cache_has_no_row_for_is_reported(tmp_path): + """`position_ids` is read at run time, so an index outside the cache is the argument + error the status enum exists for rather than a read past the buffer.""" + compiled = compile_onnx(_rotary_model(_ROTARY_OPERANDS), tmp_path).load() + feeds = _feeds(_rotary_model(_ROTARY_OPERANDS)) + feeds["position_ids"] = np.array([[0, 1, 5]], np.int64) + + with pytest.raises(HarnessError, match="status 1"): + compiled.run(feeds) + + +@pytest.mark.parametrize( + ("operands", "attributes", "message"), + [ + ( + {"V": (TensorProto.DOUBLE, (1, 2, 3, 4))}, + {}, + "this compiler attends one element type", + ), + ( + {"attn_mask": (TensorProto.INT32, (2, 3))}, + {}, + "a boolean mask or a float mask of the operands' own type", + ), + ( + {}, + {"softmax_precision": TensorProto.FLOAT16}, + "`softmax_precision` of `FLOAT16`", + ), + ( + {}, + {"softmax_precision": TensorProto.INT32}, + "`softmax_precision` of `INT32`", + ), + ({}, {"qk_matmul_output_mode": 4}, "`qk_matmul_output_mode` of 4"), + ( + { + "Q": (TensorProto.FLOAT, (1, 2, 9)), + "K": (TensorProto.FLOAT, (1, 3, 8)), + "V": (TensorProto.FLOAT, (1, 3, 8)), + }, + {"q_num_heads": 2, "kv_num_heads": 2}, + "hidden axis of 9 into 2 head(s), which does not divide it", + ), + ({}, {"q_num_heads": 3}, "states `q_num_heads` 3"), + ( + {"K": (TensorProto.FLOAT, (1, 2, 3, 6))}, + {}, + "contracts `Q` against `K` over the head size", + ), + ( + {"V": (TensorProto.FLOAT, (1, 2, 5, 4))}, + {}, + "reads `K` and `V` at the same key positions", + ), + ( + {"Q": (TensorProto.FLOAT, (1, 3, 2, 4))}, + {}, + "needs 3 query head(s) to be a multiple of 2", + ), + ( + {"past_key": (TensorProto.FLOAT, (1, 2, 2, 4))}, + {}, + "ONNX defines them as used together", + ), + ( + { + "past_key": (TensorProto.FLOAT, (1, 2, 2, 4)), + "past_value": (TensorProto.FLOAT, (1, 2, 3, 4)), + }, + {}, + "as a cache of shape [1, 2, 2, 4]", + ), + ( + {"nonpad_kv_seqlen": (TensorProto.INT64, (2,))}, + {}, + "one key length per batch item", + ), + ( + {"attn_mask": (TensorProto.FLOAT, (5, 3))}, + {}, + "does not broadcast", + ), + # Shorter than the key axis is padded out with -inf; longer is what the reference + # raises on, so there is nothing for the columns past the end to mean. + ( + {"attn_mask": (TensorProto.FLOAT, (2, 4))}, + {}, + "reads `attn_mask` along 4 key position(s), but the node attends 3", + ), + # The reference takes the triangle's extent from the mask's own query axis, so a + # mask with no such axis is a node it cannot evaluate at all. + ( + {"attn_mask": (TensorProto.FLOAT, (3,))}, + {"is_causal": 1}, + "the triangle's extent from the mask's own query axis", + ), + ], +) +def test_an_attention_node_the_compiler_cannot_serve_is_rejected( + tmp_path, operands, attributes, message +): + model = _attention_model({**_ATTENTION_OPERANDS, **operands}, **attributes) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("operands", "attributes", "message"), + [ + ({}, {"rotary_embedding_dim": 3}, "rotates the first 3 lane(s) of a head of 4"), + ({}, {"rotary_embedding_dim": 6}, "rotates the first 6 lane(s) of a head of 4"), + ({}, {"num_heads": 3}, "states `num_heads` 3"), + ( + {"X": (TensorProto.FLOAT, (1, 3, 9))}, + {"num_heads": 2}, + "hidden axis of 9 into 2 head(s), which does not divide it", + ), + ( + {"sin_cache": (TensorProto.FLOAT, (5, 1))}, + {}, + "as one pair of caches", + ), + ( + { + "cos_cache": (TensorProto.FLOAT, (5, 1)), + "sin_cache": (TensorProto.FLOAT, (5, 1)), + }, + {}, + "rotates 2 pair(s) per head", + ), + ( + { + "cos_cache": (TensorProto.FLOAT, (1, 3, 2)), + "sin_cache": (TensorProto.FLOAT, (1, 3, 2)), + }, + {}, + "reads `cos_cache` as rank 2", + ), + ( + {"position_ids": (TensorProto.INT64, (1, 1, 3))}, + {}, + "reads `position_ids` as `(batch, sequence)`", + ), + # Rank 1 is refused for the same reason rank 3 is: the reference inserts the head + # axis at position 2 of whatever it gathered, which for one axis fewer lands past + # the angles rather than before them. + ( + {"position_ids": (TensorProto.INT64, (3,))}, + {}, + "reads `position_ids` as `(batch, sequence)`", + ), + ], +) +def test_a_rotary_node_the_compiler_cannot_serve_is_rejected( + tmp_path, operands, attributes, message +): + model = _rotary_model({**_ROTARY_OPERANDS, **operands}, **attributes) + + with pytest.raises(CompileError, match=re.escape(message)): + compile_onnx(model, tmp_path) + + +def test_one_tensor_named_for_both_caches_is_checked_against_both(tmp_path): + """`past_key` and `past_value` are checked separately even when they are one tensor. + + A node may name the same tensor for both, and the two are then the same value — so a + check that pairs each operand with its expected shape has to keep the pair, not key on + the operand. `head_size` and `v_head_size` differ here, so at most one of the two can + hold, and the kernel would otherwise address 8 lanes of a cache laid out at 2. + """ + operands = { + "Q": (TensorProto.FLOAT, (1, 1, 2, 8)), + "K": (TensorProto.FLOAT, (1, 1, 2, 8)), + "V": (TensorProto.FLOAT, (1, 1, 2, 2)), + } + node = helper.make_node( + "Attention", ["Q", "K", "V", "", "cache", "cache"], ["Y"], name="node" + ) + model = _model( + [node], + [_tensor(name, *spec) for name, spec in operands.items()] + + [_tensor("cache", TensorProto.FLOAT, (1, 1, 3, 2))], + [helper.make_empty_tensor_value_info("Y")], + opset=_TRANSFORMER_OPSET, + ) + + with pytest.raises( + CompileError, match=re.escape("as a cache of shape [1, 1, 3, 8]") + ): + compile_onnx(model, tmp_path) + + +@requires_c_compiler +def test_two_losses_ignoring_different_labels_share_one_kernel(tmp_path): + """`ignore_index` decides what a kernel skips, not the code that skips it. + + Which of `reduction`, `weights` and the presence of the attribute a node carries forks + the kernel's body; the index itself is a call-site literal like every other attribute + value here, so two nodes that differ only in it share one emitted function instead of + claiming one name for two definitions. + """ + nodes = [ + helper.make_node( + "SoftmaxCrossEntropyLoss", + ["x", "t"], + [f"l{index}"], + name=f"loss{index}", + reduction="sum", + ignore_index=index, + ) + for index in (1, 2) + ] + nodes.append(helper.make_node("Add", ["l1", "l2"], ["y"], name="total")) + model = _model( + nodes, + [ + _tensor("x", TensorProto.FLOAT, (3, 5)), + _tensor("t", TensorProto.INT64, (3,)), + ], + [helper.make_empty_tensor_value_info("y")], + opset=_SCE_OPSET, + ) + feeds = { + "x": np.arange(15, dtype=np.float32).reshape(3, 5) / 3, + "t": np.array([0, 1, 2], dtype=np.int64), + } + + report, header = _compile(model, tmp_path) + + assert len([name for name in report["kernels"] if "crossentropy" in name]) == 1 + outputs = compile_onnx(model, tmp_path / "run").load().run(feeds) + expected = ReferenceEvaluator(model).run(None, feeds) + np.testing.assert_allclose(outputs["y"], expected[0], rtol=1e-6, atol=1e-6) + + +def test_a_cache_nothing_gathers_carries_its_own_positions(tmp_path): + """Without `position_ids` the caches stand as they are, which ONNX shapes + `(batch, sequence, rotary_embedding_dim / 2)` rather than by position.""" + operands = { + name: value + for name, value in _ROTARY_OPERANDS.items() + if name != "position_ids" + } + + with pytest.raises(CompileError, match="reads `cos_cache` as rank 3"): + compile_onnx(_rotary_model(operands), tmp_path) + + +# -------------------------------------------------------------------------------------- +# TfIdfVectorizer +# -------------------------------------------------------------------------------------- + +# What the op counts is settled by the conformance and differential suites, against ONNX's +# own corpus and reference evaluator. What is asserted here is that the pool becomes one +# shared kernel over `static const` tables rather than code per node, and the refusals for +# the nodes the compiler will not emit at all. +_TFIDF_OPSET = 9 +_TFIDF_POOL = { + "ngram_counts": [0, 4], + "ngram_indexes": [0, 1, 2, 3, 4, 5, 6], + "pool_int64s": [2, 3, 5, 4, 5, 6, 7, 8, 6, 7], +} + + +def _tfidf_model(*, shape=(12,), elem_type=TensorProto.INT64, nodes=1, **attributes): + counted = { + **_TFIDF_POOL, + "mode": "TF", + "min_gram_length": 1, + "max_gram_length": 2, + "max_skip_count": 0, + **attributes, + } + graph = [ + helper.make_node( + "TfIdfVectorizer", ["x"], [f"y{index}"], name=f"count{index}", **counted + ) + for index in range(nodes) + ] + return _model( + graph, + [_tensor("x", elem_type, shape)], + [helper.make_empty_tensor_value_info(f"y{index}") for index in range(nodes)], + opset=_TFIDF_OPSET, + ) + + +def test_the_pool_is_emitted_once_for_the_nodes_that_share_it(tmp_path): + """Two nodes over one pool: one kernel, and one copy of each table it walks.""" + report, header = _compile(_tfidf_model(nodes=2), tmp_path) + + assert len(_kernels(report, "tfidfvectorizer")) == 1 + for role in ("tokens", "targets", "first_edge", "edge_count", "counted_column"): + assert len(re.findall(rf"tfidfvectorizer_{role}_\w+\[\d+\] = ", header)) == 1 + + +def test_the_modes_that_read_no_weights_share_one_kernel(tmp_path): + """`TF` ignores the weights a node carries, so it emits the code a node without them does.""" + weighted = _tfidf_model(weights=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]) + + report, header = _compile(weighted, tmp_path) + + (kernel,) = _kernels(report, "tfidfvectorizer") + assert kernel.endswith("_tf_flat_int64_t") + # Past the preamble, whose footprint summary names weights as a memory category. + assert "weights" not in header.split("*/", 1)[1] + + +def test_a_string_pool_is_rejected(tmp_path): + """A string pool is matched against a string tensor, which the artifact cannot hold.""" + model = _tfidf_model(pool_strings=["a", "b", "c", "d"]) + del model.graph.node[0].attribute[ + [entry.name for entry in model.graph.node[0].attribute].index("pool_int64s") + ] + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path) + + message = str(error.value) + assert "`count0`" in message + assert "pool_strings" in message + + +@pytest.mark.parametrize( + ("attributes", "expected"), + ( + ({"mode": "COUNT"}, "`TF`, `IDF` and `TFIDF`"), + ({"min_gram_length": 0}, "`min_gram_length`"), + ({"min_gram_length": 3, "max_gram_length": 2}, "`min_gram_length`"), + ({"weights": [1.0, 2.0]}, "weight(s) for"), + # Every n-gram of the pool takes an identifier, and every identifier reads an index. + ({"ngram_indexes": [0, 1, 2]}, "`ngram_indexes` entries account for"), + ), +) +def test_a_node_the_compiler_cannot_count_from_is_rejected( + attributes, expected, tmp_path +): + with pytest.raises(CompileError) as error: + compile_onnx(_tfidf_model(**attributes), tmp_path) + + message = str(error.value) + assert "`count0`" in message + assert expected in message + + +def _tfidf_emission(shape, elem_type, width=7, **attributes): + """The registered kernel generator, run on one node directly. + + ONNX's own inference reads the rank and the index list before dispatch asks for a kernel, + and refuses what it cannot type: these are the generator's last word for a node that + arrives without inference having vetted it, and this is where they are read. + """ + node = helper.make_node( + "TfIdfVectorizer", + ["x"], + ["y"], + name="count0", + **{ + **_TFIDF_POOL, + "mode": "TF", + "min_gram_length": 1, + "max_gram_length": 2, + "max_skip_count": 0, + **attributes, + }, + ) + spec = KERNELS.select("", "TfIdfVectorizer", _TFIDF_OPSET) + context = NodeContext( + node=node, + domain="", + opset_version=_TFIDF_OPSET, + since_version=spec.since_version, + prefix="count", + inputs=(TensorRef("x", elem_type, shape, "x"),), + outputs=(TensorRef("y", TensorProto.FLOAT, (*shape[:-1], width), "y"),), + ) + return spec.generator(context) + + +@pytest.mark.parametrize( + ("shape", "elem_type", "attributes", "expected"), + ( + ((2, 3, 4), TensorProto.INT64, {}, "one token sequence or a batch"), + ((12,), TensorProto.FLOAT, {}, "int32` or `int64` tokens"), + ((12,), TensorProto.INT64, {"ngram_indexes": [0, -1]}, "`ngram_indexes` entry"), + ), +) +def test_a_node_inference_did_not_vet_is_still_refused( + shape, elem_type, attributes, expected +): + with pytest.raises(CompileError, match=re.escape(expected)): + _tfidf_emission(shape, elem_type, **attributes) diff --git a/src/python/tests/test_extra_compiler_loader.py b/src/python/tests/test_extra_compiler_loader.py new file mode 100644 index 0000000..1531e8c --- /dev/null +++ b/src/python/tests/test_extra_compiler_loader.py @@ -0,0 +1,275 @@ +"""Model loading, external-data resolution, and opset resolution for the C compiler.""" + +from __future__ import annotations + +import importlib +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +loader = pytest.importorskip("fnnx.extras.compilers.c.onnx.loader") + +MODELS_DIR = Path(__file__).parent / "models" +LINREG_MODEL = ( + MODELS_DIR / "onnx_pipeline.fnnx" / "ops_artifacts" / "linreg" / "model.onnx" +) + + +def _onnx_domain_maximum(domain: str) -> int: + """Highest opset ONNX itself reports for `domain`, independent of the loader.""" + return onnx.defs.C.schema_version_map()[domain][1] + + +def _identity_model(*opset_imports: tuple[str, int], ir_version: int | None = None): + graph = onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["x"], ["y"])], + "g", + [onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1])], + [onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1])], + ) + model = onnx.helper.make_model( + graph, + opset_imports=[ + onnx.helper.make_opsetid(domain, version) + for domain, version in opset_imports + ], + ) + if ir_version is not None: + model.ir_version = ir_version + return model + + +def _external_data_model(values, location: str): + tensor = onnx.numpy_helper.from_array(values, "w") + onnx.external_data_helper.set_external_data(tensor, location=location) + tensor.ClearField("raw_data") + graph = onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["w"], ["y"])], + "g", + [], + [ + onnx.helper.make_tensor_value_info( + "y", onnx.TensorProto.FLOAT, list(values.shape) + ) + ], + initializer=[tensor], + ) + return onnx.helper.make_model( + graph, opset_imports=[onnx.helper.make_opsetid("", 21)] + ) + + +class LoadModelTest(unittest.TestCase): + def test_loads_bundle_node_model_with_both_domains(self): + imported = { + entry.domain: entry.version + for entry in onnx.load(str(LINREG_MODEL)).opset_import + } + self.assertEqual(set(imported), {"", "ai.onnx.ml"}) + + loaded = loader.load_model(LINREG_MODEL) + + self.assertEqual(loaded.opsets, imported) + self.assertEqual( + loaded.opset_for("ai.onnx"), loaded.opset_for(loader.STANDARD_DOMAIN) + ) + + def test_accepts_in_memory_proto(self): + loaded = loader.load_model(_identity_model(("", 21))) + self.assertEqual(loaded.opsets, {loader.STANDARD_DOMAIN: 21}) + + def test_missing_file_names_the_path(self): + missing = MODELS_DIR / "does_not_exist.onnx" + with self.assertRaises(CompileError) as ctx: + loader.load_model(missing) + self.assertIn(str(missing), str(ctx.exception)) + + def test_unparseable_file_is_a_compile_error(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "broken.onnx" + path.write_bytes(b"this is not a protobuf") + with self.assertRaises(CompileError) as ctx: + loader.load_model(path) + self.assertIn(str(path), str(ctx.exception)) + + def test_ir_version_newer_than_installed_onnx(self): + model = _identity_model(("", 21), ir_version=onnx.IR_VERSION + 1) + with self.assertRaises(CompileError) as ctx: + loader.load_model(model) + message = str(ctx.exception) + self.assertIn(f"IR version {onnx.IR_VERSION + 1}", message) + self.assertIn(f"at most {onnx.IR_VERSION}", message) + self.assertIn("upgrade", message.lower()) + + +class ResolveOpsetsTest(unittest.TestCase): + def test_domain_alias_is_normalized(self): + opsets = loader.resolve_opsets(_identity_model(("ai.onnx", 17))) + self.assertEqual(opsets, {loader.STANDARD_DOMAIN: 17}) + + def test_ml_domain_is_supported(self): + opsets = loader.resolve_opsets(_identity_model(("", 21), (loader.ML_DOMAIN, 1))) + self.assertEqual(opsets[loader.ML_DOMAIN], 1) + + def test_repeated_consistent_import_is_accepted(self): + opsets = loader.resolve_opsets(_identity_model(("", 17), ("ai.onnx", 17))) + self.assertEqual(opsets, {loader.STANDARD_DOMAIN: 17}) + + def test_conflicting_imports_for_one_domain(self): + with self.assertRaises(CompileError) as ctx: + loader.resolve_opsets(_identity_model(("", 17), ("ai.onnx", 18))) + message = str(ctx.exception) + self.assertIn("17", message) + self.assertIn("18", message) + + def test_custom_domain_is_rejected(self): + with self.assertRaises(CompileError) as ctx: + loader.resolve_opsets(_identity_model(("", 21), ("com.example.ops", 1))) + self.assertIn("com.example.ops", str(ctx.exception)) + + def test_opset_newer_than_installed_onnx(self): + maximum = _onnx_domain_maximum(loader.STANDARD_DOMAIN) + with self.assertRaises(CompileError) as ctx: + loader.resolve_opsets(_identity_model(("", maximum + 1))) + message = str(ctx.exception) + self.assertIn(f"imports opset version {maximum + 1}", message) + self.assertIn(f"at most version {maximum}", message) + self.assertIn("upgrade", message.lower()) + + def test_ml_opset_newer_than_installed_onnx(self): + maximum = _onnx_domain_maximum(loader.ML_DOMAIN) + with self.assertRaises(CompileError) as ctx: + loader.resolve_opsets( + _identity_model(("", 21), (loader.ML_DOMAIN, maximum + 1)) + ) + message = str(ctx.exception) + self.assertIn(loader.ML_DOMAIN, message) + self.assertIn(f"imports opset version {maximum + 1}", message) + self.assertIn(f"at most version {maximum}", message) + + def test_invalid_opset_version(self): + with self.assertRaises(CompileError): + loader.resolve_opsets(_identity_model(("", 0))) + + def test_model_without_opset_imports(self): + model = _identity_model(("", 21)) + del model.opset_import[:] + with self.assertRaises(CompileError): + loader.resolve_opsets(model) + + def test_opset_for_unimported_domain(self): + loaded = loader.load_model(_identity_model(("", 21))) + with self.assertRaises(CompileError) as ctx: + loaded.opset_for(loader.ML_DOMAIN) + self.assertIn(loader.ML_DOMAIN, str(ctx.exception)) + + def test_max_supported_opset_matches_onnx_domain_version_map(self): + for domain in loader.SUPPORTED_DOMAINS: + self.assertEqual( + loader.max_supported_opset(domain), _onnx_domain_maximum(domain) + ) + + +class ExternalDataTest(unittest.TestCase): + def setUp(self): + self.values = np.arange(6, dtype=np.float32).reshape(2, 3) + + def test_external_tensor_is_embedded_from_the_model_directory(self): + with tempfile.TemporaryDirectory() as tmp: + (Path(tmp) / "w.bin").write_bytes(self.values.tobytes()) + path = Path(tmp) / "model.onnx" + onnx.save_model(_external_data_model(self.values, "w.bin"), str(path)) + + loaded = loader.load_model(path) + + initializer = loaded.model.graph.initializer[0] + self.assertFalse(onnx.external_data_helper.uses_external_data(initializer)) + np.testing.assert_array_equal( + onnx.numpy_helper.to_array(initializer), self.values + ) + + def test_base_dir_overrides_the_model_directory(self): + with tempfile.TemporaryDirectory() as tmp: + weights = Path(tmp) / "weights" + weights.mkdir() + (weights / "w.bin").write_bytes(self.values.tobytes()) + path = Path(tmp) / "model.onnx" + onnx.save_model(_external_data_model(self.values, "w.bin"), str(path)) + + loaded = loader.load_model(path, base_dir=weights) + + np.testing.assert_array_equal( + onnx.numpy_helper.to_array(loaded.model.graph.initializer[0]), + self.values, + ) + + def test_in_memory_proto_without_base_dir_names_the_tensor(self): + with self.assertRaises(CompileError) as ctx: + loader.load_model(_external_data_model(self.values, "w.bin")) + self.assertIn("`w`", str(ctx.exception)) + + def test_missing_external_file_names_the_directory(self): + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(CompileError) as ctx: + loader.load_model( + _external_data_model(self.values, "w.bin"), base_dir=tmp + ) + self.assertIn(tmp, str(ctx.exception)) + + def test_external_path_escaping_the_base_dir_is_rejected(self): + with tempfile.TemporaryDirectory() as tmp: + (Path(tmp) / "secret.bin").write_bytes(self.values.tobytes()) + inner = Path(tmp) / "artifacts" + inner.mkdir() + with self.assertRaises(CompileError): + loader.load_model( + _external_data_model(self.values, "../secret.bin"), base_dir=inner + ) + + def test_source_proto_is_not_mutated(self): + with tempfile.TemporaryDirectory() as tmp: + (Path(tmp) / "w.bin").write_bytes(self.values.tobytes()) + model = _external_data_model(self.values, "w.bin") + + loaded = loader.load_model(model, base_dir=tmp) + + self.assertTrue( + onnx.external_data_helper.uses_external_data(model.graph.initializer[0]) + ) + self.assertFalse( + onnx.external_data_helper.uses_external_data( + loaded.model.graph.initializer[0] + ) + ) + + +class OptionalDependencyTest(unittest.TestCase): + def test_missing_onnx_package_raises_an_actionable_error(self): + """A `ModuleNotFoundError` keeps `pytest.importorskip` skipping rather than erroring.""" + package = "fnnx.extras.compilers.c.onnx" + saved = { + name: module + for name, module in sys.modules.items() + if name == package or name.startswith(f"{package}.") + } + try: + for name in saved: + del sys.modules[name] + with mock.patch("importlib.util.find_spec", return_value=None): + with self.assertRaises(ModuleNotFoundError) as ctx: + importlib.import_module(package) + self.assertIn("fnnx[compiler]", str(ctx.exception)) + finally: + sys.modules.update(saved) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/python/tests/test_extra_compiler_ml.py b/src/python/tests/test_extra_compiler_ml.py new file mode 100644 index 0000000..9f8adaa --- /dev/null +++ b/src/python/tests/test_extra_compiler_ml.py @@ -0,0 +1,677 @@ +"""The `ai.onnx.ml` preprocessing surface: the ZipMap pass, its metadata, and the pipelines. + +What each op computes is settled by the conformance and differential suites, against ONNX's +own corpus and reference evaluator. The corpus is thin here — it carries a node test for two +of the nine ONNX-ML ops this compiler serves — so what this module adds is the coverage a +single-node sweep cannot reach: whole converted pipelines, whose expected values still come +from the reference evaluator; the graph pass and header metadata that have no op of their +own; and the errors for the models the compiler will not compile at all. +""" + +from __future__ import annotations + +import ctypes +import shutil + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +# The harness refuses to import without numpy, so this covers both dependencies. +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from onnx import TensorProto, ValueInfoProto, helper # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +from fnnx.extras.compilers.c import compile_onnx # noqa: E402 + +ML_OPSET = 5 +STANDARD_OPSET = 21 + +requires_c_compiler = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + + +def _tensor(name, elem_type, shape): + return helper.make_tensor_value_info(name, elem_type, list(shape)) + + +def _ml_node(op_type, inputs, outputs, **attributes): + return helper.make_node( + op_type, + list(inputs), + list(outputs), + name=outputs[0], + domain="ai.onnx.ml", + **attributes, + ) + + +def _model(nodes, inputs, outputs): + """A model whose intermediates carry no `value_info`, as a converter's output does.""" + graph = helper.make_graph(nodes, "ml", list(inputs), list(outputs)) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", STANDARD_OPSET), + helper.make_opsetid("ai.onnx.ml", ML_OPSET), + ], + ) + model.ir_version = 9 + return model + + +def _untyped(name): + return helper.make_empty_tensor_value_info(name) + + +def _map_output(name): + """A `ZipMap` result: the sequence of maps its schema declares, not a tensor.""" + entry = ValueInfoProto() + entry.name = name + entry.type.CopyFrom( + helper.make_sequence_type_proto( + helper.make_map_type_proto( + TensorProto.STRING, + helper.make_tensor_type_proto(TensorProto.FLOAT, []), + ) + ) + ) + return entry + + +def _matches_reference(model, feeds, tmp_path): + """Every output of the compiled artifact against the reference evaluator's own.""" + compiled = compile_onnx(model, tmp_path).load() + outputs = compiled.run(feeds) + expected = ReferenceEvaluator(model).run(None, feeds) + assert [spec.name for spec in compiled.outputs] == [ + entry.name for entry in model.graph.output + ] + for entry, want in zip(model.graph.output, expected): + got = outputs[entry.name] + assert got.dtype == want.dtype, entry.name + np.testing.assert_allclose(got, want, rtol=1e-6, atol=1e-7, err_msg=entry.name) + + +# -------------------------------------------------------------------------------------- +# Pipelines, which is what an ONNX-ML model is made of +# -------------------------------------------------------------------------------------- + + +@requires_c_compiler +def test_a_chain_of_preprocessors_compiles_and_matches_the_reference(tmp_path): + """The shape and type of every intermediate come from the compiler's own inference. + + ONNX ships no inference function for `Imputer`, `Scaler` or `Normalizer`, and a converted + model states nothing about its intermediates, so nothing but that inference gives these + tensors a type at all. + """ + model = _model( + [ + _ml_node( + "Imputer", + ["x"], + ["filled"], + imputed_value_floats=[0.5, -1.0, 2.0], + replaced_value_float=float("nan"), + ), + _ml_node( + "Scaler", + ["filled"], + ["scaled"], + offset=[1.0, 0.0, -1.0], + scale=[0.5, 2.0, 1.0], + ), + _ml_node("Normalizer", ["scaled"], ["y"], norm="L2"), + ], + [_tensor("x", TensorProto.FLOAT, [4, 3])], + [_untyped("y")], + ) + feeds = { + "x": np.array( + [[1.0, np.nan, 3.0], [np.nan, 0.0, -2.0], [0.0, 0.0, 0.0], [4.0, 5.0, 6.0]], + dtype=np.float32, + ) + } + + _matches_reference(model, feeds, tmp_path) + + +@requires_c_compiler +def test_a_feature_union_of_encoded_and_scaled_columns_matches_the_reference(tmp_path): + """The fan-out a converted `ColumnTransformer` produces: two branches, then a union.""" + model = _model( + [ + _ml_node( + "ArrayFeatureExtractor", ["x", "categorical"], ["category_column"] + ), + _ml_node( + "OneHotEncoder", + ["category_column"], + ["encoded"], + cats_int64s=[0, 1, 2], + zeros=1, + ), + helper.make_node( + "Reshape", ["encoded", "flat"], ["encoded_rows"], name="reshape" + ), + _ml_node("ArrayFeatureExtractor", ["x", "numeric"], ["numeric_columns"]), + _ml_node( + "Scaler", ["numeric_columns"], ["scaled"], offset=[1.0], scale=[0.25] + ), + _ml_node( + "FeatureVectorizer", + ["encoded_rows", "scaled"], + ["y"], + inputdimensions=[3, 2], + ), + ], + [_tensor("x", TensorProto.FLOAT, [4, 3])], + [_untyped("y")], + ) + model.graph.initializer.extend( + [ + helper.make_tensor("categorical", TensorProto.INT64, [1], [0]), + helper.make_tensor("numeric", TensorProto.INT64, [2], [1, 2]), + helper.make_tensor("flat", TensorProto.INT64, [2], [4, 3]), + ] + ) + feeds = { + "x": np.array( + [[0.0, 1.0, 2.0], [1.0, 3.0, 4.0], [2.0, 5.0, 6.0], [9.0, 7.0, 8.0]], + dtype=np.float32, + ) + } + + _matches_reference(model, feeds, tmp_path) + + +@requires_c_compiler +def test_a_label_encoder_over_integer_classes_matches_the_reference(tmp_path): + model = _model( + [ + _ml_node( + "LabelEncoder", + ["x"], + ["y"], + keys_int64s=[0, 1, 2], + values_floats=[0.25, -0.5, 1.0], + default_float=-9.0, + ) + ], + [_tensor("x", TensorProto.INT64, [2, 4])], + [_untyped("y")], + ) + feeds = {"x": np.array([[0, 1, 2, 3], [-1, 2, 1, 0]], dtype=np.int64)} + + _matches_reference(model, feeds, tmp_path) + + +@requires_c_compiler +def test_repeated_preprocessors_share_one_kernel_and_one_table(tmp_path): + """Two nodes of the same op and element type are one kernel; equal tables are one array.""" + model = _model( + [ + _ml_node("Scaler", ["x"], ["a"], offset=[1.0, 2.0], scale=[0.5, 0.5]), + _ml_node("Scaler", ["a"], ["b"], offset=[1.0, 2.0], scale=[0.5, 0.5]), + _ml_node("Scaler", ["b"], ["y"], offset=[3.0, 4.0], scale=[0.5, 0.5]), + ], + [_tensor("x", TensorProto.FLOAT, [2, 2])], + [_untyped("y")], + ) + + result = compile_onnx(model, tmp_path) + header = result.header_path.read_text(encoding="utf-8") + + assert [name for name in result.report["kernels"] if "scaler" in name] == [ + "ml_scaler_float" + ] + # Two distinct offset tables and one shared scale table, each defined exactly once. + tables = sorted( + { + line.split()[3].split("[")[0] + for line in header.splitlines() + if line.startswith("static const float ml_scaler_") + } + ) + assert len(tables) == 3 + assert result.report["memory"]["weights_bytes"] == 3 * 2 * 4 + + +@requires_c_compiler +def test_an_out_of_range_index_returns_a_nonzero_status(tmp_path): + """`ArrayFeatureExtractor` reads its columns from a run-time operand, so it checks them.""" + model = _model( + [_ml_node("ArrayFeatureExtractor", ["x", "i"], ["y"])], + [ + _tensor("x", TensorProto.FLOAT, [2, 3]), + _tensor("i", TensorProto.INT64, [2]), + ], + [_untyped("y")], + ) + compiled = compile_onnx(model, tmp_path).load() + + with pytest.raises(harness.HarnessError, match="status"): + compiled.run( + { + "x": np.zeros((2, 3), np.float32), + "i": np.array([0, 3], dtype=np.int64), + } + ) + + +@requires_c_compiler +def test_a_value_in_no_category_returns_a_nonzero_status(tmp_path): + """`zeros` cleared makes an unknown category the failure the schema prescribes.""" + model = _model( + [_ml_node("OneHotEncoder", ["x"], ["y"], cats_int64s=[1, 2], zeros=0)], + [_tensor("x", TensorProto.INT64, [3])], + [_untyped("y")], + ) + compiled = compile_onnx(model, tmp_path).load() + + assert compiled.run({"x": np.array([1, 2, 1], np.int64)})["y"].shape == (3, 2) + with pytest.raises(harness.HarnessError, match="status"): + compiled.run({"x": np.array([1, 2, 7], np.int64)}) + + +# -------------------------------------------------------------------------------------- +# The ZipMap pass and the class-label metadata it produces +# -------------------------------------------------------------------------------------- + +_LABELS = ["setosa", 'versi"colo\\r', "vir/*ginica"] + + +def _zipmap_model(labels_attribute, labels, *, extra_output=False): + """A stand-in classifier: a probability tensor, keyed by `ZipMap` into a map output.""" + nodes = [ + _ml_node("Scaler", ["x"], ["scores"], offset=[1.0, 0.0, -1.0], scale=[0.5] * 3), + _ml_node("Normalizer", ["scores"], ["probabilities"], norm="L1"), + _ml_node( + "ZipMap", + ["probabilities"], + ["output_probability"], + **{labels_attribute: labels}, + ), + ] + outputs = [_map_output("output_probability")] + if extra_output: + nodes.insert(0, helper.make_node("Identity", ["x"], ["passthrough"], name="id")) + outputs.append(_untyped("passthrough")) + return _model(nodes, [_tensor("x", TensorProto.FLOAT, [2, 3])], outputs) + + +def test_a_trailing_zipmap_is_replaced_by_the_tensor_it_reads(tmp_path): + result = compile_onnx(_zipmap_model("classlabels_strings", _LABELS), tmp_path) + + assert [entry["name"] for entry in result.report["entrypoint"]["outputs"]] == [ + "probabilities" + ] + (labels,) = result.report["class_labels"] + assert labels["tensor"] == "probabilities" + assert labels["dtype"] == "str" + assert labels["values"] == _LABELS + assert "ZipMap" not in result.header_path.read_text(encoding="utf-8") + + +def test_repeated_compiles_of_an_ml_model_are_byte_identical(tmp_path): + """A table an ONNX-ML kernel embeds is named after its contents, not after its node.""" + model = _zipmap_model("classlabels_strings", _LABELS) + + first = compile_onnx(model, tmp_path / "first") + second = compile_onnx(model, tmp_path / "second") + + assert first.header_path.read_bytes() == second.header_path.read_bytes() + assert first.report_path.read_bytes() == second.report_path.read_bytes() + + +def test_the_promoted_tensor_keeps_the_output_position_zipmap_held(tmp_path): + result = compile_onnx( + _zipmap_model("classlabels_strings", _LABELS, extra_output=True), tmp_path + ) + + assert [entry["name"] for entry in result.report["entrypoint"]["outputs"]] == [ + "probabilities", + "passthrough", + ] + + +def test_the_class_label_table_is_declared_with_a_count_macro(tmp_path): + result = compile_onnx(_zipmap_model("classlabels_int64s", [7, 8, 9]), tmp_path) + header = result.header_path.read_text(encoding="utf-8") + (labels,) = result.report["class_labels"] + + assert labels["dtype"] == "int64" + assert labels["values"] == [7, 8, 9] + assert f"#define {labels['macro']} 3" in header + assert f"extern const int64_t {labels['symbol']}[{labels['macro']}];" in header + + +def test_two_label_tables_on_one_output_get_their_own_count_macro(tmp_path): + """Each macro is derived from its table's own symbol, not from the tensor they share. + + Two tables keying one tensor would otherwise define one macro name twice, with the two + lengths — which is the header failing the `-Werror` build its own contract states. + """ + model = _zipmap_model("classlabels_strings", _LABELS) + model.graph.node.append( + _ml_node( + "ZipMap", ["probabilities"], ["output_label"], classlabels_int64s=[7, 8] + ) + ) + model.graph.output.append(_map_output("output_label")) + + result = compile_onnx(model, tmp_path) + defines = result.header_path.read_text(encoding="utf-8").splitlines() + strings, integers = result.report["class_labels"] + + assert strings["tensor"] == integers["tensor"] == "probabilities" + assert [ + line for line in defines if line.startswith(f"#define {strings['macro']} ") + ] == [f"#define {strings['macro']} 3"] + assert [ + line for line in defines if line.startswith(f"#define {integers['macro']} ") + ] == [f"#define {integers['macro']} 2"] + + +@requires_c_compiler +def test_the_promoted_tensor_holds_the_probabilities_the_labels_name(tmp_path): + """The pairing `ZipMap` stood for, read back off the tensor and the table replacing it. + + onnxruntime is the oracle here rather than the reference evaluator, which has no `ZipMap` + implementation at all: which label names which column is the one thing a run of the graph + with the node already removed cannot show. + """ + runtime = pytest.importorskip("onnxruntime") + model = _zipmap_model("classlabels_strings", _LABELS) + feeds = {"x": np.array([[1.0, 2.0, 3.0], [4.0, -5.0, 6.0]], dtype=np.float32)} + + result = compile_onnx(model, tmp_path) + probabilities = result.load().run(feeds)["probabilities"] + (rows,) = runtime.InferenceSession(model.SerializeToString()).run(None, feeds) + (labels,) = result.report["class_labels"] + + expected = np.array( + [[row[label] for label in labels["values"]] for row in rows], dtype=np.float32 + ) + assert probabilities.shape == expected.shape + np.testing.assert_allclose(probabilities, expected, rtol=1e-6, atol=1e-7) + + +@requires_c_compiler +def test_the_emitted_string_table_holds_the_labels_byte_for_byte(tmp_path): + """Read back through the built library, which is what proves the escaping is right.""" + result = compile_onnx(_zipmap_model("classlabels_strings", _LABELS), tmp_path) + compiled = result.load() + (labels,) = result.report["class_labels"] + + library = ctypes.CDLL(str(compiled.library_path)) + table = (ctypes.c_char_p * len(_LABELS)).in_dll(library, labels["symbol"]) + + assert [entry.decode("utf-8") for entry in table] == _LABELS + + +@requires_c_compiler +def test_the_int64_table_holds_the_labels_the_model_declared(tmp_path): + values = [-1, 0, 2**62] + result = compile_onnx(_zipmap_model("classlabels_int64s", values), tmp_path) + compiled = result.load() + (labels,) = result.report["class_labels"] + + library = ctypes.CDLL(str(compiled.library_path)) + table = (ctypes.c_int64 * len(values)).in_dll(library, labels["symbol"]) + + assert list(table) == values + + +def test_a_zipmap_whose_result_is_not_a_graph_output_is_rejected(tmp_path): + model = _zipmap_model("classlabels_strings", _LABELS) + del model.graph.output[:] + model.graph.output.append(_untyped("probabilities")) + + with pytest.raises(CompileError, match="sequence of maps"): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + "attributes", + [{}, {"classlabels_strings": _LABELS, "classlabels_int64s": [1, 2, 3]}], + ids=["neither", "both"], +) +def test_a_zipmap_without_exactly_one_label_list_is_rejected(tmp_path, attributes): + model = _zipmap_model("classlabels_strings", _LABELS) + (zipmap,) = [node for node in model.graph.node if node.op_type == "ZipMap"] + del zipmap.attribute[:] + for name, values in attributes.items(): + zipmap.attribute.append(helper.make_attribute(name, values)) + + with pytest.raises(CompileError, match="classlabels_strings"): + compile_onnx(model, tmp_path) + + +# -------------------------------------------------------------------------------------- +# What the compiler refuses +# -------------------------------------------------------------------------------------- + + +def test_a_string_tensor_between_ml_ops_is_rejected(tmp_path): + """`CategoryMapper` maps to or from strings whichever way it is pointed.""" + model = _model( + [ + _ml_node( + "CategoryMapper", + ["x"], + ["y"], + cats_int64s=[1, 2], + cats_strings=["a", "b"], + ) + ], + [_tensor("x", TensorProto.INT64, [3])], + [_untyped("y")], + ) + + with pytest.raises(CompileError, match="STRING"): + compile_onnx(model, tmp_path) + + +def test_string_categories_against_a_numeric_input_are_rejected(tmp_path): + model = _model( + [_ml_node("OneHotEncoder", ["x"], ["y"], cats_strings=["a", "b"])], + [_tensor("x", TensorProto.FLOAT, [3])], + [_untyped("y")], + ) + + with pytest.raises(CompileError, match="cats_int64s"): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + ("attributes", "message"), + [ + ({"keys_strings": ["a", "b"], "values_int64s": [1, 2]}, "LabelEncoder"), + ({"keys_int64s": [1, 2], "values_strings": ["a", "b"]}, "STRING"), + ], + ids=["string_keys", "string_values"], +) +def test_a_label_encoder_over_strings_is_rejected(tmp_path, attributes, message): + """Either side of the mapping being strings puts it out of reach, for its own reason.""" + model = _model( + [_ml_node("LabelEncoder", ["x"], ["y"], **attributes)], + [_tensor("x", TensorProto.INT64, [3])], + [_untyped("y")], + ) + + with pytest.raises(CompileError, match=message): + compile_onnx(model, tmp_path) + + +def test_a_coefficient_list_that_fits_no_feature_axis_is_rejected(tmp_path): + model = _model( + [_ml_node("Scaler", ["x"], ["y"], offset=[1.0, 2.0], scale=[1.0, 2.0])], + [_tensor("x", TensorProto.FLOAT, [4, 3])], + [_untyped("y")], + ) + + with pytest.raises(CompileError, match="2 `offset` value"): + compile_onnx(model, tmp_path) + + +def test_a_norm_onnx_does_not_define_is_rejected(tmp_path): + model = _model( + [_ml_node("Normalizer", ["x"], ["y"], norm="L3")], + [_tensor("x", TensorProto.FLOAT, [4, 3])], + [_untyped("y")], + ) + + with pytest.raises(CompileError, match="L1, L2, MAX"): + compile_onnx(model, tmp_path) + + +def test_imputer_setting_both_value_families_is_rejected(tmp_path): + model = _model( + [ + _ml_node( + "Imputer", + ["x"], + ["y"], + imputed_value_floats=[1.0], + imputed_value_int64s=[1], + ) + ], + [_tensor("x", TensorProto.FLOAT, [4, 3])], + [_untyped("y")], + ) + + with pytest.raises(CompileError, match="exactly one of"): + compile_onnx(model, tmp_path) + + +def test_a_width_per_input_is_required(tmp_path): + model = _model( + [_ml_node("FeatureVectorizer", ["a", "b"], ["y"], inputdimensions=[2])], + [ + _tensor("a", TensorProto.FLOAT, [4, 3]), + _tensor("b", TensorProto.FLOAT, [4, 2]), + ], + [_untyped("y")], + ) + + with pytest.raises(CompileError, match="one width per input"): + compile_onnx(model, tmp_path) + + +def test_inputs_that_disagree_on_their_row_count_are_rejected(tmp_path): + model = _model( + [_ml_node("FeatureVectorizer", ["a", "b"], ["y"], inputdimensions=[3, 2])], + [ + _tensor("a", TensorProto.FLOAT, [4, 3]), + _tensor("b", TensorProto.FLOAT, [2, 2]), + ], + [_untyped("y")], + ) + + with pytest.raises(CompileError, match="first dimension"): + compile_onnx(model, tmp_path) + + +# -------------------------------------------------------------------------------------- +# Converted scikit-learn pipelines, where the backend corpus has nothing +# -------------------------------------------------------------------------------------- + +_SAMPLES = 24 +_FEATURES = 4 + + +def _converted(estimator, data): + """The ONNX a scikit-learn converter emits for `estimator`, at `data`'s exact shape. + + The shape is stated rather than taken from a sample row: a converter leaves the batch + dimension open otherwise, and a compiled artifact is specialized to one concrete shape. + """ + skl2onnx = pytest.importorskip("skl2onnx") + data_types = pytest.importorskip("skl2onnx.common.data_types") + tensor_type = ( + data_types.Int64TensorType + if data.dtype == np.int64 + else data_types.FloatTensorType + ) + return skl2onnx.convert_sklearn( + estimator.fit(data), initial_types=[("X", tensor_type(list(data.shape)))] + ) + + +def _fitted_data(seed=0): + generator = np.random.default_rng(seed) + return generator.normal(size=(_SAMPLES, _FEATURES)).astype(np.float32) + + +@requires_c_compiler +@pytest.mark.parametrize( + "name", + ["StandardScaler", "MinMaxScaler", "MaxAbsScaler", "RobustScaler", "Normalizer"], +) +def test_a_converted_scaler_matches_the_reference(tmp_path, name): + preprocessing = pytest.importorskip("sklearn.preprocessing") + data = _fitted_data() + + model = _converted(getattr(preprocessing, name)(), data) + + _matches_reference(model, {model.graph.input[0].name: data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_imputer_matches_the_reference(tmp_path): + impute = pytest.importorskip("sklearn.impute") + data = _fitted_data() + data[3, 1] = np.nan + data[7, 2] = np.nan + + model = _converted(impute.SimpleImputer(strategy="mean"), data) + + _matches_reference(model, {model.graph.input[0].name: data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_binarizer_matches_the_reference(tmp_path): + preprocessing = pytest.importorskip("sklearn.preprocessing") + data = _fitted_data() + + model = _converted(preprocessing.Binarizer(threshold=0.25), data) + + _matches_reference(model, {model.graph.input[0].name: data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_preprocessing_pipeline_matches_the_reference(tmp_path): + pipeline = pytest.importorskip("sklearn.pipeline") + preprocessing = pytest.importorskip("sklearn.preprocessing") + impute = pytest.importorskip("sklearn.impute") + data = _fitted_data(seed=1) + data[2, 0] = np.nan + + model = _converted( + pipeline.make_pipeline( + impute.SimpleImputer(strategy="median"), + preprocessing.StandardScaler(), + preprocessing.Normalizer(norm="l2"), + ), + data, + ) + + _matches_reference(model, {model.graph.input[0].name: data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_one_hot_encoder_matches_the_reference(tmp_path): + preprocessing = pytest.importorskip("sklearn.preprocessing") + generator = np.random.default_rng(2) + data = generator.integers(0, 3, size=(_SAMPLES, 2)).astype(np.int64) + + model = _converted(preprocessing.OneHotEncoder(sparse_output=False), data) + + _matches_reference(model, {model.graph.input[0].name: data}, tmp_path) diff --git a/src/python/tests/test_extra_compiler_runtime_dims.py b/src/python/tests/test_extra_compiler_runtime_dims.py new file mode 100644 index 0000000..beb57ba --- /dev/null +++ b/src/python/tests/test_extra_compiler_runtime_dims.py @@ -0,0 +1,750 @@ +"""Bounded runtime dimensions: one artifact for a whole family of sizes. + +The oracles are the same two the rest of the compiler suite uses — the FNNX `Runtime` +(onnxruntime) for bundles and `onnx.reference.ReferenceEvaluator` for single-node models — +run at each size the compiled artifact is then executed at. Nothing here states an expected +output of its own; what *is* asserted directly is the artifact's own contract: buffers sized +for the maximum, a status code for a size outside the range, and a compile error wherever +the dimension cannot be tracked. +""" + +from __future__ import annotations + +import ctypes +import re +import shutil +import subprocess +from pathlib import Path + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError, HarnessError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +pytest.importorskip("onnxruntime") +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") +specialize = pytest.importorskip("fnnx.extras.compilers.c.onnx.specialize") + +from fnnx.extras.compilers.c import compile_bundle, compile_onnx # noqa: E402 +from fnnx.extras.compilers.c.__main__ import main as cli_main # noqa: E402 +from fnnx.runtime import Runtime # noqa: E402 +from onnx import TensorProto, helper, numpy_helper # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +OPSET = 21 +SEED = 20260727 +MAX_BATCH = 8 +STRICT_FLAGS = ("-std=c99", "-Wall", "-Wextra", "-Werror", "-Werror=vla") +C_COMPILERS = [name for name in ("gcc", "clang") if shutil.which(name)] +ALLOCATION_TOKENS = ("malloc", "calloc", "realloc", "free", "alloca") + +PIPELINE_BUNDLE = Path(__file__).parent / "models" / "onnx_pipeline.fnnx" + +pytestmark = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + + +# -------------------------------------------------------------------------------------- +# Fixtures +# -------------------------------------------------------------------------------------- + + +def _values(shape, *, seed: int = SEED): + return np.random.default_rng(seed).normal(size=shape).astype(np.float32) + + +def _tensor(name, shape, elem_type=TensorProto.FLOAT): + return helper.make_tensor_value_info(name, elem_type, list(shape)) + + +def _model(nodes, inputs, outputs, *, initializer=(), opset=OPSET, name="graph"): + return helper.make_model( + helper.make_graph(nodes, name, list(inputs), list(outputs), list(initializer)), + opset_imports=[helper.make_opsetid("", opset)], + ) + + +def _affine_model(): + """`relu(x @ W + b)` over a symbolic batch: a matmul, a broadcast add, an activation.""" + weight = _values((3, 4), seed=SEED + 1) + bias = _values((1, 4), seed=SEED + 2) + return _model( + [ + helper.make_node("MatMul", ["x", "w"], ["scored"], name="matmul"), + helper.make_node("Add", ["scored", "b"], ["biased"], name="add"), + helper.make_node("Relu", ["biased"], ["y"], name="relu"), + ], + [_tensor("x", ["batch", 3])], + [_tensor("y", [None, 4])], + initializer=[ + numpy_helper.from_array(weight, "w"), + numpy_helper.from_array(bias, "b"), + ], + ) + + +def _pipeline_outputs(feeds): + return Runtime(str(PIPELINE_BUNDLE)).compute(dict(feeds), {}) + + +@pytest.fixture(scope="module") +def pipeline_artifact(tmp_path_factory): + directory = tmp_path_factory.mktemp("runtime_dim_pipeline") + return compile_bundle( + PIPELINE_BUNDLE, directory, runtime_dims={"batch": MAX_BATCH}, prefix="rtb" + ) + + +# -------------------------------------------------------------------------------------- +# Scenario: Bounded runtime batch dimension +# -------------------------------------------------------------------------------------- + + +def test_a_pipeline_compiled_at_a_maximum_runs_at_every_smaller_batch( + pipeline_artifact, +): + model = pipeline_artifact.load() + + for batch in (1, 5, MAX_BATCH): + inputs = {"x": _values((batch, 3), seed=SEED + batch)} + computed = model.run(inputs)["y4"] + expected = np.asarray(_pipeline_outputs(inputs)["y4"]) + assert computed.shape == expected.shape + np.testing.assert_allclose(computed, expected, rtol=1e-5, atol=1e-5) + + +def test_buffers_and_the_arena_are_sized_for_the_maximum(pipeline_artifact, tmp_path): + at_max = compile_bundle( + PIPELINE_BUNDLE, tmp_path / "pinned", dim_bindings={"batch": MAX_BATCH} + ) + + assert pipeline_artifact.report["memory"] == at_max.report["memory"] + for tensor in pipeline_artifact.report["entrypoint"]["inputs"]: + assert tensor["shape"][0] == MAX_BATCH + + +def test_the_header_publishes_the_maximum_of_each_runtime_dimension(pipeline_artifact): + header = pipeline_artifact.header_path.read_text(encoding="utf-8") + + assert f"#define RTB_DIM_BATCH_MAX {MAX_BATCH}" in header + assert pipeline_artifact.report["runtime_dims"] == [ + { + "name": "batch", + "max": MAX_BATCH, + "parameter": "dim_batch", + "macro": "RTB_DIM_BATCH_MAX", + } + ] + + +def test_entrypoints_take_the_dimension_value_ahead_of_their_buffers(pipeline_artifact): + header = pipeline_artifact.header_path.read_text(encoding="utf-8") + + assert "int rtb_run(int32_t dim_batch, const float* x, float* y4);" in header + assert ( + "int rtb_node_linreg_run(int32_t dim_batch, const float* float_input, " + "float* variable);" in header + ) + + +def test_the_report_describes_which_axes_scale_with_which_dimension(pipeline_artifact): + entry = pipeline_artifact.report["entrypoint"] + + assert entry["inputs"][0]["runtime_shape"] == [ + {"dim": "batch", "coefficient": 1}, + 3, + ] + assert entry["outputs"][0]["runtime_shape"] == [ + {"dim": "batch", "coefficient": 1}, + 1, + ] + + +def test_a_node_entrypoint_runs_at_a_smaller_batch_than_the_maximum(pipeline_artifact): + model = pipeline_artifact.load() + inputs = {"x": _values((3, 3))} + + inside = np.asarray(_pipeline_outputs(inputs)["y4"]) + alone = model.run_node("linreg", {"float_input": inputs["x"]})["variable"] + + assert alone.shape == (3, 1) + # `linreg` is one of the three regressors the pipeline sums, so its own output is not + # the pipeline's; what is asserted is that the node entry works at a partial batch and + # that the whole pipeline agrees with the runtime at the same one. + np.testing.assert_allclose(model.run(inputs)["y4"], inside, rtol=1e-5, atol=1e-5) + + +def test_a_call_below_the_maximum_touches_only_the_size_it_asked_for(pipeline_artifact): + """Compute scales with the value passed, not with the capacity the buffers hold. + + Handed a buffer sized for the maximum, the artifact must write the leading rows of the + shape this call's dimension value gives and leave the rest of the capacity alone — the + same property that keeps it inside a caller's exactly-sized buffer. + """ + model = pipeline_artifact.load() + library = ctypes.CDLL(str(model.library_path)) + entry = library.rtb_run + entry.restype = ctypes.c_int + entry.argtypes = [ctypes.c_int32, ctypes.c_void_p, ctypes.c_void_p] + + values = _values((MAX_BATCH, 3)) + outputs = np.full((MAX_BATCH, 1), 1234.5, dtype=np.float32) + assert entry(2, values.ctypes.data, outputs.ctypes.data) == 0 + + expected = np.asarray(_pipeline_outputs({"x": values[:2]})["y4"]) + np.testing.assert_allclose(outputs[:2], expected, rtol=1e-5, atol=1e-5) + assert np.all(outputs[2:] == np.float32(1234.5)) + + +# -------------------------------------------------------------------------------------- +# Scenario: Out-of-range runtime dim value +# -------------------------------------------------------------------------------------- + + +def test_the_c_entrypoint_rejects_a_dimension_outside_the_compiled_range( + pipeline_artifact, +): + """The artifact's own contract, checked below the harness that also validates.""" + model = pipeline_artifact.load() + library = ctypes.CDLL(str(model.library_path)) + entry = library.rtb_run + entry.restype = ctypes.c_int + entry.argtypes = [ctypes.c_int32, ctypes.c_void_p, ctypes.c_void_p] + + values = _values((MAX_BATCH, 3)) + for batch in (0, MAX_BATCH + 1, -1): + outputs = np.full((MAX_BATCH, 1), 1234.5, dtype=np.float32) + status = entry(batch, values.ctypes.data, outputs.ctypes.data) + assert status != 0 + assert np.all(outputs == np.float32(1234.5)) + + +def test_the_harness_rejects_a_dimension_outside_the_compiled_range(pipeline_artifact): + model = pipeline_artifact.load() + + with pytest.raises(HarnessError, match=r"`batch` is 9, outside the \[1, 8\]"): + model.run({"x": _values((1, 3))}, dims={"batch": 9}) + with pytest.raises(HarnessError, match=r"`batch` is 0, outside the \[1, 8\]"): + model.run({"x": _values((1, 3))}, dims={"batch": 0}) + + +def test_the_harness_rejects_an_input_that_disagrees_with_the_stated_dimension( + pipeline_artifact, +): + model = pipeline_artifact.load() + + with pytest.raises(HarnessError, match=r"input `x` has shape \(5, 3\)"): + model.run({"x": _values((5, 3))}, dims={"batch": 4}) + + +def test_the_harness_rejects_a_dimension_the_artifact_does_not_have(pipeline_artifact): + model = pipeline_artifact.load() + + with pytest.raises(HarnessError, match="`width` is not a runtime dimension"): + model.run({"x": _values((2, 3))}, dims={"width": 2}) + + +# -------------------------------------------------------------------------------------- +# Scenario: Runtime dim that cannot be tracked +# -------------------------------------------------------------------------------------- + + +def test_a_reshape_folding_the_batch_into_a_fixed_dim_is_rejected(tmp_path): + model = _model( + [helper.make_node("Reshape", ["x", "shape"], ["y"], name="fold")], + [_tensor("x", ["batch", 4])], + [_tensor("y", [MAX_BATCH, None])], + initializer=[ + numpy_helper.from_array(np.array([MAX_BATCH, -1], dtype=np.int64), "shape") + ], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path, runtime_dims={"batch": MAX_BATCH}) + + assert "fold" in str(error.value) + assert "pin them via `dim_bindings`" in str(error.value) + assert not list(tmp_path.iterdir()) + + +def test_a_concat_along_the_dimension_with_a_constant_is_rejected(tmp_path): + model = _model( + [helper.make_node("Concat", ["x", "pad"], ["y"], axis=0, name="join")], + [_tensor("x", ["batch", 4])], + [_tensor("y", [None, 4])], + initializer=[numpy_helper.from_array(np.zeros((2, 4), np.float32), "pad")], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path, runtime_dims={"batch": MAX_BATCH}) + + assert "Node `join`" in str(error.value) + assert "constant multiple of a runtime dimension" in str(error.value) + assert "pin them via `dim_bindings`" in str(error.value) + + +def test_a_broadcast_the_dimension_cannot_be_proven_against_is_rejected(tmp_path): + """`[batch, 4] + [3, 4]` broadcasts at batch 3 and at nothing else in the range.""" + model = _model( + [helper.make_node("Add", ["x", "other"], ["y"], name="combine")], + [_tensor("x", ["batch", 4])], + [_tensor("y", [None, 4])], + initializer=[numpy_helper.from_array(np.ones((3, 4), np.float32), "other")], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path, runtime_dims={"batch": MAX_BATCH}) + + assert "combine" in str(error.value) + assert "pin them via `dim_bindings`" in str(error.value) + + +def test_a_slice_clamping_the_dimension_is_rejected(tmp_path): + """`x[:3]` is the batch below 3 and 3 above it — correct code for neither family.""" + model = _model( + [helper.make_node("Slice", ["x", "s", "e", "a"], ["y"], name="cut")], + [_tensor("x", ["batch", 4])], + [_tensor("y", [None, 4])], + initializer=[ + numpy_helper.from_array(np.array([0], np.int64), "s"), + numpy_helper.from_array(np.array([3], np.int64), "e"), + numpy_helper.from_array(np.array([0], np.int64), "a"), + ], + ) + + with pytest.raises(CompileError, match="pin them via `dim_bindings`"): + compile_onnx(model, tmp_path, runtime_dims={"batch": MAX_BATCH}) + + +def test_a_size_that_grows_faster_than_the_dimension_is_rejected(tmp_path): + """A `[batch, batch]` intermediate: its element count is the dimension squared.""" + model = _model( + [ + helper.make_node("MatMul", ["x", "z"], ["t"], name="matmul"), + helper.make_node("Relu", ["t"], ["y"], name="relu"), + ], + [_tensor("x", ["batch", 4]), _tensor("z", [4, "batch"])], + [_tensor("y", [None, None])], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path, runtime_dims={"batch": MAX_BATCH}) + + assert "Node `relu`" in str(error.value) + assert "does not scale linearly" in str(error.value) + + +@pytest.mark.parametrize("maximum", [3, 4, MAX_BATCH]) +def test_a_size_quadratic_in_the_dimension_is_rejected_at_every_maximum( + tmp_path, maximum +): + """A `[batch+1, batch+1]` intermediate, whose element count is the dimension squared. + + Two sizes are exactly what an affine reading of a literal has free parameters, so a + schedule offering only two would fit this count and lift it wrongly — code correct at + the sizes probed and silently wrong at the rest of the range. + """ + model = _model( + [ + helper.make_node("Concat", ["x", "row"], ["wide"], axis=0), + helper.make_node("Concat", ["z", "column"], ["tall"], axis=1), + helper.make_node("MatMul", ["wide", "tall"], ["square"], name="matmul"), + helper.make_node("ReduceSum", ["square"], ["y"], keepdims=0, name="total"), + ], + [_tensor("x", ["batch", 4]), _tensor("z", [4, "batch"])], + [_tensor("y", [])], + initializer=[ + numpy_helper.from_array(np.ones((1, 4), np.float32), "row"), + numpy_helper.from_array(np.ones((4, 1), np.float32), "column"), + ], + ) + + with pytest.raises(CompileError, match="does not scale linearly"): + compile_onnx(model, tmp_path, runtime_dims={"batch": maximum}) + + +def test_a_quadratic_size_is_rejected_where_only_its_own_dimension_is_small(tmp_path): + """The same count, in a dimension with two sizes above 1 beside one with four. + + Whether the emitted code may be read off the sizes above 1 is a question per dimension: + `rows` has enough of them to check a reading against, `cols` does not, so `cols` has to + be read at 1 as well however much room `rows` has to spare. + """ + model = _model( + [ + helper.make_node("Concat", ["w", "row"], ["wide"], axis=0), + helper.make_node("Concat", ["z", "column"], ["tall"], axis=1), + helper.make_node("MatMul", ["wide", "tall"], ["square"], name="matmul"), + helper.make_node("ReduceSum", ["square"], ["total"], keepdims=0), + helper.make_node("Mul", ["x", "total"], ["y"], name="scale"), + ], + [ + _tensor("x", ["rows", 4]), + _tensor("z", [4, "cols"]), + _tensor("w", ["cols", 4]), + ], + [_tensor("y", [None, 4])], + initializer=[ + numpy_helper.from_array(np.ones((1, 4), np.float32), "row"), + numpy_helper.from_array(np.ones((4, 1), np.float32), "column"), + ], + ) + + with pytest.raises(CompileError, match="does not scale linearly"): + compile_onnx(model, tmp_path, runtime_dims={"rows": 6, "cols": 3}) + + +def test_a_value_folded_from_the_dimension_is_rejected(tmp_path): + """`Shape(x)` reaching the data makes a weight the artifact would have to recompute.""" + model = _model( + [ + helper.make_node("Shape", ["x"], ["extent"], name="shape"), + helper.make_node("Cast", ["extent"], ["sized"], to=TensorProto.FLOAT), + helper.make_node("ReduceSum", ["sized"], ["total"], keepdims=0), + helper.make_node("Mul", ["x", "total"], ["y"], name="scale"), + ], + [_tensor("x", ["batch", 4])], + [_tensor("y", [None, 4])], + ) + + with pytest.raises(CompileError, match="pin them via `dim_bindings`"): + compile_onnx(model, tmp_path, runtime_dims={"batch": MAX_BATCH}) + + +def test_a_folded_value_that_only_moves_at_size_one_is_rejected(tmp_path): + """`min(batch, 2)` folded into a weight: the same at every size in range but 1. + + Constant data is what a fold took out of the graph, so it is read at every probe rather + than only where the emitted code is — a clamp is flat across most of the range and the + size it bends at is the one the code comparison leaves out. + """ + model = _model( + [ + helper.make_node("Shape", ["x"], ["extent"], name="shape"), + helper.make_node("Gather", ["extent", "first"], ["rows"], axis=0), + helper.make_node("Min", ["rows", "ceiling"], ["clamped"], name="clamp"), + helper.make_node("Cast", ["clamped"], ["scale"], to=TensorProto.FLOAT), + helper.make_node("Mul", ["x", "scale"], ["y"], name="apply"), + ], + [_tensor("x", ["batch", 4])], + [_tensor("y", [None, 4])], + initializer=[ + numpy_helper.from_array(np.array([0], np.int64), "first"), + numpy_helper.from_array(np.array([2], np.int64), "ceiling"), + ], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path, runtime_dims={"batch": MAX_BATCH}) + + assert "`clamped`" in str(error.value) + assert "pin them via `dim_bindings`" in str(error.value) + + +# -------------------------------------------------------------------------------------- +# Option validation +# -------------------------------------------------------------------------------------- + + +def test_a_dimension_cannot_be_both_bound_and_runtime(tmp_path): + with pytest.raises(CompileError, match="both bound to 4 and declared runtime"): + compile_bundle( + PIPELINE_BUNDLE, + tmp_path, + dim_bindings={"batch": 4}, + runtime_dims={"batch": MAX_BATCH}, + ) + assert not tmp_path.exists() or not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("maximum", [0, -3, True, 2.5]) +def test_a_runtime_dimension_needs_a_positive_integer_maximum(tmp_path, maximum): + with pytest.raises(CompileError, match="needs a maximum of at least 1"): + compile_bundle(PIPELINE_BUNDLE, tmp_path, runtime_dims={"batch": maximum}) + + +def test_two_runtime_dimensions_may_not_share_a_c_identifier(tmp_path): + with pytest.raises(CompileError, match="sanitize to the C identifier"): + compile_bundle(PIPELINE_BUNDLE, tmp_path, runtime_dims={"a.b": 4, "a-b": 4}) + + +def test_a_maximum_of_one_leaves_the_artifact_fully_static(tmp_path): + """The only size in range is 1, so every extent is a constant the buffers are cut to.""" + pinned = compile_bundle(PIPELINE_BUNDLE, tmp_path / "pinned", dim_bindings={}) + result = compile_bundle( + PIPELINE_BUNDLE, tmp_path / "one", runtime_dims={"batch": 1} + ) + compiled = result.load() + inputs = {"x": _values((1, 3))} + + assert result.report["memory"] == pinned.report["memory"] + np.testing.assert_allclose( + compiled.run(inputs, dims={"batch": 1})["y4"], + np.asarray(_pipeline_outputs(inputs)["y4"]), + rtol=1e-5, + atol=1e-5, + ) + + +@pytest.mark.parametrize("maximum", [2, 3]) +def test_a_small_maximum_says_so_when_the_code_cannot_be_read(tmp_path, maximum): + """Size 1 is then one of the sizes the emitted code is read at, and it reads apart.""" + with pytest.raises(CompileError) as error: + compile_bundle(PIPELINE_BUNDLE, tmp_path, runtime_dims={"batch": maximum}) + + assert "A maximum below 4 leaves too few sizes above 1" in str(error.value) + + +# -------------------------------------------------------------------------------------- +# Against the reference evaluator, at several sizes +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("batch", [1, 2, 5, MAX_BATCH]) +def test_a_compiled_graph_matches_the_reference_evaluator_at_every_size( + tmp_path, batch +): + model = _affine_model() + compiled = compile_onnx( + model, tmp_path, runtime_dims={"batch": MAX_BATCH}, prefix="affine" + ).load() + values = _values((batch, 3), seed=SEED + batch) + + computed = compiled.run(x=values)["y"] + expected = ReferenceEvaluator(model).run(None, {"x": values})[0] + + assert computed.shape == expected.shape + np.testing.assert_allclose(computed, expected, rtol=1e-5, atol=1e-6) + + +def test_repeated_calls_at_different_sizes_carry_no_state(tmp_path): + model = _affine_model() + compiled = compile_onnx( + model, tmp_path, runtime_dims={"batch": MAX_BATCH}, prefix="affine" + ).load() + evaluator = ReferenceEvaluator(model) + + sizes = (MAX_BATCH, 1, 3, MAX_BATCH, 2) + for batch in sizes: + values = _values((batch, 3), seed=SEED + batch) + np.testing.assert_allclose( + compiled.run(x=values)["y"], + evaluator.run(None, {"x": values})[0], + rtol=1e-5, + atol=1e-6, + ) + + +def test_an_axis_that_is_a_multiple_of_the_dimension_is_tracked(tmp_path): + """`Concat` along the batch of two batch-sized tensors: an axis of `2 * batch`.""" + model = _model( + [helper.make_node("Concat", ["x", "x"], ["y"], axis=0, name="stack")], + [_tensor("x", ["batch", 3])], + [_tensor("y", [None, 3])], + ) + result = compile_onnx( + model, tmp_path, runtime_dims={"batch": MAX_BATCH}, prefix="stack" + ) + + assert result.report["entrypoint"]["outputs"][0]["runtime_shape"] == [ + {"dim": "batch", "coefficient": 2}, + 3, + ] + compiled = result.load() + for batch in (1, 3, MAX_BATCH): + values = _values((batch, 3), seed=SEED + batch) + np.testing.assert_array_equal( + compiled.run(x=values)["y"], + ReferenceEvaluator(model).run(None, {"x": values})[0], + ) + + +def test_two_runtime_dimensions_are_tracked_independently(tmp_path): + model = _model( + [helper.make_node("MatMul", ["x", "z"], ["y"], name="matmul")], + [_tensor("x", ["rows", 4]), _tensor("z", [4, "cols"])], + [_tensor("y", [None, None])], + ) + compiled = compile_onnx( + model, tmp_path, runtime_dims={"rows": 6, "cols": 5}, prefix="pair" + ).load() + evaluator = ReferenceEvaluator(model) + + for rows, columns in ((1, 1), (6, 5), (4, 2), (2, 5)): + feeds = { + "x": _values((rows, 4), seed=SEED + rows), + "z": _values((4, columns), seed=SEED + columns), + } + np.testing.assert_allclose( + compiled.run(feeds)["y"], + evaluator.run(None, feeds)[0], + rtol=1e-5, + atol=1e-6, + ) + + +def test_a_size_that_is_the_product_of_two_dimensions_is_rejected(tmp_path): + """An elementwise loop over `[rows, cols]` counts `rows * cols` — linear in neither. + + The per-dimension probes move one dimension at a time, and a product agrees with a + linear reading along each of them on its own; only the probe that moves both at once + tells them apart. + """ + model = _model( + [helper.make_node("Add", ["x", "x"], ["y"], name="twice")], + [_tensor("x", ["rows", "cols"])], + [_tensor("y", [None, None])], + ) + + with pytest.raises(CompileError) as error: + compile_onnx(model, tmp_path, runtime_dims={"rows": 6, "cols": 5}) + + assert "Node `twice`" in str(error.value) + assert "does not scale linearly" in str(error.value) + + +def test_a_dimension_no_input_carries_has_to_be_passed_explicitly(tmp_path): + """`ConstantOfShape` has nothing to read the size off, so the caller states it.""" + model = _model( + [helper.make_node("Add", ["x", "x"], ["y"], name="twice")], + [_tensor("x", [2, 3])], + [_tensor("y", [2, 3])], + ) + compiled = compile_onnx( + model, tmp_path, runtime_dims={"unused": 4}, prefix="lonely" + ).load() + values = _values((2, 3)) + + with pytest.raises(HarnessError, match="has to be passed as"): + compiled.run(x=values) + np.testing.assert_array_equal( + compiled.run({"x": values}, dims={"unused": 3})["y"], values + values + ) + + +# -------------------------------------------------------------------------------------- +# The artifact's build contract +# -------------------------------------------------------------------------------------- + + +@pytest.mark.skipif(not C_COMPILERS, reason="no system C compiler available") +@pytest.mark.parametrize("compiler", C_COMPILERS) +def test_the_artifact_builds_under_the_strict_flags( + pipeline_artifact, tmp_path, compiler +): + unit = tmp_path / "implementation.c" + unit.write_text( + f"#define {pipeline_artifact.report['prefix'].upper()}_IMPLEMENTATION\n" + f'#include "{pipeline_artifact.report["header"]}"\n', + encoding="utf-8", + ) + build = subprocess.run( + [ + compiler, + *STRICT_FLAGS, + "-c", + f"-I{pipeline_artifact.header_path.parent}", + str(unit), + "-o", + str(tmp_path / "artifact.o"), + ], + capture_output=True, + text=True, + ) + + assert build.returncode == 0, build.stderr + + +def test_the_artifact_allocates_nothing_and_compiles_deterministically(tmp_path): + first = compile_bundle( + PIPELINE_BUNDLE, tmp_path / "first", runtime_dims={"batch": MAX_BATCH} + ) + second = compile_bundle( + PIPELINE_BUNDLE, tmp_path / "second", runtime_dims={"batch": MAX_BATCH} + ) + + source = first.header_path.read_text(encoding="utf-8") + for token in ALLOCATION_TOKENS: + assert not re.search(rf"\b{token}\b", source), token + assert first.header_path.read_bytes() == second.header_path.read_bytes() + assert first.report_path.read_bytes() == second.report_path.read_bytes() + + +def test_compiling_without_runtime_dimensions_is_unchanged(tmp_path): + """The default contract stays exactly what it was: no parameter, no macro, no field.""" + result = compile_bundle(PIPELINE_BUNDLE, tmp_path, prefix="plain") + header = result.header_path.read_text(encoding="utf-8") + + assert "int plain_run(const float* x, float* y4);" in header + assert "int32_t" not in header + assert "_DIM_BATCH_MAX" not in header + assert result.report["runtime_dims"] == [] + assert "runtime_shape" not in result.report["entrypoint"]["inputs"][0] + + +# -------------------------------------------------------------------------------------- +# CLI +# -------------------------------------------------------------------------------------- + + +def test_the_cli_compiles_with_a_runtime_dimension(tmp_path, capsys): + status = cli_main( + [ + str(PIPELINE_BUNDLE), + "-o", + str(tmp_path), + "--runtime-dim", + f"batch={MAX_BATCH}", + "--prefix", + "cli", + ] + ) + output = capsys.readouterr().out + + assert status == 0 + assert f"batch<={MAX_BATCH}" in output + assert f"#define CLI_DIM_BATCH_MAX {MAX_BATCH}" in (tmp_path / "cli.h").read_text( + encoding="utf-8" + ) + + +def test_the_cli_rejects_a_dimension_that_is_both_bound_and_runtime(tmp_path, capsys): + status = cli_main( + [ + str(PIPELINE_BUNDLE), + "-o", + str(tmp_path), + "--dim", + "batch=2", + "--runtime-dim", + "batch=4", + ] + ) + + assert status == 1 + assert "declared runtime" in capsys.readouterr().err + + +# -------------------------------------------------------------------------------------- +# The probe schedule itself +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("maximum", "expected"), + [ + (1, (1,)), + (2, (1, 2)), + (3, (1, 2, 3)), + (4, (1, 2, 3, 4)), + (8, (1, 2, 3, 7, 8)), + (32, (1, 2, 3, 31, 32)), + ], +) +def test_the_probe_schedule_spans_both_ends_of_the_range(maximum, expected): + dim = specialize.RuntimeDim("batch", maximum, "batch") + + assert specialize.probe_values(dim) == expected diff --git a/src/python/tests/test_extra_compiler_svm.py b/src/python/tests/test_extra_compiler_svm.py new file mode 100644 index 0000000..8073ac0 --- /dev/null +++ b/src/python/tests/test_extra_compiler_svm.py @@ -0,0 +1,680 @@ +"""The ONNX-ML support vector machines and linear models: converted models, and refusals. + +What each op computes on a single node is settled by the differential sweep against the +reference evaluator; the corpus carries no node test for any of the four. Neither reaches +where these ops actually turn up — a scikit-learn converter emits `ai.onnx.ml` opset 1 and +wraps a multi-class `SVC` in a graph of a further thirty nodes — so the parity tests below run +those converted models against onnxruntime, the second oracle, independent of both the +compiler and the reference evaluator. Neither reaches the float edges either: all four sweep +finite operands, since a dot product's summation order is the reference's and not the +kernel's, and an op with no node test has no corpus to cover the edges in its place — so the +edge cases here feed them directly. The rest of the module covers the error contracts, which +no sweep asserts. +""" + +from __future__ import annotations + +import shutil + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +# The harness refuses to import without numpy, so this covers both dependencies. +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from onnx import TensorProto, helper # noqa: E402 +from onnx.reference import ReferenceEvaluator # noqa: E402 + +from fnnx.extras.compilers.c import compile_onnx # noqa: E402 + +ML_OPSET = 5 + +requires_c_compiler = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + +_SAMPLES = 32 +_FEATURES = 4 + +# Three support vectors over three features, and the gamma/coef0/degree triple the kernel +# functions read. +_SUPPORT_VECTORS = [1.0, 2.0, 3.0, 0.0, 0.0, 1.0, -1.0, 0.5, 2.0] +_KERNEL_PARAMS = [0.5, 1.0, 3.0] + + +def _model(nodes, inputs, outputs, ml=ML_OPSET): + """A model whose intermediates carry no `value_info`, as a converter's output does.""" + model = helper.make_model( + helper.make_graph(nodes, "predictors", list(inputs), list(outputs)), + opset_imports=[helper.make_opsetid("ai.onnx.ml", ml)], + ) + model.ir_version = 10 + return model + + +def _node(op_type, outputs, attributes, results=None): + """A single-node model over a `[4, 3]` input, with `attributes` overriding the defaults.""" + return _model( + [ + helper.make_node( + op_type, + ["X"], + list(outputs), + name="predictor", + domain="ai.onnx.ml", + **attributes, + ) + ], + [helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 3])], + results or [helper.make_empty_tensor_value_info(name) for name in outputs], + ) + + +def _linear_regressor(**attributes): + declared = {"coefficients": [1.0, 0.0, -1.0], "intercepts": [0.5], **attributes} + return _node("LinearRegressor", ["Y"], declared) + + +def _linear_classifier(results=None, **attributes): + declared = { + "coefficients": [1.0, 0.0, -1.0, -0.5, 0.5, 0.25], + "intercepts": [0.5, -0.25], + "classlabels_ints": [3, 7], + **attributes, + } + return _node("LinearClassifier", ["Y", "Z"], declared, results) + + +def _svm_regressor(**attributes): + declared = {"coefficients": [1.0, 0.0, -1.0], "rho": [0.25], **attributes} + return _node("SVMRegressor", ["Y"], declared) + + +def _dropping(model, *names): + """`model` with those attributes taken off its node, which is the only way to leave an + empty one: ONNX's own builder refuses to infer a list attribute's type from no values.""" + node = model.graph.node[0] + kept = [entry for entry in node.attribute if entry.name not in names] + del node.attribute[:] + node.attribute.extend(kept) + return model + + +def _svm_classifier(**attributes): + declared = { + "coefficients": [0.5, -0.25, 0.75], + "rho": [0.25], + "classlabels_ints": [3, 7], + "vectors_per_class": [2, 1], + "support_vectors": _SUPPORT_VECTORS, + "kernel_type": "RBF", + "kernel_params": _KERNEL_PARAMS, + **attributes, + } + return _node("SVMClassifier", ["Y", "Z"], declared) + + +# -------------------------------------------------------------------------------------- +# Converted scikit-learn models, against onnxruntime +# -------------------------------------------------------------------------------------- + + +def _converted(estimator, data, target, **options): + """The ONNX a scikit-learn converter emits for `estimator`, at `data`'s exact shape.""" + skl2onnx = pytest.importorskip("skl2onnx") + data_types = pytest.importorskip("skl2onnx.common.data_types") + return skl2onnx.convert_sklearn( + estimator.fit(data, target), + initial_types=[("X", data_types.FloatTensorType(list(data.shape)))], + **options, + ) + + +def _session(model): + runtime = pytest.importorskip("onnxruntime") + runtime.set_default_logger_severity(3) + return runtime.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + + +def _matches_onnxruntime(model, feeds, tmp_path): + """Every output of the compiled artifact against the one onnxruntime computes for it.""" + compiled = compile_onnx(model, tmp_path).load() + outputs = compiled.run(feeds) + expected = _session(model).run(None, dict(feeds)) + assert [spec.name for spec in compiled.outputs] == [ + entry.name for entry in model.graph.output + ] + for entry, want in zip(model.graph.output, expected): + got, want = outputs[entry.name], np.asarray(want) + assert got.dtype == want.dtype, entry.name + assert got.shape == want.shape, entry.name + if want.dtype.kind == "f": + np.testing.assert_allclose( + got, want, rtol=1e-5, atol=1e-6, err_msg=entry.name + ) + else: + np.testing.assert_array_equal(got, want, err_msg=entry.name) + + +def _fitted_data(seed=0): + generator = np.random.default_rng(seed) + return generator.normal(size=(_SAMPLES, _FEATURES)).astype(np.float32) + + +def _regression_target(data): + return (2 * data[:, 0] + data[:, 1] - data[:, 3]).astype(np.float64) + + +@requires_c_compiler +@pytest.mark.parametrize("kernel", ["linear", "poly", "rbf", "sigmoid"]) +def test_a_converted_support_vector_regressor_matches_onnxruntime(tmp_path, kernel): + """Every kernel function ONNX defines, as scikit-learn's own converter encodes it.""" + svm = pytest.importorskip("sklearn.svm") + data = _fitted_data() + + model = _converted(svm.SVR(kernel=kernel, degree=3), data, _regression_target(data)) + + assert "SVMRegressor" in [node.op_type for node in model.graph.node] + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_one_class_support_vector_machine_matches_onnxruntime(tmp_path): + """The novelty detector, whose converter reads the score's sign downstream of the op.""" + svm = pytest.importorskip("sklearn.svm") + data = _fitted_data(seed=1) + + model = _converted(svm.OneClassSVM(), data, None) + + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +@pytest.mark.parametrize( + ("module", "name", "arguments"), + [ + ("sklearn.linear_model", "LinearRegression", {}), + ("sklearn.linear_model", "Ridge", {"alpha": 0.5}), + ("sklearn.svm", "LinearSVR", {"max_iter": 2000}), + ], +) +def test_a_converted_linear_regressor_matches_onnxruntime( + tmp_path, module, name, arguments +): + estimators = pytest.importorskip(module) + data = _fitted_data(seed=2) + + model = _converted( + getattr(estimators, name)(**arguments), data, _regression_target(data) + ) + + assert [node.op_type for node in model.graph.node] == ["LinearRegressor"] + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_multi_target_linear_regressor_matches_onnxruntime(tmp_path): + """Several targets at once, which is what the coefficient matrix's rows are.""" + linear_model = pytest.importorskip("sklearn.linear_model") + data = _fitted_data(seed=3) + targets = np.stack( + [_regression_target(data), data[:, 2].astype(np.float64)], axis=1 + ) + + model = _converted(linear_model.Ridge(alpha=0.5), data, targets) + + (node,) = model.graph.node + assert {entry.name for entry in node.attribute} >= {"targets"} + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +@pytest.mark.parametrize("classes", [2, 3], ids=["binary", "multiclass"]) +def test_a_converted_logistic_regression_matches_onnxruntime(tmp_path, classes): + """`LinearClassifier` with the two transforms scikit-learn's converter emits.""" + linear_model = pytest.importorskip("sklearn.linear_model") + data = _fitted_data(seed=4) + labels = np.digitize(data[:, 0], np.linspace(-1, 1, classes - 1)).astype(np.int64) + + model = _converted( + linear_model.LogisticRegression(max_iter=500), + data, + labels, + options={"zipmap": False}, + ) + + (node,) = [ + entry for entry in model.graph.node if entry.op_type == "LinearClassifier" + ] + (transform,) = [ + entry.s for entry in node.attribute if entry.name == "post_transform" + ] + assert transform in (b"LOGISTIC", b"SOFTMAX") + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +@pytest.mark.parametrize("classes", [2, 3], ids=["binary", "multiclass"]) +def test_a_converted_support_vector_classifier_matches_onnxruntime(tmp_path, classes): + """The pairwise scheme, and the graph the converter wraps its votes in above two classes.""" + svm = pytest.importorskip("sklearn.svm") + data = _fitted_data(seed=5) + labels = np.digitize(data[:, 0], np.linspace(-1, 1, classes - 1)).astype(np.int64) + + model = _converted( + svm.SVC(kernel="rbf", random_state=0), + data, + labels, + options={"zipmap": False}, + ) + + (node,) = [entry for entry in model.graph.node if entry.op_type == "SVMClassifier"] + assert {entry.name for entry in node.attribute} >= {"vectors_per_class"} + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_probability_classifier_matches_onnxruntime(tmp_path): + """Platt scaling, which is the one thing `prob_a` and `prob_b` are there for.""" + svm = pytest.importorskip("sklearn.svm") + data = _fitted_data(seed=6) + labels = (data[:, 0] + data[:, 1] > 0).astype(np.int64) + + model = _converted( + svm.SVC(kernel="rbf", probability=True, random_state=0), + data, + labels, + options={"zipmap": False}, + ) + + (node,) = [entry for entry in model.graph.node if entry.op_type == "SVMClassifier"] + assert {entry.name for entry in node.attribute} >= {"prob_a", "prob_b"} + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_pipeline_of_scaler_and_classifier_matches_onnxruntime(tmp_path): + """The classifier downstream of a preprocessor, whose intermediates carry no type.""" + pipeline = pytest.importorskip("sklearn.pipeline") + preprocessing = pytest.importorskip("sklearn.preprocessing") + svm = pytest.importorskip("sklearn.svm") + data = _fitted_data(seed=7) + labels = (data[:, 1] > 0).astype(np.int64) + + model = _converted( + pipeline.make_pipeline( + preprocessing.StandardScaler(), svm.SVC(kernel="rbf", random_state=0) + ), + data, + labels, + options={"zipmap": False}, + ) + + assert [node.op_type for node in model.graph.node][:2] == [ + "Scaler", + "SVMClassifier", + ] + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_classifier_keeps_its_probabilities_through_the_zipmap_pass( + tmp_path, +): + """The converter's default output is a map, which the graph pass turns back into scores. + + onnxruntime is the oracle for the pairing itself: which label names which column is the + one thing a run of the graph with `ZipMap` already removed cannot show. + """ + svm = pytest.importorskip("sklearn.svm") + data = _fitted_data(seed=8) + labels = np.digitize(data[:, 0], [-0.5, 0.5]).astype(np.int64) + model = _converted(svm.SVC(kernel="rbf", random_state=0), data, labels) + + assert "ZipMap" in [node.op_type for node in model.graph.node] + result = compile_onnx(model, tmp_path) + outputs = result.load().run({"X": data}) + predicted, rows = _session(model).run(None, {"X": data}) + (table,) = result.report["class_labels"] + + assert table["dtype"] == "int64" + np.testing.assert_array_equal(outputs["output_label"], predicted) + np.testing.assert_allclose( + outputs[table["tensor"]], + np.array( + [[row[label] for label in table["values"]] for row in rows], np.float32 + ), + rtol=1e-5, + atol=1e-6, + ) + + +@requires_c_compiler +def test_two_linear_models_of_one_shape_share_a_single_scoring_kernel(tmp_path): + """The dot product is one shared kernel per element type, however many nodes run it.""" + model = _linear_regressor(coefficients=[2.0]) + model.graph.input[0].type.tensor_type.shape.dim[1].dim_value = 1 + second = onnx.NodeProto() + second.CopyFrom(model.graph.node[0]) + second.name = "second" + second.input[0] = "Y" + second.output[0] = "Y2" + model.graph.node.append(second) + del model.graph.output[:] + model.graph.output.append(helper.make_empty_tensor_value_info("Y2")) + + result = compile_onnx(model, tmp_path) + + assert [name for name in result.report["kernels"] if "scores" in name] == [ + "predictors_ml_scores_float" + ] + + +# -------------------------------------------------------------------------------------- +# The float edges, against the reference evaluator +# -------------------------------------------------------------------------------------- + +_INFO = np.finfo(np.float32) +_EDGE_VALUES = ( + 0.0, + -0.0, + np.inf, + -np.inf, + np.nan, + _INFO.max, + -_INFO.max, + _INFO.tiny, + _INFO.smallest_subnormal, +) + + +def _edge_rows(features=3): + """One special value per row, in a column that rotates, and zero everywhere else. + + A row holding a single value is what makes its score independent of the order the + products of that row are summed in — the one thing about these ops the reference cannot + be held to, and the reason their sweeps run on finite operands. Everything else about an + edge value is the arithmetic itself: what an infinity does to a kernel function, what a + value that is not a number does to the winning column, and which side of a threshold an + overflow lands on. + """ + values = np.array(_EDGE_VALUES, np.float32) + rows = np.zeros((len(values), features), np.float32) + rows[np.arange(len(values)), np.arange(len(values)) % features] = values + return rows + + +def _with_rows(model, rows): + model.graph.input[0].type.tensor_type.shape.dim[0].dim_value = rows + return model + + +def _svm_classifier_over_classes(): + """The support vector classifier's other mode: one score per class, no support vectors.""" + return _dropping( + _svm_classifier(coefficients=[1.0, 0.0, -1.0, -0.5, 0.5, 0.25]), + "vectors_per_class", + "support_vectors", + ) + + +@requires_c_compiler +@pytest.mark.parametrize( + ("label", "builder"), + [ + ("linear_regressor", _linear_regressor), + ("linear_classifier", _linear_classifier), + ( + "linear_classifier_paired", + lambda: _linear_classifier( + coefficients=[1.0, 0.0, -1.0], + intercepts=[0.5], + post_transform="LOGISTIC", + ), + ), + ("svm_regressor", _svm_regressor), + *( + ( + f"svm_regressor_{kernel.lower()}", + lambda kernel=kernel: _svm_regressor( + n_supports=3, + support_vectors=_SUPPORT_VECTORS, + kernel_params=_KERNEL_PARAMS, + kernel_type=kernel, + ), + ) + for kernel in ("LINEAR", "POLY", "RBF", "SIGMOID") + ), + ("svm_classifier", _svm_classifier), + ("svm_classifier_linear", _svm_classifier_over_classes), + ( + "svm_classifier_probabilities", + lambda: _svm_classifier(prob_a=[-1.5], prob_b=[0.25]), + ), + ], +) +def test_a_predictor_matches_the_reference_on_the_float_edges(tmp_path, label, builder): + data = _edge_rows() + model = _with_rows(builder(), len(data)) + + outputs = compile_onnx(model, tmp_path).load().run({"X": data}) + with np.errstate(all="ignore"): + expected = ReferenceEvaluator(model).run(None, {"X": data}) + + for entry, want in zip(model.graph.output, expected): + got, want = outputs[entry.name], np.asarray(want) + assert got.dtype == want.dtype, entry.name + assert got.shape == want.shape, entry.name + if want.dtype.kind == "f": + np.testing.assert_allclose( + got, want, rtol=1e-5, atol=1e-6, equal_nan=True, err_msg=entry.name + ) + else: + np.testing.assert_array_equal(got, want, err_msg=entry.name) + + +# -------------------------------------------------------------------------------------- +# What the compiler refuses +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("builder", "attributes"), + [ + (_linear_regressor, {}), + (_linear_classifier, {}), + (_svm_regressor, {}), + (_svm_classifier, {}), + ], + ids=["linear_regressor", "linear_classifier", "svm_regressor", "svm_classifier"], +) +@pytest.mark.parametrize("shape", [[3], [2, 2, 3]], ids=["vector", "rank_3"]) +def test_an_input_that_is_not_a_matrix_is_rejected( + tmp_path, builder, attributes, shape +): + """These ops read `[N, F]`; every other rank is a model the reference cannot run either.""" + model = builder(**attributes) + del model.graph.input[0].type.tensor_type.shape.dim[:] + for extent in shape: + model.graph.input[0].type.tensor_type.shape.dim.add().dim_value = extent + + with pytest.raises(CompileError): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + "builder", [_linear_regressor, _linear_classifier], ids=["regressor", "classifier"] +) +def test_a_linear_model_without_intercepts_is_rejected(tmp_path, builder): + """The one attribute the reference and onnxruntime read differently when it is absent.""" + with pytest.raises(CompileError, match="sets no `intercepts`"): + compile_onnx(_dropping(builder(), "intercepts"), tmp_path) + + +def test_a_linear_model_with_an_intercept_per_nothing_is_rejected(tmp_path): + with pytest.raises(CompileError, match="intercept"): + compile_onnx( + _linear_regressor( + coefficients=[1.0, 0.0, -1.0, 0.5, 0.5, 0.5], + intercepts=[0.5, 0.25, 1.0], + targets=2, + ), + tmp_path, + ) + + +def test_a_coefficient_per_nothing_is_rejected(tmp_path): + with pytest.raises(CompileError, match="coefficient"): + compile_onnx(_linear_regressor(coefficients=[1.0, 0.0]), tmp_path) + + +def test_a_linear_regressor_scoring_no_targets_is_rejected(tmp_path): + with pytest.raises(CompileError, match="target"): + compile_onnx(_linear_regressor(targets=0), tmp_path) + + +def test_a_linear_classifier_with_fewer_labels_than_columns_is_rejected(tmp_path): + """A declared result wide enough to hide the disagreement is still refused.""" + model = _linear_classifier( + results=[ + helper.make_tensor_value_info("Y", TensorProto.INT64, [4]), + helper.make_tensor_value_info("Z", TensorProto.FLOAT, [4, 3]), + ], + coefficients=[1.0] * 9, + intercepts=[0.5, 0.0, -0.5], + ) + + with pytest.raises(CompileError, match="class label"): + compile_onnx(model, tmp_path) + + +@pytest.mark.parametrize( + "builder", [_linear_classifier, _svm_classifier], ids=["linear", "svm"] +) +def test_a_string_labelled_classifier_is_rejected(tmp_path, builder): + model = _dropping(builder(classlabels_strings=["a", "b"]), "classlabels_ints") + + with pytest.raises(CompileError, match="STRING"): + compile_onnx(model, tmp_path) + + +def test_a_support_vector_classifier_without_class_labels_is_rejected(tmp_path): + """A declared result, since with no labels at all nothing derives one to begin with.""" + model = _dropping(_svm_classifier(), "classlabels_ints") + model.graph.output[0].CopyFrom( + helper.make_tensor_value_info("Y", TensorProto.INT64, [4]) + ) + model.graph.output[1].CopyFrom( + helper.make_tensor_value_info("Z", TensorProto.FLOAT, [4, 2]) + ) + + with pytest.raises(CompileError, match="classlabels_ints"): + compile_onnx(model, tmp_path) + + +def test_a_support_vector_machine_without_rho_is_rejected(tmp_path): + with pytest.raises(CompileError, match="`rho`"): + compile_onnx(_dropping(_svm_regressor(), "rho"), tmp_path) + + +def test_a_kernel_onnx_does_not_define_is_rejected(tmp_path): + with pytest.raises(CompileError, match="kernel_type"): + compile_onnx(_svm_classifier(kernel_type="COSINE"), tmp_path) + + +def test_a_transform_onnx_does_not_define_is_rejected(tmp_path): + with pytest.raises(CompileError, match="post_transform"): + compile_onnx(_linear_regressor(post_transform="SIGMOID"), tmp_path) + + +def test_kernel_parameters_that_describe_less_than_a_kernel_are_rejected(tmp_path): + with pytest.raises(CompileError, match="kernel_params"): + compile_onnx(_svm_classifier(kernel_params=[0.5]), tmp_path) + + +def test_support_vectors_that_do_not_fill_the_matrix_are_rejected(tmp_path): + with pytest.raises(CompileError, match="support_vectors"): + compile_onnx( + _svm_regressor( + n_supports=3, + support_vectors=_SUPPORT_VECTORS[:-1], + kernel_params=_KERNEL_PARAMS, + ), + tmp_path, + ) + + +def test_a_support_vector_regressor_with_a_coefficient_per_nothing_is_rejected( + tmp_path, +): + with pytest.raises(CompileError, match="coefficient"): + compile_onnx(_svm_regressor(coefficients=[1.0, 0.5]), tmp_path) + + +def test_a_single_class_over_support_vectors_is_rejected(tmp_path): + """The pairwise scheme has no pairs to score, which its own reference refuses outright.""" + with pytest.raises(CompileError, match="at least two"): + compile_onnx( + _svm_classifier(classlabels_ints=[7], vectors_per_class=[3]), tmp_path + ) + + +def test_fewer_vector_counts_than_classes_is_rejected(tmp_path): + with pytest.raises(CompileError, match="vectors_per_class"): + compile_onnx(_svm_classifier(vectors_per_class=[3]), tmp_path) + + +@pytest.mark.parametrize("counts", [[-1, 4], [-4, 1]], ids=["scored", "unscored"]) +def test_a_negative_count_of_support_vectors_is_rejected(tmp_path, counts): + """Each count is a length the pairwise loops run to, and a negative one walks off the + tables they read; the reference scores such a pair as zero, which nothing can vouch for.""" + with pytest.raises(CompileError, match="negative count"): + compile_onnx(_svm_classifier(vectors_per_class=counts), tmp_path) + + +def test_too_few_coefficient_rows_for_the_class_pairs_is_rejected(tmp_path): + model = _svm_classifier( + classlabels_ints=[1, 2, 3], + vectors_per_class=[1, 1, 1], + rho=[0.25, 0.1, -0.2], + ) + + with pytest.raises(CompileError, match="row"): + compile_onnx(model, tmp_path) + + +def test_fewer_rho_than_class_pairs_is_rejected(tmp_path): + model = _svm_classifier( + classlabels_ints=[1, 2, 3], + vectors_per_class=[1, 1, 1], + coefficients=[0.5, -0.25, 0.75, 0.1, 0.2, -0.3], + ) + + with pytest.raises(CompileError, match="`rho` holds"): + compile_onnx(model, tmp_path) + + +def test_probabilities_over_more_than_two_classes_are_rejected(tmp_path): + """The one attribute combination whose two oracles disagree.""" + model = _svm_classifier( + classlabels_ints=[1, 2, 3], + vectors_per_class=[1, 1, 1], + coefficients=[0.5, -0.25, 0.75, 0.1, 0.2, -0.3], + rho=[0.25, 0.1, -0.2], + prob_a=[-1.5, 0.5, 1.0], + prob_b=[0.25, 0.0, -0.5], + ) + + with pytest.raises(CompileError, match="two classes only"): + compile_onnx(model, tmp_path) + + +def test_a_probability_without_its_pair_is_rejected(tmp_path): + with pytest.raises(CompileError, match="prob_b"): + compile_onnx(_svm_classifier(prob_a=[-1.5]), tmp_path) diff --git a/src/python/tests/test_extra_compiler_trees.py b/src/python/tests/test_extra_compiler_trees.py new file mode 100644 index 0000000..3a0f369 --- /dev/null +++ b/src/python/tests/test_extra_compiler_trees.py @@ -0,0 +1,641 @@ +"""The ONNX-ML tree ensembles: converted forests, and the models the compiler refuses. + +What each op computes on a single node is settled by the conformance suite (opset 5's two +corpus tests) and by the differential sweep against the reference evaluator. Neither reaches +where these ops actually turn up: a scikit-learn converter emits `ai.onnx.ml` opset **1**, +which the reference evaluator is not version-faithful for and therefore cannot be the oracle +of. The parity tests below run those converted forests against onnxruntime instead — the +second oracle, independent of both the compiler and the reference evaluator — and the rest of +the module covers the error contracts, which no sweep asserts. +""" + +from __future__ import annotations + +import shutil + +import pytest + +from fnnx.extras.compilers.c.errors import CompileError + +onnx = pytest.importorskip("onnx") +np = pytest.importorskip("numpy") +# The harness refuses to import without numpy, so this covers both dependencies. +harness = pytest.importorskip("fnnx.extras.compilers.c.harness") + +from onnx import TensorProto, helper # noqa: E402 + +from fnnx.extras.compilers.c import compile_onnx # noqa: E402 + +ML_OPSET = 5 + +requires_c_compiler = pytest.mark.skipif( + not any(shutil.which(name) for name in harness.COMPILER_CANDIDATES), + reason="no system C compiler available", +) + +_SAMPLES = 32 +_FEATURES = 4 + + +def _model(nodes, inputs, outputs, ml=ML_OPSET): + """A model whose intermediates carry no `value_info`, as a converter's output does.""" + model = helper.make_model( + helper.make_graph(nodes, "trees", list(inputs), list(outputs)), + opset_imports=[helper.make_opsetid("ai.onnx.ml", ml)], + ) + model.ir_version = 10 + return model + + +def _regressor(**attributes): + """A one-stump, one-target regressor, with `attributes` overriding what it declares.""" + declared = { + "n_targets": 1, + "nodes_treeids": [0, 0, 0], + "nodes_nodeids": [0, 1, 2], + "nodes_featureids": [0, 0, 0], + "nodes_modes": ["BRANCH_LEQ", "LEAF", "LEAF"], + "nodes_values": [0.5, 0.0, 0.0], + "nodes_truenodeids": [1, 0, 0], + "nodes_falsenodeids": [2, 0, 0], + "target_treeids": [0, 0], + "target_nodeids": [1, 2], + "target_ids": [0, 0], + "target_weights": [1.5, -2.5], + **attributes, + } + return _model( + [ + helper.make_node( + "TreeEnsembleRegressor", + ["X"], + ["Y"], + name="ensemble", + domain="ai.onnx.ml", + **declared, + ) + ], + [helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 3])], + [helper.make_empty_tensor_value_info("Y")], + ) + + +def _classifier(**attributes): + """A one-stump, two-class classifier, with `attributes` overriding what it declares.""" + declared = { + "classlabels_int64s": [10, 20], + "nodes_treeids": [0, 0, 0], + "nodes_nodeids": [0, 1, 2], + "nodes_featureids": [0, 0, 0], + "nodes_modes": ["BRANCH_LEQ", "LEAF", "LEAF"], + "nodes_values": [0.5, 0.0, 0.0], + "nodes_truenodeids": [1, 0, 0], + "nodes_falsenodeids": [2, 0, 0], + "class_treeids": [0, 0], + "class_nodeids": [1, 2], + "class_ids": [0, 1], + "class_weights": [1.0, 1.0], + **attributes, + } + return _model( + [ + helper.make_node( + "TreeEnsembleClassifier", + ["X"], + ["Y", "Z"], + name="ensemble", + domain="ai.onnx.ml", + **declared, + ) + ], + [helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 3])], + [ + helper.make_empty_tensor_value_info("Y"), + helper.make_empty_tensor_value_info("Z"), + ], + ) + + +# -------------------------------------------------------------------------------------- +# Converted scikit-learn forests, against onnxruntime +# -------------------------------------------------------------------------------------- + + +def _converted(estimator, data, target, **options): + """The ONNX a scikit-learn converter emits for `estimator`, at `data`'s exact shape.""" + skl2onnx = pytest.importorskip("skl2onnx") + data_types = pytest.importorskip("skl2onnx.common.data_types") + return skl2onnx.convert_sklearn( + estimator.fit(data, target), + initial_types=[("X", data_types.FloatTensorType(list(data.shape)))], + **options, + ) + + +def _session(model): + runtime = pytest.importorskip("onnxruntime") + runtime.set_default_logger_severity(3) + return runtime.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + + +def _matches_onnxruntime(model, feeds, tmp_path): + """Every output of the compiled artifact against the one onnxruntime computes for it.""" + compiled = compile_onnx(model, tmp_path).load() + outputs = compiled.run(feeds) + expected = _session(model).run(None, dict(feeds)) + assert [spec.name for spec in compiled.outputs] == [ + entry.name for entry in model.graph.output + ] + for entry, want in zip(model.graph.output, expected): + got, want = outputs[entry.name], np.asarray(want) + assert got.dtype == want.dtype, entry.name + assert got.shape == want.shape, entry.name + if want.dtype.kind == "f": + np.testing.assert_allclose( + got, want, rtol=1e-5, atol=1e-6, err_msg=entry.name + ) + else: + np.testing.assert_array_equal(got, want, err_msg=entry.name) + + +def _fitted_data(seed=0): + generator = np.random.default_rng(seed) + return generator.normal(size=(_SAMPLES, _FEATURES)).astype(np.float32) + + +def _regression_target(data): + return (2 * data[:, 0] + data[:, 1] - data[:, 3]).astype(np.float64) + + +@requires_c_compiler +@pytest.mark.parametrize( + ("module", "name", "arguments"), + [ + ("sklearn.tree", "DecisionTreeRegressor", {"max_depth": 4}), + ( + "sklearn.ensemble", + "RandomForestRegressor", + {"n_estimators": 5, "max_depth": 3}, + ), + ( + "sklearn.ensemble", + "GradientBoostingRegressor", + {"n_estimators": 5, "max_depth": 3}, + ), + ( + "sklearn.ensemble", + "ExtraTreesRegressor", + {"n_estimators": 4, "max_depth": 3}, + ), + ], +) +def test_a_converted_regressor_matches_onnxruntime(tmp_path, module, name, arguments): + """The `AVERAGE` a forest aggregates with and the `base_values` a boosted one offsets.""" + estimators = pytest.importorskip(module) + data = _fitted_data() + + model = _converted( + getattr(estimators, name)(random_state=0, **arguments), + data, + _regression_target(data), + ) + + assert [node.op_type for node in model.graph.node] == ["TreeEnsembleRegressor"] + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_multi_target_regressor_matches_onnxruntime(tmp_path): + """One leaf weighting several targets, which is what the flat leaf ranges are for. + + The converter declares a one-column result for a forest whose `n_targets` is two, which + onnxruntime pays no attention to; the compiler refuses the disagreement rather than + writing two columns into a buffer the header would size for one, and compiles the same + model once that declaration is dropped. + """ + ensemble = pytest.importorskip("sklearn.ensemble") + data = _fitted_data(seed=1) + targets = np.stack( + [_regression_target(data), data[:, 2].astype(np.float64)], axis=1 + ) + model = _converted( + ensemble.RandomForestRegressor(n_estimators=4, max_depth=3, random_state=0), + data, + targets, + ) + + with pytest.raises(CompileError, match="addresses a result of shape"): + compile_onnx(model, tmp_path / "declared") + model.graph.output[0].ClearField("type") + + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +@pytest.mark.parametrize("classes", [2, 3], ids=["binary", "multiclass"]) +def test_a_converted_classifier_matches_onnxruntime(tmp_path, classes): + ensemble = pytest.importorskip("sklearn.ensemble") + data = _fitted_data(seed=2) + labels = np.digitize(data[:, 0], np.linspace(-1, 1, classes - 1)).astype(np.int64) + + model = _converted( + ensemble.RandomForestClassifier(n_estimators=5, max_depth=3, random_state=0), + data, + labels, + options={"zipmap": False}, + ) + + assert [node.op_type for node in model.graph.node] == ["TreeEnsembleClassifier"] + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_boosted_classifier_matches_onnxruntime(tmp_path): + """Gradient boosting is where a classifier carries `base_values` and a transform.""" + ensemble = pytest.importorskip("sklearn.ensemble") + data = _fitted_data(seed=3) + labels = (data[:, 0] + data[:, 1] > 0).astype(np.int64) + + model = _converted( + ensemble.GradientBoostingClassifier( + n_estimators=5, max_depth=3, random_state=0 + ), + data, + labels, + options={"zipmap": False}, + ) + (node,) = [ + entry for entry in model.graph.node if entry.op_type == "TreeEnsembleClassifier" + ] + + assert {entry.name for entry in node.attribute} >= {"base_values", "post_transform"} + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +def test_a_converted_classifier_keeps_its_probabilities_through_the_zipmap_pass( + tmp_path, +): + """The converter's default output is a map, which the graph pass turns back into scores. + + onnxruntime is the oracle for the pairing itself: which label names which column is the + one thing a run of the graph with `ZipMap` already removed cannot show. + """ + ensemble = pytest.importorskip("sklearn.ensemble") + data = _fitted_data(seed=4) + labels = np.digitize(data[:, 1], [-0.5, 0.5]).astype(np.int64) + model = _converted( + ensemble.RandomForestClassifier(n_estimators=4, max_depth=3, random_state=0), + data, + labels, + ) + + assert "ZipMap" in [node.op_type for node in model.graph.node] + result = compile_onnx(model, tmp_path) + outputs = result.load().run({"X": data}) + predicted, rows = _session(model).run(None, {"X": data}) + (table,) = result.report["class_labels"] + + assert table["dtype"] == "int64" + np.testing.assert_array_equal(outputs["output_label"], predicted) + np.testing.assert_allclose( + outputs[table["tensor"]], + np.array( + [[row[label] for label in table["values"]] for row in rows], np.float32 + ), + rtol=1e-5, + atol=1e-6, + ) + + +@requires_c_compiler +def test_a_converted_pipeline_of_scaler_and_forest_matches_onnxruntime(tmp_path): + """The ensemble downstream of a preprocessor, whose intermediates carry no declared type.""" + pipeline = pytest.importorskip("sklearn.pipeline") + preprocessing = pytest.importorskip("sklearn.preprocessing") + ensemble = pytest.importorskip("sklearn.ensemble") + data = _fitted_data(seed=5) + + model = _converted( + pipeline.make_pipeline( + preprocessing.StandardScaler(), + ensemble.RandomForestRegressor(n_estimators=4, max_depth=3, random_state=0), + ), + data, + _regression_target(data), + ) + + assert [node.op_type for node in model.graph.node] == [ + "Scaler", + "TreeEnsembleRegressor", + ] + _matches_onnxruntime(model, {"X": data}, tmp_path) + + +@requires_c_compiler +def test_two_ensembles_of_one_shape_share_a_single_walker(tmp_path): + """The tree walk is one shared kernel per element type, however many nodes run it.""" + model = _regressor() + second = onnx.NodeProto() + second.CopyFrom(model.graph.node[0]) + second.name = "second" + second.input[0] = "Y" + second.output[0] = "Y2" + model.graph.node.append(second) + del model.graph.output[:] + model.graph.output.append(helper.make_empty_tensor_value_info("Y2")) + model.graph.input[0].type.tensor_type.shape.dim[1].dim_value = 1 + + result = compile_onnx(model, tmp_path) + + assert [name for name in result.report["kernels"] if "tree" in name] == [ + "trees_tree_aggregate_float_float" + ] + + +# -------------------------------------------------------------------------------------- +# What the compiler refuses +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "family", + ["base_values_as_tensor", "nodes_values_as_tensor", "target_weights_as_tensor"], +) +def test_a_double_precision_table_is_rejected(tmp_path, family): + """The `*_as_tensor` families, which ONNX's own reference implementation never reads.""" + table = helper.make_tensor(family, TensorProto.DOUBLE, [1], [0.5]) + + with pytest.raises(CompileError, match="_as_tensor"): + compile_onnx(_regressor(**{family: table}), tmp_path) + + +def test_a_string_labelled_classifier_is_rejected(tmp_path): + """A label tensor of strings is a run-time string whatever the ensemble computes.""" + model = _classifier(classlabels_int64s=None, classlabels_strings=["a", "b"]) + + with pytest.raises(CompileError, match="STRING"): + compile_onnx(model, tmp_path) + + +@requires_c_compiler +def test_a_converted_classifier_over_string_classes_is_rejected(tmp_path): + """Which is what a converter emits for a model fitted on string labels.""" + ensemble = pytest.importorskip("sklearn.ensemble") + data = _fitted_data(seed=6) + labels = np.where(data[:, 0] > 0, "yes", "no") + + model = _converted( + ensemble.RandomForestClassifier(n_estimators=3, max_depth=2, random_state=0), + data, + labels, + options={"zipmap": False}, + ) + + with pytest.raises(CompileError, match="STRING"): + compile_onnx(model, tmp_path) + + +def test_a_classifier_setting_both_label_families_is_rejected(tmp_path): + with pytest.raises(CompileError, match="classlabels_int64s"): + compile_onnx(_classifier(classlabels_strings=["a", "b"]), tmp_path) + + +def test_a_classifier_declaring_no_classes_is_rejected(tmp_path): + """Nothing then says what element type its labels would even be, let alone how many.""" + with pytest.raises(CompileError, match="no type"): + compile_onnx(_classifier(classlabels_int64s=None), tmp_path) + + +def test_a_single_class_label_other_than_one_is_rejected(tmp_path): + """The case ONNX's own reference implementation raises on, so nothing can vouch for it.""" + model = _classifier(classlabels_int64s=[7], class_ids=[0, 0]) + + with pytest.raises(CompileError, match="single class"): + compile_onnx(model, tmp_path) + + +def test_a_branch_test_onnx_does_not_define_is_rejected(tmp_path): + model = _regressor(nodes_modes=["BRANCH_APPROX", "LEAF", "LEAF"]) + + with pytest.raises(CompileError, match="BRANCH_APPROX"): + compile_onnx(model, tmp_path) + + +def test_a_feature_outside_the_input_is_rejected(tmp_path): + """A feature id the walker would read past the end of a row with.""" + model = _regressor(nodes_featureids=[7, 0, 0]) + + with pytest.raises(CompileError, match="feature 7"): + compile_onnx(model, tmp_path) + + +def test_a_target_outside_the_scored_targets_is_rejected(tmp_path): + model = _regressor(target_ids=[0, 3]) + + with pytest.raises(CompileError, match="target_ids"): + compile_onnx(model, tmp_path) + + +def test_a_child_no_node_defines_is_rejected(tmp_path): + model = _regressor(nodes_truenodeids=[9, 0, 0]) + + with pytest.raises(CompileError, match="node 9 of tree 0"): + compile_onnx(model, tmp_path) + + +def test_a_cycle_between_nodes_is_rejected(tmp_path): + """Which the emitted walker would otherwise loop on forever.""" + model = _regressor( + nodes_modes=["BRANCH_LEQ", "BRANCH_LEQ", "LEAF"], + nodes_truenodeids=[1, 0, 0], + nodes_falsenodeids=[2, 2, 0], + target_treeids=[0], + target_nodeids=[2], + target_ids=[0], + target_weights=[1.5], + ) + + with pytest.raises(CompileError, match="reachable more than once"): + compile_onnx(model, tmp_path) + + +def test_families_that_disagree_on_their_length_are_rejected(tmp_path): + model = _regressor(nodes_featureids=[0, 0]) + + with pytest.raises(CompileError, match="nodes_featureids"): + compile_onnx(model, tmp_path) + + +def test_a_base_value_per_nothing_is_rejected(tmp_path): + model = _regressor(base_values=[1.0, 2.0]) + + with pytest.raises(CompileError, match="base_values"): + compile_onnx(model, tmp_path) + + +def test_a_transform_onnx_does_not_define_is_rejected(tmp_path): + model = _regressor(post_transform="LOGIT") + + with pytest.raises(CompileError, match="post_transform"): + compile_onnx(model, tmp_path) + + +def test_an_aggregation_onnx_does_not_define_is_rejected(tmp_path): + model = _regressor(aggregate_function="MEDIAN") + + with pytest.raises(CompileError, match="aggregate_function"): + compile_onnx(model, tmp_path) + + +def test_a_rank_3_input_is_rejected(tmp_path): + """An ensemble reads a matrix of rows; the result is declared so that it is the rank + the kernel objects to rather than a shape nothing could infer.""" + model = _regressor() + model.graph.input[0].type.CopyFrom( + helper.make_tensor_type_proto(TensorProto.FLOAT, [2, 2, 3]) + ) + model.graph.output[0].type.CopyFrom( + helper.make_tensor_type_proto(TensorProto.FLOAT, [2, 2, 1]) + ) + + with pytest.raises(CompileError, match=r"\[N, F\]"): + compile_onnx(model, tmp_path) + + +# -------------------------------------------------------------------------------------- +# The opset-5 encoding, whose corpus tests cover the happy path +# -------------------------------------------------------------------------------------- + + +def _ensemble(**attributes): + declared = { + "n_targets": 1, + "tree_roots": [0], + "nodes_featureids": [0], + "nodes_truenodeids": [0], + "nodes_falsenodeids": [1], + "nodes_trueleafs": [1], + "nodes_falseleafs": [1], + "nodes_modes": helper.make_tensor("nodes_modes", TensorProto.UINT8, [1], [0]), + "nodes_splits": helper.make_tensor( + "nodes_splits", TensorProto.FLOAT, [1], [0.5] + ), + "leaf_targetids": [0, 0], + "leaf_weights": helper.make_tensor( + "leaf_weights", TensorProto.FLOAT, [2], [1.5, -2.5] + ), + **attributes, + } + return _model( + [ + helper.make_node( + "TreeEnsemble", + ["X"], + ["Y"], + name="ensemble", + domain="ai.onnx.ml", + **declared, + ) + ], + [helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 3])], + [helper.make_empty_tensor_value_info("Y")], + ) + + +@requires_c_compiler +def test_the_opset_5_encoding_runs_the_stump_it_describes(tmp_path): + compiled = compile_onnx(_ensemble(), tmp_path).load() + + scores = compiled.run({"X": np.array([[0.0, 0, 0], [1, 0, 0]] * 2, np.float32)})[ + "Y" + ] + + np.testing.assert_array_equal(scores, np.array([[1.5], [-2.5]] * 2, np.float32)) + + +def test_a_set_test_without_members_is_rejected(tmp_path): + model = _ensemble( + nodes_modes=helper.make_tensor("nodes_modes", TensorProto.UINT8, [1], [6]) + ) + + with pytest.raises(CompileError, match="membership_values"): + compile_onnx(model, tmp_path) + + +def test_a_membership_list_that_names_too_few_sets_is_rejected(tmp_path): + model = _ensemble( + nodes_modes=helper.make_tensor("nodes_modes", TensorProto.UINT8, [1], [6]), + membership_values=helper.make_tensor( + "membership_values", TensorProto.FLOAT, [2], [1.0, 2.0] + ), + ) + + with pytest.raises(CompileError, match="NaN-terminated"): + compile_onnx(model, tmp_path) + + +def test_a_branch_number_onnx_does_not_define_is_rejected(tmp_path): + model = _ensemble( + nodes_modes=helper.make_tensor("nodes_modes", TensorProto.UINT8, [1], [9]) + ) + + with pytest.raises(CompileError, match="nodes_modes"): + compile_onnx(model, tmp_path) + + +def test_a_root_outside_the_nodes_is_rejected(tmp_path): + with pytest.raises(CompileError, match="tree_roots"): + compile_onnx(_ensemble(tree_roots=[3]), tmp_path) + + +@pytest.mark.parametrize( + ("attributes", "message"), + [ + ({"nodes_falsenodeids": [5], "nodes_falseleafs": [0]}, "node 5 as a child"), + ({"nodes_truenodeids": [9]}, "leaf 9 as a child"), + ], + ids=["node", "leaf"], +) +def test_a_child_outside_the_family_it_addresses_is_rejected( + tmp_path, attributes, message +): + """An interior child is reached by the set-member traversal before the walker's tables + are built, so it is that traversal which has to refuse the ones no node defines.""" + with pytest.raises(CompileError, match=message): + compile_onnx(_ensemble(**attributes), tmp_path) + + +def test_a_cycle_between_opset_5_nodes_is_rejected(tmp_path): + """The traversal that lays out the set members walks these too, and stops here.""" + model = _ensemble( + nodes_featureids=[0, 0], + nodes_truenodeids=[1, 0], + nodes_falsenodeids=[0, 1], + nodes_trueleafs=[0, 0], + nodes_falseleafs=[1, 1], + nodes_modes=helper.make_tensor("nodes_modes", TensorProto.UINT8, [2], [0, 0]), + nodes_splits=helper.make_tensor( + "nodes_splits", TensorProto.FLOAT, [2], [0.5, 0.25] + ), + ) + + with pytest.raises(CompileError, match="reachable more than once"): + compile_onnx(model, tmp_path) + + +def test_a_missing_required_table_is_rejected(tmp_path): + model = _ensemble() + (node,) = model.graph.node + del node.attribute[ + next( + index + for index, entry in enumerate(node.attribute) + if entry.name == "leaf_weights" + ) + ] + + with pytest.raises(CompileError, match="leaf_weights"): + compile_onnx(model, tmp_path)