Demo 3 (GraphCast in SQL): what would make it tractable, and what the design is missing #36

Activity

1. alxmrs

   added

   documentationImprovements or additions to documentationImprovements or additions to documentation

   enhancementNew feature or requestNew feature or request

   on Oct 18, 2026
2. alxmrs commented on Oct 18, 2026

   

   alxmrs

   on Oct 18, 2026

   MemberAuthor

   More actions

   🌀 Claude (GraphCast demo, xarray-sql)

   @alxmrs asked whether these suggestions call for a new design, or whether v16 can absorb them. My answer: v16 can absorb them. Almost every suggestion fits an abstraction einfold already has. What GraphCast changes is mostly the roadmap's order, plus one real generalization of a rule.

   Already in the design: reorder or extend

   Suggestion	Where it already fits
   Narrow detection of x ⋈ W GROUP BY id, g	M2's precedent: one detected pattern, ahead of M3's general detection
   Density checked at run time on the weights	Measured facts (8.1) and S25's existence check (8.3)
   DataFusion recomputes CTEs	A correction to the DataFusion target profile (7.4, 9.3); profiles are data for exactly this reason
   Factor hoisting	A detail of how 9.1's "one expression over one operand is one factor" is realized
   Distributing a layer over the gather	Contraction order (9.3), eager aggregation (10.1), sums over UNION ALL (9.1, #17); the f + 512 offset is an affine coordinate map (8.1)
   Caching edge embeddings	8.5's structure/value split, extended from cached plans to cached results
   Memory budgets for relational-form plans	M7's slicing, moved earlier
   Epilogues (bias, activation, layer norm)	A 9.6 pipeline whose consumer sums nothing away; layer norm is an algebraic fold (sum, sum of squares, count)

   The one real generalization

   Fusing a message MLP with its sum by receiver. Section 9.6 requires the consumer's group to be the producer's group minus a dimension the consumer sums away. In GNN message passing, the consumer's group (receiver) is a function of the summed dimension (edge), through a join with the edge table. Allowing a functional dependency through a mapping relation changes one rule's precondition, not the architecture. Its test cases would be GNN message passing and the gather/scatter half of sparse attention, next to Nadaraya–Watson.

   Where it would become a redesign

   If the goal is to approach JAX (DeepMind's code runs the forecast in 15 s), not to beat the hosts by 4–14×. Most of the remaining gap is that every fold's output becomes rows before the next fold reads it, about ten times per layer. S21 already found that writing output rows bounds the win. Closing that gap means keeping a whole layer's intermediates as dense tiles across many folds. That's still einfold's frame (folds, tiles, pipelines, program mode), but it moves "pipelines of more than two folds" (section 14, M7) from deferred to central.

   Proposed v17

   1. Bring M1's dense kernels to plain SQL with a narrow forward-pass detector. Today's order is set by ddx's gradients, which emit folds explicitly; GraphCast's forward pass doesn't.
   2. Correct the DataFusion target profile: plain CTEs run once per reference.
   3. Generalize 9.6's fusion condition to functional dependencies through a mapping relation, and add the MLP (linear_0 → swish → linear_1) and GNN message passing as test cases.
   4. Bring relational-form slicing forward from M7.

   This rests on one demo that hasn't run end to end, so before rewriting the roadmap I'd run the two spikes below. Either result could firm up or overturn items 1 and 3.

   ---

   Spikes for another agent

   Both spikes use GraphCast_small's real weights and graphs. The code that fetches and tables them is in xarray-sql's benchmarks/ml/_graphcast.py, on branch graphcast-weights (xqlsystems/xarray-sql#270). xarray-sql is XQL Systems' library for querying Xarray datasets with SQL. Install it from that branch with its ml and duckdb extras, plus weathernext @ git+https://github.com/google-deepmind/graphcast, DeepMind's GraphCast package. Then use:

   • _graphcast.checkpoint(), graphs(), param_tables(), channels_table() and graph_tables(): the published weights, the three graphs, and their tables, downloaded once to ~/.cache/xarray-sql/graphcast;
   • _graphcast.mlp(net, x): one of DeepMind's MLPs in plain JAX, float64. This is the reference. Its stages match DeepMind's own model to about 1e-6 of each stage's largest value;
   • 03_graphcast.py's mlp(), messages() and grid_inputs(): the SQL as the demo writes it.

   Memory safety, required. This machine has 15 GB, and an uncapped run of the demo's encoder was OOM-killed at 13.9 GB, taking the agent session with it. Run every measurement in its own process inside a memory cap:

   systemd-run --user --scope -p MemoryMax=10G -p MemorySwapMax=0 python ...
   

   Set each host's memory limit explicitly, and report it: DataFusion's RuntimeEnvBuilder().with_fair_spill_pool(...), DuckDB's memory_limit. At a 0.5 GB limit, DataFusion spilled more than 11 GB in 20 minutes on a slice that takes 30 s unlimited, so watch spill size and free disk.

   Follow docs/spikes/README.md: one folder per spike, a report with method, results, findings and limitations, and SPDX headers.

   S28: dense kernels on GraphCast's forward MLPs

   Question. Do S21's dense kernels give their 4–14× on GraphCast's MLPs, from Arrow rows to Arrow rows? Does a dense output remove the hosts' memory problem? And how much more does fusing the MLP's epilogues (bias, swish, layer norm) add?

   Workloads, from the real weights and inputs:

   1. The grid embedding, grid2mesh_gnn/encoder_nodes_grid_nodes: 65,160 cells × [186 → 512 → 512], then layer norm, about 23 G joined pairs. Its input is the demo's grid_in table, which grid_inputs() builds.
   2. A processor message MLP, mesh_gnn/processor_edges_0_mesh: 81,900 edges × [1536 → 512 → 512], then layer norm, about 86 G pairs. Its input is each edge's state, sender and receiver concatenated. Build it from _graphcast.mirror()'s mesh state after the encoder and the embedded mesh edges, _graphcast.mlp("mesh_gnn/encoder_edges_mesh", graphs()["mesh"]["features"]). If either host can't finish within the cap, use a fixed slice of edges, the same on every method.

   Methods.

   • DataFusion and DuckDB running the demo's SQL mlp(): each matrix product alone, and the whole MLP.
   • S21's fold and GEMM kernels: each matrix product alone, with Arrow rows in and out.
   • The whole MLP in one kernel: both products, with bias, swish and layer norm applied while the output is written, so the hidden and pre-norm rows never exist. This measures epilogue and MLP fusion (the 9.6 pipeline) apart from the kernels themselves.

   Report. For each workload and method: time, peak memory, joined pairs per second, and error against _graphcast.mlp in float64.

   Prediction. DataFusion runs about 0.09 G pairs/s, about 260 s for workload 1. Its hash aggregates hold ids × 512 groups, which is where memory goes. S21 predicts 4–14× for the kernels. Fusion's further gain is unmeasured.

   S29: distributing the first layer over the gather

   Question. Does writing a message MLP's first layer as W_e·e + (W_s·s)[sender] + (W_r·r)[receiver] save in plain SQL, on both hosts, what the pair counts predict? The demo writes W · [e; s[sender]; r[receiver]] instead: concatenation by UNION ALL with offset feature indices, then one join with a 1536 × 512 matrix. And which of einfold's rules would each step of the rewrite need?

   Method. One processor round on the real mesh: 10,242 nodes, 81,900 edges. Take the inputs from _graphcast.mirror() and the edge embedding as in S28, so the SQL encoder never runs, and register them as tables.

   • (a) The demo's messages() SQL, as written.
   • (b) The rewrite, also as SQL:
     • each node's sender part, mesh ⋈ W[512:1024], and receiver part, mesh ⋈ W[1024:1536], each 10,242 × 512;
     • each edge's own part, edge_state ⋈ W[0:512];
     • the three parts summed per edge, through joins with the edge table;
     • then the bias, swish, the second layer and the layer norm, unchanged.

   Check that (a) and (b) agree, within the float sum bound of section 8.6, with _graphcast.mlp("mesh_gnn/processor_edges_0_mesh", …).

   Report. Time, peak memory and joined pairs on DataFusion and DuckDB. Then, for each step of the rewrite, say which existing rule covers it and what's missing. The steps are: the offset f + 512 as an affine reindexing of the contracted dimension; splitting W by ranges of that dimension; the sum over UNION ALL; and contracting before a many-to-one join.

   Prediction. The first layer goes from 64.4 G pairs to 26.8 G per round, and the whole message MLP from 85.9 G to 48.3 G (−44%). Across the forecast, the rewrite saves 717 G of 2,054 G pairs (35%).
3. alxmrs commented on Oct 18, 2026

   

   alxmrs

   on Oct 18, 2026

   MemberAuthor

   More actions

   🌀 Claude (GraphCast demo, xarray-sql)

   @alxmrs asked for notes on how einfold might approach JAX's time on GraphCast. Here they are. The numbers are estimates from the measurements above, not results.

   The target

   DeepMind's JAX code runs one GraphCast_small forecast in 15 s on this machine, an Intel i7-8700 with 6 cores. JAX is Google's array library, and XLA is its compiler. That forecast uses float32 and the same 2.05 trillion multiply-adds the SQL forms as joined pairs, so it ran at about 137 G multiply-adds per second.

   Two caveats. First, the 15 s was the first call of a jax.jit function, so it includes XLA's compilation, and the steady-state forecast is faster. A fair target needs a warm, repeated timing. Second, JAX's run is the floor this machine has demonstrated, not its peak. Six AVX2 cores can do several hundred GFLOP/s in float32.

   DataFusion runs the demo's SQL at about 0.09 G joined pairs per second: roughly 1,500× slower. The gap has four parts, and each needs a different fix.

   1. Arithmetic: every product as a GEMM on dense float32 tiles

   All 2.05 T pairs are two-operand products with a dense weight matrix of at most 1536 × 512 (see above). As GEMMs, these shapes are compute-bound, with K from 186 to 1536 and N = 512. That's the regime where a blocked GEMM kernel reaches a large fraction of peak, unlike S21's ddx shapes (K = 16), where writing the output dominated.

   • What it needs: M1's GEMM path, reached from plain SQL (the narrow detector in the first comment), with the weights as its dense operand.
   • Precision: float32. JAX's 15 s is float32, and float64 halves a GEMM's throughput. That makes float32 (M4) a requirement for this goal, not a later option. bfloat16, which DeepMind runs on accelerators, is a further 2× where a host supports it.
   • Weights: packed once into the kernel's layout and cached. They are the same for every forecast (section 8.5's structure/value split, applied to values that never change).

   2. Data movement: rows only at the edges of the program

   This is the largest gap, and the one the design defers. Each stage of the demo writes its output as (id, f, val) rows: two 8-byte keys per 8-byte value, where a dense float32 array spends 4 bytes per value. Each MLP writes its hidden and pre-norm layers as rows too. Per processor round, the message stage alone writes about 168M values (81,900 edges × 512 channels, materialized about four times: two hash aggregates, the layer norm's window, and the stored result). Over 16 rounds, that's about 2.7 G values, or 64 GB of row traffic at 24 bytes each. The hosts then hash it all again in the next stage's join and GROUP BY. S21 found that writing output rows bounds the win, and here it would dominate even after (1).

   JAX writes activations as dense arrays, and XLA fuses the elementwise work into the matrix products' loops. To get close, einfold needs dense residency at three scopes, from narrow to wide:

   • Epilogues. Bias, activation and layer norm run as each output tile is written (the first comment's fold epilogues).
   • The layer. linear_0 → swish → linear_1 → layer norm is one 9.6 pipeline, so the hidden layer lives in a tile, as FlashAttention keeps its scores.
   • The program. Node and edge states stay dense between statements. In program mode (7.3), a stored stage could be a table backed by an einfold dense buffer, turned into rows only if an operator outside einfold reads it. Rows would then cross into the host only at the program's boundaries: the ERA5 inputs and weights in, the forecast out. The logical model is unchanged, since every table still has its rows (principle 2); only their physical layout differs. That fits the design's statement that "layouts, tiles, and partial states are physical, and they live inside operators", extended from operators to a program's stored results.

   3. Graph joins as gathers and scatters, not hash joins

   Besides the products with weights, the joins in a message-passing round aren't contractions. They join the edge table to node states by sender or receiver (a gather), and sum messages by receiver (a scatter-add, JAX's segment_sum). As hash joins, they rehash tens of millions of rows per round. When facts say that sender and receiver are positions into the node tables (the edge table is fixed, and its keys are dense ranges), they become indexed loads and adds.

   • Gather: a join whose key is a position, with no aggregate after it. That's not a fold, so EinFold doesn't cover it today. The design has positions (8.1, coordinate maps and dictionary encoding), but no realization for a positional join without an aggregate.
   • Scatter-add: a fold whose group key is a position, the first comment's message → sum-by-receiver pipeline. It needs 9.6's condition generalized to a functional dependency through the edge table.

   4. Do less arithmetic than JAX

   einfold's algebra can do less work than JAX's program, because DeepMind's code concatenates before each edge MLP, and XLA doesn't distribute the product over the concatenation:

   • Gather distributivity (S29): 2.05 T → 1.34 T multiply-adds (−35%).
   • Caching the edge embeddings, which are constant across forecasts: −0.10 T.
   • Exact regional pruning for a region (−27% of the processor, −97% of the decoder, for the contiguous US).
   • For a rollout, the static parts are cached once for all steps.

   Even at half JAX's GEMM efficiency, a 35–50% smaller plan would bring einfold within about 2× of JAX. That makes "within 2× of JAX, and below it for regional forecasts" a realistic stretch goal, and a strong research result: a declarative model, optimized by algebra JAX's program can't express, against a hand-built array program.

   An order of work toward it

   Step	What	Estimated forecast time on this machine
   Today	Host SQL, one stored table per stage	~6 h
   M1 kernels via narrow detection (S28)	Products as GEMMs; still rows between stages	tens of minutes (S21's 4–14×)
   + epilogues, MLP pipelines	Hidden and pre-norm rows gone	lower, by the share of row traffic they carry (S28's fused variant measures it)
   + positional gather and scatter-add	Graph joins without hashing	lower again
   + program-level dense residency, float32	Rows only at the program's edges	the regime where JAX lives: seconds to a minute
   + distributivity, caching, regional pruning	Less arithmetic than JAX	toward or below JAX's time

   The first two rows are near-term and measurable now. The last three are the redesign-shaped part from the previous comment: dense residency across folds and statements. That puts M7's "pipelines of more than two folds" and program mode at the center of the roadmap, not the end.

   A spike that would bound all of this: S30, the roofline

   Before committing to that order, measure the ceiling on this machine:

   • a warm, repeated timing of DeepMind's forecast in JAX, with compilation excluded, plus a per-stage profile to see where JAX spends its time;
   • float32 and float64 GEMMs from a BLAS library, and S21's kernels, at GraphCast's shapes (65,160 × 186 × 512, 81,900 × 1536 × 512, 195,480 × 1536 × 512);
   • a gather and a scatter-add over the real edge tables, with 512 channels.

   Those three numbers tell us how close a perfect einfold could get, and which of the four gaps is worth closing first. Same rules as S28 and S29: run each measurement inside a memory cap (the machine has 15 GB), in its own process.
4. alxmrs commented on Oct 31, 2026

   

   alxmrs

   on Oct 31, 2026

   Last edited by alxmrs

   MemberAuthor

   More actions

   🌀 Claude: GraphCast on einfold 0.1.0rc2. Measured on an Intel i7-8700 with DataFusion 55.0.0 and ddxdb 0.4.6, and checked against DeepMind's model.

   It runs, and it's correct. The demo's earlier form stored each stage as a table. In that form, the encoder and one processor round ran over the full 1° grid in about 10 minutes, and each checked stage matched DeepMind's model to 1e-6 or better:

   Stage	Time
   Grid embedding (65,160 cells × [186 → 512 → 512])	59 s
   Mesh embedding	10 s
   Grid→mesh edge embedding	126 s
   Grid→mesh messages	118 s
   Mesh and grid updates	13 s + 67 s
   Mesh edge embedding	80 s
   One processor round (messages, update, edge residual)	98 s + 11 s + 2 s

   Summing these puts a full forecast at about 45–50 minutes. DeepMind's JAX takes 15 s. Peak memory reached the run's 12 GB limit, with 3.9 GB in swap.

   einfold runs the matrix products; the layer norm is now the cost. On a 2,048-cell slice of the grid embedding, the whole plan is one einfold operator at every stage:

   Up to	einfold	DataFusion alone
   first layer (SUM(x.val * w.w) … GROUP BY x.id, w.g)	0.14 s
   bias and swish	0.17 s
   second layer, plus bias	0.32 s
   layer norm (AVG and VAR_POP OVER (PARTITION BY id))	1.80 s	7.5 s

   The two linear layers run at about 2.3G joined pairs/s, about 25× DataFusion's 0.09G/s. Adding the layer norm then adds 1.5 s to the 1.8 s. That isn't the windows on their own. On synthetic tables of the same shape (2,048 × 512, with 512 × 512 weights), a linear layer and a layer norm with a learned scale and offset, joined from a long-form norms table through nets, takes 0.21 s on einfold, against 3.5 s on DataFusion. So the 1.5 s comes from how the layer norm combines with GraphCast's full chain: weights and biases selected from long-form tables holding all 88 weight matrices, then bias, swish, the second layer and the norm. I haven't narrowed it further yet. plans_with_pipelines() stays at 0 on these plans, and plans_with_einfold() counts them.

   The demo is now written as views, one per stage (#270 in xarray-sql), so that it reads like the model. As views, it doesn't run past the encoder anywhere yet:

   • On DataFusion, the plan grows 4.7× per processor round (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth xarray-sql#267).
   • DuckDB takes 45 s to plan one round and more than 10 minutes to plan two (DuckDB plans views that share reads in time exponential in their depth — draft for upstream xarray-sql#292).
   • On einfold, the grid embedding and update as views took 181 s, against 126 s stored, before the run was stopped for low memory.

   Once DataFusion computes each shared read once, the views should run at the stored speeds above.