# Autoresearch: fast exact one-directional Chamfer geometry loss

Pure exploration (all code in this folder). Autoresearch loop, run manually.
Mirrors the `playground/autoresearch_pose_align/` workflow (which reached a 957x
fused-Triton champion, rel err 0.0, PASS).

## Goal
`L_geo` (paper sec. Geometric supervision) is a one-directional Chamfer loss that
pulls each predicted Gaussian mean toward its nearest back-projected GT surface
point. It is the single most expensive loss term in training (`geo/chamfer` =
~137 ms, 82% of `loss_fn`). **Minimize wall-clock time** of one fwd+bwd while the
result stays IDENTICAL to the production `chamferdist` KNN (same loss + gradient).

    chamfer[b] = mean_i  min_j  || source[b, i] - target[b, j] ||^2

## Autoresearch config
- **Scope (modifiable):** `playground/chamfer_opt/optimized.py` only (helpers may
  be added in this folder, e.g. a Triton kernel module).
- **Metric (minimize):** sum of fwd+bwd ms across the 2v/12v/24v sweep
  (`verify.py` -> `METRIC=`).
- **Guard (must all pass, else discard):**
  - identity: candidate loss + grad match `chamferdist` within fp32 tol
    (`fwd rtol 1e-4`, `grad rtol 1e-3`)
  - memory: candidate peak GPU mem (24v) <= reference peak * 2.0
- **Direction:** lower is better.
- **Baseline:** `chamferdist` brute-force KNN, measured once in `baseline.py` ->
  `baseline.json`.

## Operating points (REAL, measured from a zipsplat-v4 forward)
| tag | ctx views | N (source) | M (target) |
|---|---|---|---|
| 6v  | 6  | 62,208  | ~254,016 |
| 12v | 12 | 124,416 | ~254,016 |
| 24v | 24 | 248,832 | ~254,016 |

N = 10368 * ctx_views (scales). M ~= num_target_views(4) * 252 * 252 ~= 254k and
is CONSTANT -- the chamfer target is the 4 TARGET views' depth (`loss_on_context
= false`), NOT the context views. An earlier synthetic sweep wrongly scaled M
with views (M = views*252*252 -> 3-6x too large); `cache_real_samples.py`
captured the real tensors and corrected this.

## Thesis
Exact KNN, but beat the brute-force CUDA kernel. Pure-torch rewrites (cdist,
matmul-expansion) are bit-identical but 20-35x SLOWER -- they materialize
`[chunk, M]` blocks (bandwidth-bound), while the kernel keeps a running min in
registers (O(N) mem). So the win must be a better KERNEL or ALGORITHM:
- **Triton K=1, D=3 brute force** -- specialized, no general-K `mink` overhead,
  running min+argmin in SRAM, optimal tiling. (exact)
- **Spatial uniform-grid pruning** -- bucket targets into voxels; each source
  searches its cell + neighbors with exact ring-expansion/fallback so it returns
  the TRUE nearest neighbor. Surface points are clustered -> big win. (exact)
- **torch.compile** the differentiable gather + squared-distance tail.
Explicitly NOT used (yet): bf16 / approximate search (separate later knob).

The gradient is always exact: search (no_grad) -> NN index -> autograd tail
`(source - nn)^2`, whose grad `2(source - nn)/N` matches the kernel.

## Setup
```
conda run -n splatfactory --no-capture-output python -m playground.chamfer_opt.baseline
```
Measures the chamferdist baseline -> `baseline.json`.

## Loop protocol (per iteration)
1. Review `results.tsv` + `git log` of this folder.
2. Make ONE change to `optimized.py` (a CONFIG knob or a new search method).
3. `conda run -n splatfactory ... python -m playground.chamfer_opt.verify`.
4. Keep if `METRIC` < champion AND `GUARD=PASS`; else surgically revert.
5. Append a row to `results.tsv`; `git add playground/chamfer_opt/ && git commit`.

GPU hygiene: time on an otherwise-idle GPU; deployment target is GH200 (96 GB),
development/measurement here is the local 4090 (24 GB).

## Results
### Phase A: synthetic sizes (M wrongly scaled with views -- see above)
| iter | method | total ms | speedup | identity | guard |
|---|---|---|---|---|---|
| 0 | chamferdist (reference) | 592.11 | 1.0x | ref | ref |
| 1 | triton_brute bn64 bm128 | 1996.22 | 0.30x | ok | PASS (slow, untuned) |
| 2 | block-tuning UP (bn 128->512) | -- | -- | ok | wrong way (spills) |
| 3 | **triton_brute bn16 bm512 w2** | **330.58** | **1.79x** | ok | PASS (champion) |

### Phase B: REAL sizes (M~254k const) + REAL geometry (real_clouds.pt)
sum over 6v/12v/24v, fwd+bwd:
| method | total ms | speedup | identity | mem | guard |
|---|---|---|---|---|---|
| chamferdist (reference) | 138.71 | 1.0x | ref | 11.0 MB | ref |
| triton_brute bn16 bm512 w2 (synthetic-tuned) | 75.64 | 1.83x | ok | 11.9 MB | PASS |
| **triton_brute bn8 bm512 w1 (real-tuned)** | **59.49** | **2.33x** | **ok** | **11.9 MB** | **PASS** |

Per-config on REAL clouds (test_real.py), bn8/bm512/w1: 6v 2.45-2.52x /
12v 2.28-2.44x / 24v 2.33-2.47x, all bit-identical. Re-tuning at the correct
M~254k (tune_real.py) found num_warps=1 + bn8 win (smaller M -> less work/program
-> fewer warps maximize occupancy); num_stages>1 hurts (more SRAM -> less
occupancy, consistent with occupancy-bound).

## Spatial grid: investigated and REJECTED (analyze_grid.py)
A uniform-grid prune could cut the reduction-bound work, but the data kills it:
- **Trained head** (real_clouds.pt): median Gaussian sees 122 candidates vs 254k
  (379x) -- BUT ~1% are stray (dist 1.9-4.5, outside the ~2.0 bbox), worst case
  136 rings -> mandatory brute-force fallback; candidate p99~3700 -> divergence.
- **Init head** (real_clouds_init.pt, untrained = start of training): the
  Gaussians are a tight blob FAR from the surface, so NN dist is large (~0.27) ->
  cell size large -> ~2489 targets/cell -> a 3x3x3 lookup grabs ~150k of 254k =
  **2x reduction (useless)**.
- The grid's benefit is INVERSELY correlated with when the chamfer matters: ~2x
  at init (chamfer critical), ~380x at convergence (chamfer least needed). Brute
  force is distribution-independent -> 2.33x holds at init AND convergence.
- DECISION: keep the brute-force champion (bn8/bm512/w1), drop the grid.

## Key findings
- Pure-torch exact KNN is 20-35x slower than chamferdist (cdist/mm @12v:
  2.5-4 s vs 118 ms); identical results (diff 0.0) but bandwidth-bound.
- No KNN lib installed (faiss/keops/pytorch3d/torch_cluster); Triton is, and is
  the sanctioned path (no install).
- **The whole cost is the SEARCH kernel.** torch.profiler @12v: chamferdist
  forward `KNearestNeighborKernelV3<float,3,1>` = 119.6 ms (99.97%), backward
  11 us. The no_grad-search + gather + squared-dist autograd tail is ~0.2 ms
  (free). So only the KNN matters.
- **chamferdist is a SPECIALIZED kernel** (`V3<float,3,1>`, i.e. D=3 K=1), not
  generic cdist -- already purpose-built for our case. Hard to beat.
- **Brute-force Triton CAN beat V3 -- with tiny BLOCK_N.** The kernel is
  spill/occupancy-bound, NOT memory-bound: bigger BLOCK_N is WORSE (bn128 549 ms
  -> bn512 15900 ms @12v search, register spill to local mem). Going the other
  way wins -- bn=16, bm=512, 2 warps = 65 ms @12v search (vs V3 116 ms). Small
  tile (16x512=32KB) -> no spill + high occupancy; wide bm -> few argmin
  reductions; 2 warps -> many concurrent programs. Champion: 1.79x overall,
  identical, less memory. (Roofline "maximize reuse via big BLOCK_N" reasoning
  was wrong here -- the limiter is occupancy, not target bandwidth.)
- ncu is NOT installed here (broken stub); used torch.profiler + roofline. nsys
  is available if a timeline is needed.
- **Champion stands: fp32 Triton bn16/bm512/w2 = 1.79x, identical, less memory.**
  A further identical win would be algorithmic (spatial-grid pruning:
  N x M -> N x local); subsampling the 762k-1.5M target points is the trivial
  big win but is NOT bit-identical (stochastic surface estimate).
- **Hybrid low-prec search + fp32 backward (test_precision.py).** The KNN search
  (assignment) only needs to pick a near-tie neighbor; the gather + (s-nn)^2
  backward runs fp32, so loss/grad are fp32-exact for the chosen neighbor. On
  REAL clouds, search-dtype vs accuracy:
    fp16: NN disagree 5-7%, loss <0.02%, grad cos **0.998-0.9999**  (sweet spot)
    bf16: NN disagree 30-33%, loss <0.6%, grad cos 0.985-0.999
  fp16 >> bf16 because coords are scene-normalized (in fp16 range) so fp16's 10
  mantissa bits beat bf16's 7. (Synthetic randn was worst-case: bf16 cos 0.874.)
- **But reduced precision buys NO speed on this kernel -- even re-tuned for it.**
  fp16/bf16 search times are within noise of fp32. Re-tuning the WHOLE block grid
  in fp16 (tune_real.py --dtype fp16) returns the SAME best config (bn8/bm512/w1)
  at 1.02x (noise). The kernel is occupancy-bound (tiny BLOCK_N, 1 warp = minimal
  footprint already), NOT memory-bound, so halving the 16-bit load bytes does
  nothing. Champion stays fp32 + identical. Low prec would only pay off in a
  memory-bound regime (GH200's 4x BW, or much larger M) -- and there fp16 (not
  bf16) is the choice.
- Block config is hardware-specific -> `tune.py` caches the best (bn,bm,warps)
  per `torch.cuda.get_device_name()`; re-run once on GH200 (Hopper: 2.3x SRAM,
  4x bandwidth -> optimum shifts; V3 also rebalances, so the 1.79x may move).
