How Systolic Array Works: A Case Study on Matrix Multiplication
Author: Jie Wang (jiewang@cs.ucla.edu)
This page gives a detailed explanation about the AutoSA generated systolic array architecture for matrix multiplication.
Generating the Systolic Array
We will use the example code in ${AUTOSA_ROOT}/autosa_tests/mm/kernel.c.
#pragma scop
for (int i = 0; i < 64; i++)
for (int j = 0; j < 64; j++) {
C[i][j] = 0;
for (int k = 0; k < K64; k++)
C[i][j] = C[i][j] + A[i][k] * B[j][k];
}
#pragma endscop
Use the following command to generate the systolic array.
./autosa ./autosa_tests/mm/kernel.c --config=./autosa_config/autosa_config.json --target=autosa_hls_c --output-dir=./autosa.tmp/output --sa-sizes="{kernel[]->space_time[3];kernel[]->array_part[16,16,16];kernel[]->latency[8,8];kernel[]->simd[2]}" --simd-info=./autosa_tests/mm/simd_info.json --hls
This will generate a \(2\times 2\) 2D systolic array as shown below.
Understanding the Systolic Array
The systolic array architecture is composed of two parts: the processing elements (PE) and the I/O network. We will explain these two components in sequence.
Processing Elements (PE)
Below is the AutoSA generated HLS code for the PE.
1/* Module Definition */
2void PE(int idx, int idy, hls::stream<A_t2> &fifo_A_in, hls::stream<A_t2> &fifo_A_out, hls::stream<B_t2> &fifo_B_in, hls::stream<B_t2> &fifo_B_out, hls::stream<float> &fifo_C_drain_out) {
3#pragma HLS INLINE OFF
4 /* Variable Declaration */
5 int p0 = idx, p1 = idy; // module id
6 A_t1 local_A[1][2];
7 #pragma HLS ARRAY_PARTITION variable=local_A dim=0 complete
8 B_t1 local_B[1][2];
9 #pragma HLS ARRAY_PARTITION variable=local_B dim=0 complete
10 C_t1 local_C[8][8];
11 #pragma HLS RESOURCE variable=local_C core=RAM_2P_BRAM
12 /* Variable Declaration */
13
14 for (ap_uint<3> c0 = 0; c0 <= 3; c0 += 1)
15 for (ap_uint<3> c1 = 0; c1 <= 3; c1 += 1) {
16 // array
17 // pe
18 // latency
19 for (ap_uint<4> c6 = 0; c6 <= 7; c6 += 1) {
20 // latency
21 for (ap_uint<4> c7 = 0; c7 <= 7; c7 += 1) {
22 #pragma HLS PIPELINE II=1
23 // simd
24 // hls_unroll
25 local_C[c7][c6] = 0;
26 }
27 }
28 for (ap_uint<3> c2 = 0; c2 <= 3; c2 += 1) {
29 // array
30 // pe
31 for (ap_uint<4> c5 = 0; c5 <= 7; c5 += 1) {
32 // latency
33 for (ap_uint<4> c6 = 0; c6 <= 7; c6 += 1) {
34 // latency
35 for (ap_uint<4> c7 = 0; c7 <= 7; c7 += 1) {
36 #pragma HLS PIPELINE II=1
37 {
38 {
39 A_t2 fifo_data;
40 fifo_data = fifo_A_in.read();
41 for (ap_uint<2> n = 0; n < 2; n++) {
42 #pragma HLS UNROLL
43 union {unsigned int ui; float ut;} u;
44 u.ui = (unsigned int)fifo_data(31, 0);
45 local_A[0][n] = u.ut;
46 fifo_data = fifo_data >> 32;
47 }
48 }
49 {
50 B_t2 fifo_data;
51 fifo_data = fifo_B_in.read();
52 for (ap_uint<2> n = 0; n < 2; n++) {
53 #pragma HLS UNROLL
54 union {unsigned int ui; float ut;} u;
55 u.ui = (unsigned int)fifo_data(31, 0);
56 local_B[0][n] = u.ut;
57 fifo_data = fifo_data >> 32;
58 }
59 }
60 // simd
61 for (ap_uint<2> c8 = 0; c8 <= 1; c8 += 1) {
62 #pragma HLS UNROLL
63 local_C[c7][c6] = (local_C[c7][c6] + (local_A[0][c8] * local_B[0][c8]));
64 }
65 if (c2 == 3 && c5 == 7)
66 fifo_C_drain_out.write(local_C[c7][c6]);
67 {
68 B_t2 fifo_data;
69 union {unsigned int ui; float ut;} u1, u0;
70 u1.ut = local_B[0][1];
71 u0.ut = local_B[0][0];
72 fifo_data = (ap_uint<32>(u1.ui), ap_uint<32>(u0.ui));
73 fifo_B_out.write(fifo_data);
74 }
75 {
76 A_t2 fifo_data;
77 union {unsigned int ui; float ut;} u1, u0;
78 u1.ut = local_A[0][1];
79 u0.ut = local_A[0][0];
80 fifo_data = (ap_uint<32>(u1.ui), ap_uint<32>(u0.ui));
81 fifo_A_out.write(fifo_data);
82 }
83 }
84 }
85 }
86 }
87 }
88 }
89}
90/* Module Definition */
In this 2D systolic array, data of matrix A are reused horizontally across PEs, data of matrix B are reused vertically. Each PE computes elements of matrix C locally. After the computation is done, final results of matrix C will be drained out to the external memory.
The PE interface (line 2) contains the following components:
Module index (
idx,idy): Indices of the PE module.FIFO (
fifo_A_in,fifo_A_out,fifo_B_in,fifo_B_out,fifo_C_drain_out): FIFOs for transfering data.
While generating this array, we applied latency hiding on the orginal loops \(i\) and \(j\) with the factor \((8,8)\), and SIMD vectorization on the loop \(k\) with a factor of 2. With the latency hiding, each PE will compute a tile of \(8\times 8\) of the matrix C. With the SIMD vectorization, at each cycle, two elements of matrix A and two elements of matrix B are required to update the local elements of matrix C.
With this knowledge, we could take a look at the local variable declarations in lines 5-11 now. Line 5 is simply storing the module indices. Lines 6-11 allocate local storage inside PEs for storing the data of matrix A, B, and C.
The rest of the code performs the computation. At each cycle, PE reads data of matrix A and B from neighbor PEs at lines 38-59 and passes the data to neighbor PEs at lines 67-82. PE performs the computation at lines 61-64. When the final results of matrix C are derived, PE writes out the final results at lines 65-66.
I/O Network
I/O network is composed of a series of I/O modules for transferring data between the external memory and PEs. We will use the I/O modules of array A as an example.
There are two types of I/O modules for array A:
Level-3 (L3) I/O modules: modules that read data from the external memory and send to the array.
Level-2 (L2) I/O modules: modules that pass data between each other. Data that belong to the PEs that the module is connected to are kept locally, the rest data are passed to the downstreaming I/O modules.
Below is the code of the L3 I/O module.
1/* Module Definition */
2void A_IO_L3_in(A_t8 *A, hls::stream<A_t8> &fifo_A_local_out) {
3#pragma HLS INLINE OFF
4 /* Variable Declaration */
5 /* Variable Declaration */
6
7 for (ap_uint<3> c0 = 0; c0 <= 3; c0 += 1)
8 for (ap_uint<3> c1 = 0; c1 <= 3; c1 += 1)
9 for (ap_uint<3> c2 = 0; c2 <= 3; c2 += 1) {
10 // array
11 // io_L3
12 for (ap_uint<2> c3 = 0; c3 <= 1; c3 += 1) {
13 // io_L2
14 for (ap_uint<4> c4 = 0; c4 <= 7; c4 += 1) {
15 // access_coalesce
16 for (ap_uint<2> c5 = 0; c5 <= 1; c5 += 1) {
17 #pragma HLS PIPELINE II=1
18 {
19 A_t8 fifo_data;
20 fifo_data = A[128*c0 + 2*c2 + 64*c3 + 8*c4 + c5];
21 fifo_A_local_out.write(fifo_data);
22 }
23 }
24 }
25 }
26 }
27}
28/* Module Definition */
In this design, we apply the array partitioning on the original loops \(i\), \(j\), and \(k\) with the factors \((16,16,16)\). The orignal loop bounds for these three loops are \((64,64,64)\). Therefore, array partitioning loops at lines 7-9 have loop bounds of \((4,4,4)\).
When transferring the data to the PEs, we will pass data through the chain of L2 I/O modules. In this design, there are two such modules. The loop for traversing the L2 I/O modules is at line L2. Inside each L2 I/O module, we will need to load the data tile required by the PEs that it is connected to.
With the array partitioning factors \((16,16,16)\), at each array partition, a sub tile of matrix A with the size \(16\times 16\) is loaded from the external memory. As this array have the dimension of \(2\times 2\), each L2 I/O module will store a tile with the size \(8\times 16\). The loops for loading the data tiles for each I/O modules can be found at lines 14-16. Note that AutoSA will pack data together to increase the I/O througput. In this case, every 8 elements are packed together. Therefore, the size of the local tile is \(8\times 2\), with a data width of 8 data elements.
Next, we will look at the L2 I/O module. The figure below shows the micro-architecture of the L2 I/O module.
L2 I/O module loads data from the upstream I/O modules, keeps the data that belong to it, and sends the rest to the downstream modules. For I/O modules with local buffers inside, AutoSA automatically applies double buffering to overlap the data transfer betwen the I/O modules and data transfer to/from PEs.
Below is the code of L2 I/O module.
1/* Module Definition */
2void A_IO_L2_in(int idx, hls::stream<A_t8> &fifo_A_in, hls::stream<A_t8> & fifo_A_out, hls::stream<A_t2> &fifo_A_local_out) {
3#pragma HLS INLINE OFF
4 /* Variable Declaration */
5 int p0 = idx; // module id
6 A_t8 local_A_ping[8][2];
7 #pragma HLS RESOURCE variable=local_A_ping core=RAM_2P_BRAM
8 A_t8 local_A_pong[8][2];
9 #pragma HLS RESOURCE variable=local_A_pong core=RAM_2P_BRAM
10 bool arb = 0;
11 bool inter_trans_en = 1;
12 bool intra_trans_en = 0;
13 int c0, c0_prev;
14 int c1, c1_prev;
15 int c2, c2_prev;
16 /* Variable Declaration */
17
18 {
19 for (ap_uint<3> c0 = 0; c0 <= 3; c0 += 1)
20 for (ap_uint<3> c1 = 0; c1 <= 3; c1 += 1)
21 for (ap_uint<3> c2 = 0; c2 <= 3; c2 += 1) {
22 // array
23 // io_L3
24 {
25 if (arb == 0) {
26 A_IO_L2_in_inter_trans(
27 /* module id */ idx,
28 /* host iter */ c0,
29 /* host iter */ c1,
30 /* host iter */ c2,
31 /* array */ local_A_pong,
32 /* fifo */ fifo_A_in,
33 /* fifo */ fifo_A_out,
34 /* enable */ inter_trans_en
35 );
36 A_IO_L2_in_intra_trans(
37 /* module id */ idx,
38 /* host iter */ c0_prev,
39 /* host iter */ c1_prev,
40 /* host iter */ c2_prev,
41 /* array */ local_A_ping,
42 /* fifo */ fifo_A_local_out,
43 /* enable */ intra_trans_en
44 );
45 } else {
46 A_IO_L2_in_inter_trans(
47 /* module id */ idx,
48 /* host iter */ c0,
49 /* host iter */ c1,
50 /* host iter */ c2,
51 /* array */ local_A_ping,
52 /* fifo */ fifo_A_in,
53 /* fifo */ fifo_A_out,
54 /* enable */ inter_trans_en
55 );
56 A_IO_L2_in_intra_trans(
57 /* module id */ idx,
58 /* host iter */ c0_prev,
59 /* host iter */ c1_prev,
60 /* host iter */ c2_prev,
61 /* array */ local_A_pong,
62 /* fifo */ fifo_A_local_out,
63 /* enable */ intra_trans_en
64 );
65 }
66 intra_trans_en = 1;
67 arb = !arb;
68 c0_prev = c0;
69 c1_prev = c1;
70 c2_prev = c2;
71 }
72 }
73 if (arb == 0) {
74 A_IO_L2_in_intra_trans(
75 /* module id */ idx,
76 /* host iter */ c0_prev,
77 /* host iter */ c1_prev,
78 /* host iter */ c2_prev,
79 /* array */ local_A_ping,
80 /* fifo */ fifo_A_local_out,
81 /* enable */ intra_trans_en
82 );
83 } else {
84 A_IO_L2_in_intra_trans(
85 /* module id */ idx,
86 /* host iter */ c0_prev,
87 /* host iter */ c1_prev,
88 /* host iter */ c2_prev,
89 /* array */ local_A_pong,
90 /* fifo */ fifo_A_local_out,
91 /* enable */ intra_trans_en
92 );
93 }
94 }
95}
96/* Module Definition */
Lines 6-9 define the double buffers inside the I/O module.
Lines 19-95 performs the double buffering to overlap the data transfer between I/O modules (defined in the function A_IO_L2_in_inter_trans) and data transfer to/from PEs (defined in the function A_IO_L2_in_intra_trans).
Please refer to the generated code for more details of the functions A_IO_L2_in_inter_trans and A_IO_L2_in_intra_trans.
The similar principles apply to the other I/O modules. Together with both the I/O modules and PEs, we have a complete functional systolic array that can be synthesized and executed on FPGAs.
Note
When adding the argument --host-serialize to the AutoSA command, the data of each array will be serialized on the host and transfered to the systolic array. AutoSA will introduce an additional I/O module for loading/writing the serialized data from/to the external memory before the original I/O modules. Feel free to try it out and compare with the code without serialization. The major benefit of using host serialization is to increase the DDR bus width and burst length to improve the effective DRAM bandwidth.