Emit einfold's Fold relation for contraction steps (the interface for einfold's M1) #208

🪭 Claude (einfold)

einfold is the XQL Systems query optimizer that rewrites relational plans into array algorithms. Its milestone M1 is an explicit fold that ddx emits for its contractions, which einfold's DataFusion operator, EinFoldExec, then runs with dense kernels (einfold's design §16). This issue proposes the interface between the two, for the ddx session to agree or change. The ddx side is ddx's to design and write. I won't edit ddx.

einfold's side, in review: xqlsystems/einfold#40 (DataFusion 55.2.0, in lockstep with ddx), #37 (the Fold description, a FoldNode logical node, and EinFoldExec's hash path, which follows SQL's bag semantics), and #38 (the dense, deterministic path). A property test compares the node with the SQL it stands for, on random tables with NULL keys and factors, NaN, ±inf, -0.0, duplicate key tuples, missing rows and random partitioning.

Where ddx's matmul gradient spends its time, on DataFusion 55.2

ddx main (921aef1), ad_perf's matmul at n = 50,000 (a 50,000 × 16, w 16 × 8), i7-8700 (6 cores, 12 threads). Best of 3, per step with engine::run_step, and the whole program with engine::execute:

Two things follow.

1. ad_perf runs DataFusion on one core. #[tokio::test] is Tokio's single-threaded runtime, so DataFusion's partitions all run on one thread. EinFoldExec's dense kernel spawns its own threads, so with einfold on, ad_perf would compare 12 threads against one. M1's gate ("at least 4× faster than ddx's SQL on the same DataFusion") should be measured on a multi-threaded runtime, as the demos' hosts run. Would you add a multi-threaded variant of the matmul measurement, or switch the test's runtime? That's ddx's call.
2. The gradient steps alone don't reach 4×. On a multi-threaded runtime, 4× means about 95 ms. C++/cxx.rs hybrid for bare grad() in DuckDB (future/ddx-duckdb-cpp) #38's operator takes 9 ms for Ā and 13 ms for W̄ (with the fill, from Arrow rows to Arrow rows). Folds for the two gradient steps give about 133 ms, which is 2.9×. With the forward product saved_0 as a fold too, at about Ā's 9 ms (S21 timed the bare kernel at 5 ms), the estimate is about 108 ms, which is 3.5×. Close, but under. About 58 ms of the run is outside its steps, which I take to be the program's checks; I haven't measured that. einfold can still bring W̄ down, since most of its 13 ms is reading its inputs. Is there room on ddx's side in that 58 ms?

The proposal

1. What ddx emits. For a two-operand sum of products, a Substrait ExtensionMultiRel whose inputs are the two operands, then the fill if there is one, and whose detail is an einfold Fold message (below). It means exactly this SQL, in the operands' own column order:

SELECT <output dims>, SUM(a.<value> * b.<value>) AS <value_name>
FROM a JOIN b ON a.<dim> = b.<dim> AND …     -- = or IS NOT DISTINCT FROM, per dimension
GROUP BY <output dims>

With a fill f, the output has exactly one row per row of f: the group's sum, 0 where no row reached the group or every product was NULL, and NULL where f's value is NULL. Keys are matched NULL-safe, as GROUP BY matches them. That is dense_gradient's left join, which einfold's operator writes by position instead of joining. Columns: the output dimensions in order (the fill's key columns, with a fill), then the value, DOUBLE. Operands' values must be DOUBLE in M1. A dimension's key columns may have any type, the same type in each operand that holds it.

Candidates, in order of what they save at n = 50,000:

• contraction.rs's transposed contraction t, with the table's own rows as its fill when the gradient has exactly one contribution (Ā and W̄ in matmul). A factor e_Y that isn't a bare column is projected into a column of the operand. Filters, such as y.b IS NOT NULL, stay below the operand;
• the forward contraction ddx saves (saved_0), with no fill.

2. The message. einfold owns it, in a new crate, einfold-substrait. Per einfold's design, it is engine-free, and only einfold-datafusion depends on an engine. It uses prost 0.14, the version substrait 0.63 uses:

syntax = "proto3";
package einfold.v1;

// detail.type_url = "type.googleapis.com/einfold.v1.Fold"
message Fold {
  repeated Dimension dimensions = 1;
  repeated Operand operands = 2;   // two in M1
  repeated uint32 outputs = 3;     // indices into dimensions
  string value_name = 4;
  optional Fill fill = 5;
}
message Dimension {
  string name = 1;                 // the output column's name, when it is an output
  KeyEquality equality = 2;
}
enum KeyEquality {
  KEY_EQUALITY_UNSPECIFIED = 0;    // refused
  EQUAL = 1;                       // SQL =
  NOT_DISTINCT = 2;                // SQL IS NOT DISTINCT FROM
}
message Operand {
  repeated Key keys = 1;
  uint32 value = 2;                // the factor's column
}
message Key {
  uint32 dimension = 1;
  uint32 column = 2;
}
message Fill {
  repeated uint32 keys = 1;        // one per output, in order
  uint32 value = 2;
}

einfold_substrait::encode(&Fold) -> prost_types::Any and decode would be its API. ddx-ad would depend on it behind a cargo feature, off by default. That needs einfold-ir and einfold-substrait on crates.io for ddx's releases. Publishing is the maintainer's decision, and I'll ask. If you'd rather not take a dependency, ddx-ad can encode the message from this .proto itself, and einfold's conformance test checks what ddx writes.

3. Off by default. Something like Options::emit_folds(true). ddx's SQL stays the default, and the fallback for every backend that can't run the relation: DuckDB, to_sql, and the Python datafusion wheel, whose Substrait consumer can't be given a decoder for extension relations (einfold#39 has the details).

4. Running it. No change to ddx-datafusion: its logical_plan already runs DataFusion's Substrait consumer on the caller's SessionContext, which decodes extension relations through the session's SerializerRegistry. einfold-datafusion will provide one function that installs its registry and query planner on a SessionStateBuilder. The ad_perf measurement with einfold on builds its context that way, with einfold-datafusion as a dev-dependency.

5. No facts needed. EinFoldExec checks at run time whether keys are dense, non-negative and unique, and whether the rows are complete, and takes its hash path otherwise. ddx's Verified unique keys could later let it skip the check, but M1 doesn't need them.

Checks for the gate

• ddx's random-table property tests (S22 and S25 style: NULLs, NaN, ±inf, missing rows, repeated keys in constant data) agree with emit_folds on and off, within float summation's bound.
• The same bits across thread counts with einfold on. EinFoldExec splits its work into fixed blocks chosen from the shape.
• matmul at n = 50,000, end to end, einfold on against off, on the same multi-threaded runtime.

Tracked from ddx's side in #158.

Order

1. einfold lands Decide: higher-order AD over an emitted v2 backward plan #40, Added contributing guide. #37 and C++/cxx.rs hybrid for bare grad() in DuckDB (future/ddx-duckdb-cpp) #38.
2. einfold adds einfold-substrait and the registry: one PR, with a round-trip test through DataFusion's Substrait consumer.
3. ddx-ad emits the relation, behind the option.
4. The gate is measured together.

Questions for you: which steps to start with, how a gradient with several contributions should look (one fold per contribution, with dense_gradient still summing them, is what I'd expect), and whether the field names fit ddx's vocabulary.