文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载本篇技术指南聚焦 AISystem 项目中编译后端模块的算子循环优化Loop Optimization。在 AI 芯片实际执行计算时算子本质上是大量嵌套 for 循环对存储数据的重复读写与重复指令SIMT/SIMD循环优化的核心目标是提升数据局部性与计算并行性从而提升整体算子性能。读完本篇你将掌握循环分块、循环展开、循环重排、循环融合、循环拆分五类变换的原理、Cache miss 推导方法与应用边界并可对照仓库中的示例代码与 TVM 调度实践深入验证。为什么算子需要循环优化在具体硬件执行计算的时候实际会大量地使用for等循环指令不断地去读取不同的数据执行重复的指令SIMT/SIMD。因此循环优化主要是为了提升数据的局部性或者计算的并行性从而提升整体算子性能——当然这二者都需要 AI 芯片硬件的支持。从硬件视角看AI 算子如卷积、矩阵乘天然具有高度规则化的多层嵌套循环结构这为优化提供了丰富的技术手段。以逐元素操作为例如 ReLU、加法Add和乘法Mul等可以通过在所有循环轴上进行迭代来执行计算即使是较为复杂的卷积Conv也可以通过七层嵌套循环实现。然而如果仅仅采用直观的原生计算方法往往会导致效率低下。在 算子手工优化 一节中我们已经借助 Roofline 模型分析了算子的计算瓶颈Compute-Bound与访存瓶颈Memory-Bound而循环优化正是针对访存与并行性两大瓶颈的核心调度手段。循环优化的两大挑战数据局部性Cache 容量有限命中率决定访存开销数据的局部性与计算机存储层次有关。计算机拥有速度由慢到快、容量由大到小的多层次存储器以最优的控制调度算法和合理的成本构成存储系统寄存器位于最上层一般在 CPU 芯片内直接参与运算速度最快但价格最高、容量最小主存存放运行中的程序与数据但其速度与 CPU 差距较大Cache高速缓冲存储器插入在主存与 CPU 之间比主存速度更快、容量更小用于匹配二者的速度差异。主存与 Cache 之间的数据迁移由硬件自动完成对程序员来说是透明的、无法直接编程控制的。在现代多核 CPU 架构中Cache 分为 L1、L2、L3 三级对计算机运算速度影响最大的是L3 Cache因为该级 Cache 的不命中将导致片外访存。L3 Cache 的大小一般在几 MB 至几十 MB 之间容量较小因此无法把所有的数据都放到 Cache 里——Cache 里应该放 CPU 最有可能会对它进行处理的数据。这里有两个著名的程序运行时局部性原理诠释了什么样的数据很可能会被 CPU 处理时间局部性CPU 处理了一个数据以后它很有可能会对它进行第二次处理。空间局部性CPU 处理了内存中某一块的数据它很可能会还对它附近的数据块进行读写操作。这两个局部性原理十分匹配循环的特点。以两层嵌套循环为例for i in range(n): for j in range(m): A[i] B[j] # 这个计算可能并没有什么意义只是为了举例时间局部性A[i]在内层循环J 循环完整结束一次后被重用每一个 A 中的元素会被重用 m 次B[j]在每一次内层循环都改变在 i 循环时重用每一个 B 中的元素会被重用 n 次。空间局部性A 中的元素在 m 次后切换到下一个元素B 中的元素每次都会切换到下一个。这两个数组具有不同的模式——A 的时间局部性更明显B 的空间局部性更明显。为什么要分析数据的局部性因为 Cache 容量小一次只能容纳一部分数据。对 Cache 而言如果 CPU 要读写的数据在 Cache 中则称为命中可以直接在 Cache 这一级别返回数据如果不在 Cache 中则未命中需要去下一级存储取数据增加了 IO 的开销。计算并行性多核线程并行与 SIMD 向量化现代 CPU 通常是多核结构可以进行线程级并行。在单核 CPU 上运行多任务时通常是模拟进行的并行通过进程调度算法如时间片轮转来同时进行多任务实际上 CPU 在一个时刻只进行了一个任务。在多核 CPU 上每一个核都可以进行计算。搭配超线程技术Intel CPU时一个核可以执行两个线程。通过将多个计算线程分配到多个核可以同时执行多线程计算实现并行加速这是 CPU 上最有效的优化方式。在 Windows 中可以通过任务管理器查看内核与逻辑处理器数量——逻辑处理器就是用超线程技术模拟的能让处理器在同一时间段内同时处理多个线程从而实现更好的多任务处理能力。向量化则是一种数据级并行优化。向量化即批量操作在计算机中常见执行模型是单指令多数据SIMDSingle Instruction Multiple Data——通过对批量数据同时进行相同计算以提高效率。向量体系结构获取在存储器中散布的数据集将多个数据元素放在大型的顺序寄存器堆即向量寄存器中对整个寄存器进行操作从而同时计算了多个数据元素。向量本身可以容纳不同大小的数据因此如果一个向量寄存器可以容纳 64 个 64 bit 元素那么也可以容纳 128 个 32 bit 元素或者 512 个 8 bit 元素。凭借这种硬件上的多样性向量化特别适合用于多媒体应用和科学计算。传统的执行方式为单指令单数据SISDSingle Instruction Single Data硬件不支持并行计算。现代 CPU 几乎都支持 SIMD 指令集如 Intel 的 SSEStreaming SIMD Extensions和 AVXAdvanced Vector Extensions系列指令集。五类循环优化方案详解循环的优化方案针对不同的数据局部性和计算并行性主要有循环分块、循环展开、循环重排、循环融合、循环拆分等。仓库在 03Compiler/04Backend/code 目录下提供了 tiling、fusion、reorder、spliting 四份可直接对照的示例文件下面逐一展开。循环分块Loop Blocking / Tiling循环分块是利用 Cache 的数据局部性进行优化的一种方法。现代 CPU 通常具有多级 CacheCache 是除 CPU 寄存器外最接近 CPU 的存储层次相比主存速度更快但容量更小。Cache 中复制有 CPU 频繁使用的数据以进行快速访问。由于 Cache 容量有限数据会在 Cache 中进行换入换出当访问的数据在 Cache 中没有时产生Cache miss会向低一级存储层次发出访问请求然后该数据存储进 Cache访问时间大大提高当访问数据就在 Cache 中时会直接使用该数据以进行复用。循环分块主要针对大型数据集进行优化——大数据集无法一次全部存入 Cache。当遍历该数据集时循环按照顺序进行访问会替换掉之前加载进 Cache 的数据导致后面的指令对之前的数据无法复用要重新加载数据产生大量 Cache miss数据复用性很差程序执行时间变长大量时间花费在载入数据上。循环分块将大数据集分成多个小块以充分进行数据复用。数据块的内存访问是一个具有高内存局部性的小邻域该数据块可以一次加载进 Cache执行完所有或者尽可能多的计算任务后才被替换出。在实现中将一层内层循环分成outer loop * inner loop然后把 outer loop 移到更外层去从而确保 inner loop 一定能满足 Cache。原始与分块后的数据存储访问模式如下以如下代码为例分析其 Cache 利用率for i in range(n): for j in range(m): A[i] B[j]假设 m 和 n 是很大的数大数组那么对于每一轮 I 循环等到B[m-1]被访问时B[1]、B[2]等已经被清出缓存了。假设每个 Cache line 可以容纳 b 个数组元素以全相联的方式管理则 A 的 Cache miss 是n/bB 的 Cache miss 是n*m/b。如果把 j 层循环进行 tile分为j_o和j_i两层循环并把j_o提到最外面循环就变成了for j_o in range(0, m, T): for i in range(n): for j_i in range(j_o, min(j_o T, m)): A[i] B[j_i]当 T 个 B 中的元素可以放进 Cache 时在 i 层循环即使 i 切换了但 B 的数据此时还在 Cache 中只有j_o循环变换时 B 中的数据才会发生 Cache miss。对于 B 来说 Cache miss 为m/T * T/b m/b但对于 A 来说Cache miss 变为了m/T * n/b nm / Tb反而增加了。那么再对 i 层循环进行 tile最终变为for i_o in range(0, n, W): for j_o in range(0, m, T): for i_i in range(i_o, min(i_o W, n)): for j_i in range(j_o, min(j_o T, m)): A[i_i] B[j_i]假设 W 个 A 中的元素也能一次性放入 CacheB 的 Cache miss 没变A 的 Cache miss 变为n/W * W/b n/b。仓库中的示例文件 tiling 正是这一变换的 C 风格伪代码实现for j 0, n for i 1, m A(i) B(j) endfor endfor # After Tiling for j_o 0, m, T: for i 0, n: for j_i j_o, j_o T: A[i] B[j_i] endfor endfor endfor从公式上推导似乎只要 T 和 W 满足能放进 Cache 的数量那么 Cache miss 与 T、W 就无关不同 T、W 的性能应该是一样的——但实际运行并不是这样。tile 大小的选择受到 Cache line 大小、Cache 关联度、数据替换策略、硬件存储架构等多个因素的共同影响分块过大时数据尚未充分利用 Cache 的重用就被替换出去从而导致 miss分块过小又会造成较大的成本开销从而掩盖带来的性能优势。tile 大小的选择目前还没有一个确定的算法目前流行的方法有基于分析的方法早期有研究者对嵌套循环和硬件存储特征进行静态分析为编译器选择合适的分块大小。这类方法主要用来解决容量失效、自干扰失效、交叉干扰失效导致的 Cache 不命中和局部性优化问题。其缺陷在于1理论分析不能完全反应实际存储的复杂过程影响程序分块性能的因素很多导致实际最优分块的性能和分析方法选择的分块性能差距较大2分析建模过程复杂、成本较高对硬件和程序布局依赖严重不具有通用性。基于经验搜索的方法将循环嵌套看作一个黑盒根据经验选择一系列不同分块大小的组合在目标机器上对这些分块组合进行自动调优并从较好的分块大小组合中选取性能最优的分块大小。其缺陷在于对多层循环进行分块时要遍历庞大的搜索空间导致时间成本过高。值得说明的是分块思想在现代 AI 编译器中已通过调度 API 高度工程化。在 TVM 开发算子 的示例中cfg.define_split(tile_b, b, num_outputs2)、cfg.define_split(tile_ci, c_i, num_outputs2)等调度原语正是把分块大小定义为可搜索的配置项交由自动调优器在目标机器上搜索最优组合——这正是上文基于经验搜索方法的 DSL 化落地。循环展开Loop Unrolling循环展开将一个循环中的多次迭代展开成多个单独的迭代以减少程序执行的开销提高代码的运行效率。在计算机执行程序的流水线中每次跳转到循环体内部都需要进行额外的指令处理和跳转操作这会增加程序的开销。而循环展开可以通过减少跳转次数、减少指令处理次数等方式降低指令分支预测的开销来优化程序性能提高程序执行的速度。通常循环展开包含以下几个步骤复制循环体 n 次使展开后循环体包括原来循环体 n 个拷贝这里的 n 一般被称为展开因子调整数组索引变量的增量以及内存读写操作的地址删除除了最后一个循环体外的所有循环体中的循环计数变量的增加操作并且修改最后一个循环体中循环计数变量的增量为原来的 n 倍删除除了最后一个循环体外的所有循环中的循环条件判断语句。例如原始循环for i in range(m): a[i] b[i]通过循环展开展开因子 4可以将其转换为以下形式for i in range(0, m-3, 4): a[i] b[i] a[i1] b[i1] a[i2] b[i2] a[i3] b[i3] for i in range(m-3, m): a[i] b[i]在展开后的循环中原本执行了 n 次循环迭代变成了执行n/4次循环展开。从优化收益上分析循环展开不仅可以减少循环开销如循环变量测试及分支语句等还提高了指令之间的并发度并且因为减少了分支语句从而减少流水线停顿提升了流水线效率。另一个角度是循环展开后可能会为其他优化如指令级并行、向量化提供更多机会。但循环展开也可能带来负面效果如果展开后循环体超过指令缓存容量会引起缓存失效造成程序性能下降循环展开会增加寄存器压力可能导致生成更多的寄存器溢出处理操作从而降低优化效果。循环展开最关键的是确定展开因子目前主要有三种方法启发式方法对循环体代码进行分析然后使用静态模型计算展开因子。分析时需要考虑循环本身减少的循环开销、循环展开与其他优化的交互关系等建立模型要充分考虑指令级并行度、流水线效率、数据局部性、指令缓存与寄存器的限制等。机器学习方法根据循环的特征将循环分类通过大量样本学习使用分类器建立循环类型和展开因子之间的映射在实际优化循环时根据循环类型确定最优展开因子。迭代编译使用不同展开因子编译生成多个版本的程序并实际运行选取运行时间最短的作为最优展开因子。比较三个方法启发式方法开销最小展开因子的选择依赖于静态模型的准确性机器学习开销次之展开因子的选择不仅依赖于提取的循环特征还需要大量样本进行训练迭代编译开销最大但在不考虑开销的情况下肯定可以找到最优展开因子。循环重排Loop Reorder循环重排序是矩阵乘法常见的优化方式指的是对程序中的循环结构重新排列顺序以优化数据访问模式特别是在 CNN 中卷积层的应用。通过改变循环的嵌套顺序或者循环内部的迭代顺序可以改善数据的局部性减少缓存失效。在矩阵乘法计算中B 是逐列访问的在行优先的存储模式下访问模式很不友好。切换内层的循环顺序可以使得所有元素按顺序读取和写入——一次计算输出的一行得到的是中间结果全部累加即可得到结果矩阵的一行最终结果这种方式利用的是内存的空间局部性。仓库中的示例文件 reorder 展示了最直观的重排变换——交换两层循环的顺序for i 1, n for j 1, m A(i,j) B(i, j) * C(i, j) endfor endfor # After Reorder for j 1, m for i 1, n A(i, j) B(i, j) * C(i, j) endfor endfor需要注意的是循环重排必须保证变换前后循环语义迭代空间等价。在 TVM 开发算子 的调度示例中s[output].reorder(x_bo, x_co, x_bi, x_ci)正是将分块后的输出通道循环重排到更外层配合compute_at将计算阶段锚定到合适的存储层级实现数据局部性与并行性的协同优化。循环融合Loop Fusion循环融合用于将多个循环合并为一个更大的循环将相邻或紧密间隔的循环融合在一起。通过合并多个循环可以减少程序中的循环次数从而减少循环开销合并循环可以减少内存访问次数提高数据局部性减少缓存未命中的可能性从而提高程序执行效率。以下是一个简单的循环融合示例。两个独立的循环# 独立的循环 for i in range(len(a)): a[i] b[i] x for i in range(len(b)): d[i] a[i] y在第一个循环中a 的值被依次写入在第二个循环中又被马上读取。当数组非常大时在第二个循环要读取a[0]时a[0]早已因为 Cache 容量的限制而被清除需要从下一级存储中读取。通过循环融合可以将这两个循环合并为一个循环# 循环融合 for i in range(len(a)): a[i] b[i] x d[i] a[i] y这样在第二个对 a 的读取语句执行的时候a 的元素还在 Cache 中。除了这种数据局部性的收益循环融合还减少了对分支跳转指令的生成。仓库中的示例文件 fusion 展示了更完整的形式——将两个循环体合并后边界上的首尾迭代单独处理for i 0, n A(i) a(i) b(i); c(i) 2 * a(i); endfor for i 1, n - 1 D(i) c(i) a(i); endfor # After fusion A(0) a(0) b(0); c(0) 2 * a(0); A(n - 1) a(n - 1) b(n - 1); c(n - 1) 2 * a(n - 1); for i 1, n - 1 A(i) a(i) b(i) c(i) 2 * a(i) D(i) c(i) a(i) endfor循环融合并不总是具有正向收益的有时反而会降低性能甚至导致错误的结果。当前后两个循环存在数据依赖关系时将它们融合可能会导致错误的结果# 第一个循环 for i in range(N): A[i] B[i] C1 # 第二个循环 for i in range(N): D[i] A[i1] C2 # 循环融合 for i in range(N): A[i] B[i] C1 D[i] A[i1] C2融合后第二个循环本来要读取的是 A 改变之后的值但现在读取的是改变之前的值导致错误的结果。当然也可以通过修改源代码的方式进行对齐peeling把公共部分融合A[0] B[0] C1 for i in range(2, N-1): A[i] B[i] C1 D[i-1] A[i] C2 D[N-1] A[N] C2循环拆分Loop Splitting拆分主要是将循环分成多个循环可以在有条件的循环中使用分为无条件循环和含条件循环。以仓库示例 spliting 的代码为例for i in range(n): A[i] a[i] b[i] c[i]2 * a[i] if(temp[i] data): d[i] a[i] # 循环拆分 for i in range(n): A[i] a[i] b[i] c[i]2 * a[i] for i in range(n): if(temp[i] data): d[i] a[i]通过拆分将包含控制流的代码独立为一个循环一部分代码只有计算可以在加速器上计算而加速器不支持的控制流部分就可以回退到 CPU 计算。循环拆分一般可以创造出更多的优化机会例如和循环融合结合# 第一个循环 for i in range(N): A[i] B[i] C1 E[i] K[i] * 2 # 第二个循环 for i in range(N): D[i] A[i1] E[i]这两个循环无法直接合并因为 A 数组在两个循环之间存在依赖关系。但是可以通过对第一个循环进行拆分把 E 数组的部分拆分出来再融合进第二个循环中# 拆分 for i in range(N): A[i] B[i] C1 # 融合 for i in range(N): E[i] K[i] * 2 D[i] A[i1] E[i]这样做可以提升数组 E 的局部性减少 Cache miss——这正体现了拆分融合组合拳的价值先拆分消除依赖阻碍再融合重建局部性。小结与思考因为 Cache 容量小的特点一次只能容纳一部分数据因此需要分析计算的时间局部性和空间局部性。循环的优化方案针对不同的数据局部性和计算并行性有循环分块、循环展开、循环重排、循环融合、循环拆分等方案。循环优化的目的是提升数据的局部性或计算的并行性以提高整体算子性能。其挑战包括数据局部性和计算并行性。在 AISystem 仓库中五类变换均有可运行的对照示例tiling、fusion、reorder、spliting且 TVM/Triton 等 DSL 已把分块、重排、融合等变换封装为可组合、可自动调优的调度原语开发者只需声明计算逻辑编译器即可在目标硬件上搜索最优的循环变换组合。赞分享文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载相关推荐TencentDB Agent Memory内存检索算法BM25与向量搜索的融合应用TencentDB Agent Memory内存检索算法BM25与向量搜索的融合应用 TencentDB Agent Memory作为团队级AI Agent内人工智能大模型AI AgentAgent 记忆后端前端MCP 服务agno Workflow 循环执行Loop Execution实战指南端条件求值、迭代上限与并行子循环agno Workflow 循环执行Loop Execution实战指南端条件求值、迭代上限与并行子循环 本指南以 cookbook/04_workflo人工智能大模型AI AgentAgent 框架多智能体工具调用RAGAgent 工作流Agent 记忆终极解决方案如何彻底重置Cursor免费试用并突破使用限制终极解决方案如何彻底重置Cursor免费试用并突破使用限制 你是否正在寻找解决Cursor AI编程助手免费试用限制的方法当遇到Youve reache开发工具CLI上一篇ExplorerPatcher完整指南免费恢复Windows经典界面体验的终极工具下一篇从安装到精通Snippai用户操作完全指南创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考