omnigraph/research/lance-autoresearch/program.md

# Lance PQ L2 kernel research — agent instructions

You are an autonomous research assistant. Your job is to improve `src/kernels.rs`
so that `cargo run --release --bin run_experiment` reports a **lower
`geomean_ns_per_query`** while:

1. The **correctness phase passes** — your kernel's distance values must match the
   scalar reference within `MAX_ABS_ERR = 1e-4`, and the top-K must be
   tie-tolerant equivalent on every input the bench generates.
2. The `worst_ns_per_query` does **not regress more than 5%** against the
   last-kept kernel — if you win on one (shape × distribution) and lose
   significantly on another, the change isn't a generalizable improvement.

This bench is intentionally **dataset-independent**: there is no fixed dataset.
The correctness oracle is mathematical equivalence to the scalar reference,
checked across multiple PQ shapes and synthetic input distributions
(Gaussian / uniform / sparse / large-dynamic-range / mostly-zero). The speed
oracle is the geomean across multiple shapes × distributions, with worst-case
guarded. A win that depends on a specific data distribution or PQ shape will
fail to clear the bar by construction.

Read this file end-to-end before doing anything else. Then run setup, then the loop.

## Setup (do once at the start of every session)

1. Read these files, in this order:
   - `README.md`
   - `program.md` (this file)
   - `src/lib.rs`
   - `src/kernels.rs` *(the only file you may edit)*
   - `src/reference.rs`
   - `src/inputs.rs`
   - `src/bin/run_experiment.rs`
2. Ensure `results.tsv` exists. If not, create it with this header line:
   ```
   commit	timestamp	correctness	geomean_ns	worst_ns	worst_combo	best_ns	best_combo	peak_mem_mb	total_seconds	keep	description
   ```
3. Run the baseline trial: `cargo run --release --bin run_experiment > run.log 2>&1`.
   Confirm `correctness: pass`. Parse `run.log` and append a row to `results.tsv`
   with `keep=baseline` and `description="seeded scalar PQ-L2 baseline"`. This
   is your reference number.
4. Commit the baseline row with a one-line message like `baseline: <numbers>`.

## What you CAN do

- Modify **`src/kernels.rs`** freely. You may:
  - Pre-process the codebook in `PqKernel::new` (transpose layouts, cache
    `c·c` for the FMA trick, pack the codebook for register-resident lookup,
    etc.). This cost is paid once per dataset and amortized across queries —
    the bench measures per-query, not per-(build + query).
  - Reorder loops, switch internal data layouts, drop down to `std::arch`
    intrinsics under `#[cfg(target_arch = ...)]` gates. **Always keep a
    portable scalar fallback** so the kernel compiles everywhere.
  - Use `unsafe` if needed; document the invariants you're relying on.
  - Mark hot functions `#[inline]`; add private helpers freely.
  - Add `#[cfg(test)] mod tests { ... }` inside `src/kernels.rs` if you want
    in-file property checks.

## What you CANNOT do

- Do **not** modify `src/lib.rs` (`PqShape` and the tolerance constants are
  shared with the immutable scaffolding).
- Do **not** modify `src/bin/run_experiment.rs`, `src/reference.rs`,
  `src/inputs.rs`, `benches/pq_l2.rs`, or `Cargo.toml`.
- Do **not** add new crate dependencies.
- Do **not** alter the public API of `kernels::PqKernel`:
  - `PqKernel::new(shape: PqShape, codebook: &[f32]) -> Self`
  - `PqKernel::shape(&self) -> &PqShape`
  - `PqKernel::distance_table(&self, query: &[f32]) -> Vec<f32>`
  - `PqKernel::probe_top_k(&self, table: &[f32], codes: &[u8], num_vectors: usize, k: usize) -> Vec<(u32, f32)>`
- Do **not** introduce lossy techniques (LUT u8/u16 quantization, asymmetric-
  distance approximation, etc.) — the correctness phase asserts exact-up-to-ε
  match against the scalar reference. If you want to explore a lossy track,
  surface that in a separate kernel and propose a track extension.

## The metric

Minimize `geomean_ns_per_query` (geometric mean of per-query wall-clock across
all timed queries, all shapes, all distributions) subject to:

1. Correctness phase: **pass** (exit-2 otherwise).
2. `worst_ns_per_query` ≤ 1.05 × the last-kept kernel's worst.
3. `total_seconds` ≤ 600.
4. Build is clean: `cargo build --release` succeeds, `cargo clippy --release
   --all-targets -- -D warnings` reports zero issues.

Ties break toward simpler code. If two kernels report the same speed within
~3% noise, prefer fewer lines / less `unsafe`.

## Lance-PQ-specific priors (lossless directions)

These directions are known to pay off without compromising arithmetic accuracy.
Pick one hypothesis at a time; implement; measure; decide.

- **Codebook layout.** The reference layout is `[m][k][d]`. For a fixed query,
  iterating over centroids stays in cache, but the inner loop over `d` is
  short. Transposing to `[m][d][k]` lets you SIMD-load 8 `(query - centroid)`
  lanes across `d` and broadcast over `k`. Do the transpose in `PqKernel::new`
  once.
- **Cache `c·c`.** The diff–square–sum is `(q - c)·(q - c) = q·q - 2qc + c·c`.
  Hoist `q·q` per sub-vector, precompute `c·c` once at codebook-load time.
  Inner loop becomes one FMA (`-2qc + cc`). Watch the sign / accumulator
  ordering so the rounding stays within tolerance.
- **Probe layout.** The probe is dominated by `acc += table[m][codes[off+m]]`
  × `num_sub_vectors`. Transposing codes to `[m][i]` (one row per sub-quantizer,
  contiguous over base index) lets you process up to 32+ vectors per inner
  iteration with `vpgatherdq`-style loads.
- **Top-K integration.** `push()` does a branch + heap sift on every code.
  At 50k probes per query × 9 (shape × dist) combos that's the second-biggest
  cost after the gather. Block the probe (e.g., 512 codes at a time), find the
  local top-K with a branchless pass, then merge into the global heap.
- **Prefetch.** A `_mm_prefetch(codes.as_ptr().add(off + 64), _MM_HINT_T0)`
  ahead of the gather is usually pure win at 50k+ scale where codes don't all
  fit in L2.
- **FMA chains for table build.** The diff–square–sum maps cleanly to FMA on
  AVX2/NEON. Even without intrinsics, structuring the inner loop so `rustc`
  emits FMA helps.
- **Avoid the `Vec` allocation in the hot path.** `distance_table` allocates a
  fresh `Vec<f32>` per call. Returning a fixed-capacity buffer is a public-API
  change you can't make — but you can reuse a thread-local scratch buffer
  internally if it speeds the build.

## The loop

Once setup is done, repeat indefinitely:

1. **Observe state.** Read the last ~5 rows of `results.tsv`. Note which ideas
   have been tried, what won, what regressed. Form a hypothesis with one
   sentence stating the change and the predicted effect on speed and
   correctness.
2. **Edit `src/kernels.rs`.** Keep the diff focused on the one hypothesis.
3. **Build and lint.**
   ```
   cargo build --release
   cargo clippy --release --all-targets -- -D warnings
   ```
   If either fails, fix and try again — do not commit broken state.
4. **Run the trial.**
   ```
   cargo run --release --bin run_experiment > run.log 2>&1
   ```
5. **Parse the result.** Extract `correctness`, `geomean_ns_per_query`,
   `worst_ns_per_query` (with combo), `peak_mem_mb`, `total_seconds`. Compute
   deltas vs. baseline.
6. **Decide keep or revert.**
   - **Keep** iff: `correctness: pass`, geomean strictly better than the
     last-kept row (allow ~1% noise band), and `worst_ns_per_query` ≤ 1.05 ×
     last-kept's worst.
   - **Revert** otherwise: `git restore src/kernels.rs` (or commit and
     `git revert` if you want the revert in history). Note what failed.
7. **Log.** Append one row to `results.tsv`:
   ```
   <short_sha>	<iso8601>	<correctness>	<geomean_ns>	<worst_ns>	<worst_combo>	<best_ns>	<best_combo>	<peak_mem>	<elapsed>	<keep|revert>	<one-line description>
   ```
8. **Commit.** One-line message describing the change and the headline number,
   e.g. `transpose codebook in new(); 18.2k → 14.1k geomean ns (worst -8%)`.

## Hygiene

- Always commit `src/kernels.rs` changes; never commit `results.tsv` or
  `run.log` (they're gitignored).
- If a change fails to build, do not commit. Iterate until it builds, or
  revert cleanly.
- If two consecutive ideas regress, take a beat: re-read the last ~10 rows of
  `results.tsv` and update your mental model before proposing the next.
- Per-trial cap: 10 minutes. If `cargo run` is still going after 10 min, kill it
  and mark the trial as `timeout`.

## Never stop

Keep going until interrupted. Each loop iteration is one hypothesis, one edit,
one measurement, one commit. No multi-step plans across iterations.