Merge pull request #6 from eseiler/misc/refine

Use two queues

Author: eseiler

include/fpgalign/search/search.hpp

@@ -4,57 +4,27 @@
 
 #pragma once
 
-#include <mutex>
-#include <string>
-#include <vector>
-
 #include <fpgalign/config.hpp>
+#include <fpgalign/contrib/slotted_cart_queue.hpp>
 #include <fpgalign/meta.hpp>
 
 namespace search
 {
 
-struct hit
-{
-    std::vector<uint64_t> bins;
-};
-
-struct wip_alignment
+struct alignment_info
 {
     size_t bin;
     size_t sequence_number;
     size_t position;
     size_t idx;
 };
 
-class alignment_vector
-{
-private:
-    std::vector<wip_alignment> data;
-    std::mutex mtx;
-
-public:
-    std::vector<wip_alignment> & get() noexcept
-    {
-        return data;
-    }
-
-    std::vector<wip_alignment> const & get() const noexcept
-    {
-        return data;
-    }
-
-    void emplace_back(wip_alignment elem)
-    {
-        std::lock_guard guard{mtx};
-        data.emplace_back(std::move(elem));
-    }
-
-    void emplace_back(wip_alignment & elem) = delete;
-};
-
 void search(config const & config);
-std::vector<hit> ibf(config const & config, meta & meta);
-std::vector<wip_alignment> fmindex(config const & config, meta & meta, std::vector<hit> hits);
+void ibf(config const & config, meta & meta, scq::slotted_cart_queue<size_t> & filter_queue);
+void fmindex(config const & config,
+             meta & meta,
+             scq::slotted_cart_queue<size_t> & filter_queue,
+             scq::slotted_cart_queue<alignment_info> & alignment_queue);
+void do_alignment(config const & config, meta & meta, scq::slotted_cart_queue<alignment_info> & alignment_queue);
 
 } // namespace search

src/CMakeLists.txt

@@ -13,6 +13,7 @@ set (FPGAlign_SOURCE_FILES
         search/ibf.cpp
         search/fmindex.cpp
         search/search.cpp
+        search/do_alignment.cpp
 )
 
 # An object library (without main) to be used in multiple targets.

src/build/fmindex.cpp

@@ -54,7 +54,7 @@ void fmindex(config const & config, meta & meta)
         {
             read_reference_into(reference, meta, i);
 
-            fmc::BiFMIndex<5> index{reference, /*samplingRate*/ 16, config.threads};
+            fmc::BiFMIndex<5> index{reference, /*samplingRate*/ 16, /*threads*/ 1u};
 
             {
                 std::ofstream os{fmt::format("{}.{}.fmindex", config.output_path.c_str(), i), std::ios::binary};

src/search/do_alignment.cpp

@@ -0,0 +1,82 @@
+// SPDX-FileCopyrightText: 2006-2025 Knut Reinert & Freie Universität Berlin
+// SPDX-FileCopyrightText: 2016-2025 Knut Reinert & MPI für molekulare Genetik
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <seqan3/alignment/cigar_conversion/cigar_from_alignment.hpp>
+#include <seqan3/alignment/pairwise/align_pairwise.hpp>
+#include <seqan3/alphabet/nucleotide/dna4.hpp>
+#include <seqan3/io/sam_file/output.hpp>
+
+#include <fpgalign/search/search.hpp>
+
+namespace search
+{
+
+using sam_out_t = seqan3::sam_file_output<seqan3::fields<seqan3::field::seq,
+                                                         seqan3::field::id,
+                                                         seqan3::field::ref_id,
+                                                         seqan3::field::ref_offset,
+                                                         seqan3::field::cigar,
+                                                         //    seqan3::field::qual,
+                                                         seqan3::field::mapq>>;
+
+void task(meta & meta, std::span<alignment_info> alignment_infos, sam_out_t & sam_out)
+{
+    static seqan3::configuration const align_config =
+        seqan3::align_cfg::method_global{seqan3::align_cfg::free_end_gaps_sequence1_leading{true},
+                                         seqan3::align_cfg::free_end_gaps_sequence2_leading{false},
+                                         seqan3::align_cfg::free_end_gaps_sequence1_trailing{true},
+                                         seqan3::align_cfg::free_end_gaps_sequence2_trailing{false}}
+        | seqan3::align_cfg::edit_scheme | seqan3::align_cfg::output_alignment{}
+        | seqan3::align_cfg::output_begin_position{} | seqan3::align_cfg::output_score{};
+
+    for (auto & [bin, sequence_number, position, idx] : alignment_infos)
+    {
+        auto & seq = meta.queries[idx].sequence();
+        auto seq_view = std::views::transform(seq,
+                                              [](seqan3::dna4 const in) -> uint8_t
+                                              {
+                                                  return in.to_rank() + 1u;
+                                              });
+        auto & seq_id = meta.queries[idx].id();
+        auto & ref = meta.references[bin][sequence_number];
+        auto & ref_id = meta.ref_ids[bin][sequence_number];
+
+        size_t const start = position - static_cast<size_t>(position != 0u);
+        size_t const length = seq.size();
+        auto it = std::ranges::next(ref.begin(), start, ref.end());
+        auto end = std::ranges::next(it, length + 1u, ref.end());
+        std::span ref_text{it, end};
+
+        for (auto && alignment : seqan3::align_pairwise(std::tie(ref_text, seq_view), align_config))
+        {
+            auto cigar = seqan3::cigar_from_alignment(alignment.alignment());
+            size_t ref_offset = alignment.sequence1_begin_position() + 2 + start;
+            size_t map_qual = 60u + alignment.score();
+
+            sam_out.emplace_back(seq,
+                                 seq_id,
+                                 ref_id,
+                                 ref_offset,
+                                 cigar,
+                                 //  record.base_qualities(),
+                                 map_qual);
+        }
+    }
+}
+
+void do_alignment(config const & config, meta & meta, scq::slotted_cart_queue<alignment_info> & alignment_queue)
+{
+    sam_out_t sam_out{config.output_path};
+
+    while (true)
+    {
+        scq::cart_future<alignment_info> cart = alignment_queue.dequeue();
+        if (!cart.valid())
+            return;
+
+        task(meta, cart.get().second, sam_out);
+    }
+}
+
+} // namespace search

src/search/fmindex.cpp

@@ -30,67 +30,48 @@ fmc::BiFMIndex<5> load_index(config const & config, size_t const id)
     return index;
 }
 
-std::vector<wip_alignment> fmindex(config const & config, meta & meta, std::vector<hit> hits)
+void fmindex(config const & config,
+             meta & meta,
+             scq::slotted_cart_queue<size_t> & filter_queue,
+             scq::slotted_cart_queue<alignment_info> & alignment_queue)
 {
-    // todo capacity
-    // each slot = 1 bin
-    // a cart is full if it has 5 elements (hits)
-    alignment_vector res;
+#pragma omp parallel num_threads(config.threads)
     {
-        scq::slotted_cart_queue<size_t> queue{
-            {.slots = meta.number_of_bins, .carts = meta.number_of_bins, .capacity = 5}};
-        size_t thread_id{};
-
-        auto get_thread = [&]()
+        while (true)
         {
-            return std::jthread(
-                [&, thread_id = thread_id++]()
+            scq::cart_future<size_t> cart = filter_queue.dequeue();
+            if (!cart.valid())
+                break;
+            auto [slot, span] = cart.get();
+            auto index = load_index(config, slot.value);
+            for (auto idx : span)
+            {
+                auto callback = [&](auto cursor, size_t)
                 {
-                    while (true)
+                    for (auto j : cursor)
                     {
-                        scq::cart_future<size_t> cart = queue.dequeue();
-                        if (!cart.valid())
-                            return;
-                        auto [slot, span] = cart.get();
-                        auto index = load_index(config, slot.value);
-                        for (auto idx : span)
-                        {
-                            auto callback = [&](auto cursor, size_t)
-                            {
-                                for (auto j : cursor)
-                                {
-                                    auto [entry, offset] = index.locate(j);
-                                    auto [seqId, pos] = entry;
-                                    res.emplace_back(wip_alignment{.bin = slot.value,
-                                                                   .sequence_number = seqId,
-                                                                   .position = pos + offset,
-                                                                   .idx = idx});
-                                }
-                            };
-
-                            auto seq_view = std::views::transform(meta.queries[idx].sequence(),
-                                                                  [](seqan3::dna4 const in) -> uint8_t
-                                                                  {
-                                                                      return in.to_rank() + 1u;
-                                                                  });
-
-                            fmc::search<true>(index, seq_view, config.errors, callback);
-                        }
+                        auto [entry, offset] = index.locate(j);
+                        auto [seqId, pos] = entry;
+                        alignment_queue.enqueue(scq::slot_id{0u},
+                                                alignment_info{.bin = slot.value,
+                                                               .sequence_number = seqId,
+                                                               .position = pos + offset,
+                                                               .idx = idx});
                     }
-                });
-        };
+                };
 
-        std::vector<std::jthread> worker(config.threads);
-        std::ranges::generate(worker, get_thread);
+                auto seq_view = std::views::transform(meta.queries[idx].sequence(),
+                                                      [](seqan3::dna4 const in) -> uint8_t
+                                                      {
+                                                          return in.to_rank() + 1u;
+                                                      });
 
-        for (auto && [idx, hit] : seqan::stl::views::enumerate(hits))
-            for (auto bin : hit.bins)
-                queue.enqueue(scq::slot_id{bin}, idx);
-
-        queue.close();
-    } // Wait for threads to finish
+                fmc::search<true>(index, seq_view, config.errors, callback);
+            }
+        }
+    }
 
-    return res.get();
+    alignment_queue.close();
 }
 
 } // namespace search

src/search/ibf.cpp

@@ -39,7 +39,7 @@ threshold::threshold get_thresholder(config const & config, meta const & meta)
 using seqfile_t = seqan3::sequence_file_input<dna4_traits, seqan3::fields<seqan3::field::id, seqan3::field::seq>>;
 using record_t = typename seqfile_t::record_type;
 
-std::vector<hit> ibf(config const & config, meta & meta)
+void ibf(config const & config, meta & meta, scq::slotted_cart_queue<size_t> & filter_queue)
 {
     seqan::hibf::interleaved_bloom_filter ibf{};
 
@@ -60,8 +60,6 @@ std::vector<hit> ibf(config const & config, meta & meta)
         return result;
     }();
 
-    std::vector<hit> hits(meta.queries.size());
-
 #pragma omp parallel num_threads(config.threads)
     {
         auto agent = ibf.membership_agent();
@@ -80,11 +78,14 @@ std::vector<hit> ibf(config const & config, meta & meta)
             hashes.assign(view.begin(), view.end());
 
             auto & result = agent.membership_for(hashes, thresholder.get(hashes.size()));
-            std::ranges::copy(result, std::back_inserter(hits[i].bins));
+            for (size_t bin : result)
+            {
+                filter_queue.enqueue(scq::slot_id{bin}, i);
+            }
         }
     }
 
-    return hits;
+    filter_queue.close();
 }
 
 } // namespace search