""" accel.py —— 第 4 讲的内核：朴素 / 分块 / 线程粗化的矩阵乘法、三种矩阵转置、CSR 稀疏矩阵-向量乘法（在 SIMT 模拟器上运行） ===================================================================== 【这个文件要回答的问题】 原课第 4 讲讲了四种 GPU 加速技术：分块（tiling）、内存并行（合并访存、bank 冲突）、稀疏矩阵乘法、cuBLAS。 前三种都是"同一个计算，换一种让线程分工和取数的方式"。这里把 slides 里的每个内核都写成模拟器能执行的版本， 让 01–04 号脚本去数它们各自省下了多少次访存。 【本文件提供】 矩阵乘法 : matmul_naive（每线程一个输出，直接读全局内存）、matmul_tiled（slides 的 MatMulTiledKernel，可选是否做越界处理、 可以故意漏掉 __syncthreads 做反面教材）、matmul_tiled_coarse（线程粗化：一个块算横向相邻的几个输出块） 转置 : transpose_naive、transpose_smem（共享内存，有 bank 冲突）、transpose_smem_padded（+1 填充，无冲突） 稀疏 : dense_to_csr、spmv_csr_thread（slides 的 SpMVCSRKernel：一个线程一行）、spmv_csr_warp（一个线程束一行） 工具 : run_matmul（分配 → 拷贝 → 发射 → 拷回，返回结果和统计）、A100_PCIE（slides 用的 A100 规格） 【约定】矩阵都是行优先、float32、方阵（转置和矩阵乘法）；线程的 x 方向对应列、y 方向对应行（第 3 讲 §2.3）。 """ import os import sys from dataclasses import dataclass import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture02")) sys.path.insert(0, os.path.join(HERE, "..", "lecture03")) from simt import Device, ceil_div, launch # noqa: E402 @dataclass(frozen=True) class Spec: """roofline 估算用的规格：名字、FP32 峰值、TF32 张量核心峰值（dense）、显存带宽（字节/秒）。""" name: str fp32: float tf32: float bw: float # 原课第 4 讲 slides 用的 A100 80GB PCIe：FP32 19.5 TFLOPS、TF32 156 TFLOPS、显存带宽 1,935 GB/s A100_PCIE = Spec("A100 80GB PCIe", 19.5e12, 156e12, 1935e9) H100_SXM = Spec("H100 SXM", 67e12, 494.7e12, 3.35e12) # ===================================================================== # 1. 矩阵乘法 # ===================================================================== def matmul_naive(t, A, B, C, N): """ slides 第 6 页的 MatMulKernel：每个线程算 C 的一个元素，循环里每一步从全局内存读 A、B 各一个数。 for (k = 0; k < N; k++) Pvalue += a[row * N + k] * b[k * N + col]; """ row = t.blockIdx.y * t.blockDim.y + t.threadIdx.y col = t.blockIdx.x * t.blockDim.x + t.threadIdx.x if row >= N or col >= N: return acc = np.float32(0) for k in range(N): acc += A[row * N + k] * B[k * N + col] C[row * N + col] = acc def matmul_tiled(t, A, B, C, N, TILE, bounds=True, sync1=True, sync2=True): """ slides 第 36 页的 MatMulTiledKernel（共享内存分块）。块大小 = TILE × TILE，每个线程算 C 的一个元素。 【思路】把 C 的这一块所需要的 A 的"一条横带"和 B 的"一条竖带"切成 TILE × TILE 的小块，一相（phase）处理一对： 1. 块内每个线程从全局内存搬一个 A 的元素、一个 B 的元素到共享内存 As、Bs 2. __syncthreads()：等大家都搬完（否则有人会读到还没写的位置） 3. 每个线程用 As 的一行和 Bs 的一列做 TILE 次乘加（只读共享内存） 4. __syncthreads()：等大家都算完（否则有人会在别人还在读时覆盖 As、Bs） 每个从全局内存读进来的数在共享内存里被用了 TILE 次 → 全局内存读取减少到 1/TILE。 【参数】bounds：True 时用 ⌈N/TILE⌉ 相、越界位置补 0（正确处理 N 不是 TILE 倍数的情况）； False 时照搬 slides：相数 = N / TILE（整数除法）、不检查越界。 sync1 / sync2：是否保留两次 __syncthreads（故意去掉一个用来演示数据竞争） 【数值】每个线程仍然按 k = 0, 1, …, N−1 的顺序累加，所以结果和 matmul_naive 逐位相同。 """ As = t.shared("As", (TILE, TILE)) Bs = t.shared("Bs", (TILE, TILE)) tx, ty = t.threadIdx.x, t.threadIdx.y row = t.blockIdx.y * TILE + ty col = t.blockIdx.x * TILE + tx acc = np.float32(0) phases = ceil_div(N, TILE) if bounds else N // TILE for ph in range(phases): ka, kb = ph * TILE + tx, ph * TILE + ty if bounds: As[ty, tx] = A[row * N + ka] if (row < N and ka < N) else np.float32(0) Bs[ty, tx] = B[kb * N + col] if (kb < N and col < N) else np.float32(0) else: As[ty, tx] = A[row * N + ka] Bs[ty, tx] = B[kb * N + col] if sync1: yield t.syncthreads() for k in range(TILE): acc += As[ty, k] * Bs[k, tx] if sync2: yield t.syncthreads() if not bounds or (row < N and col < N): C[row * N + col] = acc def matmul_tiled_coarse(t, A, B, C, N, TILE, COARSE): """ 线程粗化（thread coarsening，PMPP 6.3 节）：一个块负责横向相邻的 COARSE 个输出块，每个线程算 COARSE 个输出。 【为什么】在 matmul_tiled 里，横向相邻的块要的是同一条 A 的横带，却各自从全局内存读一遍。 粗化后 A 的小块每相只读一次，被 COARSE 个输出块共用 → A 的全局读取再减少到 1/COARSE。 【代价】每个线程要 COARSE 个累加器（更多寄存器），块数减少到 1/COARSE（并行度降低）。 【边界】要求 N 是 TILE × COARSE 的倍数（简化）。 """ As = t.shared("As", (TILE, TILE)) Bs = t.shared("Bs", (TILE, TILE)) tx, ty = t.threadIdx.x, t.threadIdx.y row = t.blockIdx.y * TILE + ty col0 = t.blockIdx.x * TILE * COARSE + tx acc = [np.float32(0)] * COARSE for ph in range(N // TILE): As[ty, tx] = A[row * N + ph * TILE + tx] # A 的小块：每相只读一次 for c in range(COARSE): col = col0 + c * TILE Bs[ty, tx] = B[(ph * TILE + ty) * N + col] yield t.syncthreads() for k in range(TILE): acc[c] += As[ty, k] * Bs[k, tx] yield t.syncthreads() for c in range(COARSE): C[row * N + col0 + c * TILE] = acc[c] def run_matmul(kernel, A, B, block, grid, *extra, **launch_kw): """ 主机端：分配 → 拷入 → 发射 → 拷回 → 释放。【返回】(C 的 NumPy 数组, LaunchStats) """ N = A.shape[0] dev = Device() dA, dB = dev.to_device(A, "A"), dev.to_device(B, "B") dC = dev.malloc(N * N, np.float32, "C") s = launch(kernel, grid, block, dA, dB, dC, N, *extra, **launch_kw) C = dev.to_host(dC, (N, N)) return C, s # ===================================================================== # 2. 转置（slides 第 25–33 页） # ===================================================================== def transpose_naive(t, m, a, c): """ slides 第 26 页的 naive_cuda_transpose： c[INDX(myCol, myRow, m)] = a[INDX(myRow, myCol, m)]; 读 a 时一个线程束读同一行的相邻元素（合并）；写 c 时相邻线程写同一列、相隔 m 个元素（不合并）。 """ my_col = t.blockDim.x * t.blockIdx.x + t.threadIdx.x my_row = t.blockDim.y * t.blockIdx.y + t.threadIdx.y if my_row < m and my_col < m: c[my_col * m + my_row] = a[my_row * m + my_col] def transpose_smem(t, m, a, c, pad=0): """ slides 第 30 页的 smem_cuda_transpose：先按行（合并地）把一块读进共享内存，再按行（合并地）写出转置后的块。 smemArray[threadIdx.x][threadIdx.y] = a[INDX(tileRow + threadIdx.y, tileCol + threadIdx.x, m)]; __syncthreads(); c[INDX(tileCol + threadIdx.y, tileRow + threadIdx.x, m)] = smemArray[threadIdx.y][threadIdx.x]; pad = 1 时就是 slides 第 33 页的版本：共享数组声明成 [TX][TY + 1]，消除 bank 冲突。 【边界】要求 m 是块大小的倍数（和 slides 一样，不做越界检查）。 """ TX, TY = t.blockDim.x, t.blockDim.y smem = t.shared("smemArray", (TX, TY + pad)) tile_col = TX * t.blockIdx.x tile_row = TY * t.blockIdx.y tx, ty = t.threadIdx.x, t.threadIdx.y smem[tx, ty] = a[(tile_row + ty) * m + tile_col + tx] yield t.syncthreads() c[(tile_col + ty) * m + tile_row + tx] = smem[ty, tx] def transpose_smem_padded(t, m, a, c): """slides 第 33 页：smemArray[THREADS_PER_BLOCK_X][THREADS_PER_BLOCK_Y + 1]。""" yield from transpose_smem(t, m, a, c, pad=1) # ===================================================================== # 3. 稀疏矩阵（slides 第 41–43 页） # ===================================================================== def dense_to_csr(M): """ 稠密矩阵 → CSR（Compressed Sparse Row，压缩稀疏行）三个数组： data[k] : 第 k 个非零元的值（按行优先顺序排列） col_index[k] : 第 k 个非零元所在的列 row_ptr[r] : 第 r 行的非零元在 data 里从哪里开始；row_ptr[r+1] − row_ptr[r] = 第 r 行的非零元个数 （长度是行数 + 1，最后一个元素 = 非零元总数） 【返回】(data float32, col_index int32, row_ptr int32) """ rows, cols = np.nonzero(M) data = M[rows, cols].astype(np.float32) row_ptr = np.zeros(M.shape[0] + 1, np.int32) np.add.at(row_ptr, rows + 1, 1) return data, cols.astype(np.int32), np.cumsum(row_ptr).astype(np.int32) def spmv_csr_thread(t, data, col_index, row_ptr, x, y, num_rows): """ slides 第 43 页的 SpMVCSRKernel：一个线程负责一行，y[row] += Σ data[elem] · x[col_index[elem]]。 【问题】相邻线程处理相邻的行，同一时刻读的是各自那一行的第 k 个非零元——位置相隔一整行的非零元个数（不合并）； 各行长度不同时，线程束要等最长的那一行（分化）。 """ row = t.blockIdx.x * t.blockDim.x + t.threadIdx.x if row < num_rows: dot = np.float32(0) for elem in range(row_ptr[row], row_ptr[row + 1]): dot += x[col_index[elem]] * data[elem] y[row] = y[row] + dot def spmv_csr_warp(t, data, col_index, row_ptr, x, y, num_rows): """ 改进（PMPP 第 14 章的"向量 CSR"）：一个线程束负责一行。32 个线程读这一行连续的 32 个非零元（合并）， 各自累加部分和，最后在共享内存里树形归约成一个数，由 0 号车道写出。 【适用】每行非零元很多时；每行只有几个非零元时，线程束里大部分车道闲着。 """ lane = t.threadIdx.x % 32 warp_in_block = t.threadIdx.x // 32 row = t.blockIdx.x * (t.blockDim.x // 32) + warp_in_block part = t.shared("part", t.blockDim.x) dot = np.float32(0) if row < num_rows: for elem in range(row_ptr[row] + lane, row_ptr[row + 1], 32): dot += x[col_index[elem]] * data[elem] part[t.threadIdx.x] = dot yield t.syncthreads() s = 16 while s > 0: if lane < s: part[t.threadIdx.x] = part[t.threadIdx.x] + part[t.threadIdx.x + s] yield t.syncthreads() s //= 2 if lane == 0 and row < num_rows: y[row] = y[row] + part[t.threadIdx.x]