reference_index_file.cuh

Go to the documentation of this file.

#pragma once
 
#include <cuddl/detail/hash.cuh>
#include <cuddl/reference_database_file.cuh>
#include <cuddl/reference_index.cuh>
 
#include <algorithm>
#include <atomic>
#include <thread>
#include <vector>
 
namespace cuddl::detail {
 
struct reference_index_digest {
    uint64_t digest = 0;
 
    Result<void> bytes(void const* data, size_t size) {
        digest = xxhash64(static_cast<uint8_t const*>(data), size, digest);
        return Ok();
    }
 
    template <typename T>
    Result<void> value(T item) {
        item = database_file_little_endian(item);
        return bytes(&item, sizeof(item));
    }
 
    template <typename T>
    Result<void> words(std::span<T const> items) {
        if constexpr (std::endian::native == std::endian::little) {
            return bytes(items.data(), items.size_bytes());
        } else {
            std::vector<T> little_endian(items.begin(), items.end());
            for (auto& item : little_endian) {
                item = database_file_little_endian(item);
            }
            return bytes(little_endian.data(), items.size_bytes());
        }
    }
};
 
inline Result<uint64_t> reference_database_digest(reference_database_file const& db) {
    reference_index_digest hash;
    CUDDL_TRY(hash.bytes("CUDDLDB-XXH64-v1", sizeof("CUDDLDB-XXH64-v1") - 1));
    auto const metadata = db.metadata();
    auto const& c = metadata.compatibility;
    CUDDL_TRY(hash.value(c.kmer_length));
    CUDDL_TRY(hash.value(c.bucket_count));
    CUDDL_TRY(hash.value(c.indexed_bucket_count));
    CUDDL_TRY(hash.value(c.score_encoder_identity));
    CUDDL_TRY(hash.value(c.exponent_bits));
    CUDDL_TRY(hash.value(c.mantissa_bits));
    CUDDL_TRY(hash.value(c.hash_identity));
    CUDDL_TRY(hash.value(c.hash_seed));
    CUDDL_TRY(hash.value(c.canonicalisation_policy));
    CUDDL_TRY(hash.value(c.blacklist_identity));
    CUDDL_TRY(hash.value(c.blacklist_version));
    CUDDL_TRY(hash.value(c.key_mask));
    CUDDL_TRY(hash.value(metadata.reference_count));
    CUDDL_TRY(hash.value(static_cast<uint64_t>(db.names().size())));
    for (auto const& name : db.names()) {
        CUDDL_TRY(hash.value(static_cast<uint64_t>(name.size())));
        CUDDL_TRY(hash.bytes(name.data(), name.size()));
    }
    CUDDL_TRY(hash.words(db.rows()));
    return hash.digest;
}
 
}  // namespace cuddl::detail
 
namespace cuddl {
 
class reference_index_file {
   public:
    template <uint32_t K, size_t BucketCount, typename Layout>
    [[nodiscard]] static Result<void> save(
        reference_index<K, BucketCount, Layout> const& index,
        reference_database<K, BucketCount, Layout> const& source,
        std::filesystem::path const& path,
        cuda::stream_ref stream
    ) try {
        return [&]() -> Result<void> {
            detail::device_scope const scope{stream};
            CUDDL_TRY(source.view(&index));
            auto database = CUDDL_TRY(snapshot(source, stream));
            auto storage = index.storage();
            auto offsets = CUDDL_TRY(
                download(index.index_offsets_.data(), index.index_offsets_.size(), stream)
            );
            auto posting_count =
                storage == index_storage::dense ? offsets.back() : index.index_posting_capacity_;
            auto postings =
                CUDDL_TRY(download(index.index_postings_.data(), posting_count, stream));
            auto keys =
                CUDDL_TRY(download(index.index_keys_.data(), index.index_keys_.size(), stream));
            auto digest = CUDDL_TRY(detail::reference_database_digest(database));
            detail::database_file_writer writer;
            writer.temporary = path.string() + ".tmp.XXXXXX";
            auto fd = ::mkstemp(writer.temporary.data());
            if (fd < 0) {
                return Err(Error::resource("cannot create index file beside: " + path.string()));
            }
            writer.output = ::fdopen(fd, "wb");
            if (!writer.output) {
                ::close(fd);
                return Err(Error::resource("cannot open temporary index file"));
            }
            CUDDL_TRY(writer.bytes("CUDDLIX\0", 8));
            CUDDL_TRY(writer.value(uint32_t{1}));
            CUDDL_TRY(writer.value(storage == index_storage::dense ? uint32_t{0} : uint32_t{1}));
            CUDDL_TRY(writer.value(digest));
            CUDDL_TRY(writer.value(static_cast<uint64_t>(postings.size())));
            CUDDL_TRY(writer.value(index.pair_work_));
            CUDDL_TRY(writer.words(offsets));
            CUDDL_TRY(writer.words(postings));
            if constexpr (std::endian::native == std::endian::little) {
                CUDDL_TRY(writer.bytes(keys.data(), keys.size() * sizeof(uint16_t)));
            } else {
                for (auto key : keys) {
                    CUDDL_TRY(writer.value(key));
                }
            }
            CUDDL_TRY(writer.value(static_cast<uint32_t>(writer.checksum)));
            if (std::fclose(std::exchange(writer.output, nullptr)) != 0) {
                return Err(Error::resource("cannot close index file"));
            }
            std::filesystem::rename(writer.temporary, path);
            writer.temporary.clear();
            return Ok();
        }();
    } catch (cuda::cuda_error const& error) {
        return Err(Error::cuda(static_cast<cudaError_t>(error.status())));
    } catch (std::system_error const& error) {
        return Err(Error::resource(error.what()));
    } catch (std::bad_alloc const& error) {
        return Err(Error::resource(error.what()));
    }
 
    template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
    [[nodiscard]] static Result<reference_index<K, BucketCount, Layout>> load(
        std::filesystem::path const& path,
        reference_database<K, BucketCount, Layout> const& database,
        cuda::stream_ref stream
    ) {
        detail::device_scope const scope{stream};
        auto source = CUDDL_TRY(snapshot(database, stream));
        auto decoded = CUDDL_TRY(read(path, std::move(source)));
        auto index = CUDDL_CUDA_TRY((reference_index<K, BucketCount, Layout>(database, stream)));
        index.index_offsets_ = CUDDL_CUDA_TRY(
            cuda::make_device_buffer<uint32_t>(stream, stream.device(), decoded.offsets_)
        );
        index.index_keys_ = CUDDL_CUDA_TRY(
            cuda::make_device_buffer<uint16_t>(stream, stream.device(), decoded.keys_)
        );
        index.index_posting_capacity_ = static_cast<size_t>(detail::indexed_posting_count(
            database.reference_count(), database.metadata().compatibility
        ));
        index.index_postings_ = CUDDL_CUDA_TRY(
            cuda::make_device_buffer<uint32_t>(
                stream, stream.device(), index.index_posting_capacity_, cuda::no_init
            )
        );
        if (!decoded.postings_.empty()) {
            CUDDL_CUDA_TRY(
                cuda::copy_bytes(
                    stream,
                    cuda::std::span{decoded.postings_.data(), decoded.postings_.size()},
                    cuda::std::span{index.index_postings_.data(), decoded.postings_.size()}
                )
            );
        }
        // The driver may still read the pageable host sources after the copies return.
        CUDDL_CUDA_TRY(stream.sync());
        CUDDL_TRY(index.build_key_directory(
            database.reference_count(),
            database.metadata().compatibility.indexed_bucket_count,
            stream
        ));
        index.indexed_ = true;
        index.pair_work_ = decoded.pair_work_;
        {
            auto const reference_count = database.reference_count();
            auto const indexed_buckets = database.metadata().compatibility.indexed_bucket_count;
            index.pair_fraction_ = reference_count == 0U || indexed_buckets == 0U
                                       ? 0.0
                                       : static_cast<double>(index.pair_work_) /
                                             (static_cast<double>(reference_count) *
                                              reference_count * indexed_buckets);
        }
        CUDDL_TRY(index.build_posting_bitmaps(database, stream));
        return index;
    }
 
   private:
    template <uint32_t K, size_t BucketCount, typename Layout>
    static Result<reference_database_file>
    snapshot(reference_database<K, BucketCount, Layout> const& source, cuda::stream_ref stream) {
        CUDDL_TRY((detail::validate_non_indexed_score_compatibility<K, BucketCount, Layout>(
            source.metadata().compatibility
        )));
        reference_database_file database;
        database.metadata_ = source.metadata();
        database.names_.assign(source.names().begin(), source.names().end());
        if (database.names_.size() != source.reference_count()) {
            return Err(Error::invalid_argument("index files require named reference rows"));
        }
        auto scores = CUDDL_CUDA_TRY(
            cuda::make_device_buffer<uint16_t>(
                stream,
                stream.device(),
                static_cast<size_t>(source.reference_count()) * BucketCount,
                cuda::no_init
            )
        );
        CUDDL_TRY(source.copy_scores_async({scores.data(), scores.size()}, stream));
        database.rows_ = CUDDL_TRY(download(scores.data(), scores.size(), stream));
        return database;
    }
 
    reference_index_file() = delete;
 
    struct decoded_index {
        reference_database_file database_;
        index_storage storage_;
        unsigned long long pair_work_{};
        std::vector<uint32_t> offsets_{};
        std::vector<uint32_t> postings_{};
        std::vector<uint16_t> keys_{};
    };
 
    [[nodiscard]] static Result<decoded_index>
    read(std::filesystem::path const& path, reference_database_file database) try {
        return [&]() -> Result<decoded_index> {
            detail::database_file_reader reader{
                std::ifstream(path, std::ios::binary | std::ios::ate)
            };
            auto size = reader.input.tellg();
            if (size < 0) {
                return Err(Error::resource("cannot open index file: " + path.string()));
            }
            reader.remaining = static_cast<uint64_t>(size);
            reader.input.seekg(0);
            char magic[8];
            CUDDL_TRY(reader.bytes(magic, sizeof(magic)));
            if (std::string_view(magic, 8) != std::string_view("CUDDLIX\0", 8)) {
                return Err(Error::invalid_argument("not a cuDDL reference index file"));
            }
            uint32_t version{}, kind{};
            CUDDL_TRY(reader.value(version));
            CUDDL_TRY(reader.value(kind));
            if (version != 1 || kind > 1) {
                return Err(Error::invalid_argument("unsupported reference index format"));
            }
            uint64_t digest{};
            CUDDL_TRY(reader.value(digest));
            auto const source = database.metadata();
            if (database.rows().size() !=
                static_cast<uint64_t>(source.reference_count) * source.compatibility.bucket_count) {
                return Err(
                    Error::invalid_argument("reference database has no complete stored rows")
                );
            }
            if (digest != CUDDL_TRY(detail::reference_database_digest(database))) {
                return Err(
                    Error::invalid_argument("reference index belongs to a different database")
                );
            }
            uint64_t count{};
            CUDDL_TRY(reader.value(count));
            unsigned long long pair_work{};
            CUDDL_TRY(reader.value(pair_work));
            auto const metadata = database.metadata();
            auto const capacity =
                detail::indexed_posting_count(metadata.reference_count, metadata.compatibility);
            auto const offset_count =
                kind == 0 ? detail::indexed_cell_count(metadata.compatibility) + 1 : 0;
            auto const key_count = kind == 1 ? capacity : 0;
            if (capacity > std::numeric_limits<uint32_t>::max() || count > capacity ||
                (kind == 1 && count != capacity) ||
                offset_count > std::numeric_limits<size_t>::max() / sizeof(uint32_t) ||
                capacity > std::numeric_limits<size_t>::max() / sizeof(uint32_t) ||
                reader.remaining != (offset_count + count) * sizeof(uint32_t) +
                                        key_count * sizeof(uint16_t) + sizeof(uint32_t)) {
                return Err(Error::invalid_argument("invalid reference index extents"));
            }
            decoded_index result{
                std::move(database), kind == 0 ? index_storage::dense : index_storage::sparse
            };
            result.pair_work_ = pair_work;
            result.offsets_.resize(static_cast<size_t>(offset_count));
            result.postings_.resize(static_cast<size_t>(count));
            result.keys_.resize(static_cast<size_t>(key_count));
            CUDDL_TRY(reader.words(result.offsets_));
            CUDDL_TRY(reader.words(result.postings_));
            CUDDL_TRY(reader.bytes(result.keys_.data(), result.keys_.size() * sizeof(uint16_t)));
            if constexpr (std::endian::native != std::endian::little) {
                for (auto& key : result.keys_) {
                    key = detail::database_file_little_endian(key);
                }
            }
            auto expected = static_cast<uint32_t>(reader.checksum);
            uint32_t checksum{};
            CUDDL_TRY(reader.value(checksum));
            if (checksum != expected) {
                return Err(Error::invalid_argument("reference index checksum mismatch"));
            }
            CUDDL_TRY(validate_payload(result));
            return result;
        }();
    } catch (std::bad_alloc const& error) {
        return Err(Error::resource(error.what()));
    }
 
    template <typename T>
    static Result<std::vector<T>> download(T const* data, size_t size, cuda::stream_ref stream) {
        std::vector<T> host(size);
        if (size != 0) {
            CUDDL_CUDA_TRY(
                cuda::copy_bytes(
                    stream, cuda::std::span{data, size}, cuda::std::span{host.data(), host.size()}
                )
            );
            CUDDL_CUDA_TRY(stream.sync());
        }
        return host;
    }
 
    static Result<void> validate_payload(decoded_index const& index) {
        auto const& offsets_ = index.offsets_;
        auto const& postings_ = index.postings_;
        auto const& keys_ = index.keys_;
        auto const m = index.database_.metadata();
        auto const& c = m.compatibility;
        auto const rows = index.database_.rows();
        auto const sparse = index.storage_ == index_storage::sparse;
        auto const key_count = static_cast<size_t>(c.key_mask) + 1;
        if (!sparse) {
            if (offsets_.front() != 0 || offsets_.back() != postings_.size()) {
                return Err(Error::invalid_argument("invalid dense index endpoints"));
            }
        }
        // Each bucket either passes or names why it fails; seen[id] == bucket marks an ID
        // already listed in that bucket.
        auto const check_bucket =
            [&](uint32_t bucket, std::vector<uint32_t>& seen, size_t& indexed) -> char const* {
            auto const score_of = [&](uint32_t id) {
                return rows[static_cast<size_t>(id) * c.bucket_count + bucket];
            };
            if (sparse) {
                auto start = static_cast<size_t>(bucket) * m.reference_count;
                for (uint32_t i = 0; i < m.reference_count; ++i) {
                    auto pos = start + i;
                    auto id = postings_[pos];
                    if (id >= m.reference_count || seen[id] == bucket ||
                        (i != 0 && keys_[pos] < keys_[pos - 1]) ||
                        (i != 0 && keys_[pos] == keys_[pos - 1] && id <= postings_[pos - 1])) {
                        return "invalid sparse index postings";
                    }
                    seen[id] = bucket;
                    auto score = score_of(id);
                    auto key =
                        score == 0 || c.key_mask == 0xffffU ? score : (score & c.key_mask) + 1U;
                    if (keys_[pos] != key) return "sparse index key does not match reference";
                }
                return nullptr;
            }
            for (uint32_t id = 0; id < m.reference_count; ++id) indexed += score_of(id) != 0;
            for (size_t key = 0; key < key_count; ++key) {
                auto cell = static_cast<size_t>(bucket) * key_count + key;
                auto begin = offsets_[cell], end = offsets_[cell + 1];
                if (begin > end || end > postings_.size()) return "invalid dense index offsets";
                for (auto pos = begin; pos < end; ++pos) {
                    auto id = postings_[pos];
                    if (id >= m.reference_count || seen[id] == bucket ||
                        (pos != begin && id <= postings_[pos - 1])) {
                        return "invalid dense index postings";
                    }
                    seen[id] = bucket;
                    auto score = score_of(id);
                    if (score == 0 || (score & c.key_mask) != key) {
                        return "dense index key does not match reference";
                    }
                }
            }
            return nullptr;
        };
        auto const workers = std::max(
            1U, std::min<unsigned>(std::thread::hardware_concurrency(), c.indexed_bucket_count)
        );
        std::atomic<uint32_t> next{0};
        std::atomic<size_t> indexed{0};
        std::atomic<char const*> failure{nullptr};
        {
            std::vector<std::jthread> pool;
            for (unsigned worker = 0; worker < workers; ++worker) {
                pool.emplace_back([&] {
                    std::vector<uint32_t> seen(
                        m.reference_count, std::numeric_limits<uint32_t>::max()
                    );
                    size_t local = 0;
                    for (uint32_t bucket; failure.load(std::memory_order_relaxed) == nullptr &&
                                          (bucket = next++) < c.indexed_bucket_count;) {
                        if (auto const error = check_bucket(bucket, seen, local)) {
                            char const* none = nullptr;
                            failure.compare_exchange_strong(none, error);
                        }
                    }
                    indexed += local;
                });
            }
        }
        if (auto const error = failure.load()) return Err(Error::invalid_argument(error));
        if (!sparse && indexed != postings_.size()) {
            return Err(Error::invalid_argument("dense index omits reference scores"));
        }
        return Ok();
    }
};
 
}  // namespace cuddl