From 0b6433bad270aa6687a68948227d235629336865 Mon Sep 17 00:00:00 2001 From: amkram Date: Mon, 13 Jul 2026 14:13:07 -0700 Subject: [PATCH] perf(index): optional uncompressed index that is mmap'd on load (no decompression pass) Adds --index-uncompressed: stores the raw capnp payload after the param header instead of ZSTD frames, so IndexReader mmaps the file and hands the bytes straight to capnp (zero-copy) rather than decompressing the whole index into a heap buffer on every load. The header's byte 26 records the format; old and compressed indexes (flag 0) are unaffected, so this is fully backward compatible and off by default. A/B (page cache warm, byte-identical placement at -t1): TB place -t8 2.13 -> 1.69s (-21%) -t16 1.82 -> 1.41s (-23%) SARS place -t8 1.49 -> 1.39s (-7%) -t16 1.31 -> 1.25s (-5%) Bigger at higher thread counts: placement compute is fast there, so the eliminated decompression pass is a larger fraction and page faults parallelize across threads. Neutral at -t1. Trade-off: larger on disk (TB 230->462MB, SARS 8->30MB) and cold first-load reads more bytes, so it targets repeated placement against a resident index -- hence opt-in. Placement output is byte-identical at -t1; the -t8 last-digit wobble is pre-existing parallel float-reduction non-determinism (the compressed index alone gives different values run-to-run). unit + e2e + examples pass. --- src/index_single_mode.cpp | 30 +++++++++++++++++++++------- src/index_single_mode.hpp | 7 ++++++- src/main.cpp | 41 ++++++++++++++++++++++++++++++--------- 3 files changed, 61 insertions(+), 17 deletions(-) diff --git a/src/index_single_mode.cpp b/src/index_single_mode.cpp index b79d4a16..6497ca9f 100644 --- a/src/index_single_mode.cpp +++ b/src/index_single_mode.cpp @@ -1535,6 +1535,7 @@ std::array encodeIndexHeader(const IndexParamsHeader& put32(20, static_cast(p.l)); h[24] = p.hpc ? 1 : 0; h[25] = p.open ? 1 : 0; + h[26] = p.uncompressed ? 1 : 0; return h; } @@ -1550,11 +1551,13 @@ bool readIndexHeader(const std::string& path, IndexParamsHeader& out) { out.l = static_cast(get32(20)); out.hpc = h[24] != 0; out.open = h[25] != 0; + out.uncompressed = h[26] != 0; return true; } } // namespace index_single_mode -void index_single_mode::IndexBuilder::writeIndex(const std::string& path, int numThreads, int zstdLevel) { +void index_single_mode::IndexBuilder::writeIndex(const std::string& path, int numThreads, int zstdLevel, + bool uncompressed) { output::step("Serializing index..."); kj::Array flatArray = capnp::messageToFlatArray(outMessage); @@ -1572,14 +1575,27 @@ void index_single_mode::IndexBuilder::writeIndex(const std::string& path, int nu ph.l = indexBuilder.getL(); ph.hpc = indexBuilder.getHpc(); ph.open = indexBuilder.getOpen(); + ph.uncompressed = uncompressed; const auto header = encodeIndexHeader(ph); - // Many independent 64 MB frames so the index inflates in parallel on load. - constexpr size_t kIndexFrameSize = 64ull * 1024 * 1024; - if (!panmap_zstd::compressToFile(data, dataSize, path, zstdLevel, numThreads, kIndexFrameSize, - header.data(), header.size())) { - output::error("failed to write compressed index to {}", path); - std::exit(1); + if (uncompressed) { + // Store the raw capnp bytes so load() can mmap the file and hand them straight to + // capnp (zero-copy, no decompression pass) -- bigger on disk, faster to load. + std::ofstream out(path, std::ios::binary | std::ios::trunc); + out.write(reinterpret_cast(header.data()), header.size()); + out.write(reinterpret_cast(data), static_cast(dataSize)); + if (!out) { + output::error("failed to write uncompressed index to {}", path); + std::exit(1); + } + } else { + // Many independent 64 MB frames so the index inflates in parallel on load. + constexpr size_t kIndexFrameSize = 64ull * 1024 * 1024; + if (!panmap_zstd::compressToFile(data, dataSize, path, zstdLevel, numThreads, kIndexFrameSize, + header.data(), header.size())) { + output::error("failed to write compressed index to {}", path); + std::exit(1); + } } output::done("Index written to " + path); diff --git a/src/index_single_mode.hpp b/src/index_single_mode.hpp index aa6c0d55..9b686410 100644 --- a/src/index_single_mode.hpp +++ b/src/index_single_mode.hpp @@ -31,6 +31,10 @@ constexpr size_t kIndexHeaderSize = 32; struct IndexParamsHeader { int32_t k = 0, s = 0, t = 0, l = 0; bool hpc = false, open = false; + // When set, the payload after the header is raw (unframed) capnp bytes rather than + // ZSTD frames, so load can mmap it and hand the bytes straight to capnp (zero-copy, + // no decompression pass). Bigger on disk, faster to load. + bool uncompressed = false; }; std::array encodeIndexHeader(const IndexParamsHeader& p); // Reads the header from the start of `path`. Returns false if absent / not this format. @@ -211,7 +215,8 @@ class IndexBuilder { // Compute 4x4 substitution spectrum from tree mutations and store in index void computeSubstitutionSpectrum(); - void writeIndex(const std::string& path, int numThreads = 0, int zstdLevel = 7); + void writeIndex(const std::string& path, int numThreads = 0, int zstdLevel = 7, + bool uncompressed = false); // Parameter getters (for testing) int getK() const { return k_; } diff --git a/src/main.cpp b/src/main.cpp index cce7eca3..55331463 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -126,6 +127,7 @@ struct Config { // Metagenomic options bool indexPacked = false; bool readPacked = false; + bool indexUncompressed = false; bool noProgress = false; size_t topOc = 1000; uint32_t maskReads = 0; @@ -196,22 +198,40 @@ struct Config { class IndexReader : public ::capnp::MessageReader { public: std::vector data; + // Held for the uncompressed path; the mmap must outlive `reader` (declared after it, + // so it is destroyed first), which points directly into the mapped bytes. + boost::iostreams::mapped_file_source mmapFile; std::unique_ptr<::capnp::FlatArrayMessageReader> reader; explicit IndexReader(const std::string& path, int numThreads = 0) : ::capnp::MessageReader(makeOptions()) { - // New-format indexes carry a small uncompressed param header before the frames; + // New-format indexes carry a small uncompressed param header before the payload; // skip it. Old-format indexes (no header) decompress from offset 0. index_single_mode::IndexParamsHeader ph; - const size_t dataOffset = - index_single_mode::readIndexHeader(path, ph) ? index_single_mode::kIndexHeaderSize : 0; - if (!panmap_zstd::decompressFromFile(path, data, numThreads, dataOffset)) { - throw std::runtime_error("Failed to decompress index: " + path); + const bool hasHeader = index_single_mode::readIndexHeader(path, ph); + const size_t dataOffset = hasHeader ? index_single_mode::kIndexHeaderSize : 0; + + const capnp::word* words = nullptr; + size_t numWords = 0; + if (hasHeader && ph.uncompressed) { + // Uncompressed index: mmap and hand the bytes straight to capnp (zero-copy, + // no decompression pass). dataOffset (32) is 8-aligned within the page-aligned + // mapping, so the word cast is properly aligned. + mmapFile.open(path); + if (!mmapFile.is_open()) { + throw std::runtime_error("Failed to mmap index: " + path); + } + words = reinterpret_cast(mmapFile.data() + dataOffset); + numWords = (mmapFile.size() - dataOffset) / sizeof(capnp::word); + } else { + if (!panmap_zstd::decompressFromFile(path, data, numThreads, dataOffset)) { + throw std::runtime_error("Failed to decompress index: " + path); + } + words = reinterpret_cast(data.data()); + numWords = data.size() / sizeof(capnp::word); } reader = std::make_unique<::capnp::FlatArrayMessageReader>( - kj::ArrayPtr(reinterpret_cast(data.data()), - data.size() / sizeof(capnp::word)), - makeOptions()); + kj::ArrayPtr(words, numWords), makeOptions()); } kj::ArrayPtr getSegment(uint id) override { return reader->getSegment(id); } @@ -417,7 +437,7 @@ bool buildIndex(const Config& cfg) { cfg.extentGuard); builder.buildIndexParallel(cfg.threads); builder.computeSubstitutionSpectrum(); - builder.writeIndex(cfg.index, cfg.threads, cfg.zstdLevel); + builder.writeIndex(cfg.index, cfg.threads, cfg.zstdLevel, cfg.indexUncompressed); auto ms = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); output::done("index", cfg.index, fmt::format("{} nodes", output::fmt_count(tg->trees[0].allNodes.size())), ms); @@ -1865,6 +1885,9 @@ int main(int argc, char** argv) { metagenomic.add_options()( "index-packed", po::bool_switch(&cfg.indexPacked), "Build packed capnp message (default false)")( "read-packed", po::bool_switch(&cfg.readPacked), "Read packed capnp message (default false)")( + "index-uncompressed", po::bool_switch(&cfg.indexUncompressed), + "Store the index uncompressed so it is mmap'd on load (no decompression pass; " + "larger on disk, faster to load; default false)")( "no-progress", po::bool_switch(&cfg.noProgress), "Disable progress bars"); po::options_description em("Metagenomic: EM");