动手实践：使用多面体工具优化循环

多面体建模是优化复杂循环嵌套的强大技术。实践指导将展示如何使用常见的多面体编译工具来优化标准矩阵乘法核心，这是机器学习 (machine learning)工作负载中的常见性能瓶颈。用户应具备可用的编译器环境，该环境具有多面体优化功能，例如启用Polly的Clang/LLVM，或者熟悉PPCG之类的独立工具或ISL库本身。

我们的目标是将一个简单、朴素的实现转换为一个高度优化的版本，借助分块等技术来改进数据局部性并实现并行化。

目标：朴素矩阵乘法

考虑标准的矩阵乘法例程 $C = A \times B$，其中A、B和C是大小为 $N \times N$ 的方阵。一个直接的C语言实现如下所示：

#define N 1024 // 示例尺寸

void matmul_naive(float C[N][N], float A[N][N], float B[N][N]) {
    for (int i = 0; i < N; ++i) {       // 循环 i
        for (int j = 0; j < N; ++j) {   // 循环 j
            C[i][j] = 0.0f;
            for (int k = 0; k < N; ++k) { // 循环 k
                C[i][j] += A[i][k] * B[k][j];
            }
        }
    }
}

这种实现存在数据局部性差的问题。对于大的N值，在最内层循环中访问B[k][j]会导致内存的跨步访问（假设行主序存储），从而引起频繁的缓存缺失。A的每个元素在j循环中被重复使用N次，B的每个元素在i循环中被重复使用N次，但朴素的循环顺序使得两者无法同时有效地利用缓存。

多面体表示回顾

如前所述，多面体工具以数学方式表示此循环嵌套。

• 迭代域： 语句所有动态实例的集合，由循环边界定义： $$S_{init}: \{ [i, j] \mid 0 \le i < N \land 0 \le j < N \} \quad \text{(初始化 } C[i][j]=0)$$ $$S_{update}: \{ [i, j, k] \mid 0 \le i < N \land 0 \le j < N \land 0 \le k < N \} \quad \text{(更新 } C[i][j]+=...)$$
• 内存访问函数： 将迭代向量 (vector)映射到内存位置：
  • S_init: $W_C: [i, j] \mapsto C[i][j]$
  • S_update: $R_A: [i, j, k] \mapsto A[i][k]$, $R_B: [i, j, k] \mapsto B[k][j]$, $R_C: [i, j, k] \mapsto C[i][j]$, $W_C: [i, j, k] \mapsto C[i][j]$
• 依赖关系： 这些工具分析语句实例之间的依赖关系（读后写、写后写等）。例如，$S_{init}[i, j]$ 和 $S_{update}[i, j, 0]$ 之间存在对C[i][j]的RAW依赖，并且在$S_{update}$的$k$循环迭代中存在对C[i][j]的WAR/WAW依赖。

使用Clang/Polly应用多面体优化

Polly是LLVM的多面体优化框架。我们可以使用Clang命令行标志调用它。假设matmul.c包含上述代码：

# 使用标准 O3 优化进行编译
clang -O3 matmul.c -o matmul_O3

# 使用 O3 和 Polly 优化进行编译
# -mllvm -polly: 启用 Polly
# -mllvm -polly-process-unprofitable: 即使启发式算法认为不划算也强制 Polly 执行（用于演示）
# -mllvm -polly-parallel: 为并行循环生成 OpenMP 代码
# -fopenmp: 链接 OpenMP 运行时库
clang -O3 -mllvm -polly -mllvm -polly-process-unprofitable -mllvm -polly-parallel -fopenmp matmul.c -o matmul_polly

Polly在内部执行以下步骤（简化视图）：

1. 检测： 识别适合多面体分析的静态控制部分（SCoPs）（我们的matmul_naive函数是一个主要适用对象）。
2. 表示： 使用ISL（整数集库）等库构建多面体表示（域、访问、依赖）。
3. 调度： 应用复杂的调度算法（例如基于Pluto的算法）以找到合法且优化的转换。这通常包括寻找多维仿射调度，以优化局部性并实现并行化。分块是常见的成果。
4. 代码生成： 根据新的调度生成优化的C代码（或LLVM IR），通常涉及大量重构的循环，如果启用并行化，可能还有OpenMP编译指示。

分析转换：分块

Polly（或类似工具）对矩阵乘法施加的典型转换是循环分块。并非遍历整个 $N \times N \times N$ 空间，计算被分解为更小的块（瓦片）。

循环可能被重构如下（使用TILE_SIZE）：

// 分块结构（简化版）
#define TILE_SIZE 32

void matmul_tiled(float C[N][N], float A[N][N], float B[N][N]) {
    #pragma omp parallel for // 外部瓦片的并行化
    for (int it = 0; it < N; it += TILE_SIZE) {
        for (int jt = 0; jt < N; jt += TILE_SIZE) {
            for (int kt = 0; kt < N; kt += TILE_SIZE) {
                // 处理一个瓦片：C[it:it+TS, jt:jt+TS] += A[it:it+TS, kt:kt+TS] * B[kt:kt+TS, jt:jt+TS]
                for (int i = it; i < it + TILE_SIZE; ++i) {
                    for (int j = jt; j < jt + TILE_SIZE; ++j) {
                        // 仅当 kt==0 时初始化 C 瓦片元素（工具处理）
                        // if (kt == 0) C[i][j] = 0.0f;
                        for (int k = kt; k < kt + TILE_SIZE; ++k) {
                            C[i][j] += A[i][k] * B[k][j];
                        }
                    }
                }
            }
        }
    }
}

注意： Polly生成的实际代码会更复杂，它会处理边界条件、瓦片内可能不同的循环顺序（例如，为了更好地复用A而采用i, k, j）、寄存器分块以及准确的初始化。#pragma omp parallel for表示最外层循环迭代（处理独立的C块）可以并行执行。

为什么分块有效？

• 时间局部性： 瓦片内的数据（例如，TILE_SIZE x TILE_SIZE大小的C块、A的相应行和B的相应列）在最内层循环（i, j, k）中被大量复用。如果TILE_SIZE选择得当（根据缓存大小），这些块在计算过程中很可能保留在缓存中。
• 空间局部性： 瓦片循环内的访问通常比原始代码更连续，尤其是在内层循环顺序优化后。
• 并行性： 处理不同的瓦片（由it、jt、kt定义）通常可以独立进行，特别是it和jt的外层循环，显现出适合多核CPU的粗粒度并行性。

观察生成的代码和性能

要检查实际的优化代码，您可以查看Clang/Polly生成的LLVM IR或汇编代码：

# 生成 LLVM IR（可读）
clang -O3 -mllvm -polly -S -emit-llvm matmul.c -o matmul_polly.ll

# 生成汇编代码
clang -O3 -mllvm -polly -S matmul.c -o matmul_polly.s

检查matmul_polly.ll或matmul_polly.s将显示重构的循环，通常与原始源代码有很大差异。您会寻找与瓦片和点循环对应的嵌套循环，以及可能与并行执行相关的OpenMP指令或函数调用。

衡量性能非常必要。使用time或更先进的性能分析工具（第9章介绍）来比较matmul_O3和matmul_polly。您应该会看到matmul_polly的明显加速，尤其是在缓存效应占主导地位的更大矩阵尺寸下。

大型矩阵乘法任务的性能比较。多面体优化通过提高缓存利用率和实现并行化，通常会带来显著的加速。实际结果在很大程度上取决于矩阵大小、硬件和工具的有效性。

后续研究和考量

• 瓦片尺寸选择： 最佳TILE_SIZE取决于缓存大小（L1、L2、L3）、寄存器可用性和内存带宽。多面体工具使用启发式算法或成本模型，但手动调整可能会带来额外收益。
• 更复杂的核： 将这些工具应用于其他基于循环的ML核心，如卷积。复杂性会增加，但原理保持不变。分析依赖关系以及工具如何重构循环。
• 工具局限： 多面体工具最适用于SCoPs：具有仿射循环边界和数组访问的静态控制流。它们可能难以处理复杂的数据依赖条件或间接内存访问。有时需要重构代码以显现可优化区域。
• 收益性： 并非所有多面体转换都有益。复杂循环结构或不完美瓦片尺寸的开销有时会超过其收益，特别是对于小问题规模。工具使用启发式算法，但理解权衡取舍很重要。

本实践练习表明了多面体编译技术对张量操作的核心性能具有显著影响。通过自动重构循环以实现局部性和并行化，这些工具提供了一个强大的机制，用于优化许多ML模型核心的计算密集型核，衔接了高级数学描述与高效硬件执行之间的距离。