# threads.parallel

> A module of threads's API.

`parallel_for` — the reason this tin exists — plus the `spawn_n` / `join_all` layer underneath it.

```mojo
from threads import parallel_for, AtomicCounter

@fieldwise_init
struct Totals(Copyable, Movable):
    var sum: Int64

def task(i: Int, mut totals: Totals) -> None:
    _ = AtomicCounter.at(Int(Pointer(to=totals.sum))).fetch_add(Int64(i))

def main() raises:
    var totals = Totals(0)
    parallel_for[task](1000, totals)
    print(totals.sum)  # 499500
```

`parallel_for` starts `min(n_tasks, num_workers or num_cpus())` kernel threads.
Each one loops on `AtomicCounter.fetch_add(1)` and runs `task(i, state)` for
every index it draws, stopping when the counter passes `n_tasks`. Every thread
is joined before the call returns, so by construction nothing the caller owns
can be destroyed while a worker still reads it.

## Two spellings of the same call

The **typed** form above takes the shared state by `ref` and the task by
`mut`. Because the state is an argument, it is alive for the whole call — the
compiler extends its lifetime across the joins, with no later mention needed —
and an immutable binding or a temporary is rejected at the call site. The
erasure to `void *` that pthread demands happens inside the library, once.

The **opaque** form is the one underneath:

```mojo
def square_chunk(i: Int, ctx: OpaquePtr) -> None:
    var cells = i64_ptr(Int(ctx))
    cells[unsafe_offset=i] = cells[unsafe_offset=i] * cells[unsafe_offset=i]

parallel_for[square_chunk](n_tasks=64, ctx=opaque_ptr(Int(data)))
```

Reach for it when the context is a hand-laid-out block of cells rather than a
struct, or when the pointer comes from somewhere the compiler cannot track
anyway (an allocation, a foreign buffer). It carries no origin: the caller
guarantees the pointee outlives the call, and `ctx` being an untracked
`Pointer` means Mojo's last-use destruction can free a local out from under
the workers with no diagnostic — see `threads.thread`.

What neither form checks: that `T` is safe to mutate from several threads at
once. Every task receives `mut` access to the same value; anything that is
not an atomic, a mutex-guarded region, or a slot only one task writes is a
data race the compiler will not see.

## The send contract, restated

- **`work` is thin and non-raising.** It is called through a C function
  pointer, so it cannot capture, and pthread has no exception channel.
- **All shared state travels through `ctx`.** One pointer, for every task.
- **The caller guarantees `ctx` outlives the call.** `parallel_for` joins
  before returning, so "outlives the call" is all that is needed — but see
  `threads.thread` on Mojo's last-use destruction if you build the context out
  of a local value.
- **No allocator assumption beyond what the tests prove.** Workers *may*
  allocate and free (`List`, `String`, `alloc`) concurrently; the
  `allocator under threads` test in `tests/test_threads.mojo` is the evidence,
  and it is run on all four CI legs.

## Scheduling

The queue is a single shared counter, so it is dynamically load balanced — a
slow task does not stall the others — at a cost of one atomic RMW per task.
Make tasks chunky: prefer `n_tasks` in the tens-to-thousands over one task per
element. `n_tasks` may exceed the worker count freely; that is the normal case.

## Reporting failures out of a task

`work` cannot raise. The pattern that fits the model is an error cell in the
context: an `AtomicFlag` plus a fixed-size array of error codes, one slot per
worker or per task.

```mojo
# ctx layout, in 64-bit cells:
#   [0]        error flag (0 = clean)
#   [1..n+1]   per-task error code
def task(i: Int, ctx: OpaquePtr) -> None:
    var flag = AtomicFlag.at(ctx, 0)
    var cells = i64_ptr(Int(ctx))
    if something_went_wrong:
        cells[1 + i] = MY_ERROR_CODE   # plain write, this task owns the slot
        flag.set()                      # release: publishes the write above
        return

# after parallel_for returns:
if flag.is_set():                       # acquire
    raise Error("task ", first_bad, " failed with ", cells[1 + first_bad])
```

The flag's release/acquire pairing is what makes the plain per-slot writes
visible to the caller; the join at the end of `parallel_for` also gives you
that, so the flag is really there to let you skip the scan when nothing failed.

## Structs

### `ThreadGroup`

```mojo
struct ThreadGroup
```

A batch of threads spawned together and joined together.

Move-only for the same reason `ThreadHandle` is: the group owns the threads
and there must be exactly one join.

Implements: `Movable`, `Sized`

#### `__init__`

```mojo
def __init__(out self)
```

An empty group.

#### `__len__`

```mojo
def __len__(self) -> Int
```

The number of threads in the group.

#### `append`

```mojo
def append(mut self, var handle: ThreadHandle)
```

Take ownership of one more thread.

Args:

- `handle`: The handle to adopt.

#### `join_all`

```mojo
def join_all(mut self)
```

Join every thread in the group.

Joins are attempted in spawn order. If one fails, the remaining threads
are still joined before the error is re-raised — leaving live threads
behind while the caller unwinds and frees their context would be far
worse than a late error.

Raises: Error: The first `pthread_join` failure, after all joins have been
attempted.

#### `pin_all`

```mojo
def pin_all(self)
```

Pin thread `i` to CPU `i % num_cpus()`.

A no-op on macOS (see `ThreadHandle.pin_to_cpu`).

Raises: Error: On Linux if the affinity call fails.

## Functions

### `join_all`

```mojo
def join_all(mut group: ThreadGroup)
```

Free-function spelling of `ThreadGroup.join_all`.

Args:

- `group`: The group to join.

Raises: Error: The first join failure.

### `spawn_n`

```mojo
def spawn_n[start: def(Pointer[UInt8, MutUntrackedOrigin]) thin -> OpaquePtr](n: Int, ctx: Pointer[UInt8, MutUntrackedOrigin], stride: Int = Int(0)) -> ThreadGroup
```

Start `n` threads running `start`, thread `i` receiving `ctx + i * stride`.

With the default `stride` of 0 every thread gets the same pointer, which is
what you want for a shared context. A non-zero stride carves the context
into equal per-thread slices without a second allocation.

If a spawn fails partway through, the threads already started are joined
before the error propagates, so the caller never has to reason about a
half-started group.

Parameters:

- `start`: The thread body — thin and non-raising.

Args:

- `n`: How many threads to start. `n <= 0` yields an empty group.
- `ctx`: Base context pointer.
- `stride`: Bytes between consecutive threads' context pointers.

Raises: Error: If any `pthread_create` fails.

### `parallel_for`

```mojo
def parallel_for[work: def(Int, Pointer[UInt8, MutUntrackedOrigin]) thin -> None](n_tasks: Int, ctx: Pointer[UInt8, MutUntrackedOrigin], num_workers: Int = Int(0))
```

Run `work(i, ctx)` for every `i` in `[0, n_tasks)`, across threads.

Starts `min(n_tasks, num_workers or num_cpus())` threads, hands out indices
from one shared atomic counter, and joins every thread before returning.
The calling thread does not run tasks; it blocks in the joins.

Parameters:

- `work`: The task body. Thin (non-capturing) and non-raising — see the
    module docstring for reporting failures out of a task.

Args:

- `n_tasks`: How many indices to hand out. Zero or negative does nothing.
- `ctx`: The shared context pointer given to every task. Must outlive the
    call; since the call joins before returning, that is automatic for
    anything the caller keeps alive across it.
- `num_workers`: Thread count. `0` (the default) means `num_cpus()`. Values
    above `n_tasks` are clamped down, because a thread with no index to
    draw is pure overhead.

Raises: Error: If a thread fails to start or to join. The worker context block
is freed on every path *after* the joins, never before.

```mojo
def parallel_for[T: AnyType, origin: MutOrigin, //, task: def(Int, mut T) thin -> None](n_tasks: Int, ref[origin] state: T, num_workers: Int = Int(0))
```

Run `task(i, state)` for every `i` in `[0, n_tasks)`, across threads.

The typed form of `parallel_for`. `state` is an argument, so it is alive
for the whole call, joins included, without the caller having to mention
it afterwards; and it is taken by mutable `ref`, so a `read` argument or a
temporary is refused at the call site. The pointer pthread carries is
built here and rebuilt in `_typed_worker`; nothing outside the library
handles it.

Parameters:

- `T`: The shared state type. Inferred from `state`.
- `origin`: The origin of `state`. Inferred.
- `task`: The task body, `def(Int, mut T) thin -> None`. Thin and
    non-raising, as in the opaque form.

Args:

- `n_tasks`: How many indices to hand out. Zero or negative does nothing.
- `state`: The value every task receives by `mut`. Guarding concurrent
    writes to it is the task's job — atomics, a `Mutex`, or per-task
    slots.
- `num_workers`: Thread count, as in the opaque form.

Raises: Error: If a thread fails to start or to join.

## Aliases

### `WorkFn`

```mojo
comptime WorkFn = def(Int, Pointer[UInt8, MutUntrackedOrigin]) thin -> None
```

A `parallel_for` task body, opaque form: the task index and the shared context pointer.

### `TaskFn`

```mojo
comptime TaskFn: __mlir_type.`!kgen.non_struct_type` = def(Int, mut T) thin -> None
```

A `parallel_for` task body, typed form: the task index and mutable access to the shared state. Every task gets the same `T`, at the same time.