LeetGPU | 02_Matrix_Multiplication

关注 LeetGPU repository，了解更多内容。

Problem Description#

Write a program that multiplies two matrices of 32-bit floating point numbers on a GPU. Given matrix $A$ of dimensions $M \times N$ and matrix $B$ of dimensions $N \times K$ , compute the product matrix $C = A \times B$ , which will have dimensions $M \times K$ . All matrices are stored in row-major format.

Implementation Requirements#

Use only native features (external libraries are not permitted).
The solve function signature must remain unchanged.
The final result must be stored in matrix C.

Example 1#

Input:

Matrix $A$ ( $2 \times 2$ ):

\begin{bmatrix} 1.0 & 2.0 \\ 3.0 & 4.0 \end{bmatrix}

Matrix $B$ ( $2 \times 2$ ):

\begin{bmatrix} 5.0 & 6.0 \\ 7.0 & 8.0 \end{bmatrix}

Output:

Matrix $C$ ( $2 \times 2$ ):

\begin{bmatrix} 19.0 & 22.0 \\ 43.0 & 50.0 \end{bmatrix}

Example 2#

Input:

Matrix $A$ ( $1 \times 3$ ):

\begin{bmatrix} 1.0 & 2.0 & 3.0 \end{bmatrix}

Matrix $B$ ( $3 \times 1$ ):

\begin{bmatrix} 4.0 \\ 5.0 \\ 6.0 \end{bmatrix}

Output:

Matrix $C$ ( $1 \times 1$ ):

\begin{bmatrix} 32.0 \end{bmatrix}

Constraints#

1 <= M, N, K <= 8192.
Performance is measured with M = 8192, N = 6144, K = 4096.

CUDA#

Approach#

CUDA 实现采用 16 x 16 的 shared-memory tiling。每个 block 包含 16 x 16 个 thread，负责计算输出矩阵 C 中一个 16 x 16 的 tile；每个 thread 最终负责一个输出元素。

矩阵乘法需要沿 N 维度完成内积。kernel 将这一维按 16 个元素分成多个 reduction tile。每轮中，block 内的 thread 协作完成两次加载：从 A 读取当前输出行对应的数据写入 A_shared，同时从 B 读取当前输出列对应的数据写入 B_shared。

数据进入 shared memory 后，第一次 __syncthreads() 确保整个 tile 已加载完成。随后，每个 thread 从 A_shared 的一行和 B_shared 的一列取数，执行 16 次乘加，并把结果累积到寄存器变量 S。第二次 __syncthreads() 确保所有 thread 都已使用完当前 tile，之后 shared memory 才能被下一轮数据覆盖。

这样，从 global memory 读取的每个输入值都可以被同一 block 中的多个 thread 复用，从而减少重复的 global-memory 访问。对于不能被 16 整除的矩阵尺寸，加载阶段把越界位置填充为 0.0f，最终写回时再检查 row < M && col < K，因此不会越界，也不会影响正确结果。

Solution#

1
#include <cuda_runtime.h>
2

3
#define BLOCK_SIZE  16
4
__global__ void matrix_multiplication_kernel(const float* A, const float* B, float* C, int M, int N,
5
                                             int K) {
6
    __shared__ float A_shared[BLOCK_SIZE][BLOCK_SIZE];
7
    __shared__ float B_shared[BLOCK_SIZE][BLOCK_SIZE];
8
    int row = threadIdx.y + blockIdx.y * blockDim.y;
9
    int col = threadIdx.x + blockIdx.x * blockDim.x;
10
    int tx = threadIdx.x;
11
    int ty = threadIdx.y;
12
    float S = 0.0f;
13
    for (int i = 0 ; i < (N + BLOCK_SIZE - 1) / BLOCK_SIZE ; i++){
14
        int col2 = tx + i * BLOCK_SIZE;
15
        int row2 = ty + i * BLOCK_SIZE;
16
        A_shared[ty][tx] = (row < M && col2 < N) ? A[row * N + col2] : 0.0f;
17
        B_shared[ty][tx] = (row2 < N && col < K) ? B[row2 * K + col] : 0.0f;
18
        __syncthreads();
19
        for (int j = 0 ; j <  BLOCK_SIZE ; j++){
20
            S += A_shared[ty][j] * B_shared[j][tx];
21
        }
22
        __syncthreads();
23
    }
24
    if(row < M && col < K){
25
        C[row * K + col] = S ;
26
    }
27
}
28

29
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
30
    dim3 threadsPerBlock(16, 16);
31
    dim3 blocksPerGrid((K + threadsPerBlock.x - 1) / threadsPerBlock.x,
32
                       (M + threadsPerBlock.y - 1) / threadsPerBlock.y);
33

34
    matrix_multiplication_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, M, N, K);
35
}

Test Code#

1
#include <cuda_runtime.h>
2

3
#include <array>
4
#include <cmath>
5
#include <iostream>
6

7
extern "C" void solve(const float* a, const float* b, float* c,
8
                      int m, int n, int k);
9

10
int main() {
11
    constexpr int m = 2;
12
    constexpr int n = 2;
13
    constexpr int k = 2;
14
    const std::array<float, 4> matrix_a = {1.0f, 2.0f, 3.0f, 4.0f};
15
    const std::array<float, 4> matrix_b = {5.0f, 6.0f, 7.0f, 8.0f};
16
    const std::array<float, 4> expected = {19.0f, 22.0f, 43.0f, 50.0f};
17
    std::array<float, 4> actual = {};
18
    constexpr size_t bytes = 4 * sizeof(float);
19

20
    float* device_a = nullptr;
21
    float* device_b = nullptr;
22
    float* device_c = nullptr;
23

24
    auto check_cuda = [](cudaError_t status, const char* operation) {
25
        if (status == cudaSuccess) {
26
            return true;
27
        }
28
        std::cerr << "Test failed. CUDA error in " << operation << ": "
29
                  << cudaGetErrorString(status) << '\n';
30
        return false;
31
    };
32

33
    bool success = check_cuda(cudaMalloc(&device_a, bytes), "cudaMalloc(device_a)") &&
34
                   check_cuda(cudaMalloc(&device_b, bytes), "cudaMalloc(device_b)") &&
35
                   check_cuda(cudaMalloc(&device_c, bytes), "cudaMalloc(device_c)") &&
36
                   check_cuda(cudaMemcpy(device_a, matrix_a.data(), bytes,
37
                                         cudaMemcpyHostToDevice), "copy matrix_a") &&
38
                   check_cuda(cudaMemcpy(device_b, matrix_b.data(), bytes,
39
                                         cudaMemcpyHostToDevice), "copy matrix_b");
40

41
    if (success) {
42
        solve(device_a, device_b, device_c, m, n, k);
43
        success = check_cuda(cudaGetLastError(), "solve") &&
44
                  check_cuda(cudaDeviceSynchronize(), "synchronize") &&
45
                  check_cuda(cudaMemcpy(actual.data(), device_c, bytes,
46
                                        cudaMemcpyDeviceToHost), "copy result");
47
    }
48

49
    cudaFree(device_a);
50
    cudaFree(device_b);
51
    cudaFree(device_c);
52

53
    if (!success) {
54
        return 1;
55
    }
56

57
    for (size_t index = 0; index < actual.size(); ++index) {
58
        if (std::fabs(actual[index] - expected[index]) > 1e-5f) {
59
            std::cerr << "Test failed at index " << index
60
                      << ". Expected: " << expected[index]
61
                      << ", Actual: " << actual[index] << '\n';
62
            return 1;
63
        }
64
    }
65

66
    std::cout << "Test passed.\n";
67
    return 0;
68
}

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	M=8192, N=6144, K=4096	10	250.884 ms	224.301 ms	302.463 ms	1643.454 GFLOPS
NVIDIA GeForce RTX 4090	PASS	M=8192, N=6144, K=4096	10	59.709 ms	59.302 ms	60.092 ms	6905.488 GFLOPS

Triton#

Approach#

Triton 版本采用直接的二维分块方式。一个 program instance 负责 C 中一个 32 x 32 的输出 tile，二维 grid 的两个 program ID 分别确定输出 tile 的行位置和列位置。

在 reduction 维度上，kernel 每次处理 32 个元素。tl.arange 生成当前 A 子块和 B 子块所需的索引，tl.load 负责加载数据，越界位置通过 mask 按零处理。随后，tl.dot 完成两个 32 x 32 子块的乘法并累加结果。所有 reduction tile 处理完毕后，tl.store 使用输出 mask 将有效结果写入 C。

与 CUDA 版本相比，这里不需要显式编写 thread index、shared memory 和同步操作；Triton 直接用 tile 级张量操作描述分块矩阵乘法，再由编译器完成底层线程映射。

Solution#

1
import torch
2
import triton
3
import triton.language as tl
4

5

6
@triton.jit
7
def matrix_multiplication_kernel(
8
    a, b, c,
9
    M, N, K,
10
    stride_am, stride_an,
11
    stride_bn, stride_bk,
12
    stride_cm, stride_ck,
13
    BLOCK_SIZE_M: tl.constexpr,
14
    BLOCK_SIZE_N: tl.constexpr,
15
    BLOCK_SIZE_K: tl.constexpr,
16
):
17
    m = tl.program_id(0)
18
    k = tl.program_id(1)
19
    MAX_N = tl.cdiv(N, BLOCK_SIZE_N)
20

21
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K), dtype=tl.float32)
22

23
    m_start = m * BLOCK_SIZE_M
24
    k_start = k * BLOCK_SIZE_K
25

26
    a_m_offset = m_start + tl.arange(0, BLOCK_SIZE_M)
27
    b_k_offset = k_start + tl.arange(0, BLOCK_SIZE_K)
28

29
    for n in tl.range(0, MAX_N):
30
        n_start = n * BLOCK_SIZE_N
31
        a_n_offset = n_start + tl.arange(0, BLOCK_SIZE_N)
32
        b_n_offset = n_start + tl.arange(0, BLOCK_SIZE_N)
33

34
        # A block: (BM, BN)
35
        a_mn_mask = (a_m_offset[:, None] < M) & (a_n_offset[None, :] < N)
36
        a_mn_ptr = a + a_m_offset[:, None] * stride_am + a_n_offset[None, :] * stride_an
37
        a_mn = tl.load(a_mn_ptr, mask=a_mn_mask, other=0.0)
38

39
        # B block: (BN, BK)
40
        b_nk_mask = (b_n_offset[:, None] < N) & (b_k_offset[None, :] < K)
41
        b_nk_ptr = b + b_n_offset[:, None] * stride_bn + b_k_offset[None, :] * stride_bk
42
        b_nk = tl.load(b_nk_ptr, mask=b_nk_mask, other=0.0)
43

44
        acc = tl.dot(a_mn, b_nk, acc, allow_tf32=False)
45

46
    c_mk_mask = (a_m_offset[:, None] < M) & (b_k_offset[None, :] < K)
47
    c_mk_ptr = c + a_m_offset[:, None] * stride_cm + b_k_offset[None, :] * stride_ck
48
    tl.store(c_mk_ptr, acc, mask=c_mk_mask)
49

50
# a, b, c are tensors on the GPU
51
def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, M: int, N: int, K: int):
52
    stride_am, stride_an = N, 1
53
    stride_bn, stride_bk = K, 1
54
    stride_cm, stride_ck = K, 1
55

56
    BLOCK_SIZE_M = 32
57
    BLOCK_SIZE_N = 32
58
    BLOCK_SIZE_K = 32
59

60
    grid = (triton.cdiv(M,BLOCK_SIZE_M), triton.cdiv(K,BLOCK_SIZE_K))
61
    matrix_multiplication_kernel[grid](
62
        a, b, c, M, N, K, stride_am, stride_an, stride_bn, stride_bk, stride_cm, stride_ck,
63
        BLOCK_SIZE_M = BLOCK_SIZE_M,
64
        BLOCK_SIZE_N = BLOCK_SIZE_N,
65
        BLOCK_SIZE_K = BLOCK_SIZE_K,
66
    )

Test Code#

1
import importlib.util
2
import sys
3
from pathlib import Path
4

5
import torch
6

7

8
def load_implementation():
9
    source_path = Path(__file__).resolve().parents[2] / "src" / "triton" / "02_Matrix_Multiplication.py"
10
    spec = importlib.util.spec_from_file_location("matrix_multiplication_triton_minimal", source_path)
11
    if spec is None or spec.loader is None:
12
        raise RuntimeError(f"Unable to load implementation: {source_path}")
13
    module = importlib.util.module_from_spec(spec)
14
    spec.loader.exec_module(module)
15
    return module
16

17

18
def main():
19
    matrix_a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="cuda")
20
    matrix_b = torch.tensor([[5.0, 6.0], [7.0, 8.0]], device="cuda")
21
    expected = torch.tensor([[19.0, 22.0], [43.0, 50.0]], device="cuda")
22
    actual = torch.empty_like(expected)
23

24
    load_implementation().solve(matrix_a, matrix_b, actual, 2, 2, 2)
25
    torch.cuda.synchronize()
26
    torch.testing.assert_close(actual, expected)
27
    print("Test passed.")
28

29

30
if __name__ == "__main__":
31
    try:
32
        main()
33
    except Exception as error:
34
        print(f"Test failed: {error}", file=sys.stderr)
35
        raise

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	M=8192, N=6144, K=4096	3	68.536 ms	63.085 ms	77.552 ms	6016.063 GFLOPS
NVIDIA GeForce RTX 4090	PASS	M=8192, N=6144, K=4096	10	17.057 ms	16.829 ms	17.418 ms	24172.676 GFLOPS

PyTorch#

Approach#

PyTorch 版本直接调用 torch.matmul(A, B, out=C)。A 和 B 分别是形状为 M x N 与 N x K 的二维 GPU tensor，计算结果直接写入调用方提供的输出 tensor C。

该实现不显式管理 grid、thread、tile 或边界条件，这些工作由 PyTorch 的 CUDA backend 完成。参数 M、N 和 K 为保持题目要求的 solve 接口而保留，实际矩阵尺寸由 tensor shape 决定。

Solution#

1
import torch
2

3
# A, B, C are tensors on the GPU
4
def solve(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, M: int, N: int, K: int):
5
    torch.matmul(A,B,out=C)

Test Code#

1
import importlib.util
2
import sys
3
from pathlib import Path
4

5
import torch
6

7

8
def load_implementation():
9
    source_path = Path(__file__).resolve().parents[2] / "src" / "pytorch" / "02_Matrix_Multiplication.py"
10
    spec = importlib.util.spec_from_file_location("matrix_multiplication_pytorch_minimal", source_path)
11
    if spec is None or spec.loader is None:
12
        raise RuntimeError(f"Unable to load implementation: {source_path}")
13
    module = importlib.util.module_from_spec(spec)
14
    spec.loader.exec_module(module)
15
    return module
16

17

18
def main():
19
    matrix_a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="cuda")
20
    matrix_b = torch.tensor([[5.0, 6.0], [7.0, 8.0]], device="cuda")
21
    expected = torch.tensor([[19.0, 22.0], [43.0, 50.0]], device="cuda")
22
    actual = torch.empty_like(expected)
23

24
    load_implementation().solve(matrix_a, matrix_b, actual, 2, 2, 2)
25
    torch.cuda.synchronize()
26
    torch.testing.assert_close(actual, expected)
27
    print("Test passed.")
28

29

30
if __name__ == "__main__":
31
    try:
32
        main()
33
    except Exception as error:
34
        print(f"Test failed: {error}", file=sys.stderr)
35
        raise

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	M=8192, N=6144, K=4096	10	29.305 ms	25.729 ms	31.689 ms	14069.752 GFLOPS
NVIDIA GeForce RTX 4090	PASS	M=8192, N=6144, K=4096	10	7.258 ms	7.151 ms	7.410 ms	56810.183 GFLOPS

References#

LeetGPU Challenge

关注 LeetGPU repository，了解更多内容。

音乐

音乐

Problem Description#

Implementation Requirements#

Example 1#

Example 2#

Constraints#

CUDA#

Approach#

Solution#

Test Code#

Test Result#

Triton#

Approach#

Solution#

Test Code#

Test Result#

PyTorch#

Approach#

Solution#

Test Code#

Test Result#

References#

文章分享

评论区

音乐

文章目录