趋近智
算子融合是一种图优化的基本方法,主要目的是减少核函数启动和中间内存传输的开销。简单的融合策略通常结合垂直相邻的逐元素操作(如 Add 后的 ReLU),有时也结合写入相同输出缓冲区的水平并行操作。然而,为了实现最佳性能,特别是在计算密度高但内存带宽通常有限的现代加速器上,需要更精细的融合技术。这些技术超越了简单的模式,依赖于复杂的分析和成本建模。
激进的融合致力于合并操作,即使其模式不是简单的线性链或纯粹的逐元素操作。其核心思想是最大化每次核函数启动完成的工作量,并最小化主内存和计算单元之间的数据移动。
显著的性能提升来自于将计算密集型操作(如卷积或矩阵乘法)与随后的逐元素操作进行融合。考虑一个常见模式:Conv -> BiasAdd -> ReLU。
将计算密集型操作(Conv2D)与随后的逐元素操作(BiasAdd, ReLU)融合到一个核函数中。中间张量
feature_map和biased_map从全局内存中消除。
在未融合的情况下,Conv2D 的结果(feature_map)被写入内存,然后由 BiasAdd 读取。其结果(biased_map)被写入内存,然后由 ReLU 读取。每一步都涉及大量的内存I/O和单独的核函数调用。将这些融合到一个核函数(融合: Conv2D_BiasAdd_ReLU)中,允许中间结果可能保留在寄存器或缓存中(如GPU共享内存或CPU L1/L2缓存),大幅减少内存带宽消耗和延迟。逐元素操作(BiasAdd, ReLU)的计算成本通常可忽略不计,与其融合时节省的内存访问时间相比。这种将计算密集型生产者与逐元素消费者结合的融合类型,非常有效。
激进的融合不限于线性链。它可以处理一个操作的输出被多个后续操作消费的情况。
将节点A与多个消费者(B,C)融合。中间结果计算一次后,在同一个核函数内立即被融合逻辑用于B和C,避免将其写入内存。
这里,操作A 产生一个中间结果,该结果被 操作B 和 操作C 使用。如果 B 和 C 是合适的候选(例如,逐元素操作),它们可以与 A 融合。融合后的核函数会计算 A 的结果,然后立即使用该结果计算 B 和 C,可能将中间数据保留在寄存器或片上缓存中。这避免了将中间张量写回内存,然后再读取两次(B一次,C一次)。这种模式有效地结合了垂直融合(A与B,A与C),并隐式处理了常见生产者情况。
决定是否激进地融合操作需要一个成本模型。盲目融合所有操作可能有害。例如,融合过多计算密集型操作可能创建一个需要比可用寄存器更多寄存器的核函数(导致寄存器溢出到局部或全局内存,这会很慢),或者如果组合的核函数代码变得过大和复杂,可能导致指令缓存未命中。
一个实用的成本模型根据几个因素估算潜在融合的性能影响:
编译器使用此模型,通常针对特定硬件目标进行参数化,以评估潜在的融合候选(操作对或组)。仅当融合配置的估算成本(例如,执行时间)预计低于未融合配置的成本时才执行融合。
不同的硬件目标需要不同的成本模型和启发式方法。内存带宽受限的GPU可能会更积极地倾向于融合,以减少片外内存访问,而拥有大容量、快速缓存的CPU可能会以不同方式优先考虑融合,也许更侧重于指令流水线利用率。
在某些情况下,将一个操作与其一个消费者融合可能是有益的,即使另一个消费者也需要原始输出。编译器可能选择为第二个消费者在单独的融合核函数中重新计算中间值,而不是强制将中间结果写入内存。如果中间操作的计算成本与主内存往返的成本相比很低,这种做法是可行的。这种策略对于计算开销小但产生大张量的操作特别适用,例如广播或某些逐元素转换。
激进的融合通常与数据布局转换(在memory-aware-layout-transformations中讨论)密切关联。有时,只有当所涉及张量的布局匹配(例如,都期望NCHW)或作为融合过程的一部分进行转换时,融合才可能或有益。例如,融合的卷积-ReLU核函数可能为了在特定目标上的效率而在内部以NHWC格式操作,即使周围的图使用NCHW,也会在核函数的序言和尾声中纳入布局更改。反之,融合操作可以简化图并消除中间张量,从而促成对之前被阻碍的后续操作进行有益的布局转换。
实施有效且激进的融合策略会带来几个工程挑战:
尽管存在这些复杂性,激进的算子融合是现代机器学习编译器和运行时中不可或缺的技术(如TensorFlow XLA, PyTorch JIT (通过NNC/Fuser), Apache TVM, ONNX Runtime, TensorRT)。其性能提升主要由内存带宽需求和核函数启动开销的大幅减少所驱动,对于在当代硬件上实现有竞争力的推理和训练速度通常非常重要。它有效地将计算图重构为更大、更适合硬件的执行单元。
这部分内容有帮助吗?
© 2026 ApX Machine Learning用心打造