通过分块和向量化优化内存访问模式,可确保单个处理器核心高效运行。但是,现代硬件主要通过大规模并发来达到高性能。单个CPU可能拥有数十个核心,而GPU则包含数千个CUDA核心或流处理器。为使用这些资源,机器学习编译器必须将循环嵌套中不同的迭代映射到并行执行单元。并行化过程将顺序循环转换为多个迭代同时执行的结构。这要求编译器识别哪些循环存在依赖关系,哪些是独立的。如果迭代 $i$ 的计算依赖于迭代 $i-1$ 的结果,那么该循环是顺序的,无法在不更改算法的情况下进行并行化。幸运的是,许多张量操作,例如逐元素相加或矩阵乘法的外层循环,都是“非常易于并行的”,这意味着它们之间没有迭代间的依赖关系。CPU 并行:分叉-合并模型在中央处理器(CPU)上,并行化通常以粗粒度进行管理。创建和管理线程的开销相对于单个加法或乘法的成本而言是很大的。因此,编译器通常将最外层循环作为CPU并行化的目标。采用的标准执行模型是分叉-合并模式。主线程执行串行代码,直到遇到并行区域(一个循环)。此时,它“分叉”出一组线程,将循环迭代分配给它们。一旦线程完成分配的任务块,它们就会“合并”回主线程。例如,一个标准的矩阵乘法循环嵌套:$$ \text{对于 } i \text{ 在 } 0..M \text{ 中}: \ \quad \text{对于 } j \text{ 在 } 0..N \text{ 中}: \ \quad \quad \text{对于 } k \text{ 在 } 0..K \text{ 中}: \ \quad \quad \quad C[i, j] += A[i, k] \times B[k, j] $$针对多核CPU的编译器很可能会并行化循环变量 $i$。如果 $M=1024$ 且硬件有4个核心,编译器会划分范围,以便线程0处理 $0..255$,线程1处理 $256..511$,依此类推。当并行循环内部的工作量足够大,足以抵消同步开销时,此策略效果最好。这里一个常见困难是 伪共享。当多个线程写入不同的变量,而这些变量碰巧位于同一缓存行时,就会出现这种情况。尽管逻辑上正确,但缓存一致性协议会迫使核心反复使彼此的缓存行失效,从而降低性能。编译器通过确保并行任务在足够间隔开的内存区域上操作,或使用线程局部累加缓冲区来缓解这个问题。GPU 并行:分层线程绑定图形处理器(GPU)使用基于单指令多线程(SIMT)的不同执行模型。与CPU不同,GPU被设计用于处理数千个轻量级线程,且上下文切换开销极小。但是,这些线程并非扁平的;它们被分层组织成块(或线程组)和网格。在机器学习编译器中,这种层次结构通过一种称为 线程绑定 的技术来解决。编译器不依赖运行时调度器来分配工作,而是显式重写循环变量以与硬件标识符对应。当一个循环被标记为GPU执行时,编译器会执行以下转换:拆分: 循环范围根据硬件限制(例如,每个块的最大线程数)拆分为外部和内部部分。绑定: 循环迭代器被替换为固有的硬件变量,例如CUDA术语中的blockIdx.x和threadIdx.x。例如,要将大小为 $N$ 的循环映射到GPU,编译器可能会转换索引计算。如果我们为每个块分配128个线程,线性索引 $i$ 将被重构为:$$i = \text{blockIdx.x} \times 128 + \text{threadIdx.x}$$这个绑定过程准确指定了每个线程处理哪个数据元素。下图说明了逻辑迭代空间如何分解并映射到GPU层次结构。digraph G { rankdir=TB; node [shape=box, style=filled, fontname="Helvetica", fontsize=10]; edge [fontname="Helvetica", fontsize=9, color="#868e96"]; subgraph cluster_logical { label = "逻辑迭代空间 (示例: 0 到 1023)"; style = dashed; color = "#adb5bd"; fontcolor = "#495057"; logical_loop [label="全局循环范围\n[0...1023]", fillcolor="#e9ecef", color="#dee2e6"]; } subgraph cluster_hardware { label = "硬件映射"; style = dashed; color = "#adb5bd"; fontcolor = "#495057"; block_0 [label="块 0\n(索引 0-127)", fillcolor="#a5d8ff", color="#74c0fc"]; block_1 [label="块 1\n(索引 128-255)", fillcolor="#a5d8ff", color="#74c0fc"]; block_n [label="块 ...", fillcolor="#a5d8ff", color="#74c0fc"]; thread_0 [label="线程 0", fillcolor="#ffc9c9", color="#ffa8a8"]; thread_127 [label="线程 127", fillcolor="#ffc9c9", color="#ffa8a8"]; } logical_loop -> block_0 [label=" 按块维度拆分 "]; logical_loop -> block_1; logical_loop -> block_n; block_0 -> thread_0 [label=" 映射到 threadIdx.x "]; block_0 -> thread_127; }映射过程将全局循环范围拆分为独立的块,这些块进一步细分为线程。这种显式绑定使得代码生成器能够生成高效的GPU核。循环扁平化与融合单个循环通常不包含足够的迭代来饱和现代GPU,现代GPU可能需要数万个活动线程才能有效隐藏内存延迟。如果一个张量操作涉及在一个小维度(例如,批量大小为32)上迭代,然后是一个特征维度(64),那么单独的循环都无法提供足够的并行度。为解决这个问题,编译器采用 循环扁平化 (或展平)技术。此技术将嵌套循环合并为一个单一的线性循环。编译器根据原始边界的乘积($32 \times 64 = 2048$)计算一个新的上限,并生成代码,使用取模和除法操作从单一线性ID中重构原始的二维索引。这会创建一个更大的统一迭代空间,可以轻松映射到硬件的扁平线程层次结构。同步与归约当线程必须通信时,并行化会变得复杂。机器学习中一个常见的情形是归约操作,例如在Softmax层中求和值或计算向量的范数。在这些情况下,线程不能简单地独立运行;它们必须合并结果。并行归约的一种标准策略是 树状归约 方法。线程通常在寄存器中执行局部归约,然后使用共享内存和屏障(例如CUDA中的__syncthreads())在块级别进行同步,最后将部分结果写入全局内存。如果编译器未能插入正确的同步屏障,就会发生竞态条件,导致非确定性的数值错误。反之,过度的同步会迫使线程空闲等待,从而降低并行化收益。先进的编译器会分析数据依赖图,仅在绝对需要时才插入屏障,通常在GPU上使用“warp shuffles”来在线程之间交换数据,而无需共享内存访问的延迟。并行化权衡盲目地并行化每个循环并不能保证性能提升。存在一些明显的权衡,成本模型必须加以评估:粒度: 在CPU上并行化内层循环通常会导致开销超出计算时间。内存带宽: 如果一个核是内存限制的(等待来自RAM的数据),增加更多线程不会加快执行速度,并可能增加缓存竞争。资源限制: GPU对每个块可用的寄存器数量有限制。使用过多的线程或过多的局部变量可能会导致寄存器溢出到全局内存,从而大幅降低速度。有效的自动调优会专门寻找“最佳点”配置,确定要启动多少线程以及如何映射特定的循环轴,以平衡算术强度与内存带宽限制。