"Function"

[Objective]

Cache 메모리 활용도를 높이기 위해 행렬을 작은 조각(Tile)으로 나누어 연산하는 "Tile Matrix Multiplication"을 Assembly 언어로 구현하기

[Background]

· 대규모 행렬 연산 시, 데이터 전체가 On-Chip Cache(SRAM)에 들어가지 않음

→ Cache에 저장되었던 데이터가 다음 Loop에서 다시 쓰이기 전에 Cache에서 삭제됨

∴ 행렬을 Cache 크기에 맞는 작은 Tile 단위로 쪼개어 연산

→ 메모리 명령어 횟수는 늘어날 수 있지만, Cache 적중률이 획기적으로 개선되어 전체적인 연산 성능이 크게 향상

1
2
3
4
5
6
7
8
9
10
11
12
13
14
// A, B, C: Operand matrices for C += A * B
// m, n, k: GEMM tile dimensions
// M, N, K: Matrix dimensions
void GEMM(int* A, int* B, int* C, size_t m, size_t n, size_t k, size_t M, size_t N, size_t K) {
    // Tile-level nested loops
    for(size_t y = 0; y < N; y += n) {
        for(size_t z = 0; z < K; z += k) {
            for(size_t x = 0; x < M; x += m) {
                // Subroutine for matrix multiplication within a tile
                MMA(&A[x * K + z], &B[z * N + y], &C[x * N + y], m, n, k, M, N, K);
            }
        }
    }
}
Colored by Color Scripter
cs

· Tile 크기인 {m, n, k} 단위로 증가 = Tile 덩어리 단위로 행렬을 순회

· &A[x * K + z] : 행렬 A에서 현재 연산할 Tile의 첫 번째 원소가 위치한 메모리 주소

· &B[z * N + y] : 행렬 B에서 현재 연산할 Tile의 첫 번째 원소가 위치한 메모리 주소

· &C[x * N + z] : 결과 값을 누적할 행렬 C 내의 해당 Tile 첫 번째 원소 주소

[Implementation]

▣ Register State

<reg_state>

· x1 : 함수 Return Address 저장

· x2 : Stack Pointer

· x8 : Frame Pointer

· x10 : 행렬 A의 Base Address

· x11 : 행렬 B의 Base Address
· x12 : 행렬 C의 Base Address

· x13 : Tile 크기 m

· x14 : Tile 크기 n

· x15 : Tile 크기 k

· x16 : 전체 행렬 크기 M

· x17 : 전체 행렬 크기 N

· x28 : 전체 행렬 크기 K (RISC-V에서 인자를 전달하는 Register는 x10~x17이므로, K값 전달을 위해 x28 Register 사용)

· x9, x18~x27 : Saved Register이므로, 함수 호출 전후로 값이 보존되어야 함

▣ Memory State

<mem_state - A>

· A 행렬은 1024 메모리 주소부터 1336 메모리 주소까지 행 방향 순서대로 나열

<mem_state - B>

· B 행렬은 2016 메모리 주소부터 2200 메모리 주소까지 행 방향 순서대로 나열

<mem_state - C>

· C 행렬은 3008 메모리 주소부터 3480 메모리 주소까지 행 방향 순서대로 나열

▣ Code Description

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
# Function: GEMM(A, B, C, m, n, k, M, N, K)
# x10 = A,      x11 = B,     x12 = C,
# x13 = m (2),  x14 = n (3), x15 = k (2),
# x16 = M (10), x17 = N (6), x28 = K (4)
GEMM:
        # Backup Registers
        addi    x2,     x2,     -104    # Allocate Space For 13 Registers(x1, x8, x9, x18~x27)
        sd      x1,     96(x2)          # Store Return Address
        sd      x8,     88(x2)          # Store Frame Pointer
        sd      x9,     80(x2)          # Store Loop y Counter
        sd      x18,    72(x2)          # Store A Base Address
        sd      x19,    64(x2)          # Store B Base Address
        sd      x20,    56(x2)          # Store C Base Address
        sd      x21,    48(x2)          # Store Tile Size m
        sd      x22,    40(x2)          # Store Tile Size n
        sd      x23,    32(x2)          # Store Tile Size k
        sd      x24,    24(x2)          # Store Matrix Size M
        sd      x25,    16(x2)          # Store Matrix Size N
        sd      x26,    8(x2)           # Store Matrix Size K
        sd      x27,    0(x2)           # Store Loop z Counter
        addi    x8,     x2,     104     # Setting Frame Pointer
 
        # Copy To Saved Register To Prevent Original Data Overwritten
        add     x18,    x10,    x0      # Copy A Base Address
        add     x19,    x11,    x0      # Copy B Base Address
        add     x20,    x12,    x0      # Copy C Base Address
        add     x21,    x13,    x0      # Copy Tile Size m
        add     x22,    x14,    x0      # Copy Tile Size n
        add     x23,    x15,    x0      # Copy Tile Size k
        add     x24,    x16,    x0      # Copy Matrix Size M
        add     x25,    x17,    x0      # Copy Matrix Size N
        add     x26,    x28,    x0      # Copy Matrix Size K
 
        # Loop y = for (y = 0; y < N; y += n)
        add     x9,     x0,     x0      # y = 0
GEMM_LOOP_y:
        bge     x9,     x25,    GEMM_DONE       # If y >= N, Finish Loop
 
        # Loop z = for (z = 0; z < K; z += k)
        add     x27,    x0,     x0      # z = 0
GEMM_LOOP_z:
        bge     x27,    x26,    GEMM_NEXT_y     # If z >= K, Next y
 
        # Loop x = for (x = 0; x < M; x += m)
        add     x5,     x0,     x0      # x = 0
GEMM_LOOP_x:
        bge     x5,     x24,    GEMM_NEXT_z     # If x >= M, Next z
 
        # Calculate A Tile Address
        mul     x6,     x5,     x26     # x * K
        add     x6,     x6,     x27     # (x * K) + z
        slli    x6,     x6,     3       # {(x * K) + z} * 8
        add     x10,    x18,    x6      # A Base Address + Offset
 
        # Calculate B Tile Address
        mul     x6,     x27,    x25     # z * N
        add     x6,     x6,     x9      # (z * N) + y
        slli    x6,     x6,     3       # {(z * N) + y} * 8
        add     x11,    x19,    x6      # B Base Address + Offset
 
        # Calculate C Tile Address
        mul     x6,     x5,     x25     # x * N
        add     x6,     x6,     x9      # (x * N) + y
        slli    x6,     x6,     3       # {(x * N) + y} * 8
        add     x12,    x20,    x6      # C Base Address + Offset
 
        # Setting For MMA
        add     x13,    x21,    x0      # Set Tile Size m
        add     x14,    x22,    x0      # Set Tile Size n
        add     x15,    x23,    x0      # Set Tile Size k
        add     x16,    x24,    x0      # Set Matrix Size M
        add     x17,    x25,    x0      # Set Matrix Size N
        add     x28,    x26,    x0      # Set Matrix Size K
 
        # Backup Loop x Counter In Stack
        addi    x2,     x2,     -8      # Allocate Space For 1 Register
        sd      x5,     0(x2)           # Store x Value To Prevent Overwritten
 
        jal     x1,     MMA             # Call Subroutine
 
        ld      x5,     0(x2)           # Restore x Value
        addi    x2,     x2,     8       # Restore Stack
 
        # Next x
        add     x5,     x5,     x21     # x += m
        jal     x0,     GEMM_LOOP_x     # Restart Loop x
 
GEMM_NEXT_z:
        # Next z
        add     x27,    x27,    x23     # z += k
        jal     x0,     GEMM_LOOP_z     # Restart Loop z
 
GEMM_NEXT_y:
        # Next y
        add     x9,     x9,     x22     # y += n
        jal     x0,     GEMM_LOOP_y     # Restart Loop y
 
GEMM_DONE:
        # Restore Register And End Function
        ld      x27,    0(x2)           # Restore Loop z Counter
        ld      x26,    8(x2)           # Restore Matrix Size K
        ld      x25,    16(x2)          # Restore Matrix Size N
        ld      x24,    24(x2)          # Restore Matrix Size M
        ld      x23,    32(x2)          # Restore Tile Size k
        ld      x22,    40(x2)          # Restore Tile Size n
        ld      x21,    48(x2)          # Restore Tile Size m
        ld      x20,    56(x2)          # Restore C Base Address
        ld      x19,    64(x2)          # Restore B Base Address
        ld      x18,    72(x2)          # Restore A Base Address
        ld      x9,     80(x2)          # Restore Loop y Counter
        ld      x8,     88(x2)          # Restore Frame Pointer
        ld      x1,     96(x2)          # Restore Return Address
        addi    x2,     x2,     104     # Restore Stack Location
        jalr    x0,     0(x1)           # End Function And Return To Return Address
cs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
# Function: MMA(A, B, C, m, n, k, M, N, K)
# x10 = A,      x11 = B,     x12 = C,
# x13 = m (2),  x14 = n (3), x15 = k (2)
# x16 = M (10), x17 = N (6), x28 = K (4)
MMA:
        # Backup Registers
        addi    x2,     x2,     -104    # Allocate Space For 13개 Registers(x1, x8, x9, x18~x27)
        sd      x1,     96(x2)          # Store Return Address
        sd      x8,     88(x2)          # Store Frame Pointer
        sd      x9,     80(x2)          # Store Loop j Counter
        sd      x18,    72(x2)          # Store A Base Address
        sd      x19,    64(x2)          # Store B Base Address
        sd      x20,    56(x2)          # Store C Base Address
        sd      x21,    48(x2)          # Store Tile Size m
        sd      x22,    40(x2)          # Store Tile Size n
        sd      x23,    32(x2)          # Store Tile Size k
        sd      x24,    24(x2)          # Store Matrix Size M
        sd      x25,    16(x2)          # Store Matrix Size N
        sd      x26,    8(x2)           # Store Matrix Size K
        sd      x27,    0(x2)           # Store Loop i Counter
        addi    x8,     x2,     104     # Setting Frame Pointer
 
        # Copy To Saved Register To Prevent Original Data Overwritten
        add     x18,    x10,    x0      # Copy A Base Address
        add     x19,    x11,    x0      # Copy B Base Address
        add     x20,    x12,    x0      # Copy C Base Address
        add     x21,    x13,    x0      # Copy Tile Size m
        add     x22,    x14,    x0      # Copy Tile Size n
        add     x23,    x15,    x0      # Copy Tile Size k
        add     x24,    x16,    x0      # Copy Matrix Size M
        add     x25,    x17,    x0      # Copy Matrix Size N
        add     x26,    x28,    x0      # Copy Matrix Size K
 
        # Loop j
        add     x9,     x0,     x0      # j = 0
MMA_LOOP_j:
        bge     x9,     x22,    MMA_DONE        # If j >= n, End Finish Loop
 
        # Loop w
        add     x6,     x0,     x0      # w = 0
MMA_LOOP_w:
        bge     x6,     x23,    MMA_NEXT_j      # If w >= k, Next j
 
        # Load B[w * N + j]
        mul     x7,     x6,     x25     # w * N
        add     x7,     x7,     x9      # (w * N) + j
        slli    x7,     x7,     3       # {(w * N) + j} * 8
        add     x7,     x19,    x7      # B[w * N + j] Address
        ld      x29,    0(x7)           # Load B[w * N + j] Value Which In Address
 
        # Loop i
        add     x27,    x0,     x0      # i = 0
MMA_LOOP_i:
        bge     x27,    x21,    MMA_NEXT_w      # If i >= m, Next w
 
        # Load A[i * K + w]
        mul     x7,     x27,    x26     # i * K
        add     x7,     x7,     x6      # (i * K) + w
        slli    x7,     x7,     3       # {(i * K) + w} * 8
        add     x7,     x18,    x7      # A[i * K + w] Address
        ld      x28,    0(x7)           # Load A[i * K + w] Value Which In Address
 
        # Load C[i * N + j]
        mul     x5,     x27,    x25     # i * N
        add     x5,     x5,     x9      # (i * N) + j
        slli    x5,     x5,     3       # {(i * N) + j} * 8
        add     x30,    x20,    x5      # C[i * N + j] Address
        ld      x31,    0(x30)          # Load C[i * N + j] Value Which In Address
 
        # A * B MAC
        mul     x7,     x28,    x29     # A * B
        add     x31,    x31,    x7      # C += A * B
 
        sd      x31,    0(x30)          # Update C Value
 
        addi    x27,    x27,    1       # i++
        jal     x0,     MMA_LOOP_i      # Restart Loop i
 
MMA_NEXT_w:
        addi    x6,     x6,     1       # w++
        jal     x0,     MMA_LOOP_w      # Restart Loop w
 
MMA_NEXT_j:
        addi    x9,     x9,     1       # j++
        jal     x0,     MMA_LOOP_j      # Restart LOOP j
 
MMA_DONE:
        # Restore Register And End Function
        ld      x27,    0(x2)           # Restore Loop i Counter
        ld      x26,    8(x2)           # Restore Matrix Size K
        ld      x25,    16(x2)          # Restore Matrix Size N
        ld      x24,    24(x2)          # Restore Matrix Size M
        ld      x23,    32(x2)          # Restore Tile Size k
        ld      x22,    40(x2)          # Restore Tile Size n
        ld      x21,    48(x2)          # Restore Tile Size m
        ld      x20,    56(x2)          # Restore C Base Address
        ld      x19,    64(x2)          # Restore B Base Address
        ld      x18,    72(x2)          # Restore A Base Address
        ld      x9,     80(x2)          # Restore Loop j Counter
        ld      x8,     88(x2)          # Restore Frame Pointer
        ld      x1,     96(x2)          # Restore Return Address
        addi    x2,     x2,     104     # Restore Stack Location
        jalr    x0,     0(x1)           # End Function And Return To Return Address
cs