Share SQL stages and dispatch validated numerical kernels

Cargo.toml

@@ -32,6 +32,7 @@ prost = { version = "0.14", optional = true }
 # enables `pyo3/extension-module` through pyproject.toml for wheel builds; it
 # must stay disabled for ordinary Cargo test binaries so they link libpython.
 pyo3 = { version = "0.29.0", features = ["abi3-py310"] }
+sqlparser = { version = "0.62", features = ["visitor"] }
 tokio = { version = "1.46.1", features = ["macros", "rt", "sync", "time"] }
 tonic = { version = "0.14", features = ["transport"], optional = true }
 

accelerators/einfold_mlx/build.py

@@ -0,0 +1,97 @@
+"""Build the pinned EinFold dense-executor hook into this Python environment.
+
+EinFold (https://github.com/xqlsystems/einfold) is a DataFusion tensor
+optimizer. Its released Python API has no external kernel interface. This
+script builds and installs a wheel from the pinned feature commit in an
+isolated checkout. Requires git, a Rust toolchain, maturin,
+and uv. Nothing is installed into the system Python.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+_REVISION = "0b2f6db892c3cbf9df2c3d775fbcf4c5654df093"
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--source", type=Path, help="reuse a clean pinned source checkout"
+    )
+    args = parser.parse_args()
+    if sys.prefix == sys.base_prefix:
+        parser.error(
+            "run this builder with the checkout's virtual-environment Python"
+        )
+    for command in ("git", "cargo", "uv"):
+        if shutil.which(command) is None:
+            parser.error(f"{command} must be on PATH")
+    with tempfile.TemporaryDirectory(prefix="xql-einfold-") as temporary:
+        work = Path(temporary)
+        source = args.source or work / "source"
+        if args.source is None:
+            subprocess.run(
+                [
+                    "git",
+                    "clone",
+                    "https://github.com/cromwellian/einfold.git",
+                    str(source),
+                ],
+                check=True,
+            )
+            subprocess.run(
+                ["git", "checkout", _REVISION], cwd=source, check=True
+            )
+        revision = subprocess.check_output(
+            ["git", "rev-parse", "HEAD"], cwd=source, text=True
+        ).strip()
+        if revision != _REVISION:
+            parser.error(
+                f"expected EinFold revision {_REVISION}, got {revision}"
+            )
+        if subprocess.check_output(
+            ["git", "status", "--porcelain"], cwd=source
+        ):
+            parser.error(
+                "the source checkout must be clean"
+            )
+        subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "maturin",
+                "build",
+                "--release",
+                "--locked",
+                "--out",
+                str(work / "wheels"),
+            ],
+            cwd=source / "python",
+            env={**os.environ, "PYO3_PYTHON": sys.executable},
+            check=True,
+        )
+        (wheel,) = (work / "wheels").glob("*.whl")
+        subprocess.run(
+            [
+                "uv",
+                "pip",
+                "install",
+                "--python",
+                sys.executable,
+                "--force-reinstall",
+                "--no-deps",
+                str(wheel),
+            ],
+            check=True,
+        )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/ml/01_attention.py

@@ -231,7 +231,15 @@ def main() -> None:
                 rows = report.time(
                     f"forward ({spelling})", host, n, host.query, FORWARD
                 )
-                report.check(assert_close("forward", matrix(rows), want))
+                report.check(
+                    assert_close(
+                        "forward",
+                        matrix(rows),
+                        want,
+                        rtol=host.rtol,
+                        atol=host.atol,
+                    )
+                )
                 report.note("operators", host.operators(FORWARD))
 
                 step = f"gradient ({spelling})"
@@ -242,7 +250,13 @@ def main() -> None:
                 if rows is None:
                     continue
                 errors = [
-                    assert_close(f"d/d{name}", matrix(rows, f"d{name}"), ref)
+                    assert_close(
+                        f"d/d{name}",
+                        matrix(rows, f"d{name}"),
+                        ref,
+                        rtol=host.rtol,
+                        atol=host.atol,
+                    )
                     for name, ref in zip("qkv", want_grads)
                 ]
                 report.check(max(errors))
@@ -303,7 +317,11 @@ def check_nadaraya_watson(host_names: Sequence[str], report: Report) -> None:
             NADARAYA_WATSON,
         )
         got = xql.to_dataset(rows, dims=["p"])["v"].values
-        report.check(assert_close("nadaraya-watson", got, want))
+        report.check(
+            assert_close(
+                "nadaraya-watson", got, want, rtol=host.rtol, atol=host.atol
+            )
+        )
 
 
 def check_einfold_fuses(host_names: Sequence[str]) -> None: