dense_packed.cuh

Go to the documentation of this file.

#pragma once
 
#include <cuda/__cmath/ceil_div.h>
 
#include <cstdint>
 
#include <cusbf/config.cuh>
#include <cusbf/device_span.cuh>
 
namespace cusbf::detail {
 
template <typename Config>
[[nodiscard]] constexpr uint64_t dense_packed_symbols_per_word() {
    return 64 / Config::symbolBits;
}
 
template <typename Config>
[[nodiscard]] constexpr uint64_t dense_packed_word_count(uint64_t num_symbols) {
    return cuda::ceil_div(num_symbols, dense_packed_symbols_per_word<Config>());
}
 
template <typename Config>
[[nodiscard]] constexpr uint64_t dense_packed_kmer_count(uint64_t num_symbols) {
    return num_symbols < Config::k ? 0 : (num_symbols - Config::k + 1);
}
 
template <typename Config>
struct DensePackedKmerInput {
    device_span<const uint64_t> words{};
    uint64_t num_symbols{};
 
    [[nodiscard]] constexpr __host__ __device__ uint64_t kmerCount() const {
        return dense_packed_kmer_count<Config>(num_symbols);
    }
};
 
template <typename Config>
[[nodiscard]] __device__ __forceinline__ uint8_t
dense_packed_symbol_at(const uint64_t* words, uint64_t symbol_index) {
    constexpr uint64_t symbols_per_word = dense_packed_symbols_per_word<Config>();
    const uint64_t word_index = symbol_index / symbols_per_word;
    const auto bit_offset =
        static_cast<unsigned>((symbol_index % symbols_per_word) * Config::symbolBits);
    return static_cast<uint8_t>((words[word_index] >> bit_offset) & Config::symbolMask);
}
 
template <typename Config>
[[nodiscard]] __device__ __forceinline__ uint8_t dense_packed_symbol_at_local(
    const uint64_t* word_tile,
    uint64_t first_word_index,
    uint64_t symbol_index
) {
    constexpr uint64_t symbols_per_word = dense_packed_symbols_per_word<Config>();
    const uint64_t local_word = (symbol_index / symbols_per_word) - first_word_index;
    const auto bit_offset =
        static_cast<unsigned>((symbol_index % symbols_per_word) * Config::symbolBits);
    return static_cast<uint8_t>((word_tile[local_word] >> bit_offset) & Config::symbolMask);
}
 
template <typename Config>
__device__ __forceinline__ bool prepare_dense_packed_tiles(
    const uint64_t* words,
    uint64_t block_start_kmer,
    uint64_t block_kmers,
    uint64_t* word_tile,
    uint8_t* sequence_tile
) {
    constexpr uint64_t symbols_per_word = dense_packed_symbols_per_word<Config>();
    const uint64_t tile_symbols = block_kmers + Config::k - 1;
    const uint64_t first_symbol = block_start_kmer;
    const uint64_t last_symbol = block_start_kmer + tile_symbols - 1;
    const uint64_t first_word = first_symbol / symbols_per_word;
    const uint64_t last_word = last_symbol / symbols_per_word;
    const uint64_t num_words = last_word - first_word + 1;
 
    for (uint64_t idx = threadIdx.x; idx < num_words; idx += Config::cudaBlockSize) {
        word_tile[idx] = words[first_word + idx];
    }
    __syncthreads();
 
    bool local_invalid_symbol = false;
    for (uint64_t idx = threadIdx.x; idx < tile_symbols; idx += Config::cudaBlockSize) {
        const uint8_t encoded =
            dense_packed_symbol_at_local<Config>(word_tile, first_word, first_symbol + idx);
        sequence_tile[idx] = encoded;
        local_invalid_symbol |= (encoded == Config::Alphabet::invalidSymbol);
    }
    return __syncthreads_count(local_invalid_symbol) == 0;
}
 
template <typename Config>
__global__ void
pack_dense_sequence_kernel(const char* sequence, uint64_t num_symbols, uint64_t* words) {
    constexpr uint64_t symbols_per_word = dense_packed_symbols_per_word<Config>();
    const uint64_t word_index = static_cast<uint64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    const uint64_t num_words = dense_packed_word_count<Config>(num_symbols);
    if (word_index >= num_words) {
        return;
    }
 
    const uint64_t first_symbol = word_index * symbols_per_word;
    uint64_t packed_word = 0;
    for (uint64_t offset = 0; offset < symbols_per_word; ++offset) {
        const uint64_t symbol_index = first_symbol + offset;
        if (symbol_index >= num_symbols) {
            break;
        }
        const uint8_t symbol =
            Config::Alphabet::encode(sequence + symbol_index * Config::symbolWidth);
        packed_word |=
            (static_cast<uint64_t>(symbol & Config::symbolMask) << (offset * Config::symbolBits));
    }
    words[word_index] = packed_word;
}
 
}  // namespace cusbf::detail