diff --git a/Cargo.lock b/Cargo.lock index 0e147163..3d69fd8e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1365,6 +1365,7 @@ dependencies = [ "arrow", "datafusion", "einfold-ir", + "ndarray", "tokio", ] @@ -1971,6 +1972,16 @@ dependencies = [ "twox-hash", ] +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "md-5" version = "0.11.0" @@ -1997,6 +2008,21 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "ndarray" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "num-bigint" version = "0.4.8" @@ -2211,6 +2237,21 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" +[[package]] +name = "portable-atomic" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" + +[[package]] +name = "portable-atomic-util" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.6" @@ -2298,6 +2339,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "recursive" version = "0.1.1" diff --git a/crates/einfold-kernels/Cargo.toml b/crates/einfold-kernels/Cargo.toml index 370a8a22..9917ad4e 100644 --- a/crates/einfold-kernels/Cargo.toml +++ b/crates/einfold-kernels/Cargo.toml @@ -17,6 +17,9 @@ rust-version.workspace = true # without a copy. No query engine. arrow = { workspace = true } einfold-ir = { workspace = true } +# Matrix products for the fused pipeline's tiles: ndarray's safe interface +# to matrixmultiply, which picks AVX2 and FMA kernels at run time. +ndarray = { version = "0.16", default-features = false, features = ["std"] } [dev-dependencies] # The property test checks the kernels against SQL that DataFusion runs. diff --git a/crates/einfold-kernels/src/algorithm.rs b/crates/einfold-kernels/src/algorithm.rs index 601d5650..aef69e80 100644 --- a/crates/einfold-kernels/src/algorithm.rs +++ b/crates/einfold-kernels/src/algorithm.rs @@ -3,16 +3,18 @@ // SPDX-License-Identifier: Apache-2.0 //! [`fold`] and [`fold_dense`], which run a fold by the best algorithm that -//! applies, and [`Algorithm`], which says which one ran. +//! applies, [`Algorithm`], which says which one ran, and [`pipeline`], which +//! runs a fold pipeline. use arrow::compute; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use einfold_ir::Fold; +use einfold_ir::{Fold, Pipeline}; +use crate::fused::Piped; use crate::options::{DomainRows, Options}; use crate::relation::Relation; -use crate::{bad, dense, discriminate, Result}; +use crate::{bad, dense, discriminate, fused, Result}; /// Which algorithm computed a fold. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -113,3 +115,24 @@ pub fn fold( batches, }) } + +/// `pipeline` over its producer's operands `q` and `k` and its values `v`, as +/// batches of `schema` (which [`pipeline_schema`](crate::pipeline_schema) +/// gives), by the fused algorithm: tile by tile, without the producer's +/// results ever existing in full. `None` when the algorithm's requirements +/// don't hold: every input complete, with one row and a non-NULL factor at +/// every position of its dimensions, keys small non-negative integers, and +/// at most one condition, between the producer's two outputs. The caller +/// then runs something else, such as the host's plan. +/// +/// The result's bits depend only on the inputs' rows, not on how they were +/// batched or ordered, nor on the number of threads. +pub fn pipeline( + pipeline: &Pipeline, + [q, k, v]: [Relation; 3], + schema: SchemaRef, + options: Options, +) -> Result> { + let threads = options.threads.max(1); + fused::pipeline(pipeline, [q.batches, k.batches, v.batches], schema, threads) +} diff --git a/crates/einfold-kernels/src/dense.rs b/crates/einfold-kernels/src/dense.rs index 440f0130..064ee7af 100644 --- a/crates/einfold-kernels/src/dense.rs +++ b/crates/einfold-kernels/src/dense.rs @@ -83,7 +83,7 @@ fn worth(cells: usize, rows: usize) -> bool { /// plus one), or `None` if it isn't an integer column, or holds a NULL, a /// negative key, or one too large for `i64` (an unsigned key of 2⁶³ or more), /// or whose extent doesn't fit (a key of `i64::MAX`). -fn keys(column: &ArrayRef) -> Option<(Int64Array, usize)> { +pub(crate) fn keys(column: &ArrayRef) -> Option<(Int64Array, usize)> { if !column.data_type().is_integer() || column.null_count() > 0 { return None; } @@ -104,7 +104,7 @@ fn keys(column: &ArrayRef) -> Option<(Int64Array, usize)> { /// column at a time, in 32 bits, which the compiler vectorizes: every /// position fits, as the dense algorithm runs only on fewer cells than /// `u32::MAX`. -fn positions( +pub(crate) fn positions( keys: &[(usize, Int64Array)], dims: &[usize], extent: &[usize], diff --git a/crates/einfold-kernels/src/fused.rs b/crates/einfold-kernels/src/fused.rs new file mode 100644 index 00000000..af56c376 --- /dev/null +++ b/crates/einfold-kernels/src/fused.rs @@ -0,0 +1,477 @@ +// SPDX-FileCopyrightText: 2026 Alexander Merose & einfold Authors +// +// SPDX-License-Identifier: Apache-2.0 + +//! EinFold's fused pipeline algorithm: a [`Pipeline`] computed tile by tile, +//! so that the producer's results never exist in full. +//! +//! It is loop fusion: the producer's loop over its groups and the consumer's +//! loop over the values it sums share one loop over pairs of tiles. That +//! applies to any pipeline whose producer groups pair the consumer's groups +//! with the dimensions it sums away, whatever the query means. Two such +//! queries: attention, in which each query position averages the values at +//! the key positions, weighted by a softmax of its dot products with the +//! keys (each exponentiated, then divided by their sum, so that the weights +//! sum to one); and a kernel smoother, which averages values at source +//! points, weighted by a function of their distance to each target point. +//! The names below follow attention, where they are familiar: the operands +//! `q`, `k` and `v`, and a query, a key and a feature. +//! +//! Every dimension has one of four roles: +//! +//! - **X**: a producer output held by its first operand, `q`, kept by the +//! pipeline (in attention, the query `i`; in a kernel smoother, the target +//! point); +//! - **E**: summed by the producer, held by both its operands (a feature, or +//! a coordinate); +//! - **J**: a producer output held by its second operand, `k`, summed by the +//! consumer and joined with the values operand `v` (a key, or a source +//! point); +//! - **Y**: the pipeline's own dimensions, held by `v` (a value's feature). +//! +//! As in the dense fold, keys are positions, each role's dimensions are +//! flattened into one position, and every operand is placed in a dense +//! matrix. For each tile of X positions, on one thread, the J positions are +//! visited in tiles, in order. Each pair of tiles computes its block of +//! producer values with a matrix product over E, turns them into scores and +//! weights, and folds the weights into the tile's state, then the next pair +//! reuses the block. Memory beyond inputs and output is one block per thread. +//! +//! **The scaled accumulator.** A row's state is `(m, z, a[y])`: a shift `m`, +//! the sum of its weights `z = Σ exp(s − m)`, and the weighted sums `a[y] = +//! Σ exp(s − m)·v[y]`. A block whose largest score is above `m` raises the +//! shift to it, and rescales `z` and `a` by `exp(m_old − m_new)` first. This +//! is Milakov and Gimelshein's online softmax ("Online normalizer calculation +//! for softmax", 2018). With [`Shift::KeptMaximum`], it equals SQL's `exp(s − +//! MAX(s))`, since the shift cancels in `a / z`. No weight exceeds 1, and +//! each term goes through one rounding of `exp` and a product per later +//! raise of the shift. With [`Shift::None`], `m` stays 0, as SQL's `exp(s)` +//! does, overflow included. +//! +//! **Masks are absent rows.** A [`Condition`] between the X and J dimensions +//! allows, for each X position, an interval of J positions, since keys are +//! positions. A pair of tiles the condition disallows entirely is never +//! visited. In a pair it disallows in part, disallowed pairs are skipped, not +//! given weight 0: `0 · NaN` is NaN, but SQL forms no row for them. An X +//! position with no allowed J position has no rows in the output, as in SQL. +//! +//! **When SQL and dense arrays agree.** Every operand must be complete: one +//! row at every position of its dimensions, and no NULL factors. Then every +//! producer group exists unless masked, both [`Normalizer`]s agree, and NaN +//! and infinities propagate as SQL's arithmetic does. Otherwise the caller +//! runs something else, such as the host's plan. +//! +//! **Determinism.** Tiles are fixed by the shape, each X tile is computed by +//! one thread, J tiles are visited in order, and each tile's matrix products +//! are the same calls, on the same shapes, on every run. So the bits don't +//! depend on the number of threads nor on the order in which rows arrive. +//! They can differ between machines, as the matrix product picks its +//! kernel, with or without fused multiply-adds, by the CPU's features. +//! +//! [`Condition`]: einfold_ir::Condition + +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array, Int64Array, RecordBatch}; +use arrow::compute; +use arrow::datatypes::SchemaRef; +use einfold_ir::{Comparison, DimensionId, Operand, Pipeline, Shift}; + +use crate::dense; +use crate::gemm::{self, Matrix}; +use crate::Result; + +/// X positions per tile, and J positions per tile. +const XT: usize = 64; +const JT: usize = 64; + +/// A pipeline's result, and how many pairs of tiles its masks let it skip. +#[derive(Clone, Debug)] +pub struct Piped { + /// The result's rows. + pub batches: Vec, + /// Pairs of an X tile and a J tile that were computed. + pub tiles_visited: usize, + /// Every pair of an X tile and a J tile. + pub tiles: usize, +} + +/// How one operand's rows are placed: its dimensions along the matrix's rows +/// and along its columns. +struct Placement<'a> { + operand: &'a Operand, + batches: &'a [RecordBatch], + rows: &'a [usize], + columns: &'a [usize], +} + +/// `pipeline` over its operands `q`, `k` and values `v`, by the fused +/// algorithm on `threads` threads, or `None` when its requirements don't +/// hold. +pub(crate) fn pipeline( + pipeline: &Pipeline, + [q, k, v]: [&[RecordBatch]; 3], + schema: SchemaRef, + threads: usize, +) -> Result> { + let producer = pipeline.producer(); + let [op_q, op_k] = producer.operands() else { + return Ok(None); + }; + let values = pipeline.values(); + let own = producer.dimensions().len(); + let n = own + pipeline.dimensions().len(); + let holds = |op: &Operand, d: usize| op.column_of(DimensionId(d)).is_some(); + let output = |d: usize| pipeline.outputs().contains(&DimensionId(d)); + let (mut xs, mut es, mut js) = (Vec::new(), Vec::new(), Vec::new()); + for d in 0..own { + let produced = producer.output_of(DimensionId(d)).is_some(); + match ( + holds(op_q, d), + holds(op_k, d), + produced, + output(d), + holds(values, d), + ) { + (true, false, true, true, false) => xs.push(d), + (true, true, false, false, false) => es.push(d), + (false, true, true, false, true) => js.push(d), + _ => return Ok(None), + } + } + let ys: Vec = (own..n).collect(); + // The output has a row for each position of the pipeline's own + // dimensions, so one left out of the outputs, to be summed away, isn't + // this algorithm's. + if !ys.iter().all(|&d| output(d)) { + return Ok(None); + } + // A condition between the one X and the one J dimension, as `x j`. + let mask = match producer.conditions() { + [] => None, + [c] => match (&xs[..], &js[..]) { + // The four comparisons `allowed` knows; one added later declines. + ([x], [j]) if (c.left.0, c.right.0) == (*x, *j) => match c.comparison { + Comparison::Less + | Comparison::LessOrEqual + | Comparison::Greater + | Comparison::GreaterOrEqual => Some(c.comparison), + _ => return Ok(None), + }, + ([x], [j]) if (c.left.0, c.right.0) == (*j, *x) => match c.comparison { + Comparison::Less => Some(Comparison::Greater), + Comparison::LessOrEqual => Some(Comparison::GreaterOrEqual), + Comparison::Greater => Some(Comparison::Less), + Comparison::GreaterOrEqual => Some(Comparison::LessOrEqual), + _ => return Ok(None), + }, + _ => return Ok(None), + }, + _ => return Ok(None), + }; + + // Each batch's keys as positions, and each dimension's extent. + let placements = [ + Placement { + operand: op_q, + batches: q, + rows: &xs, + columns: &es, + }, + // k is placed transposed, [e][j], so that the scores are q · k. + Placement { + operand: op_k, + batches: k, + rows: &es, + columns: &js, + }, + Placement { + operand: values, + batches: v, + rows: &js, + columns: &ys, + }, + ]; + let mut extent = vec![0usize; n]; + let mut all_keys = Vec::new(); + for p in &placements { + let batch_keys = dense::par_tasks(p.batches.len(), threads, |i| { + (p.operand.keys.iter()) + .map(|key| { + let (k, n) = dense::keys(p.batches[i].column(key.column.0))?; + Some((key.dimension.0, k, n)) + }) + .collect::>>() + }); + let Some(batch_keys) = batch_keys.into_iter().collect::>>() else { + return Ok(None); + }; + for &(d, _, n) in batch_keys.iter().flatten() { + extent[d] = extent[d].max(n); + } + all_keys.push(batch_keys); + } + let size = |dims: &[usize]| { + dims.iter() + .try_fold(1usize, |n, &d| n.checked_mul(extent[d])) + }; + let (Some(nx), Some(ne), Some(nj), Some(ny)) = (size(&xs), size(&es), size(&js), size(&ys)) + else { + return Ok(None); + }; + let mut matrices = Vec::new(); + for (p, batch_keys) in placements.iter().zip(all_keys) { + let keys: Vec> = (batch_keys.into_iter()) + .map(|keys| keys.into_iter().map(|(d, k, _)| (d, k)).collect()) + .collect(); + let (Some(rows), Some(columns)) = (size(p.rows), size(p.columns)) else { + return Ok(None); + }; + let Some(matrix) = place(p, &keys, &extent, (rows, columns), threads) else { + return Ok(None); + }; + matrices.push(matrix); + } + let [q, kt, v] = &matrices[..] else { + unreachable!("three placements"); + }; + if nx + .checked_mul(ny) + .is_none_or(|cells| cells >= u32::MAX as usize) + { + return Ok(None); + } + + // Each X position's allowed J positions, an interval. + let allowed = |x: usize| -> (usize, usize) { + match mask { + None => (0, nj), + Some(Comparison::GreaterOrEqual) => (0, (x + 1).min(nj)), + Some(Comparison::Greater) => (0, x.min(nj)), + Some(Comparison::LessOrEqual) => (x.min(nj), nj), + Some(_) => ((x + 1).min(nj), nj), + } + }; + let (x_tiles, j_tiles) = (nx.div_ceil(XT), nj.div_ceil(JT)); + let tiles = dense::par_tasks(x_tiles, threads, |t| { + tile(pipeline, t, (nx, ne, nj, ny), [q, kt, v], &allowed) + }); + let tiles_visited = tiles.iter().map(|t| t.visited).sum(); + + // The rows: every y for each x with an allowed J position, in position + // order. + let rows: Vec = (0..nx) + .filter(|&x| { + nj > 0 && { + let (lo, hi) = allowed(x); + lo < hi + } + }) + .collect(); + let total = rows.len() * ny; + let mut columns: Vec = Vec::with_capacity(pipeline.outputs().len() + 1); + for &DimensionId(d) in pipeline.outputs() { + // The key's place in x's or y's position, first dimension slowest. + let (in_x, role) = match xs.contains(&d) { + true => (true, &xs), + false => (false, &ys), + }; + let slot = role.iter().position(|&e| e == d).expect("an output's role"); + let stride: usize = role[slot + 1..].iter().map(|&e| extent[e]).product(); + let key_at = |p: usize| ((p / stride) % extent[d]) as i64; + let mut key = Vec::with_capacity(total); + match in_x { + true => rows + .iter() + .for_each(|&x| key.extend(std::iter::repeat_n(key_at(x), ny))), + false => { + let along: Vec = (0..ny).map(key_at).collect(); + rows.iter().for_each(|_| key.extend_from_slice(&along)); + } + } + let ty = schema.field(columns.len()).data_type(); + columns.push(compute::cast( + &(Arc::new(Int64Array::from(key)) as ArrayRef), + ty, + )?); + } + let mut out = Vec::with_capacity(total); + for &x in &rows { + let state = &tiles[x / XT]; + out.extend_from_slice(&state.out[(x % XT) * ny..(x % XT + 1) * ny]); + } + columns.push(Arc::new(Float64Array::from(out))); + Ok(Some(Piped { + batches: vec![RecordBatch::try_new(schema, columns)?], + tiles_visited, + tiles: x_tiles * j_tiles, + })) +} + +/// `p`'s rows as a dense `rows × columns` matrix of its factors, or `None` +/// unless it has exactly one row, with a non-NULL factor, at every position. +fn place( + p: &Placement, + keys: &[Vec<(usize, Int64Array)>], + extent: &[usize], + (rows, columns): (usize, usize), + threads: usize, +) -> Option> { + use arrow::array::Array; + + let cells = rows.checked_mul(columns)?; + let total: usize = p.batches.iter().map(|b| b.num_rows()).sum(); + if total != cells || cells >= u32::MAX as usize { + return None; + } + let factor = p.operand.factor?; + let positions = dense::par_tasks(p.batches.len(), threads, |i| { + let n = p.batches[i].num_rows(); + let row = dense::positions(&keys[i], p.rows, extent, n); + let column = dense::positions(&keys[i], p.columns, extent, n); + (row, column) + }); + let mut matrix = vec![0.0; cells]; + let mut held = vec![false; cells]; + for (batch, (row, column)) in p.batches.iter().zip(&positions) { + let factors = batch + .column(factor.0) + .as_any() + .downcast_ref::()?; + if factors.null_count() > 0 { + return None; + } + for ((&r, &c), &x) in row.iter().zip(column).zip(factors.values()) { + let cell = r as usize * columns + c as usize; + if std::mem::replace(&mut held[cell], true) { + return None; + } + matrix[cell] = x; + } + } + // As many rows as cells, none repeated: every cell is held. + Some(matrix) +} + +/// One X tile's averages, `[x][y]` over its positions, and how many J tiles +/// it visited. +struct Tile { + out: Vec, + visited: usize, +} + +/// X tile `t` of `pipeline`, from `q [x][e]`, `kt [e][j]` and `v [j][y]`, +/// with `allowed(x)` the interval of J positions that position `x` may join. +fn tile( + pipeline: &Pipeline, + t: usize, + (nx, ne, nj, ny): (usize, usize, usize, usize), + [q, kt, v]: [&[f64]; 3], + allowed: &(impl Fn(usize) -> (usize, usize) + Sync), +) -> Tile { + let (x0, x1) = (t * XT, ((t + 1) * XT).min(nx)); + let rows = x1 - x0; + let scale = pipeline.scale(); + let shifted = pipeline.shift() == Shift::KeptMaximum; + let mut m = vec![if shifted { f64::NEG_INFINITY } else { 0.0 }; rows]; + let mut z = vec![0.0; rows]; + let mut a = vec![0.0; rows * ny]; + let mut block = vec![0.0; XT * JT]; + let mut visited = 0; + for j0 in (0..nj).step_by(JT) { + let width = JT.min(nj - j0); + // Each row's allowed positions within this tile, relative to j0. + let spans: Vec<(usize, usize)> = (x0..x1) + .map(|x| { + let (lo, hi) = allowed(x); + (lo.clamp(j0, j0 + width) - j0, hi.clamp(j0, j0 + width) - j0) + }) + .collect(); + if spans.iter().all(|&(lo, hi)| lo >= hi) { + continue; + } + visited += 1; + let full = spans.iter().all(|&span| span == (0, width)); + let block = &mut block[..rows * width]; + block.fill(0.0); + // The producer's values for the pair: q[x0..x1] · kt[.., j0..j0+width]. + let q_tile = Matrix { + data: &q[x0 * ne..], + stride: ne, + }; + let k_tile = Matrix { + data: &kt[j0..], + stride: nj, + }; + gemm::multiply_add((rows, ne, width), q_tile, k_tile, block, width); + for (r, &(lo, hi)) in spans.iter().enumerate() { + if lo >= hi { + continue; + } + let row = &mut block[r * width..(r + 1) * width]; + let scores = &mut row[lo..hi]; + for s in scores.iter_mut() { + *s = scale.apply(*s); + } + // The largest score so far, or NaN once any score is NaN: its + // weight is NaN, and so is every average of its row, in SQL too. + let shift = match shifted { + true => (scores.iter()).fold(m[r], |m, &s| if s > m || s.is_nan() { s } else { m }), + false => 0.0, + }; + if shift == f64::NEG_INFINITY { + // Every score so far is -inf, so every weight is exp(-inf - + // shift) for the final shift, or NaN if that is -inf too: + // 0 here, rescaled to NaN below if it stays so. + scores.fill(0.0); + continue; + } + let rescale = (m[r] - shift).exp(); + let mut sum = 0.0; + for s in scores.iter_mut() { + *s = (*s - shift).exp(); + sum += *s; + } + if shifted && rescale != 1.0 { + z[r] *= rescale; + a[r * ny..(r + 1) * ny] + .iter_mut() + .for_each(|a| *a *= rescale); + } + z[r] += sum; + m[r] = shift; + } + // The weighted values: a += w · v[j0..j0+width], skipping disallowed + // pairs, which have no row. + let v_tile = Matrix { + data: &v[j0 * ny..], + stride: ny, + }; + if full { + let w = Matrix { + data: block, + stride: width, + }; + gemm::multiply_add((rows, width, ny), w, v_tile, &mut a, ny); + } else { + for (r, &(lo, hi)) in spans.iter().enumerate() { + let w = Matrix { + data: &block[r * width + lo..], + stride: width, + }; + let v_rows = Matrix { + data: &v[(j0 + lo) * ny..], + stride: ny, + }; + let hi = hi.max(lo); + gemm::multiply_add((1, hi - lo, ny), w, v_rows, &mut a[r * ny..], ny); + } + } + } + // A row whose every score was -inf has z = 0, and 0 / 0 is NaN, as + // SQL's exp(-inf - -inf) is. + for (r, z) in z.iter().enumerate() { + a[r * ny..(r + 1) * ny].iter_mut().for_each(|a| *a /= z); + } + Tile { out: a, visited } +} diff --git a/crates/einfold-kernels/src/gemm.rs b/crates/einfold-kernels/src/gemm.rs new file mode 100644 index 00000000..6cdf653d --- /dev/null +++ b/crates/einfold-kernels/src/gemm.rs @@ -0,0 +1,90 @@ +// SPDX-FileCopyrightText: 2026 Alexander Merose & einfold Authors +// +// SPDX-License-Identifier: Apache-2.0 + +//! [`multiply_add`]: matrix products for tiles that fit in cache, by +//! [`matrixmultiply`](https://docs.rs/matrixmultiply), a Rust crate of dense +//! matrix products, through the safe interface of +//! [`ndarray`](https://docs.rs/ndarray), Rust's n-dimensional arrays. It +//! picks kernels using AVX2 and FMA, x86's 256-bit vector and fused +//! multiply-add instructions, at run time when the CPU has them. + +/// A matrix in a slice, row-major, with `stride` elements from one row's +/// start to the next. +#[derive(Clone, Copy)] +pub(crate) struct Matrix<'a> { + pub(crate) data: &'a [f64], + pub(crate) stride: usize, +} + +/// `c[m×n] += a[m×k] · b[k×n]`, with `c` row-major at `c_stride`. +/// +/// Each cell is within float summation's bound of the exact sum. The order +/// of its additions is the library's, fixed for a given shape on a given +/// machine, so the same call gives the same bits. +pub(crate) fn multiply_add( + (m, k, n): (usize, usize, usize), + a: Matrix, + b: Matrix, + c: &mut [f64], + c_stride: usize, +) { + use ndarray::{linalg, ArrayView2, ArrayViewMut2, ShapeBuilder}; + // With no terms to add, or no cells to add them to, nothing changes. + if m == 0 || n == 0 || k == 0 { + return; + } + fn view(data: &[f64], rows: usize, cols: usize, stride: usize) -> ArrayView2<'_, f64> { + let len = (rows - 1) * stride + cols; + ArrayView2::from_shape((rows, cols).strides((stride, 1)), &data[..len]).expect("a matrix") + } + let (a, b) = (view(a.data, m, k, a.stride), view(b.data, k, n, b.stride)); + let len = (m - 1) * c_stride + n; + let mut c = + ArrayViewMut2::from_shape((m, n).strides((c_stride, 1)), &mut c[..len]).expect("a matrix"); + linalg::general_mat_mul(1.0, &a, &b, 1.0, &mut c); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn products_within_the_summation_bound_and_repeatable() { + // Sizes on both sides of any register block, inside larger matrices. + let (m, k, n) = (9, 5, 19); + let value = |i: usize| ((i as f64) * 0.37).sin() * 1e3; + let a: Vec = (0..m * 7).map(value).collect(); + let b: Vec = (0..k * 23).map(|i| value(i + 101)).collect(); + let c0: Vec = (0..m * 21).map(|i| value(i + 7)).collect(); + let mut c = c0.clone(); + let (am, bm) = ( + Matrix { + data: &a, + stride: 7, + }, + Matrix { + data: &b, + stride: 23, + }, + ); + multiply_add((m, k, n), am, bm, &mut c, 21); + let u = f64::EPSILON / 2.0; + for cell in 0..m * 21 { + let (i, j) = (cell / 21, cell % 21); + if j >= n { + assert_eq!(c[cell].to_bits(), c0[cell].to_bits(), "outside the product"); + continue; + } + let terms: Vec = (0..k).map(|p| a[i * 7 + p] * b[p * 23 + j]).collect(); + let exact = c0[cell] + terms.iter().sum::(); + let magnitude = c0[cell].abs() + terms.iter().map(|t| t.abs()).sum::(); + assert!((c[cell] - exact).abs() <= 2.0 * (k as f64 + 1.0) * u * magnitude); + } + // The same call gives the same bits. + let mut again = c0.clone(); + multiply_add((m, k, n), am, bm, &mut again, 21); + let bits = |v: &[f64]| v.iter().map(|x| x.to_bits()).collect::>(); + assert_eq!(bits(&c), bits(&again)); + } +} diff --git a/crates/einfold-kernels/src/lib.rs b/crates/einfold-kernels/src/lib.rs index 2c4d9025..8a78c393 100644 --- a/crates/einfold-kernels/src/lib.rs +++ b/crates/einfold-kernels/src/lib.rs @@ -18,6 +18,10 @@ //! few things SQL leaves open, such as whether `-0.0` and `0.0` are one key. //! [`Host`] says which. //! +//! [`pipeline`] computes a fold pipeline (see [`einfold_ir::pipeline`]), such +//! as attention, one tile at a time, so that its producer's results never +//! exist in full. +//! //! Nothing here depends on a query engine, so any engine that holds its data //! as Arrow arrays can run these, from DataFusion to an extension of DuckDB. @@ -26,15 +30,18 @@ use arrow::error::ArrowError; mod algorithm; mod dense; mod discriminate; +mod fused; +mod gemm; mod keys; mod options; mod relation; mod schema; -pub use algorithm::{fold, fold_dense, Algorithm, Folded}; +pub use algorithm::{fold, fold_dense, pipeline, Algorithm, Folded}; +pub use fused::Piped; pub use options::{DomainRows, FloatKeys, Host, Options}; pub use relation::Relation; -pub use schema::schema; +pub use schema::{pipeline_schema, schema}; /// The result of a kernel. pub type Result = std::result::Result; diff --git a/crates/einfold-kernels/src/schema.rs b/crates/einfold-kernels/src/schema.rs index 20d6dc36..7a16cda7 100644 --- a/crates/einfold-kernels/src/schema.rs +++ b/crates/einfold-kernels/src/schema.rs @@ -2,12 +2,13 @@ // // SPDX-License-Identifier: Apache-2.0 -//! [`schema`]: the schema of a fold's result. +//! [`schema`] and [`pipeline_schema`]: the schemas of a fold's and a +//! pipeline's results. use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use einfold_ir::{ColumnIndex, Fold, Key}; +use einfold_ir::{ColumnIndex, DimensionId, Fold, Key, Pipeline}; use crate::{bad, Result}; @@ -118,3 +119,64 @@ pub fn schema(fold: &Fold, operands: &[&Schema], domain: Option<&Schema>) -> Res )); Ok(Arc::new(Schema::new(fields))) } + +/// The schema of `pipeline`'s result over inputs of these schemas: its +/// producer's two operands, then its values operand. It has one key column +/// per output, typed as the dimension's columns are, then the average, a +/// nullable `Float64`. Columns are named by position, as [`schema`]'s are. +/// +/// Fails unless the producer's schema is valid (see [`schema`]), the values +/// operand's columns exist, its factor is `Float64`, and each of its keys has +/// the type of its dimension's other columns. +pub fn pipeline_schema(pipeline: &Pipeline, inputs: &[&Schema; 3]) -> Result { + let producer = schema(pipeline.producer(), &inputs[..2], None)?; + let values = pipeline.values(); + let column = |ColumnIndex(c): ColumnIndex| match inputs[2].fields().get(c) { + Some(f) => Ok(f.as_ref().clone()), + None => bad(format!( + "column {c} of the values operand, of {}", + inputs[2].fields().len() + )), + }; + let factor = column(values.factor.expect("checked by Pipeline::try_new"))?; + if factor.data_type() != &DataType::Float64 { + return bad(format!( + "EinFold factors must be Float64, not {}", + factor.data_type() + )); + } + // A key's type: from the producer's output column, or the values'. + let key = |d: DimensionId| -> Result { + if let Some(o) = pipeline.producer().output_of(d) { + return Ok(producer.field(o).clone()); + } + let c = values.column_of(d).expect("checked by Pipeline::try_new"); + column(c) + }; + for k in &values.keys { + let (want, got) = (key(k.dimension)?, column(k.column)?); + if want.data_type() != got.data_type() { + return bad(format!( + "dimension {}'s columns have types {} and {}", + k.dimension.0, + want.data_type(), + got.data_type() + )); + } + } + let mut fields = Vec::new(); + for &d in pipeline.outputs() { + let f = key(d)?; + fields.push(Field::new( + fields.len().to_string(), + f.data_type().clone(), + f.is_nullable(), + )); + } + fields.push(Field::new( + fields.len().to_string(), + DataType::Float64, + true, + )); + Ok(Arc::new(Schema::new(fields))) +} diff --git a/crates/einfold-kernels/tests/pipeline_vs_sql.rs b/crates/einfold-kernels/tests/pipeline_vs_sql.rs new file mode 100644 index 00000000..53c786be --- /dev/null +++ b/crates/einfold-kernels/tests/pipeline_vs_sql.rs @@ -0,0 +1,652 @@ +// SPDX-FileCopyrightText: 2026 Alexander Merose & einfold Authors +// +// SPDX-License-Identifier: Apache-2.0 + +//! Property: EinFold's fused pipeline algorithm computes what the SQL query a +//! pipeline stands for computes on DataFusion, the host, for random +//! attention-shaped pipelines over random tables. +//! +//! Each case draws tables `q(t, d, x)`, `k(t, d, x)` and `v(t, [c,] x)` with +//! random extents, and a pipeline: scores `SUM(q.x * k.x)`, scaled by a +//! random constant or not, under a random causal-style condition between +//! `q.t` and `k.t` (either way round, any of `<`, `<=`, `>`, `>=`) or none; +//! weights `exp(s - MAX(s) OVER (PARTITION BY i))` or `exp(s)`; and an average +//! normalized by the kept groups' weights (a `norm` relation of `SUM(w) … GROUP BY i`) or by the +//! joined rows' (`SUM(w * v.x) / SUM(w)`). Values are drawn so that scores +//! spread widely, and some are NaN, ±inf, or large enough that `exp(s)` +//! overflows, and some rows' scores are all `-inf`, or `-inf` beside a NaN +//! in their first tile of keys only. Rows are shuffled and split into batches at random. The +//! invariants: +//! +//! 1. **Runs when it can, and declines when it can't.** With every table +//! complete, the algorithm runs. When a table misses a row, repeats one, +//! or holds a NULL factor, it either returns `None`, so that the caller +//! runs the host's plan, or (when the gap only shrinks a dimension's +//! range) gives SQL's result, as below. +//! 2. **Same rows.** The same keys as SQL: no rows for a query position that +//! the condition leaves no key. +//! 3. **Same values.** NaN where SQL gives NaN, the same infinity, and +//! otherwise within `1e-9·max|v.x|` of SQL's value: an average of `v.x` +//! with positive weights, so a scale on which any wrong weight, missed +//! term or masked pair counted shows. +//! 4. **Deterministic.** Other batches, another order and another number of +//! threads give the same rows, bit for bit. +//! +//! Reproduce a failure with `EINFOLD_SEED=`; search longer with +//! `EINFOLD_CASES=`. + +use std::sync::Arc; + +use datafusion::arrow::array::{Array, Float64Array, Int64Array, RecordBatch}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::datasource::MemTable; +use datafusion::prelude::{SessionConfig, SessionContext}; +use einfold_ir::{ + Aggregate, ColumnIndex, Comparison, Condition, Dimension, DimensionId, Fold, Key, KeyEquality, + Normalizer, Operand, Output, Pipeline, RowValue, Scale, Shift, +}; +use einfold_kernels::{FloatKeys, Host, Options, Relation}; + +struct Rng(u64); + +impl Rng { + /// A generator for `seed`, scrambled by SplitMix64 so that nearby seeds + /// give unrelated streams, and never zero, which xorshift can't leave. + fn new(seed: u64) -> Self { + let mut z = seed.wrapping_add(0x9e37_79b9_7f4a_7c15); + z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + Rng((z ^ (z >> 31)).max(1)) + } + fn next(&mut self) -> u64 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + self.0.wrapping_mul(0x2545_f491_4f6c_dd1d) + } + fn below(&mut self, n: usize) -> usize { + (self.next() % n as u64) as usize + } + fn chance(&mut self, p: f64) -> bool { + self.unit() < p + } + fn unit(&mut self) -> f64 { + ((self.next() >> 11) as f64) / ((1u64 << 53) as f64) + } +} + +/// A table of integer keys and a nullable `DOUBLE` value, as rows. +struct Table { + schema: SchemaRef, + rows: Vec<(Vec, Option)>, +} + +impl Table { + /// Every key tuple of `extents`, with values from `value`. + fn complete(names: &[&str], extents: &[usize], mut value: impl FnMut() -> f64) -> Table { + let mut fields: Vec = (names.iter()) + .map(|n| Field::new(*n, DataType::Int64, false)) + .collect(); + fields.push(Field::new("x", DataType::Float64, true)); + let cells: usize = extents.iter().product(); + let rows = (0..cells) + .map(|mut c| { + let mut keys = vec![0; extents.len()]; + for (k, &n) in keys.iter_mut().zip(extents).rev() { + *k = (c % n) as i64; + c /= n; + } + (keys, Some(value())) + }) + .collect(); + Table { + schema: Arc::new(Schema::new(fields)), + rows, + } + } + + /// The rows shuffled, in partitions of batches of random sizes. + fn partitions(&self, rng: &mut Rng) -> Vec> { + let mut order: Vec = (0..self.rows.len()).collect(); + for i in (1..order.len()).rev() { + order.swap(i, rng.below(i + 1)); + } + let mut parts = vec![Vec::new(); 1 + rng.below(3)]; + let mut start = 0; + while start < order.len() { + let end = (start + 1 + rng.below(300)).min(order.len()); + let rows = &order[start..end]; + let keys = self.schema.fields().len() - 1; + let mut columns: Vec> = (0..keys) + .map(|c| { + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|&r| self.rows[r].0[c]), + )) as _ + }) + .collect(); + columns.push(Arc::new( + rows.iter() + .map(|&r| self.rows[r].1) + .collect::(), + )); + let part = rng.below(parts.len()); + parts[part].push(RecordBatch::try_new(Arc::clone(&self.schema), columns).unwrap()); + start = end; + } + parts + } +} + +struct Case { + q: Table, + k: Table, + v: Table, + pipeline: Pipeline, + sql: String, + /// Whether every table has one non-NULL row per key tuple. + complete: bool, + /// The largest `|v.x|`, the scale of every average. + v_scale: f64, +} + +/// A value: mostly in [-1, 1] times `spread`, sometimes special. +fn value(rng: &mut Rng, spread: f64, specials: bool) -> f64 { + if specials && rng.chance(0.004) { + return [f64::NAN, f64::INFINITY, f64::NEG_INFINITY, 1e300, -1e300][rng.below(5)]; + } + (2.0 * rng.unit() - 1.0) * spread +} + +fn case(rng: &mut Rng) -> Case { + let (nq, nk, nd) = (1 + rng.below(150), 1 + rng.below(150), 1 + rng.below(6)); + let channels = (rng.chance(0.6)).then(|| 1 + rng.below(4)); + let specials = rng.chance(0.3); + let spread = [1.0, 4.0, 12.0][rng.below(3)]; + let mut q = Table::complete(&["t", "d"], &[nq, nd], || 0.0); + let mut k = Table::complete(&["t", "d"], &[nk, nd], || 0.0); + for r in q.rows.iter_mut().chain(k.rows.iter_mut()) { + r.1 = Some(value(rng, spread, specials)); + } + let mut v = match channels { + Some(c) => Table::complete(&["t", "c"], &[nk, c], || 0.0), + None => Table::complete(&["t"], &[nk], || 0.0), + }; + for r in v.rows.iter_mut() { + r.1 = Some(value(rng, 1.0, specials)); + } + // Rows of q whose scores are -inf: every one, which SQL turns into NaN; + // or only those with the first 64 keys, beside a NaN score, with finite + // ones after. Rows are (t, d) in order, so row t·nd + d is (t, d). + let at = |t: usize, d: usize| t * nd + d; + let target = rng.below(nq); + match rng.below(6) { + 0 => { + q.rows[at(target, 0)].1 = Some(f64::NEG_INFINITY); + for j in 0..nk { + let x = &mut k.rows[at(j, 0)].1; + *x = x.map(|x| x.abs() + 0.1); + } + } + 1 if nk > 64 => { + q.rows[at(target, 0)].1 = Some(-1e300); + for j in 0..64 { + k.rows[at(j, 0)].1 = Some(1e300); + } + k.rows[at(rng.below(64), nd - 1)].1 = Some(f64::NAN); + } + _ => {} + } + // Sometimes a table that isn't complete: a missing, repeated or NULL row. + let complete = !rng.chance(0.15); + if !complete { + let t = match rng.below(3) { + 0 => &mut q, + 1 => &mut k, + _ => &mut v, + }; + let r = rng.below(t.rows.len()); + match rng.below(3) { + 0 => drop(t.rows.remove(r)), + 1 => t.rows.push(t.rows[r].clone()), + _ => t.rows[r].1 = None, + } + } + let v_scale = (v.rows.iter()) + .filter_map(|r| r.1) + .filter(|x| x.is_finite()) + .fold(0.0f64, |m, x| m.max(x.abs())); + + // Dimensions: 0 = i (q.t), 1 = e (d), 2 = j (k.t), then 3 = c (v.c). + const EQ: Dimension = Dimension { + equality: KeyEquality::Equal, + }; + let key = |d, c| Key::new(DimensionId(d), ColumnIndex(c)); + let out = |d| Output { + dimension: DimensionId(d), + domain_key: None, + }; + let ops = vec![ + Operand { + keys: vec![key(0, 0), key(1, 1)], + factor: Some(ColumnIndex(2)), + }, + Operand { + keys: vec![key(2, 0), key(1, 1)], + factor: Some(ColumnIndex(2)), + }, + ]; + let mut producer = Fold::try_new( + vec![EQ; 3], + ops, + vec![out(0), out(2)], + RowValue::Product, + Aggregate::sum(), + None, + ) + .unwrap(); + let comparisons = [ + (Comparison::Less, "<"), + (Comparison::LessOrEqual, "<="), + (Comparison::Greater, ">"), + (Comparison::GreaterOrEqual, ">="), + ]; + let mut condition = String::new(); + if rng.chance(0.7) { + let (comparison, op) = comparisons[rng.below(4)]; + let (left, right, sql) = match rng.chance(0.5) { + true => (0, 2, format!(" AND q.t {op} k.t")), + false => (2, 0, format!(" AND k.t {op} q.t")), + }; + condition = sql; + let c = Condition { + left: DimensionId(left), + comparison, + right: DimensionId(right), + }; + producer = producer.with_conditions(vec![c]).unwrap(); + } + let (scale, scale_sql) = match rng.below(3) { + 0 => (Scale::One, String::new()), + 1 => (Scale::Times(-0.7), " * -0.7".to_string()), + _ => (Scale::Over(8.0), " / 8.0".to_string()), + }; + let (shift, weight) = match rng.chance(0.7) { + true => (Shift::KeptMaximum, "exp(s - MAX(s) OVER (PARTITION BY i))"), + false => (Shift::None, "exp(s)"), + }; + let normalizer = match rng.chance(0.5) { + true => Normalizer::KeptGroups, + false => Normalizer::JoinedRows, + }; + let (values, outputs, c_col, c_group) = match channels { + Some(_) => ( + Operand { + keys: vec![key(2, 0), key(3, 1)], + factor: Some(ColumnIndex(2)), + }, + vec![DimensionId(0), DimensionId(3)], + ", v.c", + ", v.c", + ), + None => ( + Operand { + keys: vec![key(2, 0)], + factor: Some(ColumnIndex(1)), + }, + vec![DimensionId(0)], + "", + "", + ), + }; + let own = if channels.is_some() { vec![EQ] } else { vec![] }; + let pipeline = + Pipeline::try_new(producer, scale, shift, own, values, outputs, normalizer).unwrap(); + let scores = format!( + "WITH scores AS (SELECT q.t AS i, k.t AS j, SUM(q.x * k.x){scale_sql} AS s \ + FROM q JOIN k ON q.d = k.d{condition} GROUP BY q.t, k.t), \ + weights AS (SELECT i, j, {weight} AS w FROM scores)" + ); + let sql = match normalizer { + Normalizer::KeptGroups => format!( + "{scores}, norm AS (SELECT i, SUM(w) AS z FROM weights GROUP BY i) \ + SELECT weights.i{c_col}, SUM(weights.w * v.x) / norm.z \ + FROM weights JOIN v ON weights.j = v.t JOIN norm ON weights.i = norm.i \ + GROUP BY weights.i{c_group}, norm.z" + ), + _ => format!( + "{scores} SELECT weights.i{c_col}, SUM(weights.w * v.x) / SUM(weights.w) \ + FROM weights JOIN v ON weights.j = v.t GROUP BY weights.i{c_group}" + ), + }; + Case { + q, + k, + v, + pipeline, + sql, + complete, + v_scale, + } +} + +/// Rows as (keys, value), sorted by keys. +fn rows(batches: &[RecordBatch]) -> Vec<(Vec, Option)> { + let mut out = Vec::new(); + for b in batches { + let keys = b.num_columns() - 1; + let value = b + .column(keys) + .as_any() + .downcast_ref::() + .unwrap(); + for r in 0..b.num_rows() { + let k: Vec = (0..keys) + .map(|c| { + b.column(c) + .as_any() + .downcast_ref::() + .unwrap() + .value(r) + }) + .collect(); + out.push((k, value.is_valid(r).then(|| value.value(r)))); + } + } + out.sort_by(|a, b| a.0.cmp(&b.0)); + out +} + +const HOST: Host = Host { + float_keys: FloatKeys::SignedZerosEqual, +}; + +/// The kernel's rows, from the tables in a random order and batches, on a +/// random number of threads, or `None` if it declined. +fn run(case: &Case, rng: &mut Rng) -> Option> { + let batches = |t: &Table, rng: &mut Rng| -> Vec { t.partitions(rng).concat() }; + let (q, k, v) = ( + batches(&case.q, rng), + batches(&case.k, rng), + batches(&case.v, rng), + ); + let schemas = [ + case.q.schema.as_ref(), + case.k.schema.as_ref(), + case.v.schema.as_ref(), + ]; + let schema = einfold_kernels::pipeline_schema(&case.pipeline, &schemas).unwrap(); + let options = Options { + host: HOST, + threads: 1 + rng.below(12), + ..Options::default() + }; + let inputs = [ + Relation::new(&case.q.schema, &q), + Relation::new(&case.k.schema, &k), + Relation::new(&case.v.schema, &v), + ]; + let piped = einfold_kernels::pipeline(&case.pipeline, inputs, schema, options).unwrap(); + piped.map(|p| p.batches) +} + +#[tokio::test] +async fn pipeline_matches_sql() { + let seed: u64 = std::env::var("EINFOLD_SEED") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(0xa77e_0001); + let cases: u64 = match std::env::var("EINFOLD_CASES") { + Ok(n) => n.parse().unwrap(), + Err(_) if std::env::var("EINFOLD_SEED").is_ok() => 1, + Err(_) => 150, + }; + for i in 0..cases { + let seed = seed.wrapping_add(i); + let mut rng = Rng::new(seed); + let case = case(&mut rng); + let fail = |what: String| panic!("seed {seed}: {what}\nSQL: {}", case.sql); + let ours = run(&case, &mut rng); + if case.complete && ours.is_none() { + fail("declined complete inputs".into()); + } + let Some(ours) = ours else { continue }; + let again = run(&case, &mut rng).unwrap(); + let bits = |r: Vec<(Vec, Option)>| -> Vec<_> { + r.into_iter() + .map(|(k, v)| (k, v.map(f64::to_bits))) + .collect() + }; + let (ours, again) = (rows(&ours), rows(&again)); + if bits(ours.clone()) != bits(again) { + fail("two arrival orders and thread counts gave different bits".into()); + } + + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(4)); + for (name, t) in [("q", &case.q), ("k", &case.k), ("v", &case.v)] { + let mem = MemTable::try_new(Arc::clone(&t.schema), t.partitions(&mut rng)).unwrap(); + ctx.register_table(name, Arc::new(mem)).unwrap(); + } + let theirs = rows(&ctx.sql(&case.sql).await.unwrap().collect().await.unwrap()); + if ours.len() != theirs.len() { + fail(format!("{} rows, SQL has {}", ours.len(), theirs.len())); + } + let bound = 1e-9 * case.v_scale; + for ((ok, ov), (tk, tv)) in ours.iter().zip(&theirs) { + let same = match (ov, tv) { + (Some(x), Some(y)) if x.is_nan() || y.is_nan() => x.is_nan() && y.is_nan(), + (Some(x), Some(y)) if x.is_infinite() || y.is_infinite() => x == y, + (Some(x), Some(y)) => (x - y).abs() <= bound, + (None, None) => true, + _ => false, + }; + if ok != tk || !same { + fail(format!("row {ok:?} = {ov:?}, SQL has {tk:?} = {tv:?}")); + } + } + } +} + +/// The causal mask `q.t >= k.t` on a square leaves about half the pairs of +/// tiles unvisited. +#[test] +fn causal_masks_skip_half_the_tiles() { + let mut rng = Rng::new(7); + let l = 64 * 16; + let q = Table::complete(&["t", "d"], &[l, 8], || rng.unit()); + let k = Table::complete(&["t", "d"], &[l, 8], || 0.5); + let v = Table::complete(&["t"], &[l], || 1.0); + const EQ: Dimension = Dimension { + equality: KeyEquality::Equal, + }; + let key = |d, c| Key::new(DimensionId(d), ColumnIndex(c)); + let out = |d| Output { + dimension: DimensionId(d), + domain_key: None, + }; + let ops = vec![ + Operand { + keys: vec![key(0, 0), key(1, 1)], + factor: Some(ColumnIndex(2)), + }, + Operand { + keys: vec![key(2, 0), key(1, 1)], + factor: Some(ColumnIndex(2)), + }, + ]; + let causal = Condition { + left: DimensionId(0), + comparison: Comparison::GreaterOrEqual, + right: DimensionId(2), + }; + let producer = Fold::try_new( + vec![EQ; 3], + ops, + vec![out(0), out(2)], + RowValue::Product, + Aggregate::sum(), + None, + ) + .unwrap() + .with_conditions(vec![causal]) + .unwrap(); + let values = Operand { + keys: vec![key(2, 0)], + factor: Some(ColumnIndex(1)), + }; + let pipeline = Pipeline::try_new( + producer, + Scale::One, + Shift::KeptMaximum, + vec![], + values, + vec![DimensionId(0)], + Normalizer::KeptGroups, + ) + .unwrap(); + let batches = |t: &Table| t.partitions(&mut Rng::new(1)).concat(); + let (bq, bk, bv) = (batches(&q), batches(&k), batches(&v)); + let schemas = [q.schema.as_ref(), k.schema.as_ref(), v.schema.as_ref()]; + let schema = einfold_kernels::pipeline_schema(&pipeline, &schemas).unwrap(); + let inputs = [ + Relation::new(&q.schema, &bq), + Relation::new(&k.schema, &bk), + Relation::new(&v.schema, &bv), + ]; + let piped = einfold_kernels::pipeline(&pipeline, inputs, schema, Options::default()) + .unwrap() + .unwrap(); + // 16 × 16 pairs; those on and below the diagonal are visited. + assert_eq!((piped.tiles_visited, piped.tiles), (16 * 17 / 2, 256)); + // Every average of v.x = 1 is 1. + let rows = rows(&piped.batches); + assert_eq!(rows.len(), l); + assert!(rows.iter().all(|(_, x)| (x.unwrap() - 1.0).abs() < 1e-12)); +} + +/// A pipeline may leave one of its own dimensions out of its outputs, to be +/// summed away: `SELECT i, SUM(w * v.x) / norm.z ... GROUP BY i` over values +/// keyed by `(t, c)`. The fused algorithm writes a row for each position of +/// its own dimensions, so it declines such a pipeline, and the caller runs +/// SQL's plan instead. +#[test] +fn declines_a_pipeline_that_sums_its_own_dimension_away() { + let q = Table::complete(&["t", "d"], &[4, 2], || 0.5); + let k = Table::complete(&["t", "d"], &[4, 2], || 0.25); + let v = Table::complete(&["t", "c"], &[4, 3], || 1.0); + const EQ: Dimension = Dimension { + equality: KeyEquality::Equal, + }; + let key = |d, c| Key::new(DimensionId(d), ColumnIndex(c)); + let out = |d| Output { + dimension: DimensionId(d), + domain_key: None, + }; + let ops = vec![ + Operand { + keys: vec![key(0, 0), key(1, 1)], + factor: Some(ColumnIndex(2)), + }, + Operand { + keys: vec![key(2, 0), key(1, 1)], + factor: Some(ColumnIndex(2)), + }, + ]; + let producer = Fold::try_new( + vec![EQ; 3], + ops, + vec![out(0), out(2)], + RowValue::Product, + Aggregate::sum(), + None, + ) + .unwrap(); + let values = Operand { + keys: vec![key(2, 0), key(3, 1)], + factor: Some(ColumnIndex(2)), + }; + // Dimension 3, `v.c`, is the pipeline's own, and not an output. + let pipeline = Pipeline::try_new( + producer, + Scale::One, + Shift::KeptMaximum, + vec![EQ], + values, + vec![DimensionId(0)], + Normalizer::KeptGroups, + ) + .unwrap(); + let batches = |t: &Table| t.partitions(&mut Rng::new(1)).concat(); + let (bq, bk, bv) = (batches(&q), batches(&k), batches(&v)); + let schemas = [q.schema.as_ref(), k.schema.as_ref(), v.schema.as_ref()]; + let schema = einfold_kernels::pipeline_schema(&pipeline, &schemas).unwrap(); + let inputs = [ + Relation::new(&q.schema, &bq), + Relation::new(&k.schema, &bk), + Relation::new(&v.schema, &bv), + ]; + let piped = einfold_kernels::pipeline(&pipeline, inputs, schema, Options::default()).unwrap(); + assert!(piped.is_none()); +} + +/// With no queries or keys, attention has no rows, whatever the values: an +/// empty `q` and `k` give no positions of the query or of the feature, so no +/// rows, as SQL's inner joins give none. +#[test] +fn gives_no_rows_without_queries_or_keys() { + let q = Table::complete(&["t", "d"], &[0, 2], || 0.5); + let k = Table::complete(&["t", "d"], &[0, 2], || 0.25); + let v = Table::complete(&["t"], &[4], || 1.0); + const EQ: Dimension = Dimension { + equality: KeyEquality::Equal, + }; + let key = |d, c| Key::new(DimensionId(d), ColumnIndex(c)); + let out = |d| Output { + dimension: DimensionId(d), + domain_key: None, + }; + let ops = vec![ + Operand { + keys: vec![key(0, 0), key(1, 1)], + factor: Some(ColumnIndex(2)), + }, + Operand { + keys: vec![key(2, 0), key(1, 1)], + factor: Some(ColumnIndex(2)), + }, + ]; + let producer = Fold::try_new( + vec![EQ; 3], + ops, + vec![out(0), out(2)], + RowValue::Product, + Aggregate::sum(), + None, + ) + .unwrap(); + let values = Operand { + keys: vec![key(2, 0)], + factor: Some(ColumnIndex(1)), + }; + let pipeline = Pipeline::try_new( + producer, + Scale::One, + Shift::KeptMaximum, + vec![], + values, + vec![DimensionId(0)], + Normalizer::KeptGroups, + ) + .unwrap(); + let batches = |t: &Table| t.partitions(&mut Rng::new(1)).concat(); + let (bq, bk, bv) = (batches(&q), batches(&k), batches(&v)); + let schemas = [q.schema.as_ref(), k.schema.as_ref(), v.schema.as_ref()]; + let schema = einfold_kernels::pipeline_schema(&pipeline, &schemas).unwrap(); + let inputs = [ + Relation::new(&q.schema, &bq), + Relation::new(&k.schema, &bk), + Relation::new(&v.schema, &bv), + ]; + let piped = einfold_kernels::pipeline(&pipeline, inputs, schema, Options::default()).unwrap(); + // Declining is right too: the caller's SQL then gives no rows. + assert!(piped.is_none_or(|p| rows(&p.batches).is_empty())); +}