首页 / 资讯中心 / 文章详情

AI芯片软硬件协同设计:从算力到存储的效能优化

AI芯片软硬件协同设计:从算力到存储的效能优化 ★ FEATURED ARTICLE
1. 从能跑模型到跑得高效AI芯片软硬件协同设计的核心命题聊到AI芯片很多人第一反应是算力数字——多少TOPS、多少TFLOPS。但真正在一线做过部署的人都知道纸面算力和实际吞吐之间往往隔着一条巨大的鸿沟。这条鸿沟的宽度基本就由软硬件协同设计的水平决定。我见过太多团队芯片流片回来发现实际推理性能只有理论峰值的百分之十几问题不出在硬件本身而是软件栈没跟上、数据搬运成了瓶颈、算子映射一塌糊涂。这个系列写到第10篇我想把前面散落各处的点串成一条线专门讲清楚一件事AI芯片的软硬件设计到底是怎么互相牵制的以及一个合格的从业者在做架构定义、编译器适配、算子优化时脑子里应该装着哪些约束条件。不管你是做芯片前端设计的、写推理框架的还是负责把模型落到板子上的部署工程师这套协同思维都是绕不开的基本功。先给不太熟悉背景的读者一个整体图景。AI芯片的软硬件设计大致可以分成几个层次最底下是芯片的微架构包括计算阵列、片上存储层次、互联总线往上是指令集和编程模型决定了软件怎么描述计算再往上是编译器和运行时负责把高层模型翻译成芯片能执行的指令序列最上面是框架和算子库直接对接用户写的模型。每一层之间都有接口接口设计得好不好直接决定了整条链路的效率上限。我个人的经验是做AI芯片最怕的就是各扫门前雪。硬件团队闷头堆MAC阵列软件团队抱怨硬件不灵活编译器团队夹在中间两头受气。真正高效的团队一定是硬件架构师懂模型的计算特征编译器工程师懂硬件的存储带宽约束部署工程师懂编译器的调度策略。这篇就把这些交叉点一个个拆开讲。2. 计算阵列与数据复用硬件设计里最容易被低估的存储问题2.1 为什么算力堆上去性能却没涨很多初入行的同学会有一个直觉把MAC阵列从64x64扩到128x128算力翻四倍性能应该也差不多翻四倍。实际做下来往往不是这么回事。原因在于计算单元越多对数据供给的要求就越高。如果片上存储带宽跟不上计算阵列就会大量空转也就是所谓的饥饿。这里有个很朴素的计算。假设一个128x128的脉动阵列每个周期要喂进去128个激活值和128个权重那就是每周期256个数据。如果片上SRAM的读取带宽只有每周期128个字那阵列最多只能跑到一半利用率。这个账在架构定义阶段就必须算清楚而不是等流片回来再补救。我参与过的一个模拟项目里架构团队一开始把重点全放在算力上片上缓存只给了很小的容量。结果编译器团队在做算子映射时发现稍微大一点的卷积层权重根本放不下只能反复从外部存储搬运带宽瞬间成为瓶颈。后来重新调整了存储层次把权重缓存做大虽然牺牲了一点面积但端到端性能反而提升了两倍多。这个教训很典型AI芯片的性能瓶颈十有八九出在数据搬运上而不是计算本身。2.2 数据复用维度的取舍逻辑要解决数据供给问题核心思路是提高数据复用率。卷积计算里有几个天然的复用维度输入激活可以在多个输出通道间复用权重可以在多个空间位置间复用部分和可以在累加过程中复用。不同的复用策略对应不同的数据流设计常见的有权重固定、输出固定、行固定等几种。选择哪种数据流取决于你的目标模型特征。如果主要跑的是大卷积核、小通道数的网络权重固定的策略比较合适因为权重总量小可以常驻在寄存器或近端缓存里。如果跑的是小卷积核、大通道数的网络那激活值的复用就更关键输出固定的策略可能更优。这里没有银弹必须结合目标场景来定。我通常建议团队先做一轮模型画像分析统计目标模型集的卷积核尺寸分布、通道数分布、特征图大小分布然后据此选择数据流。这个分析工作看起来繁琐但能省下后面大量的返工。2.3 片上存储层次的划分经验片上存储一般分几级寄存器文件、近端缓存、全局共享缓存。寄存器文件最快但最贵通常只放当前计算最活跃的数据近端缓存放权重和部分激活全局缓存放跨计算单元共享的数据。划分的关键是容量和带宽的平衡。我的经验是寄存器文件的大小要能容纳一个计算单元在一个时间步内需要的所有操作数否则就会频繁访问近端缓存增加延迟。近端缓存的容量要能容纳一个完整算子的权重避免算子执行中途换出。全局缓存的带宽要能支撑所有计算单元同时取数的最坏情况。这些数字不是拍脑袋定的需要结合目标算子的计算强度和访存模式来算。一个实用的方法是写一个简单的性能模型把不同存储层次的容量和带宽作为参数跑一遍目标模型看瓶颈出现在哪里然后针对性调整。3. 指令集与编程模型软件怎么告诉硬件该干什么3.1 从SIMD到领域专用指令的演进早期AI加速器很多直接沿用SIMD指令集一条指令处理一组数据。这种方式通用性好但对AI计算的针对性不强。比如卷积里的乘累加操作用SIMD指令描述会引入大量额外的数据搬移和地址计算指令指令开销占比很高。后来大家开始设计领域专用指令把常见的计算模式固化下来。比如一条卷积指令直接描述一个卷积窗口的计算硬件内部自动完成数据取用和累加。这样指令密度大幅提升取指和译码的开销就降下来了。但专用指令也有代价就是灵活性下降。如果模型里出现了指令集没覆盖的算子就得退化成多条通用指令的组合性能可能断崖式下跌。所以指令集设计要在通用性和专用性之间找平衡点。我的建议是先覆盖高频算子低频算子用通用指令兜底同时预留一定的可扩展空间。3.2 编程模型对编译器友好度的考量编程模型是软件和硬件之间的契约。一个好的编程模型应该让编译器容易做优化同时让开发者容易理解。这里有个常见的误区为了追求灵活性把编程模型设计得过于底层结果编译器要处理大量细节优化难度陡增。我比较推崇的做法是分层编程模型。底层是硬件指令中间层是算子原语上层是图级别的描述。编译器在不同层次做不同的优化底层做指令调度和寄存器分配中间层做算子融合和内存规划上层做图切分和并行调度。这样每层的复杂度都可控。编程模型还要考虑的一个点是数据布局。不同的数据布局对访存效率影响巨大。比如NHWC和NCHW两种布局在不同硬件上的表现可能差好几倍。编程模型最好能抽象掉布局细节让编译器根据硬件特征自动选择而不是让开发者手动指定。3.3 指令流水与依赖处理的实操细节AI计算里指令之间的依赖关系很复杂尤其是存在数据依赖和累加依赖的时候。硬件如果处理不好就会频繁插入气泡降低流水线效率。一个实用的技巧是在指令集里显式标注依赖关系让硬件可以提前做调度。比如用记分牌或者预留站的方式把没有依赖的指令提前发射。这需要编译器和硬件配合编译器在生成指令时标注依赖信息硬件据此做动态调度。另一个技巧是软件流水。编译器把循环展开把不同迭代的指令交错排列让硬件可以并行执行。这个技术在卷积和矩阵乘里特别有效能把流水线填得很满。但软件流水对寄存器压力很大需要编译器做好寄存器分配否则会溢出到内存反而变慢。4. 编译器后端把模型图翻译成高效指令序列的关键战场4.1 算子融合的收益与边界算子融合是编译器优化里收益最直接的手段之一。把相邻的算子合并成一个可以减少中间结果的写回和读取节省大量带宽。比如卷积后面接一个激活函数完全可以融合成一个算子激活在卷积的累加结果上直接做不用把卷积结果写回内存再读出来。但融合不是无脑合。融合的前提是中间结果不需要被其他算子引用且融合后的算子能在硬件上高效执行。如果融合后算子太大超出了片上存储容量反而会导致频繁换入换出得不偿失。我的经验是融合策略要结合硬件资源来定。片上缓存能放下多大的中间结果就融合到多大。编译器在做融合决策时应该查询硬件的资源描述而不是用一个固定的融合规则。4.2 内存规划与生命周期分析内存规划是编译器后端另一个重头戏。模型里的张量有各自的生存周期编译器要分析这些周期尽量复用内存空间减少峰值内存占用。具体做法是构建一个张量生命周期图找出哪些张量的生命周期不重叠让它们共享同一块内存。这个分析在静态图模型上比较好做动态图就麻烦一些需要运行时支持。内存规划还要考虑对齐和bank冲突。片上存储通常分bank如果多个张量映射到同一个bank访问时就会冲突降低带宽。编译器在做内存分配时应该尽量把同时访问的张量分散到不同bank。4.3 针对特定硬件的调度策略调优不同硬件的调度策略差异很大。有的硬件适合大粒度的任务并行有的适合细粒度的指令并行。编译器需要根据硬件特征调整调度策略。比如对于多核架构编译器要把计算图切分成多个子图分配到不同核上同时处理好核间的数据依赖和同步。切分策略直接影响负载均衡和通信开销。切得太细通信开销大切得太粗负载不均衡。对于单核多发射架构编译器要重点做指令级并行把没有依赖的指令打包到一起发射。这需要编译器有良好的依赖分析和调度算法。调优的过程往往是迭代的。先跑一个基线版本用性能分析工具找出瓶颈然后针对性调整调度策略再跑再调。这个循环可能要重复很多轮直到性能收敛。5. 部署阶段的软硬件对齐那些流片后才暴露的问题5.1 实际带宽与理论带宽的差距来源芯片手册上写的带宽是理论峰值实际能跑到多少取决于访存模式。连续大块访问通常能接近峰值随机小块访问可能只有峰值的几分之一。部署阶段经常遇到的问题是模型里的访存模式不理想导致实际带宽远低于预期。比如某些算子的输入张量在内存里不是连续存放的访问时就要跳着读效率很低。解决办法是在编译器里做布局转换把张量重排成对硬件友好的格式。另一个常见问题是bank冲突。如果多个计算单元同时访问同一个bank就会串行化。部署时要检查关键算子的访存模式必要时调整数据布局或访问顺序。5.2 精度与性能的权衡实操AI计算里精度和性能往往是一对矛盾。用低精度比如INT8能大幅提升吞吐、降低功耗但可能损失模型精度。用高精度比如FP16精度好但性能差一些。实际部署时通常采用混合精度策略对精度敏感的层用高精度对精度不敏感的层用低精度。哪些层敏感需要做敏感度分析逐层量化看精度损失。量化本身也有很多坑。比如激活值的动态范围在不同输入下变化很大用固定的量化参数可能在某些输入上溢出或精度损失严重。这时候要用动态量化根据实际输入调整量化参数。但动态量化会增加运行时开销需要权衡。5.3 从profiling数据反推设计缺陷部署阶段最重要的工具是profiling。通过profiling能看到每个算子的耗时、访存带宽利用率、计算单元利用率等指标。这些数据能直接反映硬件设计的缺陷。比如某个算子计算单元利用率很低可能是数据供给不足说明片上缓存或带宽设计有问题。某个算子访存带宽利用率很高但计算利用率低说明是访存瓶颈需要优化数据复用。我习惯在部署阶段做一轮全面的profiling把每个算子的各项指标列出来找出最差的几个然后分析根因。这些根因往往能追溯到架构设计阶段的某个决策。把这些反馈给架构团队下一版芯片就能针对性改进。这就是软硬件协同设计的闭环。6. 协同设计的工作流让硬件和软件团队真正对话6.1 架构探索阶段的软硬件联合建模架构探索阶段最忌讳硬件团队闭门造车。正确的做法是软硬件团队一起建一个性能模型把硬件参数和软件行为都放进去快速评估不同架构方案的性能。这个模型不需要很精确但要有足够的代表性。通常用解析模型或者简化的仿真模型就够了。关键是能快速迭代一天能跑几十个方案而不是等几个月流片回来才知道好坏。模型里要包含的关键参数计算阵列规模、存储层次容量和带宽、指令发射宽度、访存延迟等。软件侧要包含目标模型的算子分布、数据复用模式等。两边一起调找到性能、面积、功耗的平衡点。6.2 编译器与硬件的接口定义编译器和硬件之间的接口是协同设计的关键。这个接口定义了编译器能看到哪些硬件资源、能发出哪些指令、能控制哪些参数。接口设计的原则是暴露必要的信息隐藏不必要的细节。编译器需要知道存储层次的大小和带宽才能做内存规划需要知道指令的延迟和吞吐才能做调度。但编译器不需要知道具体的电路实现。接口一旦定下来就不要轻易改。因为编译器团队要基于这个接口做大量开发改接口意味着大量返工。所以接口定义阶段要充分讨论把各种场景都考虑到。6.3 迭代优化的节奏把控软硬件协同设计是一个迭代过程但迭代节奏要把握好。硬件迭代周期长软件迭代周期短两者要协调。通常的做法是硬件先出一个功能模型或者FPGA原型软件团队基于这个原型做开发和优化。软件优化过程中发现的问题反馈给硬件团队硬件在下一版中改进。同时软件团队也要为硬件的未来版本预留适配空间。节奏把控的关键是找到足够好的点。不要追求一次完美而是快速迭代每轮解决最突出的问题。芯片设计周期长等不起无限次的打磨。7. 一些踩过坑之后才明白的道理做AI芯片软硬件设计这些年有几个道理是踩了坑才真正理解的。第一个是存储比计算更难。计算阵列的设计相对成熟无非是规模、数据流、精度的组合。但存储层次的设计要考虑容量、带宽、延迟、功耗、面积的多重约束还要适配各种算子的访存模式复杂度高得多。很多芯片性能不达标根子都在存储上。第二个是编译器的价值被严重低估。同样的硬件好的编译器和差的编译器性能可能差好几倍。但编译器开发周期长、见效慢很多团队不愿意投入。我的建议是从项目一开始就把编译器团队建起来和硬件团队同步推进而不是等硬件好了再补编译器。第三个是部署反馈必须回流到设计。很多团队把芯片交付给部署团队就完事了部署中发现的性能问题没有系统性地反馈给架构团队。这导致下一代芯片还是犯同样的错误。建立部署到设计的反馈闭环是持续改进的关键。第四个是不要迷信通用性。AI芯片领域专用性很强试图做一个什么模型都能高效跑的通用芯片往往什么都做不好。明确目标场景针对性地优化反而能做出有竞争力的产品。最后说一个实操层面的小技巧。在做架构评估时我习惯准备一组代表性模型覆盖目标场景的主要计算模式。每次架构调整都跑一遍这组模型看整体性能变化。这比只看单个模型靠谱得多能避免为了优化某个模型而牺牲整体的情况。这组模型要定期更新跟上模型演进的趋势。
阅读完成 · 觉得有帮助?
咨询建站