diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4aad4d5..ecf7d94 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -86,6 +86,7 @@ jobs: cmake -B build -DEDGEVECTOR_BUILD_EXAMPLES=ON cmake --build build ./build/quickstart + ./build/itq_pipeline asan-ubsan: name: ASan + UBSan diff --git a/.gitignore b/.gitignore index 8580716..8f379d1 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,9 @@ tests/wsl_results/ # Example build artifacts examples/quickstart examples/quickstart.exe +examples/itq_pipeline +examples/itq_pipeline.exe +examples/*.evrt examples/*.evec examples/*.evhg build/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 87468bb..30828d8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -35,9 +35,11 @@ install(DIRECTORY include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) option(EDGEVECTOR_BUILD_EXAMPLES "Build the EdgeVector examples" OFF) if(EDGEVECTOR_BUILD_EXAMPLES) - add_executable(quickstart examples/quickstart.cpp) - target_link_libraries(quickstart PRIVATE edgevector::edgevector) - # -march=native activates POPCNT/AVX2 (x86) or NEON codegen; correctness - # does not depend on it, speed does. - target_compile_options(quickstart PRIVATE -O2 -march=native -Wall -Wextra) + foreach(example quickstart itq_pipeline) + add_executable(${example} examples/${example}.cpp) + target_link_libraries(${example} PRIVATE edgevector::edgevector) + # -march=native activates POPCNT/AVX2 (x86) or NEON codegen; correctness + # does not depend on it, speed does. + target_compile_options(${example} PRIVATE -O2 -march=native -Wall -Wextra) + endforeach() endif() diff --git a/examples/Makefile b/examples/Makefile index 11daab0..97e5614 100644 --- a/examples/Makefile +++ b/examples/Makefile @@ -14,10 +14,15 @@ endif quickstart$(EXE): quickstart.cpp $(wildcard ../include/edgevector/*.hpp) $(CXX) $(CXXFLAGS) -I../include quickstart.cpp -o $@ -run: quickstart$(EXE) +itq_pipeline$(EXE): itq_pipeline.cpp $(wildcard ../include/edgevector/*.hpp) + $(CXX) $(CXXFLAGS) -I../include itq_pipeline.cpp -o $@ + +run: quickstart$(EXE) itq_pipeline$(EXE) $(RUNNER) ./quickstart$(EXE) + $(RUNNER) ./itq_pipeline$(EXE) clean: - rm -f quickstart$(EXE) quickstart.evec quickstart.evhg + rm -f quickstart$(EXE) itq_pipeline$(EXE) quickstart.evec quickstart.evhg \ + itq_example.evrt .PHONY: run clean diff --git a/examples/itq_pipeline.cpp b/examples/itq_pipeline.cpp new file mode 100644 index 0000000..85fcffd --- /dev/null +++ b/examples/itq_pipeline.cpp @@ -0,0 +1,168 @@ +// ============================================================================ +// EdgeVector ITQ pipeline example: when your embeddings have a decaying +// spectrum (real embeddings always do), a learned rotation makes every bit +// of the binary code informative. +// +// Build and run (from examples/): make run (runs quickstart too) +// or via CMake (from the repo root): +// cmake -B build -DEDGEVECTOR_BUILD_EXAMPLES=ON && cmake --build build +// ./build/itq_pipeline +// +// The program synthesizes anisotropic data (energy concentrated in the +// leading dimensions), trains an ItqRotation, builds two indexes over the +// SAME floats - raw sign codes vs rotated codes - and measures recall +// against exact float32 ground truth for both. It checks its own results +// (the rotated index must win) and exits non-zero otherwise; CI runs it. +// +// The two properties that make ITQ practical here: +// 1. Rotation-only (no centering/PCA), so cosine is preserved EXACTLY - +// ground truth and exact re-ranking are unchanged. +// 2. Deterministic per seed - the same training data gives bit-identical +// rotations on every platform. +// ============================================================================ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace edgevector; + +static int g_failures = 0; +static void expect(bool ok, const char* what) { + std::printf(" %s %s\n", ok ? "ok " : "FAIL", what); + if (!ok) ++g_failures; +} + +int main() { + std::printf("EdgeVector %s ITQ pipeline example\n\n", EDGEVECTOR_VERSION); + + // ------------------------------------------------------------------ + // 0. Anisotropic embedding-like data: clustered, with dimension d + // scaled by exp(-4d/dim) - a decaying spectrum. Sign quantization + // alone wastes most of its bits on the low-energy tail. + // ------------------------------------------------------------------ + const std::size_t dim = 256; + const std::uint32_t n = 5000; + const std::size_t n_queries = 40; + const std::uint32_t k = 10; + const std::size_t rb = padded_bytes(dim); + + std::mt19937 rng(42); + std::normal_distribution gauss(0.0f, 1.0f); + std::vector centers(60 * dim); + for (float& c : centers) c = gauss(rng); + std::vector scale(dim); + for (std::size_t d = 0; d < dim; ++d) + scale[d] = std::exp(-4.0f * static_cast(d) / + static_cast(dim)); + + const std::size_t total = n + n_queries; + std::vector floats(total * dim); + for (std::size_t i = 0; i < total; ++i) { + const float* c = centers.data() + (i % 60) * dim; + for (std::size_t d = 0; d < dim; ++d) + floats[i * dim + d] = (c[d] + 0.5f * gauss(rng)) * scale[d]; + } + const float* qfloats = floats.data() + static_cast(n) * dim; + + // ------------------------------------------------------------------ + // 1. Train the rotation on the corpus (offline; deterministic). + // Persist it next to your index files - queries need the same G. + // ------------------------------------------------------------------ + // A 2k-row subsample and 10 iterations train in seconds and capture the + // spectrum fine; production rotations can afford the full corpus. + ItqRotation rot(dim); + expect(rot.train(floats.data(), 2000, dim, /*iterations=*/10) == + ItqStatus::ok, + "rotation trained"); + std::printf(" quantization error %.3f -> %.3f, orthogonality %.1e\n", + rot.error_before_training(), rot.error_after_training(), + rot.orthogonality_error()); + expect(rot.save("itq_example.evrt") == ItqStatus::ok, "rotation saved"); + ItqRotation loaded(dim); + expect(loaded.load("itq_example.evrt") == ItqStatus::ok, + "rotation reloaded (validated: magic, length, orthogonality)"); + + // ------------------------------------------------------------------ + // 2. Two code blocks over the SAME floats: raw and rotated. + // ------------------------------------------------------------------ + std::vector raw_w((rb / 8) * n), itq_w((rb / 8) * n); + auto* raw_codes = reinterpret_cast(raw_w.data()); + auto* itq_codes = reinterpret_cast(itq_w.data()); + std::vector scratch(dim); + for (std::uint32_t i = 0; i < n; ++i) { + quantize(floats.data() + i * dim, dim, raw_codes + i * rb); + loaded.rotate_quantize(floats.data() + i * dim, scratch.data(), + itq_codes + i * rb); + } + + HNSWGraph raw_graph(raw_codes, rb, dim, n); + HNSWGraph itq_graph(itq_codes, rb, dim, n); + std::vector ids(n); + std::iota(ids.begin(), ids.end(), 0u); + expect(raw_graph.insert_batch(ids.data(), n, 0) == n, "raw index built"); + expect(itq_graph.insert_batch(ids.data(), n, 0) == n, "ITQ index built"); + + // ------------------------------------------------------------------ + // 3. Query both with asymmetric re-ranking; measure recall against + // exact float32 cosine ground truth. Cosine is rotation-invariant, + // so ONE ground truth serves both indexes. + // ------------------------------------------------------------------ + std::vector qrot(dim); + std::vector qw(rb / 8); + auto* qbits = reinterpret_cast(qw.data()); + ScoredResult res[10]; + std::vector> gt(n); + + std::size_t raw_hits = 0, itq_hits = 0; + for (std::size_t qi = 0; qi < n_queries; ++qi) { + const float* q = qfloats + qi * dim; + for (std::uint32_t i = 0; i < n; ++i) + gt[i] = {-cosine_similarity_f32(q, floats.data() + i * dim, dim), + i}; + std::partial_sort(gt.begin(), gt.begin() + k, gt.end()); + + quantize(q, dim, qbits); // raw pipeline: quantize the query as-is + std::uint32_t found = raw_graph.search_reranked(qbits, q, k, 100, res); + for (std::uint32_t i = 0; i < found; ++i) + for (std::size_t t = 0; t < k; ++t) + if (res[i].id == gt[t].second) { ++raw_hits; break; } + + loaded.rotate(q, qrot.data()); // ITQ pipeline: rotate, then quantize + quantize(qrot.data(), dim, qbits); + found = itq_graph.search_reranked(qbits, qrot.data(), k, 100, res); + for (std::uint32_t i = 0; i < found; ++i) + for (std::size_t t = 0; t < k; ++t) + if (res[i].id == gt[t].second) { ++itq_hits; break; } + } + const double denom = static_cast(n_queries * k); + const double raw_recall = raw_hits / denom; + const double itq_recall = itq_hits / denom; + std::printf(" recall@10 vs float32 truth: raw %.3f, ITQ %.3f\n", + raw_recall, itq_recall); + expect(itq_recall > raw_recall, "rotated codes beat raw codes"); + + // ------------------------------------------------------------------ + // 4. Exact re-ranking over the ROTATED index uses the ORIGINAL floats + // and the ORIGINAL query: orthogonal rotations preserve cosine. + // ------------------------------------------------------------------ + const float* q = qfloats; // first query again + loaded.rotate(q, qrot.data()); + quantize(qrot.data(), dim, qbits); + const std::uint32_t found = itq_graph.search_exact_reranked( + qbits, q, floats.data(), dim, k, 100, res); + expect(found == k, "exact re-rank over the rotated index works " + "with unrotated floats"); + + std::remove("itq_example.evrt"); + std::printf("\n%s\n", g_failures == 0 ? "All ITQ example checks passed." + : "ITQ EXAMPLE FAILURES - see above."); + return g_failures == 0 ? 0 : 1; +}