A Conceptual View
Contents
A Conceptual View¶
This section presents a conceptual view of computing with the Cerebras architecture. Read this before you get into the details of how to write programs with the Cerebras SDK.
The Cerebras Wafer-Scale Engine (WSE) is a wafer-parallel compute accelerator, containing hundreds of thousands of independent processing elements (PEs). The PEs are interconnected by communication links into a two-dimensional rectangular mesh on one single silicon wafer. Each PE has its own memory (used by it and no other) and its own program counter. It has its own executable code in its memory. 32-bit messages, called wavelets, can be sent to or received by neighboring PEs in a single clock cycle.
The PE also has dataflow control characteristics. An instruction can terminate the currently running thread (called a task), at which time, hardware selects a new task from among the set of tasks that constitute the PE’s code. It selects a runnable task, one that has been activated (and unblocked; we will describe this in more detail later). Incoming wavelets travel along a virtual channel, called a color. All colors transfer data on a single physical channel. The congestion of one color does not block the traffic of another color. For each color used for incoming wavelets, there may be a task that is activated by its arrival.
The Cerebras System (CS) is a self-contained rack-mounted system containing packaging, power supply, cooling and I/O for a single WSE. The CS communicates via parallel 100 Gigabit ethernet connections to a host CPU cluster. Throughout this documentation, the CS is referred to as the “device,” the host CPU cluster as the “host,” and the ethernet connections connecting the two as “host I/O”. The SDK provides mechanisms for using host I/O to move data between host and device or launch functions on the device.
The below figure gives a visual representation of the mesh of PEs that make up the WSE, and its connections to the outside world. Data is streamed onto the device via host I/O, and enters the WSE through a series of links along its edges. The programming model of the SDK abstracts away the details of these links, and allows the programmer to copy data from the host to arbitrary PEs on the device.
A processing element (PE)¶
A PE contains three key elements:
A processor. Also referred as a compute engine (CE).
A router. The router of a PE is directly connected via bidirectional links to its own CE and to the routers of the four nearest neighboring PEs in the mesh. The link to its own CE is called the RAMP, and the links to the four neighboring PEs are referred to by their cardinal directions. The router is the only communication device the PEs use to send and receive data.
The local PE memory. All of the PE’s data and code are stored within this memory. Neither the CE nor the local memory of a PE is directly accessible by other PEs.
The programming model¶
To develop code for the WSE, you write device code in CSL, and host code in Python. You then compile the device code, and run your program on either the Cerebras fabric simulator, or the actual network-attached device. The host code is responsible for copying data to and from the device, and launching discrete programs referred to as kernels.
CSL gives you full control of the WSE. To understand how to structure device code in CSL, we need to introduce a few key concepts first.
Programs and Tasks¶
A CSL program consists of one or more subprograms. Some of these are callable functions, and some are tasks. A task is a procedure that cannot be called from other code. Rather, tasks are started by the PE hardware, and run until they complete, at which point, the hardware chooses a new task to run. Tasks cannot be called, only activated. They cannot return values.
For example, in this code snippet, main_task will set the value
of the global variable result to 5.0 when it is activated and runs.
In CSL this task is represented as:
var result: f32 = 0.0; task main_task() void { result = 5.0; }
Each task is bound to a task ID, which serves as a handle for identifying the task.
Task IDs and Types of Tasks¶
The term “task identifier” or “task ID” is used to refer to a numerical value that can be associated with a task. A task ID is a number from 0 to 63. Within this range there are two properties that further distinguish a task ID: routable and activatable.
There exist three types of tasks, each with an associated task ID handle type:
The first are data tasks. These are associated with a
data_task_id, which on thewse2architecture is created from a routable identifier associated with a color.The second are local tasks. These are associated with a
local_task_id, which is created from an activatable identifier.The third are control tasks. These are associated with a
control_task_id, which can be created from any identifier, including those that are neither routable nor activatable.
Both data tasks and control tasks are a type of wavelet-triggered task, or WTT: their activation is triggerd by the arrival of a wavelet. In this introduction, we will explore what it means for a task ID to be routable or activatable, and the usage of data tasks and local tasks.
Communication¶
The WSE provides efficient, fine-grained communication between PEs. A PE must be able to quickly, within a few cycles, respond to the arrival of a wavelet, update its internal state, and send out wavelets.
On wse2, the IDs 0 to 23 are recognized by the hardware as virtual communication channels
for passing wavelets between PEs.
This range of IDs are the routable IDs, and when they are used in a routable manner,
we refer to them as colors.
Each wavelet has a 5-bit tag that encodes its color. This color determines both the wavelet’s routing through the fabric. and what task, if any, will consume the wavelet when received.
In the below code block, our task now takes an argument, named wavelet_data.
This task is an example of a data task.
The builtin CSL function @bind_data_task creates a binding between
a task ID associated with the color
of the incoming wavelet and the task main_task.
When a red wavelet arrives, the task main_task is activated,
which allows it to be selected by the task picker.
The red wavelet sits in a buffer until the task picker
selects the associated task main_task,
at which point the wavelet is moved into a register so that the task
can get instant access to that data.
Syntactically, the wavelet’s data is an argument of the task.
We also call main_task a wavelet-triggered task,
since it is activated by the arrival of a wavelet.
Note that the @bind_data_task operation occurs within a
comptime block: everything in a comptime block
is evaluated at compile-time.
// 7 is the ID of a color with some defined data routing
const red: color = @get_color(7);
// The ID for this task is created from a color
const red_task_id: data_task_id = @get_data_task_id(red);
var result: f32 = 0.0;
task main_task(wavelet_data: f32) {
result = wavelet_data;
}
comptime {
@bind_data_task(main_task, red_task_id);
}
Note
Colors can also be used in a manner such that wavelets arriving along the associated virtual communication channel do not activate a task, but are received by a construct known as a fabric DSD. We will explore this usage in the coming tutorials.
Task Activation and Control Flow¶
As we saw above, a task will become available for selection by the task picker when its associated task ID is activated.
The task in the above block was a data task bound to a data_task_id,
associated with a color which defines a route taken by wavelets
tagged with that color.
Additionally, we can create tasks which do not take wavelets as arguments,
and instead are explicitly activated by other tasks or functions.
We call these tasks local tasks, and the associated task ID type
is local_task_ID.
On wse2, we can create a local_task_ID from the range of task IDs
0 to 30. These IDs are the activatable IDs.
In the below example, main_task activates the task ID foo_task_id.
This task ID is bound to foo_task, and so activating foo_task_id
will cause foo_task to execute next.
const red: color = @get_color(7);
const red_task_id: data_task_id = @get_data_task_id(red);
const foo_task_id: local_task_id = @get_local_task_id(8);
var result: f32 = 0.0;
var sum: f32 = 0.0;
task main_task(wavelet_data: f32) {
result = wavelet_data;
@activate(foo_task_id);
}
task foo_task() {
sum += result;
}
comptime {
@bind_data_task(main_task, red_task_id);
@bind_local_task(foo_task, foo_task_id);
}
Blocking and Unblocking¶
We can additionally block a task ID to provide further control over task execution. A task must be unblocked and activated for it be scheduled by the task picker. If a task is activated while blocked, it will only run once it has become unblocked. By default, all IDs are unblocked and inactive.
In the below example, we introduce one additional task, bar_task.
We block the ID of this task at compile time, with @block(bar_task_id).
When main_task executes, it activates both foo_task_id and bar_task_id.
However, because bar_task_id is blocked, we guarantee that foo_task
executes first. When foo_task executes, in unblocks bar_task_id,
allowing bar_task to begin execution once foo_task finishes.
If bar_task_id were not blocked at compile time, then the
execution of foo_task and bar_task could occur in any order.
const red: color = @get_color(7);
const red_task_id: data_task_id = @get_data_task_id(red);
const foo_task_id: local_task_id = @get_local_task_id(8);
const bar_task_id: local_task_id = @get_local_task_id(9);
var result: f32 = 0.0;
var sum: f32 = 0.0;
task main_task(wavelet_data: f32) {
result = wavelet_data;
@activate(foo_task_id);
@activate(bar_task_id);
}
task foo_task() {
sum += result;
@unblock(bar_task_id);
}
task bar_task() {
sum *= 2.0;
}
comptime {
@block(bar_task_id);
@bind_data_task(main_task, red_task_id);
@bind_local_task(foo_task, foo_task_id);
@bind_local_task(bar_task, bar_task_id);
}
Layout¶
Layout blocks are how you connect up the multiple PEs in your program rectangle in a way your computation requires. For example, see the following 2-PE rectangle:
The diagram shows only one of the many ways you can interconnect the two PEs. You connect a PE to another PE by specifying routes for colors.
Using CSL you can define the specific colors and routes by which your rectangle is stitched up.
This configuration of colors and routes forms an essential aspect of your computation,
transforming the wavelets as they enter and pass through your rectangle.
See below an example of layout block showing a layout of two PEs in a single row:
// layout.csl
// Top-level program source
const main_color: color = @get_color(0);
layout {
@set_rectangle(2, 1); // A row containing two PEs.
@set_tile_code(0, 0, "send.csl", .{ .send_color = main_color });
@set_tile_code(1, 0, "recv.csl", .{ .recv_color = main_color });
const send_route = .{ .rx = .{ RAMP }, .tx = .{ EAST } };
const recv_route = .{ .rx = .{ WEST }, .tx = .{ RAMP } };
@set_color_config(0, 0, main_color, send_route);
@set_color_config(1, 0, main_color, recv_route);
}
The @set_tile_code() builtin CSL function is used to specify the .csl file
containing the program for the individual PE.
For example, the program send.csl contains the task description
that only the PE at the coordinate (0, 0) will perform.
Note
Exactly one
@set_rectangle()call is allowed. Zero or multiple calls are illegal.The built-in
@set_tile_code()must be called after@set_rectangle().
For example, if your rectangle contains five PEs, then you can configure each PE
with a different program by having five @set_tile_code() calls after a single
@set_rectangle(),
with each @set_tile_code() function call specified with a separate .csl file.
The @set_color_config calls assign the routing associated with main_color for each PE,
from the perspective of the PE’s router.
For instance, the router of PE (0, 0) will receive wavelets along color main_color
from the RAMP, which connects the router to the CE.
It will then transmit wavelets to the EAST, where it will be received by the router
of the neighboring PE.