5. Multiple PEs

The power of the Wafer-Scale Engine lies in its hundreds of thousands of processing elements. Now that the basics for writing a complete program using a single PE are covered, this tutorial creates your first program using multiple PEs.

Learning Objectives

After completing this tutorial, you should know how to:

Define a layout file that compiles code for multiple PEs
Copy data to and from multiple PEs on the device

Example Overview

Your program will now run on four processing elements (PE). This tutorial demonstrates the program with a simulated fabric consisting of an 11 x 3 block of PEs. For this program, each PE performs the exact same work; that is, A, x, and b are copied to each of the four PEs, the four PEs each perform a GEMV, and then the result y is copied back from each PE. pe_program.csl does not change. Only layout.csl needs to be modified to assign it to multiple PEs. The host code also needs to be modified to copy to and from multiple PEs instead of just one.

Problem Steps

Visually, this program consists of the following steps: 1. Host copies A, x, b to four PEs on device.

Diagram of the host copying A, x, and b to four PEs on the device

2. Host launches function on each PE to compute y.

Diagram of the host launching a function on each PE to compute y

3. Host copies result y from each PE.

Diagram of the host copying the result y from each PE

Modify the CSL

How does the layout file need to change to support running the program on multiple PEs?

Modify @set_rectangle to reflect the new program rectangle.
Modify the memcpy infrastructure to reflect the use of multiple PEs.
Call @set_tile_code for each coordinate inside this program rectangle.

pe_program.csl remains largely the same; it’s simply assigned to more PEs. The new layout.csl is included below, with the changes highlighted.

// matrix dimensions on each PE
param M: i16;
param N: i16;

// number of PEs in program
param width: i16;

const memcpy = @import_module("<memcpy/get_params>", .{
  .width = width,
  .height = 1
});

layout {
  // PE coordinates are (column, row)
  @set_rectangle(width, 1);
  for (@range(i16, width)) |x| {
    @set_tile_code(x, 0, "pe_program.csl", .{
      .memcpy_params = memcpy.get_params(x),
      .M = M,
      .N = N
    });
  }

  // export symbol names
  @export_name("A", [*]f32, true);
  @export_name("x", [*]f32, true);
  @export_name("b", [*]f32, true);
  @export_name("y", [*]f32, false);
  @export_name("compute", fn()void);
}

Notice that a new compile time parameter width is defined, whose value is set in the compile command. This value sets the number of PEs in the row of PEs used by the program. When <memcpy/get_params> is imported, width specifies the width of the program rectangle for which memcpy infrastructure will be generated. The height is still 1. Inside the layout block, the program rectangle is now specified with @set_rectangle(width, 1). For each of the PEs in this rectangle, @set_tile_code must be called, so this happens in a loop. The loop coordinate is the PE’s x-coordinate, or column number, which is needed to set the correct memcpy_params for each PE.

Modify the Host Code

The host code must now copy A, x and b to multiple PEs, and must copy back y from multiple PEs. Take a look at how the memcpy_h2d and memcpy_d2h calls in run.py need to change:

#!/usr/bin/env cs_python

import argparse
import json
import numpy as np

from cerebras.sdk.runtime.sdkruntimepybind import SdkRuntime, MemcpyDataType, MemcpyOrder # pylint: disable=no-name-in-module

# Read arguments
parser = argparse.ArgumentParser()
parser.add_argument('--name', help="the test compile output dir")
parser.add_argument('--cmaddr', help="IP:port for CS system")
args = parser.parse_args()

# Get matrix dimensions from compile metadata
with open(f"{args.name}/out.json", encoding='utf-8') as json_file:
  compile_data = json.load(json_file)

# Matrix dimensions
N = int(compile_data['params']['N'])
M = int(compile_data['params']['M'])

# Number of PEs in program
width = int(compile_data['params']['width'])

# Construct A, x, b
A = np.arange(M*N, dtype=np.float32)
x = np.full(shape=N, fill_value=1.0, dtype=np.float32)
b = np.full(shape=M, fill_value=2.0, dtype=np.float32)

# Calculate expected y
y_expected = A.reshape(M,N)@x + b

# Construct a runner using SdkRuntime
runner = SdkRuntime(args.name, cmaddr=args.cmaddr)

# Get symbols for A, x, b, y on device
A_symbol = runner.get_id('A')
x_symbol = runner.get_id('x')
b_symbol = runner.get_id('b')
y_symbol = runner.get_id('y')

# Load and run the program
runner.load()
runner.run()

# Copy A, x, b to device
runner.memcpy_h2d(A_symbol, np.tile(A, width), 0, 0, width, 1, M*N, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)
runner.memcpy_h2d(x_symbol, np.tile(x, width), 0, 0, width, 1, N, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)
runner.memcpy_h2d(b_symbol, np.tile(b, width), 0, 0, width, 1, M, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

# Launch the init_and_compute function on device
runner.launch('compute', nonblock=False)

# Copy y back from device
y_result = np.zeros([M*width], dtype=np.float32)
runner.memcpy_d2h(y_result, y_symbol, 0, 0, width, 1, M, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

# Stop the program
runner.stop()

# Ensure that the result matches expectations
np.testing.assert_allclose(y_result, np.tile(y_expected, width), atol=0.01, rtol=0)
print("SUCCESS!")

First, note that one more parameter is read from the compile output, width. The host code uses this to specify how many PEs it must copy tensors to and from. Now take a closer look at the memcpy_h2d calls:

runner.memcpy_h2d(A_symbol, np.tile(A, width), 0, 0, width, 1, M*N, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)
runner.memcpy_h2d(x_symbol, np.tile(x, width), 0, 0, width, 1, N, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)
runner.memcpy_h2d(b_symbol, np.tile(b, width), 0, 0, width, 1, M, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

Each of A, x, and b needs to be copied to each PE in the program rectangle. But memcpy_h2d does not perform a broadcast; it takes its input array and distributes it within the region of interest (ROI) based on the order parameter. Here, np.tile duplicates each array width times. In the first memcpy_h2d, the input array np.tile(A, width) is a 1D array formed by duplicating A width times, so the full input array’s size is M*N*width. The ROI is specified by 0, 0, width, 1, meaning the copy goes to a row of width PEs beginning at PE (0, 0). M*N elements are copied to each PE. Because the order is ROW_MAJOR, the result is that PE (0, 0) receives the first M*N elements of the tiled array, PE (1, 0) receives the next M*N elements, and so on. Thus, each PE receives an identical M*N elements corresponding to a copy of A. When y is copied back from the device, memcpy_d2h proceeds similarly:

y_result = np.zeros([M*width], dtype=np.float32)
runner.memcpy_d2h(y_result, y_symbol, 0, 0, width, 1, M, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

The output array y_result has size M*width, since each of the width PEs copies back the M elements of y. The copied-back result is tested for correctness across all PEs by comparing y_result to a tiled y_expected. See A Complete Program for an explanation of the remaining arguments.

Compile and Run the Program

This compile command adds one additional compile time parameter to specify the width of the program rectangle:

$ cslc layout.csl --fabric-dims=11,3 --fabric-offsets=4,1 --params=M:4,N:6,width:4 --memcpy --channels=1 -o out
$ cs_python run.py --name out

The run command stays the same. You should see a SUCCESS! message at the end of execution.

Exercises

In this program, each PE is computing an identical GEMV. Modify the program so that each PE receives different values for the input tensors A, x, and b, and check that the computed outputs y are correct.

Now that multiple PEs have been introduced into the program, instead of duplicating the GEMV problem between them, the next tutorial distributes the work for computing a single GEMV.

Start Here

Tutorials

Development Guides

Debugging

CSL Language Reference

Host API Reference

Learning Objectives

Example Overview

Problem Steps

Modify the CSL

Modify the Host Code

Compile and Run the Program

Exercises

Next

​Learning Objectives

​Example Overview

​Problem Steps

​Modify the CSL

​Modify the Host Code

​Compile and Run the Program

​Exercises

​Next

Learning Objectives

Example Overview

Problem Steps

Modify the CSL

Modify the Host Code

Compile and Run the Program

Exercises

Next