Host Runtime and Tensor Streaming

Once your kernel is compiled, you still need to get data onto the device, launch it, and get results back. The Cerebras SDK handles this through two pieces working together: the CSL memcpy library, which moves data across the WSE’s I/O channels, and SdkRuntime, the Python host runtime that drives it — loading programs, launching functions, and transferring data on and off the device. This page covers how both work, starting with how memcpy routes data through the fabric. The functions provided by SdkRuntime manage the data transfer to and from the host’s filesystem or memory, through the host and WSE network interfaces, and finally route the data into your kernel. This last step is implemented on the WSE itself. It connects the I/O channel entry-points — which sit in fixed locations at the edges of the WSE — to your kernel, which can have a variable size and location. These I/O channels are connected to the fabric routers at PEs on the East and West edges, spaced roughly 16 rows apart. On WSE-2, there are a total of 60 channels at each edge. On WSE-3, there are a total of 62 channels at each edge.

Diagram of I/O channels connecting the WSE's East and West edges to the host

The SDK memcpy infrastructure uses additional PEs around your kernel to route tensor data and also adds a small executable component to the kernel PEs. In addition to a halo around the kernel, the additional support PEs consume three columns on the West of the kernel and two columns on the East.

Diagram of the support PEs the memcpy infrastructure adds around a kernel

SdkRuntime supports up to 16 I/O channels, and can further reduce the I/O latency by buffer insertion on either side of the core kernel.

Set Up the memcpy Infrastructure

The memcpy infrastructure of SdkRuntime moves data on and off the device in one of two modes:

streaming mode delivers data as a sequence of wavelets. Your kernel counts the wavelets as they arrive and acts once the full tensor has been received.
copy mode writes data directly into device memory without notifying your kernel — you copy the tensor in first, then launch a kernel function to act on it.

For example, given a tensor A and a function f that transforms it, computing f(A) looks different depending on the mode:

In streaming mode, you’d define a wavelet-triggered data task that receives A and calls f once all of A has arrived;
In copy mode, you’d copy A onto the device first, then launch a kernel that calls f.

To instantiate and use the memcpy infrastructure, you’ll need to do the following:

Pass compiler flags

Pass --memcpy and --channels=k to cslc (the CSL compiler), where k is an integer between 1 and 16 specifying the number of I/O channels to use.

Set fabric dimensions and offsets

Specify --fabric-dims=dim_x,dim_y and --fabric-offsets=x,y to cslc, where width and height are your program rectangle’s dimensions, such that:

dim_x >= 7 + width
dim_y >= 2 + height
x >= 4
y >= 1

Import memcpy parameters

In your top-level layout CSL file, instantiate memcpy parameters by importing <memcpy/get_params> with an @import_module() statement. This specifies everything the infrastructure needs, including width and height of the kernel and any colors needed for streaming mode.

Pass params to the PE program

Pass the memcpy params to the PE program in the @set_tile_code call. These params are parameterized by the PE’s x coordinate in the program rectangle.

Import the memcpy module

In your PE program, instantiate the memcpy module by importing <memcpy/memcpy>.

Altogether, instantiating memcpy infrastructure in the top-level CSL file and the PE program will resemble the following example:

// in top-level CSL file
const memcpy = @import_module("<memcpy/get_params>", .{
    .width = width,
    .height = height
});

layout {
  @set_rectangle(1, 1);
  @set_tile_code(0, 0, "pe_program.csl", .{ .memcpy_params = memcpy.get_params(0) });
}

// in PE program CSL file pe_program.csl
param memcpy_params: comptime_struct;

const sys_mod = @import_module("<memcpy/memcpy>", memcpy_params);

The memcpy infrastructure reserves the following resources. The compiler and runtime cannot detect all resource conflicts, so do not use these in your own program:

Resource	Reserved IDs
Colors	21, 22, 23
Local task IDs	27, 28, 30
Control task IDs	33, 34, 35, 36, 37
Microthread	0
Input queue	0 (WSE-2 and WSE-3); additionally 1 on WSE-3
Output queue	0

Do not set or modify the routing of an input, output tensor color or kernel launch color. The compiler configures the routing pattern implicitly. If you modify those routing patterns, the behavior is undefined.

Use Streaming Mode

To use streaming mode, you must specify colors for host-to-device and device-to-host streaming. Input streaming parameters are prefixed with MEMCPYH2D_DATA_ and output streaming parameters are prefixed with MEMCPYD2H_DATA_, followed by the tensor ID (an integer in the range 1-4) and an _ID suffix — for example, MEMCPYH2D_DATA_1_ID. Unused colors should be omitted, and only four colors per direction are allowed. You can block and unblock the input tensor colors to overlap computation and communication. Here’s an example instantiation of a program in the top-level CSL file using colors for memcpy streaming:

// in top-level CSL file
// Compile-time IDs for memcpy streaming colors
param MEMCPYH2D_DATA_1_ID: i16;
param MEMCPYH2D_DATA_2_ID: i16;
param MEMCPYD2H_DATA_1_ID: i16;
param MEMCPYD2H_DATA_2_ID: i16;

// Generate colors from IDs
const MEMCPYH2D_DATA_1: color = @get_color(MEMCPYH2D_DATA_1_ID);
const MEMCPYH2D_DATA_2: color = @get_color(MEMCPYH2D_DATA_2_ID);
const MEMCPYD2H_DATA_1: color = @get_color(MEMCPYD2H_DATA_1_ID);
const MEMCPYD2H_DATA_2: color = @get_color(MEMCPYD2H_DATA_2_ID);

const memcpy = @import_module("<memcpy/get_params>", .{
     .width = width,
     .height = height,
     .MEMCPYH2D_1=MEMCPYH2D_DATA_1,
     .MEMCPYH2D_2=MEMCPYH2D_DATA_2,
     .MEMCPYD2H_1=MEMCPYD2H_DATA_1,
     .MEMCPYD2H_2=MEMCPYD2H_DATA_2
});

You must also pass the input/output tensor ID and color value pairs to cslc as parameters. Here <i> is the tensor index (1-4) and <color_id> is the numeric ID of the color to use for that tensor:

--params=MEMCPYH2D_DATA_<i>_ID:<color_id>
--params=MEMCPYD2H_DATA_<i>_ID:<color_id>

For example, to use color 1 for input tensor 1 and color 16 for output tensor 1: --params=MEMCPYH2D_DATA_1_ID:1 --params=MEMCPYD2H_DATA_1_ID:16. To stream the data into the device, you can either use a data task to read the data from the input tensor color or use a microthread (a lightweight hardware thread that can issue DSD operations without occupying the main compute engine) to read the data from a fabin_dsd. To bind a data task to an input color, call @bind_data_task at compile time:

const MEMCPYH2D_1_TASK_ID = @get_data_task_id(MEMCPYH2D_DATA_1);

comptime {
  // Task reads data on color MEMCPYH2D_DATA_1
  @bind_data_task(memcpyh2d_data_1_task, MEMCPYH2D_DATA_1_TASK_ID);
}

You can send data to an output tensor color using a fabout_dsd. For instance, assuming my_fabout_dsd and my_mem_buf_dsd are already defined:

@mov32(my_fabout_dsd, my_mem_buf_dsd, .{.async=true});

Use Copy Mode

To use copy mode to copy data to/from the device, you have to define the symbols for the tensors to be copied. For example, the following code defines a pointer ptr_A pointing to tensor A, and exports it.

// in top-level CSL file
const memcpy = @import_module("<memcpy/get_params>", .{
    .width = width,
    .height = height
});

layout {
  @set_rectangle(1, 1);
  @set_tile_code(0, 0, "pe_program.csl", .{ .memcpy_params = memcpy.get_params(0) });

  // export symbol names
  @export_name("A", [*]f32, true);
}

// in PE program CSL file pe_program.csl
param memcpy_params: comptime_struct;

const sys_mod = @import_module("<memcpy/memcpy>", memcpy_params);

var A = @zeros([4]f32);
var ptr_A : [*]f32 = &A;

comptime {
    @export_symbol(ptr_A, "A");
}

Launch Kernels

We can additionally use memcpy to launch a kernel function. The following is an example of the kernel launching protocol. This program exports two functions to the host: f1 and f2.

// in top-level CSL file
const memcpy = @import_module("<memcpy/get_params>", .{
    .width = width,
    .height = height
});

layout {
  @set_rectangle(1, 1);
  @set_tile_code(0, 0, "pe_program.csl", .{ .memcpy_params = memcpy.get_params(0) });

  // export symbol names
  @export_name("f1", fn()void);
  @export_name("f2", fn()void);
}

// in PE program CSL file pe_program.csl
param memcpy_params: comptime_struct;

const sys_mod = @import_module("<memcpy/memcpy>", memcpy_params);

fn f1() void {
  // do something
}

fn f2(my_arg: f32) void {
  // do something else
}

comptime {
  @export_symbol(f1);
  @export_symbol(f2);
}

Use Buffers

The compiler can insert buffers in the infrastructure to reduce the latency of the I/O. The buffer stores the wavelets from the I/O for one row of PEs while the core program rectangle is busy and cannot process the wavelets from the I/O. In other words, the buffer acts like a prefetch from the point of view of the computation. There are two kinds of buffers: one stores the data for host-to-device transfers, and the other stores the data for device-to-host transfers. The width of the former is configured by --width-west-buf, and the width of the latter is configured by --width-east-buf. By default, --width-west-buf=0 and --width-east-buf=0, i.e., no buffers are inserted. --width-west-buf=k means k columns of PEs are inserted to the West of the core kernel, and each PE can buffer 46 KB of data. If you have 500 PEs in a row, then 46 KB can buffer 23 wavelets per PE (recall that each wavelet holds 32 bits of data). If you want to stream or copy a tensor of size 100 per PE, then --width-west-buf=5 can buffer the whole tensor. When compiling with --width-west-buf=k and --width-east-buf=p, you must specify --fabric-offsets=x,y such that x >= 4 + k and y >= 1, and --fabric-dims=dim_x,dim_y such that dim_x >= x + width + 3 + p and dim_y >= y + height + 1, where width and height are the width and height of the program rectangle.

SdkRuntime Host API

See SdkRuntime API Reference for full documentation of the SdkRuntime Python host API. The SdkRuntime Python host API supports memory transfers and kernel launches through the functions memcpy_h2d(), memcpy_d2h() and launch(): memcpy_h2d() is used for host-to-device data transfers, memcpy_d2h() is used for device-to-host data transfers, and launch() is used for kernel launches. Each function can be a blocking or nonblocking call, depending on the parameter nonblock of the API. If blocking mode (nonblock=False) is selected, the API waits until the operation is done. Otherwise, the function returns before the operation even starts. SdkRuntime can aggregate multiple nonblocking operations together to reduce the latency. However, you must take care to avoid race conditions in nonblocking mode. For example, if you have two memcpy_d2h() calls to the same destination, the content of the destination is undefined if both operations are nonblocking.

Instantiate SdkRuntime

You’ll need to import the SdkRuntime module, as well as the MemcpyDataType and MemcpyOrder modules for specifying data type and ordering of tensors. To create an SdkRuntime object, pass the directory which contains the ELF files produced by the compiler, and the IP address of the WSE, if running on hardware, to SdkRuntime(). You can load the ELFs by load() and start the simulator or WSE with run(). After that, you can do any operation, either memory transfers or kernel launches. Finally, call stop() to shut down the simulator or WSE.

from cerebras.sdk.runtime.sdkruntimepybind import SdkRuntime
from cerebras.sdk.runtime.sdkruntimepybind import MemcpyDataType
from cerebras.sdk.runtime.sdkruntimepybind import MemcpyOrder

simulator = SdkRuntime(args.name, cmaddr=args.cmaddr)
simulator.load()
simulator.run()
# a sequence of operations
simulator.stop()

Instantiating the SdkRuntime object uses slightly different syntax if you are compiling and running an SDK program on a Wafer-Scale Cluster in appliance mode. See Running SDK on a Wafer-Scale Cluster.

memcpy_h2d() and memcpy_d2h()

The function memcpy_h2d() transfers a tensor from host to device using either streaming mode or copy mode.

Parameter	Description
`streaming`	`True` for `streaming` mode, `False` for `copy` mode.
ROI (`x, y, w, h`)	Region of interest: a subrectangle starting at `(x,y)` with size `(w, h)`. The origin `(0, 0)` is the top-left PE in your program rectangle — in absolute coordinates, the PE at the coordinates given by `--fabric-offsets`. The ROI must lie within the program rectangle.
`l`	Number of elements (wavelets) per PE.
`data_type`	For `copy` mode: `MemcpyDataType.MEMCPY_16BIT` or `MemcpyDataType.MEMCPY_32BIT`.
`order`	`MemcpyOrder.ROW_MAJOR` or `MemcpyOrder.COL_MAJOR` for the input/output tensor of the form `A[h][w][l]`.
`nonblock`	Whether the operation is blocking or nonblocking.
`dest`	The color for this host-to-device transfer if `streaming=True`, or the symbol if `streaming=False`.

memcpy_h2d(dest, src, x, y, w, h, l, streaming, data_type, order, nonblock)

Similarly, the function memcpy_d2h() transfers a tensor from device to host using either streaming mode or copy mode. The first parameter dest is the host tensor to receive the data from the device. The second parameter src is the color associated with this device-to-host transfer if streaming=True or the device symbol from which to copy if streaming=False. All other parameters are the same as memcpy_h2d().

memcpy_d2h(dest, src, x, y, w, h, l, streaming, data_type, order, nonblock)

The parameter order of memcpy_h2d() and memcpy_d2h() specifies either row-major or column-major. In both cases, the host tensor from which or to which data is copied is a 1D array of length w*h*l, where w and h are the width and height of the region of interest and l is the number of elements per PE to copy.

Row-major
Column-major

Mapping from 1D to [w][h][l], l is the fastest varying dimension — elements contiguous on a PE will be contiguous on the host.

Mapping from 1D to [w][h][l], w is the fastest varying dimension. Column-major delivers better bandwidth than row-major.

memcpy_h2d() and memcpy_d2h() support both 16-bit and 32-bit data transfer via copy mode or streaming mode. When using memcpy_h2d() for a 16-bit tensor, you must perform zero extension from 16-bit to 32-bit. When using memcpy_d2h() for a 16-bit tensor, the returned array will contain 32-bit data where the higher 16 bits are zero. You have to strip out the higher 16 bits. See the sdk_utils module documentation for utilities to help perform this data transformation.

launch()

The launch() function performs remote kernel launches of host-callable functions.

Parameter	Description
`sym`	The symbol of the host-callable function.
positional arguments	Match the arguments of the host-callable function.
`nonblock`	Keyword argument specifying whether the kernel launch is blocking or nonblocking.

For example, to launch a host-callable function my_fun with two arguments of type f32 in blocking mode, the call would look as follows:

my_fun_symbol = runner.get_symbol('my_fun')
runner.launch(my_fun_symbol, 1.0, 2.0, nonblock=False)

Start Here

Tutorials

Development Guides

Debugging

CSL Language Reference

Host API Reference

Host Runtime and Tensor Streaming

Set Up the memcpy Infrastructure

Use Streaming Mode

Use Copy Mode

Launch Kernels

Use Buffers

SdkRuntime Host API

Instantiate SdkRuntime

memcpy_h2d() and memcpy_d2h()

launch()

​Set Up the memcpy Infrastructure

​Use Streaming Mode

​Use Copy Mode

​Launch Kernels

​Use Buffers

​SdkRuntime Host API

​Instantiate SdkRuntime

​memcpy_h2d() and memcpy_d2h()

​launch()

Set Up the memcpy Infrastructure

Use Streaming Mode

Use Copy Mode

Launch Kernels

Use Buffers

SdkRuntime Host API

Instantiate SdkRuntime

memcpy_h2d() and memcpy_d2h()

launch()