NUMA

#include <vector> // `std::vector`
#include <span> // `std::span`
#include <forkunion.hpp> // `linux_numa_allocator`, `machine_topology_t`, `distributed_pool_t`
#include <numkong/spatial.h> // `nk_angular_f32`, `nk_f64_t`

namespace fu = ashvardanian::forkunion;
using floats_alloc_t = fu::linux_numa_allocator<float>;

constexpr std::size_t dimensions = 768; /// Matches most BERT-like models
static std::vector<float, floats_alloc_t> first_half(floats_alloc_t(fu::memory_domain_id_t {0}));
static std::vector<float, floats_alloc_t> second_half(floats_alloc_t(fu::memory_domain_id_t {1}));
static fu::machine_topology_t machine_topology;
static fu::distributed_pool_t distributed_pool;

/// Dynamically shards incoming vectors across 2 nodes in a round-robin fashion.
void append(std::span<float, dimensions> vector) {
    bool put_in_second = first_half.size() > second_half.size();
    if (put_in_second) second_half.insert(second_half.end(), vector.begin(), vector.end());
    else first_half.insert(first_half.end(), vector.begin(), vector.end());
}

/// On each compute domain we'll synchronize the threads
struct search_result_t {
    nk_f64_t best_distance {std::numeric_limits<nk_f64_t>::max()};
    std::size_t best_index {0};
};

inline search_result_t pick_best(search_result_t const& a, search_result_t const& b) noexcept {
    return a.best_distance < b.best_distance ? a : b;
}

/// Uses all CPU threads to search for the closest vector to the @p query.
search_result_t search(std::span<float, dimensions> query) {

    bool const need_to_spawn_threads = distributed_pool.threads_count() == 0;
    if (need_to_spawn_threads) {
        fu::status_t const harvested = machine_topology.harvest();
        assert(fu::succeeded(harvested) && "Failed to harvest NUMA topology");
        assert(machine_topology.memory_domains_count() == 2 && "Expected exactly 2 NUMA nodes");
        fu::status_t const spawned = distributed_pool.spawn(machine_topology);
        assert(fu::succeeded(spawned) && "Failed to spawn NUMA pools");
    }

    search_result_t result;
    fu::spin_mutex_t result_update; // ? Lighter `std::mutex` alternative w/out futexes

    std::size_t const first_vectors = first_half.size() / dimensions;
    std::size_t const total_vectors = first_vectors + second_half.size() / dimensions;

    auto slices = distributed_pool.for_slices(total_vectors,
        [&](fu::tasks_range_t tasks, fu::thread_in_domain_t) noexcept {

        search_result_t thread_local_result;
        for (std::size_t const task : tasks) {
            bool const in_second = task >= first_vectors;
            auto const &shard = in_second ? second_half : first_half;
            std::size_t const local_index = in_second ? task - first_vectors : task;

            nk_f64_t distance;
            nk_angular_f32(query.data(), shard.data() + local_index * dimensions, dimensions, &distance);
            thread_local_result = pick_best(thread_local_result, {distance, task});
        }

        // ! Still synchronizing over a shared mutex for brevity.
        std::lock_guard<fu::spin_mutex_t> lock(result_update);
        result = pick_best(result, thread_local_result);
    });
    slices.join();
    return result;
}