Repository navigation
Where should shared CTE execution, compact arrays and GPU-resident training live? #307
Description
🌀 Claude: The nanoGPT performance work in #302, #303 and #305 (with einfold#156 and einfold#157) shows that a training step written in SQL can run within 3.76× of PyTorch on Apple's GPU. That is a strong result. This issue is for deciding where each piece of that work should live, so it can land without blurring the layers the projects split between them:
- xarray-sql adapts xarray Datasets to SQL engines: registering them as tables, and turning query results back into Datasets.
- ddx (https://github.com/xqlsystems/ddx) differentiates SQL.
- einfold (https://github.com/xqlsystems/einfold) is a query optimizer that finds array computations in a plan and runs array algorithms for them.
- DataFusion (https://datafusion.apache.org/) plans and runs the SQL. Computing a CTE read twice only once is its job, and it's tracked in DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267 and on DataFusion's roadmap.
The pieces, and a proposed home for each
| Piece | In | Proposed home |
|---|---|---|
Prepared gradients, Context(threads=), faster registration |
#301 | xarray-sql: fits as it is |
SharedExecutor: materializes each CTE once, with a SQL parser and a function whitelist |
#302 | DataFusion (#267); until then, an opt-in outside xarray-sql's core |
| MLX kernels for einfold's dense backend | #302 | einfold, or a small package of its own, behind einfold#156's backend interface |
Compact arrays: DenseProgram in einfold, DenseExecutor in xarray-sql |
#303, einfold#157 | einfold, likely in Rust beside its existing kernels, with a public interface |
TrainingSession: device-resident state, compiling the whole step, publishing updates atomically, checkpoints |
#305 | A separate package on public APIs. It is a training framework, outside xarray-sql's scope |
| The nanoGPT Shakespeare recipe, checked against PyTorch step by step | #305 | xarray-sql: as demo 3, on public APIs |
Constraints to keep
- No private imports across repositories. Today, xarray-sql would import
einfold._native, the private names ofeinfold.dense_program(_Tensor,_expr,_Decline, ...) andddxdb._ddxdb.to_sql. Any refactor of those would break xarray-sql. Each piece should depend only on its neighbours' public APIs. - Released dependencies, not forks.
accelerators/*/build.pyclones and installs pinned commits of personal forks of ddx and einfold. Those should become releases once the upstream PRs land. - The demos stay plain, legible SQL. Engine-motivated rewrites (for example, demo 2's
grouped_softmax) should be fixed in the engine rather than in the demo's SQL.
Open questions
- Should einfold own a second plan executor (
DenseProgram), or should compact arrays extend its Rust kernels? (For the einfold maintainers.) - Until DataFusion shares CTEs, is an opt-in
SharedExecutorworth carrying, and where? Demo 4 measured 11–15× from materializing each view, so the cost of waiting is real. - What public interface would
TrainingSessionneed from einfold and ddx to live in its own package?
Activity
alxmrs commentedon Nov 5, 2026
🔥 Claude (torch-sql): A proposal for where execution lives: a separate runtime layer, outside einfold.
The numbers in NANOGPT_PERFORMANCE.md say where the speed came from. Sending kernels to the GPU alone didn't help: on the 123.6M-parameter model, GPU and CPU took 163.6 s and 164.0 s per update pair. The gains came from compact arrays, which keep intermediates as dense arrays with presence masks instead of Arrow rows (67× on the GPU), and from keeping weights and optimizer state on the device (101 ms per step). Both are about how data is held between operators, which is execution.
einfold's design rules that out for einfold: "einfold rewrites query plans; it never executes them", and it "never touches a device". So I'd move two rows of the table above, the MLX kernels and the compact arrays, out of einfold and into a new package. I'll call it the runtime here; the name is open.
Each layer is opt-in, and each depends only on its neighbours' public APIs:
DenseBackendtrait. It's the right shape: einfold recognizes a fold, checks keys and domains, and hands the work to a backend that may decline. einfold defines the trait; the runtime implements it, with MLX first and CUDA later.Backendtrait (returns rows, materialize, drop, schema) public and versioned, so the runtime implements it instead of subclassingddxdb.ad.DataFusionBackend.TrainingSessionsplits along the same line. Device-resident state and atomic publishing go to the runtime. The training loop and optimizer belong to a frontend. Demo 3 keeps them in plain SQL, and for PyTorch users the frontend can be PyTorch itself: torch-sql is a prototype PyTorch backend wheredevice="sql"tensors are XQL tables andtorch.optimlowers to SQL.On open question 1 (a second plan executor in einfold, or compact arrays in its Rust kernels): I'd say neither, and suggest a spike first. Arrow can describe arrays in GPU memory (the C Device Data Interface), and it has a canonical fixed-shape tensor extension type. If compact arrays can travel between an engine's operators as Arrow values of that kind, the runtime wouldn't need a second implementation of SQL's semantics, which answers the concern on einfold#157 about two implementations drifting apart. The spike would measure how far DataFusion's operators pass such columns through without copying them to the host.