CUDA矩阵乘法在16384×16384维度时运行时骤增问题排查

核心原因：TLB失效与未优化的内存访问模式 #

你的朴素矩阵乘法实现存在两个关键问题，在N=16384时触发了性能骤降：

1. TLB（地址转换旁路缓冲）大规模失效

   • RTX 4090的TLB缓存容量有限，每个TLB条目通常覆盖64KB的连续内存区域。当N=16384时，单个矩阵的每行刚好是16384 * 4字节 = 64KB，对应一个TLB条目。计算矩阵乘法时，每个线程需要遍历B矩阵的所有行（共16384行），这需要加载16384个TLB条目，远超TLB的缓存容量，导致大量TLB miss。每次TLB miss都会触发耗时的页表遍历，直接拉高整体运行时间。
   • 当N=8192时，每行大小是32KB，两个行共享一个TLB条目，总需求为4096个TLB条目，处于TLB缓存的有效覆盖范围内，命中率较高；而N>16384时，TLB会通过页表预取等机制适应更大的内存访问范围，命中率回升，因此耗时增长回到正常的5-7倍区间。

2. 未优化的全局内存访问模式

   • 你的代码中，计算resMatrix[col][row]时，对B矩阵的访问是按列遍历（B[i][row]），虽然同一block内的线程访问同一行的连续列时是合并访问，但跨循环迭代的访问是随机跳转到不同行的同一列，这种模式的全局内存缓存命中率极低。当N=16384时，这种低命中率结合TLB失效，放大了性能损耗。

验证方法 #

• 通过CUDA Profiler（如Nsight Systems）查看TLB miss率、全局内存命中率指标，可确认N=16384时的TLB miss率远高于其他维度。

优化方案：分块矩阵乘法（Tiling）+共享内存 #

通过将矩阵划分为小分块，利用共享内存缓存子矩阵，将不合并的全局内存访问转换为合并访问，同时大幅减少TLB miss和全局内存访问次数。优化后的代码如下：

import numpy as np
from numba import cuda

@cuda.jit
def optimized_matrix_multiply(A, B, C, n, block_size):
    # 分配共享内存，存储A和B的子块
    shared_A = cuda.shared.array(shape=(block_size, block_size), dtype=np.float32)
    shared_B = cuda.shared.array(shape=(block_size, block_size), dtype=np.float32)

    tx = cuda.threadIdx.x
    ty = cuda.threadIdx.y
    bx = cuda.blockIdx.x
    by = cuda.blockIdx.y

    # 计算当前线程负责的C矩阵元素坐标
    row = by * block_size + ty
    col = bx * block_size + tx

    if row >= n or col >= n:
        return

    sum_val = 0.0
    # 遍历所有子块
    for i in range(0, n, block_size):
        # 加载A和B的子块到共享内存
        if i + tx < n and row < n:
            shared_A[ty][tx] = A[row][i + tx]
        else:
            shared_A[ty][tx] = 0.0
        if i + ty < n and col < n:
            shared_B[ty][tx] = B[i + ty][col]
        else:
            shared_B[ty][tx] = 0.0
        # 等待所有线程加载完成
        cuda.syncthreads()

        # 计算子块内的乘法累加
        for k in range(block_size):
            sum_val += shared_A[ty][k] * shared_B[k][tx]
        # 等待所有线程计算完成，避免覆盖共享内存
        cuda.syncthreads()

    C[row][col] = sum_val

def fill_matrix(n):
    return np.random.rand(n, n).astype(np.float32), np.random.rand(n, n).astype(np.float32)

if __name__ == '__main__':
    N = 16384
    block_size = 32  # 通常设置为32或16，匹配GPU warp大小
    grid_size = (N + block_size - 1) // block_size

    matrix1, matrix2 = fill_matrix(N)
    result_matrix = np.zeros(shape=(N, N), dtype=np.float32)

    print(f"Measuring time for N = {N} and block_size = {block_size}")
    start_event = cuda.event()
    end_event = cuda.event()
    start_event.record()
    
    optimized_matrix_multiply[(grid_size, grid_size), (block_size, block_size)](matrix1, matrix2, result_matrix, N, block_size)

    end_event.record()
    end_event.synchronize()

    execution_time = cuda.event_elapsed_time(start_event, end_event)

    print(f"Execution time: {execution_time} ms")

优化效果说明 #

• 共享内存缓存：将全局内存中的子块加载到共享内存（延迟远低于全局内存），每个子块只需访问一次全局内存，大幅减少全局内存访问总量。
• 合并访问：加载共享内存时，对A和B的访问都是按行连续的，完全符合CUDA的合并访问要求，最大化带宽利用率。
• TLB压力缓解：分块后，每个线程的内存访问范围被限制在小分块内，TLB命中率显著提升，避免大规模TLB失效。

额外建议 #

• 优先使用float32类型（若代码中用float64，内存占用会翻倍，进一步加剧TLB和缓存压力）。
• 确保矩阵是连续存储的（NumPy默认是行主序，符合CUDA的内存访问偏好）。

内容的提问来源于stack exchange，提问作者Ceriana