# NUMA > ForkUnion for Mojo 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](/docs/numkong/) library — on a 2 dual-socket CPU system, the kind that powers production [USearch](https://github.com/unum-cloud/USearch/pulls) 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: ```cpp #include // `std::vector` #include // `std::span` #include // `linux_numa_allocator`, `machine_topology_t`, `distributed_pool_t` #include // `nk_angular_f32`, `nk_f64_t` namespace fu = ashvardanian::forkunion; using floats_alloc_t = fu::linux_numa_allocator; constexpr std::size_t dimensions = 768; /// Matches most BERT-like models static std::vector first_half(floats_alloc_t(fu::memory_domain_id_t {0})); static std::vector 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 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()); } ``` 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: 1. Each thread finds the best match within its "slice" of a compute domain, tracking the best distance and index in a local CPU register. 2. All threads in each compute domain atomically synchronize using a domain-local instance of `search_result_t`. 3. 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: ```cpp /// On each compute domain we'll synchronize the threads struct search_result_t { nk_f64_t best_distance {std::numeric_limits::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 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 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`, one shard per domain, and `ReplicatedArray`, 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. `from forkunion import AllocationResult, DomainAllocator, SymmetricAllocation, …` ## AllocationResult Pages placed on one memory domain, released when the last reference goes out of use. ```mojo struct AllocationResult ``` - `library` `Library`: The loader the release symbol comes from, kept alive for exactly as long as the pages. - `memory_domain_id` `MemoryDomainId`: The domain the pages sit on, replayed on release because the C free takes it too. - `pointer` `Bytes`: The base of the placed pages, aligned to at least a cache line. - `allocated_bytes` `Int`: What was actually placed, which the release has to repeat exactly. - `bytes_per_page` `Int`: The page size behind the pages, or zero where the entry point reported none. Conforms to: `Deinitable`, `Movable` ## __init__ Takes ownership of pages a placement call has already returned. ```mojo def __init__( out self, library: Library, memory_domain_id: MemoryDomainId, address: Int, allocated_bytes: Int, bytes_per_page: Int ) ``` - `library`: The loader to release through, which the result then keeps alive. - `memory_domain_id`: The domain the pages were placed on, needed again to free them. - `address`: The address the C side wrote out, not yet a typed pointer. - `allocated_bytes`: What the placement really returned, which may exceed the request. - `bytes_per_page`: The page size behind the pages, or zero where the call reported none. ## __deinit__ Hands the pages back to the domain they came from, at the byte count they were placed at. ```mojo def __deinit__(deinit self) ``` ## as_pointer The same storage, seen as the element type the caller means to put in it. ```mojo def as_pointer[dtype: DType](mut self) -> Pointer[Scalar[dtype], origin_of(self)] ``` 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. - `dtype` (comptime): The element type to read the untyped pages as; nothing is converted. ## count_of How many elements of `dtype` the placed pages hold. ```mojo def count_of[dtype: DType](self) -> Int ``` - `dtype` (comptime): The element type whose width divides the placed bytes. ## DomainAllocator Places pages on one memory domain, named by the operating system's id for it. ```mojo 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. - `library` `Library`: The loader the placement and release symbols come from. - `memory_domain_id` `MemoryDomainId`: The operating system's own id for the domain, never a dense index into the topology. Conforms to: `Copyable`, `Deinitable`, `Equatable`, `ImplicitlyCopyable`, `Movable` ## __eq__ Two allocators match when they name the same domain; the loader reference is not compared. ```mojo def __eq__(self, other: Self) -> Bool ``` - `other`: The allocator whose domain id is compared with this one's. ## at An allocator for a domain the machine actually has. ```mojo def at(library: Library, memory_domain_id: MemoryDomainId) -> Self ``` - `library`: The loader every allocation from here is placed and released through. - `memory_domain_id`: An id from [`Topology.memory_domain_id_at_index`](https://www.unum.cloud/docs/forkunion/mojo/topology.md#memory_domain_id_at_index), never a dense index. Raises: When the id is negative, which is how a domain the OS names none of comes back. ## allocate Exactly `bytes` on this domain. ```mojo def allocate(self, bytes: Int) -> AllocationResult ``` The page size is not reported by this entry point, so the result carries zero for it. - `bytes`: The exact size to place, with no page rounding added on top of it. Raises: When the request is zero bytes, or the domain cannot back the placement. ## allocate_at_least At least `minimum_bytes`, using the largest suitable page size, reporting both. ```mojo def allocate_at_least(self, minimum_bytes: Int) -> AllocationResult ``` - `minimum_bytes`: The floor on the request; huge-page rounding may hand back more. Raises: When the request is zero bytes, or the domain cannot back the placement. ## allocate_for Room for `count` elements of `dtype`. ```mojo def allocate_for[dtype: DType](self, count: Int) -> AllocationResult ``` - `dtype` (comptime): The element type whose width scales the request. - `count`: How many elements to make room for, not a byte count. Raises: When the count times the element width would wrap, or the domain refuses the placement. ## allocate_for_at_least Room for at least `minimum_count` elements of `dtype`, on the largest suitable page. ```mojo def allocate_for_at_least[dtype: DType](self, minimum_count: Int) -> AllocationResult ``` - `dtype` (comptime): The element type whose width scales the request. - `minimum_count`: The floor on elements; read [`count_of`](#count_of) back for what the pages really fit. Raises: When the count times the element width would wrap, or the domain refuses the placement. ## SymmetricAllocation One mapping striped across every memory domain at a uniform stride. ```mojo struct SymmetricAllocation ``` Backs both the replicated and the sharded views: they differ only in what the per-domain run means, not in how the storage is laid out. - `library` `Library`: The loader the unmap goes through, kept alive until it does. - `base` `Bytes`: The start of the whole mapping; the run for domain `d` begins `d` strides in. - `stride_bytes` `Int`: The page-aligned distance between runs, which rounding may lift above what was asked. - `memory_domains_count` `Int`: How many runs the mapping got, which a machine with no NUMA API collapses to one. - `total_bytes` `Int`: The whole mapping's size, which the unmap needs rather than one run's. - `bytes_per_page` `Int`: The page size the mapping used, which is what drove the stride rounding. Conforms to: `Deinitable`, `Movable` ## __init__ Maps one run per memory domain and binds each run to the domain it belongs to. ```mojo def __init__(out self, topology: Topology, bytes_per_domain: Int) ``` - `topology`: Passed straight to the C mapping, which reads the domain count off it. - `bytes_per_domain`: The usable floor per run; the stride is page-rounded up from it. Raises: When the request is zero bytes, or the mapping cannot be placed. ## __deinit__ Unmaps the whole range in one call, which is why the total and not a stride is kept. ```mojo def __deinit__(deinit self) ``` ## on_memory_domain The run of storage living on one memory domain, seen as the caller's element type. ```mojo def on_memory_domain[dtype: DType]( mut self, domain: MemoryDomain ) -> Pointer[Scalar[dtype], origin_of(self)] ``` - `dtype` (comptime): The element type to read the run as; nothing is converted. - `domain`: A dense memory-domain index, counted from zero, not an operating-system id. ## ReplicatedArray The same `length` elements held once per memory domain, so every reader is local. ```mojo struct ReplicatedArray[dtype: DType] ``` Write each replica, then let every compute domain read the copy nearest to it. - `dtype` (comptime): The element type every replica holds. - `storage` `SymmetricAllocation`: The single mapping behind every replica, one run per domain. - `length` `Int`: Elements in one replica, which is the same for all of them. Conforms to: `Deinitable`, `Movable` ## __init__ Adopts a mapping whose runs are already sized for one replica each. ```mojo def __init__(out self, *, var storage: SymmetricAllocation, length: Int) ``` - `storage` (var): The mapping to take over, one run per domain. - `length`: Elements per replica, which the caller already sized the runs for. ## new A replica of `length` elements on every memory domain. ```mojo def new(topology: Topology, length: Int) -> Self ``` - `topology`: Says how many domains there are to replicate onto. - `length`: Elements in each copy, not the total held across the domains. Raises: When the length times the element width would wrap, or the mapping cannot be placed. ## memory_domains_count How many replicas the mapping produced. ```mojo def memory_domains_count(self) -> Int ``` ## replica One domain's copy, all `length` elements of it. ```mojo def replica(mut self, domain: MemoryDomain) -> Pointer[Scalar[dtype], origin_of(self)] ``` - `domain`: A dense memory-domain index, counted from zero, not an operating-system id. ## at One element of one domain's copy. ```mojo def at(mut self, domain: MemoryDomain, index: Int) -> ref[self] Scalar[dtype] ``` - `domain`: Which copy to reach into; each holds the same `length` elements. - `index`: Position within that copy, below `length` and unchecked. ## ShardLocation Where a logical index lives: which memory domain, and where inside that domain's shard. ```mojo struct ShardLocation ``` - `memory_domain` `MemoryDomain`: Which shard holds the element, as a dense index. - `local_index` `Int`: Position inside that shard, which is not the logical index it came from. Conforms to: `Copyable`, `Deinitable`, `Equatable`, `ImplicitlyCopyable`, `Movable`, `RegisterPassable`, `TrivialRegisterPassable` ## ShardedArray `length` elements split across the memory domains, each domain owning one contiguous run. ```mojo struct ShardedArray[dtype: DType] ``` - `dtype` (comptime): The element type every shard holds. - `storage` `SymmetricAllocation`: The single mapping behind every shard, one run per domain. - `length` `Int`: Elements across all the shards together, not per shard. - `segment` `Int`: Elements one shard has room for, by ceiling division, so only the last one can be short. Conforms to: `Deinitable`, `Movable` ## __init__ Adopts a mapping whose runs are already sized for one shard each. ```mojo def __init__(out self, *, var storage: SymmetricAllocation, length: Int, segment: Int) ``` - `storage` (var): The mapping to take over, one run per domain. - `length`: Elements across all the shards together. - `segment`: Elements each shard has room for, which is what fixes where an index lands. ## new `length` elements striped across every memory domain. ```mojo def new(topology: Topology, length: Int) -> Self ``` - `topology`: Says how many domains the elements are split over. - `length`: Elements in total; ceiling division settles the per-domain run. Raises: When the per-shard count times the element width would wrap, or the mapping fails. ## memory_domains_count How many shards the mapping produced. ```mojo def memory_domains_count(self) -> Int ``` ## shard One domain's run of elements. ```mojo def shard(mut self, domain: MemoryDomain) -> Pointer[Scalar[dtype], origin_of(self)] ``` - `domain`: A dense memory-domain index, counted from zero, not an operating-system id. ## length_on_memory_domain How many elements one domain's shard actually holds; the last one may be short. ```mojo def length_on_memory_domain(self, domain: MemoryDomain) -> Int ``` - `domain`: Which shard to measure, as a dense index. ## location_of Which domain holds a logical index, and where inside that domain's shard. ```mojo def location_of(self, logical_index: Int) -> ShardLocation ``` Requires a non-empty array - `segment` is the divisor and is zero when empty. - `logical_index`: A position in the whole array, not one within a single shard. ## logical_index_of The inverse: a domain and a local offset back to a logical index. ```mojo def logical_index_of(self, domain: MemoryDomain, local_index: Int) -> Int ``` - `domain`: The shard the offset is measured inside, as a dense index. - `local_index`: Position within that shard, below `segment`. ## default_domain_allocator An allocator for the machine's first memory domain, which every machine has. ```mojo def default_domain_allocator(topology: Topology) -> DomainAllocator ``` - `topology`: Consulted only to resolve index zero to an id, and not held afterwards. Raises: When the machine cannot be described well enough to name even its first domain. ## local_domain_allocator An allocator for the memory domain nearest a compute domain - "run here, allocate here". ```mojo def local_domain_allocator(topology: Topology, domain: ComputeDomain) -> DomainAllocator ``` This is the pairing the index-versus-id distinction exists to protect: [`local_memory_of`](https://www.unum.cloud/docs/forkunion/mojo/topology.md#local_memory_of) answers a dense index, and only [`memory_domain_id_at_index`](https://www.unum.cloud/docs/forkunion/mojo/topology.md#memory_domain_id_at_index) turns it into an allocator's id. - `topology`: Consulted for both halves of that resolution, and not held afterwards. - `domain`: A dense compute-domain index, counted from zero - where the work will run. Raises: When the machine has no such compute domain, or names no memory domain beside it.