reference_database.cuh

Go to the documentation of this file.

#pragma once
 
#include <cuda_runtime.h>
#include <cub/device/device_radix_sort.cuh>
#include <cub/device/device_scan.cuh>
#include <cub/device/device_segmented_sort.cuh>
#include <cub/device/device_select.cuh>
#include <cub/device/device_transform.cuh>
#include <cuda/algorithm>
#include <cuda/buffer>
#include <cuda/iterator>
#include <cuda/memory_pool>
#include <cuda/std/cstdint>
#include <cuda/std/functional>
#include <cuda/std/optional>
#include <cuda/std/utility>
#include <cuda/stream>
 
#include <algorithm>
#include <array>
#include <atomic>
#include <bit>
#include <cstddef>
#include <functional>
#include <limits>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <span>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <cuddl/detail/hash.cuh>
#include <cuddl/detail/kernels.cuh>
#include <cuddl/detail/pinned_pool.cuh>
#include <cuddl/detail/reference_deltas.cuh>
#include <cuddl/detail/register.cuh>
#include <cuddl/device_span.cuh>
#include <cuddl/error.hpp>
#include <cuddl/pairwise_counts.cuh>
 
namespace cuddl {
 
template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
class reference_index;
 
namespace detail {
using device_scope = ::cuda::__ensure_current_context;
}  // namespace detail
 
enum class index_storage { automatic, dense, sparse };
 
struct score_compatibility {
    uint8_t kmer_length{};
    uint16_t bucket_count{};
    uint16_t indexed_bucket_count{};
    uint32_t score_encoder_identity{};
    uint8_t exponent_bits{};
    uint8_t mantissa_bits{};
    uint32_t hash_identity{};
    uint64_t hash_seed{};
    uint8_t canonicalisation_policy{};
    uint64_t blacklist_identity{};
    uint8_t blacklist_version{};
    uint16_t key_mask{};
 
    template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
    [[nodiscard]] static constexpr score_compatibility current() noexcept {
        static_assert(K <= std::numeric_limits<uint8_t>::max());
        static_assert(BucketCount <= std::numeric_limits<uint16_t>::max());
        return {
            .kmer_length = K,
            .bucket_count = BucketCount,
            .indexed_bucket_count = BucketCount,
            .score_encoder_identity = 1U,
            .exponent_bits = Layout::exponent_bits,
            .mantissa_bits = Layout::mantissa_bits,
            .hash_identity = 1U,
            .hash_seed = detail::seed,
            .canonicalisation_policy = 1U,
            .blacklist_identity = 0U,
            .blacklist_version = 0U,
            // Genomic winner scores virtually never set the top bit, so folding it halves the
            // dense index at no practical cost in selectivity.
            .key_mask = 0x7fffU,
        };
    }
 
    friend bool operator==(score_compatibility const&, score_compatibility const&) = default;
};
 
struct reference_database_metadata {
    score_compatibility compatibility{};
    uint32_t reference_count{};
 
    friend bool operator==(reference_database_metadata const&, reference_database_metadata const&) =
        default;
};
 
struct reference_search_result {
    uint32_t reference_id{};
    pairwise_counts counts{};
 
    friend bool operator==(reference_search_result const&, reference_search_result const&) =
        default;
};
 
struct batch_search_result {
    uint32_t query_id{};
    uint32_t reference_id{};
    pairwise_counts counts{};
 
    friend bool operator==(batch_search_result const&, batch_search_result const&) = default;
};
 
struct batch_result_tile {
    packed_pairwise_counts const* results{};
    uint32_t const* match_counts{};
    uint32_t const* pass_bits{};
    uint32_t const* word_offsets{};
    uint32_t first_query_id{};
    uint32_t query_count{};
    uint32_t reference_count{};
    uint32_t bucket_count{};
    bool upper_triangle{};
 
    [[nodiscard]] __host__ __device__ constexpr uint32_t words_per_query() const noexcept {
        return (reference_count + 31U) / 32U;
    }
 
    [[nodiscard]] __host__ __device__ uint64_t slot_count() const noexcept {
        if (query_count == 0U) return 0U;
        auto const last = query_count - 1U;
        if (upper_triangle) {
            return first_query_id + last + 1U >= reference_count
                       ? slot(last, reference_count)
                       : slot(last, reference_count - 1U) + 1U;
        }
        return static_cast<uint64_t>(query_count) * reference_count;
    }
 
    [[nodiscard]] __host__ __device__ bool
    passed(uint32_t query_id, uint32_t reference_id) const noexcept {
        if (upper_triangle && reference_id < query_id) {
            cuda::std::swap(query_id, reference_id);
        }
        if (query_id < first_query_id || query_id - first_query_id >= query_count ||
            reference_id >= reference_count || (upper_triangle && reference_id <= query_id)) {
            return false;
        }
        if (pass_bits == nullptr) return true;
        return ((pass_bits[word(query_id, reference_id)] >> (reference_id % 32U)) & 1U) != 0U;
    }
 
    [[nodiscard]] __host__ __device__ cuda::std::optional<batch_search_result>
    find(uint32_t query_id, uint32_t reference_id) const noexcept {
        if (!passed(query_id, reference_id)) return cuda::std::nullopt;
        if (upper_triangle && reference_id < query_id) {
            auto const stored = results[position(reference_id, query_id)].unpack(bucket_count);
            return batch_search_result{query_id, reference_id, stored.mirrored()};
        }
        return result(query_id, reference_id, results[position(query_id, reference_id)]);
    }
 
    [[nodiscard]] __host__ __device__ uint32_t
    match_count(uint32_t query_id, uint32_t reference_id) const noexcept {
        if (upper_triangle && reference_id < query_id) {
            cuda::std::swap(query_id, reference_id);
        }
        return match_counts[position(query_id, reference_id)];
    }
 
    template <typename F>
    __host__ __device__ void for_each_passing(F&& f) const {
        for_each_passing(0U, query_count, f);
    }
 
    template <typename F>
    __host__ __device__ void for_each_passing(uint32_t begin, uint32_t end, F&& f) const {
        auto next = begin < end ? stored_before(begin) : 0U;
        for (uint32_t query = begin; query < end; ++query) {
            auto const query_id = first_query_id + query;
            if (pass_bits == nullptr) {
                // Every slot passes, and slots are in this order.
                for (auto reference = upper_triangle ? query_id + 1U : 0U;
                     reference < reference_count;
                     ++reference) {
                    f(result(query_id, reference, results[next++]));
                }
                continue;
            }
            auto const words = words_per_query();
            for (uint32_t w = 0; w < words; ++w) {
                for (auto bits = pass_bits[static_cast<size_t>(query) * words + w]; bits != 0U;
                     bits &= bits - 1U) {
                    auto const reference =
                        w * 32U + static_cast<uint32_t>(cuda::std::countr_zero(bits));
                    auto const at = word_offsets != nullptr ? next++ : slot(query, reference);
                    f(result(query_id, reference, results[at]));
                }
            }
        }
    }
 
    [[nodiscard]] __host__ __device__ uint64_t passing_before(uint32_t query) const noexcept {
        if (query >= query_count) return count();
        if (pass_bits == nullptr || word_offsets != nullptr) return stored_before(query);
        auto const words = static_cast<size_t>(query) * words_per_query();
        uint64_t total = 0U;
        for (size_t w = 0; w < words; ++w) {
            total += static_cast<uint64_t>(cuda::std::popcount(pass_bits[w]));
        }
        return total;
    }
 
    [[nodiscard]] __host__ __device__ uint64_t count() const noexcept {
        if (pass_bits == nullptr) return slot_count();
        auto const words = static_cast<size_t>(query_count) * words_per_query();
        if (words == 0U) return 0U;
        if (word_offsets != nullptr) {
            return static_cast<uint64_t>(word_offsets[words - 1U]) +
                   static_cast<uint64_t>(cuda::std::popcount(pass_bits[words - 1U]));
        }
        uint64_t total = 0U;
        for (size_t w = 0; w < words; ++w) {
            total += static_cast<uint64_t>(cuda::std::popcount(pass_bits[w]));
        }
        return total;
    }
 
   private:
    [[nodiscard]] __host__ __device__ batch_search_result
    result(uint32_t query_id, uint32_t reference_id, packed_pairwise_counts counts) const noexcept {
        return {
            .query_id = query_id,
            .reference_id = reference_id,
            .counts = counts.unpack(bucket_count),
        };
    }
 
    [[nodiscard]] __host__ __device__ size_t
    word(uint32_t query_id, uint32_t reference_id) const noexcept {
        return static_cast<size_t>(query_id - first_query_id) * words_per_query() +
               reference_id / 32U;
    }
 
    [[nodiscard]] __host__ __device__ uint64_t
    position(uint32_t query_id, uint32_t reference_id) const noexcept {
        if (word_offsets != nullptr) {
            auto const w = word(query_id, reference_id);
            auto const below = (1U << (reference_id % 32U)) - 1U;
            return static_cast<uint64_t>(word_offsets[w]) +
                   static_cast<uint64_t>(cuda::std::popcount(pass_bits[w] & below));
        }
        return slot(query_id - first_query_id, reference_id);
    }
 
    [[nodiscard]] __host__ __device__ uint64_t
    slot(uint32_t query, uint32_t reference_id) const noexcept {
        return upper_triangle ? detail::batch_result_slot<true>(
                                    query, reference_id, first_query_id, reference_count
                                )
                              : detail::batch_result_slot<false>(
                                    query, reference_id, first_query_id, reference_count
                                );
    }
 
    [[nodiscard]] __host__ __device__ uint64_t stored_before(uint32_t query) const noexcept {
        if (pass_bits == nullptr) {
            return slot(query, upper_triangle ? first_query_id + query + 1U : 0U);
        }
        if (word_offsets == nullptr || reference_count == 0U) return 0U;
        return word_offsets[static_cast<size_t>(query) * words_per_query()];
    }
};
 
class host_batch_results {
   public:
    explicit host_batch_results(cuda::stream_ref stream)
        : host_batch_results(
              stream,
              detail::device_pinned_pool(stream.device(), detail::pinned_placement::device_node)
          ) {}
 
    host_batch_results(cuda::stream_ref stream, cuda::pinned_memory_pool_ref pool)
        : pool_(pool),
          results_(stream, pool, 0, cuda::no_init),
          match_counts_(stream, pool, 0, cuda::no_init),
          pass_bits_(stream, pool, 0, cuda::no_init),
          word_offsets_(stream, pool, 0, cuda::no_init) {}
 
    host_batch_results(host_batch_results const&) = delete;
    host_batch_results& operator=(host_batch_results const&) = delete;
 
    __host__ host_batch_results(host_batch_results&& other) noexcept
        : pool_(other.pool_),
          results_(std::move(other.results_)),
          match_counts_(std::move(other.match_counts_)),
          pass_bits_(std::move(other.pass_bits_)),
          word_offsets_(std::move(other.word_offsets_)),
          count_(std::exchange(other.count_, 0U)),
          has_match_counts_(other.has_match_counts_),
          layout_(other.layout_) {}
 
    __host__ ~host_batch_results() {}
 
    [[nodiscard]] std::vector<batch_search_result> passing() const {
        std::vector<batch_search_result> kept;
        kept.reserve(static_cast<size_t>(count_));
        tile().for_each_passing([&](batch_search_result const& result) { kept.push_back(result); });
        return kept;
    }
 
    // A span into a temporary copy would dangle, so rvalues are rejected.
    [[nodiscard]] std::span<uint32_t const> passing_match_counts() const& noexcept {
        return {match_counts_.data(), has_match_counts_ ? static_cast<size_t>(count_) : 0U};
    }
    void passing_match_counts() const&& = delete;
 
    [[nodiscard]] batch_result_tile tile() const& noexcept {
        auto view = layout_;
        view.results = results_.data();
        view.match_counts = has_match_counts_ ? match_counts_.data() : nullptr;
        view.pass_bits = layout_.pass_bits == nullptr ? nullptr : pass_bits_.data();
        view.word_offsets = layout_.pass_bits == nullptr ? nullptr : word_offsets_.data();
        return view;
    }
    void tile() const&& = delete;
 
   private:
    template <typename T>
    using pinned_buffer = cuda::buffer<T, cuda::mr::host_accessible, cuda::mr::device_accessible>;
 
    template <typename T>
    void reserve(pinned_buffer<T>& buffer, size_t count, cuda::stream_ref stream) {
        if (buffer.size() < count) buffer = pinned_buffer<T>(stream, pool_, count, cuda::no_init);
    }
 
    friend Result<void> download(batch_result_tile const&, cuda::stream_ref, host_batch_results&);
 
    cuda::pinned_memory_pool_ref pool_;
    pinned_buffer<packed_pairwise_counts> results_;
    pinned_buffer<uint32_t> match_counts_;
    pinned_buffer<uint32_t> pass_bits_;
    pinned_buffer<uint32_t> word_offsets_;
    uint64_t count_{};
    bool has_match_counts_{};
    batch_result_tile layout_{};
};
 
[[nodiscard]] inline Result<void>
download(batch_result_tile const& tile, cuda::stream_ref stream, host_batch_results& host) {
    detail::device_scope const scope{stream};
    host.layout_ = tile;
    host.layout_.word_offsets = nullptr;
    host.has_match_counts_ = tile.match_counts != nullptr;
    if (tile.pass_bits == nullptr) {
        auto const slots = static_cast<size_t>(tile.slot_count());
        host.count_ = slots;
        host.reserve(host.results_, slots, stream);
        if (slots != 0U) {
            CUDDL_CUDA_TRY(
                cuda::copy_bytes(
                    stream,
                    cuda::std::span{tile.results, slots},
                    cuda::std::span{host.results_.data(), slots}
                )
            );
        }
        if (host.has_match_counts_) {
            host.reserve(host.match_counts_, slots, stream);
            if (slots != 0U) {
                CUDDL_CUDA_TRY(
                    cuda::copy_bytes(
                        stream,
                        cuda::std::span{tile.match_counts, slots},
                        cuda::std::span{host.match_counts_.data(), slots}
                    )
                );
            }
        }
        CUDDL_CUDA_TRY(stream.sync());
        return Ok();
    }
 
    auto const words = static_cast<size_t>(tile.query_count) * tile.words_per_query();
    host.reserve(host.pass_bits_, words, stream);
    host.reserve(host.word_offsets_, words, stream);
    if (words == 0U) {
        host.count_ = 0U;
        CUDDL_CUDA_TRY(stream.sync());
        return Ok();
    }
    auto const popcounts = cuda::make_transform_iterator(tile.pass_bits, detail::word_popcount{});
    size_t scan_bytes = 0;
    CUDDL_CUDA_TRY(
        cub::DeviceScan::ExclusiveSum(
            nullptr,
            scan_bytes,
            popcounts,
            host.word_offsets_.data(),
            static_cast<int64_t>(words),
            stream.get()
        )
    );
    auto scan_storage = CUDDL_CUDA_TRY(
        cuda::make_device_buffer<uint8_t>(stream, stream.device(), scan_bytes, cuda::no_init)
    );
    CUDDL_CUDA_TRY(
        cub::DeviceScan::ExclusiveSum(
            scan_storage.data(),
            scan_bytes,
            popcounts,
            host.word_offsets_.data(),
            static_cast<int64_t>(words),
            stream.get()
        )
    );
    CUDDL_CUDA_TRY(
        cuda::copy_bytes(
            stream,
            cuda::std::span{tile.pass_bits, words},
            cuda::std::span{host.pass_bits_.data(), words}
        )
    );
    CUDDL_CUDA_TRY(stream.sync());
    host.count_ = static_cast<uint64_t>(host.word_offsets_.data()[words - 1U]) +
                  static_cast<uint64_t>(cuda::std::popcount(host.pass_bits_.data()[words - 1U]));
    auto const count = static_cast<size_t>(host.count_);
    host.reserve(host.results_, count, stream);
    if (host.has_match_counts_) host.reserve(host.match_counts_, count, stream);
    if (count == 0U) {
        CUDDL_CUDA_TRY(stream.sync());
        return Ok();
    }
    // Passing records are gathered into bounded device staging, so each chunk crosses the bus as
    // one bulk copy rather than as scattered small writes.
    constexpr size_t staging_records = size_t{1} << 23U;
    auto const staged = std::min(count, staging_records);
    auto staging = CUDDL_CUDA_TRY(
        cuda::make_device_buffer<packed_pairwise_counts>(
            stream, stream.device(), staged, cuda::no_init
        )
    );
    auto staging_matches = CUDDL_CUDA_TRY(
        cuda::make_device_buffer<uint32_t>(
            stream, stream.device(), tile.match_counts == nullptr ? 0U : staged, cuda::no_init
        )
    );
    auto const kernel = tile.upper_triangle
                            ? detail::gather_passing_kernel<true, packed_pairwise_counts>
                            : detail::gather_passing_kernel<false, packed_pairwise_counts>;
    auto const* offsets = host.word_offsets_.data();
    for (size_t first_word = 0U; first_word < words;) {
        auto const base = offsets[first_word];
        // The last word whose results all fit the staging buffer ends the chunk.
        auto const last_word = static_cast<size_t>(
            std::upper_bound(offsets + first_word + 1U, offsets + words, base + staged - 32U) -
            offsets
        );
        auto const end = last_word == words ? count : static_cast<size_t>(offsets[last_word]);
        kernel<<<
            detail::warp_grid_blocks(last_word - first_word),
            detail::block_size,
            0,
            stream.get()>>>(
            tile.results,
            tile.match_counts,
            tile.pass_bits,
            host.word_offsets_.data(),
            tile.first_query_id,
            tile.reference_count,
            first_word,
            last_word,
            base,
            staging.data(),
            tile.match_counts == nullptr ? nullptr : staging_matches.data()
        );
        CUDDL_CUDA_TRY(cudaGetLastError());
        auto const records = end - base;
        if (records != 0U) {
            CUDDL_CUDA_TRY(
                cuda::copy_bytes(
                    stream,
                    cuda::std::span{staging.data(), records},
                    cuda::std::span{host.results_.data() + base, records}
                )
            );
            if (tile.match_counts != nullptr) {
                CUDDL_CUDA_TRY(
                    cuda::copy_bytes(
                        stream,
                        cuda::std::span{staging_matches.data(), records},
                        cuda::std::span{host.match_counts_.data() + base, records}
                    )
                );
            }
        }
        first_word = last_word;
    }
    CUDDL_CUDA_TRY(stream.sync());
    return Ok();
}
 
[[nodiscard]] inline Result<host_batch_results>
download(batch_result_tile const& tile, cuda::stream_ref stream) {
    host_batch_results host(stream);
    CUDDL_TRY(download(tile, stream, host));
    return Result<host_batch_results>::ok(std::move(host));
}
 
namespace detail {
 
inline constexpr size_t visit_chunk_records = size_t{1} << 20U;
 
struct visit_chunk {
    size_t first{};
    size_t last{};
    size_t first_record{};
    size_t end_record{};
};
 
}  // namespace detail
 
template <typename F>
[[nodiscard]] Result<void>
for_each_passing(batch_result_tile const& tile, cuda::stream_ref stream, F&& f) {
    detail::device_scope const scope{stream};
    using pinned_counts = cuda::
        buffer<packed_pairwise_counts, cuda::mr::host_accessible, cuda::mr::device_accessible>;
    auto const memory =
        detail::device_pinned_pool(stream.device(), detail::pinned_placement::device_node);
    constexpr auto chunk_records = detail::visit_chunk_records;
    auto const threshold = tile.pass_bits != nullptr;
    auto const words = static_cast<size_t>(tile.query_count) * tile.words_per_query();
 
    // Chunk plan, and for threshold tiles the host copies of the pass bits and word offsets.
    std::vector<detail::visit_chunk> chunks;
    std::vector<uint32_t> pass_bits;
    std::vector<uint32_t> offsets;
    cuda::device_buffer<uint32_t> device_offsets(
        stream, cuda::device_default_memory_pool(stream.device())
    );
    if (!threshold) {
        size_t record = 0U;
        detail::visit_chunk chunk{};
        for (uint32_t query = 0; query < tile.query_count; ++query) {
            auto const query_id = tile.first_query_id + query;
            auto const row = static_cast<size_t>(
                tile.upper_triangle
                    ? (tile.reference_count > query_id + 1U ? tile.reference_count - query_id - 1U
                                                            : 0U)
                    : tile.reference_count
            );
            if (record + row - chunk.first_record > chunk_records && query != chunk.first) {
                chunk.last = query;
                chunk.end_record = record;
                chunks.push_back(chunk);
                // NVCC 13.4 mislowers assignments that skip aggregate members.
                chunk.first = query;
                chunk.first_record = record;
            }
            record += row;
        }
        chunk.last = tile.query_count;
        chunk.end_record = record;
        if (record != 0U) chunks.push_back(chunk);
    } else if (words != 0U) {
        device_offsets = CUDDL_CUDA_TRY(
            cuda::make_device_buffer<uint32_t>(stream, stream.device(), words, cuda::no_init)
        );
        auto const popcounts =
            cuda::make_transform_iterator(tile.pass_bits, detail::word_popcount{});
        size_t scan_bytes = 0;
        CUDDL_CUDA_TRY(
            cub::DeviceScan::ExclusiveSum(
                nullptr,
                scan_bytes,
                popcounts,
                device_offsets.data(),
                static_cast<int64_t>(words),
                stream.get()
            )
        );
        auto scan_storage = CUDDL_CUDA_TRY(
            cuda::make_device_buffer<uint8_t>(stream, stream.device(), scan_bytes, cuda::no_init)
        );
        CUDDL_CUDA_TRY(
            cub::DeviceScan::ExclusiveSum(
                scan_storage.data(),
                scan_bytes,
                popcounts,
                device_offsets.data(),
                static_cast<int64_t>(words),
                stream.get()
            )
        );
        pass_bits.resize(words);
        offsets.resize(words);
        CUDDL_CUDA_TRY(
            cuda::copy_bytes(
                stream, cuda::std::span{tile.pass_bits, words}, cuda::std::span{pass_bits}
            )
        );
        CUDDL_CUDA_TRY(cuda::copy_bytes(stream, device_offsets, cuda::std::span{offsets}));
        CUDDL_CUDA_TRY(stream.sync());
        auto const total = static_cast<size_t>(offsets.back()) +
                           static_cast<size_t>(cuda::std::popcount(pass_bits.back()));
        for (size_t first = 0U; first < words;) {
            auto const base = static_cast<size_t>(offsets[first]);
            // Whole words only: a word holds at most 32 results.
            auto const last = static_cast<size_t>(
                std::upper_bound(
                    offsets.begin() + static_cast<std::ptrdiff_t>(first) + 1,
                    offsets.end(),
                    static_cast<uint32_t>(base + chunk_records - 32U)
                ) -
                offsets.begin()
            );
            auto const end = last == words ? total : static_cast<size_t>(offsets[last]);
            if (end != base) chunks.push_back({first, last, base, end});
            first = last;
        }
    }
    if (chunks.empty()) {
        CUDDL_CUDA_TRY(stream.sync());
        return Ok();
    }
 
    auto staging = CUDDL_CUDA_TRY(
        cuda::make_device_buffer<packed_pairwise_counts>(
            stream, stream.device(), threshold ? chunk_records : 0U, cuda::no_init
        )
    );
    std::array<pinned_counts, 2> buffers{
        pinned_counts(stream, memory, chunk_records, cuda::no_init),
        pinned_counts(stream, memory, chunk_records, cuda::no_init),
    };
    std::array<cuda::event, 2> ready{cuda::event(stream.device()), cuda::event(stream.device())};
    auto const enqueue = [&](size_t k) -> Result<void> {
        auto const& chunk = chunks[k];
        auto const records = chunk.end_record - chunk.first_record;
        packed_pairwise_counts const* source = tile.results + chunk.first_record;
        if (threshold) {
            auto const kernel = tile.upper_triangle
                                    ? detail::gather_passing_kernel<true, packed_pairwise_counts>
                                    : detail::gather_passing_kernel<false, packed_pairwise_counts>;
            kernel<<<
                detail::warp_grid_blocks(chunk.last - chunk.first),
                detail::block_size,
                0,
                stream.get()>>>(
                tile.results,
                nullptr,
                tile.pass_bits,
                device_offsets.data(),
                tile.first_query_id,
                tile.reference_count,
                chunk.first,
                chunk.last,
                static_cast<uint32_t>(chunk.first_record),
                staging.data(),
                nullptr
            );
            CUDDL_CUDA_TRY(cudaGetLastError());
            source = staging.data();
        }
        CUDDL_CUDA_TRY(
            cuda::copy_bytes(
                stream,
                cuda::std::span{source, records},
                cuda::std::span{buffers[k % 2U].data(), records}
            )
        );
        ready[k % 2U].record(stream);
        return Ok();
    };
 
    CUDDL_TRY(enqueue(0U));
    for (size_t k = 0; k < chunks.size(); ++k) {
        // Buffer (k + 1) % 2 last held chunk k - 1, which the host has finished visiting.
        if (k + 1U < chunks.size()) CUDDL_TRY(enqueue(k + 1U));
        CUDDL_CUDA_TRY(ready[k % 2U].sync());
        auto const* packed = buffers[k % 2U].data();
        auto const& chunk = chunks[k];
        size_t next = 0U;
        auto const emit = [&](uint32_t query_id, uint32_t reference_id) {
            f(batch_search_result{
                .query_id = query_id,
                .reference_id = reference_id,
                .counts = packed[next++].unpack(tile.bucket_count),
            });
        };
        if (!threshold) {
            for (auto query = static_cast<uint32_t>(chunk.first); query < chunk.last; ++query) {
                auto const query_id = tile.first_query_id + query;
                for (auto reference = tile.upper_triangle ? query_id + 1U : 0U;
                     reference < tile.reference_count;
                     ++reference) {
                    emit(query_id, reference);
                }
            }
            continue;
        }
        auto const words_per_query = tile.words_per_query();
        auto query = static_cast<uint32_t>(chunk.first / words_per_query);
        auto column = static_cast<uint32_t>(chunk.first % words_per_query);
        auto w = chunk.first;
        while (w < chunk.last) {
            // A transfer chunk may start or end inside a query row.
            auto const row_end = std::min(chunk.last, w + words_per_query - column);
            auto const query_id = tile.first_query_id + query;
            for (auto base = column * 32U; w < row_end; ++w, base += 32U) {
                for (auto bits = pass_bits[w]; bits != 0U; bits &= bits - 1U) {
                    emit(query_id, base + static_cast<uint32_t>(cuda::std::countr_zero(bits)));
                }
            }
            ++query;
            column = 0U;
        }
    }
    return Ok();
}
 
struct batch_search_requirements {
    uint32_t maximum_pair_count{};
    size_t counter_bytes{};
    size_t candidate_bytes{};
    size_t temporary_bytes{};
    size_t workspace_bytes{};
    size_t result_bytes{};
    size_t match_count_bytes{};
 
    friend bool operator==(batch_search_requirements const&, batch_search_requirements const&) =
        default;
};
 
struct search_options {
    uint32_t minimum_matches = 5;
};
 
struct all_to_all_shard {
    uint32_t index = 0U;
    uint32_t count = 1U;
};
 
namespace detail {
 
struct sparse_index_key {
    uint16_t mask;
    template <typename Row>
    __host__ __device__ uint16_t operator()(Row row) const noexcept {
        auto const score = static_cast<uint16_t>(row);
        return score == 0U || mask == 0xffffU ? score : static_cast<uint16_t>((score & mask) + 1U);
    }
};
struct sparse_reference_id {
    uint32_t reference_count;
    __host__ __device__ uint32_t operator()(uint32_t i) const noexcept {
        return i % reference_count;
    }
};
struct sparse_segment_offset {
    uint32_t reference_count;
    __host__ __device__ int64_t operator()(uint32_t bucket) const noexcept {
        return static_cast<int64_t>(bucket) * reference_count;
    }
};
 
template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
[[nodiscard]] inline Result<void> validate_score_compatibility(
    score_compatibility const& compatibility
) {
    if (compatibility.kmer_length != K) {
        return Err(Error::invalid_argument("k-mer length does not match the database type"));
    }
    if (compatibility.bucket_count != BucketCount) {
        return Err(Error::invalid_argument("bucket count does not match the database type"));
    }
    if (compatibility.score_encoder_identity == 0U ||
        compatibility.exponent_bits != Layout::exponent_bits ||
        compatibility.mantissa_bits != Layout::mantissa_bits) {
        return Err(Error::invalid_argument("score encoding does not match the register layout"));
    }
    if (compatibility.hash_identity == 0U) {
        return Err(Error::invalid_argument("hash identity must be specified"));
    }
    if (compatibility.canonicalisation_policy == 0U) {
        return Err(Error::invalid_argument("canonicalisation policy must be specified"));
    }
    return Ok();
}
 
template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
[[nodiscard]] inline Result<void> validate_non_indexed_score_compatibility(
    score_compatibility const& compatibility
) {
    CUDDL_TRY((validate_score_compatibility<K, BucketCount, Layout>(compatibility)));
    if (compatibility.indexed_bucket_count != BucketCount) {
        return Err(Error::invalid_argument("non-indexed builds require every bucket"));
    }
    if (compatibility.key_mask != 0xffffU && compatibility.key_mask != 0x7fffU) {
        return Err(Error::invalid_argument("non-indexed builds require a 16-bit or 15-bit key"));
    }
    return Ok();
}
 
template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
[[nodiscard]] inline Result<void> validate_indexed_score_compatibility(
    score_compatibility const& compatibility
) {
    CUDDL_TRY((validate_score_compatibility<K, BucketCount, Layout>(compatibility)));
    auto const full_bucket_count = static_cast<uint32_t>(BucketCount);
    if (compatibility.indexed_bucket_count != full_bucket_count &&
        compatibility.indexed_bucket_count != full_bucket_count / 2U) {
        return Err(Error::invalid_argument("indexed builds require a full or half bucket count"));
    }
    if (compatibility.key_mask != 0xffffU && compatibility.key_mask != 0x7fffU) {
        return Err(Error::invalid_argument("indexed builds require a 16-bit or 15-bit key mask"));
    }
    return Ok();
}
 
[[nodiscard]] inline constexpr uint64_t indexed_cell_count(
    score_compatibility const& compatibility
) noexcept {
    return static_cast<uint64_t>(compatibility.indexed_bucket_count) *
           (static_cast<uint64_t>(compatibility.key_mask) + 1U);
}
[[nodiscard]] inline constexpr uint64_t
indexed_posting_count(uint32_t reference_count, score_compatibility const& compatibility) noexcept {
    return static_cast<uint64_t>(reference_count) * compatibility.indexed_bucket_count;
}
 
[[nodiscard]] inline constexpr index_storage
smaller_index_storage(uint32_t reference_count, score_compatibility const& compatibility) noexcept {
    auto const dense = (indexed_cell_count(compatibility) + 1U) * sizeof(uint32_t);
    auto const sparse = indexed_posting_count(reference_count, compatibility) * sizeof(uint16_t) +
                        static_cast<uint64_t>(sparse_directory_entries(reference_count)) *
                            compatibility.indexed_bucket_count * sizeof(uint32_t);
    return sparse < dense ? index_storage::sparse : index_storage::dense;
}
 
[[nodiscard]] inline uintptr_t align_up(uintptr_t address, size_t alignment) noexcept {
    return (address + alignment - 1U) & ~(static_cast<uintptr_t>(alignment) - 1U);
}
 
// ponytail: one crossover measured on an RTX 5070 Ti with 2048-bucket RefSeq/synthetic data;
// make it per-device if other GPUs disagree.
constexpr double index_pair_fraction_limit = 0.03;
 
constexpr uint32_t batch_query_tile_count = 128U;
 
[[nodiscard]] constexpr uint32_t batch_query_tile_size(
    uint32_t reference_count,
    uint32_t query_count,
    uint32_t tile_limit = batch_query_tile_count
) noexcept {
    if (reference_count == 0U || query_count == 0U) {
        return 0U;
    }
    auto const count_capacity = std::numeric_limits<uint32_t>::max() / reference_count;
    return std::min(query_count, std::min(tile_limit, std::max(1U, count_capacity)));
}
 
[[nodiscard]] constexpr uint32_t
all_to_all_tile_queries(uint32_t reference_count, uint32_t parts = 1U) noexcept {
    constexpr uint32_t pair_budget = 1U << 25U;
    constexpr uint64_t tiles_per_part = 4U;
    auto const limit = reference_count == 0U
                           ? batch_query_tile_count
                           : std::max(batch_query_tile_count, pair_budget / reference_count);
    auto const whole = batch_query_tile_size(reference_count, reference_count, limit);
    if (parts < 2U || whole == 0U || whole >= reference_count) return whole;
    auto const per_part = static_cast<uint32_t>(reference_count / (parts * tiles_per_part));
    return std::min(whole, std::max(batch_query_tile_count, per_part));
}
 
[[nodiscard]] inline uint32_t
external_batch_query_limit(uint32_t reference_count, cuda::device_ref device) {
    constexpr size_t bytes_per_pair = sizeof(packed_pairwise_counts) + 3U * sizeof(uint32_t);
    auto const budget = device.attribute(cuda::device_attributes::total_global_memory) / 4U;
    auto const pairs = budget / bytes_per_pair;
    auto const queries = reference_count == 0U ? pairs : pairs / reference_count;
    return static_cast<uint32_t>(
        std::clamp<size_t>(queries, batch_query_tile_count, std::numeric_limits<uint32_t>::max())
    );
}
 
}  // namespace detail
 
class all_to_all_tile_queue {
   public:
    struct tile {
        uint32_t first_query_id{};
        uint32_t query_count{};
    };
 
    [[nodiscard]] static Result<all_to_all_tile_queue> make(
        uint32_t reference_count,
        all_to_all_shard shard = {},
        uint32_t searchers = 1U,
        std::function<uint64_t()> next_tile = {}
    ) {
        if (shard.count == 0U || shard.index >= shard.count) {
            return Err(Error::invalid_argument("all-to-all shard index must be below its count"));
        }
        if (searchers == 0U || shard.count > std::numeric_limits<uint32_t>::max() / searchers) {
            return Err(Error::invalid_argument("all-to-all searcher count is out of range"));
        }
        all_to_all_tile_queue queue;
        queue.reference_count_ = reference_count;
        queue.tile_queries_ =
            detail::all_to_all_tile_queries(reference_count, shard.count * searchers);
        auto const count = queue.tile_queries_ == 0U
                               ? 0U
                               : (reference_count + queue.tile_queries_ - 1U) / queue.tile_queries_;
        for (uint32_t t = 0U; t < count; ++t) {
            auto const lane = t % shard.count;
            auto const owner = (t / shard.count) % 2U == 0U ? lane : shard.count - 1U - lane;
            if (next_tile || owner == shard.index) queue.tiles_.push_back(t);
        }
        queue.next_tile_ = std::move(next_tile);
        return Result<all_to_all_tile_queue>::ok(std::move(queue));
    }
 
    all_to_all_tile_queue(all_to_all_tile_queue&& other) noexcept
        : tiles_(std::move(other.tiles_)),
          next_(other.next_.load()),
          next_tile_(std::move(other.next_tile_)),
          cancelled_(other.cancelled_.load()),
          tile_queries_(other.tile_queries_),
          reference_count_(other.reference_count_) {}
 
    [[nodiscard]] std::optional<tile> claim() noexcept {
        if (cancelled_.load(std::memory_order_relaxed)) return std::nullopt;
        auto const i = next_tile_ ? next_tile_() : next_.fetch_add(1U, std::memory_order_relaxed);
        if (i >= tiles_.size()) return std::nullopt;
        auto const first = tiles_[i] * tile_queries_;
        return tile{first, std::min(tile_queries_, reference_count_ - first)};
    }
 
    void cancel() noexcept {
        cancelled_.store(true);
    }
 
    [[nodiscard]] bool shared() const noexcept {
        return static_cast<bool>(next_tile_);
    }
 
    [[nodiscard]] uint32_t tile_queries() const noexcept {
        return tile_queries_;
    }
 
    [[nodiscard]] size_t size() const noexcept {
        return tiles_.size();
    }
 
    [[nodiscard]] uint32_t reference_count() const noexcept {
        return reference_count_;
    }
 
    [[nodiscard]] size_t position(uint32_t first_query_id) const noexcept {
        return static_cast<size_t>(
            std::lower_bound(tiles_.begin(), tiles_.end(), first_query_id / tile_queries_) -
            tiles_.begin()
        );
    }
 
   private:
    all_to_all_tile_queue() = default;
 
    std::vector<uint32_t> tiles_;
    std::atomic<size_t> next_{0U};
    std::function<uint64_t()> next_tile_;
    std::atomic<bool> cancelled_{false};
    uint32_t tile_queries_{};
    uint32_t reference_count_{};
};
 
namespace detail {
 
template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
class reference_database_view {
    static_assert(K >= 1 && K <= 31);
    static_assert(BucketCount >= (size_t{1} << 11) && BucketCount <= (size_t{1} << 13));
    static_assert((BucketCount & (BucketCount - 1)) == 0);
 
   public:
    using score_type = uint16_t;
    using layout_type = Layout;
    using result_type = reference_search_result;
    using batch_result_type = packed_pairwise_counts;
 
    __host__ __device__ constexpr reference_database_view(
        device_span<uint32_t const> planes,
        reference_database_metadata metadata,
        device_span<uint32_t const> index_offsets = {},
        device_span<uint32_t const> index_postings = {},
        bool indexed = false,
        device_span<uint16_t const> index_keys = {},
        device_span<uint32_t const> key_directory = {},
        double index_pair_fraction = 0.0,
        device_span<uint32_t const> bitmap_cells = {},
        device_span<uint32_t const> posting_bitmaps = {},
        reference_delta_view deltas = {}
    ) noexcept
        : planes_(planes),
          metadata_(metadata),
          index_offsets_(index_offsets),
          index_postings_(index_postings),
          index_keys_(index_keys),
          key_directory_(key_directory),
          bitmap_cells_(bitmap_cells),
          posting_bitmaps_(posting_bitmaps),
          deltas_(deltas),
          index_pair_fraction_(index_pair_fraction),
          indexed_(indexed) {}
 
    [[nodiscard]] __host__ __device__ constexpr device_span<uint32_t const>
    planes() const noexcept {
        return planes_;
    }
 
    [[nodiscard]] constexpr reference_database_metadata metadata() const noexcept {
        return metadata_;
    }
 
    [[nodiscard]] constexpr uint32_t reference_count() const noexcept {
        return metadata_.reference_count;
    }
 
    [[nodiscard]] constexpr bool has_index() const noexcept {
        return indexed_;
    }
 
    [[nodiscard]] static constexpr size_t persistent_row_bytes(uint32_t reference_count) noexcept {
        return static_cast<size_t>(reference_count) * BucketCount * sizeof(score_type);
    }
 
    [[nodiscard]] constexpr size_t persistent_row_bytes() const noexcept {
        return planes_.size_bytes();
    }
 
    [[nodiscard]] constexpr size_t persistent_index_bytes() const noexcept {
        return index_offsets_.size_bytes() + index_postings_.size_bytes() +
               index_keys_.size_bytes();
    }
 
    [[nodiscard]] static constexpr size_t single_query_workspace_bytes(
        uint32_t reference_count
    ) noexcept {
        return reference_count == 0U ? 0U : query_plane_bytes(1U) + plane_alignment - 1U;
    }
 
    [[nodiscard]] constexpr size_t single_query_workspace_bytes() const noexcept {
        return single_query_workspace_bytes(metadata_.reference_count);
    }
 
    [[nodiscard]] static constexpr uint32_t single_query_result_count(
        uint32_t reference_count
    ) noexcept {
        return reference_count;
    }
 
    [[nodiscard]] constexpr uint32_t single_query_result_count() const noexcept {
        return metadata_.reference_count;
    }
 
    [[nodiscard]] Result<size_t> indexed_single_query_workspace_bytes(
        cuda::stream_ref stream
    ) const {
        CUDDL_TRY(validate_index_storage());
        if (metadata_.reference_count == 0U) {
            return size_t{0};
        }
 
        size_t selection_bytes = 0;
        auto const ids = cuda::make_counting_iterator(uint32_t{0});
        auto const selection = cuda_try(
            cub::DeviceSelect::If(
                nullptr,
                selection_bytes,
                ids,
                static_cast<uint32_t*>(nullptr),
                static_cast<uint32_t*>(nullptr),
                static_cast<int64_t>(metadata_.reference_count),
                detail::minimum_match_predicate{nullptr, 1U},
                stream.get()
            )
        );
        if (!selection) {
            return Err(selection.error());
        }
        // The counting pass stages two words per indexed bucket in the selection storage.
        selection_bytes = std::max(
            selection_bytes,
            static_cast<size_t>(metadata_.compatibility.indexed_bucket_count) * 2U *
                sizeof(uint32_t)
        );
 
        constexpr size_t alignment_slack = plane_alignment - 1U + 3U + 255U;
        auto const arrays_bytes =
            query_plane_bytes(1U) +
            static_cast<size_t>(metadata_.reference_count) * 2U * sizeof(uint32_t);
        if (selection_bytes > std::numeric_limits<size_t>::max() - arrays_bytes - alignment_slack) {
            return Err(Error::resource("indexed single-query workspace size overflows"));
        }
        return arrays_bytes + alignment_slack + selection_bytes;
    }
 
    [[nodiscard]] Result<cuddl::batch_search_requirements>
    batch_search_requirements(uint32_t query_count, cuda::stream_ref stream) const {
        auto const pair_count =
            CUDDL_TRY(dense_batch_pair_count(external_batch_tile(query_count, stream)));
        return make_batch_requirements(
            pair_count, pair_count, false, pair_count / reference_divisor()
        );
    }
 
    [[nodiscard]] Result<cuddl::batch_search_requirements>
    indexed_batch_search_requirements(uint32_t query_count, cuda::stream_ref stream) const {
        CUDDL_TRY(validate_index_storage());
        auto const pair_count =
            CUDDL_TRY(dense_batch_pair_count(external_batch_tile(query_count, stream)));
        return make_batch_requirements(
            pair_count, pair_count, true, pair_count / reference_divisor()
        );
    }
 
    [[nodiscard]] Result<cuddl::batch_search_requirements> all_to_all_search_requirements() const {
        auto const slot = CUDDL_TRY(all_to_all_slot_pairs());
        if (slot > std::numeric_limits<uint32_t>::max() / 2U) {
            return Err(Error::resource("all-to-all tile exceeds 32-bit result counts"));
        }
        return make_batch_requirements(
            0U, pipeline_slots(all_to_all_tile_count()) * slot, false, 0U
        );
    }
 
    [[nodiscard]] static constexpr uint32_t all_to_all_result_capacity(
        uint32_t reference_count
    ) noexcept {
        if (reference_count < 2U) {
            return 0U;
        }
        auto const query_count = detail::all_to_all_tile_queries(reference_count);
        auto const pair_count = static_cast<uint64_t>(query_count) *
                                (2ULL * reference_count - query_count - 1ULL) / 2ULL;
        auto const tiles = (reference_count + query_count - 1U) / query_count;
        return static_cast<uint32_t>(pair_count * pipeline_slots(tiles));
    }
 
    [[nodiscard]] Result<cuddl::batch_search_requirements>
    indexed_all_to_all_search_requirements() const {
        auto requirements = CUDDL_TRY(indexed_all_to_all_tile_search_requirements(
            0U, detail::all_to_all_tile_queries(metadata_.reference_count)
        ));
        auto const exhaustive = CUDDL_TRY(all_to_all_search_requirements());
        requirements.maximum_pair_count =
            std::max(requirements.maximum_pair_count, exhaustive.maximum_pair_count);
        requirements.result_bytes = std::max(requirements.result_bytes, exhaustive.result_bytes);
        requirements.match_count_bytes =
            std::max(requirements.match_count_bytes, exhaustive.match_count_bytes);
        return requirements;
    }
 
    [[nodiscard]] Result<void> search_async(
        device_span<score_type const> query,
        score_compatibility const& query_compatibility,
        device_span<uint8_t> workspace,
        device_span<result_type> results,
        cuda::stream_ref stream
    ) const {
        device_scope const scope{stream};
        auto const expected_scores = static_cast<size_t>(metadata_.reference_count) * BucketCount;
        if (!rows_match_metadata(expected_scores)) {
            return Err(Error::invalid_argument("database extent does not match its metadata"));
        }
        if (query.size() != BucketCount || query.data() == nullptr) {
            return Err(Error::invalid_argument("query must contain one complete score row"));
        }
        CUDDL_TRY(
            (detail::validate_score_compatibility<K, BucketCount, Layout>(query_compatibility))
        );
        if (query_compatibility != metadata_.compatibility) {
            return Err(Error::invalid_argument("query construction metadata is incompatible"));
        }
        if (workspace.size_bytes() < single_query_workspace_bytes()) {
            return Err(Error::resource("single-query workspace is too small"));
        }
        if (results.size() < metadata_.reference_count) {
            return Err(Error::resource("single-query result capacity is too small"));
        }
        if (metadata_.reference_count == 0U) {
            return Ok();
        }
        if (results.data() == nullptr) {
            return Err(Error::invalid_argument("result storage must be device accessible"));
        }
 
        if (workspace.data() == nullptr) {
            return Err(Error::invalid_argument("single-query workspace must be device accessible"));
        }
        auto* query_planes = CUDDL_TRY(stage_query_planes(query, 0U, 1U, workspace, stream));
        detail::single_query_planes_kernel<BucketCount>
            <<<detail::warp_grid_blocks(metadata_.reference_count),
               detail::block_size,
               0,
               stream.get()>>>(
                query_planes,
                planes_.data(),
                metadata_.reference_count,
                nullptr,
                nullptr,
                results.data()
            );
        return cuda_try(cudaGetLastError());
    }
 
    [[nodiscard]] Result<void> search_indexed_async(
        device_span<score_type const> query,
        score_compatibility const& query_compatibility,
        device_span<uint8_t> workspace,
        device_span<result_type> results,
        device_span<uint32_t> result_count,
        search_options options,
        cuda::stream_ref stream
    ) const {
        device_scope const scope{stream};
        auto const expected_scores = static_cast<size_t>(metadata_.reference_count) * BucketCount;
        if (!rows_match_metadata(expected_scores)) {
            return Err(Error::invalid_argument("database extent does not match its metadata"));
        }
        CUDDL_TRY(validate_index_storage());
        if (query.size() != BucketCount || query.data() == nullptr) {
            return Err(Error::invalid_argument("query must contain one complete score row"));
        }
        CUDDL_TRY(
            (detail::validate_score_compatibility<K, BucketCount, Layout>(query_compatibility))
        );
        if (query_compatibility != metadata_.compatibility) {
            return Err(Error::invalid_argument("query construction metadata is incompatible"));
        }
        if (options.minimum_matches > metadata_.compatibility.indexed_bucket_count) {
            return Err(Error::invalid_argument("minimum matches exceeds indexed bucket count"));
        }
        if (results.size() < metadata_.reference_count) {
            return Err(Error::resource("indexed result capacity is too small"));
        }
        if (result_count.empty()) {
            return Err(Error::resource("indexed result count capacity is too small"));
        }
        if (result_count.data() == nullptr) {
            return Err(Error::invalid_argument("result count storage must be device accessible"));
        }
        if (metadata_.reference_count != 0U && results.data() == nullptr) {
            return Err(Error::invalid_argument("result storage must be device accessible"));
        }
 
        if (options.minimum_matches == 0U) {
            CUDDL_TRY(search_async(query, query_compatibility, workspace, results, stream));
            return write_batch_result_count(metadata_.reference_count, result_count, stream);
        }
        if (metadata_.reference_count == 0U) {
            return write_batch_result_count(0U, result_count, stream);
        }
 
        auto const required_workspace = CUDDL_TRY(indexed_single_query_workspace_bytes(stream));
        if (workspace.size_bytes() < required_workspace) {
            return Err(Error::resource("indexed single-query workspace is too small"));
        }
        if (workspace.data() == nullptr) {
            return Err(Error::invalid_argument("indexed workspace must be device accessible"));
        }
 
        auto const workspace_begin = reinterpret_cast<uintptr_t>(workspace.data());
        auto const workspace_end = workspace_begin + workspace.size_bytes();
        if (workspace_end < workspace_begin) {
            return Err(Error::resource("indexed workspace address range overflows"));
        }
        auto* query_planes = CUDDL_TRY(stage_query_planes(query, 0U, 1U, workspace, stream));
        auto address = detail::align_up(
            reinterpret_cast<uintptr_t>(query_planes) + query_plane_bytes(1U), alignof(uint32_t)
        );
        auto* match_counts = reinterpret_cast<uint32_t*>(address);
        address += static_cast<size_t>(metadata_.reference_count) * sizeof(uint32_t);
        address = detail::align_up(address, alignof(uint32_t));
        auto* candidate_ids = reinterpret_cast<uint32_t*>(address);
        address += static_cast<size_t>(metadata_.reference_count) * sizeof(uint32_t);
        address = detail::align_up(address, 256U);
        if (address > workspace_end) {
            return Err(Error::resource("indexed workspace layout exceeds its capacity"));
        }
        auto* selection_workspace = reinterpret_cast<void*>(address);
        auto selection_bytes = static_cast<size_t>(workspace_end - address);
 
        if (!indexed_) {
            detail::count_plane_matches_kernel<BucketCount>
                <<<detail::warp_grid_blocks(metadata_.reference_count),
                   detail::block_size,
                   0,
                   stream.get()>>>(
                    query_planes,
                    1U,
                    planes_.data(),
                    metadata_.reference_count,
                    metadata_.compatibility.indexed_bucket_count,
                    metadata_.compatibility.key_mask,
                    match_counts
                );
            CUDDL_CUDA_TRY(cudaGetLastError());
        } else {
            CUDDL_CUDA_TRY(
                cuda::fill_bytes(
                    stream,
                    cuda::std::span{match_counts, static_cast<size_t>(metadata_.reference_count)},
                    0
                )
            );
            // The posting ranges live in the selection workspace, which is free until the
            // selection below runs on the same stream.
            auto const indexed_bucket_count = metadata_.compatibility.indexed_bucket_count;
            auto* range_begin = static_cast<uint32_t*>(selection_workspace);
            auto* range_end_scan = range_begin + indexed_bucket_count;
            detail::index_posting_ranges_kernel<<<
                1,
                detail::posting_range_block_size,
                0,
                stream.get()>>>(
                query.data(),
                index_offsets_.data(),
                index_keys_.empty() ? nullptr : index_keys_.data(),
                metadata_.reference_count,
                indexed_bucket_count,
                metadata_.compatibility.key_mask,
                range_begin,
                range_end_scan
            );
            CUDDL_CUDA_TRY(cudaGetLastError());
            auto const counting_blocks =
                static_cast<uint32_t>(
                    stream.device().attribute(cuda::device_attributes::multiprocessor_count)
                ) *
                static_cast<uint32_t>(stream.device().attribute(
                    cuda::device_attributes::max_threads_per_multiprocessor
                )) /
                detail::block_size;
            detail::count_balanced_index_matches_kernel<<<
                counting_blocks,
                detail::block_size,
                0,
                stream.get()>>>(
                range_begin,
                range_end_scan,
                indexed_bucket_count,
                index_postings_.data(),
                match_counts
            );
            CUDDL_CUDA_TRY(cudaGetLastError());
        }
 
        auto const ids = cuda::make_counting_iterator(uint32_t{0});
        CUDDL_CUDA_TRY(
            cub::DeviceSelect::If(
                selection_workspace,
                selection_bytes,
                ids,
                candidate_ids,
                result_count.data(),
                static_cast<int64_t>(metadata_.reference_count),
                detail::minimum_match_predicate{match_counts, options.minimum_matches},
                stream.get()
            )
        );
 
        detail::single_query_planes_kernel<BucketCount>
            <<<detail::warp_grid_blocks(metadata_.reference_count),
               detail::block_size,
               0,
               stream.get()>>>(
                query_planes,
                planes_.data(),
                metadata_.reference_count,
                candidate_ids,
                result_count.data(),
                results.data()
            );
        return cuda_try(cudaGetLastError());
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_batch_async(
        device_span<score_type const> queries,
        score_compatibility const& query_compatibility,
        uint32_t query_id_offset,
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        cuda::stream_ref stream
    ) const {
        device_scope const scope{stream};
        auto const query_count =
            CUDDL_TRY(validate_batch_queries(queries, query_compatibility, query_id_offset));
        auto const requirements = CUDDL_TRY(batch_search_requirements(query_count, stream));
        CUDDL_TRY(validate_batch_outputs(requirements, results, result_match_counts));
        CUDDL_TRY(validate_workspace(requirements, workspace));
        // Result storage holds one full-size tile; two half-size tiles alternate in it.
        auto const storage_queries = external_batch_tile(query_count, stream);
        auto const tile_size = storage_queries < 2U ? storage_queries : storage_queries / 2U;
        auto const tiles = tile_size == 0U ? 0U : (query_count + tile_size - 1U) / tile_size;
        auto const slot = static_cast<size_t>(tile_size) * metadata_.reference_count;
        uint32_t next_tile = 0U;
        return pipeline_tiles(
            [&]() -> std::optional<uint32_t> {
                if (next_tile >= tiles) return std::nullopt;
                return next_tile++;
            },
            slot,
            results,
            result_match_counts,
            [&](uint32_t tile,
                device_span<batch_result_type> slot_results,
                device_span<uint32_t> slot_matches,
                cuda::stream_ref work) {
                auto const first_query_id = tile * tile_size;
                return launch_batch_exhaustive<false>(
                    queries,
                    first_query_id,
                    std::min(tile_size, query_count - first_query_id),
                    query_id_offset + first_query_id,
                    workspace,
                    slot_results,
                    slot_matches,
                    work
                );
            },
            on_tile,
            stream
        );
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_batch_indexed_async(
        device_span<score_type const> queries,
        score_compatibility const& query_compatibility,
        uint32_t query_id_offset,
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        search_options options,
        cuda::stream_ref stream
    ) const {
        device_scope const scope{stream};
        auto const query_count =
            CUDDL_TRY(validate_batch_queries(queries, query_compatibility, query_id_offset));
        auto const requirements = CUDDL_TRY(indexed_batch_search_requirements(query_count, stream));
        CUDDL_TRY(validate_indexed_batch_inputs(
            requirements, workspace, results, result_match_counts, options
        ));
        auto const tile_size = external_batch_tile(query_count, stream);
        for (uint32_t first_query_id = 0U; tile_size != 0U && first_query_id < query_count;
             first_query_id += tile_size) {
            auto const tile_query_count = std::min(tile_size, query_count - first_query_id);
            on_tile(CUDDL_TRY(
                launch_batch_indexed<false>(
                    queries.data() + static_cast<size_t>(first_query_id) * BucketCount,
                    first_query_id,
                    tile_query_count,
                    query_id_offset + first_query_id,
                    workspace,
                    results,
                    result_match_counts,
                    options,
                    stream
                )
            ));
        }
        return Ok();
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_all_to_all_async(
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        cuda::stream_ref stream,
        all_to_all_shard shard = {}
    ) const {
        auto tiles = CUDDL_TRY(all_to_all_tile_queue::make(metadata_.reference_count, shard));
        return search_all_to_all_async(
            workspace, results, std::forward<OnTile>(on_tile), result_match_counts, stream, tiles
        );
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_all_to_all_async(
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        cuda::stream_ref stream,
        all_to_all_tile_queue& tiles
    ) const {
        device_scope const scope{stream};
        CUDDL_TRY(validate_stored_rows());
        CUDDL_TRY(validate_tile_queue(tiles));
        auto const requirements = CUDDL_TRY(all_to_all_search_requirements());
        CUDDL_TRY(validate_batch_outputs(requirements, results, result_match_counts));
        std::vector<all_to_all_tile_queue::tile> claimed;
        return pipeline_tiles(
            [&]() -> std::optional<uint32_t> {
                auto const tile = tiles.claim();
                if (!tile) return std::nullopt;
                claimed.push_back(*tile);
                return static_cast<uint32_t>(claimed.size() - 1U);
            },
            CUDDL_TRY(all_to_all_slot_pairs()),
            results,
            result_match_counts,
            [&](uint32_t index,
                device_span<batch_result_type> slot_results,
                device_span<uint32_t> slot_matches,
                cuda::stream_ref work) {
                auto const tile = claimed[index];
                return launch_batch_exhaustive<true>(
                    {},
                    tile.first_query_id,
                    tile.query_count,
                    tile.first_query_id,
                    workspace,
                    slot_results,
                    slot_matches,
                    work
                );
            },
            on_tile,
            stream
        );
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_all_to_all_indexed_async(
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        search_options options,
        cuda::stream_ref stream,
        all_to_all_shard shard = {}
    ) const {
        auto tiles = CUDDL_TRY(all_to_all_tile_queue::make(metadata_.reference_count, shard));
        return search_all_to_all_indexed_async(
            workspace,
            results,
            std::forward<OnTile>(on_tile),
            result_match_counts,
            options,
            stream,
            tiles
        );
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_all_to_all_indexed_async(
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        search_options options,
        cuda::stream_ref stream,
        all_to_all_tile_queue& tiles
    ) const {
        device_scope const scope{stream};
        CUDDL_TRY(validate_stored_rows());
        CUDDL_TRY(validate_tile_queue(tiles));
        auto const requirements = CUDDL_TRY(indexed_all_to_all_search_requirements());
        CUDDL_TRY(validate_indexed_batch_inputs(
            requirements, workspace, results, result_match_counts, options
        ));
        // Every tile reuses the same storage, so the next tile is claimed once the previous
        // tile's consumers are done; tiles claimed from a shared queue then go to whichever
        // device frees up first.
        std::optional<cuda::event> consumed;
        while (true) {
            if (consumed) CUDDL_CUDA_TRY(consumed->sync());
            auto const tile = tiles.claim();
            if (!tile) break;
            on_tile(CUDDL_TRY(
                launch_batch_indexed<true>(
                    nullptr,
                    tile->first_query_id,
                    tile->query_count,
                    tile->first_query_id,
                    workspace,
                    results,
                    result_match_counts,
                    options,
                    stream
                )
            ));
            if (!consumed) consumed.emplace(CUDDL_CUDA_TRY(cuda::event(stream.device())));
            CUDDL_CUDA_TRY(consumed->record(stream));
        }
        return Ok();
    }
 
   private:
    [[nodiscard]] static constexpr uint32_t pipeline_slots(uint32_t tiles) noexcept {
        return tiles < 2U ? 1U : 2U;
    }
 
    [[nodiscard]] constexpr uint32_t all_to_all_tile_count() const noexcept {
        auto const tile_size = detail::all_to_all_tile_queries(metadata_.reference_count);
        return tile_size == 0U ? 0U : (metadata_.reference_count + tile_size - 1U) / tile_size;
    }
 
    [[nodiscard]] Result<void> validate_tile_queue(all_to_all_tile_queue const& tiles) const {
        if (tiles.reference_count() != metadata_.reference_count) {
            return Err(Error::invalid_argument("tile queue was made for another reference count"));
        }
        return Ok();
    }
 
    [[nodiscard]] Result<uint32_t> all_to_all_slot_pairs() const {
        return all_to_all_pair_count(
            0U, detail::all_to_all_tile_queries(metadata_.reference_count)
        );
    }
 
    template <typename Next, typename Launch, typename OnTile>
    [[nodiscard]] Result<void> pipeline_tiles(
        Next&& next,
        size_t slot,
        device_span<batch_result_type> results,
        device_span<uint32_t> match_counts,
        Launch&& launch,
        OnTile&& on_tile,
        cuda::stream_ref stream
    ) const {
        auto const slot_of = [&](uint32_t tile) {
            auto const offset = static_cast<size_t>(tile % 2U) * slot;
            return std::pair{
                device_span<batch_result_type>{results.data() + offset, slot},
                match_counts.empty() ? match_counts
                                     : device_span<uint32_t>{match_counts.data() + offset, slot}
            };
        };
        auto const first = next();
        if (!first) return Ok();
        auto ahead = next();
        if (!ahead) {
            auto const [slot_results, slot_matches] = slot_of(*first);
            on_tile(CUDDL_TRY(launch(*first, slot_results, slot_matches, stream)));
            return Ok();
        }
        auto work = CUDDL_CUDA_TRY(cuda::stream(stream.device()));
        std::array<cuda::event, 2> searched{
            CUDDL_CUDA_TRY(cuda::event(stream.device())),
            CUDDL_CUDA_TRY(cuda::event(stream.device())),
        };
        std::array<batch_result_tile, 2> pending{};
        auto const enqueue = [&](uint32_t tile) -> Result<void> {
            // The slot's previous tile is consumed once the caller's stream reaches this point.
            CUDDL_CUDA_TRY(work.wait(stream));
            auto const [slot_results, slot_matches] = slot_of(tile);
            pending[tile % 2U] = CUDDL_TRY(launch(tile, slot_results, slot_matches, work));
            CUDDL_CUDA_TRY(searched[tile % 2U].record(work));
            return Ok();
        };
        auto consumed = CUDDL_CUDA_TRY(cuda::event(stream.device()));
        CUDDL_TRY(enqueue(*first));
        for (auto tile = *first;; ++tile) {
            bool const more = ahead.has_value();
            if (more) CUDDL_TRY(enqueue(*ahead));
            CUDDL_CUDA_TRY(stream.wait(searched[tile % 2U]));
            on_tile(pending[tile % 2U]);
            if (!more) return Ok();
            // Tile n + 2 takes this tile's slot, so it is claimed once this tile's consumers are
            // done, while tile n + 1 keeps the device busy.
            CUDDL_CUDA_TRY(consumed.record(stream));
            CUDDL_CUDA_TRY(consumed.sync());
            ahead = next();
        }
    }
 
    [[nodiscard]] Result<cuddl::batch_search_requirements>
    indexed_all_to_all_tile_search_requirements(
        uint32_t first_query_id,
        uint32_t query_count
    ) const {
        CUDDL_TRY(validate_index_storage());
        auto const dense_pair_count = CUDDL_TRY(dense_batch_pair_count(query_count));
        auto const result_pair_count =
            CUDDL_TRY(all_to_all_pair_count(first_query_id, query_count));
        return make_batch_requirements(dense_pair_count, result_pair_count, true, query_count);
    }
 
    [[nodiscard]] Result<void> validate_stored_rows() const {
        auto const expected_scores = static_cast<size_t>(metadata_.reference_count) * BucketCount;
        if (!rows_match_metadata(expected_scores)) {
            return Err(Error::invalid_argument("database extent does not match its metadata"));
        }
        return Ok();
    }
 
    [[nodiscard]] Result<void> validate_index_storage() const {
        CUDDL_TRY((detail::validate_indexed_score_compatibility<K, BucketCount, Layout>(
            metadata_.compatibility
        )));
        if (!indexed_) return Ok();
        auto const cell_count = detail::indexed_cell_count(metadata_.compatibility);
        auto const expected_postings = static_cast<size_t>(
            detail::indexed_posting_count(metadata_.reference_count, metadata_.compatibility)
        );
        auto const sparse = index_offsets_.empty();
        auto const valid_keys =
            sparse ? index_keys_.size() == expected_postings &&
                         (expected_postings == 0U || index_keys_.data() != nullptr)
                   : index_keys_.empty() &&
                         index_offsets_.size() == static_cast<size_t>(cell_count + 1U) &&
                         index_offsets_.data() != nullptr;
        if (!indexed_ || !valid_keys || index_postings_.size() != expected_postings ||
            (expected_postings != 0U && index_postings_.data() == nullptr)) {
            return Err(Error::invalid_argument("database has no valid retrieval index"));
        }
        return Ok();
    }
 
    [[nodiscard]] uint32_t
    external_batch_tile(uint32_t query_count, cuda::stream_ref stream) const {
        return detail::batch_query_tile_size(
            metadata_.reference_count,
            query_count,
            detail::external_batch_query_limit(metadata_.reference_count, stream.device())
        );
    }
 
    [[nodiscard]] Result<uint32_t> dense_batch_pair_count(uint32_t query_count) const {
        auto const pair_count = static_cast<uint64_t>(query_count) * metadata_.reference_count;
        if (pair_count > std::numeric_limits<uint32_t>::max()) {
            return Err(Error::resource("batch pair count exceeds 32-bit workspace IDs"));
        }
        return static_cast<uint32_t>(pair_count);
    }
 
    [[nodiscard]] Result<uint32_t>
    all_to_all_pair_count(uint32_t first_query_id, uint32_t query_count) const {
        if (first_query_id > metadata_.reference_count ||
            query_count > metadata_.reference_count - first_query_id) {
            return Err(Error::invalid_argument("all-to-all query tile is outside the database"));
        }
        if (query_count == 0U) {
            return 0U;
        }
 
        uint64_t left = query_count;
        uint64_t right = 2U * static_cast<uint64_t>(metadata_.reference_count) -
                         2U * first_query_id - query_count - 1U;
        if ((left & 1U) == 0U) {
            left /= 2U;
        } else {
            right /= 2U;
        }
        auto const pair_count = left * right;
        if (pair_count > std::numeric_limits<uint32_t>::max()) {
            return Err(Error::resource("all-to-all tile exceeds 32-bit result counts"));
        }
        return static_cast<uint32_t>(pair_count);
    }
 
    static constexpr size_t plane_alignment = 16U;
 
    [[nodiscard]] static constexpr size_t query_plane_bytes(uint32_t query_count) noexcept {
        return static_cast<size_t>(query_count) * BucketCount * sizeof(uint16_t);
    }
 
    [[nodiscard]] constexpr uint32_t reference_divisor() const noexcept {
        return metadata_.reference_count == 0U ? 1U : metadata_.reference_count;
    }
 
    [[nodiscard]] constexpr bool compares_planes() const noexcept {
        return !indexed_ ||
               (bitmap_cells_.empty() && index_pair_fraction_ > detail::index_pair_fraction_limit);
    }
 
    [[nodiscard]] Result<cuddl::batch_search_requirements> make_batch_requirements(
        uint32_t dense_pair_count,
        uint32_t maximum_pair_count,
        bool indexed,
        uint32_t staged_queries
    ) const {
        cuddl::batch_search_requirements requirements{
            .maximum_pair_count = maximum_pair_count,
            .result_bytes = static_cast<size_t>(maximum_pair_count) * sizeof(batch_result_type),
            .match_count_bytes = static_cast<size_t>(maximum_pair_count) * sizeof(uint32_t),
        };
        auto const staging_bytes = query_plane_bytes(staged_queries);
        if (!indexed || dense_pair_count == 0U) {
            // The refinement's cell counter follows the staged query planes.
            if (maximum_pair_count != 0U) {
                requirements.counter_bytes = sizeof(uint32_t);
                requirements.workspace_bytes =
                    staging_bytes + plane_alignment - 1U + sizeof(uint32_t);
            }
            return requirements;
        }
 
        auto const tile_query_count = dense_pair_count / reference_divisor();
        auto const word_count = static_cast<size_t>(tile_query_count) *
                                detail::candidate_bit_words(metadata_.reference_count);
        requirements.candidate_bytes = word_count * sizeof(uint32_t);
        if (!compares_planes()) {
            requirements.counter_bytes = static_cast<size_t>(dense_pair_count) * sizeof(uint32_t);
            if (!index_keys_.empty()) {
                // Sparse counting first stages the tile's posting ranges.
                requirements.temporary_bytes = static_cast<size_t>(tile_query_count) *
                                               metadata_.compatibility.indexed_bucket_count *
                                               sizeof(uint2);
            }
            if (!bitmap_cells_.empty()) {
                requirements.temporary_bytes = static_cast<size_t>(tile_query_count) *
                                               (metadata_.compatibility.indexed_bucket_count *
                                                    (sizeof(uint2) + sizeof(uint32_t)) +
                                                sizeof(uint32_t));
            }
        } else {
            // The match filter and refinement reuse one cell counter.
            requirements.counter_bytes = sizeof(uint32_t);
        }
        constexpr size_t alignment_slack = alignof(uint32_t) - 1U + plane_alignment - 1U + 255U;
        requirements.workspace_bytes = requirements.counter_bytes + requirements.candidate_bytes +
                                       staging_bytes + requirements.temporary_bytes +
                                       alignment_slack;
        return requirements;
    }
 
    [[nodiscard]] static Result<void> validate_workspace(
        cuddl::batch_search_requirements const& requirements,
        device_span<uint8_t> workspace
    ) {
        if (workspace.size_bytes() < requirements.workspace_bytes) {
            return Err(Error::resource("batch workspace is too small"));
        }
        if (requirements.workspace_bytes != 0U && workspace.data() == nullptr) {
            return Err(Error::invalid_argument("batch workspace must be device accessible"));
        }
        return Ok();
    }
 
    [[nodiscard]] Result<uint32_t*> stage_query_planes(
        device_span<score_type const> queries,
        size_t first_row,
        uint32_t count,
        device_span<uint8_t> workspace,
        cuda::stream_ref stream
    ) const {
        auto const address =
            detail::align_up(reinterpret_cast<uintptr_t>(workspace.data()), plane_alignment);
        if (address + query_plane_bytes(count) >
            reinterpret_cast<uintptr_t>(workspace.data()) + workspace.size_bytes()) {
            return Err(Error::resource("query staging exceeds the workspace"));
        }
        auto* planes = reinterpret_cast<uint32_t*>(address);
        detail::build_score_planes_kernel<BucketCount>
            <<<detail::warp_grid_blocks(static_cast<size_t>(count) * (BucketCount / 32U)),
               detail::block_size,
               0,
               stream.get()>>>(queries.data() + first_row * BucketCount, count, planes);
        CUDDL_CUDA_TRY(cudaGetLastError());
        return planes;
    }
 
    [[nodiscard]] Result<uint32_t> validate_batch_queries(
        device_span<score_type const> queries,
        score_compatibility const& query_compatibility,
        uint32_t query_id_offset
    ) const {
        CUDDL_TRY(validate_stored_rows());
        if (queries.size() % BucketCount != 0U) {
            return Err(Error::invalid_argument("query tile must contain complete score rows"));
        }
        if (!queries.empty() && queries.data() == nullptr) {
            return Err(Error::invalid_argument("query tile must be device accessible"));
        }
        auto const query_count = queries.size() / BucketCount;
        if (query_count > std::numeric_limits<uint32_t>::max()) {
            return Err(Error::resource("query tile exceeds stable 32-bit IDs"));
        }
        auto const count = static_cast<uint32_t>(query_count);
        if (count != 0U && query_id_offset > std::numeric_limits<uint32_t>::max() - (count - 1U)) {
            return Err(Error::resource("query IDs exceed the stable 32-bit range"));
        }
        if (auto const validation =
                detail::validate_score_compatibility<K, BucketCount, Layout>(query_compatibility);
            !validation) {
            return Err(validation.error());
        }
        if (query_compatibility != metadata_.compatibility) {
            return Err(Error::invalid_argument("query construction metadata is incompatible"));
        }
        return count;
    }
 
    [[nodiscard]] Result<void> validate_batch_outputs(
        cuddl::batch_search_requirements const& requirements,
        device_span<batch_result_type> results,
        device_span<uint32_t> result_match_counts
    ) const {
        if (results.size() < requirements.maximum_pair_count) {
            return Err(Error::resource("batch result capacity is too small"));
        }
        if (requirements.maximum_pair_count != 0U && results.data() == nullptr) {
            return Err(Error::invalid_argument("batch results must be device accessible"));
        }
        if (!result_match_counts.empty()) {
            if (result_match_counts.size() < requirements.maximum_pair_count) {
                return Err(Error::resource("batch match-count capacity is too small"));
            }
            if (result_match_counts.data() == nullptr) {
                return Err(Error::invalid_argument("batch match counts must be device accessible"));
            }
        }
        return Ok();
    }
 
    [[nodiscard]] Result<void> validate_indexed_batch_inputs(
        cuddl::batch_search_requirements const& requirements,
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        device_span<uint32_t> result_match_counts,
        search_options options
    ) const {
        CUDDL_TRY(validate_index_storage());
        if (options.minimum_matches > metadata_.compatibility.indexed_bucket_count) {
            return Err(Error::invalid_argument("minimum matches exceeds indexed bucket count"));
        }
        if (workspace.size_bytes() < requirements.workspace_bytes) {
            return Err(Error::resource("indexed batch workspace is too small"));
        }
        if (requirements.workspace_bytes != 0U && workspace.data() == nullptr) {
            return Err(
                Error::invalid_argument("indexed batch workspace must be device accessible")
            );
        }
        return validate_batch_outputs(requirements, results, result_match_counts);
    }
 
    [[nodiscard]] static Result<void> write_batch_result_count(
        uint32_t count,
        device_span<uint32_t> result_count,
        cuda::stream_ref stream
    ) {
        return cuda_try(
            cub::DeviceTransform::Transform(
                cuda::make_constant_iterator(count),
                result_count.data(),
                1,
                cuda::std::identity{},
                stream
            )
        );
    }
 
    template <bool AllToAll>
    [[nodiscard]] Result<batch_result_tile> launch_batch_exhaustive(
        device_span<score_type const> queries,
        uint32_t first_query,
        uint32_t query_count,
        uint32_t query_id_offset,
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        device_span<uint32_t> result_match_counts,
        cuda::stream_ref stream
    ) const {
        uint32_t const* query_planes = nullptr;
        if constexpr (AllToAll) {
            query_planes = planes_.data() + static_cast<size_t>(first_query) * (BucketCount / 2U);
        } else {
            query_planes =
                CUDDL_TRY(stage_query_planes(queries, first_query, query_count, workspace, stream));
        }
        // The cell counter follows the staged query planes; see make_batch_requirements.
        auto const counter_address =
            detail::align_up(reinterpret_cast<uintptr_t>(workspace.data()), plane_alignment) +
            query_plane_bytes(AllToAll ? 0U : query_count);
        if (counter_address + sizeof(uint32_t) >
            reinterpret_cast<uintptr_t>(workspace.data()) + workspace.size_bytes()) {
            return Err(Error::resource("exhaustive batch workspace is too small"));
        }
        auto* cell_counter = reinterpret_cast<uint32_t*>(counter_address);
        CUDDL_CUDA_TRY(cudaMemsetAsync(cell_counter, 0, sizeof(uint32_t), stream.get()));
        if (deltas_.child_count != 0U && result_match_counts.empty() && corrects_deltas(stream)) {
            // Bases are refined exactly; every child pair then follows from its base pair.
            CUDDL_TRY((launch_refine<detail::refine_candidates::reference_mask, AllToAll>(
                query_planes,
                query_id_offset,
                query_count,
                static_cast<uint32_t>(BucketCount),
                uint16_t{0xffffU},
                deltas_.base_bits,
                results,
                result_match_counts,
                stream,
                0U,
                nullptr,
                cell_counter
            )));
            auto const kernel =
                detail::apply_reference_deltas_kernel<BucketCount, batch_result_type, AllToAll>;
            CUDDL_CUDA_TRY(cudaFuncSetAttribute(
                reinterpret_cast<void const*>(kernel),
                cudaFuncAttributeMaxDynamicSharedMemorySize,
                static_cast<int>(detail::delta_correction_shared_bytes<BucketCount>)
            ));
            kernel<<<
                static_cast<uint32_t>(
                    stream.device().attribute(cuda::device_attributes::multiprocessor_count)
                ),
                detail::delta_correction_warps * 32U,
                detail::delta_correction_shared_bytes<BucketCount>,
                stream.get()>>>(
                query_planes,
                query_id_offset,
                query_count,
                metadata_.reference_count,
                deltas_,
                results.data()
            );
            CUDDL_CUDA_TRY(cudaGetLastError());
        } else {
            // Over every bucket with the full key, index match counts are the equal buckets.
            CUDDL_TRY((launch_refine<detail::refine_candidates::all, AllToAll>(
                query_planes,
                query_id_offset,
                query_count,
                static_cast<uint32_t>(BucketCount),
                uint16_t{0xffffU},
                nullptr,
                results,
                result_match_counts,
                stream,
                0U,
                nullptr,
                cell_counter
            )));
        }
        return make_tile(
            results, result_match_counts, nullptr, query_id_offset, query_count, AllToAll
        );
    }
 
    [[nodiscard]] static bool corrects_deltas(cuda::stream_ref stream) {
        return static_cast<size_t>(stream.device().attribute(
                   cuda::device_attributes::max_shared_memory_per_block_optin
               )) >= detail::delta_correction_shared_bytes<BucketCount>;
    }
 
    [[nodiscard]] batch_result_tile make_tile(
        device_span<batch_result_type> results,
        device_span<uint32_t> result_match_counts,
        uint32_t const* pass_bits,
        uint32_t first_query_id,
        uint32_t query_count,
        bool upper_triangle
    ) const noexcept {
        return {
            .results = results.data(),
            .match_counts = result_match_counts.empty() ? nullptr : result_match_counts.data(),
            .pass_bits = pass_bits,
            .first_query_id = first_query_id,
            .query_count = query_count,
            .reference_count = metadata_.reference_count,
            .bucket_count = static_cast<uint32_t>(BucketCount),
            .upper_triangle = upper_triangle,
        };
    }
 
    template <detail::refine_candidates Candidates, bool UpperTriangle>
    [[nodiscard]] Result<void> launch_refine(
        uint32_t const* query_planes,
        uint32_t query_id_offset,
        uint32_t query_count,
        uint32_t indexed_bucket_count,
        uint16_t key_mask,
        uint32_t const* candidate_bits,
        device_span<batch_result_type> results,
        device_span<uint32_t> result_match_counts,
        cuda::stream_ref stream,
        uint32_t minimum_matches = 0U,
        uint32_t* pass_bits = nullptr,
        uint32_t* cell_counter = nullptr
    ) const {
        auto const counting = !result_match_counts.empty() || pass_bits != nullptr;
        if constexpr (
            (Candidates == detail::refine_candidates::all ||
             Candidates == detail::refine_candidates::reference_mask) &&
            detail::match_filter_supported<BucketCount>
        ) {
            if (!counting) {
                return launch_all_pairs<UpperTriangle>(
                    query_planes,
                    query_id_offset,
                    query_count,
                    results,
                    stream,
                    cell_counter,
                    Candidates == detail::refine_candidates::reference_mask ? candidate_bits
                                                                            : nullptr
                );
            }
        }
        auto const select_kernel = [&]<uint32_t Group>() {
            return counting ? detail::refine_batch_bitmap_kernel<
                                  BucketCount,
                                  true,
                                  batch_result_type,
                                  Candidates,
                                  UpperTriangle,
                                  Group>
                            : detail::refine_batch_bitmap_kernel<
                                  BucketCount,
                                  false,
                                  batch_result_type,
                                  Candidates,
                                  UpperTriangle,
                                  Group>;
        };
        auto kernel = select_kernel.template operator()<detail::bitmap_refine_group<BucketCount>>();
        auto shared_bytes = detail::bitmap_refine_dynamic_bytes<BucketCount>;
        if constexpr (detail::bitmap_refine_staged<BucketCount>) {
            // On Hopper the densest query groups, not L2 row reads, set the kernel's time: the
            // 115,223-genome threshold search took 416 ms on a GH200 with groups of 8 queries,
            // 438 ms with 24 and 445 ms with 32. An RTX 5070 Ti is fastest with 24.
            constexpr uint32_t hopper_group =
                cuda::std::min(8U, detail::bitmap_refine_group<BucketCount>);
            if (stream.device().attribute(cuda::device_attributes::compute_capability_major) == 9) {
                kernel = select_kernel.template operator()<hopper_group>();
                shared_bytes = hopper_group * BucketCount * sizeof(uint16_t);
            }
        }
        CUDDL_CUDA_TRY(cudaFuncSetAttribute(
            reinterpret_cast<void const*>(kernel),
            cudaFuncAttributeMaxDynamicSharedMemorySize,
            static_cast<int>(shared_bytes)
        ));
        int resident_blocks = 0;
        CUDDL_CUDA_TRY(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
            &resident_blocks, kernel, detail::bitmap_refine_block_size, shared_bytes
        ));
        auto const multiprocessors = static_cast<uint32_t>(
            stream.device().attribute(cuda::device_attributes::multiprocessor_count)
        );
        kernel<<<
            multiprocessors* static_cast<uint32_t>(resident_blocks > 0 ? resident_blocks : 1),
            detail::bitmap_refine_block_size,
            shared_bytes,
            stream.get()>>>(
            query_planes,
            query_id_offset,
            query_count,
            planes_.data(),
            metadata_.reference_count,
            indexed_bucket_count,
            key_mask,
            candidate_bits,
            results.data(),
            result_match_counts.empty() ? nullptr : result_match_counts.data(),
            minimum_matches,
            pass_bits,
            cell_counter
        );
        return cuda_try(cudaGetLastError());
    }
 
    [[nodiscard]] Result<uint16_t const*> row_major_rows(cuda::stream_ref stream) const {
        static std::mutex mutex;
        // Never freed: the buffers would outlive the CUDA context at exit.
        static auto* const decoded = new std::map<uint32_t const*, uint16_t*>();
        std::lock_guard lock(mutex);
        auto& rows = (*decoded)[planes_.data()];
        if (rows == nullptr) {
            auto const count = metadata_.reference_count;
            CUDDL_CUDA_TRY(cudaMalloc(&rows, static_cast<size_t>(count) * BucketCount * 2U));
            detail::decode_score_planes_kernel<BucketCount>
                <<<detail::warp_grid_blocks(static_cast<size_t>(count) * (BucketCount / 32U)),
                   detail::block_size,
                   0,
                   stream.get()>>>(planes_.data(), count, rows);
            CUDDL_CUDA_TRY(cudaGetLastError());
        }
        return Result<uint16_t const*>::ok(rows);
    }
 
    template <bool UpperTriangle>
    [[nodiscard]] Result<void> launch_all_pairs(
        uint32_t const* query_planes,
        uint32_t query_id_offset,
        uint32_t query_count,
        device_span<batch_result_type> results,
        cuda::stream_ref stream,
        uint32_t* cell_counter,
        uint32_t const* reference_mask
    ) const {
        if constexpr (detail::ablations_enabled) {
            bool const naive = detail::ablated("lop3-kernel");
            if (naive || detail::ablated("bit-slicing")) {
                auto const* rows = CUDDL_TRY(row_major_rows(stream));
                auto const kernel = naive ? detail::row_major_all_pairs_kernel<
                                                BucketCount,
                                                batch_result_type,
                                                UpperTriangle,
                                                false>
                                          : detail::row_major_all_pairs_kernel<
                                                BucketCount,
                                                batch_result_type,
                                                UpperTriangle,
                                                true>;
                auto const multiprocessors = static_cast<uint32_t>(
                    stream.device().attribute(cuda::device_attributes::multiprocessor_count)
                );
                kernel<<<multiprocessors * 8U, detail::block_size, 0, stream.get()>>>(
                    rows,
                    query_id_offset,
                    query_count,
                    metadata_.reference_count,
                    results.data(),
                    reference_mask
                );
                return cuda_try(cudaGetLastError());
            }
        }
        constexpr auto group = detail::bitmap_refine_group<BucketCount>;
        constexpr auto shared_bytes = detail::bitmap_refine_dynamic_bytes<BucketCount>;
        auto const kernel =
            detail::refine_all_pairs_kernel<BucketCount, batch_result_type, UpperTriangle, group>;
        CUDDL_CUDA_TRY(cudaFuncSetAttribute(
            reinterpret_cast<void const*>(kernel),
            cudaFuncAttributeMaxDynamicSharedMemorySize,
            static_cast<int>(shared_bytes)
        ));
        auto const multiprocessors = static_cast<uint32_t>(
            stream.device().attribute(cuda::device_attributes::multiprocessor_count)
        );
        kernel<<<multiprocessors, detail::all_pairs_block_size, shared_bytes, stream.get()>>>(
            query_planes,
            query_id_offset,
            query_count,
            planes_.data(),
            metadata_.reference_count,
            results.data(),
            cell_counter,
            reference_mask
        );
        return cuda_try(cudaGetLastError());
    }
 
    template <bool UpperTriangle>
    [[nodiscard]] Result<void> launch_match_filter(
        uint32_t const* query_planes,
        uint32_t query_id_offset,
        uint32_t query_count,
        uint32_t indexed_bucket_count,
        uint32_t minimum_matches,
        uint32_t* candidate_bits,
        uint32_t* cell_counter,
        cuda::stream_ref stream
    ) const {
        constexpr auto group = detail::bitmap_refine_group<BucketCount>;
        constexpr auto shared_bytes = detail::bitmap_refine_dynamic_bytes<BucketCount>;
        auto const kernel = detail::match_filter_kernel<BucketCount, UpperTriangle, group>;
        CUDDL_CUDA_TRY(cudaFuncSetAttribute(
            reinterpret_cast<void const*>(kernel),
            cudaFuncAttributeMaxDynamicSharedMemorySize,
            static_cast<int>(shared_bytes)
        ));
        int resident_blocks = 0;
        CUDDL_CUDA_TRY(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
            &resident_blocks, kernel, detail::match_filter_block_size, shared_bytes
        ));
        auto const multiprocessors = static_cast<uint32_t>(
            stream.device().attribute(cuda::device_attributes::multiprocessor_count)
        );
        kernel<<<
            multiprocessors* static_cast<uint32_t>(resident_blocks > 0 ? resident_blocks : 1),
            detail::match_filter_block_size,
            shared_bytes,
            stream.get()>>>(
            query_planes,
            query_id_offset,
            query_count,
            planes_.data(),
            metadata_.reference_count,
            indexed_bucket_count,
            minimum_matches,
            candidate_bits,
            cell_counter
        );
        return cuda_try(cudaGetLastError());
    }
 
    template <bool AllToAll>
    [[nodiscard]] Result<batch_result_tile> launch_batch_indexed(
        score_type const* external_queries,
        uint32_t first_query,
        uint32_t query_count,
        uint32_t query_id_offset,
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        device_span<uint32_t> result_match_counts,
        search_options options,
        cuda::stream_ref stream
    ) const {
        auto const requirements = CUDDL_TRY(make_batch_requirements(
            CUDDL_TRY(dense_batch_pair_count(query_count)), 0U, true, query_count
        ));
        auto const workspace_begin = reinterpret_cast<uintptr_t>(workspace.data());
        auto const workspace_end = workspace_begin + workspace.size_bytes();
        if (workspace_end < workspace_begin) {
            return Err(Error::resource("indexed batch workspace address range overflows"));
        }
        auto const word_count = static_cast<size_t>(query_count) *
                                detail::candidate_bit_words(metadata_.reference_count);
        auto address = detail::align_up(workspace_begin, alignof(uint32_t));
        auto* match_counts = reinterpret_cast<uint32_t*>(address);
        address += requirements.counter_bytes;
        auto* candidate_bits = reinterpret_cast<uint32_t*>(address);
        address += requirements.candidate_bytes;
        address = detail::align_up(address, plane_alignment);
        auto* staging = reinterpret_cast<void*>(address);
        address += query_plane_bytes(query_count);
        address = detail::align_up(address, 256U);
        if (address + requirements.temporary_bytes > workspace_end) {
            return Err(Error::resource("indexed batch workspace layout exceeds its capacity"));
        }
        auto* temporary_workspace = reinterpret_cast<void*>(address);
        constexpr uint32_t warp_width = 32;
        constexpr uint32_t warps_per_block = detail::block_size / warp_width;
        auto const multiprocessors = static_cast<uint32_t>(
            stream.device().attribute(cuda::device_attributes::multiprocessor_count)
        );
        auto const tile_pair_count = query_count * metadata_.reference_count;
        auto const reference_count = metadata_.reference_count;
        auto const indexed_bucket_count = metadata_.compatibility.indexed_bucket_count;
        auto const key_mask = metadata_.compatibility.key_mask;
        auto const result_tile = make_tile(
            results, result_match_counts, candidate_bits, query_id_offset, query_count, AllToAll
        );
        if (word_count == 0U) {
            return result_tile;
        }
 
        // External queries arrive as scores and are converted to bit-planes; all-to-all queries
        // are database planes, decoded to scores only for the index lookups.
        uint16_t const* query_scores = external_queries;
        uint32_t const* query_planes = nullptr;
        if constexpr (AllToAll) {
            query_planes = planes_.data() + static_cast<size_t>(first_query) * (BucketCount / 2U);
            if (!compares_planes()) {
                auto* decoded = static_cast<uint16_t*>(staging);
                detail::decode_score_planes_kernel<BucketCount>
                    <<<detail::warp_grid_blocks(
                           static_cast<size_t>(query_count) * (BucketCount / warp_width)
                       ),
                       detail::block_size,
                       0,
                       stream.get()>>>(query_planes, query_count, decoded);
                CUDDL_CUDA_TRY(cudaGetLastError());
                query_scores = decoded;
            }
        } else {
            auto* converted = static_cast<uint32_t*>(staging);
            detail::build_score_planes_kernel<BucketCount>
                <<<detail::warp_grid_blocks(
                       static_cast<size_t>(query_count) * (BucketCount / warp_width)
                   ),
                   detail::block_size,
                   0,
                   stream.get()>>>(external_queries, query_count, converted);
            CUDDL_CUDA_TRY(cudaGetLastError());
            query_planes = converted;
        }
 
        auto const refine_selected = [&](bool thresholded) -> Result<void> {
            CUDDL_CUDA_TRY(cudaMemsetAsync(match_counts, 0, sizeof(uint32_t), stream.get()));
            return launch_refine<detail::refine_candidates::bitmap, AllToAll>(
                query_planes,
                query_id_offset,
                query_count,
                indexed_bucket_count,
                key_mask,
                candidate_bits,
                results,
                result_match_counts,
                stream,
                thresholded ? options.minimum_matches : 0U,
                thresholded ? candidate_bits : nullptr,
                match_counts
            );
        };
 
        if (compares_planes()) {
            if constexpr (AllToAll) {
                // The filter and refinement kernels skip words wholly below the triangle, so
                // those words must already read as no pair.
                CUDDL_CUDA_TRY(
                    cudaMemsetAsync(candidate_bits, 0, word_count * sizeof(uint32_t), stream.get())
                );
            }
            // A low-plane bound selects candidates; the refinement counts their matches exactly
            // and keeps, in the same bitmap, those meeting the threshold.
            if constexpr (detail::match_filter_supported<BucketCount>) {
                if (!detail::ablated("match-filter")) {
                    CUDDL_CUDA_TRY(
                        cudaMemsetAsync(match_counts, 0, sizeof(uint32_t), stream.get())
                    );
                    CUDDL_TRY((launch_match_filter<AllToAll>(
                        query_planes,
                        query_id_offset,
                        query_count,
                        indexed_bucket_count,
                        options.minimum_matches,
                        candidate_bits,
                        match_counts,
                        stream
                    )));
                    CUDDL_TRY(refine_selected(true));
                    return result_tile;
                }
            }
            // Every passing pair's summary lands at its slot as the pass bitmap is written.
            CUDDL_TRY((launch_refine<detail::refine_candidates::all, AllToAll>(
                query_planes,
                query_id_offset,
                query_count,
                indexed_bucket_count,
                key_mask,
                nullptr,
                results,
                result_match_counts,
                stream,
                options.minimum_matches,
                candidate_bits
            )));
            return result_tile;
        }
 
        auto const bits_from_counts = [&]() -> Result<void> {
            detail::candidate_bits_from_counts_kernel<<<
                detail::warp_grid_blocks((word_count + warp_width - 1U) / warp_width),
                detail::block_size,
                0,
                stream.get()>>>(
                match_counts,
                query_count,
                reference_count,
                options.minimum_matches,
                query_id_offset,
                AllToAll,
                candidate_bits
            );
            return cuda_try(cudaGetLastError());
        };
        if (detail::uses_tiled_index_counts(reference_count, query_count, indexed_bucket_count)) {
            auto const tile = detail::index_tile_references(
                reference_count,
                query_count,
                multiprocessors,
                static_cast<uint32_t>(stream.device().attribute(
                    cuda::device_attributes::max_shared_memory_per_block_optin
                ))
            );
            auto const tiles = (reference_count + tile - 1U) / tile;
            // Sparse ranges are resolved bucket-major first: a block's lookups all search one
            // bucket's keys, so their shared upper search levels stay in L1.
            uint2* ranges = nullptr;
            uint32_t counter_bits = 0U;
            if (!index_keys_.empty()) {
                ranges = static_cast<uint2*>(temporary_workspace);
                detail::sparse_batch_posting_ranges_kernel<BucketCount>
                    <<<indexed_bucket_count, detail::block_size, 0, stream.get()>>>(
                        query_scores,
                        0U,
                        query_count,
                        index_keys_.data(),
                        key_directory_.data(),
                        reference_count,
                        indexed_bucket_count,
                        key_mask,
                        ranges
                    );
                CUDDL_CUDA_TRY(cudaGetLastError());
            }
            auto const hybrid = !bitmap_cells_.empty();
            uint32_t* starts = nullptr;
            uint32_t* counts = nullptr;
            if (hybrid) {
                ranges = static_cast<uint2*>(temporary_workspace);
                auto const range_count = static_cast<size_t>(query_count) * indexed_bucket_count;
                starts = reinterpret_cast<uint32_t*>(ranges + range_count);
                counts = starts + range_count;
                detail::hybrid_posting_ranges_kernel<BucketCount>
                    <<<query_count, detail::block_size, 0, stream.get()>>>(
                        query_scores,
                        index_offsets_.empty() ? nullptr : index_offsets_.data(),
                        bitmap_cells_.data(),
                        static_cast<uint32_t>(bitmap_cells_.size()),
                        detail::candidate_bit_words(reference_count),
                        indexed_bucket_count,
                        key_mask,
                        ranges,
                        starts,
                        counts
                    );
                CUDDL_CUDA_TRY(cudaGetLastError());
                counter_bits = std::max<uint32_t>(1U, std::bit_width(options.minimum_matches));
            }
            auto const count_postings = [&]<uint32_t Bits>() -> Result<void> {
                CUDDL_CUDA_TRY(cudaFuncSetAttribute(
                    reinterpret_cast<void const*>(
                        detail::count_batch_index_tile_kernel<BucketCount, Bits, uint16_t>
                    ),
                    cudaFuncAttributeMaxDynamicSharedMemorySize,
                    static_cast<int>(tile / 2U * sizeof(uint32_t))
                ));
                detail::count_batch_index_tile_kernel<BucketCount, Bits>
                    <<<query_count * tiles,
                       detail::index_tile_block_size,
                       tile / 2U * sizeof(uint32_t),
                       stream.get()>>>(
                        query_scores,
                        0U,
                        index_offsets_.data(),
                        index_postings_.data(),
                        reference_count,
                        indexed_bucket_count,
                        key_mask,
                        tile,
                        match_counts,
                        index_keys_.empty() ? nullptr : index_keys_.data(),
                        ranges,
                        options.minimum_matches,
                        query_id_offset,
                        AllToAll,
                        candidate_bits,
                        hybrid ? posting_bitmaps_.data() : nullptr,
                        starts,
                        counts
                    );
                return cuda_try(cudaGetLastError());
            };
            switch (counter_bits) {
                case 0:
                    CUDDL_TRY((count_postings.template operator()<0>()));
                    break;
                case 1:
                    CUDDL_TRY((count_postings.template operator()<1>()));
                    break;
                case 2:
                    CUDDL_TRY((count_postings.template operator()<2>()));
                    break;
                case 3:
                    CUDDL_TRY((count_postings.template operator()<3>()));
                    break;
                case 4:
                    CUDDL_TRY((count_postings.template operator()<4>()));
                    break;
                case 5:
                    CUDDL_TRY((count_postings.template operator()<5>()));
                    break;
                case 6:
                    CUDDL_TRY((count_postings.template operator()<6>()));
                    break;
                case 7:
                    CUDDL_TRY((count_postings.template operator()<7>()));
                    break;
                case 8:
                    CUDDL_TRY((count_postings.template operator()<8>()));
                    break;
                case 9:
                    CUDDL_TRY((count_postings.template operator()<9>()));
                    break;
                case 10:
                    CUDDL_TRY((count_postings.template operator()<10>()));
                    break;
                case 11:
                    CUDDL_TRY((count_postings.template operator()<11>()));
                    break;
                case 12:
                    CUDDL_TRY((count_postings.template operator()<12>()));
                    break;
                case 13:
                    CUDDL_TRY((count_postings.template operator()<13>()));
                    break;
                case 14:
                    CUDDL_TRY((count_postings.template operator()<14>()));
                    break;
                case 15:
                    CUDDL_TRY((count_postings.template operator()<15>()));
                    break;
                case 16:
                    CUDDL_TRY((count_postings.template operator()<16>()));
                    break;
            }
        } else {
            CUDDL_CUDA_TRY(
                cuda::fill_bytes(
                    stream, cuda::std::span{match_counts, static_cast<size_t>(tile_pair_count)}, 0
                )
            );
            auto const query_buckets = static_cast<size_t>(query_count) * indexed_bucket_count;
            constexpr auto cells_per_block = warps_per_block * detail::index_match_cells_per_warp;
            auto const bucket_blocks = static_cast<uint32_t>(
                std::min<size_t>((query_buckets + cells_per_block - 1U) / cells_per_block, 65535U)
            );
            detail::count_batch_index_matches_kernel<BucketCount>
                <<<bucket_blocks, detail::block_size, 0, stream.get()>>>(
                    query_scores,
                    0U,
                    query_count,
                    index_offsets_.data(),
                    index_postings_.data(),
                    reference_count,
                    indexed_bucket_count,
                    key_mask,
                    match_counts,
                    index_keys_.data()
                );
            CUDDL_CUDA_TRY(cudaGetLastError());
            CUDDL_TRY(bits_from_counts());
        }
        CUDDL_TRY(refine_selected(false));
        return result_tile;
    }
 
   private:
    [[nodiscard]] __host__ __device__ constexpr bool rows_match_metadata(
        size_t expected_scores
    ) const noexcept {
        return planes_.size() * 2U == expected_scores &&
               (expected_scores == 0U || planes_.data() != nullptr);
    }
 
    device_span<uint32_t const> planes_;
    reference_database_metadata metadata_;
    device_span<uint32_t const> index_offsets_;
    device_span<uint32_t const> index_postings_;
    device_span<uint16_t const> index_keys_;
    device_span<uint32_t const> key_directory_;
    device_span<uint32_t const> bitmap_cells_;
    device_span<uint32_t const> posting_bitmaps_;
    reference_delta_view deltas_;
    double index_pair_fraction_{};
    bool indexed_{};
};
 
}  // namespace detail
 
template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
class reference_database {
    friend class reference_index_file;
    template <uint32_t, size_t, typename>
    friend class reference_index;
    friend class reference_database_file;
    static_assert(K >= 1 && K <= 31);
    static_assert(BucketCount >= (size_t{1} << 11) && BucketCount <= (size_t{1} << 13));
    static_assert((BucketCount & (BucketCount - 1)) == 0);
 
   public:
    using layout_type = Layout;
    using score_type = uint16_t;
    using result_type = reference_search_result;
    using batch_result_type = packed_pairwise_counts;
 
    reference_database(reference_database const&) = delete;
    reference_database& operator=(reference_database const&) = delete;
 
    // An explicit body keeps NVCC from inferring device-side buffer destruction.
    __host__ ~reference_database() {}  // NOLINT(modernize-use-equals-default)
 
    reference_database(reference_database&& other) noexcept
        : planes_(std::move(other.planes_)),
          device_(other.device_),
          metadata_(std::exchange(other.metadata_, {})),
          names_(std::move(other.names_)),
          identity_(std::move(other.identity_)),
          deltas_(std::move(other.deltas_)) {}
 
    reference_database& operator=(reference_database&& other) noexcept {
        if (this != &other) {
            planes_ = std::move(other.planes_);
            device_ = other.device_;
            metadata_ = std::exchange(other.metadata_, {});
            names_ = std::move(other.names_);
            identity_ = std::move(other.identity_);
            deltas_ = std::move(other.deltas_);
        }
        return *this;
    }
 
    [[nodiscard]] static Result<reference_database> build_async(
        device_span<score_type const> rows,
        score_compatibility compatibility,
        cuda::stream_ref stream
    ) {
        detail::device_scope const scope{stream};
        auto const reference_count = CUDDL_TRY(validate_rows(rows, compatibility));
        auto database = CUDDL_CUDA_TRY(reference_database(stream));
        database.metadata_ = {.compatibility = compatibility, .reference_count = reference_count};
        if (!rows.empty()) {
            database.planes_ = CUDDL_CUDA_TRY(
                cuda::make_device_buffer<uint32_t>(
                    stream, stream.device(), rows.size() / 2U, cuda::no_init
                )
            );
            detail::build_score_planes_kernel<BucketCount>
                <<<detail::
                       warp_grid_blocks(static_cast<size_t>(reference_count) * (BucketCount / 32U)),
                   detail::block_size,
                   0,
                   stream.get()>>>(rows.data(), reference_count, database.planes_.data());
            CUDDL_CUDA_TRY(cudaGetLastError());
            // Rows are compared as base references plus changed buckets; see
            // detail::reference_delta_view. Building it waits for the stream.
            if (!detail::ablated("deltas")) {
                database.deltas_ = CUDDL_TRY(
                    detail::build_reference_deltas<BucketCount>(
                        rows.data(), reference_count, stream
                    )
                );
            }
        }
        return Result<reference_database>::ok(std::move(database));
    }
 
    [[nodiscard]] Result<void>
    copy_scores_async(device_span<score_type> scores, cuda::stream_ref stream) const {
        detail::device_scope const scope{stream};
        auto const count = static_cast<size_t>(metadata_.reference_count) * BucketCount;
        if (scores.size() < count) {
            return Err(Error::resource("score destination is too small"));
        }
        if (count == 0U) {
            return Ok();
        }
        if (scores.data() == nullptr) {
            return Err(Error::invalid_argument("score destination must be device accessible"));
        }
        detail::decode_score_planes_kernel<BucketCount>
            <<<detail::warp_grid_blocks(
                   static_cast<size_t>(metadata_.reference_count) * (BucketCount / 32U)
               ),
               detail::block_size,
               0,
               stream.get()>>>(planes_.data(), metadata_.reference_count, scores.data());
        return cuda_try(cudaGetLastError());
    }
 
    [[nodiscard]] std::span<std::string const> names() const noexcept {
        return names_;
    }
 
    [[nodiscard]] reference_database_metadata metadata() const noexcept {
        return metadata_;
    }
 
    [[nodiscard]] uint32_t reference_count() const noexcept {
        return metadata_.reference_count;
    }
 
    [[nodiscard]] cuda::device_ref device() const noexcept {
        return device_;
    }
 
    [[nodiscard]] static constexpr size_t persistent_row_bytes(uint32_t reference_count) noexcept {
        return view_type::persistent_row_bytes(reference_count);
    }
 
    [[nodiscard]] size_t persistent_row_bytes() const noexcept {
        return view().persistent_row_bytes();
    }
 
    [[nodiscard]] static constexpr size_t single_query_workspace_bytes(
        uint32_t reference_count
    ) noexcept {
        return view_type::single_query_workspace_bytes(reference_count);
    }
 
    [[nodiscard]] size_t single_query_workspace_bytes() const noexcept {
        return view().single_query_workspace_bytes();
    }
 
    [[nodiscard]] static constexpr uint32_t all_to_all_result_capacity(
        uint32_t reference_count
    ) noexcept {
        return view_type::all_to_all_result_capacity(reference_count);
    }
 
    [[nodiscard]] static constexpr uint32_t single_query_result_count(
        uint32_t reference_count
    ) noexcept {
        return view_type::single_query_result_count(reference_count);
    }
 
    [[nodiscard]] uint32_t single_query_result_count() const noexcept {
        return view().single_query_result_count();
    }
 
    [[nodiscard]] Result<void> search_async(
        device_span<score_type const> query,
        score_compatibility const& query_compatibility,
        device_span<uint8_t> workspace,
        device_span<result_type> results,
        cuda::stream_ref stream
    ) const {
        return view().search_async(query, query_compatibility, workspace, results, stream);
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_batch_async(
        device_span<score_type const> queries,
        score_compatibility const& query_compatibility,
        uint32_t query_id_offset,
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        cuda::stream_ref stream
    ) const {
        return view().search_batch_async(
            queries,
            query_compatibility,
            query_id_offset,
            workspace,
            results,
            std::forward<OnTile>(on_tile),
            result_match_counts,
            stream
        );
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_all_to_all_async(
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        cuda::stream_ref stream,
        all_to_all_shard shard = {}
    ) const {
        return view().search_all_to_all_async(
            workspace, results, std::forward<OnTile>(on_tile), result_match_counts, stream, shard
        );
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_all_to_all_async(
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        cuda::stream_ref stream,
        all_to_all_tile_queue& tiles
    ) const {
        return view().search_all_to_all_async(
            workspace, results, std::forward<OnTile>(on_tile), result_match_counts, stream, tiles
        );
    }
 
    [[nodiscard]] Result<size_t> search_workspace_bytes(
        cuda::stream_ref stream,
        reference_index<K, BucketCount, Layout> const* index = nullptr
    ) const {
        auto search_view = CUDDL_TRY(view(index));
        return search_view.indexed_single_query_workspace_bytes(stream);
    }
 
    [[nodiscard]] Result<cuddl::batch_search_requirements> batch_search_requirements(
        uint32_t query_count,
        cuda::stream_ref stream,
        reference_index<K, BucketCount, Layout> const* index = nullptr
    ) const {
        auto search_view = CUDDL_TRY(view(index));
        return search_view.indexed_batch_search_requirements(query_count, stream);
    }
 
    [[nodiscard]] Result<cuddl::batch_search_requirements> all_to_all_search_requirements(
        reference_index<K, BucketCount, Layout> const* index = nullptr
    ) const {
        auto search_view = CUDDL_TRY(view(index));
        return search_view.indexed_all_to_all_search_requirements();
    }
 
    [[nodiscard]] Result<void> search_async(
        device_span<score_type const> query,
        score_compatibility const& query_compatibility,
        device_span<uint8_t> workspace,
        device_span<result_type> results,
        device_span<uint32_t> result_count,
        search_options options,
        cuda::stream_ref stream,
        reference_index<K, BucketCount, Layout> const* index = nullptr
    ) const {
        auto search_view = CUDDL_TRY(view(index));
        return search_view.search_indexed_async(
            query, query_compatibility, workspace, results, result_count, options, stream
        );
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_batch_async(
        device_span<score_type const> queries,
        score_compatibility const& query_compatibility,
        uint32_t query_id_offset,
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        search_options options,
        cuda::stream_ref stream,
        reference_index<K, BucketCount, Layout> const* index = nullptr
    ) const {
        auto search_view = CUDDL_TRY(view(index));
        return search_view.search_batch_indexed_async(
            queries,
            query_compatibility,
            query_id_offset,
            workspace,
            results,
            std::forward<OnTile>(on_tile),
            result_match_counts,
            options,
            stream
        );
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_all_to_all_async(
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        search_options options,
        cuda::stream_ref stream,
        reference_index<K, BucketCount, Layout> const* index = nullptr,
        all_to_all_shard shard = {}
    ) const {
        auto search_view = CUDDL_TRY(view(index));
        return search_view.search_all_to_all_indexed_async(
            workspace,
            results,
            std::forward<OnTile>(on_tile),
            result_match_counts,
            options,
            stream,
            shard
        );
    }
 
    template <typename OnTile>
    [[nodiscard]] Result<void> search_all_to_all_async(
        device_span<uint8_t> workspace,
        device_span<batch_result_type> results,
        OnTile&& on_tile,
        device_span<uint32_t> result_match_counts,
        search_options options,
        cuda::stream_ref stream,
        reference_index<K, BucketCount, Layout> const* index,
        all_to_all_tile_queue& tiles
    ) const {
        auto search_view = CUDDL_TRY(view(index));
        return search_view.search_all_to_all_indexed_async(
            workspace,
            results,
            std::forward<OnTile>(on_tile),
            result_match_counts,
            options,
            stream,
            tiles
        );
    }
 
   private:
    using view_type = detail::reference_database_view<K, BucketCount, Layout>;
 
    [[nodiscard]] view_type view() const noexcept {
        return view_type(
            {planes_.data(), planes_.size()},
            metadata_,
            {},
            {},
            false,
            {},
            {},
            0.0,
            {},
            {},
            delta_view()
        );
    }
 
    [[nodiscard]] Result<view_type> view(
        reference_index<K, BucketCount, Layout> const* index
    ) const {
        if (index == nullptr) return view();
        if (!identity_ || index->identity_ != identity_ || !index->indexed_) {
            return Err(Error::invalid_argument("index does not belong to this database"));
        }
        auto const offset_count = index->index_offsets_.size();
        auto const offsets =
            device_span<uint32_t const>{index->index_offsets_.data(), offset_count};
        auto const postings = device_span<uint32_t const>{
            index->index_postings_.data(), index->index_posting_capacity_
        };
        return view_type(
            {planes_.data(), planes_.size()},
            metadata_,
            offsets,
            postings,
            index->indexed_,
            index->index_keys_,
            index->key_directory_,
            detail::ablated("index-fallback") ? 0.0 : index->pair_fraction_,
            {index->bitmap_cells_.data(), index->bitmap_count_},
            index->posting_bitmaps_,
            delta_view()
        );
    }
 
    [[nodiscard]] detail::reference_delta_view delta_view() const noexcept {
        return deltas_ ? deltas_->view() : detail::reference_delta_view{};
    }
 
    explicit reference_database(cuda::stream_ref stream)
        : planes_(stream, cuda::device_default_memory_pool(stream.device())),
          device_(stream.device()) {}
 
    [[nodiscard]] static Result<uint32_t>
    validate_rows(device_span<score_type const> rows, score_compatibility compatibility) {
        CUDDL_TRY(
            (detail::validate_indexed_score_compatibility<K, BucketCount, Layout>(compatibility))
        );
        if (rows.size() % BucketCount != 0U) {
            return Err(Error::invalid_argument("row extent must contain complete rows"));
        }
        if (!rows.empty() && rows.data() == nullptr) {
            return Err(Error::invalid_argument("rows must be device accessible"));
        }
        auto const reference_count = rows.size() / BucketCount;
        if (reference_count > std::numeric_limits<uint32_t>::max()) {
            return Err(Error::resource("reference count exceeds stable 32-bit IDs"));
        }
        return static_cast<uint32_t>(reference_count);
    }
 
    cuda::device_buffer<uint32_t> planes_;
    cuda::device_ref device_{0};
    reference_database_metadata metadata_{};
    std::vector<std::string> names_;
    // A shared identity prevents accepting an index for a different or recycled row allocation.
    std::shared_ptr<char const> identity_ = std::make_shared<char>(0);
    std::optional<detail::reference_deltas> deltas_;
};
 
}  // namespace cuddl
 
#include <cuddl/reference_index.cuh>