> ## Documentation Index > Fetch the complete documentation index at: https://sdk.cerebras.ai/llms.txt > Use this file to discover all available pages before exploring further. # 5. Multiple PEs > Scale your CSL program to run across multiple processing elements by configuring a multi-PE layout and copying data to and from each PE. 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? 1. Modify `@set_rectangle` to reflect the new program rectangle. 2. Modify the `memcpy` infrastructure to reflect the use of multiple PEs. 3. 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. ```csl theme={"languages":{"custom":["/languages/csl-tmlanguage.json"]}} // matrix dimensions on each PE param M: i16; param N: i16; // number of PEs in program param width: i16; const memcpy = @import_module("", .{ .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 `` 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: ```python theme={"languages":{"custom":["/languages/csl-tmlanguage.json"]}} #!/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: ```python theme={"languages":{"custom":["/languages/csl-tmlanguage.json"]}} 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: ```python theme={"languages":{"custom":["/languages/csl-tmlanguage.json"]}} 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](/csl/tutorials/gemv-01-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: ```bash theme={"languages":{"custom":["/languages/csl-tmlanguage.json"]}} $ 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. ## Next 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. This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.