Parallel workers

Each candidate economy has its own parameters and returns its own moments. Give independent economies to separate workers: a worker can start its next job while the others are still solving theirs.

What is removed?

The wait for one independent economy to finish before another can begin.

What remains?

Every economy must still be solved in full. Setup, the final join and uneven job lengths can limit the gain.

# Each call solves one independent economy.
results = pool.map(solve_economy, candidates)
moments = collect_in_candidate_order(results)

Try the worker illustration above ↑

Vectorization

Put firms in rows and markets in columns. One array expression applies the same operation across the matrix; a sum along the firm axis still gives a separate total for each market.

Given

One economy’s firm values, arranged as M rows and J market columns, with inactive firms masked out.

Change

Follow the film from separate market calls to one matrix operation. Trace one column through the firm sum and normalization in the code below.

Check

The sum runs down firms, not across markets. Every nonempty market’s shares sum to one; each market keeps its own convergence state.

Record

Label the shapes of the input, market totals and output. Explain which repeated calls disappear and which arithmetic remains.

# weights has shape (firms, markets).
market_totals = weights.sum(axis=0, keepdims=True)
shares = weights / market_totals

Vectorization changes how the work is written. This array code still runs on one CPU core; workers and GPUs are separate steps, in lessons 01 and 03.

Time it: python lessons/04_parallel/loop_vs_vectorized.py solves the same 2,000 markets in one call and in one call per market, with 5 to 80 firms per market, in about 30 seconds.

GPU batching

Stack compatible economies into a B × M × J array. A GPU operation can then work across many economies and markets, with a separate active mask for the parts that still need to converge.

Given

Fixed candidate rows and saved draws. Economies with the same M = floor(N) can share a B × M × J batch.

Change

Follow the film as separate economy matrices enter one batch. Give each launch more work, and keep track of which parameters and result belong to which economy.

Check

A finished market stops updating while others continue. Every economy still goes through both periods, every exit round and every root solve. The batch must fit in GPU memory.

Record

Draw the three axes and mark the active mask. Propose two batch sizes to compare, with the same rows and checks, and a clock that covers the whole solve.

# B economies × M firms × J markets.
state = stack_compatible_economies(economies)
state = update_active_markets(state, active_mask)

A larger batch helps while it fits in memory and keeps the GPU busy. Compare batch sizes by timing the whole solve.

Kernel fusion

Optional lesson: its code uses GPU kernels written in Triton. The core lessons are 01, 02, 03, 06 and 08.

A chain of small GPU operations can write a temporary array to device memory, then read it back in the next operation. A fused kernel combines compatible operations so intermediate values can stay inside the kernel.

Given

The same input array and a chain of compatible operations. The unfused version writes an intermediate array to device memory.

Change

Trace the separate operations in the film, then the fused kernel. Cross out the temporary-array write/read and launch that fusion can remove.

Check

The required arithmetic and reductions still happen. Retain the intended precision, damping, masks and residual checks; compare outputs before timing.

Record

A before/after diagram of launches and memory traffic. Name the saving for the share update alone.

# Separate operations materialize a temporary array.
temporary = scale * x
y = temporary + offset

# A compatible fused kernel computes the same expression.
y = fused_scale_and_add(x, scale, offset)

Fusion changes how the work runs and how much it reads and writes in memory.

Resident iteration

Optional lesson: its code uses GPU kernels written in Triton. The core lessons are 01, 02, 03, 06 and 08.

When the host checks every small iteration, it may repeatedly wait for the GPU and launch the next step. A resident inner loop keeps state and compatible stopping checks on the device through that sequence.

Given

The same share state, FP32 predictor and FP64 correction budget as the mixed backend.

Change

Follow one correction step in each film lane. Move the correction loop and its stopping state onto the GPU, reducing host checks and relaunches.

Check

Each step still uses the previous state. Preserve each market’s first stopping step, masks, iteration cap, failure path and independent final FP64 residual.

Record

Mark where state lives and where the host waits in both versions. Name one removed handoff and one dependency that remains.

# Inside the resident kernel, with the original stopping rules:
while any(active) and iteration < iteration_cap:
    state = update_active_markets(state, active)
    active = original_remaining_work_mask(state)
    iteration += 1
return state, convergence_status

Only the inner share loop moves onto the GPU; the output root, the exit rounds and the rest of the economy run as before. Each update still needs the one before it.