CS267 · 第 16 课

稠密线性代数

从高斯消元的 5 步优化到 LU 分解的 BLAS3 形式，再到通信避免算法（CALU、TSLU、TSQR）和 GPU 实现，系统讲解稠密线性代数的并行优化方法论。

优化目标

稠密线性代数算法的优化遵循三个原则：

1. 最小化通信（数据移动）：减少 words 和 messages
2. 保持 flop 数不变（或尽可能少增）：不为了减少通信而增加算术复杂度
3. 保证正确性：数值稳定性不能妥协

从高斯消元到 LU 分解

求解 Ax = b 的经典方法是高斯消元（Gaussian Elimination, GE）。通过 5 步优化，GE 演变为高效的 BLAS3 LU 分解。

第 1 步：提取除法

// 原始高斯消元：对每行做除法
for (k = 0; k < n; k++)
    for (i = k+1; i < n; i++) {
        A[i][k] /= A[k][k];  // 除法
        for (j = k+1; j < n; j++)
            A[i][j] -= A[i][k] * A[k][j];
    }

第 2 步：跳过已知的零

// 利用 A[k+1:n, k] 被覆盖为乘子（multipliers）的事实
// L = tril(A,-1) + I, U = triu(A)

第 3 步：存储乘子

// 乘子 L_ik = A[i][k] 存储在 A[i][k] 位置
// 这样 L 和 U 共用存储空间

第 4 步：拆分循环

// 将更新分为两步：
for (k = 0; k < n; k++) {
    // Step 1: 计算乘子（列操作）
    for (i = k+1; i < n; i++)
        A[i][k] /= A[k][k];
    // Step 2: 更新 trailing submatrix
    for (i = k+1; i < n; i++)
        for (j = k+1; j < n; j++)
            A[i][j] -= A[i][k] * A[k][j];
}

第 5 步：BLAS 表述

// 转化为 BLAS 操作
for (k = 0; k < n; k++) {
    // SCAL: 列缩放
    scal(n-k-1, 1.0/A[k][k], &A[k+1][k]);
    // GER: rank-1 更新（BLAS2）
    ger(n-k-1, n-k-1, -1.0, &A[k+1][k], 1, &A[k][k+1], 1, &A[k+1][k+1]);
}

部分主元高斯消元（GEPP）

为处理小主元，引入部分主元（partial pivoting）：在第 k 列的下方找最大元素，交换行。

// GEPP
for (k = 0; k < n; k++) {
    // 找主元
    pivot = argmax_{i≥k} |A[i][k]|
    swap rows k and pivot;
    if (A[k][k] == 0) continue; // singular
    // 消元
    for (i = k+1; i < n; i++) {
        A[i][k] /= A[k][k];
        for (j = k+1; j < n; j++)
            A[i][j] -= A[i][k] * A[k][j];
    }
}

GEPP 产生 PA = LU 分解，其中 P 是置换矩阵。数值稳定性：增长因子通常 O(1)，最坏情况 O(2^(n-1))（极罕见）。

分块 GEPP：从 BLAS2 到 BLAS3

关键技巧：延迟更新（delayed updates）。将一次 rank-1 更新累积为 rank-b 更新。

// Blocked GEPP
for (k = 0; k < n; k += b) {
    // Panel factorization (BLAS2, 必须立即做，确保稳定性)
    for (kk = k; kk < min(k+b, n); kk++) {
        pivot = argmax_{i≥kk} |A[i][kk]|
        swap rows;
        for (i = kk+1; i < n; i++) A[i][kk] /= A[kk][kk];
    }
    // 累积乘子为三角矩阵 T (b×b)
    // Trailing matrix update (BLAS3!)
    A[k+b:n, k+b:n] -= L_21 * U_12;  // GEMM
}

[L11  0 ]   [U11 U12]
[L21 L22] = [ 0  A22']

         ↑panel    ↑trailing update (GEMM)

分块 GEPP 的效率

高斯消元的通信下界

LU 分解（无 pivoting 或全 pivoting）的通信下界与矩阵乘法相同：

#words = Ω(n³ / √M)
#messages = Ω(n³ / M^(3/2))

LAPACK GEPP 的次优性

标准 LAPACK 的 DGETRF（GEPP）在三种情况下未达到通信下界：

Cache-Oblivious 递归 LU

Toledo（1997）提出的递归 LU 算法天然适应任意 cache 层次：

// Recursive LU: factor A into L and U
// A is n x n, split into 2x2 block
void recursive_lu(double *A, int n) {
    if (n <= BASE_CASE) {
        dgetrf_(A, n);  // 调用 LAPACK 基础情况
        return;
    }
    int n1 = n/2, n2 = n - n1;
    // [A11 A12]   [L11  0 ] [U11 U12]
    // [A21 A22] = [L21 L22] [ 0  U22]

    recursive_lu(A11, n1);           // Factor A11 = L11*U11
    trsm("R", A11, A12);             // Solve L11*U12 = A12
    trsm("L", A11, A21);             // Solve L21*U11 = A21
    gemm(A21, A12, A22, -1, 1);      // A22 -= L21*U12
    recursive_lu(A22, n2);           // Recurse on Schur complement
}

高斯消元的并行数据布局

GE 的并行布局选择比矩阵乘法更复杂，因为 pivoting 引入了行间的顺序依赖。

1D 布局

• Row-cyclic：每行分配给一个处理器。简单但每次 pivot 需要广播整行。
• Column-cyclic：每列分配给一个处理器。适合 panel factorization。

2D 布局

• 2D block-cyclic：标准 ScaLAPACK 选择。
• Block-skewed：在 block-cyclic 中加入 skew，使 pivot 行在处理器间轮转，减少 idle time。

分布式 GE 与 ScaLAPACK 的次优性

ScaLAPACK 的 PDGEF 实现了分布式 GEPP，但在大规模系统上次优：

// ScaLAPACK PDGETRF 的通信模式
for (k = 0; k < n; k += NB) {
    // Panel factorization: O(log P) messages
    // Broadcast pivot/column to all processors in column
    // Trailing update: GEMM on local blocks
    // 但每次迭代都需要 O(n/NB) 次独立通信
}

#messages = O(n · log P) 或 O(n²/√P · log P)

通信避免 LU（CALU）

Communication-Avoiding LU（CALU）通过锦标赛选主元（tournament pivoting）将 panel 分解的通信批量化：

// CALU 的核心思想：将多步 panel 合并
// 1. 将 k 列打包为一个 "tall skinny" panel
// 2. 在 panel 内并行做 GE（用 tournament pivoting）
// 3. 一次性通信，更新 trailing matrix

// Tournament pivoting: 每轮选出 k 个候选 pivot
// 通过 k x k 的 "tournament" 比较选出最佳
for (round = 0; round < log P; round++) {
    // 每对处理器交换候选 pivot
    // 本地比较，选出胜者
    // 最终得到 k 个全局最优 pivot
}

TSLU：Tall Skinny LU

对于 tall-skinny 矩阵（m >> n，如 m = 10⁶, n = 100），标准 LU 效率极低。TSLU 是为此设计的专用算法：

// TSLU for m x n matrix (m >> n)
// 1. 将矩阵分为 P 个大小为 (m/P) x n 的块
// 2. 每处理器本地做 QR（或 LU）得到局部 R_i
// 3. 将所有 R_i 堆叠，做 tournament pivoting
// 4. 一次通信选出全局 pivot，更新 trailing

// 通信量：O(n² log P) words, O(log² P) messages
// 远低于标准 LU 的 O(mn²/√M)

TSQR：Tall Skinny QR

TSQR 是 TSLU 的 QR 版本，通过二叉树 reduction 实现高效并行：

// TSQR for m x n matrix (m >> n)
// 1. 每处理器本地 QR: A_i = Q_i * R_i
// 2. 将 R_i 配对，做 QR of stacked [R_i; R_j]
// 3. 递归直到只剩一个 R

// 树形 reduction（假设 P=4）：
// Level 0: 4 local QR → 4 R_i (n×n)
// Level 1: 2 QR of [R0;R1], [R2;R3] → 2 R
// Level 2: 1 QR of [R01; R23] → final R

性能结果

通信避免算法在实际系统中的加速效果：

最优算法总结

LU 在 GPU 上的实现

Volkov 在 SC08 的论文 "Benchmarking GPUs to tune dense linear algebra" 展示了 GPU 上的高性能 LU 实现：

// Volkov-style GPU GEMM kernel (simplified)
__global__ void volkov_gemm(float *A, float *B, float *C, int N) {
    // 每个 block 处理 TILE x TILE 的 C 块
    // 每个 thread 使用寄存器累加多个元素
    float reg[REGS]; // 大量寄存器
    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;

    #pragma unroll
    for (int k = 0; k < N; k += TILE) {
        // 协作加载 A 和 B 的 tile 到 shared memory
        __syncthreads();
        // 每个线程计算多个 C 元素（寄存器密集）
        #pragma unroll
        for (int i = 0; i < REGS; i++) {
            reg[i] += As[ty+i][k_local] * As[tx][k_local];
        }
        __syncthreads();
    }
    // 写回 C
}

多核 DAG 调度

现代多核系统使用 DAG（有向无环图）调度来动态分配任务：

Cholesky 的 DAG

// Cholesky DAG (n=4, b=1)
// 每个节点是一个任务，边表示依赖
POTRF(1,1)
  ├── TRSM(2,1) ── POTRF(2,2)
  │                 ├── TRSM(3,2) ── POTRF(3,3)
  │                 │                 └── TRSM(4,3) ── POTRF(4,4)
  │                 └── SYRK(3,2) ── TRSM(3,2) ...
  └── SYRK(2,1) ── GEMM(3,2,1) ...

课程项目方向

基于本课程内容的潜在课程项目：

• 在 GPU 上实现通信避免 Cholesky 分解
• 使用 auto-tuning 优化 GEMM micro-kernel
• 实现 CALU 并比较与 ScaLAPACK 的性能
• 在 RISC-V 向量扩展上优化 BLAS3
• 使用 task-based 运行时（Legion、PaRSEC）实现分布式 LU

本课要点

知识检查