diff --git a/.gitignore b/.gitignore index 8999c64c..dab5339f 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,10 @@ CHANGELOG.md .direnv # Pyodide wheels copied in to test the browser demo locally docs/demo/*.whl +# SQL training checkpoints; retain the accompanying text metrics/logs. +out-shakespeare-sql/ +benchmarks/apple/shakespeare-matched-*/last.npz +benchmarks/apple/shakespeare-matched-*/best.npz +benchmarks/apple/shakespeare-resident-*/last.npz +benchmarks/apple/shakespeare-compiled-*/last.npz +benchmarks/apple/shakespeare-compiled-*/best.npz diff --git a/Cargo.lock b/Cargo.lock index 864aa8e5..040fcbdc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3728,6 +3728,7 @@ dependencies = [ "prost", "pyo3", "pyo3-build-config", + "sqlparser", "tokio", "tonic", ] diff --git a/Cargo.toml b/Cargo.toml index 482e10ee..300b5c9e 100644 --- a/Cargo.toml +++ b/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 } diff --git a/README.md b/README.md index 42fe05a7..26c5b819 100644 --- a/README.md +++ b/README.md @@ -289,6 +289,26 @@ I want to give a special thanks to the following folks and institutions: opinion. +## Apple silicon tensor acceleration + +The optional [Apple silicon integration](accelerators/einfold_mlx/README.md) +connects [EinFold](https://github.com/xqlsystems/einfold), a tensor optimizer +for SQL, to [MLX](https://github.com/ml-explore/mlx), Apple's array framework. +It keeps SQL planning and gradients from +[ddx](https://github.com/xqlsystems/ddx), which differentiates SQL queries. +Float64 CPU execution is the default; +float32 GPU execution requires explicit opt-in. See the +[local M5 Ultra measurements](benchmarks/apple/README.md) for end-to-end +results, numerical errors, dispatch thresholds and limitations. +The [shared-CTE nanoGPT comparison](benchmarks/apple/nanogpt-shared.md) +combines cached shared SQL stages with the same CPU/GPU kernels and checks +larger models against JAX. +The experimental [compact execution path](benchmarks/apple/dense-execution.md) +retains arrays across forward and backward stages, reuses plans, and +executes supported operators through MLX. Its report includes checked +five-step GPT-2-size training, memory bounds and remaining native-framework +performance gaps. + ## License ``` diff --git a/accelerators/ddx_math/README.md b/accelerators/ddx_math/README.md new file mode 100644 index 00000000..0a892402 --- /dev/null +++ b/accelerators/ddx_math/README.md @@ -0,0 +1,44 @@ +# Exact GELU for SQL training + +[ddx](https://github.com/xqlsystems/ddx) differentiates SQL query plans. +Its released scalar rules do not include the error function, `erf`, used +by [nanoGPT](https://github.com/karpathy/nanoGPT)'s exact GELU activation. +The pinned feature commit adds `erf'(x) = 2/sqrt(pi) * exp(-x*x)` to the scalar +registry and retains `erf` in derivative plans. The rule composes with +the chain rule and higher derivatives; it does not substitute a model's +handwritten backward pass. + +The builder uses the [ddx feature fork](https://github.com/cromwellian/ddx), pinned to commit +`953b78ecfc8bce2b1c1a71f0d4c7300444d00cc7`, based on release commit +`af71683f6678bb093e600f948c212490a4d8b74b`. Historical benchmark wheels used +`ddxdb==0.4.6+xqlmath.3`; the contribution retains upstream version 0.4.6. +The feature commit also exposes +`Prepared.values`: metadata for the primal outputs computed by prepared +derivative programs. These outputs can be read during a run before its +temporary tables are released. This lets the GPU training session reuse +the actual loss from backward preparation instead of executing a separate +forward query. It does not alter existing derivative rules. + +Engines must +register an `erf(DOUBLE)` implementation; `xarray_sql.math.register_erf` +does this for DataFusion and marks its equivalent array primitive as safe +for EinFold CPU and MLX GPU execution. Replacing the SQL UDF removes that +permission. This contribution does not add an erf implementation to DuckDB. + +From the xarray-sql directory, after installing the normal ML dependencies +and the [compact EinFold executor](../einfold_mlx/README.md): + +```bash +source .venv/bin/activate +uv pip install cmake maturin +python accelerators/ddx_math/build.py +``` + +Requires git, a Rust toolchain, CMake and uv on PATH. The builder clones +an isolated checkout, validates the pinned feature revision, builds its +locked wheel and installs it in the active virtual environment. No patch +is applied. Use `--source /path/to/ddx` for a clean checkout already at the pin. `uv sync` may +replace locally built wheels; rerun the optional builders if necessary. + +The public SQL numerical checks are in `tests/test_math.py`. The ddx contribution +also adds a Rust check of the chain rule and second derivative. diff --git a/accelerators/ddx_math/build.py b/accelerators/ddx_math/build.py new file mode 100644 index 00000000..a2bd6041 --- /dev/null +++ b/accelerators/ddx_math/build.py @@ -0,0 +1,90 @@ +"""Build ddx's exact error-function derivative in an isolated checkout. + +ddx (https://github.com/xqlsystems/ddx) differentiates SQL. The pinned +feature commit adds erf's analytic derivative, including higher derivatives, and +exposes prepared primal output metadata for loss reuse. +Requires git, cargo, CMake, maturin and uv on PATH. +""" + +import argparse +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +_REVISION = "953b78ecfc8bce2b1c1a71f0d4c7300444d00cc7" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--source", type=Path, help="reuse a clean pinned checkout" + ) + args = parser.parse_args() + if sys.prefix == sys.base_prefix: + parser.error("use the repository virtual environment's Python") + for command in ("git", "cargo", "cmake", "uv"): + if shutil.which(command) is None: + parser.error(f"{command} must be on PATH") + with tempfile.TemporaryDirectory(prefix="xql-ddx-math-") as temporary: + source = args.source or Path(temporary) / "source" + if args.source is None: + subprocess.run( + [ + "git", + "clone", + "https://github.com/cromwellian/ddx.git", + str(source), + ], + check=True, + ) + subprocess.run( + ["git", "checkout", _REVISION], cwd=source, check=True + ) + if ( + subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=source, text=True + ).strip() + != _REVISION + ): + parser.error(f"expected ddx revision {_REVISION}") + if subprocess.check_output( + ["git", "status", "--porcelain"], cwd=source + ): + parser.error("source checkout must be clean") + wheels = Path(temporary) / "wheels" + subprocess.run( + [ + sys.executable, + "-m", + "maturin", + "build", + "--release", + "--locked", + "--out", + str(wheels), + ], + cwd=source / "python/ddxdb", + env=os.environ | {"PYO3_PYTHON": sys.executable}, + check=True, + ) + (wheel,) = wheels.glob("*.whl") + subprocess.run( + [ + "uv", + "pip", + "install", + "--python", + sys.executable, + "--no-deps", + "--reinstall", + str(wheel), + ], + check=True, + ) + + +if __name__ == "__main__": + main() diff --git a/accelerators/einfold_mlx/README.md b/accelerators/einfold_mlx/README.md new file mode 100644 index 00000000..d1406abf --- /dev/null +++ b/accelerators/einfold_mlx/README.md @@ -0,0 +1,361 @@ +# EinFold kernel bridge for Apple silicon + +[EinFold](https://github.com/xqlsystems/einfold) recognizes tensor operations +in Apache DataFusion SQL plans. [MLX](https://github.com/ml-explore/mlx) is +Apple's array framework for Apple silicon. This implements Option B: + +```text +SQL → DataFusion logical plan → ddx rewrite, when requested + → EinFold detection and SQL input validation + → float64 CPU matrix multiplication or float32 MLX GPU kernels + → EinFold row/domain reconstruction → DataFusion → Arrow/xarray +``` + +[ddx](https://github.com/xqlsystems/ddx), the SQL automatic differentiation +engine, constructs every gradient, VJP and JVP. Logical plans remain normal +DataFusion plans until physical planning, so ddx can serialize them through +Substrait. MLX only evaluates arithmetic inside those plans. + +The released EinFold Python API has no external executor hook. The optional +builder pins the feature commit `9feb8e5a7efaecf6322afc436d0cac60814ef842` in the [EinFold feature fork](https://github.com/cromwellian/einfold), +which contains the native bridge and compact array module. Its tested base +is `39b25210d3b39027ef91584383fb114dfa8a4a48` (Python 0.2.0). Historical +benchmark wheels were labeled `0.2.0+xqlmlx.6`; the contributed source retains +version 0.2.0. This feature is not in the ordinary PyPI wheel yet. The +xarray-sql Rust extension supplies SQL parsing and verified ordered Arrow +providers used by shared execution. + +## Owned state and complete training graphs + +`xarray_sql.training.TrainingSession` adds an explicit ownership boundary +for complete, finite, nonnullable rectangular state tables. It imports +weights and optimizer moments once, then retains MLX arrays between SQL +updates. Ordinary `DenseExecutor` queries retain their existing live-input +and Arrow snapshot behavior. Use a dedicated Context for a session and +read its current state through `session.query` or `session.export`. + +`session.prepare_step` receives the original SQL gradient statements, +ordinary SQL intermediate/optimizer statements and a mapping from their +outputs to persistent state tables. It evaluates gradients, clipping and +all parameter updates in one invocation. The optional +[ddx math/primal build](../ddx_math/README.md) exposes the already computed +primal loss. No independent loss-forward call is needed. + +With `compile=True`, the session traces these SQL/ddx array programs into +one cached `mx.compile` function, with mutable numerical state passed as +explicit array arguments. Schemas, coordinate layouts, view definitions +and session settings guard reuse. Structural checks still run or must be +proven from the coordinate/count facts; a check that requires an +unsupported dynamic host decision declines compilation. Device finite, +NULL and complete-domain checks return a status scalar before all new +state tables are published together. Failed updates retain the previous +state and RNG counter. + +The pinned EinFold integration supplies device tensor leaves, finite integer +domain metadata, GPU index joins and embedding scatter grouping. Declare +bounded integer inputs with `session.index_domains` before binding them. +Actual input values are checked against those domains. Numerical index +arithmetic requires an explicit floating cast; unsupported integer/domain +behavior declines instead of silently changing SQL arithmetic. Schema or +coordinate changes require a new session. Explicit `bind` replacements +copy caller inputs and preserve ownership; exports are independent Arrow +snapshots. Numerical host conversion during compilation is refused. + +`session.dropout` generates masks from an explicit seed and step key on +GPU. Save the seed and next step with the checkpoint. Complete training +results can retain masks and raw gradients for correctness inspection; +these exports are outside benchmark timers. Ordinary steps read only the +validation and loss scalars. Full state is exported for checkpoints. + +From the xarray-sql directory, the matching Shakespeare recipe uses: + +```sh +MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/ml/03_shakespeare.py \ + --compile-training --gpu-dropout --steps 100 +``` + +Omit `--steps` for the recipe's full training run. `--resident` exercises +the intermediate implementation with device state and shared primal loss, +without whole-step compilation. `--compile-training` without GPU dropout +keeps NumPy-generated masks as explicit inputs, useful for comparison with +the earlier deterministic trajectory. CPU resident execution uses +`--resident --backend dense_cpu --cache-gib 32`. + +## Compact execution across forward and backward stages + +`xarray_sql.dense.DenseExecutor` is an opt-in implementation of +SQL → DataFusion logical plans → EinFold array programs → NumPy/MLX. +It retains shared CTE and ddx stage values as private arrays with explicit +coordinate axes. Numerical joins become gathers or contractions, window +reductions broadcast, and reduced UNION branches accumulate gradients. +Coordinates expand into scalar Arrow columns at public output boundaries. +The executor orders compact axes before expansion where their exact domains +prove the requested SQL ordering. + +```python +from einfold.datafusion import EinFold +from xarray_sql import Context +from xarray_sql.dense import DenseExecutor + +ctx = Context(threads=30).with_extensions(EinFold()) +# Register input and parameter Datasets on ctx. +executor = DenseExecutor(ctx) # float64 CPU by default +# Explicitly permit float32 GPU arithmetic, with MLX_ENABLE_TF32=0 at launch: +# executor = DenseExecutor(ctx, precision="float32") +executor.register_view("model", model_select_sql) +loss = executor.prepare(loss_select_sql) +step = executor.prepare_all(gradient_update_sqls) +loss_rows = loss() +updated_parameter_rows = step() +print(executor.statistics()) +``` + +ddx differentiates the original SQL. This path does not use MLX automatic +differentiation or replace a model with framework code. Logical programs +and compiled elementwise functions are reused across steps; compiled +functions receive every numerical leaf as an argument. Input NumPy buffers +are imported afresh each call, preserving in-place writes and registration +changes. Private schema-only providers serve ddx's named reads; they are +never scanned to produce numerical results. Validation checks evaluate +actual rows. Public outputs are independent Arrow snapshots, and input +parameters change only when the caller writes an update. + +The compact source contract is a rectangular Dataset with unique integer +or string coordinates and finite float64 numerical buffers. The initial +implementation supports a closed set of arithmetic, CASE/NULL operations, +gathers, grouped reductions, full-partition windows, axis joins and domain +left joins. Unsupported layouts or operations restart the call through +the existing shared SQL executor before any parameter is changed. +`fallback=False` reports unsupported cases directly and is used by the +benchmarks. Custom scalar-function overrides stay in SQL. GPU precision is +explicit; it retains the existing `3e-5` benchmark tolerance and the +float32 AdamW limitations documented in the measurements. + +The `.5` integration also retains multiplicative factors through aliases +and zero-valued CASE branches when reducing a conditional product. This +keeps attention backward masks inside contractions instead of expanding a +batch × heads × context × context × channels temporary. Exceptional +arithmetic retains the normal SQL fallback. Exact GELU is available with +the optional [ddx erf build](../ddx_math/README.md) and +`xarray_sql.math.register_erf(ctx)`; arbitrary UDF overrides remain in SQL. +The [Shakespeare trainer](../../benchmarks/ml/03_shakespeare.py) uses +these features for dropout and AdamW at nanoGPT's character-model size. + +GPU intermediates remain MLX arrays through their supported operations. +Compiled pointwise regions use `mx.compile`; contractions use MLX matrix +operations and embedding gradients use scatter reductions. Exact counts +use integer arithmetic, with bounded int32 scatter where Metal requires +it. GPU readbacks occur at SQL validation/output boundaries. Caches retain +program structure across calls; numerical graphs are released after each +call. This API still imports and exports the full parameter state each +training call. An API that owns persistent GPU training state is a future +extension with a different ownership and synchronization contract. + +The computed-value cache defaults to 8 GiB per invocation; adjust it with +`max_cached_bytes`. Source buffers remain available throughout the call; +computed values may be evicted and recomputed. `max_array_bytes` defaults +to a conservative 1 GiB estimate for a numerical temporary. These are +array/cache bounds, not an RSS limit: source tables, public Arrow results, +working allocations and driver memory are additional. The benchmark +runner enforces its own process limit and accepts `--dense-cache-gib`. +Exact coordinate and nonnullable-domain facts keep row and NULL masks +compact, including UNION branches whose domains differ on one axis. + +The [dense execution report](../../benchmarks/apple/dense-execution.md) +contains measured speedups, memory, copy counters, failed attempts and +remaining overheads. The +[research and execution plan](dense-execution-plan.md) records the source +research and implementation sequence. + +## Install and run + +Requires Apple silicon, a working Metal GPU, Python 3.10+, Rust 1.94+, +Git, uv and maturin. In the xarray-sql checkout: + +```sh +uv sync --extra mlx --extra ml +.venv/bin/python accelerators/einfold_mlx/build.py +MLX_ENABLE_TF32=0 .venv/bin/python your_program.py +``` + +The builder checks out the pinned feature commit into a temporary directory, +builds its locked wheel, and installs only into the interpreter's environment. +No patch is applied and no EinFold source is duplicated in this repository. +A later `uv sync` can replace the custom wheel; rerun the builder afterward. +Use `--source /path/to/einfold` to build a clean checkout already at the pin. +For supplied benchmarks, `.venv/bin/python -m benchmarks.apple.benchmark` +runs this checkout directly. Directly running an ML demo needs +`PYTHONPATH="$PWD"` unless xarray-sql is installed as editable. + +The pinned EinFold wheel includes `einfold.dense_program`, used by the +optional `DenseExecutor` to retain compact arrays across SQL stages. + +```python +from xarray_sql import Context, ad +from xarray_sql.mlx import MLX + +accelerator = MLX(precision="float32") +ctx = Context(threads=30).with_extensions(accelerator) +ctx.from_dataset("x", dataset) +rows = ctx.sql(statement).to_arrow_table() +print(accelerator.statistics()) + +# One SQL gradient query can use the same preparation API as forward SQL. +single_step = ctx.prepare(gradient_statement) +rows = single_step(x=new_dataset_values) + +# Several parameter updates can share one loss and its ddx programs. +step = ad.prepare(ctx, [gradient_statement, update_statement]) +gradient_rows, updated_rows = step(x=new_dataset_values) +``` + +`Context.with_extensions()` returns another xarray-aware handle, sharing +Dataset metadata and prepared-query invalidation with the original handle. +Always use the returned handle, which carries the extension's codecs. +`threads=` controls DataFusion's target parallel partitions. The default is +all 30 CPU cores on the measured M5 Ultra; 60 partitions modestly helped +plain DataFusion attention but slightly slowed EinFold in the confirmation +run. The [overhead report](../../benchmarks/apple/README.md#training-overheads-and-cpu-parallelism) +contains both results. + +## Numerical and SQL contracts + +- `MLX()` defaults to **float64 CPU** matrix multiplication through NumPy; + attention keeps EinFold's original CPU implementation. No implicit + conversion of double precision scientific data is enabled by default. +- `MLX(precision="float32")` explicitly permits single precision GPU + arithmetic, including intermediate values used by ddx. Outputs keep the + SQL Arrow dtype. **It does not preserve float64 rounding or error bounds.** +- Start Python with `MLX_ENABLE_TF32=0`. On M5, MLX otherwise uses arithmetic + with less precision than float32 for matrix operations. The extension + requires this setting and does not mutate process-wide MLX settings. + See [MLX's numerical precision documentation](https://ml-explore.github.io/mlx/build/html/usage/precision.html). +- EinFold retains key equality, bag semantics, NULLs, row existence, + grouping, domain defaults, and output mask reconstruction. Duplicate, + sparse, non-positional and unsupported inputs keep EinFold/DataFusion + fallback behavior. Non-finite factors stay on the CPU. Very large or tiny + finite factors also decline GPU execution. +- GPU attention is limited to complete finite inputs, product scores, + shifted exponentials and normalized weights. Unshifted exponentials, + general score expressions and other pipelines use existing CPU kernels. + Key intervals come from EinFold's normalized mask conditions. Masked + infinities never enter GPU arithmetic. +- There is no input-value cache: repeated queries and in-place parameter + writes read current values. Backend ownership is session-scoped; ordinary + DataFusion and EinFold contexts can coexist with accelerated contexts. +- Gradient definitions and typed ddx refusals are unchanged. `ad.prepare` + reuses ddx programs while checking catalog/schema/settings changes; + it materializes final results so later calls cannot overwrite them. + It intentionally uses ddx's regular preparation, which supports table + replacement, rather than experimental in-place physical-plan reuse. + +## Dispatch and memory + +The native side declines small requests before copying buffers into Python +(default `min_ops=2**20` scalar products). In float32 mode, GPU dispatch +starts at `min_gpu_ops=2**33`; smaller products use float64 CPU matrix +multiplication and smaller attention uses EinFold's CPU tiles. This avoids +GPU dispatch regressions on small workloads. Both limits can be adjusted; +`min_ops=0, min_gpu_ops=0` forces eligible computations through the GPU for +correctness testing or threshold experiments. + +Attention uses at most 1024 queries per tile, with a smaller tile when the +backend's working-memory estimate requires it. Masks and temporary score +storage grow with the number of keys and tile width, rather than the full +number of query/key pairs. Each tile trims keys outside the union of its +allowed intervals. Rectangular and causal intervals use MLX's mask-free or +causal kernels when their ranges match exactly; other masks remain explicit. Backend buffers are bounded by a configurable +512 MiB estimate and reserved through EinFold's DataFusion memory budget. +This estimate is not a hard limit on the GPU driver's allocator or its +kernel cache. All GPU results synchronize before Arrow conversion. + +`statistics()` counts completed `cpu_matmul`, `gpu_matmul`, and +`gpu_attention` kernels, plus requests declined inside Python. Requests +rejected by native thresholds are not counted. Planning a GPU-capable +operator alone does not increment a completion counter. + +## Validation + +See [the measured report](../../benchmarks/apple/README.md), the standalone +benchmark, raw timings and CPU profile in `benchmarks/apple`. + +```sh +MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python -m pytest -m 'not integration' +MLX_ENABLE_TF32=0 .venv/bin/python -m pytest tests/test_mlx.py +``` + +The native contribution also passes EinFold's kernel and randomized invariant +suite, except for the existing plain-DataFusion `nan_compares_by_its_sign` +assumption on this ARM Mac. That test expects `inf - inf` to produce a +negative NaN. Here plain DataFusion produces a positive NaN, without any +extension installed. The test remains unchanged in the contribution and was +documented in the measured report; it is not silently removed. + +## Shared CTE execution + +The historical `0.2.0+xqlmlx.3` build added safe prepared-plan resets and multi-key +attention detection. Batch and head keys must match in the scores, values, +softmax partition and normalizer. The kernel validates complete positional +domains before treating those groups as independent attention problems. + +`xarray_sql.shared.SharedExecutor` materializes reachable CTEs read more than +once and caches each stage's physical plan. It writes new stage values into +owned Arrow buffers. A changed coordinate domain, row count or schema +invalidates plans that read the replaced buffer. Dynamic filter pushdown is +disabled on this explicit executor's context, because its cached filters +would otherwise retain values from an earlier training step. + +This reset also runs when DataFusion's FFI wrapper replaces children: +DataFusion 55 does not forward a separate reset callback through that +wrapper. Ordinary `Context.prepare` conservatively retains only the +optimized logical plan when dynamic filters are enabled. + +```python +from xarray_sql import Context +from xarray_sql.mlx import MLX +from xarray_sql.shared import SharedExecutor + +accelerator = MLX(precision="float32", min_ops=0, min_gpu_ops=0) +ctx = Context().with_extensions(accelerator) +# Register parameter/input Datasets on ctx first. +executor = SharedExecutor(ctx, inline_ctes=["scores", "weights", "norm"]) +executor.register_view("model", model_select_sql) +loss = executor.prepare(loss_select_sql) +step = executor.prepare_all(gradient_update_sqls) +``` + +Virtual views retain their original SQL. ddx receives those definitions, +not the forward executor's materialized tables, and builds the backward +pass before its own stages execute. Numerical work in both paths passes +through DataFusion and EinFold before the MLX callback. Returned Arrow +tables are independent snapshots. + +DataFusion 55's Substrait consumer cannot apply a projected read to an +expanded view. When a ddx stage reads a deferred view, the executor binds +the same program through ddx's SQL writer instead. It retains the original +SQL differentiation and avoids substituting forward materializations into +the derivative's lineage. + +Join-only ddx stages remain logical views until a consumer reduces their +rows, avoiding storage of the full contraction join. Very large cotangent +sums (at least 32 UNION branches) materialize their reductions sequentially +to limit concurrent join builds. Compound reduction stages rebind transient +ddx providers on each call. They retain output storage, while ordinary +forward and backward stages also retain physical plans. + +An exact, bounded coordinate cache shares immutable integer/string buffers; +floating point values always have separate writable storage. Every cache +hit checks all values and NULLs. The default budget is 256 MiB, adjustable +with `max_key_cache_bytes`. Canonically ordered stage tables expose verified +ordering to DataFusion without an additional sort. Coordinate/domain/schema +changes invalidate dependent plans. See the +[shared nanoGPT measurements](../../benchmarks/apple/nanogpt-shared.md) +for measured benefits and remaining large-model memory limits. + +This API explicitly supports deterministic numerical SELECTs with flat, +nonrecursive WITH clauses. It refuses unknown functions and nested WITH +scopes. Callers identify CTEs that should remain inside a fused region with +`inline_ctes`. Stage boundaries still exchange Arrow data: this does not +implement a GPU-resident training graph or replace SQL derivatives with +MLX autodiff. The native xarray-sql extension must also be rebuilt for its +SQL-parser helper (`maturin develop --release` in the checkout). diff --git a/accelerators/einfold_mlx/build.py b/accelerators/einfold_mlx/build.py new file mode 100644 index 00000000..0e6b7a62 --- /dev/null +++ b/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 = "9feb8e5a7efaecf6322afc436d0cac60814ef842" + + +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() diff --git a/accelerators/einfold_mlx/dense-execution-plan.md b/accelerators/einfold_mlx/dense-execution-plan.md new file mode 100644 index 00000000..d56412f1 --- /dev/null +++ b/accelerators/einfold_mlx/dense-execution-plan.md @@ -0,0 +1,300 @@ +# Dense execution of SQL training programs + +Implementation and validation results are available in the +[dense execution report](../../benchmarks/apple/dense-execution.md). +The local executor now retains compact forward/backward arrays, reuses +structural plans, runs the supported program on MLX, preserves the original +ddx differentiation, and has completed checked five-step GPU training at +context 1024. Default APIs still return Arrow snapshots. An explicit API +for persistent GPU-owned training state remains future work; long-context +CPU cache sensitivity and the dense AdamW refusal are reported as limits. + +Research and execution plan, November 1, 2026. The requested execution +threshold is 06:00 America/Los_Angeles (14:00 UTC); research began before +it, and implementation may proceed now that this time has passed. + +## Evidence and objective + +The current SQL → DataFusion → EinFold → MLX bridge accelerates individual +matrix and attention kernels, but retains relational intermediates. The +123,563,520-parameter model (12 layers, width 768, 12 heads, vocabulary +50257, batch 1, context 16) takes 163.65 s per warm SQL GPU step. Across +five steps, counters record 131 s planning, 169 s stage storage and 504 s +execution, with 7,148 physical plans built. Native JAX CPU takes 0.334 s +eager, 0.394 s with separate loss/backward calls and 0.136 s compiled; +upstream nanoGPT's PyTorch model takes 0.095 s CPU float64 and 0.024 s MPS +float32. All parameter updates were checked. These are short-context +training results; the existing context-1024 comparison is forward-only. + +See [SQL measurements](../../benchmarks/apple/nanogpt-shared.md), +[native measurements](../../benchmarks/apple/nanogpt-frameworks.md) and +their linked raw records. They are immutable baselines for this work. + +The objective is to execute validated tensor-shaped regions of the actual +SQL and ddx programs on compact arrays, keeping intermediates dense and +avoiding per-operator Arrow reconstruction. ddx remains the source of +derivatives. The implementation must work on plans and axis relationships, +without recognizing a benchmark's name or substituting a handwritten GPT. + +## Research findings + +[EinFold](https://github.com/xqlsystems/einfold) is a tensor optimizer for +SQL. Its [design](https://github.com/xqlsystems/einfold/blob/main/docs/design.md) +already describes exact shape/layout facts in a side table, separate +symbolic and numerical passes, and optimization of a whole program of +plans. Its current DataFusion operators still collect Arrow inputs and +return relational Arrow rows. This supports adding a dense program path +without changing the SQL surface, while device-specific execution remains +in the host integration. A design aspiration is not an implemented API. + +[ddx](https://github.com/xqlsystems/ddx) differentiates SQL. Its +[program representation](https://github.com/xqlsystems/ddx/blob/main/crates/ddx-ad/src/program.rs) +contains input tables, validation checks, forward steps, backward steps +and output tables. Each step has a name and an unbound Substrait plan; +reads can be bound to schemas. The Python backend currently exposes +per-step `compile`, `materialize_compiled` and existence checks. Those +hooks permit an incremental dense backend. A whole-program preparation +hook is a possible later extension, not a precondition for useful progress. + +[DataFusion](https://datafusion.apache.org/library-user-guide/extending-operators.html) +supports custom logical/physical nodes through optimizer rules. A dense +region can therefore be represented internally while preserving its +ordinary relational input/output contract. One physical region should +contain many array operations, rather than returning scalar Arrow rows +between every pair of operators. + +[MLX](https://ml-explore.github.io/mlx/build/html/usage/unified_memory.html) +is Apple's array framework for unified CPU/GPU memory. The important costs +here are allocation, packing, conversion and synchronization; CPU/GPU +crossings on this machine are not necessarily discrete-device PCIe copies. +[Its conversion documentation](https://ml-explore.github.io/mlx/build/html/usage/numpy.html) +states that `mx.array` copies, and float64 NumPy input converts to float32. +Buffer/DLPack interop has ownership and synchronization requirements. +Unified memory alone does not make our current bridge zero-copy. + +[MLX compilation](https://ml-explore.github.io/mlx/build/html/usage/compile.html) +can eliminate common work and fuse operations. Shape/type/input changes +can cause recompilation. Build functions over explicit array arguments; +avoid capturing mutable parameter values in cached closures. Defer +evaluation to region/output boundaries, with explicit synchronization for +timing. Execute ddx's derivative expressions; do not use MLX autodiff to +replace them. Keep `MLX_ENABLE_TF32=0` as required by the +[precision documentation](https://ml-explore.github.io/mlx/build/html/usage/precision.html). + +[Arrow's fixed-shape tensor type](https://arrow.apache.org/docs/format/CanonicalExtensions.html#fixed-shape-tensor) +is useful for tensor interchange but changes a scalar field into a tensor +field. It cannot silently replace a public SQL table of scalar values. +Use a private dense representation and reconstruct the ordinary rows at +real relational/output boundaries. The +[Arrow C interface](https://arrow.apache.org/docs/format/CDataInterface.html) +can share owned Arrow buffers across runtimes, but does not prove tensor +shape, uniqueness, absence of NULLs, or MLX compatibility. + +## Recommended design + +Use a program-scoped dense store and a small array-operation graph. +Preserve existing public SQL/Arrow schemas and a relational fallback. + +### Compact values and exact domains + +A dense value contains its array/buffer, ordered axes, dtype, shape, +strides/view mapping and lifetime owner. Store coordinate dictionaries +once per axis/domain, rather than once per tensor element. Column lineage +maps SQL fields to axis coordinates, indexed gathers or scalar array +expressions. Attach exact facts to the plan node where they are valid. + +Initial eligibility is deliberately narrow: verified rectangular domains, +unique keys and supported finite numerical inputs. Arbitrary coordinate +labels may use exact dictionaries. Unsupported duplicate, NULL, sparse or +value-dependent domains retain relational execution until their handling +is proved. Shape is insufficient evidence of completeness or uniqueness. + +Cache structural proofs under coordinate/schema epochs. Refresh numerical +values under separate value epochs. Ordinary mutable NumPy/Arrow inputs +are imported once per call, so direct external writes remain visible. +Cross-call value caching is allowed only for explicitly owned, versioned +state with a defined write path. Do not infer that an untracked external +array has stayed unchanged because its schema or pointer is unchanged. + +### Dense lowering of backward subgraphs + +Lower a closed set of verified operations: reads, aliases, projections, +views, elementwise arithmetic/CASE, broadcasts, gathers, reductions, +contractions and required domain-fill operations. Use field lineage and +axis relationships from actual DataFusion/Substrait plans. Avoid SQL text +regexes and neural-network-specific substitutions. + +Recognize shared-axis join plus grouped product as a contraction before +forming its Cartesian products. For example, the tied embedding gradient +has the form `X.T @ dlogits`: context 16, width 768 and vocabulary 50257 +represent 617,558,016 possible product contributions, but only a +38,597,376-element output. Retain that matrix formulation through normalized +and residual operands. Extend existing EinFold fold/chain detection where +that removes a dominant missed pattern; use a general dense operation +graph for the surrounding maps and reductions. + +Priority patterns are reconstructed layer-normalized operands, residual +joins, null-safe key joins over proven nonnullable axes, gradient domain +left joins with defaults, gather/scatter contributions for embeddings, +and sums of many reduced gradient contributions. A UNION is a bag union, +not elementwise addition unless a following reduction and domain proof +establish that equivalence. Group existence and all-NULL SUM behavior must +be retained. Masks describe absent rows, with no multiplication of masked +NaN/Inf values by zero. + +### Stable bindings and compiled programs + +Bind intermediate reads to stable logical slots rather than new physical +table providers on each run. Each slot has an immutable structural +descriptor and a current value generation. The graph cache key includes +schema, coordinate/domain facts, function semantics/settings, dtype and +precision mode. Changing only parameter values updates bindings; changed +domains or semantics invalidate affected regions. + +Start through ddx's existing backend hooks: compiled steps may hold dense +plans and `materialize_compiled` may store a lazy internal value. Named +read schemas must remain real and fallback scans must materialize current +rows on demand. Existence checks can use exact domain facts only when +their result is proved; otherwise evaluate the check or cross a relational +boundary. Public methods returning Arrow still return complete Arrow +snapshots. + +If per-step hooks prevent adequate fusion, add an optional ddx +whole-program compilation interface using its existing program data. +Existing backends keep their current behavior. Preserve validation checks, +typed refusals and the ordering of externally visible effects. + +### CPU first, then resident MLX regions + +Implement the dense path on NumPy float64 first to distinguish relational +lowering problems from GPU precision/device behavior. Then execute the +same operation graph with MLX arrays in explicitly permitted float32 mode. +Keep maximal supported regions as arrays; convert at imports, unsupported +relational boundaries and requested public outputs. Avoid `np.asarray`, +`.item()` and `mx.eval` inside each individual numerical operator. + +The current slice-based native backend API forces packing into f64 vectors +and returning f64 vectors per call. The new path needs an owned array/graph +interface with stable handles and views, while retaining the existing +callback as a fallback. Keep EinFold's IR hardware-neutral and MLX objects +in the host bridge. Profile before adding custom C++/Metal kernels. + +Only after the compatible Arrow-output path is correct, consider an +explicit resident training-state API. It must state who owns parameters, +how writes are versioned and when host snapshots are produced. Existing +`Context.prepare`/`SharedExecutor` APIs must not silently change their +return types or ignore external writes. Compare resident compute and +Arrow-materializing API time separately. + +### Memory and fallback + +Use last-use analysis to release intermediate arrays and accumulate +gradient contributions as they become available. Bound allocations using +known shapes, include conversion/output costs and avoid retaining every +branch's large result. MLX's lazy graph must also be bounded and evaluated +at deliberate points rather than growing across unlimited training steps. + +Failed runtime guards must fall back before committing parameter updates. +Retain an unchanged input snapshot or otherwise transactional ownership +until guards pass. An unsupported region is a normal relational boundary, +not permission to change NULL/duplicate/default semantics. Use float64 CPU +as the default exact-precision path. Do not weaken tolerances or silently +alter AdamW to hide the existing float32 parity failure. + +## Execution sequence + +1. **Expose the actual bottlenecks.** Add opt-in per-stage timings and + dense/fallback reasons, row/byte counts, coordinate packing costs, + compiled-plan counts and synchronization/conversion counters. Profile + one representative backward pass, not repeated five-step large runs. + Save top offenders and minimal SQL reproductions. Native opaque nodes + require internal metrics; their outer node count is not an operator + count comparable with plain DataFusion. +2. **Build dense storage and lowering on CPU.** Add exact domains, compact + arrays and validated operator lowering. First cover ordinary matrix + gradients and reconstructed normalized/residual operands. Route those + through actual SQL/ddx plans. Keep unsupported operators on fallback. + Test different names, shapes and coordinate labels to prevent accidental + model-specific behavior. +3. **Integrate stable staged execution.** Connect dense values to ddx's + compiled backend hooks and stable slots. Share intermediates and + eliminate per-step replanning for unchanged shapes. Verify repeated + updates, domain changes, empty/NULL results and snapshots. Preserve + dynamic-filter safeguards and fresh join state for relational regions. +4. **Execute dense regions with MLX.** Compile supported graphs over + explicit parameters, retain arrays between their operators, evaluate + at required boundaries and record all copies/casts/synchronizations. + Add explicit owned resident state only after the compatible path works. +5. **Expand supported backward patterns.** Use the measured remaining + fallback list to add domain fill, scatter/reduction, nonlinear maps and + broader contraction recognition. Recheck memory liveness and accuracy + after each coherent change; avoid an unbounded optimizer rewrite search. +6. **Validate and benchmark.** Run focused public-result tests, then the + full local suite and native tests appropriate to changed crates. Start + with tiny/213K five-step cases and normalized-gradient reproductions, + then 10.7M and 123.6M context-16 five-step comparisons. Attempt larger + contexts progressively only after short-context correctness and memory + are established. Run hosts serially without builds/profilers during + timings. Compare unchanged JAX/PyTorch controls; keep failed/bounded runs + out of speedup tables. Include DuckDB controls where bounded runs + complete. Write measured speedups and remaining limitations. + +These are ordered milestones, not permission to stop after a fast kernel +microbenchmark. Adapt implementation choices to measured blockers while +preserving the scope and correctness requirements. + +## Completion criteria + +- SQL and ddx remain the source of the computation and derivatives. +- Supported dense backward paths avoid product-sized relational + intermediates and per-element coordinate storage between their operators. +- Steady-state structural planning is independent of training-step count + for unchanged domains; value changes remain visible. +- GPU work is verified by execution counters. No per-operator host + materialization/synchronization remains inside a supported dense region. +- Tests cover values, dimensions, coordinates, attributes, bags, NULLs, + empty groups, defaults, masked exceptional values, grad/VJP/JVP, typed + refusals, re-registration and repeated updates. Existing tolerances stay. +- Reports include cold/compile/setup and warm times, import/export costs, + memory, actual device work, fallback coverage and all five checked steps. + +An initial engineering target is at least a 10× reduction of the current +123M warm-step time, followed by reducing the remaining gap to native +array execution. This is a target, not a measured or promised speedup. +Separate compatible Arrow API costs from optional resident-state costs. + +## Resume details + +Work directly in `/Users/cromwellian/xql/xarray-sql`, preserving existing +uncommitted work. Read its `AGENTS.md`; when modifying the EinFold patch, +read `/private/tmp/xql-einfold/AGENTS.md` and the relevant design sections. +The user explicitly authorizes local implementation of this plan after +06:00 Pacific. No publishing, merging or messages to maintainers are +requested. Keep design changes and the reproducible local patch together. + +The checkout uses `.venv` Python 3.12.14, DataFusion Python 55.0.0, +ddxdb 0.4.6 and custom EinFold `0.2.0+xqlmlx.3`, pinned to +`39b25210d3b39027ef91584383fb114dfa8a4a48`. Native Rust DataFusion is 55.2.0. +Use `PYTHONPATH=$PWD`; installed Python sources can differ. Do not run +`uv sync`, which can replace the custom wheel. Metal access requires the +appropriate execution permissions. Use `MLX_ENABLE_TF32=0` before import. + +EinFold sources/build cache: `/private/tmp/xql-einfold`; clean pinned tree: +`/private/tmp/xql-einfold-clean`. Isolated Rust toolchain: +`PATH=/private/tmp/xql-cargo/bin:$PATH`, +`RUSTUP_HOME=/private/tmp/xql-rustup`, +`CARGO_HOME=/private/tmp/xql-cargo`. Build EinFold from its `python` +directory with maturin, locked dependencies and the checkout's Python. +Retain all new native changes in the EinFold repository and verify that the pinned build +applies cleanly to the pinned tree. Existing xarray-sql native builds must +retain Flight support. No build processes remain paused. + +The upstream nanoGPT model and isolated PyTorch 2.15.0 environment are in +`/private/tmp/xql-upstream-nanogpt` and `/private/tmp/xql-framework-venv`. +The latter needs the checkout's `.venv/lib/python3.12/site-packages` on +`PYTHONPATH` for JAX/NumPy. Baseline scripts are +`benchmarks/apple/run_nanogpt_sizes.py` and +`benchmarks/apple/compare_nanogpt_frameworks.py`. The last full local suite +passed 637 tests (277 skipped, 102 integration deselected); native suites +passed 60, with four upstream ignored tests and one unchanged plain- +DataFusion ARM NaN-sign assumption excluded. Keep those qualifications. diff --git a/accelerators/einfold_mlx/resident-training-plan.md b/accelerators/einfold_mlx/resident-training-plan.md new file mode 100644 index 00000000..7ac1e8ec --- /dev/null +++ b/accelerators/einfold_mlx/resident-training-plan.md @@ -0,0 +1,132 @@ +# Persistent GPU execution of SQL training + +The execution changes in this plan are implemented in +[`TrainingSession`](../../xarray_sql/training.py) and the optional +EinFold `.6` / ddx `.3` builds. The original plan below records their +implementation order and acceptance criteria. Whole-step compilation, +persistent state, shared primal loss, GPU dropout and explicit scalar +validation are available in the Shakespeare trainer. Final validation and +measurements are recorded in the resident execution report. + +[ddx](https://github.com/xqlsystems/ddx) differentiates SQL; +[EinFold](https://github.com/xqlsystems/einfold) lowers numerical query +plans to arrays; [MLX](https://github.com/ml-explore/mlx) executes the +supported array arithmetic on Apple silicon. The goal is to retain this +pipeline while keeping numerical training state in MLX between updates. + +## Current boundaries + +The [Shakespeare trainer](../../benchmarks/ml/03_shakespeare.py) calls a +loss query, exports all gradients, registers them again, exports the +clipping coefficient, then exports and registers weights and AdamW +moments. Each `DenseExecutor._invoke` clears its transient tensor store +and creates a new computed-value cache. These are deliberate properties +of the current live-Dataset/Arrow snapshot API. + +The current [measurements](../../benchmarks/apple/shakespeare-training.md) +are 2.024 s per warm SQL GPU update versus 0.02935 s in native PyTorch MPS +for the 10.7M-parameter recipe. This plan does not attribute the entire +difference to exports; Python plan interpretation, memory allocation, +kernel dispatch and numerical operation selection also remain. + +Apple silicon uses [unified memory](https://ml-explore.github.io/mlx/build/html/usage/unified_memory.html). +The costs to eliminate are conversion, packing, allocation and forced +synchronization, rather than assuming every CPU/GPU boundary transfers +data over PCIe. + +## Implementation order + +1. **Retain device tables within a complete training step.** Add an + opt-in internal execution path returning compact tables whose numerical + columns are MLX arrays. Execute gradients, global clipping and every + existing SQL AdamW statement in one invocation before exporting results. + Bind temporary gradient/state names to these actual device tables. + Schema providers supply planning metadata and must never substitute for + numerical evaluation or required validation. This first phase can still + export the final updated state, establishing correctness independently + of persistence. + +2. **Expose the primal value from ddx's backward program.** In the pinned + ddx release, `BackwardProgram.value` already identifies the loss output; + the Python binding exposes it. Read that output before transient tables + are released, together with the existing gradients. If the high-level + SQL preparation API does not expose the corresponding output handle, + add that bridge or execute the public backward program directly. + Do not assume adding an independent SELECT to a batch automatically + shares its forward computation. The loss and backward must use the same + saved forward values and dropout masks. + +3. **Add explicit persistent state ownership.** A new training-state + object owns weight and moment arrays, dtype/layout metadata, coordinate + domains, version and RNG state. Import the initial arrays once. Bind + the next update's SQL table names directly to the preceding update's + device outputs. Retain persistent state separately from transient + activations/cotangents; release temporary graphs after each step. + Publish all updated tables together only after checks succeed, retaining + the previous version on failure. Do not silently change the existing + API's handling of caller edits to NumPy buffers. Explicit replacement, + snapshot/export and close operations define the new ownership boundary. + +4. **Compile the complete numerical program.** Lower the prepared + forward/backward/optimizer DAG to stable array instructions once rather + than walking hundreds of DataFusion plans each update. Compile that + program with explicit weights, moments, batch, learning-rate/step scalars + and RNG inputs. Outputs include next state and the primal loss. + Cache by SQL/program identity, schema, coordinate/layout facts, shapes, + precision and UDF versions; invalidate when these change. Changing + scalar values must not trigger a compilation or capture stale weights. + Whole-program compilation can fuse pointwise work and share expressions; + it does not imply that every GEMM becomes one GPU kernel. + +5. **Keep masks, reductions and checks on the device.** Generate dropout + masks on MLX with checkpointed RNG state. Keep an explicit-mask mode for + exact comparisons against the current trainer/native oracle. Compute + the global gradient norm and all AdamW arithmetic on GPU. Prove static + coordinate uniqueness/domain properties from immutable metadata where + possible. Otherwise execute checks on actual device values and read a + scalar status before publishing an update. Preserve errors, NULL and + row-domain behavior; do not skip validation because providers have zero + rows. Unsupported plans decline before persistent state changes. + +6. **Make synchronization explicit.** Evaluate the step's device outputs + together and release references to previous computational histories. + Read the loss when logging rather than before backward. Export complete + arrays only for checkpoints, explicit inspection and correctness checks. + A required scalar validation result may still synchronize once per + update. Timed benchmarks must wait for device completion; an asynchronous + submission time is not an update time. + +MLX documents [whole-training-graph compilation and explicit state](https://ml-explore.github.io/mlx/build/html/usage/compile.html#compiling-training-graphs), +including RNG state for dropout. Its examples use MLX differentiation; +this executor would instead compile the array operations produced from +ddx's SQL derivatives. The same compilation mechanism can be used without +replacing the derivative source. MLX also documents how +[scalar extraction forces evaluation](https://ml-explore.github.io/mlx/build/html/usage/lazy_evaluation.html). + +## Contract and acceptance criteria + +- SQL remains the source of the forward model, clipping and optimizer; + ddx remains the source of derivatives. Dense-only assumptions, precision + and ownership are explicit. Ordinary Arrow-returning SQL APIs continue + to produce independent snapshots. +- Loss, raw gradients, weights, both moments and evaluation results match + the existing float64 CPU and native reference checks at the established + float32 tolerance. Exports for checks occur outside the performance timer. +- Tests cover mask reuse, fresh parameters after every update, all-table + atomic publication, invalid row/domain checks, UDF/schema invalidation, + independent snapshots, explicit host replacement and checkpoint resume. +- After initialization, the resident step performs no parameter-, + moment- or gradient-sized host exports/imports and creates no Arrow + gradient/update tables. Batch data and optional scalar control/status + are the remaining routine boundaries. Count both bytes and sync events. +- Report staged timings: current baseline; combined device step; reused + primal loss; persistent state; compiled program. Measure compilation + separately. Compare five checked steps and at least 100 steady-state + updates, including synchronized duration and peak memory. Confirm bounded + memory and unchanged SQL/gradient semantics before claiming speedups. + +If a substantial gap remains after these boundaries are removed, profile +the resulting compiled execution. Candidate work includes backward +contraction scheduling, common expressions, activation liveness and fused +normalization/attention kernels selected from proven plan patterns. +Native MPS parity cannot be predicted from the current measurements. diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py index 44f5c3cc..ce4edbb1 100644 --- a/benchmarks/ml/01_attention.py +++ b/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: diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py index 3dc477c4..d8e4ad1b 100644 --- a/benchmarks/ml/02_nanogpt.py +++ b/benchmarks/ml/02_nanogpt.py @@ -85,6 +85,7 @@ from __future__ import annotations import pathlib +import time import urllib.request import jax @@ -167,11 +168,33 @@ def layer_norm(x: str, out: str, gain: str, layer: int | None = None) -> str: GELU_SCALE = np.sqrt(2 / np.pi) -def attention(layer: int, x: str, out: str) -> str: +def attention( + layer: int, x: str, out: str, *, dropout=False, grouped_softmax=False +) -> str: """nanoGPT's ``CausalSelfAttention``: each position of ``x`` attends to itself and the positions before it, in each head, and ``out`` is the heads' outputs, projected back to the residual stream.""" l = f"l{layer}" + mask = ( + f"JOIN {l}_adrop d ON a.seq=d.seq AND a.h=d.h AND a.i=d.i AND a.j=d.j" + if dropout + else "" + ) + scale = " * d.scale" if dropout else "" + projected = f"{out}_raw" if dropout else out + weights = f""" +{l}_weights AS ( + SELECT seq, h, i, j, + exp(s - MAX(s) OVER (PARTITION BY seq, h, i)) AS p + FROM {l}_scores),""" + if grouped_softmax: + weights = f""" +{l}_maximum AS ( + SELECT seq, h, i, ddx_stop_gradient(MAX(s)) AS m FROM {l}_scores GROUP BY seq, h, i), +{l}_weights AS ( + SELECT a.seq, a.h, a.i, a.j, exp(a.s - m.m) AS p + FROM {l}_scores a JOIN {l}_maximum m + ON a.seq=m.seq AND a.h=m.h AND a.i=m.i),""" projections = "".join( f""" {l}_{name} AS ( @@ -188,30 +211,37 @@ def attention(layer: int, x: str, out: str) -> str: FROM {l}_q q JOIN {l}_k k ON q.seq = k.seq AND q.h = k.h AND q.c = k.c WHERE q.pos >= k.pos GROUP BY q.seq, q.h, q.pos, k.pos), -{l}_weights AS ( - -- The max only keeps exp() from overflowing: it cancels out of p / z. - SELECT seq, h, i, j, - exp(s - MAX(s) OVER (PARTITION BY seq, h, i)) AS p - FROM {l}_scores), +{weights} {l}_norm AS ( SELECT seq, h, i, SUM(p) AS z FROM {l}_weights GROUP BY seq, h, i), {l}_heads AS ( - SELECT a.seq, a.i AS pos, a.h, v.c, SUM(a.p / n.z * v.val) AS val + SELECT a.seq, a.i AS pos, a.h, v.c, SUM(a.p / n.z * v.val{scale}) AS val FROM {l}_weights a JOIN {l}_v v ON a.seq = v.seq AND a.h = v.h AND a.j = v.pos JOIN {l}_norm n ON a.seq = n.seq AND a.h = n.h AND a.i = n.i + {mask} GROUP BY a.seq, a.i, a.h, v.c), -{out} AS ( +{projected} AS ( SELECT a.seq, a.pos, p.e, SUM(a.val * p.w) AS val FROM {l}_heads a JOIN wo p ON a.h = p.h AND a.c = p.c WHERE p.layer = {layer} - GROUP BY a.seq, a.pos, p.e),""" + GROUP BY a.seq, a.pos, p.e),""" + ( + _dropout(projected, out, f"{l}_rdrop") if dropout else "" + ) -def mlp(layer: int, x: str, out: str) -> str: +def mlp( + layer: int, x: str, out: str, *, exact_gelu=False, dropout=False +) -> str: """nanoGPT's ``MLP``: each position of ``x`` through a hidden layer four times as wide, and GELU, back to the residual stream as ``out``.""" l = f"l{layer}" + activation = ( + "0.5 * val * (1 + erf(val / sqrt(2.0)))" + if exact_gelu + else f"0.5 * val * (1 + tanh({GELU_SCALE} * (val + 0.044715 * val * val * val)))" + ) + projected = f"{out}_raw" if dropout else out return f""" {l}_fc AS ( SELECT x.seq, x.pos, p.f, SUM(x.val * p.w) AS val @@ -220,14 +250,22 @@ def mlp(layer: int, x: str, out: str) -> str: GROUP BY x.seq, x.pos, p.f), {l}_hidden AS ( SELECT seq, pos, f, - 0.5 * val * (1 + tanh({GELU_SCALE} * (val + 0.044715 * val * val * val))) - AS val + {activation} AS val FROM {l}_fc), -{out} AS ( +{projected} AS ( SELECT x.seq, x.pos, p.e, SUM(x.val * p.w) AS val FROM {l}_hidden x JOIN w_proj p ON x.f = p.f WHERE p.layer = {layer} - GROUP BY x.seq, x.pos, p.e),""" + GROUP BY x.seq, x.pos, p.e),""" + ( + _dropout(projected, out, f"{l}_mdrop") if dropout else "" + ) + + +def _dropout(x, out, mask): + return f""" +{out} AS ( + SELECT x.seq, x.pos, x.e, x.val * d.scale AS val + FROM {x} x JOIN {mask} d ON x.seq=d.seq AND x.pos=d.pos AND x.e=d.e),""" def residual(x: str, y: str, out: str) -> str: @@ -238,38 +276,71 @@ def residual(x: str, y: str, out: str) -> str: FROM {x} x JOIN {y} y ON x.seq = y.seq AND x.pos = y.pos AND x.e = y.e),""" -def block(layer: int, x: str, out: str) -> str: +def block( + layer: int, + x: str, + out: str, + *, + exact_gelu=False, + dropout=False, + grouped_softmax=False, +) -> str: """nanoGPT's ``Block``: ``mid = x + attention(ln_1(x))``, then ``out = mid + mlp(ln_2(mid))``.""" l = f"l{layer}" return "".join( [ layer_norm(x, f"{l}_ln1", "ln_1", layer), - attention(layer, f"{l}_ln1", f"{l}_attn"), + attention( + layer, + f"{l}_ln1", + f"{l}_attn", + dropout=dropout, + grouped_softmax=grouped_softmax, + ), residual(x, f"{l}_attn", f"{l}_mid"), layer_norm(f"{l}_mid", f"{l}_ln2", "ln_2", layer), - mlp(layer, f"{l}_ln2", f"{l}_mlp"), + mlp( + layer, + f"{l}_ln2", + f"{l}_mlp", + exact_gelu=exact_gelu, + dropout=dropout, + ), residual(f"{l}_mid", f"{l}_mlp", out), ] ) -def model() -> str: +def model( + *, exact_gelu=False, dropout=False, view="logits", grouped_softmax=False +) -> str: """The model, as the view ``logits``: for each position of each sequence in ``batch``, a score for every character that could come next. A DataFusion view reads the tables that existed when it was created (DuckDB's looks them up by name), so it is recreated whenever the batch or the parameters are.""" blocks = "".join( - block(layer, f"x{layer}", f"x{layer + 1}") for layer in range(LAYERS) + block( + layer, + f"x{layer}", + f"x{layer + 1}", + exact_gelu=exact_gelu, + dropout=dropout, + grouped_softmax=grouped_softmax, + ) + for layer in range(LAYERS) ) + embedded = "embed_raw" if dropout else "x0" + embed_dropout = _dropout(embedded, "x0", "embed_drop") if dropout else "" return f""" -CREATE OR REPLACE VIEW logits AS -WITH x0 AS ( +CREATE OR REPLACE VIEW {view} AS +WITH {embedded} AS ( SELECT b.seq, b.pos, te.e, te.w + pe.w AS val FROM batch b JOIN wte te ON b.tok = te.tok JOIN wpe pe ON b.pos = pe.pos AND te.e = pe.e), +{embed_dropout} {blocks} {layer_norm(f"x{LAYERS}", "final", "ln_f").rstrip(",")} SELECT x.seq, x.pos, te.tok, SUM(x.val * te.w) AS val @@ -460,15 +531,7 @@ def init_params(vocab: int, rng) -> dict[str, np.ndarray]: """nanoGPT's initialization, from GPT-2: weights normal with standard deviation 0.02, or 0.02/√(2 × layers) for the two projections that add to the residual stream; layer norm gains 1.""" - sizes = { - "tok": vocab, - "pos": CONTEXT, - "layer": LAYERS, - "e": EMBED, - "h": HEADS, - "c": HEAD, - "f": HIDDEN, - } + sizes = _parameter_sizes(vocab) params = {} for name, dims in PARAMS.items(): shape = [sizes[d] for d in dims] @@ -481,6 +544,18 @@ def init_params(vocab: int, rng) -> dict[str, np.ndarray]: return params +def _parameter_sizes(vocab): + return { + "tok": vocab, + "pos": CONTEXT, + "layer": LAYERS, + "e": EMBED, + "h": HEADS, + "c": HEAD, + "f": HIDDEN, + } + + def as_table(name: str, w: np.ndarray, m=None, v=None) -> xr.Dataset: """A parameter as a table: its weights ``w``, and Adam's moments ``m`` and ``v`` when they are given.""" @@ -527,19 +602,62 @@ def main(): global LAYERS, EMBED, HEADS, HEAD, HIDDEN, CONTEXT parser = cli(__doc__.splitlines()[0], sizes=[4], size_name="batch") parser.add_argument("--steps", type=int, default=5, help="training steps") + parser.add_argument( + "--learning-rate", + type=float, + default=LEARNING_RATE, + help="SGD learning rate (AdamW uses its schedule)", + ) + parser.add_argument( + "--preset", + choices=["tiny", "demo", "shakespeare", "gpt2"], + default="tiny", + help="model dimensions; explicit size flags override them", + ) + parser.add_argument( + "--vocab-size", + type=int, + help="pad vocabulary (GPT-2 preset defaults to 50257; character inputs)", + ) + parser.add_argument( + "--forward-only", + action="store_true", + help="benchmark loss without parameter updates", + ) parser.add_argument( "--optimizer", choices=["sgd", "adamw"], default="sgd", help="SGD at a fixed learning rate (default), or nanoGPT's AdamW", ) - parser.add_argument("--layers", type=int, default=LAYERS, help="n_layer") - parser.add_argument("--embed", type=int, default=EMBED, help="n_embd") - parser.add_argument("--heads", type=int, default=HEADS, help="n_head") - parser.add_argument( - "--context", type=int, default=CONTEXT, help="block_size" - ) + parser.add_argument("--layers", type=int, help="n_layer") + parser.add_argument("--embed", type=int, help="n_embd") + parser.add_argument("--heads", type=int, help="n_head") + parser.add_argument("--context", type=int, help="block_size") args = parser.parse_args() + if not np.isfinite(args.learning_rate) or args.learning_rate <= 0: + parser.error("--learning-rate must be finite and positive") + preset = { + "tiny": (1, 16, 2, 16), + "demo": (1, 128, 4, 64), + "shakespeare": (6, 384, 6, 256), + "gpt2": (12, 768, 12, 1024), + }[args.preset] + for name, value in zip(("layers", "embed", "heads", "context"), preset): + if getattr(args, name) is None: + setattr(args, name, value) + if ( + min( + args.steps, + args.layers, + args.embed, + args.heads, + args.context, + *args.sizes, + ) + <= 0 + ): + parser.error("steps, model dimensions and batch sizes must be positive") if args.embed % args.heads: parser.error("--embed must be a multiple of --heads") LAYERS, EMBED, HEADS, CONTEXT = ( @@ -554,14 +672,43 @@ def main(): report = Report("nanogpt", "batch", args.reps) text = shakespeare() vocab = len(set(text)) + if args.vocab_size is None and args.preset == "gpt2": + args.vocab_size = 50257 + if args.vocab_size is not None: + if args.vocab_size < vocab: + parser.error( + "--vocab-size must cover the Shakespeare character vocabulary" + ) + vocab = args.vocab_size lines = len(model().splitlines()) + len(loss_statement().splitlines()) + sizes = _parameter_sizes(vocab) count = sum( - v.size for v in init_params(vocab, np.random.default_rng(0)).values() + int(np.prod([sizes[d] for d in dims])) for dims in PARAMS.values() ) + metadata = { + "layers": LAYERS, + "embed": EMBED, + "heads": HEADS, + "context": CONTEXT, + "vocab": vocab, + "parameters": count, + "optimizer": args.optimizer, + "sgd_learning_rate": args.learning_rate if not adamw else None, + "forward_only": args.forward_only, + } print( f" {LAYERS}-layer GPT, vocab {vocab}, {count:,} parameters: " f"the loss is {lines} lines of SQL" ) + print( + f" steps={args.steps}, optimizer={args.optimizer}, " + f"context={CONTEXT}, heads={HEADS}, embed={EMBED}, " + f"forward_only={args.forward_only}" + ) + if vocab != len(set(text)): + print( + " Padded vocabulary: inputs are character IDs, without GPT-2 BPE tokenization" + ) for size in args.sizes: init = init_params(vocab, np.random.default_rng(0)) @@ -572,18 +719,34 @@ def main(): ) for host in hosts(args.hosts): + load_seconds = 0.0 + model_seconds = 0.0 + + def load_table(name, values, host=host): + nonlocal load_seconds + start = time.perf_counter() + host.load(name, values) + load_seconds += time.perf_counter() - start + reference = dict(init) - moments = { - n: (np.zeros_like(w), np.zeros_like(w)) for n, w in init.items() - } + moments = ( + { + n: (np.zeros_like(w), np.zeros_like(w)) + for n, w in init.items() + } + if adamw + else {} + ) for name, values in init.items(): state = moments[name] if adamw else () - host.load(name, as_table(name, values, *state)) - host.load("schedule", lr) + load_table(name, as_table(name, values, *state)) + load_table("schedule", lr) for step, batch in data: - host.load("batch", batch) + load_table("batch", batch) + start = time.perf_counter() host.execute(model()) + model_seconds += time.perf_counter() - start tok, nxt = batch["tok"].values, batch["next"].values rows = report.time( @@ -595,10 +758,28 @@ def main(): ) want = jax_loss(reference, tok, nxt) report.check( - assert_close("loss", rows["l"][0].as_py(), want, rtol=1e-9) + assert_close( + "loss", + rows["l"][0].as_py(), + want, + rtol=host.rtol, + atol=host.atol, + ) ) report.note("loss", round(float(want), 6)) report.note("operators", host.operators(loss_statement())) + report.note("rtol", f"{host.rtol:.0e}") + report.note("atol", f"{host.atol:.0e}") + for key, value in metadata.items(): + report.note(key, value) + for key, value in host.measurements().items(): + report.note(key, value) + if args.forward_only: + report.note("load_seconds", load_seconds) + report.note("model_seconds", model_seconds) + report._write_csv(args.csv) + continue + report._write_csv(args.csv) train = f"train ({args.optimizer}), step {step}" if not host.differentiates: @@ -614,7 +795,7 @@ def main(): host, size, host.query_all, - step_statements(LEARNING_RATE), + step_statements(args.learning_rate), ) if updated is None: # Every later step starts from this one's weights. @@ -626,7 +807,7 @@ def main(): ) else: reference = { - name: w - LEARNING_RATE * grads[name] + name: w - args.learning_rate * grads[name] for name, w in reference.items() } @@ -641,12 +822,21 @@ def main(): f"{name}.{column}, step {step}", as_array(name, rows, column), values, + rtol=host.rtol, + atol=host.atol, ) ) - host.load( + load_table( name, xql.to_dataset(rows, dims=list(PARAMS[name])) ) report.check(max(errors)) + report.note("load_seconds", load_seconds) + report.note("model_seconds", model_seconds) + for key, value in metadata.items(): + report.note(key, value) + for key, value in host.measurements().items(): + report.note(key, value) + report._write_csv(args.csv) report.finish(args.csv) diff --git a/benchmarks/ml/03_shakespeare.py b/benchmarks/ml/03_shakespeare.py new file mode 100644 index 00000000..946c2a9a --- /dev/null +++ b/benchmarks/ml/03_shakespeare.py @@ -0,0 +1,714 @@ +"""Train nanoGPT's Shakespeare character model with SQL and ddx gradients. + +nanoGPT (https://github.com/karpathy/nanoGPT) supplies the training recipe +and prepared character data. The model, loss, gradient clipping and AdamW +updates execute SQL; EinFold (https://github.com/xqlsystems/einfold) lowers +the plans to compact CPU arrays or MLX (https://github.com/ml-explore/mlx) +GPU arrays. Python samples batches and dropout masks and writes checkpoints. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import math +import os +import resource +import subprocess +import time +from pathlib import Path + +import ddxdb +import numpy as np +import xarray as xr +from _hosts import DenseCPU, DenseGPU +from _shakespeare_torch import _TorchReference + +import xarray_sql as xql +from xarray_sql.math import register_erf +from xarray_sql.training import TrainingSession + +_SPEC = importlib.util.spec_from_file_location( + "_shakespeare_sql_model", Path(__file__).with_name("02_nanogpt.py") +) +_MODEL = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODEL) + +# The maximum cancels out of cross entropy. Spell its reduction as a join +# to avoid repeatedly traversing a shared model graph during window lowering. +# Holding this shared shift constant removes only a cancelling derivative. +_MODEL.LOSS = """ +predicted AS (SELECT seq,pos,tok,val FROM logits), +maximum AS (SELECT seq,pos,ddx_stop_gradient(MAX(val)) AS m FROM predicted GROUP BY seq,pos), +log_sum_exp AS ( + SELECT p.seq,p.pos,MAX(m.m)+ln(SUM(exp(p.val-m.m))) AS lse + FROM predicted p JOIN maximum m ON p.seq=m.seq AND p.pos=m.pos + GROUP BY p.seq,p.pos), +loss AS ( + SELECT AVG(s.lse-l.val) AS l FROM log_sum_exp s + JOIN batch b ON s.seq=b.seq AND s.pos=b.pos + JOIN predicted l ON l.seq=b.seq AND l.pos=b.pos AND l.tok=b.next) +""" + + +def _parser(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--backend", choices=["dense_cpu", "dense_gpu"], default="dense_gpu" + ) + p.add_argument( + "--data-dir", + type=Path, + default=Path.home() / "nanoGPT/data/shakespeare_char", + ) + p.add_argument("--out-dir", type=Path, default=Path("out-shakespeare-sql")) + p.add_argument( + "--resume", + type=Path, + help="resume an NPZ checkpoint including moments and RNGs", + ) + p.add_argument( + "--steps", + type=int, + help="cap this run's updates without shortening the LR schedule", + ) + p.add_argument( + "--verify-nanogpt", + type=Path, + help="check every gradient, update and moment against this nanoGPT checkout", + ) + for name, default in { + "batch_size": 12, + "block_size": 256, + "n_layer": 6, + "n_head": 6, + "n_embd": 384, + "gradient_accumulation_steps": 1, + "max_iters": 2000, + "lr_decay_iters": 2000, + "warmup_iters": 100, + "eval_interval": 200, + "eval_iters": 20, + "log_interval": 10, + "seed": 1337, + }.items(): + p.add_argument( + "--" + name, + "--" + name.replace("_", "-"), + type=int, + default=default, + ) + for name, default in { + "dropout": 0.2, + "learning_rate": 1e-3, + "min_lr": 1e-4, + "beta1": 0.9, + "beta2": 0.99, + "weight_decay": 0.1, + "grad_clip": 1.0, + }.items(): + p.add_argument( + "--" + name, + "--" + name.replace("_", "-"), + type=float, + default=default, + ) + p.add_argument("--cache-gib", type=float, default=16) + p.add_argument( + "--resident", + action="store_true", + help="keep SQL training state as owned device arrays", + ) + p.add_argument( + "--compile-training", + action="store_true", + help="compile the complete SQL-derived device step", + ) + p.add_argument( + "--gpu-dropout", + action="store_true", + help="generate dropout masks inside the device training step", + ) + p.add_argument( + "--profile-memory", + action="store_true", + help="sample current RSS outside training timers at log intervals", + ) + return p + + +def _learning_rate(it, args): + if it < args.warmup_iters: + return args.learning_rate * (it + 1) / (args.warmup_iters + 1) + if it > args.lr_decay_iters: + return args.min_lr + ratio = (it - args.warmup_iters) / (args.lr_decay_iters - args.warmup_iters) + return args.min_lr + 0.5 * (1 + math.cos(math.pi * ratio)) * ( + args.learning_rate - args.min_lr + ) + + +def _data(args): + # Match prepare.py's uint16 files and sorted character vocabulary without + # executing the pickle in meta.pkl. Check that the files encode input.txt. + text = (args.data_dir / "input.txt").read_text() + chars = sorted(set(text)) + stoi = {c: i for i, c in enumerate(chars)} + encoded = np.fromiter((stoi[c] for c in text), dtype=np.uint16) + split = int(len(encoded) * 0.9) + data = {} + for name, expected in ( + ("train", encoded[:split]), + ("val", encoded[split:]), + ): + values = np.memmap( + args.data_dir / f"{name}.bin", dtype=np.uint16, mode="r" + ) + if not np.array_equal(values, expected): + raise ValueError( + f"{name}.bin differs from nanoGPT's character preparation of input.txt" + ) + if len(values) <= args.block_size: + raise ValueError(f"{name} split is shorter than the context") + data[name] = values + return chars, data + + +def _batch(values, args, rng): + starts = rng.integers(len(values) - args.block_size, size=args.batch_size) + tok = np.stack([values[i : i + args.block_size] for i in starts]).astype( + np.int64 + ) + nxt = np.stack( + [values[i + 1 : i + args.block_size + 1] for i in starts] + ).astype(np.int64) + return xr.Dataset( + {"tok": (("seq", "pos"), tok), "next": (("seq", "pos"), nxt)}, + coords={ + "seq": np.arange(args.batch_size), + "pos": np.arange(args.block_size), + }, + ) + + +def _masks(args, rng): + def mask(dims, shape): + scale = (rng.random(shape) >= args.dropout).astype(np.float64) / ( + 1 - args.dropout + ) + return xr.Dataset( + {"scale": (dims, scale)}, + coords={d: np.arange(n) for d, n in zip(dims, shape)}, + ) + + shape = (args.batch_size, args.block_size, args.n_embd) + result = {"embed_drop": mask(("seq", "pos", "e"), shape)} + for layer in range(args.n_layer): + result[f"l{layer}_adrop"] = mask( + ("seq", "h", "i", "j"), + (args.batch_size, args.n_head, args.block_size, args.block_size), + ) + for suffix in ("rdrop", "mdrop"): + result[f"l{layer}_{suffix}"] = mask(("seq", "pos", "e"), shape) + return result + + +def _clip_sql(args): + if args.grad_clip == 0: + return "SELECT 1.0 AS c" + every_grad = " UNION ALL ".join( + f"SELECT g FROM {t}_grad" for t in _MODEL.PARAMS + ) + return f"""WITH norm AS (SELECT sqrt(SUM(g*g)) AS n FROM ({every_grad})), + coefficient AS (SELECT {args.grad_clip}/(n+1e-6) AS c FROM norm) + SELECT CASE WHEN c < 1.0 THEN c ELSE 1.0 END AS c FROM coefficient""" + + +def _update_sql(name, args): + dims = _MODEL.PARAMS[name] + keys = ", ".join(f"p.{d}" for d in dims) + same = lambda alias: " AND ".join(f"p.{d}={alias}.{d}" for d in dims) + decay = 0.0 if name.startswith("ln_") else args.weight_decay + return f"""WITH moments AS ( + SELECT {keys}, p.w, + {args.beta1}*s.m + (1.0-{args.beta1})*g.g*c.c AS m, + {args.beta2}*s.v + (1.0-{args.beta2})*power(g.g*c.c,2) AS v + FROM {name} p JOIN {name}_grad g ON {same("g")} + JOIN {name}_state s ON {same("s")} CROSS JOIN clip c) + SELECT {keys}, + p.w*(1.0-t.lr*{decay}) - t.lr*(p.m/t.b1)/(sqrt(p.v/t.b2)+1e-8) AS w, + p.m, p.v FROM moments p CROSS JOIN train_step t ORDER BY {keys}""" + + +def _save( + path, + params, + moments, + iteration, + args, + chars, + data_rng, + mask_rng, + best, + skip_eval=False, +): + metadata = { + "iteration": iteration, + "config": vars(args) + | { + "bias": False, + "dtype": "float32" if args.backend == "dense_gpu" else "float64", + }, + "chars": chars, + "data_rng": data_rng.bit_generator.state, + "mask_rng": mask_rng.bit_generator.state, + "best_val_loss": best, + "skip_eval": skip_eval, + } + arrays = { + f"{name}_{col}": value + for name, w in params.items() + for col, value in zip(("w", "m", "v"), (w, *moments[name])) + } + temporary = path.with_suffix(".tmp") + with temporary.open("wb") as handle: + np.savez( + handle, + metadata=np.array(json.dumps(metadata, default=str)), + **arrays, + ) + temporary.replace(path) + + +def main(): + parser = _parser() + args = parser.parse_args() + if args.gpu_dropout: + args.resident = True + if args.backend != "dense_gpu": + parser.error("--gpu-dropout requires dense_gpu") + if args.compile_training: + args.resident = True + if args.backend != "dense_gpu": + parser.error("--compile-training requires dense_gpu") + if ( + min( + args.batch_size, + args.block_size, + args.n_layer, + args.n_head, + args.n_embd, + args.eval_interval, + args.eval_iters, + args.log_interval, + ) + < 1 + or args.n_embd % args.n_head + ): + parser.error( + "model sizes and evaluation/logging intervals must be positive, with n_embd divisible by n_head" + ) + if args.gradient_accumulation_steps != 1: + parser.error( + "this recipe currently supports gradient_accumulation_steps=1" + ) + if ( + not 0 <= args.beta1 < 1 + or not 0 <= args.beta2 < 1 + or not all( + math.isfinite(v) and v >= 0 + for v in ( + args.learning_rate, + args.min_lr, + args.weight_decay, + args.grad_clip, + ) + ) + or not math.isfinite(args.cache_gib) + or args.cache_gib <= 0 + ): + parser.error( + "require betas in [0,1), finite nonnegative optimizer values and a positive cache size" + ) + if ( + not 0 <= args.dropout < 1 + or not 0 <= args.warmup_iters < args.lr_decay_iters + or args.max_iters < 0 + or (args.steps is not None and args.steps < 1) + ): + parser.error( + "require 0 <= dropout < 1, 0 <= warmup_iters < lr_decay_iters and positive steps" + ) + if args.backend == "dense_gpu" and os.environ.get("MLX_ENABLE_TF32") != "0": + parser.error( + "GPU execution requires MLX_ENABLE_TF32=0 before Python starts" + ) + if "erf" not in ddxdb.supported_functions(): + parser.error( + "exact GELU needs the local ddx math build: python accelerators/ddx_math/build.py" + ) + _MODEL.LAYERS, _MODEL.EMBED, _MODEL.HEADS = ( + args.n_layer, + args.n_embd, + args.n_head, + ) + _MODEL.HEAD, _MODEL.HIDDEN, _MODEL.CONTEXT = ( + args.n_embd // args.n_head, + 4 * args.n_embd, + args.block_size, + ) + chars, data = _data(args) + data_rng, mask_rng = ( + np.random.default_rng(args.seed), + np.random.default_rng(args.seed + 1), + ) + params = _MODEL.init_params(len(chars), np.random.default_rng(args.seed)) + moments = { + name: (np.zeros_like(w), np.zeros_like(w)) for name, w in params.items() + } + iteration, best, skip_eval = 0, math.inf, False + if args.resume: + with np.load(args.resume, allow_pickle=False) as saved: + meta = json.loads(str(saved["metadata"])) + ignored = { + "steps", + "resume", + "out_dir", + "data_dir", + "verify_nanogpt", + "cache_gib", + "max_iters", + "profile_memory", + "resident", + "compile_training", + } + for key, value in vars(args).items(): + if ( + key not in ignored + and meta["config"].get(key, False) != value + ): + parser.error(f"checkpoint configuration differs for {key}") + if meta["chars"] != chars: + parser.error("checkpoint vocabulary differs from input.txt") + for name, w in params.items(): + for col in ("w", "m", "v"): + if saved[f"{name}_{col}"].shape != w.shape: + parser.error( + f"checkpoint shape differs for {name}_{col}" + ) + params[name] = saved[f"{name}_w"].copy() + moments[name] = tuple( + saved[f"{name}_{c}"].copy() for c in ("m", "v") + ) + iteration, best, skip_eval = ( + meta["iteration"], + meta["best_val_loss"], + meta["skip_eval"], + ) + data_rng.bit_generator.state = meta["data_rng"] + mask_rng.bit_generator.state = meta["mask_rng"] + if args.verify_nanogpt and args.resume: + parser.error( + "reference verification starts from scratch; resume preserves SQL optimizer state" + ) + os.environ["XQL_DENSE_CACHE_BYTES"] = str(int(args.cache_gib * (1 << 30))) + host = DenseGPU() if args.backend == "dense_gpu" else DenseCPU() + register_erf(host.ctx) + for name, w in params.items(): + host.load(name, _MODEL.as_table(name, w)) + host.load( + name + "_state", + _MODEL.as_table(name, w, *moments[name]).drop_vars("w"), + ) + # Schema registration must not advance the checkpointed RNG streams. + host.load("batch", _batch(data["train"], args, np.random.default_rng(0))) + for name, value in _masks(args, np.random.default_rng(0)).items(): + host.load(name, value) + host.execute( + _MODEL.model(exact_gelu=True, dropout=True, grouped_softmax=True) + ) + # Register the second view directly: both graphs persist across evaluation. + eval_model = _MODEL.model( + exact_gelu=True, view="eval_logits", grouped_softmax=True + ) + host.executor.register_view("eval_logits", eval_model.split(" AS\n", 1)[1]) + resident = None + resident_step = None + if args.resident: + resident = TrainingSession( + host.ctx, + precision=host.dense_precision, + max_cached_bytes=int(args.cache_gib * (1 << 30)), + inline_ctes=host.executor.inline_ctes, + ) + for name, query in host.executor._views.items(): + resident.executor.register_view(name, query) + for name in host._tables: + resident.bind(name) + if args.compile_training: + resident.index_domains( + "batch", {"tok": ("wte", "tok"), "next": ("wte", "tok")} + ) + resident.bind("batch") + if args.gpu_dropout: + resident.dropout( + [name for name in host._tables if name.endswith("drop")], + probability=args.dropout, + seed=args.seed + 1, + step=iteration, + ) + updates_sql = { + name + "_next": _update_sql(name, args) for name in params + } + assignments = {} + for name, dims in _MODEL.PARAMS.items(): + assignments[name] = (name + "_next", [*dims, "w"]) + assignments[name + "_state"] = (name + "_next", [*dims, "m", "v"]) + resident_step = resident.prepare_step( + gradients=dict( + zip( + (name + "_grad" for name in params), + _MODEL.gradient_statements(), + ) + ), + stages={"clip": _clip_sql(args)}, + updates=updates_sql, + assignments=assignments, + compile=args.compile_training, + ) + + def load(name, value): + if resident is None: + host.load(name, value) + else: + resident.bind(name, value) + + def download(): + if resident is None: + return + tables = resident.export(*params, *(name + "_state" for name in params)) + for name in params: + params[name] = _MODEL.as_array(name, tables[name]) + state = xql.to_dataset( + tables[name + "_state"], dims=list(_MODEL.PARAMS[name]) + ) + moments[name] = (state.m.values, state.v.values) + + train_loss = host.executor.prepare(_MODEL.loss_statement()) + eval_loss = host.executor.prepare( + _MODEL.loss_statement().replace("FROM logits", "FROM eval_logits") + ) + gradients = host.executor.prepare_all(_MODEL.gradient_statements()) + clip = host.executor.prepare(_clip_sql(args)) + updates = { + name: host.executor.prepare(_update_sql(name, args)) for name in params + } + reference = ( + _TorchReference(args, params, len(chars)) + if args.verify_nanogpt + else None + ) + args.out_dir.mkdir(parents=True, exist_ok=True) + print( + json.dumps( + { + "backend": args.backend, + "parameters": sum(w.size for w in params.values()), + "vocabulary": len(chars), + "config": vars(args), + }, + default=str, + ), + flush=True, + ) + # train.py tests iter_num > max_iters after each update, so iteration + # max_iters is included. --steps provides an exact short-run update cap. + end = args.max_iters + 1 + if args.steps is not None: + end = min(end, iteration + args.steps) + with (args.out_dir / "metrics.jsonl").open( + "a" if args.resume else "w" + ) as log: + while iteration < end: + if iteration % args.eval_interval == 0 and not skip_eval: + losses = {} + for split in ("train", "val"): + measured = [] + for _ in range(args.eval_iters): + batch = _batch(data[split], args, data_rng) + load("batch", batch) + value = float( + ( + eval_loss() + if resident is None + else resident.query( + _MODEL.loss_statement().replace( + "FROM logits", "FROM eval_logits" + ) + ) + )["l"][0].as_py() + ) + if reference: + reference.evaluate(batch, value) + measured.append(value) + losses[split] = float(np.mean(measured)) + event = dict(kind="eval", iteration=iteration, **losses) + print(json.dumps(event), flush=True) + log.write(json.dumps(event) + "\n") + log.flush() + if losses["val"] < best: + best = losses["val"] + if iteration > 0: + download() + _save( + args.out_dir / "best.npz", + params, + moments, + iteration, + args, + chars, + data_rng, + mask_rng, + best, + skip_eval=True, + ) + skip_eval = False + batch = _batch(data["train"], args, data_rng) + masks = {} if args.gpu_dropout else _masks(args, mask_rng) + load("batch", batch) + for name, value in masks.items(): + load(name, value) + lr = _learning_rate(iteration, args) + load( + "train_step", + xr.Dataset( + { + "lr": xr.DataArray(lr), + "b1": xr.DataArray(1 - args.beta1 ** (iteration + 1)), + "b2": xr.DataArray(1 - args.beta2 ** (iteration + 1)), + } + ), + ) + started = time.perf_counter() + if resident is None: + loss = float(train_loss()["l"][0].as_py()) + grad_rows = gradients() + grads = { + name: xql.to_dataset(rows, dims=list(_MODEL.PARAMS[name])) + for name, rows in zip(params, grad_rows) + } + for name, ds in grads.items(): + host.load(name + "_grad", ds) + host.load("clip", xql.to_dataset(clip(), dims=[])) + # Compute every update before replacing any current weight. + next_tables = { + name: xql.to_dataset(run(), dims=list(_MODEL.PARAMS[name])) + for name, run in updates.items() + } + for name, ds in next_tables.items(): + params[name] = ds.w.values + moments[name] = (ds.m.values, ds.v.values) + host.load(name, ds[["w"]]) + host.load(name + "_state", ds[["m", "v"]]) + else: + result = resident_step() + loss = result.loss_value() + seconds = time.perf_counter() - started + if reference: + if resident is not None: + if args.gpu_dropout: + mask_tables = result.export( + *( + name + for name in host._tables + if name.endswith("drop") + ) + ) + masks = { + name: xql.to_dataset( + rows, + dims=[ + f for f in rows.column_names if f != "scale" + ], + ) + for name, rows in mask_tables.items() + } + grad_tables = result.export( + *(name + "_grad" for name in params) + ) + grads = { + name: xql.to_dataset( + grad_tables[name + "_grad"], + dims=list(_MODEL.PARAMS[name]), + ) + for name in params + } + download() + reference.step(batch, masks, lr, loss, grads, params, moments) + event = { + "kind": "train", + "iteration": iteration, + "loss": loss, + "lr": lr, + "seconds": seconds, + "tokens_per_second": args.batch_size + * args.block_size + / seconds, + "verified": reference is not None, + } + if resident is not None: + event["device_step"] = result.measurements + if args.profile_memory and ( + iteration % args.log_interval == 0 or iteration + 1 == end + ): + event["rss_bytes"] = ( + int( + subprocess.check_output( + ["ps", "-p", str(os.getpid()), "-o", "rss="], + text=True, + ).strip() + ) + * 1024 + ) + event["peak_rss_bytes"] = resource.getrusage( + resource.RUSAGE_SELF + ).ru_maxrss + if ( + resident is not None + and resident.executor._runtime is not None + ): + event["mlx_memory"] = resident.executor._runtime.memory() + log.write(json.dumps(event) + "\n") + log.flush() + if iteration % args.log_interval == 0 or iteration + 1 == end: + print(json.dumps(event), flush=True) + iteration += 1 + download() + _save( + args.out_dir / "last.npz", + params, + moments, + iteration, + args, + chars, + data_rng, + mask_rng, + best, + ) + print( + json.dumps( + { + "checkpoint": str(args.out_dir / "last.npz"), + "next_iteration": iteration, + "measurements": host.measurements() + if resident is None + else resident.statistics(), + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index ee632f96..651dd8eb 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -1,5 +1,13 @@ # Machine learning in SQL +The optional `dense_cpu` and `dense_gpu` hosts retain compact arrays across +SQL and gradient stages using the local +[EinFold integration](../../accelerators/einfold_mlx/README.md). +`dense_gpu` explicitly permits float32 arithmetic and requires +`MLX_ENABLE_TF32=0`. The hosts use identical model SQL, ddx derivatives and +existing reference checks. See the +[dense execution measurements](../apple/dense-execution.md). + These scripts write machine-learning models as SQL over tables that [xarray-sql](../../README.md) registers from xarray Datasets, and time the same SQL on several engines. Two related projects shape them: @@ -11,7 +19,7 @@ same SQL on several engines. Two related projects shape them: are its [target demos](https://github.com/xqlsystems/einfold/blob/main/docs/demos.md), and the demos run on DataFusion with einfold as one of their engines. -Each script checks its results against the same model in +The first two scripts check their results against the same model in [JAX](https://github.com/jax-ml/jax), Google's library for differentiable array programs. @@ -19,6 +27,7 @@ array programs. |---|---|---|---| | 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, the core of a transformer, as a view (the model) in einfold's window and `GROUP BY` spellings, and a loss that reads it | the naive plan, which stores every score, against [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan, which einfold finds from the same SQL; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention: the [Nadaraya–Watson estimator](https://en.wikipedia.org/wiki/Kernel_regression#Nadaraya%E2%80%93Watson_kernel_regression), an average of nearby points weighted by a Gaussian kernel, which einfold must fuse with the same rules | | 02 | [`02_nanogpt.py`](02_nanogpt.py) | [nanoGPT](https://github.com/karpathy/nanoGPT)'s GPT-2 model, at 1 layer and 4,416 parameters, as a view, and a cross-entropy loss that reads it, on Tiny Shakespeare | the loss on each host; training steps with SGD or, as nanoGPT trains, AdamW, every parameter table checked after every step | +| 03 | [`03_shakespeare.py`](03_shakespeare.py) | nanoGPT's 6-layer, 384-channel, 6-head Shakespeare character model with context 256 and batch 12 | training from scratch with exact GELU, dropout, SQL AdamW, train/validation evaluation and resumable checkpoints; optional checks of every gradient, weight and optimizer moment against the original PyTorch model | The SQL favors legibility over speed. Each model reads like its definition, and the slow parts are left slow on purpose: they are what ddx, xarray-sql @@ -29,7 +38,7 @@ these models fast is tracked in [#285](https://github.com/xqlsystems/xarray-sql/issues/285), which lists places to start. -Every result is checked before it is timed into the table. Each script +Every result in the original demos is checked before it is timed into the table. Each demo computes the same model in JAX from the same inputs and parameters, and a mismatch fails the run. @@ -40,12 +49,99 @@ uv run benchmarks/ml/01_attention.py uv run benchmarks/ml/02_nanogpt.py --steps 5 ``` -Each script carries [inline `uv` metadata](https://docs.astral.sh/uv/guides/scripts/) +### Shakespeare training matching nanoGPT's train.py + +`03_shakespeare.py` is a training loop with checkpoints and held-out +validation. Its defaults match `train_shakespeare_char.py` with these +overrides: batch 12, context 256, layers/heads/channels 6/6/384, gradient +accumulation 1, `max_iters=2000`, `lr_decay_iters=2000`, warmup 100, +evaluation every 200 iterations over 20 batches per split. Inherited +settings also match: dropout 0.2, exact GELU, no biases, tied embeddings, +AdamW with learning rate 1e-3 to 1e-4, betas 0.9/0.99, epsilon 1e-8, +weight decay 0.1 excluding layer norm gains, and global gradient clipping +at 1. The learning rate uses `lr*(iteration+1)/(warmup_iters+1)` during +warmup, as in the local nanoGPT source. There are **10,745,088 parameters** +with the 65-character vocabulary; nanoGPT prints 10.65M when excluding +position embeddings. + +Prepare the dataset once in a nanoGPT checkout with +`python data/shakespeare_char/prepare.py`. The SQL script reads its +`input.txt`, `train.bin` and `val.bin`, verifies the same sorted vocabulary +and 90/10 split, and samples shifted next-character targets without +crossing split boundaries. It initializes weights once per invocation +and keeps updating them across iterations. + +Use this checkout's existing virtual environment with the optional +[EinFold build](../../accelerators/einfold_mlx/README.md) and +[ddx erf build](../../accelerators/ddx_math/README.md). Avoid `uv run` +for this script, since environment synchronization can replace those local +wheels. From the xarray-sql directory: + +```sh +MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/ml/03_shakespeare.py \ + --backend dense_gpu \ + --data-dir /Users/cromwellian/nanoGPT/data/shakespeare_char \ + --out-dir out-shakespeare-sql +``` + +Use `--backend dense_cpu --cache-gib 32` for the measured compact float64 CPU path. GPU uses MLX +float32. Both differentiate the original SQL with ddx and execute the +model, global gradient clipping and AdamW arithmetic as SQL. They cache +query programs across updates. AdamW moments live in separate SQL tables +so the model's differentiated parameter column remains `w`. The maximum +used for softmax stabilization is spelled with `GROUP BY` and joins to +avoid repeated window-lowering traversal through the shared model graph. +Its maximum is held constant with `ddx_stop_gradient`: subtracting any +shared constant leaves softmax and cross entropy unchanged, so this +removes a cancelling derivative through the maximum without changing the +model's loss or parameter gradients. + +All size and recipe options accept nanoGPT's underscore spelling and +hyphens, for example `--batch_size=12` and `--eval-iters=20`. The loop +matches `train.py`'s termination condition: `max_iters=2000` includes +iterations 0 through 2000, or 2001 optimizer updates. Add `--steps 5` to +run exactly five updates while retaining the 2000-iteration LR schedule. + +Evaluation runs before the update at iteration 0, then every 200 +iterations, with dropout disabled on 20 train and 20 validation batches. +`metrics.jsonl` records losses, learning rates and synchronized SQL +training times. Times include forward loss, backward gradients, clipping, +AdamW and result registration; sampling masks/batches, evaluation, +checkpoint I/O and optional reference checks are outside the timer. + +`last.npz` contains all weights, optimizer moments, next iteration, both +RNG states, vocabulary and configuration. `best.npz` saves an improved +validation result after iteration 0. Resume with the same configuration: + +```sh +MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/ml/03_shakespeare.py \ + --resume out-shakespeare-sql/last.npz +``` + +For an independent correctness check, install PyTorch in the execution +environment and add `--steps 5 --verify-nanogpt /Users/cromwellian/nanoGPT`. +This imports the original `model.py`, retains exact GELU and PyTorch +AdamW, and compares evaluation/training loss, every raw gradient, every +updated weight and both moments on every step. The oracle uses explicit +attention and the same externally sampled Bernoulli dropout masks; +reference checks do not replace SQL computation. GPU comparisons use +the existing 3e-5 absolute/relative tolerance and CPU comparisons 1e-8. + +Random initialization and sampling use NumPy's RNG rather than PyTorch's, +so the same integer seed alone does not reproduce a separate stock +`train.py` trajectory. The oracle loads identical initial arrays and +masks to test the arithmetic. Checkpoints use this script's safe NPZ +format, rather than nanoGPT's PyTorch checkpoint format. This script +currently supports gradient accumulation 1 and training without biases. +The [measured comparison](../apple/shakespeare-training.md) records the +verified five-step CPU/GPU runs and validation results. + +The first two scripts carry [inline `uv` metadata](https://docs.astral.sh/uv/guides/scripts/) and runs against the xarray-sql in this checkout, with the `ml` extra. The options are the same in each: ```shell ---hosts datafusion duckdb # the engines to run on (default: all) +--hosts datafusion duckdb # default: datafusion einfold duckdb --sizes 64 256 # problem sizes: tokens, or batch size --reps 3 # timed repetitions; the median is reported --csv results.csv # also write the results table @@ -105,6 +201,58 @@ in its plan. | `einfold` | `xarray_sql.Context`, with einfold as its query planner | ddx, whose steps einfold plans too | | `duckdb` | DuckDB, through `xql.register` | ddx, with the `duckdb` extra of ddxdb, which pins the DuckDB it supports | +The optional Apple silicon comparison adds explicit hosts (the three +defaults above stay the same): + +| Host | Execution | +|---|---| +| `datafusion_prepared` | Retain forward query plans, ddx programs, model views and parameter buffers | +| `cpu_optimized` | Prepared path with EinFold and the float64 CPU matrix backend, using its normal thresholds | +| `cpu_forced` | Same, with the backend threshold set to zero for a direct small-kernel comparison | +| `gpu_optimized` | Prepared path with [MLX](https://github.com/ml-explore/mlx), Apple's array framework, and measured automatic GPU dispatch thresholds | +| `gpu_forced` | Same, forcing eligible dense operations to GPU to measure dispatch overhead | +| `shared_einfold` | Shared forward CTEs and reusable ddx stages, with EinFold's CPU kernels | +| `shared_cpu` | Shared execution with the float64 CPU matrix backend | +| `shared_gpu` | Shared execution with explicitly enabled float32 MLX kernels | + +These hosts need the [optional native bridge](../../accelerators/einfold_mlx/README.md). +GPU hosts explicitly permit float32 arithmetic and check JAX with +`rtol=atol=3e-5`; float64 hosts retain `1e-9`. Completed CPU/GPU kernel +counts are recorded, so an automatic-dispatch run that stays on CPU is +visible. Dataset loading and model-view creation times are also recorded +cumulatively on training rows. The loss/training times keep their original +scope and exclude JAX reference work and those loading/view times. + +```sh +MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/ml/02_nanogpt.py \ + --steps 5 --hosts datafusion einfold duckdb datafusion_prepared \ + cpu_optimized cpu_forced gpu_optimized gpu_forced --csv comparison.csv +``` + +The [M5 Ultra five-step comparison](../apple/nanogpt-5-steps.md) reports +cold and warm costs separately, checks every loss and weight update, and +includes a forced GPU run with actual kernel counts. + +The [shared-execution report](../apple/nanogpt-shared.md) adds larger models. +`--preset demo` uses one layer, 128 channels, four heads and context 64; +`--preset shakespeare` uses six layers, 384 channels, six heads and context +256, following nanoGPT's character-model dimensions. `--preset gpt2` uses +12 layers, 768 channels, 12 heads, context 1024 and vocabulary 50257. +Explicit dimension flags override presets. Inputs remain Shakespeare +character IDs: padding the vocabulary does not add GPT-2 tokenization, +pretrained weights, biases or dropout. `--forward-only` checks losses +without running gradients or updates. + +```sh +MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/apple/run_nanogpt_sizes.py \ + --preset shakespeare --steps 5 --batch 1 --learning-rate 0.01 \ + --hosts shared_cpu shared_gpu --prefix local-shakespeare +``` + +This runner starts hosts serially, checkpoints checked CSV rows and records +commands, exit status and sampled resident memory. A failed or bounded run +is not reported as a completed training benchmark. + A host that cannot differentiate still runs every forward query, and its gradient rows say `no ddx backend`. A gradient that ddx, or the engine it plans with, refuses is reported as `refused:` with the error, and a diff --git a/benchmarks/ml/_bench.py b/benchmarks/ml/_bench.py index dcb044e8..3888c8fb 100644 --- a/benchmarks/ml/_bench.py +++ b/benchmarks/ml/_bench.py @@ -24,7 +24,7 @@ from dataclasses import dataclass, field import numpy as np -from _hosts import HOSTS, Host +from _hosts import DEFAULT_HOSTS, HOSTS, Host def cli(description: str, sizes: Sequence[int], size_name: str = "n"): @@ -33,9 +33,9 @@ def cli(description: str, sizes: Sequence[int], size_name: str = "n"): parser.add_argument( "--hosts", nargs="+", - default=list(HOSTS), + default=list(DEFAULT_HOSTS), choices=list(HOSTS), - help="engines to run on (default: all)", + help="engines to run on (default: datafusion einfold duckdb)", ) parser.add_argument( "--sizes", @@ -156,6 +156,15 @@ def note(self, key: str, value) -> None: """Attach a measurement besides time to the last row.""" self.rows[-1][key] = value + def _write_csv(self, path): + """Checkpoint completed checks during long model runs.""" + if path: + columns = list(dict.fromkeys(k for row in self.rows for k in row)) + with open(path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=columns) + writer.writeheader() + writer.writerows(self.rows) + def finish(self, csv_path: str | None = None) -> None: """Print the table, and write it as CSV if asked. Exits with an error if a host refused a step it wasn't expected to.""" @@ -166,10 +175,7 @@ def finish(self, csv_path: str | None = None) -> None: cells = [_cell(row.get(c)) for c in columns] print(f"| {' | '.join(cells)} |") if csv_path: - with open(csv_path, "w", newline="") as f: - writer = csv.DictWriter(f, fieldnames=columns) - writer.writeheader() - writer.writerows(self.rows) + self._write_csv(csv_path) print(f"\nwrote {csv_path}", file=sys.stderr) if self.unexpected: sys.exit( diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py index b002b79b..1d6ea982 100644 --- a/benchmarks/ml/_hosts.py +++ b/benchmarks/ml/_hosts.py @@ -27,13 +27,21 @@ A gradient that ddx (or the engine under it) refuses to plan raises one of the host's ``refusals``, and the demos report that step with the reason. -Adding a host (DuckDB with a GPU extension, say) is one class here and -one entry in ``HOSTS``. +Optional Apple silicon hosts use MLX (https://github.com/ml-explore/mlx), +Apple's array framework, through xarray-sql's native dense-kernel bridge. +They pair prepared queries with reusable parameter buffers. Forced variants +lower kernel thresholds to measure small operations on each backend. + +Adding a host is one class here and one entry in ``HOSTS``. ``DEFAULT_HOSTS`` +keeps optional device dependencies out of ordinary runs. """ from __future__ import annotations import json +import os +import re +from collections import Counter from collections.abc import Sequence import duckdb @@ -45,6 +53,21 @@ import xarray_sql as xql +try: + from xarray_sql.mlx import MLX +except ImportError: + MLX = None # Optional Apple silicon dependency and Metal device. + +try: + from xarray_sql.shared import SharedExecutor +except ImportError: + SharedExecutor = None + +try: + from xarray_sql.dense import DenseExecutor +except ImportError: + DenseExecutor = None + class Host: """An engine that runs the demos' SQL. Subclasses fill in the rest.""" @@ -52,6 +75,11 @@ class Host: name = "" differentiates = False refusals: tuple[type[Exception], ...] = () + rtol = 1e-9 + atol = 1e-9 + + def measurements(self) -> dict: + return {} def load(self, name: str, data: xr.Dataset, **options) -> None: raise NotImplementedError @@ -114,7 +142,262 @@ class EinFold(DataFusion): def __init__(self): super().__init__() self.einfold = EinFoldExtension() - self.ctx.with_extensions(self.einfold) + self.ctx = self.ctx.with_extensions(self.einfold) + + def operators(self, sql): + return None # Native physical stages cross the FFI as opaque nodes. + + +class DataFusionPrepared(DataFusion): + """Cached SQL plans and ddx programs over reusable parameter buffers. + + Data is copied on first registration so in-place training updates never + mutate the caller's initial weights or the batches used by other hosts. + A schema/coordinate change refreshes the tables, views and query caches. + """ + + name = "datafusion_prepared" + + def __init__(self): + super().__init__() + for kind in ("", "_join", "_topk", "_aggregate"): + self.ctx.sql( + f"SET datafusion.optimizer.enable{kind}_dynamic_filter_pushdown = false" + ).collect() + self._tables: set[str] = set() + self._queries: dict = {} + self._programs: dict = {} + self._last_view: str | None = None + self._counts: Counter[str] = Counter() + self.accelerator = None + + def load(self, name, data, **options): + if name in self._tables: + try: + self.ctx._write({name: data}) + except ValueError: + pass + else: + self._counts["inplace_updates"] += 1 + return + self._queries.clear() + self._programs.clear() + self._last_view = None + super().load(name, data.copy(deep=True), **options) + self._tables.add(name) + self._counts["table_registers"] += 1 + + def execute(self, sql): + # The GPT recreates this same view after loading each batch. When + # tables retain their buffers, its existing plan reads current values. + view = sql.lstrip().upper().startswith("CREATE OR REPLACE VIEW ") + if view and sql == self._last_view: + return + super().execute(sql) + self._queries.clear() + self._programs.clear() + self._last_view = sql if view else None + + def query(self, sql): + if xql.ad.calls_ad(sql): + return self.query_all([sql])[0] + if sql not in self._queries: + self._queries[sql] = self.ctx.prepare(sql) + self._counts["query_preparations"] += 1 + return self._queries[sql]() + + def query_all(self, sqls): + key = tuple(sqls) + if key not in self._programs: + self._programs[key] = xql.ad.prepare(self.ctx, key) + self._counts["gradient_preparations"] += 1 + return self._programs[key]() + + def measurements(self): + values = dict(self._counts) + if self.accelerator is not None: + counters = self.accelerator.statistics() + values.update( + { + key: counters.get(key, 0) + for key in ( + "cpu_matmul", + "gpu_matmul", + "gpu_attention", + "declined", + ) + } + ) + return values + + def operators(self, sql): + return None if self.accelerator is not None else super().operators(sql) + + +class CPUOptimized(DataFusionPrepared): + """Prepared execution plus the float64 CPU dense-kernel bridge.""" + + name = "cpu_optimized" + precision = "float64" + force_kernels = False + + def __init__(self): + super().__init__() + if MLX is None: + raise ImportError( + "this host needs the Apple silicon MLX extra and native hook" + ) + options = {"min_ops": 0, "min_gpu_ops": 0} if self.force_kernels else {} + self.accelerator = MLX(precision=self.precision, **options) + self.ctx = self.ctx.with_extensions(self.accelerator) + + +class CPUForced(CPUOptimized): + """Force eligible small products through the float64 CPU backend.""" + + name = "cpu_forced" + force_kernels = True + + +class GPUOptimized(CPUOptimized): + """Prepared execution with the measured automatic GPU thresholds.""" + + name = "gpu_optimized" + precision = "float32" + rtol = 3e-5 + atol = 3e-5 + + +class GPUForced(GPUOptimized): + """Force eligible operations onto MLX to expose small-kernel overhead.""" + + name = "gpu_forced" + force_kernels = True + + +class SharedEinFold(DataFusionPrepared): + """Shared forward CTEs and reusable ddx physical plans through EinFold.""" + + name = "shared_einfold" + precision = None + + def __init__(self): + super().__init__() + if SharedExecutor is None: + raise ImportError("rebuild xarray-sql for shared SQL execution") + extension = EinFoldExtension() + if self.precision is not None: + if MLX is None: + raise ImportError( + "shared CPU/GPU kernels require the MLX extra" + ) + extension = self.accelerator = MLX( + precision=self.precision, + min_ops=0, + min_gpu_ops=0, + max_working_bytes=4 << 30, + ) + self.ctx = self.ctx.with_extensions(extension) + self.executor = SharedExecutor(self.ctx, max_key_cache_bytes=4 << 30) + + def load(self, name, data, **options): + # Large stacked weights remain lazy Arrow streams. One layer per + # partition lets a block's constant layer filter prune the others. + if "layer" in data.dims and "chunks" not in options: + options["chunks"] = {"layer": 1} + super().load(name, data, **options) + + def execute(self, sql): + view = re.fullmatch( + r"\s*CREATE OR REPLACE VIEW (\w+) AS\s*(.*)", + sql, + re.IGNORECASE | re.DOTALL, + ) + if view: + name, query = view.groups() + if sql != self._last_view: + # Preserve the score/softmax chain as one EinFold region. + self.executor.inline_ctes = frozenset( + re.findall( + r"\b(l\d+_(?:scores|weights|norm))\s+AS", + query, + re.IGNORECASE, + ) + ) + self.executor.register_view(name, query) + self._queries.clear() + self._programs.clear() + self._last_view = sql + return + super().execute(sql) + + def query(self, sql): + if xql.ad.calls_ad(sql): + return self.query_all([sql])[0] + if sql not in self._queries: + self._queries[sql] = self.executor.prepare(sql) + self._counts["query_preparations"] += 1 + return self._queries[sql]() + + def query_all(self, sqls): + key = tuple(sqls) + if key not in self._programs: + self._programs[key] = self.executor.prepare_all(sqls) + self._counts["gradient_preparations"] += 1 + return self._programs[key]() + + def operators(self, sql): + # The FFI planner exposes each stage as an opaque node. Its node + # count cannot be compared with ordinary DataFusion's operators. + return None + + def measurements(self): + return super().measurements() | self.executor.statistics() + + +class SharedCPU(SharedEinFold): + """Shared execution with double-precision dense CPU kernels.""" + + name = "shared_cpu" + precision = "float64" + + +class SharedGPU(SharedEinFold): + """Shared execution with explicitly enabled MLX float32 GPU kernels.""" + + name = "shared_gpu" + precision = "float32" + rtol = 3e-5 + atol = 3e-5 + + +class DenseCPU(SharedEinFold): + """Compact double-precision arrays across SQL and ddx stages.""" + + name = "dense_cpu" + dense_precision = "float64" + + def __init__(self): + super().__init__() + if DenseExecutor is None: + raise ImportError("install EinFold's optional dense program module") + self.executor = DenseExecutor( + self.ctx, + precision=self.dense_precision, + fallback=False, + max_cached_bytes=int( + os.environ.get("XQL_DENSE_CACHE_BYTES", str(8 << 30)) + ), + ) + + +class DenseGPU(DenseCPU): + """Compact MLX arrays across SQL and ddx stages, with float32 arithmetic.""" + + name = "dense_gpu" + dense_precision = "float32" + rtol = 3e-5 + atol = 3e-5 class DuckDB(Host): @@ -159,4 +442,22 @@ def _count(node: dict) -> int: return 1 + sum(_count(child) for child in node.get("children", [])) -HOSTS = {host.name: host for host in (DataFusion, EinFold, DuckDB)} +DEFAULT_HOSTS = (DataFusion.name, EinFold.name, DuckDB.name) +HOSTS = { + host.name: host + for host in ( + DataFusion, + EinFold, + DuckDB, + DataFusionPrepared, + CPUOptimized, + CPUForced, + GPUOptimized, + GPUForced, + SharedEinFold, + SharedCPU, + SharedGPU, + DenseCPU, + DenseGPU, + ) +} diff --git a/benchmarks/ml/_shakespeare_torch.py b/benchmarks/ml/_shakespeare_torch.py new file mode 100644 index 00000000..c9c742e8 --- /dev/null +++ b/benchmarks/ml/_shakespeare_torch.py @@ -0,0 +1,214 @@ +"""Optional numerical oracle using the original nanoGPT model and AdamW. + +nanoGPT (https://github.com/karpathy/nanoGPT) runs the same batches, +initial weights and externally supplied Bernoulli dropout masks. Only +dropout sampling and the fused attention selection are adapted, allowing +each SQL gradient and optimizer state to be checked independently. +""" + +import importlib.util +import sys + +import numpy as np + +try: + import torch +except ImportError: + torch = None + + +class _TorchReference: + def __init__(self, args, params, vocab): + if torch is None: + raise ImportError( + "--verify-nanogpt requires PyTorch in this Python environment" + ) + spec = importlib.util.spec_from_file_location( + "_shakespeare_native_model", args.verify_nanogpt / "model.py" + ) + native = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = native + spec.loader.exec_module(native) + torch.set_num_threads(8) + self.args = args + self.dtype = ( + torch.float32 if args.backend == "dense_gpu" else torch.float64 + ) + self.tolerance = 3e-5 if self.dtype == torch.float32 else 1e-8 + self.model = native.GPT( + native.GPTConfig( + block_size=args.block_size, + vocab_size=vocab, + n_layer=args.n_layer, + n_head=args.n_head, + n_embd=args.n_embd, + dropout=args.dropout, + bias=False, + ) + ).to(dtype=self.dtype) + self.masks = {} + owner = self + + class _Mask(torch.nn.Module): + def __init__(self, name): + super().__init__() + self.name = name + + def forward(self, x): + return x * owner.masks[self.name] if self.training else x + + self.model.transformer.drop = _Mask("embed_drop") + for layer, block in enumerate(self.model.transformer.h): + block.attn.flash = False + block.attn.register_buffer( + "bias", + torch.tril(torch.ones(args.block_size, args.block_size)).view( + 1, 1, args.block_size, args.block_size + ), + ) + block.attn.attn_dropout = _Mask(f"l{layer}_adrop") + block.attn.resid_dropout = _Mask(f"l{layer}_rdrop") + block.mlp.dropout = _Mask(f"l{layer}_mdrop") + with torch.no_grad(): + for name in params: + for layer, target, encode in self._parts(name): + source = ( + params[name] if layer is None else params[name][layer] + ) + target.copy_( + torch.as_tensor(encode(source), dtype=self.dtype) + ) + decay = [p for p in self.model.parameters() if p.dim() >= 2] + gains = [p for p in self.model.parameters() if p.dim() < 2] + self.optimizer = torch.optim.AdamW( + [ + {"params": decay, "weight_decay": args.weight_decay}, + {"params": gains, "weight_decay": 0.0}, + ], + lr=args.learning_rate, + betas=(args.beta1, args.beta2), + eps=1e-8, + foreach=False, + fused=False, + ) + + def _parts(self, name): + if name in ("wte", "wpe", "ln_f"): + yield None, self.model.transformer[name].weight, lambda x: x + return + for layer, block in enumerate(self.model.transformer.h): + if name in ("wq", "wk", "wv"): + start = ("wq", "wk", "wv").index(name) * self.args.n_embd + target = block.attn.c_attn.weight[ + start : start + self.args.n_embd + ] + yield ( + layer, + target, + lambda x: x.reshape(self.args.n_embd, self.args.n_embd).T, + ) + elif name == "wo": + yield ( + layer, + block.attn.c_proj.weight, + lambda x: x.reshape(self.args.n_embd, self.args.n_embd).T, + ) + elif name in ("w_fc", "w_proj"): + attr = "c_fc" if name == "w_fc" else "c_proj" + yield layer, getattr(block.mlp, attr).weight, lambda x: x.T + else: + yield layer, getattr(block, name).weight, lambda x: x + + def _arrays(self, name, column=None): + def value(parameter): + if column is None: + return parameter.detach().numpy() + if column == "grad": + return parameter.grad.detach().numpy() + return self.optimizer.state[parameter][column].detach().numpy() + + if name in ("wte", "wpe", "ln_f"): + return value(self.model.transformer[name].weight) + result = [] + for block in self.model.transformer.h: + if name in ("wq", "wk", "wv"): + start = ("wq", "wk", "wv").index(name) * self.args.n_embd + result.append( + value(block.attn.c_attn.weight)[ + start : start + self.args.n_embd + ].T.reshape( + self.args.n_embd, + self.args.n_head, + self.args.n_embd // self.args.n_head, + ) + ) + elif name == "wo": + result.append( + value(block.attn.c_proj.weight).T.reshape( + self.args.n_head, + self.args.n_embd // self.args.n_head, + self.args.n_embd, + ) + ) + elif name in ("w_fc", "w_proj"): + result.append( + value( + getattr( + block.mlp, "c_fc" if name == "w_fc" else "c_proj" + ).weight + ).T + ) + else: + result.append(value(getattr(block, name).weight)) + return np.stack(result) + + def _close(self, actual, expected, label): + np.testing.assert_allclose( + actual, + expected, + atol=self.tolerance, + rtol=self.tolerance, + err_msg=label, + ) + + def evaluate(self, batch, loss): + self.model.eval() + with torch.no_grad(): + _, expected = self.model( + torch.from_numpy(batch.tok.values), + torch.from_numpy(batch.next.values), + ) + self._close(loss, float(expected), "evaluation loss") + + def step(self, batch, masks, lr, loss, grads, params, moments): + self.model.train() + self.masks = { + name: torch.as_tensor(ds.scale.values, dtype=self.dtype) + for name, ds in masks.items() + } + self.optimizer.zero_grad(set_to_none=True) + _, expected = self.model( + torch.from_numpy(batch.tok.values), + torch.from_numpy(batch.next.values), + ) + self._close( + loss, float(expected.detach()), "training loss with dropout" + ) + expected.backward() + for name in params: + self._close( + grads[name].g.values, + self._arrays(name, "grad"), + f"{name} gradient", + ) + if self.args.grad_clip != 0: + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), self.args.grad_clip + ) + for group in self.optimizer.param_groups: + group["lr"] = lr + self.optimizer.step() + for name in params: + self._close(params[name], self._arrays(name), f"{name} weight") + for values, col in zip(moments[name], ("exp_avg", "exp_avg_sq")): + self._close(values, self._arrays(name, col), f"{name} {col}") diff --git a/pyproject.toml b/pyproject.toml index d7b32771..49242e07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,12 @@ dependencies = [ ] [project.optional-dependencies] +# Apple silicon GPU kernels need the native hook built by +# accelerators/einfold_mlx/build.py in addition to these dependencies. +mlx = [ + "einfold[datafusion]>=0.2,<0.3; sys_platform == 'darwin' and platform_machine == 'arm64'", + "mlx>=0.32.4,<0.33; sys_platform == 'darwin' and platform_machine == 'arm64'", +] # Lazy, chunked results from `to_dataset(chunks=...)` need a chunk manager. dask = [ "dask>=2024.8.0", @@ -143,7 +149,7 @@ dev = [ "ddxdb[duckdb]>=0.4.6,<0.5; sys_platform != 'emscripten'", # TODO(xqlsystems/einfold#118): install einfold on Windows too, once # it publishes a Windows wheel; until then it would build from source. - "einfold[datafusion]>=0.1.0rc2,<0.2; sys_platform != 'emscripten' and sys_platform != 'win32'", + "einfold[datafusion]>=0.1.0rc2,<0.3; sys_platform != 'emscripten' and sys_platform != 'win32'", # jaxlib publishes no wheels for Intel macOS, so the demos skip there. "jax; sys_platform != 'darwin' or platform_machine != 'x86_64'", ] diff --git a/src/lib.rs b/src/lib.rs index 003d8707..2fd796b3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -43,6 +43,7 @@ #[cfg(feature = "flight")] mod flight; +mod shared; use std::collections::{HashMap, HashSet}; use std::ffi::CString; @@ -1656,6 +1657,8 @@ impl LazyArrowStreamTable { #[pymodule] fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_function(wrap_pyfunction!(shared::shared_ctes, m)?)?; + m.add_class::()?; #[cfg(feature = "flight")] m.add_class::()?; Ok(()) diff --git a/src/shared.rs b/src/shared.rs new file mode 100644 index 00000000..160a81f9 --- /dev/null +++ b/src/shared.rs @@ -0,0 +1,137 @@ +//! Parse an explicitly staged SELECT with the same SQL parser as DataFusion. + +use std::ops::ControlFlow; +use std::sync::Arc; + +use arrow::array::RecordBatch; +use arrow::pyarrow::FromPyArrow; +use datafusion::catalog::TableProvider; +use datafusion::common::Column; +use datafusion::datasource::MemTable; +use datafusion::logical_expr::Expr as DFExpr; +use datafusion_ffi::table_provider::FFI_TableProvider; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; +use sqlparser::ast::{Expr, ObjectName, Query, Statement, Visit, Visitor}; +use sqlparser::dialect::GenericDialect; +use sqlparser::parser::Parser; + +#[derive(Default)] +struct References { + tables: Vec, + functions: Vec, + nested_with: bool, +} + +impl Visitor for References { + type Break = (); + + fn pre_visit_query(&mut self, query: &Query) -> ControlFlow<()> { + self.nested_with |= query.with.is_some(); + ControlFlow::Continue(()) + } + + fn pre_visit_relation(&mut self, name: &ObjectName) -> ControlFlow<()> { + if name.0.len() == 1 { + if let Some(id) = name.0[0].as_ident() { + self.tables.push(if id.quote_style.is_some() { + id.value.clone() + } else { + id.value.to_lowercase() + }); + } + } + ControlFlow::Continue(()) + } + + fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<()> { + if let Expr::Function(f) = expr { + self.functions.push(f.name.to_string().to_lowercase()); + } + ControlFlow::Continue(()) + } +} + +type Part = (String, String, Vec, Vec); + +/// Return top-level CTEs and the final SELECT, with their direct references. +#[pyfunction] +#[pyo3(signature = (sql, allow_nested=false))] +pub(crate) fn shared_ctes(sql: &str, allow_nested: bool) -> PyResult> { + let mut statements = Parser::parse_sql(&GenericDialect {}, sql) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + if statements.len() != 1 { + return Err(PyValueError::new_err( + "shared execution requires one SELECT", + )); + } + let Statement::Query(mut query) = statements.remove(0) else { + return Err(PyValueError::new_err("shared execution requires a SELECT")); + }; + let mut parts = Vec::new(); + if let Some(with) = query.with.take() { + if with.recursive { + return Err(PyValueError::new_err("recursive CTEs cannot be staged")); + } + for cte in with.cte_tables { + if !cte.alias.columns.is_empty() { + return Err(PyValueError::new_err("CTE column aliases cannot be staged")); + } + let id = cte.alias.name; + let name = if id.quote_style.is_some() { + id.value + } else { + id.value.to_lowercase() + }; + parts.push(part(name, &cte.query, allow_nested)?); + } + } + parts.push(part(String::new(), &query, allow_nested)?); + Ok(parts) +} + +fn part(name: String, query: &Query, allow_nested: bool) -> PyResult { + let mut refs = References::default(); + let _ = query.visit(&mut refs); + if refs.nested_with && !allow_nested { + return Err(PyValueError::new_err("nested WITH scopes cannot be staged")); + } + Ok((name, query.to_string(), refs.tables, refs.functions)) +} + +/// An internal Arrow stage with a verified ascending, NULL-last key order. +/// The Python stage executor verifies or sorts these keys before registration +/// and never mutates key buffers while a plan reads this provider. +#[pyclass(module = "xarray_sql._native", name = "_SortedArrowTable")] +pub(crate) struct SortedArrowTable { + table: Arc, +} + +#[pymethods] +impl SortedArrowTable { + #[new] + fn new(batch: &Bound<'_, PyAny>, keys: Vec) -> PyResult { + let batch = RecordBatch::from_pyarrow_bound(batch)?; + let ordering = keys + .into_iter() + .map(|key| DFExpr::Column(Column::from_name(key)).sort(true, false)) + .collect(); + let table = MemTable::try_new(batch.schema(), vec![vec![batch]]) + .map_err(|e| PyValueError::new_err(e.to_string()))? + .with_sort_order(vec![ordering]); + Ok(Self { + table: Arc::new(table), + }) + } + + fn __datafusion_table_provider__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, + ) -> PyResult> { + let codec = crate::ffi_logical_codec_from_pycapsule(session)?; + let provider = FFI_TableProvider::new_with_ffi_codec(self.table.clone(), true, None, codec); + PyCapsule::new_with_value(py, provider, c"datafusion_table_provider") + } +} diff --git a/tests/test_ad.py b/tests/test_ad.py index 0f868bc7..df3d2675 100644 --- a/tests/test_ad.py +++ b/tests/test_ad.py @@ -136,9 +136,54 @@ def test_a_mention_in_a_literal_is_not_a_call(): assert rows == [{"s": "grad(x)", "jvp(v)": 1}] -def test_prepare_refuses_a_gradient(problem): - ctx, *_ = problem - with pytest.raises(NotImplementedError, match="grad"): - ctx.prepare( - f"WITH {LOSS} SELECT feature, grad(loss, w) AS g FROM weights" +def test_context_prepare_runs_gradient_steps_and_reads_writes(problem): + ctx, x, w, y = problem + step = ctx.prepare( + f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature" + ) + old = step() + original = old["w"].to_numpy().copy() + replacement = xr.Dataset( + {"w": ("feature", w.copy() * 2)}, coords={"feature": np.arange(3)} + ) + got = step(weights=replacement) + np.testing.assert_allclose( + got["w"], 2 * x.T @ (x @ replacement.w.values - y) + ) + np.testing.assert_array_equal(old["w"], original) + with pytest.raises(ValueError, match="coordinates"): + step(weights=replacement.assign_coords(feature=[3, 4, 5])) + + +def test_prepared_gradients_read_new_values_and_keep_old_results(problem): + ctx, x, w, y = problem + statements = [ + f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature", + f"WITH {LOSS} SELECT feature, w - 0.01 * grad(loss, w) AS w FROM weights ORDER BY feature", + ] + run = ad.prepare(ctx, statements) + [first, _] = run() + original = first["w"].to_numpy().copy() + for scale in (0.5, -1.0, 2.0): + values = xr.Dataset( + {"w": ("feature", w.copy() * scale)}, + coords={"feature": np.arange(3)}, ) + gradient, step = run(weights=values) + want = 2 * x.T @ (x @ values.w.values - y) + np.testing.assert_allclose(gradient["w"], want) + np.testing.assert_allclose(step["w"], values.w.values - 0.01 * want) + np.testing.assert_array_equal(first["w"], original) + # Regular ddx preparation resolves a replacement table by name. + ctx.deregister_table("weights") + ctx.from_dataset("weights", values * 2) + [gradient, _] = run() + np.testing.assert_allclose( + gradient["w"], 2 * x.T @ (x @ (values.w.values * 2) - y) + ) + + +def test_prepared_statements_require_a_sequence(problem): + ctx, *_ = problem + with pytest.raises(TypeError, match="list"): + ad.prepare(ctx, "SELECT 1") diff --git a/tests/test_dense.py b/tests/test_dense.py new file mode 100644 index 00000000..43287ab3 --- /dev/null +++ b/tests/test_dense.py @@ -0,0 +1,283 @@ +"""Compact execution preserves SQL rows and live Dataset inputs.""" + +import os + +import numpy as np +import pyarrow as pa +import pytest +import xarray as xr +from datafusion import udf + +from xarray_sql import Context + +pytest.importorskip("ddxdb") +pytest.importorskip("einfold.dense_program") +from ddxdb import UnsupportedExpression +from einfold.datafusion import EinFold + +from xarray_sql.dense import DenseExecutor + +_PRECISIONS = [ + "float64", + pytest.param( + "float32", + marks=pytest.mark.skipif( + os.environ.get("MLX_ENABLE_TF32") != "0", + reason="GPU tests require MLX_ENABLE_TF32=0", + ), + ), +] + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_conditional_contraction_stays_within_array_bound(precision): + ctx = Context() + left = np.arange(64, dtype=float).reshape(8, 8) / 64 + right = np.arange(64, dtype=float).reshape(8, 8) / 32 + ctx.from_dataset("a", xr.Dataset({"v": (("i", "k"), left)})) + ctx.from_dataset("b", xr.Dataset({"v": (("k", "j"), right)})) + executor = DenseExecutor( + ctx, precision=precision, fallback=False, max_array_bytes=2048 + ) + rows = executor.prepare("""WITH products AS ( + SELECT a.i,b.j,a.k,a.v*b.v AS p FROM a JOIN b ON a.k=b.k) + SELECT i,j,SUM(CASE WHEN i>=j THEN p ELSE 0.0 END) AS x, + SUM(CASE WHEN i>=j THEN 0.0 ELSE p END) AS y + FROM products GROUP BY i,j ORDER BY i,j""")() + result = left @ right + i, j = np.indices(result.shape) + np.testing.assert_allclose( + np.asarray(rows["x"]).reshape(8, 8), np.where(i >= j, result, 0) + ) + np.testing.assert_allclose( + np.asarray(rows["y"]).reshape(8, 8), np.where(i >= j, 0, result) + ) + + +@pytest.fixture(params=_PRECISIONS) +def executor(request): + ctx = Context().with_extensions(EinFold()) + ctx.from_dataset( + "a", + xr.Dataset( + {"v": (("i", "j"), np.arange(6, dtype=float).reshape(2, 3))}, + coords={"i": [-2, 7], "j": [8, 2, 5]}, + ), + ) + ctx.from_dataset( + "b", + xr.Dataset( + {"w": (("j", "k"), np.arange(6, dtype=float).reshape(3, 2))}, + coords={"j": [5, 8, 2], "k": [9, 1]}, + ), + ) + return DenseExecutor(ctx, precision=request.param, fallback=False) + + +@pytest.mark.parametrize( + "sql", + [ + "SELECT a.i,b.k,SUM(a.v*b.w) AS v FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k ORDER BY a.i,b.k", + "SELECT i,AVG(v) AS m,COUNT(v) AS n FROM a WHERE v>1 GROUP BY i ORDER BY i", + "SELECT SUM(v) AS s,AVG(v) AS m,COUNT(v) AS n FROM a WHERE v<0", + "SELECT i,SUM(v) AS s FROM a WHERE v<0 GROUP BY i ORDER BY i", + "SELECT i,j,MAX(v) OVER(PARTITION BY i) AS m FROM a ORDER BY i,j", + "SELECT i,j,i/3 AS q,i%3 AS r FROM a ORDER BY i,j", + "SELECT i,j,(v-3.5)%2.0 AS r FROM a ORDER BY i,j", + "SELECT i,j,power(2,63) AS a,power(2,-1) AS b FROM a ORDER BY i,j", + "SELECT i,j,(CASE WHEN v>2 THEN TRUE END) AND FALSE AS a,(CASE WHEN v>2 THEN TRUE END) OR TRUE AS b FROM a ORDER BY i,j", + "WITH t AS (SELECT i,CASE WHEN v>2 THEN v END AS v FROM a) SELECT i,SUM(v) AS s,COUNT(v) AS n FROM t GROUP BY i ORDER BY i", + "WITH t AS (SELECT i,v FROM a UNION ALL SELECT i,2*v AS v FROM a) SELECT i,SUM(v) AS s,COUNT(v) AS n FROM t GROUP BY i ORDER BY i", + "SELECT a.i,a.j,b.w FROM a LEFT JOIN (SELECT j,k,w FROM b WHERE k=1) b ON a.j=b.j ORDER BY a.i,a.j", + "SELECT a.i,a.j,c.k FROM a LEFT JOIN (SELECT j,k,w FROM b WHERE k=1 AND w>1) b ON a.j=b.j JOIN (SELECT j,k FROM b WHERE k=9) c ON b.j=c.j ORDER BY a.i,a.j,c.k", + "SELECT a.i,a.j,c.k FROM a LEFT JOIN (SELECT j,k,w FROM b WHERE k=1 AND w>1) b ON a.j=b.j JOIN (SELECT j,k FROM b WHERE k=9) c ON b.j IS NOT DISTINCT FROM c.j ORDER BY a.i,a.j,c.k", + ], +) +def test_dense_sql_rows_match_datafusion(executor, sql): + actual = executor.prepare(sql)() + expected = executor.ctx.sql(sql).to_arrow_table() + assert actual.column_names == expected.column_names + for a, e in zip(actual.columns, expected.columns): + assert a.null_count == e.null_count + assert a.to_pylist() == e.to_pylist() + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_dense_gather_gradient_and_snapshots(precision): + ctx = Context().with_extensions(EinFold()) + weights = xr.Dataset( + {"v": (("tok", "e"), np.arange(6, dtype=float).reshape(3, 2))}, + coords={"tok": [8, -1, 4]}, + ) + ctx.from_dataset("w", weights) + ctx.from_dataset("batch", xr.Dataset({"tok": ("pos", [4, 4, 8])})) + executor = DenseExecutor(ctx, precision=precision, fallback=False) + executor.register_view( + "model", "SELECT b.pos,w.e,w.v FROM batch b JOIN w ON b.tok=w.tok" + ) + run = executor.prepare( + "WITH loss AS (SELECT SUM(v*v)/2 AS l FROM model) SELECT tok,e,grad(loss,v) AS g FROM w ORDER BY tok,e" + ) + first = run() + assert first.to_pydict() == { + "tok": [-1, -1, 4, 4, 8, 8], + "e": [0, 1, 0, 1, 0, 1], + "g": [0.0, 0.0, 8.0, 10.0, 0.0, 1.0], + } + weights.v.values[:] += 2 + second = run() + assert second["g"].to_pylist() == [0.0, 0.0, 12.0, 14.0, 2.0, 3.0] + assert first["g"].to_pylist() == [0.0, 0.0, 8.0, 10.0, 0.0, 1.0] + + +def test_dense_falls_back_for_overridden_scalar_function(executor): + ctx = executor.ctx + ctx.register_udf( + udf( + lambda a: pa.array(np.asarray(a) + 100), + [pa.float64()], + pa.float64(), + "immutable", + "exp", + ) + ) + run = DenseExecutor(ctx).prepare( + "SELECT i,j,exp(v) AS v FROM a ORDER BY i,j" + ) + assert run().equals( + ctx.sql("SELECT i,j,exp(v) AS v FROM a ORDER BY i,j").to_arrow_table() + ) + ctx.register_udf( + udf( + lambda a: pa.array(np.asarray(a) + 200), + [pa.float64()], + pa.float64(), + "immutable", + "exp", + ) + ) + assert run().equals( + ctx.sql("SELECT i,j,exp(v) AS v FROM a ORDER BY i,j").to_arrow_table() + ) + + +def test_dense_coordinate_cast_retains_its_integer_width(): + ctx = Context().with_extensions(EinFold()) + ctx.from_dataset( + "a", + xr.Dataset( + {"v": ("i", [1.0, 2.0])}, + coords={"i": np.array([2000000000, 2000000001], dtype=np.int32)}, + ), + ) + sql = "SELECT i,CAST(i AS BIGINT)*2 AS j FROM a ORDER BY i" + assert ( + DenseExecutor(ctx, fallback=False) + .prepare(sql)() + .equals(ctx.sql(sql).to_arrow_table()) + ) + + +def test_dense_base_log_fallback_never_writes_an_input(): + ctx = Context().with_extensions(EinFold()) + values = np.array([1.0, 2.0, 3.0]) + ctx.from_dataset("x", xr.Dataset({"v": ("i", values)})) + sql = "SELECT i,log(2.,v) AS v FROM x ORDER BY i" + expected = ctx.sql(sql).to_arrow_table() + assert DenseExecutor(ctx).prepare(sql)().equals(expected) + np.testing.assert_array_equal(values, [1.0, 2.0, 3.0]) + + +def test_dense_integer_sum_uses_the_sql_accumulator_width(): + ctx = Context().with_extensions(EinFold()) + ctx.from_dataset( + "x", + xr.Dataset( + {"v": ("i", np.array([2000000000, 2000000001], dtype=np.int32))} + ), + ) + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_dense_null_and_empty_filters_do_not_evaluate_deleted_rows(precision): + ctx = Context().with_extensions(EinFold()) + ctx.from_dataset("x", xr.Dataset({"v": ("i", [1.0, 2.0, 3.0])})) + executor = DenseExecutor(ctx, precision=precision, fallback=False) + for sql in [ + "SELECT i,v FROM x WHERE i=NULL ORDER BY i", + "SELECT i,1/(i-i) AS v FROM x WHERE i<0 ORDER BY i", + ]: + assert executor.prepare(sql)().equals(ctx.sql(sql).to_arrow_table()) + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_dense_integer_averages_with_repeated_group_keys(precision): + ctx = Context().with_extensions(EinFold()) + ctx.from_dataset( + "x", xr.Dataset({"key": ("i", [1, 1, 2, 2]), "v": ("i", [1, 2, 3, 4])}) + ) + sql = "SELECT key,AVG(v) AS m,SUM(v) AS s FROM x GROUP BY key ORDER BY key" + assert ( + DenseExecutor(ctx, precision=precision, fallback=False) + .prepare(sql)() + .equals(ctx.sql(sql).to_arrow_table()) + ) + sql = "SELECT SUM(v) AS s,AVG(v) AS m FROM x" + assert ( + DenseExecutor(ctx, fallback=False) + .prepare(sql)() + .equals(ctx.sql(sql).to_arrow_table()) + ) + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_dense_unsupported_backward_region_restarts_with_original_sql( + precision, +): + ctx = Context().with_extensions(EinFold()) + values = np.arange(6.0, dtype=float).reshape(2, 3) + ctx.from_dataset("a", xr.Dataset({"v": (("i", "j"), values)})) + query = DenseExecutor(ctx, precision=precision).prepare( + "WITH t AS (SELECT CAST(i AS DOUBLE) AS i,CAST(j AS DOUBLE) AS j,SUM(v*v) AS s FROM a GROUP BY CAST(i AS DOUBLE),CAST(j AS DOUBLE)),loss AS (SELECT SUM(s) AS l FROM t) SELECT i,j,grad(loss,v) AS g FROM a ORDER BY i,j" + ) + assert query()["g"].to_pylist() == [0.0, 2.0, 4.0, 6.0, 8.0, 10.0] + values[:] += 1 + assert query()["g"].to_pylist() == [2.0, 4.0, 6.0, 8.0, 10.0, 12.0] + + +def test_dense_directional_derivatives_and_typed_refusal(executor): + ctx = executor.ctx + tangent = "WITH tangent AS (SELECT j,k,CAST(k AS DOUBLE) AS w FROM b) " + rows = executor.prepare( + tangent + + "SELECT a.i,b.k,jvp(SUM(a.v*b.w),b.w,tangent) AS p FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k ORDER BY i,k" + )() + assert rows.to_pydict() == { + "i": [-2, -2, 7, 7], + "k": [1, 9, 1, 9], + "p": [3.0, 27.0, 12.0, 108.0], + } + ctx.from_dataset( + "cotangent", + xr.Dataset( + {"v": (("i", "k"), [[1.0, 2.0], [3.0, 4.0]])}, + coords={"i": [-2, 7], "k": [9, 1]}, + ), + ) + predicted = "WITH pred AS (SELECT a.i,b.k,SUM(a.v*b.w) AS v FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k) " + rows = executor.prepare( + predicted + "SELECT j,k,vjp(pred,w,cotangent) AS g FROM b ORDER BY j,k" + )() + assert rows.to_pydict() == { + "j": [2, 2, 5, 5, 8, 8], + "k": [1, 9, 1, 9, 1, 9], + "g": [18.0, 13.0, 24.0, 17.0, 12.0, 9.0], + } + rows = executor.prepare( + "WITH loss AS (SELECT SUM(v*ddx_stop_gradient(v)) AS l FROM a) SELECT i,j,grad(loss,v) AS g FROM a ORDER BY i,j" + )() + assert rows["g"].to_pylist() == [1.0, 2.0, 0.0, 4.0, 5.0, 3.0] + with pytest.raises(UnsupportedExpression, match="GROUP BY key"): + executor.prepare( + "WITH buckets AS (SELECT ROUND(v) AS k,SUM(v*v) AS s FROM a GROUP BY ROUND(v)),loss AS (SELECT SUM(k*s) AS l FROM buckets) SELECT i,j,grad(loss,v) AS g FROM a" + )() diff --git a/tests/test_math.py b/tests/test_math.py new file mode 100644 index 00000000..f858422c --- /dev/null +++ b/tests/test_math.py @@ -0,0 +1,99 @@ +"""Error-function SQL values and analytic gradients match the array paths.""" + +import math +import os + +import numpy as np +import pyarrow as pa +import pytest +import xarray as xr +from datafusion import udf + +from xarray_sql import Context +from xarray_sql.math import register_erf + + +def test_erf_sql_nulls_and_replacement(): + ctx = Context() + register_erf(ctx) + rows = ctx.sql( + "SELECT erf(0.0) AS zero, erf(CAST(NULL AS DOUBLE)) AS missing" + ).to_arrow_table() + assert rows.to_pydict() == {"zero": [0.0], "missing": [None]} + ctx.register_udf( + udf( + lambda x: pa.array(np.asarray(x) + 10), + [pa.float64()], + pa.float64(), + "immutable", + name="erf", + ) + ) + assert ctx.sql("SELECT erf(1.0) AS v").to_arrow_table().to_pydict() == { + "v": [11.0] + } + + +def test_erf_replacement_is_preserved_by_dense_fallback(): + pytest.importorskip("einfold.dense_program") + from xarray_sql.dense import DenseExecutor + + ctx = Context() + register_erf(ctx) + ctx.from_dataset("a", xr.Dataset({"v": ("i", [0.0, 1.0])})) + executor = DenseExecutor(ctx) + run = executor.prepare("SELECT i,erf(v) AS x FROM a ORDER BY i") + np.testing.assert_allclose(run()["x"].to_numpy(), [0.0, math.erf(1.0)]) + ctx.register_udf( + udf( + lambda x: pa.array(np.asarray(x) + 10), + [pa.float64()], + pa.float64(), + "immutable", + name="erf", + ) + ) + assert run().to_pydict() == {"i": [0, 1], "x": [10.0, 11.0]} + + +@pytest.mark.parametrize( + "precision", + [ + "sql", + "float64", + pytest.param( + "float32", + marks=pytest.mark.skipif( + os.environ.get("MLX_ENABLE_TF32") != "0", + reason="GPU tests require MLX_ENABLE_TF32=0", + ), + ), + ], +) +def test_exact_gelu_sql_gradient(precision): + ddx = pytest.importorskip("ddxdb") + if "erf" not in ddx.supported_functions(): + pytest.skip("requires the optional ddx erf patch") + ctx = Context() + register_erf(ctx) + values = np.array([-4.0, -1.0, 0.0, 0.5, 2.0]) + ctx.from_dataset( + "a", + xr.Dataset({"v": ("i", values)}, coords={"i": np.arange(len(values))}), + ) + sql = """WITH loss AS (SELECT SUM(0.5*v*(1.0+erf(v/sqrt(2.0)))) AS l FROM a) + SELECT i, grad(loss,v) AS g FROM a ORDER BY i""" + if precision == "sql": + rows = ctx.sql(sql).to_arrow_table() + else: + pytest.importorskip("einfold.dense_program") + from xarray_sql.dense import DenseExecutor + + executor = DenseExecutor(ctx, precision=precision, fallback=False) + rows = executor.prepare(sql)() + expected = 0.5 * ( + 1 + np.vectorize(math.erf)(values / np.sqrt(2)) + ) + values * np.exp(-(values**2) / 2) / np.sqrt(2 * np.pi) + np.testing.assert_allclose( + rows["g"].to_numpy(), expected, atol=1e-6, rtol=1e-6 + ) diff --git a/tests/test_ml_demos.py b/tests/test_ml_demos.py index b63ac452..06e72b21 100644 --- a/tests/test_ml_demos.py +++ b/tests/test_ml_demos.py @@ -1,7 +1,9 @@ """The ML demos in ``benchmarks/ml`` run, at their smallest sizes. -Each demo checks its results against JAX and exits with an error on a -mismatch, so this runs each one and checks that it succeeds. It needs the +The original demos check their results against JAX and exit on a +mismatch. The Shakespeare trainer's end-to-end tests, using prepared +temporary data, live in test_shakespeare.py; here its CLI is checked. +The original demos need the ``ml`` extra, ddx's DuckDB backend, einfold and JAX, which the ``dev`` dependency group installs. @@ -24,6 +26,7 @@ ["--steps", "1"], ["--steps", "1", "--optimizer", "adamw"], ], + "03_shakespeare.py": [["--help"]], } diff --git a/tests/test_mlx.py b/tests/test_mlx.py new file mode 100644 index 00000000..df02dfb9 --- /dev/null +++ b/tests/test_mlx.py @@ -0,0 +1,391 @@ +"""Public SQL and ddx contracts of the optional Apple silicon executor.""" + +import os + +import numpy as np +import pyarrow as pa +import pytest +import xarray as xr +from datafusion import SessionConfig + +pytest.importorskip("mlx.core", exc_type=ImportError) +if os.environ.get("MLX_ENABLE_TF32") != "0": + pytest.skip( + "GPU tests require MLX_ENABLE_TF32=0 at launch", allow_module_level=True + ) +pytest.importorskip("einfold") +pytest.importorskip("ddxdb") + +from ddxdb import UnsupportedExpression + +from xarray_sql import Context, ad +from xarray_sql.mlx import MLX + +PRODUCT = """ +SELECT a.i, b.k, SUM(a.v * b.v) AS v +FROM a JOIN b ON a.j = b.j GROUP BY a.i, b.k ORDER BY a.i, b.k +""" + + +def _matrix(value, dims): + return xr.Dataset( + {"v": (dims, value)}, + coords={d: np.arange(n) for d, n in zip(dims, value.shape)}, + attrs={"experiment": "accelerated"}, + ) + + +@pytest.fixture(params=["float64", "float32"]) +def accelerated(request): + try: + extension = MLX( + precision=request.param, min_ops=0, min_gpu_ops=0, query_block=3 + ) + except ImportError as exc: + pytest.skip(str(exc)) + ctx = Context().with_extensions(extension) + return ctx, extension, 3e-5 if request.param == "float32" else 1e-12 + + +def _same(got, want, tolerance): + assert got.schema == want.schema + assert got.column_names == want.column_names + for g, w in zip(got.columns, want.columns): + if pa.types.is_floating(w.type): + assert g.null_count == w.null_count + assert g.is_null().to_pylist() == w.is_null().to_pylist() + np.testing.assert_allclose( + g.to_numpy(), + w.to_numpy(), + rtol=tolerance, + atol=tolerance, + equal_nan=True, + ) + gv, wv = g.to_numpy(), w.to_numpy() + zero = (gv == 0) & (wv == 0) + np.testing.assert_array_equal( + np.signbit(gv[zero]), np.signbit(wv[zero]) + ) + else: + assert g.equals(w) + + +def test_dense_shapes_row_order_and_current_values(accelerated): + ctx, extension, tolerance = accelerated + rng = np.random.default_rng(412) + for m, k, n in [(3, 7, 5), (17, 3, 1), (1, 9, 13), (23, 11, 19)]: + a, b = rng.normal(size=(m, k)), rng.normal(size=(k, n)) + ctx.deregister_table("a") + ctx.deregister_table("b") + ctx.from_dataset("a", _matrix(a, ("i", "j"))) + ctx.from_dataset("b", _matrix(b, ("j", "k"))) + frame = ctx.sql(PRODUCT) + for scale in (1.0, -0.75): + a[:] *= scale + got = frame.to_dataset(dims=["i", "k"], template="a") + np.testing.assert_allclose( + got.v, a @ b, rtol=tolerance, atol=tolerance + ) + assert got.attrs == {"experiment": "accelerated"} + device = "gpu" if extension.precision == "float32" else "cpu" + assert extension.statistics()[f"{device}_matmul"] > 0 + + +@pytest.mark.parametrize( + "case", + [ + "null", + "duplicate", + "missing", + "infinity", + "nan", + "empty", + "keys", + "signed_zero", + "extreme", + ], +) +def test_sql_bags_nulls_special_values_and_empty_inputs(accelerated, case): + ctx, _, tolerance = accelerated + a = {"i": [0, 0, 1, 1], "j": [0, 1, 0, 1], "v": [1.0, 2.0, 3.0, 4.0]} + b = {"j": [0, 0, 1, 1], "k": [0, 1, 0, 1], "v": [0.25, 0.5, 0.75, 1.0]} + if case == "null": + a["v"] = [None, None, 3.0, 4.0] + elif case == "duplicate": + for values in a.values(): + values.append(values[0]) + elif case == "missing": + for values in b.values(): + values.pop() + elif case in ("infinity", "nan"): + a["v"][0] = np.inf if case == "infinity" else np.nan + elif case == "empty": + a = {key: [] for key in a} + elif case == "signed_zero": + a["v"] = [0.0, -0.0, -0.0, 0.0] + elif case == "extreme": + a["v"] = [1e300, -1e300, 1e-300, -1e-300] + elif case == "keys": + a["j"] = [None, -1, None, -1] + b["j"] = [None, None, -1, -1] + plain = Context() + for name, data in [("a", a), ("b", b)]: + batch = pa.record_batch( + data, + schema=pa.schema( + [ + (key, pa.float64() if key == "v" else pa.int64()) + for key in data + ] + ), + ) + for con in (ctx, plain): + con.register_record_batches(name, [[batch]]) + _same( + ctx.sql(PRODUCT).to_arrow_table(), + plain.sql(PRODUCT).to_arrow_table(), + tolerance, + ) + + +def test_prepared_forward_reads_writes_and_keeps_metadata(accelerated): + ctx, _, tolerance = accelerated + a = _matrix(np.ones((5, 3)), ("i", "j")) + b = _matrix(np.ones((3, 4)), ("j", "k")) + ctx.from_dataset("a", a) + ctx.from_dataset("b", b) + run = ctx.prepare(PRODUCT) + for scale in (1.0, 2.0, -0.5): + got = run(a=_matrix(np.full((5, 3), scale), ("i", "j"))) + np.testing.assert_allclose( + got["v"], np.full(20, 3 * scale), atol=tolerance + ) + + +ATTENTION = """ +WITH scores AS ( + SELECT q.i, k.j, SUM(q.v * k.v) / 2.0 AS s + FROM q JOIN k ON q.d = k.d WHERE {mask} GROUP BY q.i, k.j), +weights AS ( + SELECT i, j, exp(s - MAX(s) OVER (PARTITION BY i)) AS w FROM scores), +norm AS (SELECT i, SUM(w) AS z FROM weights GROUP BY i) +SELECT w.i, v.d, SUM(w.w / n.z * v.v) AS v +FROM weights w JOIN v ON w.j = v.j JOIN norm n ON w.i = n.i +GROUP BY w.i, v.d ORDER BY w.i, v.d +""" + + +@pytest.mark.parametrize( + "mask", ["q.i >= k.j", "q.i > k.j + 1", "q.i >= k.j AND q.i < k.j + 3"] +) +def test_attention_masks_are_absent_rows(accelerated, mask): + ctx, extension, tolerance = accelerated + plain = Context() + rng = np.random.default_rng(52) + datasets = [ + (name, _matrix(rng.normal(size=(9, 4)), dims)) + for name, dims in [ + ("q", ("i", "d")), + ("k", ("j", "d")), + ("v", ("j", "d")), + ] + ] + for con in (ctx, plain): + for name, data in datasets: + con.from_dataset(name, data) + sql = ATTENTION.format(mask=mask) + _same( + ctx.sql(sql).to_arrow_table(), + plain.sql(sql).to_arrow_table(), + tolerance, + ) + if extension.precision == "float32": + assert extension.statistics()["gpu_attention"] > 0 + # A non-finite value in a masked row must not contaminate earlier rows. + datasets[-1][1]["v"].values[-1, 0] = np.inf + _same( + ctx.sql(sql).to_arrow_table(), + plain.sql(sql).to_arrow_table(), + tolerance, + ) + + +def test_ddx_grad_vjp_jvp_stop_gradient_and_repeated_steps(accelerated): + ctx, _extension, tolerance = accelerated + rng = np.random.default_rng(3) + a, b = rng.normal(size=(7, 4)), rng.normal(size=(4, 3)) + ctx.from_dataset("a", _matrix(a, ("i", "j"))) + ctx.from_dataset("b", _matrix(b, ("j", "k"))) + loss = """WITH pred AS ( + SELECT a.i, b.k, SUM(a.v * b.v) AS p FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k), + loss AS (SELECT SUM(p*p)/2 AS l FROM pred) """ + select = "SELECT j,k,grad(loss,v) AS g FROM b ORDER BY j,k" + step = ctx.prepare(loss + select) + for scale in (1.0, -0.25, 2.0): + new = _matrix(b * scale, ("j", "k")) + rows = step(b=new) + np.testing.assert_allclose( + rows["g"].to_numpy().reshape(4, 3), + a.T @ (a @ new.v.values), + rtol=tolerance, + atol=tolerance, + ) + # JVP uses ddx's tangent relation and the same accelerated contractions. + tangent = _matrix(np.ones((4, 3)), ("j", "k")) + ctx.from_dataset("tangent", tangent) + got = ctx.sql("""SELECT a.i,b.k,jvp(SUM(a.v*b.v),b.v,tangent) AS g + FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k ORDER BY i,k""").to_arrow_table() + np.testing.assert_allclose( + got["g"].to_numpy().reshape(7, 3), + a @ np.ones((4, 3)), + rtol=tolerance, + atol=tolerance, + ) + ctx.sql( + "CREATE VIEW pred AS SELECT a.i,b.k,SUM(a.v*b.v) AS p FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k" + ) + ctx.sql("CREATE VIEW cotangent AS SELECT i,k,1.0 AS p FROM pred") + got = ctx.sql( + "WITH pred AS (SELECT a.i,b.k,SUM(a.v*b.v) AS p FROM a " + "JOIN b ON a.j=b.j GROUP BY a.i,b.k) " + "SELECT j,k,vjp(pred,v,cotangent) AS g FROM b ORDER BY j,k" + ).to_arrow_table() + np.testing.assert_allclose( + got["g"].to_numpy().reshape(4, 3), + a.T @ np.ones((7, 3)), + rtol=tolerance, + atol=tolerance, + ) + got = ctx.sql("""WITH loss AS (SELECT SUM(v*ddx_stop_gradient(v)) AS l FROM b) + SELECT j,k,grad(loss,v) AS g FROM b ORDER BY j,k""").to_arrow_table() + np.testing.assert_allclose( + got["g"].to_numpy().reshape(4, 3), b, rtol=tolerance, atol=tolerance + ) + + +def test_backend_threshold_retains_cpu_sql(): + try: + extension = MLX(precision="float32", min_ops=10**12) + except ImportError as exc: + pytest.skip(str(exc)) + ctx = Context(SessionConfig().with_target_partitions(2)).with_extensions( + extension + ) + ctx.from_dataset("a", _matrix(np.ones((3, 4)), ("i", "j"))) + ctx.from_dataset("b", _matrix(np.ones((4, 5)), ("j", "k"))) + assert ctx.sql(PRODUCT).to_arrow_table()["v"].to_pylist() == [4.0] * 15 + assert extension.statistics() == {} + + +def test_batched_products_across_partitions(accelerated): + ctx, _, tolerance = accelerated + rng = np.random.default_rng(985) + a, b = rng.normal(size=(3, 5, 7)), rng.normal(size=(3, 7, 4)) + ctx.from_dataset("a", _matrix(a, ("batch", "i", "j")), chunks={"batch": 1}) + ctx.from_dataset("b", _matrix(b, ("batch", "j", "k")), chunks={"batch": 1}) + got = ctx.sql("""SELECT a.batch,a.i,b.k,SUM(a.v*b.v) AS v + FROM a JOIN b ON a.batch=b.batch AND a.j=b.j + GROUP BY a.batch,a.i,b.k ORDER BY a.batch,a.i,b.k""").to_arrow_table() + np.testing.assert_allclose( + got["v"].to_numpy().reshape(3, 5, 4), + a @ b, + rtol=tolerance, + atol=tolerance, + ) + + +@pytest.mark.parametrize("mask", ["q.i >= k.j", "q.i < k.j + 2", "TRUE"]) +@pytest.mark.parametrize("per_head", [True, False]) +def test_attention_keeps_batches_and_heads_independent( + accelerated, mask, per_head +): + ctx, extension, tolerance = accelerated + plain = Context() + rng = np.random.default_rng(542) + for name, shape, dims in [ + ("q", (2, 3, 7, 4), ("seq", "h", "i", "d")), + ("k", (2, 3, 9, 4), ("seq", "h", "j", "d")), + ("v", (2, 3, 9, 5), ("seq", "h", "j", "e")), + ]: + data = _matrix(rng.normal(size=shape), dims) + ctx.from_dataset(name, data, chunks={"seq": 1}) + plain.from_dataset(name, data, chunks={"seq": 1}) + norm_keys = "seq,h,i" if per_head else "i" + norm_join = "a.seq=n.seq AND a.h=n.h AND a.i=n.i" if per_head else "a.i=n.i" + sql = f""" + WITH scores AS ( + SELECT q.seq,q.h,q.i,k.j,SUM(q.v*k.v)/2 AS s FROM q JOIN k + ON q.seq=k.seq AND q.h=k.h AND q.d=k.d WHERE {mask} + GROUP BY q.seq,q.h,q.i,k.j), + weights AS ( + SELECT seq,h,i,j,exp(s-MAX(s) OVER (PARTITION BY seq,h,i)) AS p + FROM scores), + norm AS (SELECT {norm_keys},SUM(p) AS z FROM weights GROUP BY {norm_keys}) + SELECT a.seq,a.h,a.i,v.e,SUM(a.p/n.z*v.v) AS v FROM weights a + JOIN norm n ON {norm_join} + JOIN v ON a.seq=v.seq AND a.h=v.h AND a.j=v.j + GROUP BY a.seq,a.h,a.i,v.e ORDER BY a.seq,a.h,a.i,v.e + """ + _same( + ctx.sql(sql).to_arrow_table(), + plain.sql(sql).to_arrow_table(), + tolerance, + ) + if extension.precision == "float32" and per_head: + assert extension.statistics().get("gpu_attention", 0) > 0 + + +def test_ddx_unsupported_expression_still_raises(accelerated): + ctx, _, _ = accelerated + ctx.from_dataset( + "weights", xr.Dataset({"w": ("i", [1.2, 2.4])}, coords={"i": [0, 1]}) + ) + with pytest.raises(UnsupportedExpression, match="GROUP BY key"): + ctx.sql("""WITH buckets AS ( + SELECT ROUND(w) AS k,SUM(w*w) AS s FROM weights GROUP BY ROUND(w)), + loss AS (SELECT SUM(k*s) AS l FROM buckets) + SELECT i,grad(loss,w) AS g FROM weights ORDER BY i""").collect() + + +def test_attention_backward_uses_ddx_derivatives(accelerated): + ctx, _, tolerance = accelerated + rng = np.random.default_rng(123) + q, k, v = (rng.normal(size=(7, 4)) for _ in range(3)) + for name, values, dims in [ + ("q", q, ("i", "d")), + ("k", k, ("j", "d")), + ("v", v, ("j", "d")), + ]: + ctx.from_dataset(name, _matrix(values, dims)) + scores = q @ k.T / 2 + scores[np.triu_indices(7, 1)] = -np.inf + weights = np.exp(scores - np.max(scores, axis=1, keepdims=True)) + weights /= weights.sum(axis=1, keepdims=True) + output = weights @ v + dw = output @ v.T + ds = weights * (dw - (dw * weights).sum(axis=1, keepdims=True)) + want = [ds @ k / 2, ds.T @ q / 2, weights.T @ output] + model = ATTENTION.format(mask="q.i >= k.j") + statements = [ + f"WITH attention AS ({model}), loss AS (SELECT SUM(v*v)/2 AS l FROM attention) " + f"SELECT {key},d,grad(loss,v) AS v FROM {name} ORDER BY {key},d" + for name, key in [("q", "i"), ("k", "j"), ("v", "j")] + ] + for frame, expected in zip(ad.sql_all(ctx, statements), want): + np.testing.assert_allclose( + frame.to_arrow_table()["v"].to_numpy().reshape(7, 4), + expected, + rtol=tolerance, + atol=tolerance, + ) + + +def test_low_backend_memory_limit_retains_cpu_sql(): + extension = MLX( + precision="float32", min_ops=0, min_gpu_ops=0, max_working_bytes=1 + ) + ctx = Context().with_extensions(extension) + ctx.from_dataset("a", _matrix(np.ones((3, 4)), ("i", "j"))) + ctx.from_dataset("b", _matrix(np.ones((4, 5)), ("j", "k"))) + assert ctx.sql(PRODUCT).to_arrow_table()["v"].to_pylist() == [4.0] * 15 + assert extension.statistics() == {} diff --git a/tests/test_prepare.py b/tests/test_prepare.py index 8cc356bc..1b8b2b53 100644 --- a/tests/test_prepare.py +++ b/tests/test_prepare.py @@ -9,6 +9,11 @@ from xarray_sql import Context +try: + from einfold.datafusion import EinFold +except ImportError: + EinFold = None + # The build side of the join, `a`, is what a reused hash join would keep. JOIN = """ SELECT b.i, SUM(t.v2 * b.w) AS s @@ -17,6 +22,26 @@ """ +@pytest.mark.parametrize("extension", [False, True]) +def test_join_key_data_columns_can_change_their_range(extension): + ctx = Context() + if extension: + if EinFold is None: + pytest.skip("requires the optional EinFold planner") + ctx = ctx.with_extensions(EinFold()) + a = xr.Dataset({"key": ("i", np.arange(128)), "v": ("i", np.ones(128))}) + b = xr.Dataset({"key": ("j", np.arange(256)), "w": ("j", np.arange(256.0))}) + ctx.from_dataset("a", a) + ctx.from_dataset("b", b) + sql = "SELECT a.i,b.j,a.v+b.w AS s FROM a JOIN b ON a.key=b.key ORDER BY a.i,b.j" + query = ctx.prepare(sql) + for start in (0, 128, 16): + updated = a.copy(deep=True) + updated["key"][:] = np.arange(start, start + 128) + rows = query(a=updated) + assert rows.equals(ctx.sql(sql).to_arrow_table()) + + def _a(v): return xr.Dataset( {"v": (("i", "j"), v)}, @@ -262,3 +287,50 @@ def test_a_quoted_table_replaced_in_sql_counts_too(): ctx.sql('CREATE OR REPLACE TABLE "Mixed" AS VALUES (1), (2)').collect() with pytest.raises(ValueError, match="registered again"): query() + + +def test_extension_handle_keeps_dataset_metadata_and_invalidation(ctx): + want = ctx.sql("SELECT * FROM a").to_dataset() + derived = ctx.with_extensions() + query = derived.prepare("SELECT i, j, v FROM a ORDER BY i, j") + xr.testing.assert_identical( + derived.sql("SELECT * FROM a").to_dataset(), want + ) + ctx.deregister_table("a") + with pytest.raises(ValueError, match="registered again"): + query() + + +@pytest.mark.parametrize( + "chunks", [{"i": 1, "j": 1}, {"i": 1}, {"i": 4, "j": 5}] +) +def test_fine_chunk_registration_keeps_rows_metadata_and_writes(chunks): + values = np.arange(129 * 33, dtype=np.float64).reshape(129, 33) + ds = _a(values) + ds.attrs["experiment"] = "fine chunks" + ctx = Context() + ctx.from_dataset("a", ds, chunks=chunks) + sql = "SELECT i,j,v FROM a ORDER BY i,j" + xr.testing.assert_identical(ctx.sql(sql).to_dataset(), ds) + run = ctx.prepare(sql) + replacement = ds * 2 + rows = run(a=replacement) + np.testing.assert_array_equal(rows["v"], replacement.v.values.reshape(-1)) + + +@pytest.mark.parametrize("threads", [1, 60]) +def test_thread_target_is_explicit_and_keeps_xarray_support(threads): + ctx = Context(threads=threads) + assert ctx.sql( + "SELECT value FROM information_schema.df_settings " + "WHERE name='datafusion.execution.target_partitions'" + ).to_pylist() == [{"value": str(threads)}] + ctx.from_dataset("a", _a(np.ones((4, 3)))) + assert ctx.prepare("SELECT SUM(v) AS s FROM a")().to_pylist() == [ + {"s": 12.0} + ] + + +def test_thread_target_must_be_positive(): + with pytest.raises(ValueError, match="positive"): + Context(threads=0) diff --git a/tests/test_shakespeare.py b/tests/test_shakespeare.py new file mode 100644 index 00000000..8ef6edd8 --- /dev/null +++ b/tests/test_shakespeare.py @@ -0,0 +1,121 @@ +"""Real SQL training can resume without changing sampled data or updates.""" + +import importlib.util +import json +import os +import sys +from pathlib import Path + +import numpy as np +import pytest + +ddx = pytest.importorskip("ddxdb") +pytest.importorskip("einfold.dense_program") +pytestmark = pytest.mark.skipif( + "erf" not in ddx.supported_functions(), + reason="requires the optional ddx erf patch", +) + + +@pytest.mark.parametrize( + "execution", + [ + "arrow", + "resident", + pytest.param( + "compiled", + marks=pytest.mark.skipif( + os.environ.get("MLX_ENABLE_TF32") != "0", reason="requires GPU" + ), + ), + ], +) +def test_shakespeare_sql_checkpoint_resume(tmp_path, monkeypatch, execution): + scripts = Path(__file__).resolve().parents[1] / "benchmarks/ml" + monkeypatch.syspath_prepend(str(scripts)) + spec = importlib.util.spec_from_file_location( + "_test_shakespeare", scripts / "03_shakespeare.py" + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + text = "To be, or not to be.\n" * 30 + vocab = sorted(set(text)) + ids = np.array([vocab.index(c) for c in text], dtype=np.uint16) + (tmp_path / "input.txt").write_text(text) + split = int(len(ids) * 0.9) + ids[:split].tofile(tmp_path / "train.bin") + ids[split:].tofile(tmp_path / "val.bin") + + def run(out, steps, resume=None): + args = [ + "03_shakespeare.py", + "--backend", + "dense_gpu" if execution == "compiled" else "dense_cpu", + "--data-dir", + str(tmp_path), + "--out-dir", + str(out), + "--steps", + str(steps), + "--batch_size", + "2", + "--block_size", + "4", + "--n_layer", + "1", + "--n_head", + "2", + "--n_embd", + "8", + "--eval_iters", + "1", + "--eval_interval", + "2", + ] + if execution == "resident": + args += ["--resident"] + elif execution == "compiled": + args += ["--compile-training", "--gpu-dropout"] + if resume: + args += ["--resume", str(resume)] + monkeypatch.setattr(sys, "argv", args) + module.main() + + full, resumed = tmp_path / "full", tmp_path / "resumed" + run(full, 3) + run(resumed, 2) + run(resumed, 1, resumed / "last.npz") + with ( + np.load(full / "last.npz", allow_pickle=False) as expected, + np.load(resumed / "last.npz", allow_pickle=False) as actual, + ): + for key in expected.files: + if key != "metadata": + if execution == "compiled": + np.testing.assert_allclose( + actual[key], + expected[key], + atol=3e-6, + rtol=3e-6, + err_msg=key, + ) + else: + np.testing.assert_array_equal( + actual[key], expected[key], err_msg=key + ) + a, b = ( + json.loads(str(actual["metadata"])), + json.loads(str(expected["metadata"])), + ) + assert a["iteration"] == b["iteration"] == 3 + assert a["data_rng"] == b["data_rng"] + assert a["mask_rng"] == b["mask_rng"] + assert a["best_val_loss"] == b["best_val_loss"] + events = [ + json.loads(line) + for line in (full / "metrics.jsonl").read_text().splitlines() + ] + assert [e["iteration"] for e in events if e["kind"] == "eval"] == [0, 2] + assert [e["lr"] for e in events if e["kind"] == "train"] == [ + 1e-3 * (i + 1) / 101 for i in range(3) + ] diff --git a/tests/test_shared.py b/tests/test_shared.py new file mode 100644 index 00000000..07a36905 --- /dev/null +++ b/tests/test_shared.py @@ -0,0 +1,296 @@ +"""Shared SELECT stages retain SQL row semantics and original ddx lineage.""" + +import numpy as np +import pyarrow as pa +import pytest +import xarray as xr + +from xarray_sql import Context + +pytest.importorskip("ddxdb") +einfold = pytest.importorskip("einfold._native") +pytestmark = pytest.mark.skipif( + not getattr(einfold, "RESET_ON_DECODE", False), + reason="requires the optional EinFold prepared-plan patch", +) + +from einfold.datafusion import EinFold + +from xarray_sql.shared import SharedExecutor + +try: + from xarray_sql.dense import DenseExecutor +except ImportError: + DenseExecutor = None + +_EXECUTORS = [ + SharedExecutor, + pytest.param( + DenseExecutor, + marks=pytest.mark.skipif( + DenseExecutor is None, reason="requires compact EinFold programs" + ), + ), +] + + +@pytest.fixture(params=_EXECUTORS) +def executor(request): + ctx = Context().with_extensions(EinFold()) + ctx.from_dataset("w", xr.Dataset({"v": ("i", [1.0, 2.0, 3.0])})) + return request.param(ctx) + + +def test_shared_join_refreshes_build_side_and_returns_snapshots(executor): + sql = """ + WITH square AS (SELECT i, v*v AS s FROM w), + doubled AS (SELECT a.i, a.s + b.s AS s FROM square a + JOIN square b ON a.i = b.i) + SELECT * FROM doubled ORDER BY i + """ + query = executor.prepare(sql) + first = query() + assert first.to_pydict() == {"i": [0, 1, 2], "s": [2.0, 8.0, 18.0]} + for values in ([3.0, 2.0, 1.0], [5.0, 6.0, 7.0]): + rows = query(w=xr.Dataset({"v": ("i", values)})) + assert rows.equals(executor.ctx.sql(sql).to_arrow_table()) + assert first.to_pydict()["s"] == [2.0, 8.0, 18.0] + + +def test_virtual_view_gradients_use_original_dependencies(executor): + executor.register_view("model", "SELECT i, v*v AS s FROM w") + loss = "WITH loss AS (SELECT SUM(s) AS l FROM model)" + forward = executor.prepare("SELECT SUM(s) AS l FROM model") + step = executor.prepare_all( + [ + loss + " SELECT i, grad(loss, v) AS g FROM w ORDER BY i", + loss + " SELECT i, v - 0.1*grad(loss, v) AS v FROM w ORDER BY i", + ] + ) + for values in ([1.0, 2.0, 3.0], [4.0, -2.0, 0.0]): + data = xr.Dataset({"v": ("i", values)}) + assert forward(w=data)["l"].to_pylist() == [np.dot(values, values)] + gradient, updated = step() + np.testing.assert_allclose(gradient["g"], 2 * np.array(values)) + np.testing.assert_allclose(updated["v"], 0.8 * np.array(values)) + + +def test_gradients_through_reused_ctes_in_a_virtual_view(executor): + executor.register_view( + "model", + """ + WITH squared AS (SELECT i, v*v AS a FROM w), + mean AS (SELECT AVG(a) AS m FROM squared), + centered AS (SELECT i, power(a-m, 2) AS s FROM squared CROSS JOIN mean), + total AS (SELECT SUM(s) AS t FROM centered) + SELECT i, s+t AS s FROM centered CROSS JOIN total + """, + ) + step = executor.prepare_all( + [ + ( + "WITH loss AS (SELECT SUM(s) AS l FROM model) " + "SELECT i, grad(loss,v) AS g FROM w ORDER BY i" + ) + ] + ) + for values in ([1.0, 2.0, 3.0], [-2.0, 4.0, 1.0]): + v = np.array(values) + [rows] = step(w=xr.Dataset({"v": ("i", v)})) + np.testing.assert_allclose(rows["g"], 16 * v * (v * v - np.mean(v * v))) + + +def test_registration_and_coordinate_changes_refresh_plans(executor): + query = executor.prepare( + "WITH t AS (SELECT i, v FROM w WHERE v > 0) " + "SELECT a.i, a.v+b.v AS v FROM t a JOIN t b ON a.i=b.i ORDER BY a.i" + ) + assert query().num_rows == 3 + for coords, values in (([0, 1, 2], [-1.0, 2.0, 4.0]), ([8, 9], [5.0, 6.0])): + executor.ctx.deregister_table("w") + executor.ctx.from_dataset( + "w", xr.Dataset({"v": ("i", values)}, coords={"i": coords}) + ) + assert query().to_pydict() == { + "i": [i for i, v in zip(coords, values) if v > 0], + "v": [2 * v for v in values if v > 0], + } + + +def test_empty_stages_can_gain_and_lose_rows_after_inplace_writes(executor): + query = executor.prepare( + "WITH t AS (SELECT i,v FROM w WHERE v < 0) " + "SELECT a.i,a.v+b.v AS v FROM t a JOIN t b ON a.i=b.i ORDER BY a.i" + ) + for values in ( + [1.0, 2.0, 3.0], + [2.0, 3.0, 4.0], + [-1.0, -2.0, 3.0], + [1.0, 2.0, 3.0], + ): + rows = query(w=xr.Dataset({"v": ("i", values)})) + assert rows.to_pydict() == { + "i": [i for i, v in enumerate(values) if v < 0], + "v": [2 * v for v in values if v < 0], + } + + +def test_changed_domains_refresh_every_shared_stage(executor): + query = executor.prepare(""" + WITH a AS (SELECT i,v FROM w WHERE v < 0), + b AS (SELECT x.i,x.v+y.v AS v FROM a x JOIN a y ON x.i=y.i), + c AS (SELECT x.i,x.v+y.v AS v FROM b x JOIN b y ON x.i=y.i) + SELECT x.i,x.v+y.v AS v FROM c x JOIN c y ON x.i=y.i ORDER BY x.i + """) + for values in ( + [1.0, 2.0, 3.0], + [-1.0, -2.0, 3.0], + [1.0, -3.0, -4.0], + [1.0, 2.0, 3.0], + ): + rows = query(w=xr.Dataset({"v": ("i", values)})) + assert rows.to_pydict() == { + "i": [i for i, v in enumerate(values) if v < 0], + "v": [8 * v for v in values if v < 0], + } + + +@pytest.mark.parametrize("executor_type", _EXECUTORS) +def test_shared_stages_keep_nulls_and_duplicate_rows(executor_type): + ctx = Context().with_extensions(EinFold()) + ctx.register_record_batches( + "w", + [ + [ + pa.record_batch( + { + "i": [0, 0, 1], + "v": [1.0, None, 3.0], + } + ) + ] + ], + ) + executor = executor_type(ctx) + sql = ( + "WITH t AS (SELECT i, v FROM w) " + "SELECT a.i, a.v AS a, b.v AS b FROM t a JOIN t b ON a.i=b.i " + "ORDER BY a.i, a.v NULLS FIRST, b.v NULLS FIRST" + ) + query = executor.prepare(sql) + for _ in range(2): + assert query().equals(ctx.sql(sql).to_arrow_table()) + + +def test_optimizer_scalar_functions_remain_available(executor): + rows = executor.prepare( + "WITH norm AS (SELECT sqrt(SUM(v*v)) AS n FROM w), " + "clip AS (SELECT least(1.0, 1.0/(n+1e-6)) AS c FROM norm) " + "SELECT i, v*clip.c AS v FROM w CROSS JOIN clip ORDER BY i" + )() + np.testing.assert_allclose( + rows["v"], np.array([1.0, 2.0, 3.0]) / (np.sqrt(14) + 1e-6) + ) + + +def test_ordered_stages_preserve_quoted_names_and_window_results(executor): + sql = """ + WITH t AS (SELECT i AS "axis.name", v FROM w), + doubled AS (SELECT a."axis.name", a.v+b.v AS v FROM t a JOIN t b + ON a."axis.name"=b."axis.name") + SELECT "axis.name", SUM(v) OVER ( + ORDER BY "axis.name" ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS s + FROM doubled ORDER BY "axis.name" + """ + query = executor.prepare(sql) + for values in ([3.0, 1.0, 2.0], [5.0, -1.0, 4.0]): + rows = query(w=xr.Dataset({"v": ("i", values)})) + assert rows.equals(executor.ctx.sql(sql).to_arrow_table()) + + +def test_many_gradient_contributions_share_a_bounded_reduction_graph(executor): + terms = " UNION ALL ".join( + f"SELECT i,power(v,{p}) AS y FROM w" for p in range(2, 42) + ) + step = executor.prepare_all( + [ + ( + f"WITH terms AS ({terms}), loss AS (SELECT SUM(y) AS l FROM terms) " + "SELECT i,grad(loss,v) AS g FROM w ORDER BY i" + ) + ] + ) + for values in ([1.0, 2.0, 3.0], [2.0, -1.0, 0.0]): + v = np.array(values) + [rows] = step(w=xr.Dataset({"v": ("i", v)})) + np.testing.assert_allclose( + rows["g"], sum(p * v ** (p - 1) for p in range(2, 42)) + ) + + +@pytest.mark.parametrize("executor_type", _EXECUTORS) +def test_matrix_gradients_through_a_reconstructed_normalized_operand( + executor_type, +): + rng = np.random.default_rng(48) + a, b, g = ( + rng.normal(size=(4, 3)), + rng.normal(size=(3, 5)), + rng.normal(size=3), + ) + ctx = Context().with_extensions(EinFold()) + ctx.from_dataset("a", xr.Dataset({"v": (("i", "j"), a)})) + ctx.from_dataset("b", xr.Dataset({"w": (("j", "k"), b)})) + ctx.from_dataset("g", xr.Dataset({"gain": ("j", g)})) + executor = executor_type(ctx) + executor.register_view( + "normalized", + """ + WITH mean AS (SELECT i,AVG(v) AS m FROM a GROUP BY i), + variance AS (SELECT a.i,AVG(power(a.v-mean.m,2)) AS v + FROM a JOIN mean ON a.i=mean.i GROUP BY a.i) + SELECT a.i,a.j,(a.v-mean.m)/sqrt(variance.v+1e-5)*g.gain AS v + FROM a JOIN mean ON a.i=mean.i JOIN variance ON a.i=variance.i + JOIN g ON a.j=g.j + """, + ) + step = executor.prepare_all( + [ + """ + WITH predicted AS (SELECT x.i,b.k,SUM(x.v*b.w) AS v + FROM normalized x JOIN b ON x.j=b.j GROUP BY x.i,b.k), + loss AS (SELECT SUM(v*v)/2 AS l FROM predicted) + SELECT j,k,grad(loss,w) AS g FROM b ORDER BY j,k + """ + ] + ) + centered = a - a.mean(axis=1, keepdims=True) + normalized = ( + centered + / np.sqrt(np.mean(centered**2, axis=1, keepdims=True) + 1e-5) + * g + ) + for scale in (1.0, -0.25): + current = b * scale + expected = normalized.T @ (normalized @ current) + [rows] = step(b=xr.Dataset({"w": (("j", "k"), current)})) + np.testing.assert_allclose( + rows["g"].to_numpy().reshape(3, 5), + expected, + rtol=1e-12, + atol=1e-12, + ) + + +@pytest.mark.parametrize( + "sql", + [ + "WITH t AS (SELECT random() AS x) SELECT * FROM t", + "WITH RECURSIVE t AS (SELECT 1) SELECT * FROM t", + "WITH t AS (WITH s AS (SELECT 1) SELECT * FROM s) SELECT * FROM t", + "SELECT 1; SELECT 2", + ], +) +def test_unsupported_staging_fails_explicitly(executor, sql): + with pytest.raises(ValueError): + executor.prepare(sql) diff --git a/tests/test_training.py b/tests/test_training.py new file mode 100644 index 00000000..1f89934a --- /dev/null +++ b/tests/test_training.py @@ -0,0 +1,171 @@ +"""Owned SQL training state preserves updates, failures and snapshots.""" + +import os + +import numpy as np +import pytest +import xarray as xr + +from xarray_sql import Context + +pytest.importorskip("ddxdb") +pytest.importorskip("einfold.dense_program") +from xarray_sql.training import TrainingSession + +_PRECISIONS = [ + "float64", + pytest.param( + "float32", + marks=pytest.mark.skipif( + os.environ.get("MLX_ENABLE_TF32") != "0", + reason="GPU requires MLX_ENABLE_TF32=0", + ), + ), +] + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_training_state_reuses_primal_and_owns_inputs(precision): + ctx = Context() + initial = xr.Dataset({"v": ("i", [1.0, 2.0, 3.0])}, coords={"i": [2, 4, 7]}) + ctx.from_dataset("w", initial) + session = TrainingSession(ctx, precision=precision) + session.bind("w") + initial.v.values[:] = 99 + step = session.prepare_step( + gradients={ + "g": "WITH loss AS (SELECT SUM(v*v)/2 AS l FROM w) SELECT i,grad(loss,v) AS g FROM w ORDER BY i" + }, + stages={}, + updates={ + "u": "SELECT p.i,p.v-0.1*g.g AS v FROM w p JOIN g ON p.i=g.i ORDER BY p.i" + }, + assignments={"w": ("u", ["i", "v"])}, + ) + before = session.export("w")["w"] + for k in range(3): + result = step() + np.testing.assert_allclose( + float(result.loss), 7 * 0.9 ** (2 * k), rtol=1e-6 + ) + np.testing.assert_allclose( + session.export("w")["w"]["v"], + np.array([1, 2, 3]) * 0.9 ** (k + 1), + rtol=1e-6, + ) + assert before.to_pydict() == {"i": [2, 4, 7], "v": [1.0, 2.0, 3.0]} + session.bind( + "w", xr.Dataset({"v": ("i", [4.0, 5.0, 6.0])}, coords={"i": [2, 4, 7]}) + ) + np.testing.assert_allclose(float(step().loss), 38.5, rtol=1e-6) + session.close() + with pytest.raises(RuntimeError, match="closed"): + step() + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_training_failure_keeps_all_state(precision): + ctx = Context() + for name, values in (("a", [1.0, 2.0]), ("b", [3.0, 4.0])): + ctx.from_dataset(name, xr.Dataset({"v": ("i", values)})) + session = TrainingSession(ctx, precision=precision) + session.bind("a") + session.bind("b") + before = session.export("a", "b") + step = session.prepare_step( + gradients={ + "g": "WITH loss AS (SELECT SUM(v*v) AS l FROM a) SELECT i,grad(loss,v) AS g FROM a ORDER BY i" + }, + stages={}, + updates={ + "good": "SELECT i,v+1 AS v FROM a ORDER BY i", + "bad": "SELECT i,ln(-v) AS v FROM b ORDER BY i", + }, + assignments={"a": ("good", ["i", "v"]), "b": ("bad", ["i", "v"])}, + ) + with pytest.raises(ValueError, match="finite|exceptional"): + step() + after = session.export("a", "b") + assert all(after[name].equals(before[name]) for name in before) + assert session.version == 0 + + +@pytest.mark.skipif( + os.environ.get("MLX_ENABLE_TF32") != "0", reason="requires GPU" +) +def test_compiled_training_uses_new_tokens_and_learning_rates(): + ctx = Context() + weights = np.arange(1, 7, dtype=float).reshape(3, 2) + ctx.from_dataset("w", xr.Dataset({"v": (("tok", "e"), weights)})) + session = TrainingSession(ctx) + session.bind("w") + session.bind("batch", xr.Dataset({"tok": ("pos", [2, 2, 0])})) + session.index_domains("batch", {"tok": ("w", "tok")}) + session.bind("batch") + session.bind("rate", xr.Dataset({"lr": xr.DataArray(0.1)})) + step = session.prepare_step( + gradients={ + "g": """WITH model AS ( + SELECT b.pos,w.e,w.v FROM batch b JOIN w ON b.tok=w.tok), + loss AS (SELECT SUM(v*v)/2 AS l FROM model) + SELECT tok,e,grad(loss,v) AS g FROM w ORDER BY tok,e""" + }, + stages={}, + updates={ + "u": "SELECT w.tok,w.e,w.v-r.lr*g.g AS v FROM w JOIN g ON w.tok=g.tok AND w.e=g.e CROSS JOIN rate r ORDER BY w.tok,w.e" + }, + assignments={"w": ("u", ["tok", "e", "v"])}, + compile=True, + ) + for tokens, lr in (([2, 2, 0], 0.1), ([1, 0, 1], 0.03), ([0, 0, 2], 0.07)): + session.bind("batch", xr.Dataset({"tok": ("pos", tokens)})) + session.bind("rate", xr.Dataset({"lr": xr.DataArray(lr)})) + expected_loss = np.sum(weights[tokens] ** 2) / 2 + weights = ( + weights - lr * np.bincount(tokens, minlength=3)[:, None] * weights + ) + result = step() + assert result.measurements["gpu_readback_bytes"] == 0 + np.testing.assert_allclose(float(result.loss), expected_loss, rtol=3e-6) + rows = session.export("w")["w"] + np.testing.assert_allclose( + np.asarray(rows["v"]).reshape(3, 2), weights, rtol=3e-6 + ) + before = session.export("w")["w"] + with pytest.raises(ValueError, match="domain"): + session.bind("batch", xr.Dataset({"tok": ("pos", [9, 0, 1])})) + assert session.export("w")["w"].equals(before) + + +def test_training_rejects_schema_and_coordinate_changes(): + ctx = Context() + ctx.from_dataset("w", xr.Dataset({"v": ("i", [1.0, 2.0])})) + session = TrainingSession(ctx, precision="float64") + session.bind("w") + before = session.export("w")["w"] + for ds in ( + xr.Dataset({"other": ("i", [1.0, 2.0])}), + xr.Dataset({"v": ("i", [1.0, 2.0])}, coords={"i": [4, 5]}), + ): + with pytest.raises(ValueError, match="schema|coordinate"): + session.bind("w", ds) + assert session.export("w")["w"].equals(before) + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_loss_snapshot_cannot_modify_scalar_state(precision): + ctx = Context() + ctx.from_dataset("w", xr.Dataset({"v": ("i", [2.0])})) + session = TrainingSession(ctx, precision=precision) + session.bind("w") + step = session.prepare_step( + gradients={ + "g": "WITH loss AS (SELECT SUM(v) AS l FROM w) SELECT grad(loss,v) AS g FROM w" + }, + stages={}, + updates={"u": "SELECT i,v FROM w"}, + assignments={"w": ("u", ["i", "v"])}, + ) + result = step() + result.loss[...] = 99 + assert session.export("w")["w"].to_pydict() == {"i": [0], "v": [2.0]} diff --git a/uv.lock b/uv.lock index 50583113..508b9263 100644 --- a/uv.lock +++ b/uv.lock @@ -622,14 +622,14 @@ wheels = [ [[package]] name = "einfold" -version = "0.1.0rc2" +version = "0.2.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/76/e9/b534c6660e42ae62f3bd430611b299f406833cd82757e39f74d0161290e6/einfold-0.1.0rc2.tar.gz", hash = "sha256:82df6a9e2e1954cc3ba90982b61a1236b0451e84a5f80ab4e7707c0e46f7950e", size = 224290, upload-time = "2026-10-30T02:08:18.251Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3d/03/7ea71305a4aa642c26224e598c49923a1f074659d479589c744c304a3eea/einfold-0.2.0.tar.gz", hash = "sha256:fe4e6970e8e04a53b0f4d180375e160a27f015b1f4f38294d94338468fa159c5", size = 266243, upload-time = "2026-11-01T01:35:47.288Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9a/19/986bd602078909cfa9bba3133d6f9cf5372c0e9f49de891aeeb657334187/einfold-0.1.0rc2-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:907bd43840bd3f0006436705a1f9a62939c2e8410e10911753c72975ce8edba0", size = 43853038, upload-time = "2026-10-30T02:07:57.966Z" }, - { url = "https://files.pythonhosted.org/packages/46/56/3045c7efff5a649f8f209d95175785d7ddb7a279925f2033627a62df0645/einfold-0.1.0rc2-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:1041aebdd7293dac48b680addfebb6bd90870e220a56263fe7fb43da4f3b1217", size = 42234075, upload-time = "2026-10-30T02:08:00.816Z" }, - { url = "https://files.pythonhosted.org/packages/d2/bf/9d73c11836034eceea3057e246e89beadc0b090b11c4bb88865e357102af/einfold-0.1.0rc2-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:55d10e500ef934322e1108fa24b85da37ce21b152a9b3dd0364bf604e93fb0e0", size = 48676482, upload-time = "2026-10-30T02:08:03.849Z" }, - { url = "https://files.pythonhosted.org/packages/86/63/8dc06cb2ffc0e8bc982b83d1ffd60b9bba30f7a5990c0f11254f9faa5825/einfold-0.1.0rc2-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:98cc8014c4e1f26101ecd492e9881ff3bafdddcf130692fa97a600a3f8047d6d", size = 47269964, upload-time = "2026-10-30T02:08:06.821Z" }, + { url = "https://files.pythonhosted.org/packages/97/a9/e3b22d401da30443d6aa42b4bcf21fce4736669a6a0b59368bc07f59b4fa/einfold-0.2.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:a1c4aa1f4177c0079088bfe2c4c9fc50fb7a5f34af7cbc8251bd4376d8740014", size = 43892338, upload-time = "2026-11-01T01:35:30.093Z" }, + { url = "https://files.pythonhosted.org/packages/03/f3/8e3125e25165eee2042f7da09806a980b1703d9daec51a0fdfcf22eba9a1/einfold-0.2.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:e0572d2282bb283012b814429ade84f395c8c72732ee2989c9d40093e9bf3031", size = 42425356, upload-time = "2026-11-01T01:35:34.403Z" }, + { url = "https://files.pythonhosted.org/packages/c5/69/a19afa99bcc6dfab3ef57411cda214e3aedd1eae8a8cfa7fd0fa1f7df754/einfold-0.2.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0d2dedfdbbc6f67f8864554e73c9a9718d85ae27c7e0cfc8ce547afddbb64239", size = 48717033, upload-time = "2026-11-01T01:35:39.06Z" }, + { url = "https://files.pythonhosted.org/packages/31/cf/86eb4757245705709f78d66982d90a6bff47ebf294d0777e4f6e8c3810ac/einfold-0.2.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:c2458223021307aaf3bff213954a1b7cad301eab51af16f6e24774dc0505a2d8", size = 47308031, upload-time = "2026-11-01T01:35:44.234Z" }, ] [package.optional-dependencies] @@ -1573,6 +1573,53 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5a/ff/bda40387b5c5c64254595f4d81a12351770856acc5de4e6d43606a31f161/ml_dtypes-0.6.0-cp315-cp315t-win_arm64.whl", hash = "sha256:f6cb525101b6b903779188c1e9e9490c343b455ab822883e02cf01e5547338d2", size = 572666, upload-time = "2026-08-13T14:14:38.993Z" }, ] +[[package]] +name = "mlx" +version = "0.32.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mlx-metal" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/52/73/f2b18f4bbcfeacd6880e68f2854609d6d2d305e5676586451a6c26664061/mlx-0.32.4-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:60ba41564e8b5018c9d8a5f8ce7a5aaf70c96d8dd98c1a54b854cbada362d471", size = 671799, upload-time = "2026-10-14T21:15:37.06Z" }, + { url = "https://files.pythonhosted.org/packages/7f/f0/5a722aa39986cf274438f0acc22e90304cdafdb7265045cd893924a5608f/mlx-0.32.4-cp310-cp310-macosx_15_0_arm64.whl", hash = "sha256:e635385a4f778b248ea6f61fbeb285b6bb14a1891d75cfbfd5af12e010298308", size = 671801, upload-time = "2026-10-14T21:15:39.02Z" }, + { url = "https://files.pythonhosted.org/packages/95/27/9069458b0079820b860678f28fcb6bd0a3998e17db35440fc9b2847d2145/mlx-0.32.4-cp310-cp310-macosx_26_0_arm64.whl", hash = "sha256:6dd2d7948c3708a62a1b13c32b03b5a94031e0cd3b7f5c5a0343e5a4a3b3bc41", size = 671791, upload-time = "2026-10-14T21:15:40.697Z" }, + { url = "https://files.pythonhosted.org/packages/7b/fd/f7af63951d6f4d39b47bbda462f2ed0341ce1293967de36ce9bbd4b836a9/mlx-0.32.4-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:dd0c9070d9c362fe280a0224806b1b47719b154984735041ebcc907e597c08ca", size = 671764, upload-time = "2026-10-14T21:15:48.321Z" }, + { url = "https://files.pythonhosted.org/packages/56/1d/e23a150073e07212c23358b978af0bfd21b87f0b97937ec1f300cee346ce/mlx-0.32.4-cp311-cp311-macosx_15_0_arm64.whl", hash = "sha256:fbb3fcf1363a631d05aa5adcb30d07655158be07e09339c1fe62f7267b3a78f5", size = 671761, upload-time = "2026-10-14T21:15:49.94Z" }, + { url = "https://files.pythonhosted.org/packages/c7/9b/3e5cfd79fd468dbc7f516e5fa7341419541804c5f8beefe04053326e25e9/mlx-0.32.4-cp311-cp311-macosx_26_0_arm64.whl", hash = "sha256:dfd5269398bb64372137d8c354530b6e81159ce0e92fd1ed7e9140b294fd939a", size = 671758, upload-time = "2026-10-14T21:15:51.859Z" }, + { url = "https://files.pythonhosted.org/packages/69/b0/e726ce51e1b445e4ed08c2d0e036a01c7532f4d1a856328da546a6f883bd/mlx-0.32.4-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:ff1d7be54758ea3a875b8fe54eff13bed17ad9566505e301f5d7851bf869680f", size = 667597, upload-time = "2026-10-14T21:16:00.688Z" }, + { url = "https://files.pythonhosted.org/packages/6d/92/7a999dcf6460525a31e555905cdca2f835ec0d786c528e969e04cd00d158/mlx-0.32.4-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:f0604cb3015091f285c79c7b7dbfcabcfa9948248728c7bab6e0263513c3d3fe", size = 667599, upload-time = "2026-10-14T21:16:02.467Z" }, + { url = "https://files.pythonhosted.org/packages/af/44/90bbdc3cdd120e1181f7c4f617ff996530159f2ea3ba3c08b83b08314d21/mlx-0.32.4-cp312-cp312-macosx_26_0_arm64.whl", hash = "sha256:7928869076db6267ab3ebdb253b2c80be1e738b3032c783170bd1d59d460c578", size = 667622, upload-time = "2026-10-14T21:16:04.219Z" }, + { url = "https://files.pythonhosted.org/packages/f6/25/b570700cb3d7b37477a865382880a6549b48b9efd2dd356dbc16eaaa01d1/mlx-0.32.4-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:c7375f2d6acd8c72f413d018266ae06387cb9caab31f475f3ff1c5600a7cb4c5", size = 667579, upload-time = "2026-10-14T21:16:12.885Z" }, + { url = "https://files.pythonhosted.org/packages/1c/ad/1377a281f1b81f42929a3887b5225d0a38fba2d804ac1039011a223d2ab2/mlx-0.32.4-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:d46dd07f57a65bae00ba7da239ee76ebf87e3202f042808d209b56be3140e416", size = 667585, upload-time = "2026-10-14T21:16:14.597Z" }, + { url = "https://files.pythonhosted.org/packages/e8/b1/33ca08f84550fd512f28b84055ff00631ef2ce83011ac2bbc171ffec7bbb/mlx-0.32.4-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:ef130e3da895c85374b11a6c389acb122893f2ffcbed826b8b62fd45a20b7ad6", size = 667575, upload-time = "2026-10-14T21:16:16.31Z" }, + { url = "https://files.pythonhosted.org/packages/98/85/f020b6392f04e37e9d2263daa194061736801ad51af06b1f11a253643976/mlx-0.32.4-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:43a6469797f09a985f8f035df63cd072105d5f16bfbe27d82b542065c988c3dd", size = 692145, upload-time = "2026-10-14T21:16:24.976Z" }, + { url = "https://files.pythonhosted.org/packages/df/db/95168f9492b284042628952778bfc2342c0f381ac2434b3640b80fa2f292/mlx-0.32.4-cp313-cp313t-macosx_15_0_arm64.whl", hash = "sha256:986f001874378d0a4c1f919817b7e331be96bbebf3993f5182459bcfc6ac0116", size = 692143, upload-time = "2026-10-14T21:16:26.705Z" }, + { url = "https://files.pythonhosted.org/packages/fc/81/c452f266d36c5708182ce079c2a6b6beca5c3405279216b28510368d80a7/mlx-0.32.4-cp313-cp313t-macosx_26_0_arm64.whl", hash = "sha256:733d28fc6b200910211924cd8d5deef86095cd085a0bf64502943c80b22942d5", size = 692095, upload-time = "2026-10-14T21:16:28.371Z" }, + { url = "https://files.pythonhosted.org/packages/99/2f/ed874a2de4ad30b47f27960373880aa24b53d8963ea81125e12e460340c2/mlx-0.32.4-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:f1e6ef441e1e2ce54159022a8fa5597f8ea4f7acf278752049c4f774bbd0501d", size = 668901, upload-time = "2026-10-14T21:16:36.971Z" }, + { url = "https://files.pythonhosted.org/packages/12/c6/1e15326918233e1ee74f55bcc5577f6cce1c4257e5a1bbadcdb30b589dca/mlx-0.32.4-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:83adc889ae89057fd9f292ca1a9ee56da43f181fccc7ae48a060d7aadc179eaa", size = 668899, upload-time = "2026-10-14T21:16:39.057Z" }, + { url = "https://files.pythonhosted.org/packages/95/2f/f84f5f16000adbdb1c53583f53da045165013ebb62f6943c169714be011f/mlx-0.32.4-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:1dd56491b44a2bfa217f60da66a5c9c889add7ad9dc2cbcf526389ea31d83bee", size = 668817, upload-time = "2026-10-14T21:16:40.738Z" }, + { url = "https://files.pythonhosted.org/packages/e3/7c/c90d57df7775421fe43f7a235f6cee85a67c4aecc7325051b23fff58335b/mlx-0.32.4-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:f6cbda9cb8c33c9f8673a2b69fb19bdab75b5032d4e47a42a0e7b36db1c3ebcc", size = 692106, upload-time = "2026-10-14T21:16:50.019Z" }, + { url = "https://files.pythonhosted.org/packages/6e/07/90e629a9667ec35b2aa8c84e0e096ff36fdd58731dc9d182f5ee59a488d0/mlx-0.32.4-cp314-cp314t-macosx_15_0_arm64.whl", hash = "sha256:4f95afac25431585b34846d472ea0ad5db83fc97fb014d62c70082456f6b6b5c", size = 692107, upload-time = "2026-10-14T21:16:51.978Z" }, + { url = "https://files.pythonhosted.org/packages/d0/20/bdb26b961fc8f0d54e6b36dd5360449f0cf8591d14cee10470c14ce1e154/mlx-0.32.4-cp314-cp314t-macosx_26_0_arm64.whl", hash = "sha256:63ac6fde8a55d688448fc25eb32b645276501e0a30ff961a72958dfd0a4065df", size = 692103, upload-time = "2026-10-14T21:16:53.963Z" }, + { url = "https://files.pythonhosted.org/packages/74/d7/6abf128b1740d8be1d711d8ded6ccd4714c53d8af8b6267ac6aae285c60e/mlx-0.32.4-cp315-cp315-macosx_14_0_arm64.whl", hash = "sha256:4f8e640fd3afcc19da343c5a769743f361573b8db98cd739e295a387921486a5", size = 669191, upload-time = "2026-10-14T21:17:02.83Z" }, + { url = "https://files.pythonhosted.org/packages/43/b7/3b9bcbb9a5d32025d032fb933b7c1a92bf1a4e22a796ff3b3ef7c2a7a53c/mlx-0.32.4-cp315-cp315-macosx_15_0_arm64.whl", hash = "sha256:6c6cf406fe093b4ea21876949e0d21f75384900f3df543e8411447ee37cdd394", size = 669189, upload-time = "2026-10-14T21:17:04.66Z" }, + { url = "https://files.pythonhosted.org/packages/45/10/1164007669c9e015e1124d1387ba03079ffb6650224cbc01aaba1a43ebfd/mlx-0.32.4-cp315-cp315-macosx_26_0_arm64.whl", hash = "sha256:869c4e2d99ef622cfb39e6581afe5d3286e0c01161e03ad580561db28bcd3992", size = 669097, upload-time = "2026-10-14T21:17:06.647Z" }, + { url = "https://files.pythonhosted.org/packages/a3/21/8b1b2f90bb3fa27d4bb5d0bbfb07c7f210844da6bb89f88476ab0ba84f50/mlx-0.32.4-cp315-cp315t-macosx_14_0_arm64.whl", hash = "sha256:95830850ababacf86d4e1e11b8cfa062a14b12c0c5f332dad1d4f99e1679a162", size = 692273, upload-time = "2026-10-14T21:17:16.639Z" }, + { url = "https://files.pythonhosted.org/packages/a5/90/a5f916a6937481cc5ebab3f461d4a3459175d0ddf7e0339ca15b97a42fc8/mlx-0.32.4-cp315-cp315t-macosx_15_0_arm64.whl", hash = "sha256:09da967a668de8b7033de3d74930efcf36ef993ede93173f51f11452a1c09b22", size = 692274, upload-time = "2026-10-14T21:17:18.385Z" }, + { url = "https://files.pythonhosted.org/packages/47/3e/7371d6df935a531dfcd1307c28d11d8308f8304bec25d0575a36e6c451e7/mlx-0.32.4-cp315-cp315t-macosx_26_0_arm64.whl", hash = "sha256:1e52124110303abb7513e6575f59c4c1d32e721c1a45f7a2ae200e8cc492ef13", size = 692240, upload-time = "2026-10-14T21:17:20.391Z" }, +] + +[[package]] +name = "mlx-metal" +version = "0.32.4" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/77/24/7e5b7ed5d28abacab4577c7a8dd1fca6e9058504f5f95ae4333c09420e2e/mlx_metal-0.32.4-py3-none-macosx_14_0_arm64.whl", hash = "sha256:8019b2db622599267d66cc168b2f0dbf5af929790a6c3b3913e00eabc42a7947", size = 44958619, upload-time = "2026-10-14T21:14:45.727Z" }, + { url = "https://files.pythonhosted.org/packages/54/8c/9754dd6c4db485037f59cc8be6495738e14cdeb8bf19d5bddb5e96da890f/mlx_metal-0.32.4-py3-none-macosx_15_0_arm64.whl", hash = "sha256:9f0a3cc26d89724434d9684932fc6f3e4e83669d2a895d858ed2c8a55eb84e8b", size = 44942974, upload-time = "2026-10-14T21:14:50.012Z" }, + { url = "https://files.pythonhosted.org/packages/2a/d6/4340ed6e184f71693d7f531c9ae55d677451351245760d71f8dd714dcbe5/mlx_metal-0.32.4-py3-none-macosx_26_0_arm64.whl", hash = "sha256:514df92d41a286895d2564b4b18f4da45096bd7a4a0c47e1042f345ad5a7f936", size = 78106861, upload-time = "2026-10-14T21:14:56.184Z" }, +] + [[package]] name = "multidict" version = "6.6.3" @@ -3132,6 +3179,10 @@ geo = [ ml = [ { name = "ddxdb", marker = "sys_platform != 'emscripten'" }, ] +mlx = [ + { name = "einfold", extra = ["datafusion"], marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, + { name = "mlx", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, +] polars = [ { name = "polars" }, ] @@ -3175,8 +3226,10 @@ requires-dist = [ { name = "datafusion", marker = "sys_platform != 'emscripten'", specifier = "==55.0.0" }, { name = "ddxdb", marker = "sys_platform != 'emscripten' and extra == 'ml'", specifier = ">=0.4.6,<0.5" }, { name = "duckdb", marker = "extra == 'duckdb'", specifier = ">=1.4.0" }, + { name = "einfold", extras = ["datafusion"], marker = "platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'mlx'", specifier = ">=0.2,<0.3" }, { name = "gcsfs", marker = "extra == 'test'" }, { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'" }, + { name = "mlx", marker = "platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'mlx'", specifier = ">=0.32.4,<0.33" }, { name = "polars", marker = "extra == 'polars'", specifier = ">=1.33" }, { name = "pre-commit", marker = "extra == 'dev'" }, { name = "pyarrow", specifier = ">=16.0.0" }, @@ -3190,12 +3243,12 @@ requires-dist = [ { name = "xarray-sql", extras = ["docs"], marker = "extra == 'dev'" }, { name = "zensical", marker = "extra == 'docs'" }, ] -provides-extras = ["adbc", "dask", "dev", "docs", "duckdb", "geo", "ml", "polars", "test"] +provides-extras = ["adbc", "dask", "dev", "docs", "duckdb", "geo", "ml", "mlx", "polars", "test"] [package.metadata.requires-dev] dev = [ { name = "ddxdb", extras = ["duckdb"], marker = "sys_platform != 'emscripten'", specifier = ">=0.4.6,<0.5" }, - { name = "einfold", extras = ["datafusion"], marker = "sys_platform != 'emscripten' and sys_platform != 'win32'", specifier = ">=0.1.0rc2,<0.2" }, + { name = "einfold", extras = ["datafusion"], marker = "sys_platform != 'emscripten' and sys_platform != 'win32'", specifier = ">=0.1.0rc2,<0.3" }, { name = "jax", marker = "platform_machine != 'x86_64' or sys_platform != 'darwin'" }, { name = "maturin", specifier = ">=1.9.1" }, { name = "py-spy", specifier = ">=0.4.0" }, diff --git a/xarray_sql/_dense_mlx.py b/xarray_sql/_dense_mlx.py new file mode 100644 index 00000000..1ffe642f --- /dev/null +++ b/xarray_sql/_dense_mlx.py @@ -0,0 +1,365 @@ +"""MLX array operations for EinFold's compact numerical programs. + +[MLX](https://github.com/ml-explore/mlx) evaluates the arithmetic chosen by +[EinFold](https://github.com/xqlsystems/einfold). Compiled functions take +every numerical leaf as an argument, so a cached function never captures +parameter values from an earlier training step. Coordinates stay on the CPU. +""" + +from __future__ import annotations + +import os +import string +from collections import Counter + +import mlx.core as mx +import numpy as np +from einfold.dense_program import _OPS, _align, _Decline, _op + + +def _remainder(a, b): + result = mx.remainder(a, b) + return mx.where((result != 0) & ((a < 0) != (b < 0)), result - b, result) + + +_FUNCTIONS = { + "+": mx.add, + "-": mx.subtract, + "*": mx.multiply, + "/": mx.divide, + "%": _remainder, + "=": mx.equal, + "!=": mx.not_equal, + "<": mx.less, + "<=": mx.less_equal, + ">": mx.greater, + ">=": mx.greater_equal, + "and": mx.logical_and, + "or": mx.logical_or, + "not": mx.logical_not, + "neg": mx.negative, + "where": mx.where, + "abs": mx.abs, + "sqrt": mx.sqrt, + "exp": mx.exp, + "erf": mx.erf, + "ln": mx.log, + "log": mx.log10, + "tanh": mx.tanh, + "isnan": mx.isnan, + "iszero": lambda x: x == 0, + "signum": mx.sign, + "sign": mx.sign, + "sin": mx.sin, + "cos": mx.cos, + "power": mx.power, + "pow": mx.power, + "floor": mx.floor, + "ceil": mx.ceil, + "minimum": mx.minimum, + "maximum": mx.maximum, + "float64": lambda x: x.astype(mx.float32), + "int64": lambda x: x.astype(mx.int64), + "bool": lambda x: x.astype(mx.bool_), +} + + +class _Runtime: + def __init__(self): + if os.environ.get("MLX_ENABLE_TF32") != "0": + raise ValueError("launch Python with MLX_ENABLE_TF32=0") + self.compiled = {} + self.counts = Counter() + self.tracing = False + + def owns(self, array): + return isinstance(array, mx.array) + + def boolean(self, value): + self.counts["gpu_status_readbacks"] += 1 + return bool(value.item()) + + def status(self, values, flags): + status = mx.array(True) + for flag in flags: + status = mx.logical_and(status, mx.all(self._mx(flag))) + for value in values: + status = mx.logical_and( + status, mx.all(mx.isfinite(self._mx(value))) + ) + return status + + def validate(self, values, flags): + status = self.status(values, flags) + mx.eval(values, status) + if not self.boolean(status): + raise ValueError( + "training update must be rectangular, nonnullable and finite" + ) + + def compile_training(self, function): + self.counts["gpu_training_compilations"] += 1 + return mx.compile(function) + + def finish_training(self, arrays, status): + self.counts["gpu_training_executions"] += 1 + mx.eval(arrays, status) + if not self.boolean(status): + raise ValueError( + "training update must be rectangular, nonnullable and finite" + ) + + def indices(self, values): + if values.size and ( + int(values.min()) < -(1 << 63) or int(values.max()) >= (1 << 63) + ): + raise _Decline("device index is outside exact int64 range") + self.counts["gpu_input_arrays"] += 1 + self.counts["gpu_input_bytes"] += values.nbytes + return mx.array(values, dtype=mx.int64) + + def random_key(self, seed, step): + return mx.bitwise_xor( + mx.random.key(seed), mx.array([0, step], dtype=mx.uint32) + ) + + def dropout(self, key, shapes, probability): + keys = mx.random.split(key, len(shapes)) + return [ + mx.random.bernoulli(1 - probability, shape=shape, key=k).astype( + mx.float32 + ) + / (1 - probability) + for k, shape in zip(keys, shapes) + ] + + def memory(self): + return { + "active_bytes": mx.get_active_memory(), + "cache_bytes": mx.get_cache_memory(), + "peak_bytes": mx.get_peak_memory(), + } + + def copy(self, value): + return mx.array(value) + + def numpy(self, value): + if self.owns(value): + if self.tracing: + raise _Decline( + "compiled numerical graph requires a host conversion" + ) + self.counts["gpu_readbacks"] += 1 + self.counts["gpu_readback_bytes"] += value.size * value.itemsize + return np.asarray(value) + + def array(self, value): + array = np.asarray(value) + if array.dtype.kind == "f": + # ddx uses the extreme float64 value as a masked MAX sentinel. + # It stays beyond every supported finite float32 input. An + # observable infinity is rejected at the final SQL boundary. + if array.ndim == 0 and abs(array) == np.finfo(np.float64).max: + return mx.array(float("-inf") if array < 0 else float("inf")) + finite = np.isfinite(array) + if (array.ndim and not finite.all()) or np.any( + np.abs(array[finite]) > np.finfo(np.float32).max + ): + raise _Decline( + f"GPU input outside finite float32 range: {array.dtype}, shape={array.shape}, finite={finite.sum()}/{array.size}" + ) + self.counts["gpu_input_arrays"] += 1 + self.counts["gpu_input_bytes"] += array.nbytes + return mx.array(array, dtype=mx.float32) + return array + + def _mx(self, value): + return value if self.owns(value) else mx.array(value) + + def pointwise(self, tensor): + leaves, instructions, seen = [], [], {} + + def describe(node): + if node in seen: + return seen[node] + if node.kind in _OPS and node.cached is None and len(seen) < 64: + children = tuple(describe(a) for a in node.args) + instruction = node.kind, children + else: + leaves.append(_align(node.get(), node.axes, tensor.axes)) + instruction = "leaf", len(leaves) - 1 + at = len(instructions) + seen[node] = at + instructions.append(instruction) + return at + + describe(tensor) + signature = tuple(instructions) + if not any(self.owns(a) for a in leaves): + return None + if signature not in self.compiled: + + def execute(*args): + values = [] + for kind, children in signature: + if kind == "leaf": + values.append(args[children]) + else: + values.append( + _FUNCTIONS[kind](*(values[c] for c in children)) + ) + return values[-1] + + self.compiled[signature] = mx.compile(execute) + self.counts["gpu_fused_programs"] += 1 + self.counts["gpu_fused_executions"] += 1 + return self.compiled[signature](*(self._mx(a) for a in leaves)) + + def take(self, data, indexes, axis): + return ( + mx.take(self._mx(data), self._mx(indexes), axis=axis) + if self.owns(data) or self.owns(indexes) + else np.take(data, indexes, axis=axis) + ) + + def indexes(self, domain, values): + if domain.dtype.kind not in "iu" or not len(domain): + raise _Decline("GPU index domain must contain integers") + if int(domain.min()) < -(1 << 63) or int(domain.max()) >= (1 << 63): + raise _Decline("GPU index domain is outside exact int64 range") + domain = domain.astype(np.int64, copy=False) + order = np.argsort(domain) + sorted_domain = mx.array(domain[order]) + at = mx.minimum(mx.searchsorted(sorted_domain, values), len(domain) - 1) + indexes = mx.take(mx.array(order), at) + return indexes, mx.take(mx.array(domain), indexes) == values + + def index(self, data, indexes): + return ( + self._mx(data)[tuple(self._mx(i) for i in indexes)] + if self.owns(data) or any(self.owns(i) for i in indexes) + else data[tuple(indexes)] + ) + + def contract(self, arrays, labels, output): + converted = [self._mx(a) for a in arrays] + if not any(mx.issubdtype(a.dtype, mx.floating) for a in converted): + for array in arrays: + if self.owns(array): + if array.dtype != mx.bool_: + raise _Decline( + "GPU integer contraction requires a count mask" + ) + elif np.any((array < -1) | (array > 1)): + raise _Decline( + "GPU integer contraction requires unit factors" + ) + dimensions = { + label: size + for a, names in zip(converted, labels) + for label, size in zip(names, a.shape) + } + used = tuple(sorted(dimensions)) + value = mx.array(1, dtype=mx.int64) + for a, names in zip(converted, labels): + ordered = [n for n in used if n in names] + a = a.transpose(tuple(names.index(n) for n in ordered)) + a = a.reshape( + tuple(dimensions[n] if n in names else 1 for n in used) + ) + value = value * a.astype(mx.int64) + value = mx.sum( + value, + axis=tuple(i for i, n in enumerate(used) if n not in output), + ) + remaining = [n for n in used if n in output] + return value.transpose(tuple(remaining.index(n) for n in output)) + equation = ( + ",".join( + "".join(string.ascii_letters[i] for i in label) + for label in labels + ) + + "->" + + "".join(string.ascii_letters[i] for i in output) + ) + # Mask/count operands are promoted before contraction; a boolean + # reduction otherwise counts presence rather than SQL rows. + arrays = [a.astype(mx.float32) for a in converted] + self.counts["gpu_contractions"] += 1 + return mx.einsum(equation, *arrays) + + def broadcast(self, data, source, target, shape): + return mx.broadcast_to(_align(self._mx(data), source, target), shape) + + def extreme(self, kind, data, source, mask, mask_axes, used, kept): + shape = tuple(len(a) for a in used) + if any(size == 0 for size in shape): + raise _Decline("empty GPU extremum") + data = self.broadcast(data, source, used, shape) + mask = self.broadcast(mask, mask_axes, used, shape) + data = mx.where( + mask, data, -float("inf") if kind == "max" else float("inf") + ) + axes = tuple(i for i, a in enumerate(used) if a not in kept) + return (mx.max if kind == "max" else mx.min)(data, axis=axes) + + def scatter(self, data, indexes, size, *, count_bound): + data = self._mx(data) + integer = data.dtype == mx.int64 + if integer: + if count_bound > np.iinfo(np.int32).max: + raise _Decline("GPU integer scatter exceeds exact int32 range") + data = data.astype(mx.int32) + result = mx.zeros((data.shape[0], size), dtype=data.dtype) + self.counts["gpu_scatter_reductions"] += 1 + result = result.at[ + mx.arange(data.shape[0])[:, None], mx.array(indexes) + ].add(data) + return result.astype(mx.int64) if integer else result + + def merge(self, parts, keys, key, axes, *, flags=None): + result = mx.zeros(tuple(len(a) for a in axes), dtype=mx.float32) + for part in parts: + source = tuple(part.columns[k].tensor.axis for k in keys) + indexes = [ + np.searchsorted(a.values, s.values) + for a, s in zip(axes, source) + ] + shape = tuple(len(a) for a in source) + if flags is not None: + mask = ( + part.present + if flags == "rows" + else _op("and", part.present, part.columns[key].valid) + ) + data = mx.ones(shape, dtype=mx.float32) + else: + col = part.columns[key] + value = col.tensor.get() + if self.owns(value) and not mx.issubdtype( + value.dtype, mx.floating + ): + raise _Decline( + "GPU integer partial aggregate requires relational execution" + ) + mask = _op("and", part.present, col.valid) + data = self.broadcast(value, col.tensor.axes, source, shape) + valid = self.broadcast(mask.get(), mask.axes, source, shape) + data = mx.where(valid, data, 0) + if axes: + grid = [ + mx.array( + i.reshape( + tuple( + len(i) if j == at else 1 + for j in range(len(axes)) + ) + ) + ) + for at, i in enumerate(indexes) + ] + result = result.at[tuple(grid)].add(data) + else: + result = result + data + self.counts["gpu_union_reductions"] += 1 + return result diff --git a/xarray_sql/ad.py b/xarray_sql/ad.py index d0ff2c8f..eb12555a 100644 --- a/xarray_sql/ad.py +++ b/xarray_sql/ad.py @@ -43,7 +43,11 @@ [sql_all][xarray_sql.ad.sql_all] runs several statements that take ``grad`` of the same loss for the price of one backward pass, which is -how a training step updates each of a model's parameter tables. For the +how a training step updates each of a model's parameter tables. For repeated +training steps, [prepare][xarray_sql.ad.prepare] retains ddx's +programs between calls and materializes each statement's result. +[Context.prepare][xarray_sql.Context.prepare] handles one such statement; +``ad.prepare`` handles several statements sharing a loss. For the programs underneath (``grad``, ``vjp``, ``jvp`` and ``run``, to build a backward pass once and run it every step), use ``ddxdb.ad`` directly; they take a ``Context`` like any other ``SessionContext``. @@ -52,12 +56,13 @@ from __future__ import annotations import re +import threading from collections.abc import Sequence from datafusion import SessionContext from ddxdb import ad as _ddx -__all__ = ["STOP_GRADIENT", "calls_ad", "register", "sql", "sql_all"] +__all__ = ["STOP_GRADIENT", "calls_ad", "prepare", "register", "sql", "sql_all"] STOP_GRADIENT = _ddx.STOP_GRADIENT """The SQL name of the stop-gradient function, ``ddx_stop_gradient``.""" @@ -138,3 +143,44 @@ def sql_all(ctx: SessionContext, statements: Sequence[str]) -> list: frames: list = _ddx.sql_all(ctx, list(statements)) wrap = getattr(ctx, "_wrap", None) return [wrap(f) for f in frames] if wrap else frames + + +def prepare(ctx: SessionContext, statements: Sequence[str]): + """Prepare SQL gradient statements for repeated execution. + + ddx retains the differentiation programs and checks table schemas, + views and session settings on each call. This avoids rebuilding those + programs in a training loop while allowing tables to be registered + again. Several statements can share a loss, as with ``sql_all``:: + + step = xql.ad.prepare(ctx, [ + f"WITH loss AS ({LOSS}) " + "SELECT i, val - 0.01 * grad(loss, val) AS val FROM w ORDER BY i" + ]) + [rows] = step() + + On a ``Context``, a call may first write Dataset values into registered + tables, using the same validation as ``Context.prepare``:: + + [rows] = step(w=new_weights) + + Returns a callable producing one materialized ``pyarrow.Table`` per + statement. Results remain valid after later calls. Writes and execution + are serialized for calls to this prepared function. Other users of the + context must coordinate concurrent writes themselves. + """ + if isinstance(statements, str): + raise TypeError("statements is a list of SQL statements") + run = _ddx.prepare(ctx, list(statements)) + lock = threading.Lock() + + def execute(**tables): + with lock: + if tables: + write = getattr(ctx, "_write", None) + if write is None: + raise TypeError("Dataset writes require xarray_sql.Context") + write(tables) + return [frame.to_arrow_table() for frame in run()] + + return execute diff --git a/xarray_sql/dense.py b/xarray_sql/dense.py new file mode 100644 index 00000000..f3ef3685 --- /dev/null +++ b/xarray_sql/dense.py @@ -0,0 +1,277 @@ +"""Experimental compact execution of SQL numerical programs. + +[EinFold](https://github.com/xqlsystems/einfold) lowers DataFusion plans to +array operations. Its optional ``dense_program`` module retains compact +arrays between shared SQL and [ddx](https://github.com/xqlsystems/ddx) +automatic differentiation stages. Public results remain Arrow tables. +The original SQL executor handles unsupported plans without changing SQL +or gradient definitions. Install the patch in ``accelerators/einfold_mlx``. +""" + +from __future__ import annotations + +import os +import time +from typing import Any + +import pyarrow as pa +from datafusion import SessionContext +from ddxdb import ad as _ddx +from ddxdb import prepare as _prepare_ad +from einfold.dense_program import ( + _CACHE, + _EXECUTION, + _MAX_BYTES, + _OPS, + DenseProgram, + _ComputedCache, + _Decline, + _read_dataset, +) + +from xarray_sql.ad import calls_ad +from xarray_sql.shared import SharedExecutor, _Backend, _Graph, _snapshot + +__all__ = ["DenseExecutor"] + +_runtime_factory: type[Any] | None = None +if os.environ.get("MLX_ENABLE_TF32") == "0": + try: + from xarray_sql._dense_mlx import _Runtime + except ImportError: + pass + else: + _runtime_factory = _Runtime + + +class _Compiled: + def __init__(self, owner, frame): + self.schema = frame.schema() + self.program = DenseProgram(frame.logical_plan(), self.schema) + self.epoch = owner._epoch + owner._counts["dense_programs"] += 1 + + def run(self, owner): + started = time.perf_counter() + result = self.program.run(owner._read) + owner._counts["dense_executions"] += 1 + owner._times["dense_lowering_seconds"] += time.perf_counter() - started + return result + + def store(self, owner, name, table): + owner._store[name] = table + owner._stub(name, self.schema) + + +class _DenseBackend(_Backend): + def compile(self, plan): + return _Compiled(self.owner, _ddx._dataframe(self.ctx, plan)) + + def materialize_compiled(self, name, plan, compiled): + compiled.store(self.owner, name, compiled.run(self.owner)) + + def returns_rows_compiled(self, plan, compiled): + table = compiled.run(self.owner) + # Checks may have arbitrary relational shapes; materialize their + # actual values, never infer emptiness from schema-only providers. + return table.arrow(compiled.schema).num_rows > 0 + + def drop_table(self, name): + self.owner._store.pop(name, None) + self.owner._stubs.discard(name) + super().drop_table(name) + + +class _DenseGraph(_Graph): + def run(self, stack=(), *, rebind=False): + owner, aliases = self.owner, {} + for name in sorted(self.tables & owner._views.keys()): + if name in stack: + raise ValueError("cyclic virtual views cannot be staged") + graph = owner._view_graph(name) + table = graph.run((*stack, name)) + graph.compiled[""].store(owner, graph.names[""], table) + aliases[name] = f'SELECT * FROM "{graph.names[""]}"' + for name in self.order: + compiled = self.compiled.get(name) + if compiled is None or compiled.epoch != owner._epoch: + compiled = self.compiled[name] = _Compiled( + owner, + SessionContext.sql(owner.ctx, self._sql(name, aliases)), + ) + table = compiled.run(owner) + if name: + compiled.store(owner, self.names[name], table) + return table + + +class DenseExecutor(SharedExecutor): + """Execute supported SQL regions using compact arrays across ddx stages. + + Input Datasets are read afresh on every call, including in-place + writes. All final Arrow tables are realized before returning, and no + parameter array is mutated during execution. ``fallback=False`` is a + diagnostic mode that reports unsupported regions immediately. + """ + + def __init__( + self, + ctx, + *, + precision="float64", + fallback=True, + max_array_bytes=1 << 30, + max_cached_bytes=8 << 30, + **options, + ): + if precision not in {"float64", "float32"}: + raise ValueError("precision must be float64 or float32") + if precision == "float32" and os.environ.get("MLX_ENABLE_TF32") != "0": + raise ValueError("launch Python with MLX_ENABLE_TF32=0") + if precision == "float32" and _runtime_factory is None: + raise ImportError("MLX is required for float32 GPU execution") + if max_array_bytes <= 0: + raise ValueError("max_array_bytes must be positive") + if max_cached_bytes <= 0: + raise ValueError("max_cached_bytes must be positive") + super().__init__(ctx, **options) + self._max_array_bytes = max_array_bytes + self._max_cached_bytes = max_cached_bytes + self._runtime = ( + _runtime_factory() + if precision == "float32" and _runtime_factory is not None + else None + ) + self._backend = _DenseBackend(self) + self._store = {} + self._stubs = set() + self._fallback = SharedExecutor(ctx, **options) if fallback else None + + def register_view(self, name, sql): + super().register_view(name, sql) + if self._fallback is not None: + self._fallback.register_view(name, sql) + + def _view_graph(self, name): + if name not in self._graphs: + self._graphs[name] = _DenseGraph(self, self._views[name]) + return self._graphs[name] + + def _stub(self, name, schema): + if name not in self._stubs: + empty = pa.RecordBatch.from_arrays( + [pa.array([], type=f.type) for f in schema], schema=schema + ) + self.ctx.deregister_table(name) + self.ctx.register_record_batches(name, [[empty]]) + self._stubs.add(name) + + def _read(self, name): + if name not in self._store: + if name not in self.ctx._registered_datasets: + raise _Decline( + f"source {name} has no rectangular Dataset contract" + ) + self._store[name] = _read_dataset( + name, self.ctx._registered_datasets[name] + ) + return self._store[name] + + def _invoke(self, execute, fallback, tables): + with self._lock: + started = time.perf_counter() + self.ctx._write(tables) + self._refresh() + self._store.clear() + token = _EXECUTION.set(self._runtime) + limit_token = _MAX_BYTES.set(self._max_array_bytes) + cache = _ComputedCache(self._max_cached_bytes) + cache_token = _CACHE.set(cache) + try: + overridden = ( + self.ctx._custom_udfs - self.ctx._trusted_math_udfs + ) & ( + set(_OPS) + | { + "coalesce", + "nullif", + "sum", + "avg", + "min", + "max", + "count", + } + ) + if overridden: + raise _Decline( + "custom scalar functions require SQL execution" + ) + return execute() + except _Decline as error: + self._counts["dense_fallbacks"] += 1 + self._last_decline = str(error) + if fallback is None: + raise + # Bind current function definitions when restarting in SQL. + self._fallback._epoch += 1 + self._fallback._backend._catalog = [] + return fallback() + finally: + self._store.clear() + _EXECUTION.reset(token) + _MAX_BYTES.reset(limit_token) + _CACHE.reset(cache_token) + self._counts["dense_cache_evictions"] += cache.evictions + self._counts["dense_peak_cached_bytes"] = max( + self._counts["dense_peak_cached_bytes"], cache.peak + ) + self._times["dense_call_seconds"] += ( + time.perf_counter() - started + ) + + def prepare(self, sql): + if calls_ad(sql): + run = self.prepare_all([sql]) + return lambda **tables: run(**tables)[0] + graph = _DenseGraph(self, sql) + fallback = ( + self._fallback.prepare(sql) if self._fallback is not None else None + ) + + def execute(): + result = graph.run() + return _snapshot(result.arrow(graph.compiled[""].schema)) + + return lambda **tables: self._invoke(execute, fallback, tables) + + def prepare_all(self, statements): + prepared = _prepare_ad(self._backend, statements) + finals = {} + fallback = ( + self._fallback.prepare_all(statements) + if self._fallback is not None + else None + ) + + def execute(): + with prepared.run() as rewritten: + results = [] + for sql in rewritten.statements: + if sql not in finals or finals[sql].epoch != self._epoch: + finals[sql] = _Compiled( + self, SessionContext.sql(self.ctx, sql) + ) + compiled = finals[sql] + results.append( + _snapshot(compiled.run(self).arrow(compiled.schema)) + ) + return results + + return lambda **tables: self._invoke(execute, fallback, tables) + + def statistics(self): + return ( + super().statistics() + | (dict(self._runtime.counts) if self._runtime is not None else {}) + | {"dense_last_decline": getattr(self, "_last_decline", None)} + ) diff --git a/xarray_sql/math.py b/xarray_sql/math.py new file mode 100644 index 00000000..2a9c1e1d --- /dev/null +++ b/xarray_sql/math.py @@ -0,0 +1,40 @@ +"""Optional SQL mathematical primitives with equivalent array lowering. + +DataFusion (https://datafusion.apache.org/) supplies most numerical SQL +functions. This module supplies the error function for exact GELU, the +activation used by nanoGPT (https://github.com/karpathy/nanoGPT). +""" + +import math + +import numpy as np +import pyarrow as pa +from datafusion import udf + +try: + from scipy.special import erf as _erf +except ImportError: + _erf = np.vectorize(math.erf, otypes=[np.float64]) + +__all__ = ["register_erf"] + + +def _arrow_erf(values): + return pa.array( + _erf(values.to_numpy(zero_copy_only=False)), + mask=values.is_null().to_numpy(zero_copy_only=False), + type=pa.float64(), + ) + + +def register_erf(ctx): + """Install SQL ``erf(DOUBLE)`` and its equivalent CPU/GPU primitive. + + Differentiation also requires a ddxdb build with an erf rule; see + ``accelerators/ddx_math``. Replacing this UDF removes its trusted + primitive status so custom SQL behavior is never silently bypassed. + """ + ctx.register_udf( + udf(_arrow_erf, [pa.float64()], pa.float64(), "immutable", name="erf") + ) + ctx._trusted_math_udfs.add("erf") diff --git a/xarray_sql/mlx.py b/xarray_sql/mlx.py new file mode 100644 index 00000000..c81134b4 --- /dev/null +++ b/xarray_sql/mlx.py @@ -0,0 +1,258 @@ +"""Optional Apple silicon execution for SQL tensor computations. + +[EinFold](https://github.com/xqlsystems/einfold) recognizes tensor operations +in DataFusion plans. This extension supplies its validated dense matrices to +[MLX](https://github.com/ml-explore/mlx), Apple's array framework. SQL parsing, +row existence, NULL handling and [ddx](https://github.com/xqlsystems/ddx) +gradient construction remain in their original engines. + +Build the native hook with ``python accelerators/einfold_mlx/build.py``. +Then install it on a context before planning queries:: + + # Launch Python with MLX_ENABLE_TF32=0 before importing MLX. + from xarray_sql import Context + from xarray_sql.mlx import MLX + + accelerator = MLX(precision="float32") + ctx = Context().with_extensions(accelerator) + +``float32`` explicitly permits reduced precision for dense arithmetic on +the GPU, including arithmetic in ddx's forward and backward SQL programs. +The default ``float64`` uses NumPy's CPU matrix multiplication and EinFold's +CPU attention; MLX's Metal GPU does not implement double precision. Neither +mode replaces ddx with MLX autodiff. Unsupported shapes and exceptional +values retain EinFold's CPU implementation. No table values are cached. +""" + +from __future__ import annotations + +import os +import threading +from collections import Counter +from typing import Literal + +import mlx.core as mx +import numpy as np +import pyarrow as pa +from einfold.datafusion import EinFold + +__all__ = ["MLX"] + + +class MLX: + """EinFold extension with an optional MLX GPU arithmetic backend. + + Args: + precision: ``float64`` keeps double precision on the CPU. + ``float32`` permits rounding inputs and dense results to single + precision on the GPU. Results retain their SQL Arrow dtype. + min_ops: Minimum scalar products before crossing the native/Python + boundary. Smaller computations use the original CPU kernel. + min_gpu_ops: Minimum scalar products for GPU execution. Smaller + dense products use float64 CPU matrix multiplication; smaller + attention uses EinFold's CPU tiles. + max_working_bytes: Upper bound on estimated extra backend buffers. + Larger computations retain the CPU kernel. This is separate + from DataFusion's own memory reservation. + query_block: Maximum queries in an attention block. Bounds mask + storage independently of the number of queries. + + The native hook is supplied by the pinned feature build in + ``accelerators/einfold_mlx``. An ordinary EinFold installation raises + an actionable error instead of silently claiming GPU execution. + """ + + def __init__( + self, + *, + precision: Literal["float64", "float32"] = "float64", + min_ops: int = 1 << 20, + min_gpu_ops: int = 1 << 33, + max_working_bytes: int = 512 << 20, + query_block: int = 1024, + ) -> None: + if precision not in ("float64", "float32"): + raise ValueError("precision must be 'float64' or 'float32'") + if precision == "float32" and os.environ.get("MLX_ENABLE_TF32") != "0": + raise ValueError( + "Launch Python with MLX_ENABLE_TF32=0 for full float32 " + "arithmetic; MLX reads this process setting on first use" + ) + if ( + min_ops < 0 + or min_gpu_ops < 0 + or max_working_bytes <= 0 + or not 1 <= query_block <= 1024 + ): + raise ValueError("invalid operation, memory or block limit") + self.precision = precision + self._backend = _Backend( + precision, min_ops, min_gpu_ops, max_working_bytes, query_block + ) + try: + self._extension = EinFold( + dense_backend=self._backend.run, + dense_min_ops=min_ops, + dense_max_bytes=max_working_bytes, + ) + except TypeError as exc: + raise ImportError( + "EinFold needs the native dense backend hook; run " + "'.venv/bin/python accelerators/einfold_mlx/build.py' " + "from the xarray-sql checkout" + ) from exc + + def __datafusion_session_components__(self, ctx): + return self._extension.__datafusion_session_components__(ctx) + + def __datafusion_session_planner__(self, ctx, fallback): + return self._extension.__datafusion_session_planner__(ctx, fallback) + + def statistics(self) -> dict[str, int]: + """Completed CPU/GPU kernels and declined backend requests. + + Counts describe execution, including repeated collections, rather + than the presence of a node in an explained plan. + """ + return self._backend.statistics() + + +class _Backend: + def __init__( + self, + precision: Literal["float64", "float32"], + min_ops: int, + min_gpu_ops: int, + max_working_bytes: int, + query_block: int, + ) -> None: + self.precision = precision + self.min_ops = min_ops + self.min_gpu_ops = min_gpu_ops + self.max_working_bytes = max_working_bytes + self.query_block = query_block + self._counts: Counter[str] = Counter() + self._lock = threading.Lock() + + def statistics(self) -> dict[str, int]: + with self._lock: + return dict(self._counts) + + def _count(self, key: str) -> None: + with self._lock: + self._counts[key] += 1 + + def _decline(self): + self._count("declined") + + def run(self, kind, columns, shape, scale, ranges): + # The native side has already checked positional keys, duplicates, + # SQL NULLs and the domain. It also refuses non-finite factors. + w, x, z, y = shape + ops = w * x * z * y if kind == "matmul" else w * z * (x + y) + cells = sum(len(c) for c in columns) + output = w * x * y if kind == "matmul" else w * y + if ops < self.min_ops or (cells + output) * 32 > self.max_working_bytes: + return self._decline() + arrays = [c.to_numpy(zero_copy_only=True) for c in columns] + if self.precision == "float64" or ops < self.min_gpu_ops: + if kind != "matmul": + return self._decline() + a, b = arrays + result = np.matmul(a.reshape(w, x, z), b.reshape(w, z, y)) + device = "cpu" + else: + # Keep overflow, underflow and exceptional SQL arithmetic on + # the CPU. In particular, masked infinities must stay absent + # rows instead of participating in a zero-weighted product. + maximum = [] + for a in arrays: + magnitudes = np.abs(a) + hi = magnitudes.max(initial=0.0) + if ( + not np.isfinite(hi) + or hi > 1e15 + or np.any((magnitudes > 0) & (magnitudes < 1e-15)) + ): + return self._decline() + maximum.append(hi) + if kind == "attention": + mode, constant = scale + multiplier = constant if mode == 1 else 1.0 + if mode == 2: + if constant == 0: + return self._decline() + multiplier = 1.0 / constant + # Scores too close to single precision's range, including + # infinity, keep the original double precision algorithm. + bound = maximum[0] * maximum[1] * x * abs(multiplier) + if not np.isfinite(bound) or bound > 1e30: + return self._decline() + block_size = self.query_block + if kind == "attention": + available = self.max_working_bytes - (cells + output) * 32 + block_size = min(block_size, available // (z * 16)) + if block_size < 1: + return self._decline() + with mx.stream(mx.gpu): + inputs = [mx.array(a.astype(np.float32)) for a in arrays] + if kind == "matmul": + a, b = inputs + result = mx.matmul(a.reshape(w, x, z), b.reshape(w, z, y)) + mx.eval(result) + result = np.asarray(result).astype(np.float64) + else: + q, kt, v = inputs + q, k, v = ( + q.reshape(w, x), + kt.reshape(x, z).T, + v.reshape(z, y), + ) + # Query tiles bound the mask and temporary score storage. + # Ranges encode all of EinFold's normalized key masks, + # including strict, offset and sliding-window masks. + result = np.zeros((w, y), dtype=np.float64) + for start in range(0, w, block_size): + stop = min(w, start + block_size) + lo, hi = np.asarray(ranges[start:stop]).T + nonempty = lo < hi + if not nonempty.any(): + continue + key_start = int(lo[nonempty].min()) + key_stop = int(hi[nonempty].max()) + key_count = key_stop - key_start + causal_hi = np.clip( + np.arange(stop - start) + + key_count + - (stop - start) + + 1, + 0, + key_count, + ) + if np.all(lo == key_start) and np.all(hi == key_stop): + allowed = None + elif np.all(lo == key_start) and np.array_equal( + hi - key_start, causal_hi + ): + # MLX aligns a causal mask at the lower right. + # Its last query then sees this slice's last key. + allowed = "causal" + else: + keys = mx.arange(key_start, key_stop)[None, :] + allowed = (keys >= mx.array(lo[:, None])) & ( + keys < mx.array(hi[:, None]) + ) + block = mx.fast.scaled_dot_product_attention( + q[None, None, start:stop, :], + k[None, None, key_start:key_stop, :], + v[None, None, key_start:key_stop, :], + scale=multiplier, + mask=allowed, + ) + mx.eval(block) + result[start:stop] = np.asarray(block)[0, 0] + device = "gpu" + if not np.isfinite(result).all(): + return self._decline() + self._count(f"{device}_{kind}") + return pa.array(result.reshape(-1), type=pa.float64()) diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py index 8ec88bcf..db7c82c9 100644 --- a/xarray_sql/reader.py +++ b/xarray_sql/reader.py @@ -11,6 +11,7 @@ from __future__ import annotations import itertools +import math import os import sys from collections.abc import Callable, Iterable, Iterator @@ -21,9 +22,9 @@ import xarray as xr from .df import ( + DEFAULT_BATCH_SIZE, Block, Chunks, - DEFAULT_BATCH_SIZE, _block_len, _block_metadata, _block_slices_from_resolved, @@ -296,11 +297,10 @@ def _in_memory_partitions( # isn't contiguous in memory, and flattening it would copy. More # partitions than cores only adds overhead, since an in-memory table has # no partitions to prune. + # Count the Cartesian chunk grid without constructing slices that this + # in-memory path never reads, even for very finely chunked inputs. count = min( - sum( - 1 - for _ in _block_slices_from_resolved(ds, resolve_chunks(ds, chunks)) - ), + math.prod(len(axis) for axis in resolve_chunks(ds, chunks).values()), os.cpu_count() or 1, ) # `per_partition` batches of `size` rows each make `count` partitions, @@ -324,7 +324,7 @@ def read_xarray_table( _iteration_callback: ( Callable[[Block, list[str] | None], None] | None ) = None, -) -> "LazyArrowStreamTable": +) -> LazyArrowStreamTable: """Create a lazy DataFusion table from an xarray Dataset. This is the simplest way to register xarray data with DataFusion. diff --git a/xarray_sql/shared.py b/xarray_sql/shared.py new file mode 100644 index 00000000..40376cf7 --- /dev/null +++ b/xarray_sql/shared.py @@ -0,0 +1,708 @@ +"""Explicit shared-stage execution for deterministic SQL training graphs. + +[EinFold](https://github.com/xqlsystems/einfold) recognizes array operations +inside DataFusion plans. This executor materializes reused common table +expressions (CTEs) between those plans. An installed +[MLX](https://github.com/ml-explore/mlx) extension can execute the numerical +kernels on Apple's GPU. Stage boundaries currently exchange Arrow tables. + +[ddx](https://github.com/xqlsystems/ddx) differentiates the original SQL, +including virtual view definitions, before choosing its own materialization +boundaries. Forward staging never replaces the query ddx differentiates. + +This experimental API accepts SELECTs with flat, nonrecursive WITH clauses +and the deterministic numerical functions listed below. It requires the +native feature build in ``accelerators/einfold_mlx``. Tables may be updated through +``Context._write`` or registered again between calls; calls on this executor +are serialized. Coordinate/schema changes rebuild physical plans. +""" + +from __future__ import annotations + +import threading +import time +import uuid +from collections import Counter, OrderedDict +from collections.abc import Sequence + +import einfold._native as _einfold +import numpy as np +import pyarrow as pa +from datafusion import SessionContext +from ddxdb import UnsupportedExpression +from ddxdb import ad as _ddx +from ddxdb import prepare as _prepare_ad +from ddxdb._ddxdb import to_sql as _to_sql + +from xarray_sql._native import _SortedArrowTable +from xarray_sql._native import shared_ctes as _shared_ctes +from xarray_sql.ad import calls_ad +from xarray_sql.sql import _copy, _tables_read + +__all__ = ["SharedExecutor"] + +_FUNCTIONS = frozenset( + [ + "abs", + "avg", + "ceil", + "coalesce", + "count", + "ddx_stop_gradient", + "exp", + "erf", + "floor", + "greatest", + "least", + "ln", + "log", + "max", + "min", + "nullif", + "pow", + "power", + "round", + "sign", + "sin", + "cos", + "sqrt", + "sum", + "tanh", + "row_number", + ] +) +_DYNAMIC_FILTERS = tuple( + f"datafusion.optimizer.enable{kind}_dynamic_filter_pushdown" + for kind in ("", "_join", "_topk", "_aggregate") +) + + +def _quote(name): + return '"' + name.replace('"', '""') + '"' + + +def _snapshot(rows): + return pa.Table.from_arrays( + [ + pa.concat_arrays(c.chunks) + if c.chunks + else pa.array([], type=c.type) + for c in rows.columns + ], + schema=rows.schema, + ) + + +def _canonical(rows, keys): + """Avoid sorting the already ordered positional output of dense kernels.""" + if rows.num_rows < 2: + return rows + tied = None + for name, _ in keys: + column = rows[name] + if column.null_count or not pa.types.is_integer(column.type): + return rows.sort_by(keys) + values = column.to_numpy(zero_copy_only=False) + descending = values[1:] < values[:-1] + if tied is not None: + descending &= tied + if descending.any(): + return rows.sort_by(keys) + equal = values[1:] == values[:-1] + tied = equal if tied is None else tied & equal + if not np.any(tied): + return rows + return rows + + +def _join_only(plan): + """Join rows can remain logical until a consumer reduces them.""" + joined = False + stack = [plan] + allowed = { + "Join", + "CrossJoin", + "Projection", + "Filter", + "SubqueryAlias", + "TableScan", + } + while stack: + node = stack.pop() + kind = type(node.to_variant()).__name__ + if kind not in allowed: + return False + joined |= kind in {"Join", "CrossJoin"} + stack.extend(node.inputs()) + return joined + + +class _Compiled: + def __init__(self, owner, frame, *, deferred=False, graph=None): + self.frame, self.schema = frame, frame.schema() + self.deferred = deferred + self.graph = graph + self.plan = None + self.epoch = owner._epoch + self.batch = None + self.inputs = { + name: owner._versions.get(name, 0) + for name in _tables_read(frame.logical_plan()) + } + if not deferred and graph is None: + self._physical(owner) + + def _physical(self, owner): + start = time.perf_counter() + self.plan = self.frame.execution_plan() + owner._counts["physical_plans"] += 1 + owner._times["planning_seconds"] += time.perf_counter() - start + + def valid(self, owner): + return self.epoch == owner._epoch and all( + owner._versions.get(name, 0) == generation + for name, generation in self.inputs.items() + ) + + def run(self, owner): + if self.graph is not None: + # ddx recreates transient view providers each pass. Rebind the + # compound stage to those providers while retaining its buffers. + return self.graph.run(rebind=True) + if self.plan is None: + self._physical(owner) + start = time.perf_counter() + plan = _copy(owner.ctx, self.plan) + batches = [ + b.to_pyarrow() + for p in range(plan.partition_count) + for b in owner.ctx.execute(plan, p) + ] + # Physical simplification can prove a boolean nonnullable while + # the logical schema still permits NULL. Keep the declared schema + # when collecting, without copying or converting column values. + batches = [ + pa.RecordBatch.from_arrays(b.columns, schema=self.schema) + for b in batches + ] + owner._counts["plan_executions"] += 1 + owner._times["execution_seconds"] += time.perf_counter() - start + return pa.Table.from_batches(batches, self.schema) + + def store(self, owner, name, rows): + started = time.perf_counter() + keys = [ + (f.name, "ascending") + for f in rows.schema + if pa.types.is_integer(f.type) or pa.types.is_string(f.type) + ] + if keys: + rows = _canonical(rows, keys) + arrays = [c.combine_chunks() for c in rows.columns] + compatible = ( + self.batch is not None + and self.batch.schema == rows.schema + and self.batch.num_rows == rows.num_rows + ) + if compatible and self.batch is not None: + # Integer/string keys must retain both their domain and order: + # cached scan statistics were derived from these coordinates. + compatible = all( + new.null_count == 0 and old.null_count == 0 + if pa.types.is_floating(new.type) + else new.equals(old) + for new, old in zip(arrays, self.batch.columns) + ) + if compatible and self.batch is not None: + for new, old in zip(arrays, self.batch.columns): + if pa.types.is_floating(new.type): + width = new.type.byte_width + a = memoryview(new.buffers()[1])[ + new.offset * width : (new.offset + len(new)) * width + ] + memoryview(old.buffers()[1])[:] = a + owner._counts["stage_buffer_writes"] += 1 + else: + if self.batch is not None: + if self.batch.num_rows != rows.num_rows: + reason = "rows" + elif self.batch.schema != rows.schema: + reason = "schema" + else: + reason = "keys_or_nulls" + owner._counts[f"stage_replacements_{reason}"] += 1 + owner._versions[name] = owner._versions.get(name, 0) + 1 + # concat_arrays always owns its value buffers, including a + # single input chunk. No stage overwrites its producer's data. + self.batch = pa.RecordBatch.from_arrays( + [owner._intern_key(a) for a in arrays], schema=rows.schema + ) + owner._counts["stage_buffer_allocations"] += 1 + owner.ctx.deregister_table(name) + if keys: + owner.ctx.register_table( + name, _SortedArrowTable(self.batch, [k for k, _ in keys]) + ) + else: + owner.ctx.register_record_batches(name, [[self.batch]]) + owner._times["stage_storage_seconds"] += time.perf_counter() - started + + +class _Backend(_ddx.DataFusionBackend): + def __init__(self, owner): + super().__init__(owner.ctx) + self.owner = owner + self._catalog: list = [] + self._changes = 0 + self._deferred: dict[str, _Compiled] = {} + + def select_all(self, name): + deferred = self._deferred.get(name) + if deferred is None: + return super().select_all(name) + # ddx asks for a named Read with the saved table's schema. A view + # would expand to its entire definition in Substrait. Serialize a + # schema-only table, then restore the view before any execution. + empty = pa.RecordBatch.from_arrays( + [pa.array([], type=f.type) for f in deferred.schema], + schema=deferred.schema, + ) + self.ctx.deregister_table(name) + try: + self.ctx.register_record_batches(name, [[empty]]) + return super().select_all(name) + finally: + self.ctx.deregister_table(name) + self.ctx.register_view(name, deferred.frame) + + def drop_table(self, name): + self._deferred.pop(name, None) + super().drop_table(name) + + def materialize(self, name, plan): + super().materialize(name, plan) + self.owner._versions[name] = self.owner._versions.get(name, 0) + 1 + + def view_sql(self, name): + if name in self.owner._views: + return self.owner._views[name] + return super().view_sql(name) + + def _bind_views(self, sql): + # ddx synthesizes nested WITH queries while extracting shared CTEs. + # Rejecting those probes makes it abandon staging and inline the + # entire model. The restriction on user forward graphs is separate. + parts = _shared_ctes(sql, allow_nested=True) + defined = {name for name, *_ in parts} + needed = { + ref + for _, _, refs, _ in parts + for ref in refs + if ref in self.owner._views and ref not in defined + } + if not needed: + return sql + # ddx extracts reused CTEs from large views. A simple view may be + # left to the engine; bind its original SELECT in that case too. + definitions = ", ".join( + f"{_quote(name)} AS ({self.owner._views[name]})" + for name in sorted(needed) + ) + return f'WITH {definitions} SELECT * FROM ({sql}) AS "__xql_original"' + + def plan_sql(self, sql): + return super().plan_sql(self._bind_views(sql)) + + def plan_key(self, sql): + return super().plan_key(self._bind_views(sql)) + + def session_key(self): + # Internal shared stages may be created while preparing a program. + # They are execution state, rather than changes to its public inputs. + try: + settings = ( + SessionContext.sql( + self.ctx, + "SELECT name,value FROM information_schema.df_settings ORDER BY name", + ) + .to_arrow_table() + .to_pylist() + ) + except ValueError: + return None + catalog = [] + for c in sorted(self.ctx.catalog_names()): + for s in sorted(self.ctx.catalog(c).schema_names()): + if s == "information_schema": + continue + schema = self.ctx.catalog(c).schema(s) + for t in sorted(schema.table_names()): + if t.startswith(("__ddx_", "__xql_shared_")): + continue + table = schema.table(t) + path = ".".join(_quote(p) for p in (c, s, t)) + value = ( + self.ctx.table(path).logical_plan() + if table.kind == "view" + else table.schema + ) + catalog.append((path, value)) + if catalog != self._catalog: + self._catalog, self._changes = catalog, self._changes + 1 + return repr((self._changes, settings, self.owner._views)) + + def compile(self, plan): + # DataFusion 55's Substrait consumer applies Read projections only + # to TableScan roots; a deferred view expands to a join instead. + # ddx's SQL writer preserves the same plan and handles those reads. + frame = ( + SessionContext.sql(self.ctx, _to_sql(plan, "datafusion")) + if self._deferred + else _ddx._dataframe(self.ctx, plan) + ) + stack = [frame.logical_plan()] + unions = 0 + while stack: + node = stack.pop() + inputs = node.inputs() + if type(node.to_variant()).__name__ == "Union": + unions += max(0, len(inputs) - 1) + stack.extend(inputs) + if unions >= 32: + # Large cotangent sums have many independently reduced UNION + # branches. Materialize those reductions sequentially; running + # all their join build sides together can exhaust memory. + graph = self._stage_graph(plan, reductions=True) + if graph is not None: + self.owner._counts["union_reduction_graphs"] += 1 + return _Compiled(self.owner, frame, graph=graph) + deferred = _join_only(frame.logical_plan()) + if deferred: + try: + parts = _shared_ctes( + _to_sql(plan, "datafusion"), allow_nested=True + ) + deferred = all(set(fns) <= _FUNCTIONS for _, _, _, fns in parts) + except (UnsupportedExpression, ValueError): + deferred = False + return _Compiled(self.owner, frame, deferred=deferred) + + def _stage_graph(self, plan, **options): + try: + sql = _to_sql(plan, "datafusion") + return _Graph(self.owner, sql, **options) + except (UnsupportedExpression, ValueError): + return None + + def _refresh_compiled(self, plan, compiled): + if compiled.valid(self.owner): + return False + fresh = self.compile(plan) + batch = compiled.batch + if ( + compiled.graph is not None + and fresh.graph is not None + and compiled.graph.sql == fresh.graph.sql + ): + fresh.graph = compiled.graph + compiled.__dict__.update(fresh.__dict__) + compiled.batch = batch + return True + + def materialize_compiled(self, name, plan, compiled): + changed = self._refresh_compiled(plan, compiled) + if compiled.deferred: + if changed: + self.owner._versions[name] = ( + self.owner._versions.get(name, 0) + 1 + ) + self._deferred[name] = compiled + self.ctx.deregister_table(name) + self.ctx.register_view(name, compiled.frame) + self.owner._counts["deferred_stages"] += 1 + return + rows = compiled.run(self.owner) + compiled.store(self.owner, name, rows) + compiled.epoch = self.owner._epoch + + def returns_rows_compiled(self, plan, compiled): + self._refresh_compiled(plan, compiled) + return compiled.run(self.owner).num_rows > 0 + + +class _Graph: + def __init__(self, owner, sql, *, reductions=False): + self.owner = owner + self.reductions = reductions + self.sql = sql + self.parts = _shared_ctes(sql) + self.defs = {name: (query, refs) for name, query, refs, _ in self.parts} + if len(self.defs) != len(self.parts): + raise ValueError("duplicate CTE names cannot be staged") + for _, _, _, functions in self.parts: + unknown = set(functions) - _FUNCTIONS + if unknown: + raise ValueError( + "shared execution requires deterministic numerical SQL; " + f"unsupported functions: {', '.join(sorted(unknown))}" + ) + needed = set() + + def visit(name, stack): + if name in stack: + raise ValueError("cyclic CTE dependencies cannot be staged") + if name in needed: + return + needed.add(name) + for dep in self.defs[name][1]: + if dep in self.defs: + visit(dep, stack | {name}) + + visit("", set()) + counts = Counter( + dep + for name in needed + for dep in self.defs[name][1] + if dep in needed + ) + self.boundaries = { + name + for name, count in counts.items() + if count > 1 and name not in owner.inline_ctes + } + if reductions: + self.boundaries.update( + name + for name, _, _, functions in self.parts + if name + and name in needed + and set(functions) & {"sum", "avg", "min", "max", "count"} + ) + self.names = { + name: owner._prefix + uuid.uuid4().hex + for name in self.boundaries | {""} + } + self.order = [name for name, *_ in self.parts if name in self.names] + self.compiled = {} + self.tables = { + dep + for name in needed + for dep in self.defs[name][1] + if dep not in self.defs + } + + def _sql(self, target, aliases): + definitions = dict(aliases) + + def include(name): + if name in definitions: + return + if name in self.boundaries and name != target: + definitions[name] = f"SELECT * FROM {_quote(self.names[name])}" + return + for dep in self.defs[name][1]: + if dep in self.defs: + include(dep) + definitions[name] = self.defs[name][0] + + for dep in self.defs[target][1]: + if dep in self.defs: + include(dep) + prefix = ", ".join( + f"{_quote(name)} AS ({query})" + for name, query in definitions.items() + ) + return (f"WITH {prefix} " if prefix else "") + self.defs[target][0] + + def run(self, stack=(), *, rebind=False): + owner = self.owner + aliases = {} + stages = 0 + for name in sorted(self.tables & owner._views.keys()): + if name in stack: + raise ValueError("cyclic virtual views cannot be staged") + graph = owner._view_graph(name) + rows = graph.run((*stack, name)) + stages += graph.stages + graph.compiled[""].store(owner, graph.names[""], rows) + aliases[name] = f"SELECT * FROM {_quote(graph.names[''])}" + for name in self.order: + compiled = self.compiled.get(name) + if rebind or compiled is None or not compiled.valid(owner): + frame = SessionContext.sql(owner.ctx, self._sql(name, aliases)) + fresh = _Compiled(owner, frame) + if compiled is not None: + fresh.batch = compiled.batch + compiled = self.compiled[name] = fresh + rows = compiled.run(owner) + stages += 1 + if name: + compiled.store(owner, self.names[name], rows) + compiled.epoch = owner._epoch + self.stages = stages + if self.reductions: + owner._counts["reduction_stage_executions"] += stages + else: + owner._counts["forward_stages"] = stages + return rows + + +class SharedExecutor: + """Cache shared SQL stages and ddx programs over live parameter tables. + + Install EinFold or MLX on ``ctx`` before constructing this executor. + Register model SQL with ``register_view``; it stays as SQL until ddx + has differentiated it. ``prepare`` returns a function producing an + Arrow table, and ``prepare_all`` produces a list of Arrow tables. + + ``inline_ctes`` names CTEs to keep inside consumers, for example an + attention score/softmax chain that EinFold can fuse as one operator. + Only reused, reachable CTEs outside that set become stage boundaries. + + Physical plans retain input buffers, with fresh join state each run. + Re-registering a public table invalidates those plans. Returned results + are independent snapshots. Changes to virtual views require preparing + forward queries again. Do not modify context settings while executing. + """ + + def __init__( + self, + ctx, + *, + inline_ctes: Sequence[str] = (), + max_key_cache_bytes: int = 256 << 20, + ): + if not getattr(_einfold, "RESET_ON_DECODE", False): + raise ImportError( + "rebuild accelerators/einfold_mlx for safe plan reuse" + ) + self.ctx = ctx + if max_key_cache_bytes < 0: + raise ValueError("max_key_cache_bytes must be nonnegative") + self._max_key_bytes = max_key_cache_bytes + self._key_bytes = 0 + self._keys: OrderedDict[tuple, list] = OrderedDict() + self.inline_ctes = frozenset(inline_ctes) + self._views: dict[str, str] = {} + self._graphs: dict[str, _Graph] = {} + self._counts: Counter[str] = Counter() + self._times: Counter[str] = Counter() + self._epoch = 0 + self._versions: dict[str, int] = {} + self._prefix = "__xql_shared_" + uuid.uuid4().hex + "_" + self._lock = threading.RLock() + self._catalog: tuple | None = None + for key in _DYNAMIC_FILTERS: + SessionContext.sql(ctx, f"SET {key} = false").collect() + self._backend = _Backend(self) + + def _intern_key(self, array): + size = array.nbytes + if ( + not ( + pa.types.is_integer(array.type) + or pa.types.is_string(array.type) + ) + or size > self._max_key_bytes + ): + return pa.concat_arrays([array]) + # The fingerprint is only an index: equality checks every value and + # NULL before sharing a buffer. Owned integer/string buffers are + # immutable; floating point values always get separate storage. + fingerprint = ( + str(array.type), + len(array), + array.null_count, + tuple( + array[i].as_py() for i in (0, len(array) // 2, len(array) - 1) + ) + if len(array) + else (), + ) + candidates = self._keys.get(fingerprint, []) + for old in candidates: + if array.equals(old): + self._keys.move_to_end(fingerprint) + self._counts["key_cache_hits"] += 1 + self._counts["key_copy_bytes_avoided"] += size + return old + owned = pa.concat_arrays([array]) + while self._keys and self._key_bytes + size > self._max_key_bytes: + _, evicted = self._keys.popitem(last=False) + self._key_bytes -= sum(a.nbytes for a in evicted) + # Eviction may have removed this fingerprint's bucket. + self._keys.setdefault(fingerprint, []).append(owned) + self._keys.move_to_end(fingerprint) + self._key_bytes += size + return owned + + def register_view(self, name: str, sql: str): + """Keep a deterministic SELECT as a virtual view for execution and ddx.""" + with self._lock: + _Graph(self, sql) # Validate before changing the view catalog. + if self._views.get(name) != sql: + self._views[name] = sql + self._graphs.clear() + self._epoch += 1 + + def _view_graph(self, name): + if name not in self._graphs: + self._graphs[name] = _Graph(self, self._views[name]) + return self._graphs[name] + + def _refresh(self): + catalog = tuple( + sorted( + (name, generation) + for name, generation in self.ctx._deregistrations.items() + if not name.startswith(("__ddx_", "__xql_shared_")) + ) + ) + if catalog != self._catalog: + self._catalog = catalog + self._epoch += 1 + + def prepare(self, sql: str): + """Prepare one SELECT; differentiation always uses its original SQL.""" + if calls_ad(sql): + run = self.prepare_all([sql]) + return lambda **tables: run(**tables)[0] + graph = _Graph(self, sql) + + def execute(**tables): + with self._lock: + self.ctx._write(tables) + self._refresh() + return _snapshot(graph.run()) + + return execute + + def prepare_all(self, statements: Sequence[str]): + """Prepare ddx statements sharing a backward pass over original views.""" + prepared = _prepare_ad(self._backend, statements) + finals: dict[str, _Compiled] = {} + + def execute(**tables): + with self._lock: + self.ctx._write(tables) + self._refresh() + with prepared.run() as rewritten: + results = [] + for sql in rewritten.statements: + compiled = finals.get(sql) + if compiled is None or not compiled.valid(self): + compiled = finals[sql] = _Compiled( + self, SessionContext.sql(self.ctx, sql) + ) + results.append(_snapshot(compiled.run(self))) + return results + + return execute + + def statistics(self): + """Executed stages, physical-plan builds and cumulative timings.""" + return ( + dict(self._counts) + | dict(self._times) + | {"key_cache_bytes": self._key_bytes} + ) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index f64166a4..bf801903 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -1,6 +1,8 @@ import re import time import uuid +from collections.abc import Callable +from types import ModuleType import numpy as np import pyarrow as pa @@ -9,7 +11,6 @@ from datafusion.catalog import Schema from datafusion.expr import TableScan from datafusion.plan import ExecutionPlan, LogicalPlan -from types import ModuleType from . import cftime as cft from .df import ( @@ -38,14 +39,30 @@ except ImportError: # pragma: no cover - depends on the environment _ad = None +_DYNAMIC_FILTERS = tuple( + f"datafusion.optimizer.enable{kind}_dynamic_filter_pushdown" + for kind in ("", "_join", "_topk", "_aggregate") +) + class Context(SessionContext): """A datafusion `SessionContext` that also supports `xarray.Dataset`s. Also available under its original name, ``XarrayContext``. + ``threads=`` sets DataFusion's target parallel partitions, overriding + that setting in an explicit ``config``. The default is the CPU count. + It is a concurrency target, rather than a fixed number of OS threads. """ - def __init__(self, config: SessionConfig | None = None, *args, **kwargs): + def __init__( + self, + config: SessionConfig | None = None, + *args, + threads: int | None = None, + **kwargs, + ): + if threads is not None and threads < 1: + raise ValueError("threads must be positive") if config is None: # Split a query's work across partitions only where enough rows # reach it, by DataFusion's estimate. DataFusion otherwise splits @@ -65,6 +82,8 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs): "true", ) ) + if threads is not None: + config = config.with_target_partitions(threads) super().__init__(config, *args, **kwargs) # Track registered xarray Datasets so XarrayDataFrame can recover # defaults (dimension_columns) and metadata (var/dataset attrs, @@ -78,6 +97,8 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs): # registered again from Python was deregistered first; SQL that # replaces or drops a table counts too (see sql()). self._deregistrations: dict[str, int] = {} + self._custom_udfs: set[str] = set() + self._trusted_math_udfs: set[str] = set() # With pyproj installed, every context speaks CRS out of the box: # reproject(x, y, src_crs, dst_crs), à la PostGIS ST_Transform. if _proj is not None: @@ -87,6 +108,41 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs): if _ad is not None: _ad.register(self) + def register_udf(self, udf): + """Register a scalar function and track overrides for optional lowering.""" + result = super().register_udf(udf) + self._custom_udfs.add(udf.name.lower()) + self._trusted_math_udfs.discard(udf.name.lower()) + return result + + def register_udaf(self, udaf): + """Track aggregate-function overrides for optional numerical lowering.""" + result = super().register_udaf(udaf) + self._custom_udfs.add(udaf.name.lower()) + self._trusted_math_udfs.discard(udaf.name.lower()) + return result + + def register_udwf(self, udwf): + """Track window-function overrides for optional numerical lowering.""" + result = super().register_udwf(udwf) + self._custom_udfs.add(udwf.name.lower()) + self._trusted_math_udfs.discard(udwf.name.lower()) + return result + + def with_extensions(self, *extensions) -> "Context": + """Install DataFusion extension bundles, retaining xarray support. + + Returns another handle on the same session, with this Context's + Dataset metadata and prepared-query invalidation state shared. + Extension codecs belong to the returned handle, so use it for + subsequent queries, as with DataFusion's own method. + """ + session = super().with_extensions(*extensions) + derived = object.__new__(type(self)) + derived.__dict__.update(self.__dict__) + derived.ctx = session.ctx + return derived + def from_dataset( self, name: str, @@ -214,9 +270,13 @@ def _from_dataset( table = read_xarray_table( input_table, chunks, coord_arrays=coord_arrays ) + register(table_name, table) + elif schema is None: + # Uniform Datasets can register their shared Arrow buffers + # directly, without a temporary catalog entry and lookup. + self.register_record_batches(table_name, partitions) else: - table = self._in_memory_table(partitions) - register(table_name, table) + register(table_name, self._in_memory_table(partitions)) self._maybe_register_cftime_udf(input_table) return self @@ -324,16 +384,11 @@ def prepare(self, query: str) -> "PreparedQuery": A [PreparedQuery][xarray_sql.PreparedQuery]; call it to run the query. - Raises: - NotImplementedError: If the query calls ``grad``, ``vjp`` or - ``jvp``, which only ``sql()`` computes for now. + Queries calling ``grad``, ``vjp`` or ``jvp`` retain ddx's programs + through [xarray_sql.ad.prepare][]. Their first call prepares the + programs; later calls check the catalog and settings and refresh + them if necessary. """ - if _ad is not None and _ad.calls_ad(query): - # TODO(#294): prepare the query's programs with ddx. - raise NotImplementedError( - "prepare() doesn't compute grad, vjp or jvp yet; run the " - "query with sql(), or several with xarray_sql.ad.sql_all()" - ) return PreparedQuery(self, query) def _write(self, writes: dict[str, xr.Dataset]) -> None: @@ -415,10 +470,19 @@ class PreparedQuery: query again. A call raises if a table the query reads has been registered again since it was prepared. Functions such as ``now()`` are evaluated once, when the query is prepared. + + A query calling ``grad``, ``vjp`` or ``jvp`` instead retains ddx's + programs. Its first call prepares them, and subsequent calls refresh + them when schemas, views or settings change. Tables may be registered + again under the same names, as with [xarray_sql.ad.prepare][]. """ def __init__(self, ctx: Context, query: str): self._ctx = ctx + self._gradient: Callable[..., list[pa.Table]] | None = None + if _ad is not None and _ad.calls_ad(query): + self._gradient = _ad.prepare(ctx, [query]) + return start = time.perf_counter() df = SessionContext.sql(ctx, query) plan = df.execution_plan() @@ -441,11 +505,13 @@ def __init__(self, ctx: Context, query: str): } # A plan keeps state from a run (a hash join keeps the table it # built), so each call runs a fresh copy of it, decoded from its - # serialized form. Copying a plan copies its in-memory tables, so - # over large ones it's slower than planning again; a plan that reads - # a lazy table, or is nested deeper than protobuf decodes, can't be - # copied at all. Either way, each call plans the query again. + # serialized form. The default codec copies in-memory tables, so + # over large ones it can be slower than planning again; extensions + # can supply a codec that resets state while retaining input buffers. + # A plan the session's codecs cannot copy is planned again each call. self._plan: ExecutionPlan | None = None + if not _safe_physical_reuse(ctx): + return start = time.perf_counter() try: _copy(ctx, plan) @@ -464,6 +530,9 @@ def __call__(self, **tables: xr.Dataset) -> pa.Table: Returns: The query's rows. """ + if self._gradient is not None: + (rows,) = self._gradient(**tables) + return rows for name in tables: if name not in self._registered: raise ValueError( @@ -495,7 +564,22 @@ def _copy(ctx: SessionContext, plan: ExecutionPlan) -> ExecutionPlan: # TODO(#283): reset the plan's state with datafusion-python's binding of # reset_plan_states, once it has one, rather than decode it again. # `plan` itself never runs, so it has no state to leave behind. - return ExecutionPlan.from_bytes(ctx, plan.to_bytes()) + return ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx)) + + +def _safe_physical_reuse(ctx: SessionContext) -> bool: + """Dynamic filters retain value-dependent bounds even after a reset.""" + names = ", ".join(repr(name) for name in _DYNAMIC_FILTERS) + try: + rows = SessionContext.sql( + ctx, + f"SELECT value FROM information_schema.df_settings WHERE name IN ({names})", + ).to_arrow_table() + except ValueError: + return False + return rows.num_rows == len(_DYNAMIC_FILTERS) and all( + value == "false" for value in rows["value"].to_pylist() + ) def _tables_read(plan: LogicalPlan) -> set[str]: diff --git a/xarray_sql/training.py b/xarray_sql/training.py new file mode 100644 index 00000000..7cb56c03 --- /dev/null +++ b/xarray_sql/training.py @@ -0,0 +1,702 @@ +"""Owned array state for training programs written in SQL. + +EinFold (https://github.com/xqlsystems/einfold) executes the numerical +plans and ddx (https://github.com/xqlsystems/ddx) supplies every gradient. +This opt-in session keeps complete, nonnullable rectangular tables as +arrays between calls. Use a dedicated Context; ordinary query APIs retain +their independent live-Dataset and Arrow snapshot contracts. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import pyarrow as pa +from datafusion import SessionContext +from ddxdb import prepare as _prepare +from einfold.dense_program import ( + _align, + _axes, + _Axis, + _Column, + _const, + _Decline, + _expr, + _is, + _read_dataset, + _Sorted, + _Table, + _Tensor, +) + +from xarray_sql.dense import ( + DenseExecutor, + _Compiled, + _DenseBackend, + _DenseGraph, +) +from xarray_sql.shared import _snapshot + +__all__ = ["TrainingResult", "TrainingSession"] + + +def _compact(table): + if isinstance(table, _Sorted): + sorts, table = table.sorts, table.table + if not isinstance(table, _Table): + raise _Decline("training requires a rectangular ordered result") + order = [] + for expr in sorts: + column = _expr(expr.expr(), table) + axis = column.tensor.axis + if axis is None or not _is(column.valid, True): + raise _Decline( + "training order requires nonnullable coordinate keys" + ) + if axis in order: + continue + indexes = np.argsort(axis.values, kind="stable") + if not expr.ascending(): + indexes = indexes[::-1] + if not np.array_equal(indexes, np.arange(len(axis))): + replacement = _Axis(axis.values[indexes]) + table = table.take(axis, indexes, (replacement,)) + axis = replacement + order.append(axis) + table = _Table( + _axes(order, table.axes), table.columns.items(), table.present + ) + if not isinstance(table, _Table): + raise _Decline("training requires a rectangular result") + return table + + +@dataclass(frozen=True) +class _DeviceTable: + table: _Table + schema: pa.Schema + + def select(self, names): + return _DeviceTable( + _Table( + self.table.axes, + [((None, n), self.table.resolve(None, n)) for n in names], + self.table.present, + ), + pa.schema([self.schema.field(n) for n in names]), + ) + + +class _TrainingBackend(_DenseBackend): + def returns_rows_compiled(self, plan, compiled): + table = _compact(compiled.run(self.owner)) + if any(len(axis) == 0 for axis in table.axes): + return False + flags = table.present.get() + runtime = self.owner._runtime + if runtime is not None and runtime.owns(flags): + if self.owner._tracing: + raise _Decline( + "compiled training needs static structural checks" + ) + return runtime.boolean(flags.any()) + return bool(np.any(flags)) + + +class TrainingResult: + """A completed update's loss and immutable device table snapshots. + + Full arrays reach Arrow only through ``export``. The loss may be an + MLX scalar; converting it to a Python number synchronizes that scalar. + """ + + def __init__(self, session, loss, outputs, version, measurements=None): + self._session, self._outputs = session, outputs + runtime = session.executor._runtime + self.loss = ( + runtime.copy(loss) + if runtime is not None + else np.array(loss, copy=True) + ) + self.version = version + self.measurements = measurements or {} + + def export(self, *names): + """Export independent Arrow snapshots of selected intermediate tables.""" + return self._session._export({n: self._outputs[n] for n in names}) + + def loss_value(self): + """Read the loss scalar, recording this explicit synchronization.""" + runtime = self._session.executor._runtime + if runtime is not None and runtime.owns(self.loss): + runtime.counts["gpu_loss_readbacks"] += 1 + self.measurements["gpu_loss_readbacks"] = ( + self.measurements.get("gpu_loss_readbacks", 0) + 1 + ) + return float(self.loss) + + +class TrainingSession: + """Own numerical tables and execute SQL updates without Arrow intermediates. + + ``bind`` imports a private copy. Later caller writes do not change it; + call ``bind`` again for explicit replacement. All persisted tables + must retain complete coordinate domains, non-NULL finite values and + their original schemas. An update replaces every assigned table + together after validation. A failed update leaves state unchanged. + """ + + def __init__( + self, ctx, *, precision="float32", max_cached_bytes=16 << 30, **options + ): + self.executor = DenseExecutor( + ctx, + precision=precision, + fallback=False, + max_cached_bytes=max_cached_bytes, + **options, + ) + self.executor._backend = _TrainingBackend(self.executor) + self.executor._tracing = False + self._tables = {} + self._closed = False + self.version = 0 + self._queries = {} + self._index_domains = {} + self._dropouts = () + self._random_step = 0 + + def _invoke(self, function): + if self._closed: + raise RuntimeError("training session is closed") + + def execute(): + self.executor._store.update( + {n: t.table for n, t in self._tables.items()} + ) + return function() + + return self.executor._invoke(execute, None, {}) + + def _seal(self, table, schema): + table = _compact(table) + if len(schema) != len(table.columns): + raise _Decline("training result schema changed") + columns = [] + for field, column in zip(schema, table.columns.values()): + tensor = column.tensor + value = tensor.get() + if tensor.axis is not None or tensor.host_only: + runtime = self.executor._runtime + if runtime is not None and runtime.owns(value): + raise _Decline( + "device-dependent columns cannot be frozen as static coordinates" + ) + frozen = _Tensor( + "array", + tensor.axes, + value=np.asarray(value).copy(), + axis=tensor.axis, + ) + else: + frozen = _Tensor("device", tensor.axes, value=value) + frozen.domain = tensor.domain + columns.append(((None, field.name), _Column(frozen, column.valid))) + return _DeviceTable(_Table(table.axes, columns, table.present), schema) + + def _values_and_flags(self, tables): + values, flags = [], [] + for item in tables: + table = item.table + flags.append(table.present.get()) + for col in table.columns.values(): + flags.append(col.valid.get()) + if col.tensor.axis is None: + values.append(col.tensor.get()) + return values, flags + + def _validate(self, tables): + runtime = self.executor._runtime + values, flags = self._values_and_flags(tables) + if runtime is not None: + runtime.validate(values, flags) + elif not all(np.all(f) for f in flags) or not all( + np.isfinite(v).all() for v in values + ): + raise ValueError( + "training update must be rectangular, nonnullable and finite" + ) + + def bind(self, name, dataset=None): + """Import or explicitly replace a table, independent of caller buffers.""" + with self.executor._lock: + return self._bind(name, dataset) + + def _bind(self, name, dataset): + ctx = self.executor.ctx + if dataset is None: + dataset = ctx._registered_datasets[name] + private = dataset.copy(deep=True) + if any(size == 0 for size in private.sizes.values()): + raise ValueError( + "training tables require nonempty coordinate domains" + ) + if any( + tuple(v.dims) != tuple(private.sizes) + for v in private.data_vars.values() + ): + raise ValueError( + "owned training columns must share the complete table dimensions" + ) + if name not in ctx._registered_datasets: + ctx.deregister_table(name) + ctx.from_dataset(name, private) + schema = ctx.table(name).schema() + source = _read_dataset(name, private) + if schema.names != [n for _, n in source.columns] or any( + f.type != pa.from_numpy_dtype(c.tensor.value.dtype) + for f, c in zip(schema, source.columns.values()) + ): + raise ValueError("replacement changes the training table schema") + + def execute(): + item = self._seal(source, schema) + for column, (reference, dim) in self._index_domains.get( + name, {} + ).items(): + domain = ( + self._tables[reference].table.resolve(None, dim).tensor.axis + ) + col = item.table.resolve(None, column) + values = col.tensor.get() + if ( + domain is None + or values.dtype.kind not in "iu" + or not np.isin(values, domain.values).all() + ): + raise ValueError( + "index values are outside the declared coordinate domain" + ) + tensor = _Tensor( + "device", + col.tensor.axes, + value=self.executor._runtime.indices(values), + ) + tensor.domain = domain + item.table.columns[None, column] = _Column(tensor, col.valid) + self._validate([item]) + return item + + item = self._invoke(execute) + old = self._tables.get(name) + if old is not None and ( + not old.schema.equals(item.schema) + or len(old.table.axes) != len(item.table.axes) + or any( + not np.array_equal(a.values, b.values) + for a, b in zip(old.table.axes, item.table.axes) + ) + ): + raise ValueError( + "replacement changes the training table schema or coordinate domains" + ) + self._tables[name] = item + + def index_domains(self, name, columns): + """Declare GPU integer input columns bounded by owned coordinate axes. + + The next ``bind`` validates the actual integer values and uploads + them as dynamic device inputs. ``columns`` maps each field to a + ``(reference_table, coordinate)`` pair. + """ + if self.executor._runtime is None: + raise ValueError( + "device index domains require float32 GPU execution" + ) + self._index_domains[name] = dict(columns) + self.executor._epoch += 1 + + def dropout(self, names, *, probability, seed, step=0): + """Generate owned mask tables on GPU using an explicit seed and step. + + Each named table must have one numerical field. The counter advances + only after a successful update. Save seed and next counter with the + checkpoint; identical inputs reproduce masks after resume. + """ + if self.executor._runtime is None or not 0 <= probability < 1: + raise ValueError( + "GPU dropout requires float32 execution and probability in [0,1)" + ) + if not 0 <= seed < 2**32 or not 0 <= step < 2**32: + raise ValueError("dropout seed and step must be uint32 values") + for name in names: + columns = [ + c + for c in self._tables[name].table.columns.values() + if c.tensor.axis is None + ] + if len(columns) != 1: + raise ValueError( + "each dropout table needs exactly one numerical field" + ) + self._dropouts = tuple(sorted(names)) + self._dropout_probability, self._random_seed, self._random_step = ( + probability, + seed, + step, + ) + self.executor._epoch += 1 + + def _make_masks(self, key): + templates = [self._tables[name] for name in self._dropouts] + cols = [ + next(c for c in t.table.columns.values() if c.tensor.axis is None) + for t in templates + ] + arrays = self.executor._runtime.dropout( + key, + [tuple(len(a) for a in c.tensor.axes) for c in cols], + self._dropout_probability, + ) + results = {} + for name, item, col, array in zip( + self._dropouts, templates, cols, arrays + ): + columns = [ + ( + k, + _Column( + _Tensor("device", c.tensor.axes, value=array), + _const(True), + ) + if c is col + else c, + ) + for k, c in item.table.columns.items() + ] + results[name] = _DeviceTable( + _Table(item.table.axes, columns), item.schema + ) + return results + + def prepare_step( + self, + *, + gradients, + stages, + updates, + assignments, + loss_column="l", + compile=False, + ): + """Prepare one shared ddx backward followed by ordinary SQL statements. + + ``gradients``, ``stages`` and ``updates`` map result names to SQL. + ``assignments`` maps each owned state table to ``(result, columns)``. + All update SELECTs read the same old state. The primal loss is read + from ddx's saved output, without an independent forward query. + """ + step = _TrainingStep( + self, gradients, stages, updates, assignments, loss_column + ) + step.compile = compile + return step + + def query(self, sql): + """Evaluate SQL against current owned state and export its result.""" + if sql not in self._queries: + self._queries[sql] = _DenseGraph(self.executor, sql) + graph = self._queries[sql] + + def execute(): + table = graph.run() + return _snapshot(table.arrow(graph.compiled[""].schema)) + + return self._invoke(execute) + + def _export(self, tables): + def execute(): + self.executor._counts["training_exports"] += len(tables) + return { + n: _snapshot(t.table.arrow(t.schema)) for n, t in tables.items() + } + + return self._invoke(execute) + + def export(self, *names): + """Export independent Arrow snapshots of owned tables.""" + return self._export( + {n: self._tables[n] for n in (names or self._tables)} + ) + + def close(self): + """Release this session's owned arrays.""" + with self.executor._lock: + self._tables.clear() + self._closed = True + + def statistics(self): + return self.executor.statistics() | {"training_version": self.version} + + +class _TrainingStep: + def __init__( + self, session, gradients, stages, updates, assignments, loss_column + ): + self.session, self.assignments, self.loss_column = ( + session, + dict(assignments), + loss_column, + ) + self.gradients, self.stages, self.updates = ( + dict(gradients), + dict(stages), + dict(updates), + ) + self.prepared = _prepare( + session.executor._backend, list(gradients.values()) + ) + self.compiled = {} + + def _run(self, sql): + owner = self.session.executor + if sql not in self.compiled or self.compiled[sql].epoch != owner._epoch: + self.compiled[sql] = _Compiled( + owner, SessionContext.sql(owner.ctx, sql) + ) + compiled = self.compiled[sql] + return self.session._seal(compiled.run(owner), compiled.schema) + + def _store(self, name, item): + owner = self.session.executor + owner._store[name] = item.table + owner._stub(name, item.schema) + + def _build(self): + session = self.session + results = {} + with self.prepared.run() as rewritten: + if not hasattr(self.prepared, "values"): + raise ImportError( + "rebuild accelerators/ddx_math for prepared primal outputs" + ) + values = [ + v + for v in self.prepared.values + if tuple(v.columns) == (self.loss_column,) + ] + if len(values) != 1: + raise ValueError( + "training requires one unambiguous scalar loss output" + ) + loss_table = self._run(values[0].query()) + for name, sql in zip(self.gradients, rewritten.statements): + results[name] = self._run(sql) + self._store(name, results[name]) + for name, sql in self.stages.items(): + results[name] = self._run(sql) + self._store(name, results[name]) + for name, sql in self.updates.items(): + results[name] = self._run(sql) + next_state = { + name: results[result].select(columns) + for name, (result, columns) in self.assignments.items() + } + for name, item in next_state.items(): + old = session._tables[name] + if old.schema.names != item.schema.names or any( + a.type != b.type for a, b in zip(old.schema, item.schema) + ): + raise ValueError("update changes a state table schema") + if len(old.table.axes) != len(item.table.axes) or any( + not np.array_equal(a.values, b.values) + for a, b in zip(old.table.axes, item.table.axes) + ): + raise ValueError("update changes a state coordinate domain") + shape = tuple(len(a) for a in item.table.axes) + columns = [] + for key, col in item.table.columns.items(): + if col.tensor.axis is None: + value = col.tensor.get() + runtime = session.executor._runtime + value = ( + runtime.broadcast( + value, col.tensor.axes, item.table.axes, shape + ) + if runtime is not None + else np.broadcast_to( + _align(value, col.tensor.axes, item.table.axes), + shape, + ) + ) + col = _Column( + _Tensor("device", item.table.axes, value=value), + col.valid, + ) + columns.append((key, col)) + next_state[name] = _DeviceTable( + _Table(item.table.axes, columns, item.table.present), + item.schema, + ) + return loss_table, results, next_state + + def _pack(self, tables): + arrays, layouts = [], {} + for name, item in tables.items(): + columns = [] + for key, col in item.table.columns.items(): + t = col.tensor + if t.axis is None and not t.host_only: + index = len(arrays) + arrays.append(t.get()) + columns.append((key, t.axes, t.domain, index, None)) + else: + columns.append((key, t.axes, t.domain, None, col)) + layouts[name] = (item.table.axes, item.schema, columns) + return arrays, layouts + + @staticmethod + def _unpack(arrays, layouts): + tables = {} + for name, (axes, schema, columns) in layouts.items(): + values = [] + for key, local, domain, index, static in columns: + if index is None: + value = static + else: + tensor = _Tensor("device", local, value=arrays[index]) + tensor.domain = domain + value = _Column(tensor, _const(True)) + values.append((key, value)) + tables[name] = _DeviceTable(_Table(axes, values), schema) + return tables + + def _compiled_build(self): + session, owner = self.session, self.session.executor + runtime = owner._runtime + if runtime is None: + raise ValueError("compiled training requires float32 MLX execution") + sources = { + n: t + for n, t in sorted(session._tables.items()) + if n not in session._dropouts + } + for item in sources.values(): + if any( + c.tensor.axis is None and c.tensor.host_only + for c in item.table.columns.values() + ): + raise ValueError( + "compiled integer inputs require explicit index_domains" + ) + arrays, layouts = self._pack(sources) + if session._dropouts: + arrays.append( + runtime.random_key(session._random_seed, session._random_step) + ) + guard = (owner._epoch, owner._backend.session_key()) + if getattr(self, "_compiled_guard", None) != guard: + + def function(*args): + bound = self._unpack( + args[:-1] if session._dropouts else args, layouts + ) + masks = ( + session._make_masks(args[-1]) if session._dropouts else {} + ) + owner._store.clear() + owner._store.update({n: t.table for n, t in bound.items()}) + owner._store.update({n: t.table for n, t in masks.items()}) + owner._tracing = True + runtime.tracing = True + try: + loss, results, next_state = self._build() + results.update(masks) + packed, self._output_layouts = self._pack( + { + ("loss",): loss, + **{("result", n): t for n, t in results.items()}, + **{("state", n): t for n, t in next_state.items()}, + } + ) + values, flags = session._values_and_flags( + [loss, *results.values()] + ) + status = runtime.status(values, flags) + return packed, status + finally: + owner._tracing = False + runtime.tracing = False + + self._function = runtime.compile_training(function) + packed, status = self._function(*arrays) + runtime.finish_training(packed, status) + restored = self._unpack(packed, self._output_layouts) + loss = restored[("loss",)] + results = {k[1]: v for k, v in restored.items() if k[0] == "result"} + next_state = {k[1]: v for k, v in restored.items() if k[0] == "state"} + owner._refresh() + self._compiled_guard = (owner._epoch, owner._backend.session_key()) + return loss, results, next_state + + def __call__(self): + session = self.session + + def execute(): + counters = session.executor.statistics() + if session._dropouts and session._random_step >= 2**32: + raise ValueError("GPU dropout counter is exhausted") + if getattr(self, "compile", False): + loss_table, results, next_state = self._compiled_build() + else: + masks = ( + session._make_masks( + session.executor._runtime.random_key( + session._random_seed, session._random_step + ) + ) + if session._dropouts + else {} + ) + session.executor._store.update( + {n: t.table for n, t in masks.items()} + ) + loss_table, results, next_state = self._build() + results.update(masks) + session._validate([loss_table, *results.values()]) + # Validation proved these full domains have no NULLs. + next_state = { + n: _DeviceTable( + _Table( + t.table.axes, + [ + (k, _Column(c.tensor, _const(True))) + for k, c in t.table.columns.items() + ], + ), + t.schema, + ) + for n, t in next_state.items() + } + loss = loss_table.table.resolve(None, self.loss_column).tensor.get() + session._tables.update(next_state) + session.version += 1 + session._random_step += 1 + session.executor._counts["training_steps"] += 1 + measured = session.executor.statistics() + delta = { + name: measured.get(name, 0) - counters.get(name, 0) + for name in ( + "gpu_readback_bytes", + "gpu_input_bytes", + "gpu_status_readbacks", + "training_exports", + "gpu_training_compilations", + ) + } + return TrainingResult( + session, loss, results, session.version, delta + ) + + return session._invoke(execute)