cuddl::top_k_search Class Reference

Keeps each genome's k most similar genomes from search tiles on the device. More...

Public Member Functions
Result< void >	add (batch_result_tile const &tile, cuda::stream_ref stream)
	Merges tile's passing pairs into the running neighbours, on stream.
Result< std::vector< batch_search_result > >	download (cuda::stream_ref stream) const
	Each genome's neighbours, best first, genome by genome; waits for stream.
uint32_t	k () const noexcept

Static Public Member Functions
static Result< top_k_search >	make (uint32_t genome_count, uint32_t reference_count, uint32_t bucket_count, uint32_t k, cuda::stream_ref stream, uint32_t first_genome_id=0U)

Detailed Description

Keeps each genome's k most similar genomes from search tiles on the device.

Pass add as the tile consumer of search_batch_async or search_all_to_all_async, for either the exhaustive or the indexed search, in place of downloading every passing pair. A genome's neighbours rank by weighted k-mer identity, the smaller ID first on a tie; a genome with fewer than k passing pairs keeps all of them.

Batch searches keep each query's best references, and all-to-all searches each database row's best other rows, in both directions of the stored upper triangle. Each genome keeps a min-heap of its k best. Per tile, one warp per query row merges the row's passing pairs into its genome's heap; for all-to-all tiles the pass bits are then transposed, and one warp per later genome merges its column of mirrored pairs. Each merging warp holds its genome's heap in shared memory and alone updates it, so no update takes a lock, and the work follows the passing pairs.

A collector sees only the tiles given to it. Replicas sharing an all-to-all tile queue each keep their own, and select_top_k merges their downloads.

Definition at line 317 of file top_k.cuh.

Member Function Documentation

add()

Result< void > cuddl::top_k_search::add ( batch_result_tile const &  tile, cuda::stream_ref  stream  )

inline

Merges tile's passing pairs into the running neighbours, on stream.

The tile is one a search passes to its consumer, with results in their device slots.

Definition at line 371 of file top_k.cuh.

                                                                                         {
        detail::device_scope const scope{stream};
        if (tile.reference_count != reference_count_) {
            return Err(Error::invalid_argument("tile belongs to a database of another size"));
        }
        if (tile.query_count > genome_count_ || tile.first_query_id < first_genome_id_ ||
            tile.first_query_id - first_genome_id_ > genome_count_ - tile.query_count) {
            return Err(Error::invalid_argument("tile rows lie outside the collected genomes"));
        }
        if (tile.word_offsets != nullptr) {
            return Err(Error::invalid_argument("top-k search takes device tiles, not host copies"));
        }
        if (tile.upper_triangle && (first_genome_id_ != 0U || genome_count_ != reference_count_)) {
            return Err(
                Error::invalid_argument("an all-to-all tile needs a collector of every row")
            );
        }
        if (tile.query_count == 0U || reference_count_ == 0U) return Ok();
        detail::top_k_slots const slots{keys_.data(), counts_.data(), k_, first_genome_id_};
        auto const thresholded = tile.pass_bits != nullptr;
        auto const heap_bytes = detail::top_k_heap_words(k_) * sizeof(uint64_t);
        auto const rows = tile.upper_triangle
                              ? (thresholded ? detail::top_k_rows_kernel<true, true>
                                             : detail::top_k_rows_kernel<false, true>)
                              : (thresholded ? detail::top_k_rows_kernel<true, false>
                                             : detail::top_k_rows_kernel<false, false>);
        auto const shared = warps_ * heap_bytes;
        rows<<<(tile.query_count + warps_ - 1U) / warps_, warps_ * 32U, shared, stream.get()>>>(
            tile, slots
        );
        CUDDL_CUDA_TRY(cudaGetLastError());
        if (!tile.upper_triangle) return Ok();
        // Genomes after the tile's first row take its rows before them as candidates.
        auto const first_genome = tile.first_query_id + 1U;
        if (first_genome >= reference_count_) return Ok();
        auto const first_word = first_genome / 32U;
        auto const row_words = (tile.query_count + 31U) / 32U;
        auto const needed = static_cast<size_t>(row_words) * reference_count_;
        if (columns_.size() < needed) {
            columns_ = CUDDL_CUDA_TRY(
                cuda::make_device_buffer<uint32_t>(stream, stream.device(), needed, cuda::no_init)
            );
        }
        // One warp per 32-by-32 block of the pass bits from the first genome's word on.
        auto const bit_blocks =
            static_cast<uint64_t>(row_words) * (tile.words_per_query() - first_word);
        auto const warps_per_block = detail::block_size / 32U;
        (thresholded ? detail::top_k_transpose_kernel<true>
                     : detail::top_k_transpose_kernel<false>)<<<
            static_cast<uint32_t>((bit_blocks + warps_per_block - 1U) / warps_per_block),
            detail::block_size,
            0,
            stream.get()>>>(tile, columns_.data(), first_word);
        CUDDL_CUDA_TRY(cudaGetLastError());
        auto const genomes = reference_count_ - first_genome;
        detail::top_k_columns_kernel<<<
            (genomes + warps_ - 1U) / warps_,
            warps_ * 32U,
            shared,
            stream.get()>>>(tile, slots, columns_.data(), first_genome);
        CUDDL_CUDA_TRY(cudaGetLastError());
        return Ok();
    }

Here is the call graph for this function:

download()

Result< std::vector< batch_search_result > > cuddl::top_k_search::download ( cuda::stream_ref  stream ) const

inline

Each genome's neighbours, best first, genome by genome; waits for stream.

Definition at line 436 of file top_k.cuh.

                                                                                               {
        detail::device_scope const scope{stream};
        auto const slots = keys_.size();
        std::vector<uint64_t> keys(slots), counts(slots);
        if (slots != 0U) {
            // Rows are heaps on the device; sorting them there leaves the host a linear pass.
            auto sorted_keys = CUDDL_CUDA_TRY(
                cuda::make_device_buffer<uint64_t>(stream, stream.device(), slots, cuda::no_init)
            );
            auto sorted_counts = CUDDL_CUDA_TRY(
                cuda::make_device_buffer<uint64_t>(stream, stream.device(), slots, cuda::no_init)
            );
            auto const offsets = cuda::make_transform_iterator(
                cuda::counting_iterator<int64_t>{0}, detail::top_k_row_offset{k_}
            );
            CUDDL_CUDA_TRY(
                cub::DeviceSegmentedSort::SortPairsDescending(
                    keys_.data(),
                    sorted_keys.data(),
                    counts_.data(),
                    sorted_counts.data(),
                    static_cast<int64_t>(slots),
                    static_cast<int64_t>(genome_count_),
                    offsets,
                    offsets + 1,
                    stream
                )
            );
            CUDDL_CUDA_TRY(cuda::copy_bytes(stream, sorted_keys, cuda::std::span{keys}));
            CUDDL_CUDA_TRY(cuda::copy_bytes(stream, sorted_counts, cuda::std::span{counts}));
        }
        CUDDL_CUDA_TRY(stream.sync());
        std::vector<batch_search_result> rows;
        rows.reserve(slots);
        for (size_t slot = 0U; slot < slots; ++slot) {
            if (keys[slot] == 0U) continue;
            rows.push_back({
                .query_id = first_genome_id_ + static_cast<uint32_t>(slot / k_),
                .reference_id = ~static_cast<uint32_t>(keys[slot]),
                .counts = packed_pairwise_counts{counts[slot]}.unpack(bucket_count_),
            });
        }
        return rows;
    }

Here is the call graph for this function:

k()

uint32_t cuddl::top_k_search::k ( ) const

inlinenoexcept

Definition at line 481 of file top_k.cuh.

                                              {
        return k_;
    }

Here is the caller graph for this function:

make()

static Result< top_k_search > cuddl::top_k_search::make ( uint32_t  genome_count, uint32_t  reference_count, uint32_t  bucket_count, uint32_t  k, cuda::stream_ref  stream, uint32_t  first_genome_id = 0U  )

inlinestatic

Parameters

genome_count	Queries of a batch search, or rows of an all-to-all search.
reference_count	Database rows, the length of every result row.
bucket_count	Buckets per sketch, from which downloads recover both-empty counts.
first_genome_id	ID of the first query of a batch search; zero for all-to-all.

Blocks use as many warps as their heaps fit the shared memory stream's device gives a block. Fails when one heap does not fit.

Definition at line 328 of file top_k.cuh.

      {
        if (k == 0U) return Err(Error::invalid_argument("top-k search needs k of at least 1"));
        detail::device_scope const scope{stream};
        int shared_limit = 0;
        CUDDL_CUDA_TRY(cudaDeviceGetAttribute(
            &shared_limit, cudaDevAttrMaxSharedMemoryPerBlockOptin, stream.device().get()
        ));
        auto const heap_bytes = detail::top_k_heap_words(k) * sizeof(uint64_t);
        auto const warps = static_cast<uint32_t>(
            std::min<size_t>(detail::top_k_warps, static_cast<size_t>(shared_limit) / heap_bytes)
        );
        if (warps == 0U) {
            return Err(Error::invalid_argument("top-k heap exceeds this GPU's shared memory"));
        }
        auto const bytes = warps * heap_bytes;
        CUDDL_TRY(allow_shared(detail::top_k_rows_kernel<false, false>, bytes));
        CUDDL_TRY(allow_shared(detail::top_k_rows_kernel<true, false>, bytes));
        CUDDL_TRY(allow_shared(detail::top_k_rows_kernel<false, true>, bytes));
        CUDDL_TRY(allow_shared(detail::top_k_rows_kernel<true, true>, bytes));
        CUDDL_TRY(allow_shared(detail::top_k_columns_kernel, bytes));
        auto const slots = static_cast<size_t>(genome_count) * k;
        return CUDDL_CUDA_TRY(top_k_search(
            genome_count,
            reference_count,
            bucket_count,
            k,
            first_genome_id,
            warps,
            cuda::make_device_buffer<uint64_t>(stream, stream.device(), slots, uint64_t{0U}),
            cuda::make_device_buffer<uint64_t>(stream, stream.device(), slots, cuda::no_init),
            cuda::make_device_buffer<uint32_t>(stream, stream.device(), 0U, cuda::no_init)
        ));
    }

Here is the call graph for this function:

---

The documentation for this class was generated from the following file:

• top_k.cuh