omnigraph/research/lance-autoresearch/program.md

# Lance PQ L2 kernel research — agent instructions

You are an autonomous research assistant. Your job is to improve the kernel(s) in
`src/kernels.rs` so that `cargo run --release --bin run_experiment` reports a
**lower `geomean_ns_per_query`** while keeping **`recall_at_10` within 0.005 of
the seeded baseline** (and never below the hard floor 0.50).

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/bin/run_experiment.rs`
   - `src/fixture.rs`
2. Confirm fixtures are present. SIFT1M lives under `~/.cache/lance-autoresearch/`.
   If it's missing, the bench will fall back to a deterministic synthetic dataset
   — that's fine for the loop; mention it in your log. If you want SIFT1M, run
   `bash scripts/prepare_fixtures.sh` (one-time, ~5–10 min, ~250 MB download).
3. Ensure `results.tsv` exists. If not, create it with this header line:
   ```
   commit	timestamp	source	num_base	recall_at_10	geomean_ns_per_query	peak_mem_mb	total_seconds	keep	description
   ```
4. Run the baseline trial: `cargo run --release --bin run_experiment > run.log 2>&1`.
   Parse `run.log` and append a row to `results.tsv` with `keep=baseline`,
   `description="seeded scalar PQ-L2 baseline"`. This is your reference number.
5. Commit the baseline row with a one-line message like `baseline: <numbers>`.

## What you CAN do

- Modify **`src/kernels.rs`** freely. You may:
  - Reorder loops, change iteration order over codes or sub-vectors.
  - Switch to SIMD via `std::arch` (`x86_64::_mm256_*`, `aarch64::neon::*`),
    behind `#[cfg(target_arch = "...")]` gates. Always keep a portable scalar
    fallback so the kernel compiles everywhere.
  - Reshape internal data: transpose the codebook, pack the distance LUT into
    `u8`/`u16` for `pshufb`-style lookup, group codes for SIMD gather.
  - Use `unsafe` if needed; document the invariants you're relying on.
  - Mark hot functions `#[inline]` or split them; add private helpers freely.
- Add `#[cfg(test)] mod tests { ... }` inside `src/kernels.rs` if you want
  property checks against the scalar path.

## What you CANNOT do

- Do **not** modify `src/lib.rs` (changes `DIM` / `NUM_SUB_VECTORS` / `NUM_CENTROIDS` /
  `TOP_K` — these pin the fixture geometry).
- Do **not** modify `src/bin/run_experiment.rs`, `src/reference.rs`, `src/fixture.rs`,
  `benches/pq_l2.rs`, `scripts/prepare_fixtures.sh`, or `Cargo.toml`.
- Do **not** add new crate dependencies (the bench's external surface is intentionally
  minimal — only `anyhow`, plus `criterion` as a dev-dep).
- Do **not** delete or alter the public API of `kernels.rs`:
  - `pub type DistanceTable`
  - `pub fn compute_distance_table_l2(query: &[f32], codebook: &[f32]) -> DistanceTable`
  - `pub fn probe_pq_l2_top_k(table: &DistanceTable, codes: &[u8], num_vectors: usize, out: &mut TopKHeap)`
  - `pub struct TopKHeap` with `new() / push / into_sorted`

## The metric

Minimize `geomean_ns_per_query` (geometric mean of per-query wall-clock from the
benched queries, rounded to a u64 ns) subject to:

1. `recall_at_10 >= baseline_recall_at_10 - 0.005`
2. `recall_at_10 >= 0.50` (hard floor; below this the bench exits non-zero)
3. `total_seconds <= 600`
4. Build is clean: `cargo build --release` succeeds, `cargo clippy --release -- -D warnings`
   reports zero issues. (Run `cargo clippy --release` before each commit.)

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

## Lance-PQ-specific priors

These are the directions known to pay off on this kernel shape. Don't pursue all
of them at once — pick one hypothesis, implement, measure, decide.

- **Codebook layout for the table-build step.** 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 (8 floats). An `[m][d][k]` transpose can let
  you SIMD-load 8 `(query - centroid)` lanes across `d` and broadcast over `k`.
- **LUT packing for the probe step.** The probe is dominated by `acc +=
  table[m][codes[off+m]]` × 16. Two well-known tricks:
  - Pack each `table[m]` row into 256 × `f16` or 256 × `u8` (quantized post-build)
    to fit the LUT in cache and enable `vpgatherdq` / `pshufb`.
  - Reorder code storage to `[m][i]` (transpose codes by sub-quantizer) so each
    `m` step is a contiguous gather over up to 32 vectors at once.
- **Top-K integration.** `push()` does a branch + heap sift on every code; for a
  1M-row probe this is the second-biggest cost after the gather. Consider:
  - Skip the heap entirely when the running `acc` is already `> current_max`
    (early termination, but only if your accumulator order makes that cheap).
  - Block the probe (e.g., 1024 codes at a time), find the local top-K with a
    branchless scan, 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 1M scale where codes don't all fit in L2.
- **FMA in the table build.** The diff–square–sum sequence is
  `(q - c)·(q - c)` per element — that's `(q*q) - 2qc + c*c`. You can hoist
  `q*q` once per sub-vector and precompute `c*c` once at codebook-load time
  (if you cache it as a side table), reducing the inner loop to one FMA.
  But: caching `c*c` requires a one-time setup step, which has to live in
  `kernels.rs` since you cannot touch the fixture; either lazy-init via
  `OnceLock<Vec<f32>>` or rebuild every call (probably not worth it).

## 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 recall.
2. **Edit `src/kernels.rs`.** Keep the diff focused on the one hypothesis.
3. **Build and lint.** Run:
   ```
   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 `recall_at_10`, `geomean_ns_per_query`,
   `peak_mem_mb`, `total_seconds` from `run.log`. Compute the deltas vs. baseline.
6. **Decide keep or revert.**
   - **Keep** iff: recall within tolerance, speed strictly better than the
     last-kept row (allow ~1% noise band), and total time within budget.
   - **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>	<source>	<num_base>	<recall>	<geomean_ns>	<peak_mem>	<elapsed>	<keep|revert>	<one-line description>
   ```
8. **Commit.** Use a one-line message describing the change and the headline
   number, e.g. `transpose codebook; 184k → 142k ns/query (recall 0.94)`.

## 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.