BloomFilter.cuh

Go to the documentation of this file.

#pragma once
 
#include <cuda/__cmath/ceil_div.h>
#include <cuda_runtime.h>
 
#include <cuda/std/bit>
#include <cuda/std/span>
#include <cuda/stream>
 
#include <cub/warp/warp_reduce.cuh>
 
#include <thrust/copy.h>
#include <thrust/detail/execution_policy.h>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/fill.h>
#include <thrust/transform_reduce.h>
 
#include <algorithm>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string_view>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "Alphabet.cuh"
#include "device_span.cuh"
#include "Fastx.hpp"
#include "hashutil.cuh"
#include "helpers.cuh"
 
namespace cusbf {
 
template <
    uint16_t K_,
    uint16_t S_,
    uint16_t M_,
    uint64_t HashCount_ = 4,
    uint64_t CudaBlockSize_ = 256,
    Alphabet Alphabet_ = DnaAlphabet>
struct Config {
    using Alphabet = Alphabet_;
 
    static constexpr uint16_t k = K_;
    static constexpr uint16_t m = M_;
    static constexpr uint16_t s = S_;
    static constexpr uint64_t hashCount = HashCount_;
    static constexpr uint64_t alphabetSize = Alphabet::symbolCount;
    static constexpr uint64_t symbolWidth = Alphabet::symbolWidth;
    static constexpr uint64_t symbolBits = cuda::std::bit_width(alphabetSize - 1);
    static constexpr uint64_t symbolMask = (uint64_t{1} << symbolBits) - 1;
    static constexpr uint64_t filterBlockBits = 256;
    static constexpr uint64_t cudaBlockSize = CudaBlockSize_;
 
    static constexpr uint64_t wordBits = 64;
    static constexpr uint64_t blockWordCount = filterBlockBits / wordBits;
    static constexpr uint64_t minimizerSpan = k - m + 1;
    static constexpr uint64_t findereSpan = k - s + 1;
    static constexpr uint64_t insertGroupSize = blockWordCount;
    static constexpr uint64_t queryGroupSize = 1;
    static constexpr uint64_t maxRunKmers = cudaBlockSize;
 
    static_assert(k > 0, "k must be positive");
    static_assert(symbolWidth > 0, "alphabet symbolWidth must be positive");
    static_assert(m > 0 && m <= k, "m must satisfy 0 < m <= k");
    static_assert(s > 0 && s <= k, "s must satisfy 0 < s <= k");
    static_assert(k * symbolBits <= 64, "k-mer must fit in one packed uint64_t");
    static_assert(m * symbolBits <= 64, "m-mer must fit in one packed uint64_t");
    static_assert(s * symbolBits <= 64, "s-mer must fit in one packed uint64_t");
    static_assert(hashCount > 0, "At least one Bloom hash is required");
    static_assert(hashCount <= 16, "This implementation provides 16 multiplicative salts");
    static_assert(filterBlockBits >= wordBits, "Filter block must contain at least one word");
    static_assert(
        cuda::std::has_single_bit(filterBlockBits),
        "Filter block size must be a power of two"
    );
    static_assert(filterBlockBits % wordBits == 0, "Filter block size must align to the word size");
    static_assert(blockWordCount <= 32, "At most one warp may cooperate on a filter block");
    static_assert(
        cuda::std::has_single_bit(blockWordCount),
        "blockWordCount must be a power of two"
    );
    static_assert(insertGroupSize <= 32, "insertGroupSize must fit in one warp");
    static_assert(queryGroupSize <= 32, "queryGroupSize must fit in one warp");
    static_assert(
        cuda::std::has_single_bit(insertGroupSize),
        "insertGroupSize must be a power of two"
    );
    static_assert(
        cuda::std::has_single_bit(queryGroupSize),
        "queryGroupSize must be a power of two"
    );
    static_assert(
        hashCount >= blockWordCount,
        "Sectorized layout requires hashCount >= blockWordCount"
    );
    static_assert(
        hashCount % blockWordCount == 0,
        "Hash count must distribute evenly across shard words"
    );
    static_assert(cudaBlockSize % 32 == 0, "CUDA block size must be a multiple of one warp");
    static_assert(
        cudaBlockSize % insertGroupSize == 0,
        "cudaBlockSize must divide insertGroupSize"
    );
    static_assert(cudaBlockSize % queryGroupSize == 0, "cudaBlockSize must divide queryGroupSize");
};
 
template <typename Config>
class Filter;
 
namespace detail {
 
template <typename T>
struct BitwiseOr {
    __host__ __device__ __forceinline__ T operator()(T lhs, T rhs) const {
        return lhs | rhs;
    }
};
 
template <typename Config>
struct SequenceKmerInput;
 
inline constexpr uint32_t kContainsSequenceStride = 4;
 
template <typename Config>
__global__ void containsSequenceKmersKernel(
    SequenceKmerInput<Config> input,
    device_span<const typename Filter<Config>::Shard> shards,
    device_span<uint8_t> output
);
 
template <typename Config>
__device__ __forceinline__ bool prepareSequenceHashTiles(
    const char* sequence,
    uint64_t blockStartKmer,
    uint64_t blockKmers,
    uint8_t* sequenceTile
);
 
template <typename Config>
__global__ void insertSequenceKmersKernel(
    SequenceKmerInput<Config> input,
    device_span<typename Filter<Config>::Shard> shards
);
 
inline constexpr uint64_t kInvalidHash = std::numeric_limits<uint64_t>::max();
template <uint64_t Index>
struct SaltLiteral;
 
template <>
struct SaltLiteral<0> {
    static constexpr uint64_t value = 0x9E37'79B9'7F4A'7C15ULL;
};
template <>
struct SaltLiteral<1> {
    static constexpr uint64_t value = 0xC2B2'AE3D'27D4'EB4FULL;
};
template <>
struct SaltLiteral<2> {
    static constexpr uint64_t value = 0x1656'67B1'9E37'79F9ULL;
};
template <>
struct SaltLiteral<3> {
    static constexpr uint64_t value = 0x85EB'CA77'C2B2'AE63ULL;
};
template <>
struct SaltLiteral<4> {
    static constexpr uint64_t value = 0x27D4'EB2F'1656'67C5ULL;
};
template <>
struct SaltLiteral<5> {
    static constexpr uint64_t value = 0x94D0'49BB'1331'11EFULL;
};
template <>
struct SaltLiteral<6> {
    static constexpr uint64_t value = 0xBF58'476D'1CE4'E5B9ULL;
};
template <>
struct SaltLiteral<7> {
    static constexpr uint64_t value = 0xD6E8'FEB8'6659'FD93ULL;
};
template <>
struct SaltLiteral<8> {
    static constexpr uint64_t value = 0xA076'1D64'78BD'642FULL;
};
template <>
struct SaltLiteral<9> {
    static constexpr uint64_t value = 0xE703'7ED1'A0B4'28DBULL;
};
template <>
struct SaltLiteral<10> {
    static constexpr uint64_t value = 0x8EBC'6AF0'9C88'C6E3ULL;
};
template <>
struct SaltLiteral<11> {
    static constexpr uint64_t value = 0x5899'65CC'7537'4CC3ULL;
};
template <>
struct SaltLiteral<12> {
    static constexpr uint64_t value = 0x1D8E'4E27'C47D'124FULL;
};
template <>
struct SaltLiteral<13> {
    static constexpr uint64_t value = 0xEB44'9C93'FBBE'A6B5ULL;
};
template <>
struct SaltLiteral<14> {
    static constexpr uint64_t value = 0xDB4F'0B91'75AE'2165ULL;
};
template <>
struct SaltLiteral<15> {
    static constexpr uint64_t value = 0xBBE0'56FD'ADE1'4B91ULL;
};
 
template <uint64_t Index>
[[nodiscard]] __host__ __device__ __forceinline__ constexpr uint64_t multiplicativeSaltLiteral() {
    static_assert(Index < 16, "Salt index out of range");
    return SaltLiteral<Index>::value;
}
 
template <typename Config, typename Fn, uint64_t... HashIndices>
__host__ __device__ __forceinline__ void
forEachHashIndexImpl(Fn&& fn, std::index_sequence<HashIndices...>) {
    (fn(std::integral_constant<uint64_t, HashIndices>{}), ...);
}
 
template <typename Config, typename Fn>
__host__ __device__ __forceinline__ void forEachHashIndex(Fn&& fn) {
    forEachHashIndexImpl<Config>(
        static_cast<Fn&&>(fn), std::make_index_sequence<Config::hashCount>{}
    );
}
 
template <typename Config, uint64_t Length>
[[nodiscard]] __host__ __device__ __forceinline__ constexpr uint64_t packedWindowMask() {
    if constexpr (Length * Config::symbolBits >= 64) {
        return std::numeric_limits<uint64_t>::max();
    } else {
        return (uint64_t{1} << (Config::symbolBits * Length)) - 1;
    }
}
 
template <typename Config, uint64_t WindowLength, uint64_t K>
[[nodiscard]] __host__ __device__ __forceinline__ constexpr uint64_t
extractPackedSubwindow(uint64_t packedKmer, uint64_t start) {
    static_assert(WindowLength <= K, "WindowLength must not exceed K");
    return (packedKmer >> (Config::symbolBits * (K - (start + WindowLength)))) &
           packedWindowMask<Config, WindowLength>();
}
 
__device__ __forceinline__ void atomicOrWord(uint64_t* ptr, uint64_t value) {
    atomicOr(reinterpret_cast<unsigned long long*>(ptr), static_cast<unsigned long long>(value));
}
 
}  // namespace detail
 
template <typename Config>
class Filter {
   private:
    struct PreparedRecordRange {
        uint64_t recordIndex{};
        uint64_t inputOffset{};
        uint64_t outputOffset{};
        uint64_t size{};
        uint64_t validKmers{};
    };
    struct PreparedRecordBatch {
        std::string sequence;
        std::vector<PreparedRecordRange> records;
    };
    struct FastxChunkAssembly {
        explicit FastxChunkAssembly(size_t reservedBytes) {
            sequence.reserve(reservedBytes);
        }
 
        void appendRecord(std::string_view recordSequence) {
            ranges.push_back(
                RecordRange{
                    static_cast<uint64_t>(sequence.size()),
                    static_cast<uint64_t>(recordSequence.size()),
                }
            );
            sequence.append(recordSequence);
        }
 
        [[nodiscard]] bool empty() const {
            return ranges.empty();
        }
 
        [[nodiscard]] bool reachedTarget(size_t targetBytes) const {
            return sequence.size() >= targetBytes;
        }
 
        [[nodiscard]] uint64_t recordCount() const {
            return static_cast<uint64_t>(ranges.size());
        }
 
        void clear() {
            sequence.clear();
            ranges.clear();
        }
 
        std::string sequence;
        std::vector<RecordRange> ranges;
    };
 
   public:
    struct alignas(32) Shard {
        static constexpr uint64_t wordCount = Config::blockWordCount;
        static constexpr uint64_t wordBits = Config::wordBits;
        static constexpr int wordBitsLog2 = cuda::std::bit_width(wordBits) - 1;
        static constexpr uint64_t wordMask = (1ULL << wordBitsLog2) - 1;
        static constexpr int hashShift = 64 - wordBitsLog2;
        static constexpr uint64_t sliceWidth = 64 / Config::hashCount;
        static constexpr bool useBitSlicing = sliceWidth >= wordBitsLog2;
 
        uint64_t words[wordCount];
 
        template <uint64_t HashIndex>
        [[nodiscard]] constexpr __host__ __device__ static uint64_t sectorizedBitAddress(
            uint64_t baseHash
        ) {
            static_assert(HashIndex < Config::hashCount, "Hash index out of range");
            // When there are enough bits in a 64-bit hash to give each hash
            // index its own slice, avoid the extra multiply and use
            // bit-slicing instead.
            if constexpr (useBitSlicing) {
                return (baseHash >> (sliceWidth * HashIndex)) & wordMask;
            } else {
                const uint64_t mixed = baseHash * detail::multiplicativeSaltLiteral<HashIndex>();
                return mixed >> hashShift;
            }
        }
 
        __device__ __forceinline__ static void sectorizedHashToMasks(
            uint64_t baseHash,
            uint64_t& mask0,
            uint64_t& mask1,
            uint64_t& mask2,
            uint64_t& mask3
        ) {
            detail::forEachHashIndex<Config>(
                [&]<uint64_t HashIndex>(std::integral_constant<uint64_t, HashIndex>) {
                    constexpr uint64_t s = HashIndex % Config::blockWordCount;
                    const uint64_t bitPos = sectorizedBitAddress<HashIndex>(baseHash);
                    const uint64_t bit = uint64_t{1} << bitPos;
                    // clang-format off
                    if      constexpr (s == 0) mask0 |= bit;
                    else if constexpr (s == 1) mask1 |= bit;
                    else if constexpr (s == 2) mask2 |= bit;
                    else                       mask3 |= bit;
                    // clang-format on
                }
            );
        }
    };
 
    static_assert(Config::blockWordCount == 4, "Filter only supports the fused 256-bit shard path");
    static_assert(
        Config::queryGroupSize == 1,
        "Fused path expects Theta=1 independent query mapping"
    );
    static_assert(
        Config::insertGroupSize == Config::blockWordCount,
        "Fused path expects horizontal insert mapping across shard words"
    );
 
    explicit Filter(uint64_t requestedFilterBits)
        : numShards_(
              cuda::std::bit_ceil(
                  std::max<uint64_t>(
                      1,
                      cuda::ceil_div(requestedFilterBits, Config::filterBlockBits)
                  )
              )
          ),
          filterBits_(numShards_ * Config::filterBlockBits),
          d_shards_(numShards_) {
        clear();
    }
 
    Filter(const Filter&) = delete;
    Filter& operator=(const Filter&) = delete;
    Filter(Filter&&) = default;
    Filter& operator=(Filter&&) = default;
    ~Filter() = default;
 
    [[nodiscard]] uint64_t
    insertSequence(std::string_view sequence, cuda::stream_ref stream = cudaStream_t{}) {
        if (recordSymbolCount(sequence.size()) < Config::k) {
            return 0;
        }
 
        const uint64_t totalKmers = recordKmerCount(sequence.size());
        const auto d_sequence = stagedSequenceView({sequence.data(), sequence.size()}, stream);
        launchInsertSequence(d_sequence, stream);
        stream.sync();
        return totalKmers;
    }
 
    [[nodiscard]] uint64_t insertSequenceDevice(
        device_span<const char> d_sequence,
        cuda::stream_ref stream = cudaStream_t{}
    ) {
        const uint64_t totalKmers = sequenceKmerCount(d_sequence);
        if (totalKmers == 0) {
            return 0;
        }
 
        launchInsertSequence(d_sequence, stream);
        return totalKmers;
    }
 
    [[nodiscard]] FastxInsertReport
    insertRecordBatch(RecordBatchView batch, cuda::stream_ref stream = cudaStream_t{}) {
        const PreparedRecordBatch prepared = prepareRecordBatch(batch);
        FastxInsertReport report;
        report.recordsIndexed = prepared.records.size();
        for (const PreparedRecordRange& record : prepared.records) {
            report.indexedBases += record.size;
            report.insertedKmers += record.validKmers;
        }
        if (!prepared.sequence.empty()) {
            (void)insertSequence(prepared.sequence, stream);
        }
        return report;
    }
 
    [[nodiscard]] FastxInsertReport insertFastx(
        std::istream& input,
        double fillFraction = 0.7,
        cuda::stream_ref stream = cudaStream_t{}
    ) {
        return insertFastxStream(input, "<stream>", fillFraction, stream);
    }
 
    [[nodiscard]] FastxInsertReport insertFastxFile(
        std::string_view path,
        double fillFraction = 0.7,
        cuda::stream_ref stream = cudaStream_t{}
    ) {
        auto input = detail::openFastxFile(path);
        return insertFastxStream(*input, path, fillFraction, stream);
    }
 
    void containsSequenceDevice(
        device_span<const char> d_sequence,
        device_span<uint8_t> d_output,
        cuda::stream_ref stream = cudaStream_t{}
    ) const {
        if (sequenceKmerCount(d_sequence) == 0) {
            return;
        }
 
        launchContainsSequence(d_sequence, d_output, stream);
    }
 
    [[nodiscard]] std::vector<uint8_t>
    containsSequence(std::string_view sequence, cuda::stream_ref stream = cudaStream_t{}) const {
        if (recordSymbolCount(sequence.size()) < Config::k) {
            return {};
        }
 
        std::vector<uint8_t> output(recordKmerCount(sequence.size()));
 
        const auto d_sequence = stagedSequenceView({sequence.data(), sequence.size()}, stream);
        ensureResultCapacity(output.size());
        launchContainsSequence(
            d_sequence,
            device_span<uint8_t>{thrust::raw_pointer_cast(d_resultBuffer_.data()), output.size()},
            stream
        );
        CUSBF_CUDA_CALL(cudaMemcpyAsync(
            output.data(),
            thrust::raw_pointer_cast(d_resultBuffer_.data()),
            output.size() * sizeof(uint8_t),
            cudaMemcpyDeviceToHost,
            stream.get()
        ));
 
        stream.sync();
        return output;
    }
 
    [[nodiscard]] FastxQueryReport
    queryRecordBatch(RecordBatchView batch, cuda::stream_ref stream = cudaStream_t{}) const {
        return queryRecordBatch(batch, [](const RecordQueryView&) {}, stream);
    }
 
    template <typename Consumer>
    [[nodiscard]] FastxQueryReport queryRecordBatch(
        RecordBatchView batch,
        Consumer&& consume,
        cuda::stream_ref stream = cudaStream_t{}
    ) const {
        return queryPreparedRecordBatch(prepareRecordBatch(batch), batch.sequence, consume, stream);
    }
 
    [[nodiscard]] FastxQueryReport queryFastx(
        std::istream& input,
        double fillFraction = 0.7,
        cuda::stream_ref stream = cudaStream_t{}
    ) const {
        return queryFastxStream(input, "<stream>", fillFraction, stream);
    }
 
    [[nodiscard]] FastxQueryReport queryFastxFile(
        std::string_view path,
        double fillFraction = 0.7,
        cuda::stream_ref stream = cudaStream_t{}
    ) const {
        auto input = detail::openFastxFile(path);
        return queryFastxStream(*input, path, fillFraction, stream);
    }
 
    template <typename Consumer>
    [[nodiscard]] FastxQueryReport queryFastxRecords(
        std::istream& input,
        Consumer&& consume,
        double fillFraction = 0.7,
        cuda::stream_ref stream = cudaStream_t{}
    ) const {
        return queryFastxRecordsStream(input, "<stream>", consume, fillFraction, stream);
    }
 
    template <typename Consumer>
    [[nodiscard]] FastxQueryReport queryFastxFileRecords(
        std::string_view path,
        Consumer&& consume,
        double fillFraction = 0.7,
        cuda::stream_ref stream = cudaStream_t{}
    ) const {
        auto input = detail::openFastxFile(path);
        return queryFastxRecordsStream(*input, path, consume, fillFraction, stream);
    }
 
    [[nodiscard]] FastxDetailedQueryReport queryFastxDetailed(
        std::istream& input,
        double fillFraction = 0.7,
        cuda::stream_ref stream = cudaStream_t{}
    ) const {
        return queryFastxDetailedStream(input, "<stream>", fillFraction, stream);
    }
 
    [[nodiscard]] FastxDetailedQueryReport queryFastxFileDetailed(
        std::string_view path,
        double fillFraction = 0.7,
        cuda::stream_ref stream = cudaStream_t{}
    ) const {
        auto input = detail::openFastxFile(path);
        return queryFastxDetailedStream(*input, path, fillFraction, stream);
    }
 
    void clear(cuda::stream_ref stream = cudaStream_t{}) {
        CUSBF_CUDA_CALL(cudaMemsetAsync(
            thrust::raw_pointer_cast(d_shards_.data()),
            0,
            d_shards_.size() * sizeof(Shard),
            stream.get()
        ));
 
        stream.sync();
    }
 
    [[nodiscard]] float loadFactor() const {
        const auto* wordsBegin =
            reinterpret_cast<const uint64_t*>(thrust::raw_pointer_cast(d_shards_.data()));
        const uint64_t totalWords = numShards_ * Config::blockWordCount;
        const uint64_t setBits = thrust::transform_reduce(
            thrust::device,
            wordsBegin,
            wordsBegin + totalWords,
            [] __device__(uint64_t w) -> uint64_t { return cuda::std::popcount(w); },
            uint64_t{0},
            cuda::std::plus<uint64_t>()
        );
        return static_cast<float>(setBits) / static_cast<float>(filterBits_);
    }
 
    [[nodiscard]] uint64_t filterBits() const {
        return filterBits_;
    }
 
    [[nodiscard]] uint64_t numShards() const {
        return numShards_;
    }
 
   private:
    uint64_t numShards_{};
    uint64_t filterBits_{};
 
    thrust::device_vector<Shard> d_shards_;
    mutable thrust::device_vector<char> d_sequence_;
    mutable thrust::device_vector<uint8_t> d_resultBuffer_;
 
    [[nodiscard]] uint64_t sizeBytes() const {
        return numShards() * sizeof(Shard);
    }
 
    [[nodiscard]] static uint64_t recordSymbolCount(uint64_t bases) {
        return bases / Config::symbolWidth;
    }
 
    [[nodiscard]] static uint64_t recordKmerCount(uint64_t bases) {
        const uint64_t symbols = recordSymbolCount(bases);
        return symbols < Config::k ? 0 : symbols - Config::k + 1;
    }
 
    [[nodiscard]] static uint64_t validRecordKmerCount(std::string_view sequence) {
        if (recordSymbolCount(sequence.size()) < Config::k) {
            return 0;
        }
 
        uint64_t invalidSymbols = 0;
        for (uint64_t i = 0; i < Config::k; ++i) {
            invalidSymbols += Config::Alphabet::encode(sequence.data() + i * Config::symbolWidth) ==
                              Config::Alphabet::invalidSymbol;
        }
 
        uint64_t validKmers = invalidSymbols == 0 ? 1 : 0;
        for (uint64_t start = 1; start < recordKmerCount(sequence.size()); ++start) {
            invalidSymbols -=
                Config::Alphabet::encode(sequence.data() + (start - 1) * Config::symbolWidth) ==
                Config::Alphabet::invalidSymbol;
            invalidSymbols += Config::Alphabet::encode(
                                  sequence.data() + (start + Config::k - 1) * Config::symbolWidth
                              ) == Config::Alphabet::invalidSymbol;
            validKmers += invalidSymbols == 0;
        }
        return validKmers;
    }
 
    static void appendRecordBoundary(std::string& sequence) {
        const uint64_t remainder = sequence.size() % Config::symbolWidth;
        if (remainder != 0) {
            sequence.append(
                Config::symbolWidth - remainder, static_cast<char>(Config::Alphabet::separator)
            );
        }
        sequence.append(Config::symbolWidth, static_cast<char>(Config::Alphabet::separator));
    }
 
    static void validateRecordBatch(RecordBatchView batch) {
        uint64_t nextOffset = 0;
        for (const RecordRange& record : batch.records) {
            if (record.sequenceOffset < nextOffset) {
                throw std::invalid_argument(
                    "record batch ranges must be ordered and non-overlapping"
                );
            }
            if (record.sequenceOffset > batch.sequence.size() ||
                record.sequenceBytes > batch.sequence.size() - record.sequenceOffset) {
                throw std::invalid_argument("record batch range exceeds the source sequence");
            }
            if (record.sequenceOffset % Config::symbolWidth != 0 ||
                record.sequenceBytes % Config::symbolWidth != 0) {
                throw std::invalid_argument(
                    "record batch ranges must align to the configured alphabet symbol width"
                );
            }
            nextOffset = record.sequenceOffset + record.sequenceBytes;
        }
    }
 
    static void appendPreparedRecord(
        std::string& output,
        std::vector<PreparedRecordRange>& ranges,
        uint64_t recordIndex,
        uint64_t inputOffset,
        std::string_view recordSequence
    ) {
        if (!output.empty()) {
            appendRecordBoundary(output);
        }
        const uint64_t outputOffset = recordSymbolCount(output.size());
        output.append(recordSequence);
        ranges.push_back(
            PreparedRecordRange{
                recordIndex,
                inputOffset,
                outputOffset,
                static_cast<uint64_t>(recordSequence.size()),
                validRecordKmerCount(recordSequence),
            }
        );
    }
 
    [[nodiscard]] static PreparedRecordBatch prepareRecordBatch(RecordBatchView batch) {
        validateRecordBatch(batch);
 
        PreparedRecordBatch prepared;
        prepared.sequence.reserve(
            batch.sequence.size() + batch.records.size() * Config::symbolWidth
        );
        prepared.records.reserve(batch.records.size());
 
        for (uint64_t recordIndex = 0; recordIndex < batch.records.size(); ++recordIndex) {
            const RecordRange& record = batch.records[recordIndex];
            appendPreparedRecord(
                prepared.sequence,
                prepared.records,
                recordIndex,
                record.sequenceOffset,
                batch.sequence.substr(record.sequenceOffset, record.sequenceBytes)
            );
        }
        return prepared;
    }
 
    [[nodiscard]] static RecordBatchView
    makeBatchView(const std::string& sequence, const std::vector<RecordRange>& ranges) {
        return RecordBatchView{
            sequence,
            cuda::std::span<const RecordRange>{ranges.data(), ranges.size()},
        };
    }
 
    static void accumulateInsertReport(FastxInsertReport& total, const FastxInsertReport& chunk) {
        total.recordsIndexed += chunk.recordsIndexed;
        total.indexedBases += chunk.indexedBases;
        total.insertedKmers += chunk.insertedKmers;
    }
 
    static void accumulateQueryReport(FastxQueryReport& total, const FastxQueryReport& chunk) {
        total.recordsQueried += chunk.recordsQueried;
        total.queriedBases += chunk.queriedBases;
        total.queriedKmers += chunk.queriedKmers;
        total.positiveKmers += chunk.positiveKmers;
    }
 
    [[nodiscard]] static size_t fastxChunkTargetBytes(double fillFraction) {
        size_t freeBytes = 0;
        size_t totalBytes = 0;
        CUSBF_CUDA_CALL(cudaMemGetInfo(&freeBytes, &totalBytes));
        return static_cast<size_t>(static_cast<double>(freeBytes) * fillFraction);
    }
 
    template <typename Consumer>
    [[nodiscard]] FastxQueryReport queryPreparedRecordBatch(
        const PreparedRecordBatch& batch,
        std::string_view inputSequence,
        Consumer&& consume,
        cuda::stream_ref stream
    ) const {
        FastxQueryReport report;
        report.recordsQueried = batch.records.size();
        for (const PreparedRecordRange& record : batch.records) {
            report.queriedBases += record.size;
            report.queriedKmers += record.validKmers;
        }
 
        const auto hits = containsSequence(batch.sequence, stream);
        for (const PreparedRecordRange& record : batch.records) {
            const uint64_t kmers = recordKmerCount(record.size);
            const auto sequence = inputSequence.substr(
                static_cast<size_t>(record.inputOffset), static_cast<size_t>(record.size)
            );
            if (kmers == 0) {
                consume(
                    RecordQueryView{
                        record.recordIndex,
                        sequence,
                        record.size,
                        record.validKmers,
                        0,
                        cuda::std::span<const uint8_t>{},
                    }
                );
                continue;
            }
 
            const auto* hitBegin = hits.data() + static_cast<ptrdiff_t>(record.outputOffset);
            const auto hitSpan =
                cuda::std::span<const uint8_t>{hitBegin, static_cast<size_t>(kmers)};
            const auto positiveKmers =
                static_cast<uint64_t>(std::count(hitSpan.begin(), hitSpan.end(), uint8_t{1}));
            report.positiveKmers += positiveKmers;
            consume(
                RecordQueryView{
                    record.recordIndex,
                    sequence,
                    record.size,
                    record.validKmers,
                    positiveKmers,
                    hitSpan,
                }
            );
        }
        return report;
    }
 
    [[nodiscard]] FastxInsertReport insertFastxStream(
        std::istream& input,
        std::string_view sourceName,
        double fillFraction,
        cuda::stream_ref stream
    ) {
        detail::FastxReader reader(input, sourceName);
        detail::FastxRecord record;
        FastxInsertReport report;
 
        const auto chunkTargetBytes = fastxChunkTargetBytes(fillFraction);
        FastxChunkAssembly chunk(chunkTargetBytes);
 
        auto flush = [&]() {
            if (chunk.empty()) {
                return;
            }
            accumulateInsertReport(
                report, insertRecordBatch(makeBatchView(chunk.sequence, chunk.ranges), stream)
            );
            chunk.clear();
        };
 
        while (reader.nextRecord(record)) {
            chunk.appendRecord(record.sequence);
            if (chunk.reachedTarget(chunkTargetBytes)) {
                flush();
            }
        }
 
        flush();
        return report;
    }
 
    [[nodiscard]] FastxQueryReport queryFastxStream(
        std::istream& input,
        std::string_view sourceName,
        double fillFraction,
        cuda::stream_ref stream
    ) const {
        return queryFastxRecordsStream(
            input, sourceName, [](const FastxRecordView&) {}, fillFraction, stream
        );
    }
 
    template <typename Consumer>
    [[nodiscard]] FastxQueryReport queryFastxRecordsStream(
        std::istream& input,
        std::string_view sourceName,
        Consumer&& consume,
        double fillFraction,
        cuda::stream_ref stream
    ) const {
        detail::FastxReader reader(input, sourceName);
        detail::FastxRecord record;
        FastxQueryReport report;
 
        const auto chunkTargetBytes = fastxChunkTargetBytes(fillFraction);
        FastxChunkAssembly chunk(chunkTargetBytes);
        std::vector<detail::FastxRecord> records;
        uint64_t recordIndexBase = 0;
 
        auto flush = [&]() {
            if (chunk.empty()) {
                return;
            }
            const FastxQueryReport chunkReport = queryRecordBatch(
                makeBatchView(chunk.sequence, chunk.ranges),
                [&](const RecordQueryView& recordView) {
                    const detail::FastxRecord& fastxRecord =
                        records[static_cast<size_t>(recordView.recordIndex)];
                    consume(
                        FastxRecordView{
                            recordIndexBase + recordView.recordIndex,
                            fastxRecord.header,
                            fastxRecord.sequence,
                            recordView.queriedBases,
                            recordView.queriedKmers,
                            recordView.positiveKmers,
                            recordView.hits,
                        }
                    );
                },
                stream
            );
            accumulateQueryReport(report, chunkReport);
            recordIndexBase += chunk.recordCount();
            chunk.clear();
            records.clear();
        };
 
        while (reader.nextRecord(record)) {
            chunk.appendRecord(record.sequence);
            records.push_back(std::move(record));
            if (chunk.reachedTarget(chunkTargetBytes)) {
                flush();
            }
        }
 
        flush();
        return report;
    }
 
    [[nodiscard]] FastxDetailedQueryReport queryFastxDetailedStream(
        std::istream& input,
        std::string_view sourceName,
        double fillFraction,
        cuda::stream_ref stream
    ) const {
        FastxDetailedQueryReport report;
        report.summary = queryFastxRecordsStream(
            input,
            sourceName,
            [&report](const FastxRecordView& record) {
                report.records.push_back(
                    FastxDetailedQueryRecord{
                        record.recordIndex,
                        std::string(record.header),
                        std::string(record.sequence),
                        record.queriedBases,
                        record.queriedKmers,
                        record.positiveKmers,
                        std::vector<uint8_t>(record.hits.begin(), record.hits.end()),
                    }
                );
            },
            fillFraction,
            stream
        );
        return report;
    }
 
    void ensureSequenceCapacity(uint64_t bases) const {
        if (bases > d_sequence_.size()) {
            d_sequence_.resize(bases);
        }
    }
 
    void ensureResultCapacity(uint64_t kmers) const {
        if (kmers > d_resultBuffer_.size()) {
            d_resultBuffer_.resize(kmers);
        }
    }
 
    void stageSequence(cuda::std::span<const char> sequence, cuda::stream_ref stream) const {
        ensureSequenceCapacity(sequence.size());
        CUSBF_CUDA_CALL(cudaMemcpyAsync(
            thrust::raw_pointer_cast(d_sequence_.data()),
            sequence.data(),
            sequence.size_bytes(),
            cudaMemcpyHostToDevice,
            stream.get()
        ));
    }
 
    [[nodiscard]] static uint64_t sequenceKmerCount(device_span<const char> d_sequence) {
        return detail::SequenceKmerInput<Config>{d_sequence}.kmerCount();
    }
 
    [[nodiscard]] device_span<const char>
    stagedSequenceView(cuda::std::span<const char> sequence, cuda::stream_ref stream) const {
        stageSequence(sequence, stream);
        return device_span<const char>{
            thrust::raw_pointer_cast(d_sequence_.data()), sequence.size()
        };
    }
 
    void launchInsertSequence(device_span<const char> d_sequence, cuda::stream_ref stream) {
        const auto input = detail::SequenceKmerInput<Config>{d_sequence};
        const uint64_t numKmers = input.kmerCount();
        if (numKmers == 0) {
            return;
        }
        const uint64_t gridSize = cuda::ceil_div(numKmers, Config::cudaBlockSize);
 
        detail::insertSequenceKmersKernel<Config>
            <<<gridSize, Config::cudaBlockSize, 0, stream.get()>>>(
                input, device_span<Shard>{thrust::raw_pointer_cast(d_shards_.data()), numShards_}
            );
        CUSBF_CUDA_CALL(cudaGetLastError());
    }
 
    void launchContainsSequence(
        device_span<const char> d_sequence,
        device_span<uint8_t> d_output,
        cuda::stream_ref stream
    ) const {
        const auto input = detail::SequenceKmerInput<Config>{d_sequence};
        const uint64_t numKmers = input.kmerCount();
        const uint64_t gridSize =
            cuda::ceil_div(numKmers, Config::cudaBlockSize * detail::kContainsSequenceStride);
 
        detail::containsSequenceKmersKernel<Config>
            <<<gridSize, Config::cudaBlockSize, 0, stream.get()>>>(
                input,
                device_span<const Shard>{thrust::raw_pointer_cast(d_shards_.data()), numShards_},
                d_output
            );
        CUSBF_CUDA_CALL(cudaGetLastError());
    }
};
 
namespace detail {
 
template <typename Config>
struct SequenceKmerInput {
    device_span<const char> sequence;
 
    [[nodiscard]] constexpr __host__ __device__ uint64_t kmerCount() const {
        const uint64_t symbols = sequence.size() / Config::symbolWidth;
        return symbols < Config::k ? 0 : (symbols - Config::k + 1);
    }
 
    [[nodiscard]] constexpr __host__ __device__ uint64_t smerCount() const {
        const uint64_t symbols = sequence.size() / Config::symbolWidth;
        return symbols < Config::s ? 0 : (symbols - Config::s + 1);
    }
};
 
template <typename Config>
[[nodiscard]] __device__ __forceinline__ uint64_t packedKmerMinimizerHash(uint64_t packedKmer) {
    uint64_t minimizerHash = kInvalidHash;
    _Pragma("unroll")
    for (uint64_t offset = 0; offset < Config::minimizerSpan; ++offset) {
        const uint64_t packedMmer =
            extractPackedSubwindow<Config, Config::m, Config::k>(packedKmer, offset);
        minimizerHash = min(minimizerHash, detail::minimizerHash64(packedMmer));
    }
    return minimizerHash;
}
 
template <typename Config>
[[nodiscard]] __device__ __forceinline__ uint64_t
packedKmerSmerHash(uint64_t packedKmer, uint64_t start) {
    const uint64_t packedSmer =
        extractPackedSubwindow<Config, Config::s, Config::k>(packedKmer, start);
    return detail::hash64(packedSmer);
}
 
template <typename Config>
__device__ __forceinline__ void
loadShardWords4(const typename Filter<Config>::Shard* shards, uint64_t shardIndex, uint64_t* w) {
#if __CUDA_ARCH__ >= 1000
    detail::load256BitGlobalNC(shards[shardIndex].words, w[0], w[1], w[2], w[3]);
#else
    detail::load128BitGlobalNC(shards[shardIndex].words + 0, w[0], w[1]);
    detail::load128BitGlobalNC(shards[shardIndex].words + 2, w[2], w[3]);
#endif
}
 
template <typename Config, uint64_t K>
__device__ __forceinline__ uint64_t packKmerFromTile(const uint8_t* tile, uint64_t start) {
    uint64_t packed = 0;
    _Pragma("unroll")
    for (uint64_t i = 0; i < K; ++i) {
        packed = (packed << Config::symbolBits) | (tile[start + i] & Config::symbolMask);
    }
    return packed;
}
 
template <typename Config, uint64_t K>
__device__ __forceinline__ uint64_t advancePackedKmer(uint64_t packed, uint8_t newBase) {
    return ((packed << Config::symbolBits) | (newBase & Config::symbolMask)) &
           packedWindowMask<Config, K>();
}
 
template <typename Config>
__device__ __forceinline__ bool
sectorizedContainsPackedKmer(uint64_t packedKmer, const uint64_t* w) {
    bool present = true;
    _Pragma("unroll")
    for (uint64_t smerOffset = 0; smerOffset < Config::findereSpan; ++smerOffset) {
        const uint64_t smerHash = packedKmerSmerHash<Config>(packedKmer, smerOffset);
        detail::forEachHashIndex<Config>(
            [&]<uint64_t HashIndex>(std::integral_constant<uint64_t, HashIndex>) {
                constexpr uint64_t s = HashIndex % Config::blockWordCount;
                const uint64_t bitPos =
                    Filter<Config>::Shard::template sectorizedBitAddress<HashIndex>(smerHash);
                present &= ((w[s] >> bitPos) & 1) != 0;
            }
        );
    }
    return present;
}
 
template <typename Config>
__device__ __forceinline__ bool kmerIsValid(const uint8_t* tile, uint64_t start) {
    _Pragma("unroll")
    for (uint64_t i = 0; i < Config::k; ++i) {
        if (tile[start + i] == Config::Alphabet::invalidSymbol) {
            return false;
        }
    }
    return true;
}
 
template <typename Config>
__device__ __forceinline__ bool prepareSequenceHashTiles(
    const char* sequence,
    uint64_t blockStartKmer,
    uint64_t blockKmers,
    uint8_t* sequenceTile
) {
    const uint64_t tileBases = blockKmers + Config::k - 1;
 
    bool localInvalidBase = false;
    for (uint64_t idx = threadIdx.x; idx < tileBases; idx += Config::cudaBlockSize) {
        const uint8_t encodedBase =
            Config::Alphabet::encode(sequence + (blockStartKmer + idx) * Config::symbolWidth);
        sequenceTile[idx] = encodedBase;
        localInvalidBase |= (encodedBase == Config::Alphabet::invalidSymbol);
    }
    return __syncthreads_count(localInvalidBase) == 0;
}
 
template <typename Config>
__global__ __launch_bounds__(Config::cudaBlockSize, 6) void containsSequenceKmersKernel(
    SequenceKmerInput<Config> input,
    device_span<const typename Filter<Config>::Shard> shards,
    device_span<uint8_t> output
) {
    // Each thread handles this many consecutive k-mers to amortise packing
    constexpr uint32_t kStride = kContainsSequenceStride;
    constexpr uint64_t sequenceTileBases = Config::cudaBlockSize * kStride + Config::k - 1;
 
    __shared__ uint8_t sequenceTile[sequenceTileBases];
 
    const uint64_t numKmers = input.kmerCount();
    const uint64_t blockStartKmer =
        static_cast<uint64_t>(blockIdx.x) * Config::cudaBlockSize * kStride;
    if (blockStartKmer >= numKmers) {
        return;
    }
 
    const uint64_t blockKmers = min(Config::cudaBlockSize * kStride, numKmers - blockStartKmer);
 
    const bool blockAllValid = prepareSequenceHashTiles<Config>(
        input.sequence.data(), blockStartKmer, blockKmers, sequenceTile
    );
 
    const uint64_t threadOffset = static_cast<uint64_t>(threadIdx.x) * kStride;
    if (threadOffset >= blockKmers) {
        return;
    }
 
    // Bitmask: bit s set = k-mer at offset s is valid.
    uint32_t kmerValidMask = 0;
    _Pragma("unroll")
    for (uint32_t s = 0; s < kStride; ++s) {
        if ((threadOffset + s) < blockKmers) {
            kmerValidMask |= (1u << s);
        }
    }
 
    if (!blockAllValid) {
        _Pragma("unroll")
        for (uint32_t s = 0; s < kStride; ++s) {
            if (!(kmerValidMask & (1u << s))) {
                continue;
            }
            const uint64_t localIdx = threadOffset + s;
            if (!kmerIsValid<Config>(sequenceTile, localIdx)) {
                kmerValidMask &= ~(1u << s);
            }
        }
    }
 
    // Always pack from position 0.  Sliding propagates the packed value forward
    // invalid bases from earlier k-mers are simply shifted out.
    uint64_t packedKmer = packKmerFromTile<Config, Config::k>(sequenceTile, threadOffset);
 
    for (uint32_t s = 0; s < kStride; ++s) {
        const uint64_t localIdx = threadOffset + s;
        if (localIdx >= blockKmers) {
            break;
        }
 
        const uint64_t kmerIndex = blockStartKmer + localIdx;
 
        if (s > 0) {
            packedKmer = advancePackedKmer<Config, Config::k>(
                packedKmer, sequenceTile[localIdx + Config::k - 1]
            );
        }
 
        if (!(kmerValidMask & (1u << s))) {
            output[kmerIndex] = 0;
            continue;
        }
 
        const uint64_t minimizerHash = packedKmerMinimizerHash<Config>(packedKmer);
 
        // Warp-level shard sharing.
        const auto shardIdx = static_cast<uint32_t>(minimizerHash & (shards.size() - 1));
        const uint32_t peers = __match_any_sync(0xFFFFFFFFu, shardIdx);
        const int leader = __ffs(static_cast<int>(peers)) - 1;
 
        uint64_t w[4];
        if (static_cast<int>(threadIdx.x & 31u) == leader) {
            loadShardWords4<Config>(shards.data(), shardIdx, w);
        }
        w[0] = __shfl_sync(peers, w[0], leader);
        w[1] = __shfl_sync(peers, w[1], leader);
        w[2] = __shfl_sync(peers, w[2], leader);
        w[3] = __shfl_sync(peers, w[3], leader);
 
        const bool present = sectorizedContainsPackedKmer<Config>(packedKmer, w);
        output[kmerIndex] = present;
    }
}
 
template <typename Config>
__global__ void insertSequenceKmersKernel(
    SequenceKmerInput<Config> input,
    device_span<typename Filter<Config>::Shard> shards
) {
    constexpr uint64_t sequenceTileBases = Config::cudaBlockSize + Config::k - 1;
    constexpr uint32_t warpSize = 32;
    constexpr uint32_t warpsPerBlock = Config::cudaBlockSize / warpSize;
 
    using WarpReduceWord = cub::WarpReduce<uint64_t>;
 
    __shared__ uint8_t sequenceTile[sequenceTileBases];
    __shared__ typename WarpReduceWord::TempStorage reduceStorage[warpsPerBlock][4];
 
    const uint64_t numKmers = input.kmerCount();
    const uint64_t blockStartKmer = static_cast<uint64_t>(blockIdx.x) * Config::cudaBlockSize;
    if (blockStartKmer >= numKmers) {
        return;
    }
 
    const uint64_t blockKmers = min(Config::cudaBlockSize, numKmers - blockStartKmer);
    const auto localKmerIndex = static_cast<uint64_t>(threadIdx.x);
    const bool inRange = localKmerIndex < blockKmers;
 
    const bool blockAllValid = prepareSequenceHashTiles<Config>(
        input.sequence.data(), blockStartKmer, blockKmers, sequenceTile
    );
 
    // Avoid early returns so all warp lanes can participate in the segmented
    // warp reductions below.
    bool active = inRange;
 
    if (active && !blockAllValid) {
        active = kmerIsValid<Config>(sequenceTile, localKmerIndex);
    }
 
    // Inactive threads keep zero masks and a per-lane sentinel shard index so
    // contiguous run detection naturally splits around them.
    uint64_t minimizerHash = 0;
    uint64_t wordMask0 = 0;
    uint64_t wordMask1 = 0;
    uint64_t wordMask2 = 0;
    uint64_t wordMask3 = 0;
 
    if (active) {
        const uint64_t packedKmer =
            packKmerFromTile<Config, Config::k>(sequenceTile, localKmerIndex);
        minimizerHash = packedKmerMinimizerHash<Config>(packedKmer);
 
        uint64_t h_s = packedKmerSmerHash<Config>(packedKmer, 0);
        Filter<Config>::Shard::sectorizedHashToMasks(
            h_s, wordMask0, wordMask1, wordMask2, wordMask3
        );
        _Pragma("unroll")
        for (uint64_t smerOffset = 1; smerOffset < Config::findereSpan; ++smerOffset) {
            h_s = packedKmerSmerHash<Config>(packedKmer, smerOffset);
            Filter<Config>::Shard::sectorizedHashToMasks(
                h_s, wordMask0, wordMask1, wordMask2, wordMask3
            );
        }
    }
 
    // Warp-local segmented reductions: contiguous threads sharing the same
    // shard merge their masks so only the run head issues the atomicOrs.
    const auto shardIdx =
        static_cast<uint32_t>(active ? (minimizerHash & (shards.size() - 1)) : ~threadIdx.x);
 
    const uint32_t lane = threadIdx.x & (warpSize - 1);
    const uint32_t warpIdx = threadIdx.x / warpSize;
    const uint32_t prevShardIdx = __shfl_up_sync(0xffffffff, shardIdx, 1);
    const bool runHead = (lane == 0) || (shardIdx != prevShardIdx);
    const BitwiseOr<uint64_t> bitwiseOr{};
 
    wordMask0 = WarpReduceWord(reduceStorage[warpIdx][0])
                    .HeadSegmentedReduce(wordMask0, runHead, bitwiseOr);
    wordMask1 = WarpReduceWord(reduceStorage[warpIdx][1])
                    .HeadSegmentedReduce(wordMask1, runHead, bitwiseOr);
    wordMask2 = WarpReduceWord(reduceStorage[warpIdx][2])
                    .HeadSegmentedReduce(wordMask2, runHead, bitwiseOr);
    wordMask3 = WarpReduceWord(reduceStorage[warpIdx][3])
                    .HeadSegmentedReduce(wordMask3, runHead, bitwiseOr);
 
    if (runHead && active) {
        auto& shard = shards[shardIdx];
        if (wordMask0 != 0) {
            atomicOrWord(&shard.words[0], wordMask0);
        }
        if (wordMask1 != 0) {
            atomicOrWord(&shard.words[1], wordMask1);
        }
        if (wordMask2 != 0) {
            atomicOrWord(&shard.words[2], wordMask2);
        }
        if (wordMask3 != 0) {
            atomicOrWord(&shard.words[3], wordMask3);
        }
    }
}
 
}  // namespace detail
 
}  // namespace cusbf