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