# NUMA

> ForkUnion for Zig

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 <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());
}
```

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<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.

`const forkunion = @import("forkunion");`

`const AllocationResult = @import("forkunion").AllocationResult;`

`const DomainAllocator = @import("forkunion").DomainAllocator;`

`const ReplicatedArray = @import("forkunion").ReplicatedArray;`

`const ShardedArray = @import("forkunion").ShardedArray;`

## AllocationResult

Result of a memory-domain allocation; carries only the OS id, so it can outlive the topology.

```zig
pub const AllocationResult = struct
```

- `memory_domain_id` `MemoryDomainId`
- `ptr` `[*]u8`
- `allocated_bytes` `usize`
- `bytes_per_page` `usize`

## asSlice

Returns the allocated memory as a slice

```zig
pub fn asSlice(self: AllocationResult) []u8
```

## free

Frees the allocation

```zig
pub fn free(self: AllocationResult) void
```

## DomainAllocator

Allocator bound to a single memory domain, compatible with Zig's `std.mem.Allocator` interface.

```zig
pub const DomainAllocator = struct
```

This type is the allocation API: build one for a memory domain's OS id - from
[`Topology.memoryDomainIdAtIndex`](https://www.unum.cloud/docs/forkunion/zig/topology.md#memoryDomainIdAtIndex) - and it both hands out [`AllocationResult`](#AllocationResult)s through
[`allocateAtLeast`](#allocateAtLeast)/[`allocate`](#allocate) and backs a `std.mem.Allocator` through [`allocator`](#allocator).

- `memory_domain_id` `MemoryDomainId`

## init

Binds to the memory domain named by @p memory_domain_id, which must name a real domain - ask [`isValid`](https://www.unum.cloud/docs/forkunion/zig/topology.md#isValid) first, since an out-of-range lookup answers `.none`.

```zig
pub fn init(memory_domain_id: MemoryDomainId) Self
```

## allocateAtLeast

Allocates at least @p minimum_bytes on this domain at optimal page size, or null on failure.

```zig
pub fn allocateAtLeast(
    self: Self,
    minimum_bytes: usize,
) ?AllocationResult
```

## allocate

Allocates exactly @p bytes on this domain, or null on failure.

```zig
pub fn allocate(self: Self, bytes: usize) ?[*]u8
```

## allocator

```zig
pub fn allocator(self: *Self) Allocator
```

## ReplicatedArray

One full length-`len` copy of a sequence per domain, so every thread reads a node-local replica.

```zig
pub fn ReplicatedArray(comptime T: type) type
```

A thin owner of one symmetric mapping the allocator stripes across the nodes - replica `d` lives
at `base + d * stride_bytes` and holds `len` elements. Raw uninitialized storage: the caller
fills every replica and keeps them coherent, mirroring the C++ `replicated_array`. `T` must be
plain-old-data, since the container runs no constructors or destructors.

- `base_bytes` `?[*]u8`
- `stride_bytes` `usize`
- `domains` `usize`
- `total_bytes` `usize`
- `len` `usize`

## init

Allocates one uninitialized length-`n` replica per memory domain.

```zig
pub fn init(topology: Topology, n: usize) Error!Self
```

## *

Unmaps the whole striped mapping.

```zig
pub fn deinit(self: *Self) void
```

## is*

Whether the container holds no elements.

```zig
pub fn isEmpty(self: Self) bool
```

## base

The typed base pointer of the first replica's storage.

```zig
pub fn base(self: Self) [*]T
```

## memoryDomainsCount

The number of per-domain replicas.

```zig
pub fn memoryDomainsCount(self: Self) usize
```

## strideBytes

The page-aligned byte distance between consecutive replicas.

```zig
pub fn strideBytes(self: Self) usize
```

## onMemoryDomain

The whole replica living on `memory_domain`.

```zig
pub fn onMemoryDomain(self: Self, memory_domain: MemoryDomain) []T
```

## at

One element of the replica on `memory_domain`.

```zig
pub fn at(
    self: Self,
    memory_domain: MemoryDomain,
    local_index: usize,
) *T
```

## ShardedArray

A sequence partitioned across memory domains as contiguous segments, each element stored once.

```zig
pub fn ShardedArray(comptime T: type) type
```

Each domain owns a contiguous logical segment of `segment` elements, so element `i` lives at `{i / segment, i % segment}` - see [`locationOf`](#locationOf) - and a scan of a domain's shard is sequential in
memory. Backed by one symmetric mapping; the trailing shard may be short. `T` must be
plain-old-data, mirroring [`ReplicatedArray`](#ReplicatedArray).

- `base_bytes` `?[*]u8`
- `stride_bytes` `usize`
- `domains` `usize`
- `total_bytes` `usize`
- `len` `usize`
- `segment` `usize`

## init

Allocates uninitialized storage for `n` elements partitioned round-robin across domains.

```zig
pub fn init(topology: Topology, n: usize) Error!Self
```

## base

The typed base pointer of the first shard's storage.

```zig
pub fn base(self: Self) [*]T
```

## memoryDomainsCount

The number of shards, one per memory domain.

```zig
pub fn memoryDomainsCount(self: Self) usize
```

## strideBytes

The page-aligned byte distance between consecutive shards.

```zig
pub fn strideBytes(self: Self) usize
```

## lengthOnMemoryDomain

How many elements the shard on `memory_domain` holds - a trailing shard may be shorter.

```zig
pub fn lengthOnMemoryDomain(
    self: Self,
    memory_domain: MemoryDomain,
) usize
```

## locationOf

The memory domain and local index that store logical element `logical_index`.

```zig
pub fn locationOf(
    self: Self,
    logical_index: usize,
) struct { memory_domain: MemoryDomain, local_index: usize }
```

## logicalIndexOf

The logical index at `local_index` on `memory_domain` - the inverse of [`locationOf`](#locationOf).

```zig
pub fn logicalIndexOf(
    self: Self,
    memory_domain: MemoryDomain,
    local_index: usize,
) usize
```

## onMemoryDomain

The whole shard living on `memory_domain`.

```zig
pub fn onMemoryDomain(self: Self, memory_domain: MemoryDomain) []T
```

## at

The single home of a logical element.

```zig
pub fn at(
    self: Self,
    memory_domain: MemoryDomain,
    local_index: usize,
) *T
```

## allocate

Allocates exactly the requested bytes on a domain; id from [`Topology.memoryDomainIdAtIndex`](https://www.unum.cloud/docs/forkunion/zig/topology.md#memoryDomainIdAtIndex).

```zig
pub fn allocate(
    memory_domain_id: MemoryDomainId,
    bytes: usize,
) ?[*]u8
```

## allocateAtLeast

Allocates on a memory domain with optimal page size; id from [`Topology.memoryDomainIdAtIndex`](https://www.unum.cloud/docs/forkunion/zig/topology.md#memoryDomainIdAtIndex).

```zig
pub fn allocateAtLeast(
    memory_domain_id: MemoryDomainId,
    minimum_bytes: usize,
) ?AllocationResult
```