6. Routes and Fabric DSDs

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

Learning Objectives

After completing this tutorial, you should know how to:

Use fabric DSDs fabout_dsd and fabin_dsd to send and receive data between PEs
Utilize asynchronous builtin operations on fabric DSDs
Define a local task which is activated by a local_task_id

Example Overview

Your program will run on two processing elements (PE). This tutorial demonstrates the program with a simulated fabric consisting of a 9 x 3 block of PEs. The program will first copy b into the left PE’s y array. Then, it will copy the left half of A’s columns into the left PE, and the right half of A’s columns into the right PE. Similarly, it will copy the first N/2 elements of x into the left PE, and the last N/2 elements of x into the right PE. Each PE will then compute A*x for its local pieces of A and x. Thus, both PEs perform a matrix-vector product for an M x N/2 matrix. The PEs will increment their local y arrays by this result. The left PE then sends its y array to the right PE, and the right PE increments its local y array by the received values. Because the left y array contained the contribution from b, the final summed y on the right PE is the GEMV result. The host then copies y off of the right PE.

Problem Steps

Visually, this program consists of the following steps: 1. Host copies b into y array of left PE.

Diagram of the host copying b into the y array of the left PE

2. Host copies left N/2 columns of A to left PE, right N/2 columns to right PE.

Diagram of the host copying the left N/2 columns of A to the left PE and the right N/2 columns to the right PE

3. Host copies first N/2 elements of x to left PE, last N/2 elements to right PE.

Diagram of the host copying the first N/2 elements of x to the left PE and the last N/2 elements to the right PE

4. Host launches function to compute GEMV.

Diagram of the host launching a function to compute the GEMV

5. Each PE increments local y by local portion of matrix-vector product Ax.

Diagram of each PE incrementing local y by its local portion of the matrix-vector product Ax

6. Left PE sends local y to right PE, and right PE increments y by received values.

Diagram of the left PE sending local y to the right PE, which increments y by the received values

7. Right PE now contains final result y. Host copies back y from right PE.

Diagram of the right PE holding the final result y, which the host then copies back

Write the CSL

What needs to change in the layout to distribute the GEMV between two PEs?

Define several new parameters for the two PE programs. This includes a pe_id, used to differentiate between the left and right PEs, and a color, which is used to route data between the PEs.
Set the color configuration on both PEs for the color used to send the left PE’s y array to the right PE.

Take a look at the new layout.csl, included below.

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

// Colors
const send_color: color = @get_color(0); // Color used to send/recv data between PEs

// This example only uses 2 PEs
const memcpy = @import_module("<memcpy/get_params>", .{
  .width = 2,
  .height = 1,
});

layout {
  // PE coordinates are (column, row)
  @set_rectangle(2, 1);

  // Left PE (0, 0)
  @set_tile_code(0, 0, "pe_program.csl", .{
    .memcpy_params = memcpy.get_params(0),
    .M = M,
    .N_per_PE = N / 2,
    .pe_id = 0,
    .send_color = send_color
  });

  // Left PE sends its result to the right
  @set_color_config(0, 0, send_color, .{.routes = .{ .rx = .{RAMP}, .tx = .{EAST} }});

  // Right PE (1, 0)
  @set_tile_code(1, 0, "pe_program.csl", .{
    .memcpy_params = memcpy.get_params(1),
    .M = M,
    .N_per_PE = N / 2,
    .pe_id = 1,
    .send_color = send_color
  });

  // Right PE receives result of left PE
  @set_color_config(1, 0, send_color, .{.routes = .{ .rx = .{WEST}, .tx = .{RAMP} }});

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

There are two @set_tile_code calls, one for the left PE (0, 0), and one for the right PE (1, 0). Both PEs take a new parameter, N_per_PE, equal to N / 2. This is the number of columns of A that each PE will receive and operate on. Both PEs also receive as a parameter a pe_id: the left PE has pe_id 0, and the right PE has pe_id 1. pe_program.csl, covered next, shows how pe_id parameterizes the behavior of the program. There are also two @set_color_config calls, to set the configuration of send_color on each PE:

@set_color_config(0, 0, send_color, .{.routes = .{ .rx = .{RAMP}, .tx = .{EAST} }});
...
@set_color_config(1, 0, send_color, .{.routes = .{ .rx = .{WEST}, .tx = .{RAMP} }});

The router of each PE has five directions: RAMP, NORTH, SOUTH, EAST, WEST. The cardinal directions refer to the routers of neighboring PEs: NORTH is the PE directly above the current PE, and so on. RAMP refers to the connection between a PE’s router and its compute element (CE). When setting a route for a color on a given PE, the receive rx and transmit tx fields are from the perspective of the router. Thus, receiving from the RAMP means that the compute element is sending data up to the fabric, where it can then be transmitted across the fabric. For the left PE (0, 0), send_color will send up the PE’s RAMP to the fabric, and then transmit data to the EAST. For the right PE (1, 0), send_color will receive data from the WEST on the fabric (i.e., from the left PE), and then transmit it down the RAMP to its compute element. Now take a look at the new pe_program.csl, included below.

param memcpy_params;

// Matrix dimensions
param M: i16;
param N_per_PE: i16;

// ID of PE (0 is left, 1 is right)
param pe_id: i16;

// Colors
param send_color: color; // Color used to send/recv data between PEs

// Queue IDs
const send_color_oq = @get_output_queue(2);
const send_color_iq = @get_input_queue(2);

// Task ID used by a local task to unblock cmd stream
const exit_task_id: local_task_id = @get_local_task_id(9);

// memcpy module provides infrastructure for copying data
// and launching functions from the host
const sys_mod = @import_module("<memcpy/memcpy>", memcpy_params);

// 48 kB of global memory contain A, x, y
var A: [M*N_per_PE]f32; // A is stored column major
var x: [N_per_PE]f32;
var y: [M]f32;

// DSDs for accessing A, b, y
// A_dsd accesses column of A
var A_dsd = @get_dsd(mem1d_dsd, .{ .base_address = &A, .extent = M });
var y_dsd = @get_dsd(mem1d_dsd, .{ .base_address = &y, .extent = M });

// ptrs to A, x, b, y will be advertised as symbols to host
var A_ptr: [*]f32 = &A;
var x_ptr: [*]f32 = &x;
var y_ptr: [*]f32 = &y;

// Compute gemv
fn gemv() void {
  // Loop over all columns of A
  for (@range(i16, N_per_PE)) |i| {
    // Calculate contribution to A*x from ith column of A, ith elem of x
    @fmacs(y_dsd, y_dsd, A_dsd, x[i]);
    // Move A_dsd to next column of A
    A_dsd = @increment_dsd_offset(A_dsd, M, f32);
  }
}

fn send_right() void {
  const out_dsd = @get_dsd(fabout_dsd, if (@is_arch("wse3")) .{
                    .extent = M, .output_queue = send_color_oq
                  } else .{
                    .fabric_color = send_color, .extent = M,
                    .output_queue = send_color_oq
                  });
  // After fmovs is done, activate exit_task to unblock cmd_stream
  @fmovs(out_dsd, y_dsd, .{ .async = true, .activate = exit_task_id });
}

fn recv_left() void {
  const in_dsd = @get_dsd(fabin_dsd, .{
                   .extent = M,
                   .input_queue = send_color_iq
                 });
  // After fadds is done, activate exit_task to unblock cmd stream
  @fadds(y_dsd, y_dsd, in_dsd, .{ .async = true, .activate = exit_task_id });
}

// Call gemv function and send/ receive partial result y
fn compute() void {
  gemv();
  if (pe_id == 0) {
    send_right();
  } else {
    recv_left();
  }
}

task exit_task() void {
  sys_mod.unblock_cmd_stream();
}

comptime {
  // When exit_task_id is activated, exit_task will execute
  @bind_local_task(exit_task, exit_task_id);

  // On WSE-3, the output queue must be explicitly bound to a color.
  // Input queues must be bound to a color on both architectures.
  @initialize_queue(send_color_oq, if (@is_arch("wse3")) .{ .color = send_color } else .{});
  @initialize_queue(send_color_iq, .{ .color = send_color });

  @export_symbol(A_ptr, "A");
  @export_symbol(x_ptr, "x");
  @export_symbol(y_ptr, "y");
  @export_symbol(compute);
}

In addition to the new parameters N_per_PE, pe_id, and send_color, this code also introduces exit_task_id, the first value of type local_task_id, covered a bit later. The A array now has size M*N_per_PE instead of M*N, since each PE only stores half the columns. To make the data transfer easier, A is now stored column-major instead of row-major. Notice that A_dsd now accesses M contiguous elements, instead of M elements strided by the row size, since it’s now stored column-major. The gemv function operates almost identically to before, except it only loops over N_per_PE columns instead of N columns. Since A is now column-major, @increment_dsd_offset must increment by the length of an entire column instead of by one element. Note that on the left PE, y already contains the elements of b before gemv executes.

Fabric DSDs and Async Operations

The compute function, which is called from the host, first calls gemv to compute the local contribution to y on each PE. Then, the left PE calls send_right, while the right PE calls recv_left. send_right defines a fabout_dsd, which is used to send wavelets to the fabric along the color send_color. Note that this fabout_dsd is given the extent M, since the goal is to send the M elements of y along the fabric. On WSE-3, a fabric DSD is bound to a color indirectly via its queue, so .fabric_color is omitted from the WSE-3 form; on WSE-2 the color is specified directly on the DSD. The @fmovs operation copies the M elements accessed by y_dsd into out_dsd. The .async = true field makes this operation asynchronous. The .activate field specifies a local_task_id to activate when this operation completes. When this operation completes, exit_task_id will be activated. recv_left defines a fabin_dsd to receive the wavelets sent along send_color. The @fadds operation here increments the right PE’s y_dsd by the elements received in in_dsd. Thus, after this operation, y_dsd contains the final GEMV result. This builtin also executes asynchronously, and activates exit_task_id when complete.

Whenever using fabric DSDs in builtin operations, always make these operations execute asynchronously. Using fabric DSDs synchronously can result in poor performance or deadlocks.

Tasks and Activatable Task IDs

Now, what does activating exit_task_id do? In the comptime block, the @bind_local_task builtin binds exit_task_id to the task exit_task. When exit_task_id is activated, exit_task, which unblocks the memcpy command stream, executes. This task must execute on both PEs before control is returned to the host.

Write the Host Code

The host code must:

Copy b into the left PE’s y array
Copy the left halves of A and x to the left PE, and the right halves to the right PE
After the device kernel completes, copy y back from the right PE

The following sections explain some features of the new run.py, shown below.

#!/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'])

# Construct A, x, b
A = np.arange(M*N, dtype=np.float32).reshape(M,N)
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@x + b

# Size of N dimension on each PE
N_per_PE = N // 2

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

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

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

# Copy b into y of PE (0, 0)
runner.memcpy_h2d(y_symbol, b, 0, 0, 1, 1, M, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

# Copy A in column major format
# PE (0, 0) gets first N/2 columns; PE (1, 0) gets last N/2 columns
runner.memcpy_h2d(A_symbol, A.transpose().ravel(), 0, 0, 2, 1, M*N_per_PE, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

# PE (0, 0) gets first N/2 elements; PE (1, 0) gets last N/2 elements
runner.memcpy_h2d(x_symbol, x, 0, 0, 2, 1, N_per_PE, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

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

# Copy y back from PE (1, 0)
y_result = np.zeros([M], dtype=np.float32)
runner.memcpy_d2h(y_result, y_symbol, 1, 0, 1, 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, y_expected, atol=0.01, rtol=0)
print("SUCCESS!")

Copy `b` into `y` of the Left PE

b is copied into y of the left PE here:

runner.memcpy_h2d(y_symbol, b, 0, 0, 1, 1, M, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

Notice that the ROI is a single PE, located at (0, 0) in the program rectangle. The right PE (1, 0) is omitted from this memcpy call.

Copy `A` and `x`

A and x are copied to the device as follows:

runner.memcpy_h2d(A_symbol, A.transpose().ravel(), 0, 0, 2, 1, M*N_per_PE, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

runner.memcpy_h2d(x_symbol, x, 0, 0, 2, 1, N_per_PE, streaming=False,
  order=MemcpyOrder.ROW_MAJOR, data_type=MemcpyDataType.MEMCPY_32BIT, nonblock=False)

Notice that the ROI is now both PEs, so the memcpy calls copy data into both the left and right PE. Because A is now stored column-major on the PEs, the A matrix is transposed, then flattened to a 1D array with ravel(). Each PE gets M*N_per_PE elements, so each PE gets N_per_PE columns of A. Similarly, each PE gets N_per_PE elements of x.

Copy Back the Result

y is copied back from the right PE as follows:

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

Notice that the ROI now begins at (1, 0), and contains a single PE. Thus, this memcpy call copies back the M elements of y only from the right PE. Once this call is complete, as in the previous tutorials, the received result is checked for correctness.

Compile and Run the Program

Since this program only uses two PEs, the simulated fabric dimensions are adjusted accordingly:

$ cslc layout.csl --fabric-dims=9,3 --fabric-offsets=4,1 --params=M:4,N:6 --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

Instead of using two PEs along the same row to compute this GEMV, try using two PEs along the same column.

Stay tuned for more tutorials!

Start Here

Tutorials

Development Guides

Debugging

CSL Language Reference

Host API Reference

Learning Objectives

Example Overview

Problem Steps

Write the CSL

Fabric DSDs and Async Operations

Tasks and Activatable Task IDs

Write the Host Code

Copy `b` into `y` of the Left PE

Copy `A` and `x`

Copy Back the Result

Compile and Run the Program

Exercises

Next

​Learning Objectives

​Example Overview

​Problem Steps

​Write the CSL

​Fabric DSDs and Async Operations

​Tasks and Activatable Task IDs

​Write the Host Code

​Copy b into y of the Left PE

​Copy A and x

​Copy Back the Result

​Compile and Run the Program

​Exercises

​Next

Learning Objectives

Example Overview

Problem Steps

Write the CSL

Fabric DSDs and Async Operations

Tasks and Activatable Task IDs

Write the Host Code

Copy `b` into `y` of the Left PE

Copy `A` and `x`

Copy Back the Result

Compile and Run the Program

Exercises

Next