from forkunion import AllocationResult, DomainAllocator, SymmetricAllocation, …4.0.0-rc.1 · d927445f
Placing memory on a chosen domain is implemented on Linux via the mbind syscall, on FreeBSD via domainset policies, and on Windows via VirtualAllocExNuma, and degrades to a single domain elsewhere.
The portable domain_allocator_t alias picks the right backend per platform - linux_numa_allocator_t, freebsd_numa_allocator_t, windows_numa_allocator_t, or a plain aligned fallback - each an STL-compatible allocator bound to one memory_domain_id_t.
Let’s say you are building a Big Data application — brute-forcing Vector Search with the NumKong library — on a 2 dual-socket CPU system, the kind that powers production USearch deployments.
The first part of that program may be responsible for sharding the incoming stream of data between distinct memory regions.
That part, in our simple example will be single-threaded:
The concurrent part would involve spawning threads adjacent to every memory pool to find the best search_result_t.
The primary search function, in ideal world would look like this:
Each thread finds the best match within its “slice” of a compute domain, tracking the best distance and index in a local CPU register.
All threads in each compute domain atomically synchronize using a domain-local instance of search_result_t.
The main thread collects aggregates of partial results from all compute domains.
That is, however, overly complicated to implement.
Such tree-like hierarchical reductions are optimal in a theoretical sense.
Still, assuming the relative cost of spin-locking once at the end of a thread scope and the complexity of organizing the code, the more straightforward path is better.
A minimal example would look like this:
/// On each compute domain we'll synchronize the threadsstruct 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;}
In a dream world, we would call distributed_pool.for_n, but there is no clean way to make the scheduling processes aware of the data distribution in an arbitrary application, so that’s left to the user.
The for_slices helper pairs each run with a fu::thread_in_domain_t naming its compute domain, but slices follow the thread count rather than the shard boundary, so the callback above picks the shard per task.
For more flexibility around building higher-level low-latency systems, there are unsafe APIs expecting you to manually “join” the broadcasted calls: unsafe_for_threads returns a generation token, is_complete polls it without blocking, and unsafe_join blocks until that generation completes.
The manual two-vector sharding above is what the symmetric allocators automate: symmetric_memory_allocator_t - fu_allocate_symmetric in C - maps one virtual range striped a slice per memory domain, and Rust, Zig, and Mojo wrap it as ShardedArray<T>, one shard per domain, and ReplicatedArray<T>, a full copy per domain for read-mostly data.
To place and coordinate pools by hand, spawn_on - fu_pool_spawn_on in C - pins a pool to a single compute domain, while locate_thread_in and threads_count_in map a global thread index to its domain and count the workers living there.
API
StructAllocationResultPages placed on one memory domain, released when the last reference goes out of use.4
The returned pointer carries this allocation's origin, so the compiler keeps the pages
alive for as long as the pointer is reachable. Without that, Mojo would release the
allocation after its last named use - which is this call - and the pointer would dangle
before the caller ever wrote through it.
Parameters
dtypeDType
comptime
The element type to read the untyped pages as; nothing is converted.
Returns
A pointer borrowing this allocation, so the pages cannot be released beneath it.
count_ofHow many elements of dtype the placed pages hold.
def count_of[dtype: DType](self) -> Int
Parameters
dtypeDType
comptime
The element type whose width divides the placed bytes.
Returns
The count the pages fit, above what was asked for whenever page rounding rounded up.
StructDomainAllocatorPlaces pages on one memory domain, named by the operating system's id for it.6
struct DomainAllocator
Holds the id and a loader reference and nothing else, so it is cheap to copy and free of the
topology that produced the id.
This is the pairing the index-versus-id distinction exists to protect: local_memory_of
answers a dense index, and only memory_domain_id_at_index turns it into an allocator's id.