sketch.cuh

Go to the documentation of this file.

#pragma once
 
#include <cub/version.cuh>
#include <cuda/buffer>
#include <cuda/memory_pool>
#include <cuda/std/cstdint>
#include <cuda/stream>
 
#include <algorithm>
#include <cmath>
#include <memory>
#include <optional>
#include <utility>
#include <vector>
 
#include <cuddl/detail/pinned_pool.cuh>
#include <cuddl/detail/sketch_view.cuh>
#include <cuddl/error.hpp>
#include <cuddl/hybrid_cardinality.cuh>
#include <cuddl/pairwise_counts.cuh>
 
namespace cuddl {
 
static_assert(
    CUB_VERSION == 300600,
    "cuDDL requires CUB 3.6.0 from CCCL revision 486de1c44daf7d1a343cbbf0f5477e9bea8ad600; "
    "use the Meson cccl subproject."
);
 
namespace detail {
 
struct cardinality_scratch {
    uint64_t empty;
    double estimate;
};
 
struct sketch_scratch {
    cardinality_scratch cardinality;
    hybrid_cardinality_estimates hybrid;
    pairwise_summary summary;
};
 
}  // namespace detail
 
template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
class sketch {
   public:
    using layout_type = Layout;
    using register_type = uint32_t;
 
    explicit sketch(
        cuda::stream_ref stream,
        std::optional<std::shared_ptr<device_blacklist const>> blacklist = std::nullopt
    )
        : blacklist_(std::move(blacklist)),
          storage_(
              cuda::make_device_buffer<register_type>(
                  stream,
                  stream.device(),
                  BucketCount,
                  register_type{0}
              )
          ),
          mapped_scratch_(
              stream,
              detail::device_pinned_pool(stream.device(), detail::pinned_placement::device_node),
              1,
              cuda::no_init
          ) {
        if (blacklist_) {
            if (!*blacklist_) throw std::invalid_argument("blacklist owner must not be null");
            CUDDL_UNWRAP((*blacklist_)->validate(K, BucketCount, stream));
        }
        stream.sync();
    }
 
    sketch(sketch const&) = delete;
    sketch& operator=(sketch const&) = delete;
    sketch(sketch&&) noexcept = default;
    sketch& operator=(sketch&&) noexcept = default;
 
    [[nodiscard]] device_span<register_type> data() noexcept {
        added_ = std::nullopt;
        return {storage_.data(), BucketCount};
    }
 
    [[nodiscard]] device_span<register_type const> data() const noexcept {
        return {storage_.data(), BucketCount};
    }
 
    [[nodiscard]] static constexpr size_t bucket_count() noexcept {
        return BucketCount;
    }
 
    [[nodiscard]] static constexpr uint32_t kmer_length() noexcept {
        return K;
    }
 
    [[nodiscard]] std::optional<uint64_t> added() const noexcept {
        return added_;
    }
 
    [[nodiscard]] Result<void> clear_async(cuda::stream_ref stream) const noexcept {
        if (auto const result = view().clear_async(stream); !result) {
            return result;
        }
        added_ = 0;
        return {};
    }
 
    [[nodiscard]] Result<void> clear(cuda::stream_ref stream) const noexcept {
        if (auto const result = clear_async(stream); !result) {
            return result;
        }
        return cuda_try([&] { stream.sync(); });
    }
 
    [[nodiscard]] Result<void> assign_async(
        device_span<register_type const> registers,
        cuda::stream_ref stream,
        kmer_blacklist const& source_blacklist
    ) const noexcept {
        if (source_blacklist.identity() != blacklist_identity() ||
            (!source_blacklist.keys().empty() && source_blacklist.kmer_length() != K)) {
            return Err(Error::invalid_argument("assigned registers use a different blacklist"));
        }
        if (registers.size() != BucketCount) {
            return Err(Error::invalid_argument("sketch assign needs one register per bucket"));
        }
        auto* const destination = storage_.data();
        if (auto const result = cuda_try([&] {
                cuda::copy_bytes(
                    stream,
                    cuda::std::span{registers.data(), registers.size()},
                    cuda::std::span{destination, BucketCount}
                );
            });
            !result) {
            return result;
        }
        added_ = std::nullopt;
        return {};
    }
 
    [[nodiscard]] Result<void>
    add_async(device_span<uint64_t const> input, cuda::stream_ref stream) const noexcept {
        if (auto const result = view().add_async(input, stream); !result) {
            return result;
        }
        if (added_.has_value()) {
            *added_ += input.size();
        }
        return {};
    }
 
    [[nodiscard]] Result<void>
    add(device_span<uint64_t const> input, cuda::stream_ref stream) const noexcept {
        if (auto const result = add_async(input, stream); !result) {
            return result;
        }
        return cuda_try([&] { stream.sync(); });
    }
 
    [[nodiscard]] Result<void>
    add_sequence_async(device_span<char const> sequence, cuda::stream_ref stream) const noexcept {
        if (auto const result = view().add_sequence_async(sequence, stream); !result) {
            return result;
        }
        if (added_.has_value() && sequence.size() >= K) {
            *added_ += sequence.size() - K + 1;
        }
        return {};
    }
 
    [[nodiscard]] Result<void>
    add_sequence(device_span<char const> sequence, cuda::stream_ref stream) const noexcept {
        if (auto const result = add_sequence_async(sequence, stream); !result) {
            return result;
        }
        return cuda_try([&] { stream.sync(); });
    }
 
    template <bool IncludeCardinality = false>
    [[nodiscard]] Result<void> summary_async(
        sketch const& other,
        pairwise_summary& output,
        cuda::stream_ref stream
    ) const noexcept {
        if (blacklist_identity() != other.blacklist_identity()) {
            return Err(Error::invalid_argument("sketch blacklist mismatch"));
        }
        return view().template summary_async<IncludeCardinality>(other.view(), output, stream);
    }
 
    [[nodiscard]] Result<void> compare_async(
        sketch const& other,
        pairwise_summary& output,
        cuda::stream_ref stream
    ) const noexcept {
        return summary_async(other, output, stream);
    }
 
    template <bool IncludeCardinality = false>
    [[nodiscard]] Result<pairwise_summary>
    summary(sketch const& other, cuda::stream_ref stream) const {
        auto* const output = &mapped_scratch_.data()->summary;
        if (auto const result = summary_async<IncludeCardinality>(other, *output, stream);
            !result) {
            return Err(result.error());
        }
        if (auto const result = cuda_try([&] { stream.sync(); }); !result) {
            return Err(result.error());
        }
        return mapped_scratch_.data()->summary;
    }
 
    [[nodiscard]] Result<pairwise_summary>
    compare(sketch const& other, cuda::stream_ref stream) const {
        return summary(other, stream);
    }
 
    [[nodiscard]] static std::optional<double> wkid(pairwise_summary const& summary) noexcept {
        auto const equal = summary.counts.equal;
        auto const divisor = equal + std::min(summary.counts.lower, summary.counts.higher);
        if (divisor == 0U) {
            return std::nullopt;
        }
        return static_cast<double>(equal) / static_cast<double>(std::max(divisor, 6U));
    }
 
    [[nodiscard]] static std::optional<double> ani(pairwise_summary const& summary) noexcept {
        auto const identity = wkid(summary);
        if (!identity) {
            return std::nullopt;
        }
        return std::pow(*identity, 1.0 / static_cast<double>(K));
    }
 
    [[nodiscard]] static std::optional<double> containment(
        pairwise_summary const& summary
    ) noexcept {
        auto const divisor = summary.counts.equal + summary.counts.higher;
        if (divisor == 0U) {
            return std::nullopt;
        }
        return static_cast<double>(summary.counts.equal) / static_cast<double>(divisor);
    }
 
    [[nodiscard]] static std::optional<double> completeness(
        pairwise_summary const& summary
    ) noexcept {
        auto const divisor = summary.counts.equal + summary.counts.lower;
        if (divisor == 0U) {
            return std::nullopt;
        }
        auto const value = static_cast<double>(summary.counts.equal + summary.counts.higher) /
                           static_cast<double>(divisor);
        return std::clamp(value, 0.0, 1.0);
    }
 
    [[nodiscard]] Result<void> cardinality_async(
        uint64_t* empty_out,
        double* estimate_out,
        cuda::stream_ref stream
    ) const noexcept {
        return view().cardinality_async(empty_out, estimate_out, stream);
    }
 
    [[nodiscard]] Result<double> cardinality(cuda::stream_ref stream) const {
        auto* const output = &mapped_scratch_.data()->cardinality;
        if (auto const result = view().cardinality_async(&output->empty, &output->estimate, stream);
            !result) {
            return Err(result.error());
        }
        if (auto const result = cuda_try([&] { stream.sync(); }); !result) {
            return Err(result.error());
        }
        return capped(mapped_scratch_.data()->cardinality.estimate);
    }
 
    [[nodiscard]] Result<hybrid_cardinality_estimates> hybrid_cardinality(
        cuda::stream_ref stream
    ) const {
        auto* const output = &mapped_scratch_.data()->hybrid;
        if (auto const result = view().hybrid_cardinality_async(output, stream); !result) {
            return Err(result.error());
        }
        if (auto const result = cuda_try([&] { stream.sync(); }); !result) {
            return Err(result.error());
        }
        auto estimates = mapped_scratch_.data()->hybrid;
        estimates.bbtools = capped(estimates.bbtools);
        estimates.paper = capped(estimates.paper);
        return estimates;
    }
 
   private:
    using view_type = detail::sketch_view<K, BucketCount, Layout>;
 
    [[nodiscard]] view_type view() const noexcept {
        return view_type(
            {storage_.data(), BucketCount}, blacklist_ ? (*blacklist_)->view() : cuda::std::nullopt
        );
    }
 
    [[nodiscard]] double capped(double estimate) const noexcept {
        return added_.has_value() ? std::min(estimate, static_cast<double>(*added_)) : estimate;
    }
 
    [[nodiscard]] uint64_t blacklist_identity() const noexcept {
        return blacklist_ ? (*blacklist_)->source().identity() : 0;
    }
 
    std::optional<std::shared_ptr<device_blacklist const>> blacklist_;
    mutable cuda::device_buffer<register_type> storage_;
    mutable cuda::
        buffer<detail::sketch_scratch, cuda::mr::host_accessible, cuda::mr::device_accessible>
            mapped_scratch_;
    mutable std::optional<uint64_t> added_ = 0;
};
 
}  // namespace cuddl