Writing parallel code

Vx has two kinds of parallel code. Tasks do different work at the same time, and Tasks and channels explains them. An across loop does the same work on each element of a tensor, on the blocks and lanes of a GPU, and Parallel code on a GPU explains it. This chapter explains how to choose between the tools, and how to check that a parallel program is correct. It is for programmers who come from Rust with tokio, from Go, from C with threads, or from CUDA.

Start with the program in order

Write the program without tasks first, and keep it. The version in order is the reference: run both versions on the same input, and compare the results. A parallel program that is fast and wrong is easy to write. A comparison finds the problem in the first run.

For an across loop, the version in order is the same program without the word across.

Choose the shape of the work

Most parallel programs have one of four shapes:

• The same work on each element. Each pass of a loop works on its own elements of a tensor. Write it as an across loop, as Loops across blocks and lanes shows. This shape needs no tasks.
• Split the data. Each task takes one part of the input and gives one result. The tasks borrow the input from the function around them, and join() gives each result. The first program of Tasks and channels has this shape.
• Workers take jobs. When jobs take different times, an even split leaves some tasks idle. Give each worker the next job number from an atomic counter instead.
• A pipeline. Each stage is a task, and a channel connects each stage to the next one. A full channel makes a fast stage wait for a slow one.

In this program, four workers take job numbers from a counter. Each worker sends a result with its job number on a channel (book_workers_take_jobs_from_a_counter.vx):

import core::convert;
import core::sync::atomic;
import std::sync::channel;
import std::vec;

struct Done {
  job : i32,
  value : i64,
}

impl Copy for Done {}

fn worker(jobs : i32, next : &AtomicI64, tx : Sender<Done>) -> void {
  loop {
    // Each fetch_add gives a number that no other worker gets.
    let job = next.fetch_add(1, AtomicOrdering::Relaxed);
    if job >= i64::from(jobs) {
      return;
    }
    let value = job * job;
    match tx.send(Done {
      job : job as i32, value : value
    }) {
      Result<void, SendError<Done>>::Ok() => {}
      Result<void, SendError<Done>>::Err(_e) => {
        return;
      }
    }
  }
}

fn main() -> i32 {
  let jobs = 10;
  let mut results = Vec<i64>::new();
  let mut runs = Vec<i32>::new();
  for _i in 0..jobs {
    results.push(0);
    runs.push(0);
  }
  let next = AtomicI64::new(0);
  let (tx, rx) = channel<Done>(4);
  tasks {
    for _w in 0..4 {
      let mine = tx.clone();
      spawn {
        worker(jobs, &next, mine);
      };
    }
    drop(tx);
    let mut open = true;
    while open {
      match rx.recv() {
        Result<Done, RecvError>::Ok(d) => {
          // The job number says where the result goes, whatever order the results arrive in.
          results.set(d.job, d.value);
          runs.set(d.job, runs.get(d.job) + 1);
        }
        Result<Done, RecvError>::Err(_e) => {
          open = false;
        }
      }
    }
  };
  print!("squares");
  let mut once = true;
  for i in 0..jobs {
    print!(" {}", results.get(i));
    once = once && runs.get(i) == 1;
  }
  println!("");
  println!("each job ran once: {}", once);
  return 0;
}

It prints:

squares 0 1 4 9 16 25 36 49 64 81
each job ran once: true

The results arrive in any order. The job number puts each result in its place, so the output is the same in each run.

Choose the tool

Use a channel when you are not sure. A channel moves each value to one owner, so two tasks never change the same data. A task can borrow the locals of the function around it, so most programs need no Arc.

Loops across blocks and lanes

A word after the range of a loop, across blocks or across lanes(n), makes its passes run at the same time on a GPU. A block runs on one part of the chip, and a lane is one thread of a block. This program doubles each element of a tensor, and adds the number of its row, y (book_gpu_blocks_and_lanes.vx):

fn main() -> i32 {
  let mut grid = transfer(Tensor<f32, [64, 300]>::fill(1.0), Memory::HBM);
  spawn on(Topology::Device) {
    for y in 0..64 across blocks {
      for x in 0..300 across lanes(256) {
        grid[y][x] = grid[y][x] * 2.0 + (y as f32);
      }
    }
  }
  let back = transfer(grid, Memory::CPU_DRAM);
  println!("corner {}, last row {}", back[0][0], back[63][299]);
  return 0;
}

It prints corner 2, last row 65. Give the compiler a machine file with a GPU:

vxc blocks.vx --machine "$HOME/.vx/current/fleet/h100-sxm.vx" --host default

On a computer with no GPU, each kernel runs on the threads of the CPU, and the result is the same.

An across loop is easier to get right than tasks:

• The compiler checks that two passes cannot write the same element. When they can, it gives the error E6032. So an across loop has no data race for a test to find.
• The meaning of the program does not depend on across. Remove each across, and the loops run one pass at a time with the same result.
• The lanes of a block can share a tile in Memory::SMEM. The end of each lane loop waits for all the lanes of the block, so the next lane loop can read what any lane wrote.

Parallel code on a GPU explains the rules of across, tiles, the lanes:: routines, atomics on elements, and copies in the background.

Keep the result independent of the order

The order in which tasks run changes from run to run. Do not let the result depend on it:

• Give each result a number, and put it in its place, as the program above does.
• When jobs can conflict, run each job on a copy of the data from before the jobs started. Then go through the jobs in order, and run again each job that read data an earlier job changed. Parallel blockchain executors work this way. The starter project parallel-batch shows it on a batch of transfers.
• Do not add floating-point numbers in the order that they arrive. The sum can change in the last bits. Store the numbers by job, and add them in job order.
• In an across loop, lanes::sum adds in the same order each time. lanes::sum_unordered lets the compiler choose the order, so the last bits of its sum can differ from lanes::sum.

Choose a memory order

• A counter that only gives out numbers needs Relaxed.
• To give data to another task through a flag, store the flag with Release, and load it with Acquire. Memory orders has a program that does this.
• Use SeqCst only when two or more flags must be seen in one order by every task.

By default, Vx compiles a Relaxed load as an Acquire load. With -fweak-relaxed-atomics, the load stays Relaxed. So a program that needs Acquire but asks for Relaxed can work on Vx and fail elsewhere. Ask for the order that the program needs.

Waits that do not end

• A recv() loop ends only after the program drops each Sender. Drop the Sender that made the clones, as the program above does with drop(tx).
• join() and lock() cannot be cancelled. Do not hold a Mutex while a task waits for a channel.
• When each task waits and nothing can wake one, the runtime stops the program with vx: deadlock: .... It does not hang.

Test it

Run each test in these ways:

• With the default runtime, which gives each task an OS thread.
• With VX_TASK_RUNTIME=shared, which runs the tasks on a few OS threads.
• With VX_TASK_SEED=1, 2 and more. The tasks then run on one thread, and the seed chooses the order. The same seed gives the same order, so a failure happens again in the next run.
• With vxc --sanitize=thread, which finds data races through raw pointers that the compiler cannot see.

For an across loop, also run the program with CUDA_VISIBLE_DEVICES=. Then it runs the CPU copy of each kernel, and the result must be the same as on the GPU.

Coming from tokio

Vx has no async. Why Vx has no async explains this.

Coming from CUDA

The compiler generates PTX for an NVIDIA GPU. An Apple GPU runs block and lane loops, shared memory, and atomics on i32 and u32.

Calling parallel Vx from C or Rust

A pub fn can start tasks, and C or Rust can call it as any other pub fn. The program that calls it starts no threads. Link the object with the Vx runtime library, libvx_std_core, which runs the tasks:

vxc work.vx --action emit-obj -o work.o
cc host.c work.o -L"$HOME/.vx/current/lib" -lvx_std_core -Wl,-rpath,"$HOME/.vx/current/lib" -lpthread -lm -o host

Calling Vx from C says which types a pub fn can take.