基于AXI交叉开关的轻量级多播扩展:为AI加速器打破内存墙

发布时间:2026/5/25 18:44:38

基于AXI交叉开关的轻量级多播扩展:为AI加速器打破内存墙 1. 项目概述为什么片上多播对现代AI加速器至关重要如果你最近几年深度参与过AI芯片或者高性能计算芯片的设计尤其是那些动辄集成数百个计算核心PE的加速器那你一定对“内存墙”和“数据搬运能耗”这两个词深恶痛绝。我们费尽心思把算力堆上去结果发现大量的时间和功耗都花在了等数据和搬数据上PE们经常处于“饥饿”状态。传统的解决方案比如增大片上缓存LLC、提升内存带宽固然有效但成本高昂且边际效益递减。这时一个更“聪明”的思路是减少不必要的数据搬运。在典型的AI工作负载尤其是矩阵乘法GEMM和张量运算中存在大量的数据复用模式。例如一个权重矩阵的同一列数据可能需要被广播到所有负责计算不同输出通道或不同输出位置的核心组。在传统的单播Unicast互连架构下这意味着源端需要发起N次完全相同的数据传输占用N倍的链路带宽和缓冲区资源效率极其低下。多播Multicast技术就是为了解决这个问题而生。它允许一次写事务Write Transaction将数据同时送达多个目标从设备。想象一下在公司的全员大会上CEO只需要讲一次话所有员工都能听到而不是让CEO对每个员工单独讲一遍——这就是多播的价值。在芯片内部实现高效、灵活且低成本的多播支持是释放大规模并行AI加速器性能潜力的关键一环。然而给现有的片上互连“打补丁”支持多播并非易事。许多商业方案如某些GPU的SM-to-SM网络对此讳莫如深而学术界的研究又往往聚焦于支持缓存一致性的片上网络NoC其协议复杂、开销大并不适合通常采用软件管理暂存器SPM的非一致性AI加速器。此外如何在不破坏现有标准协议如AXI兼容性的前提下以最小的硬件开销实现多播是一个极具挑战性的工程问题。本文要探讨的正是我们基于一个成熟的开源AXI交叉开关XBAR设计进行的一次轻量级多播功能扩展实践。我们的目标很明确设计一个既能保持AXI协议向后兼容又能在面积和时序上开销极小的多播方案并将其集成到一个真实的288核AI加速器Occamy中用实际的矩阵乘法性能提升来证明其价值。最终我们在一个16x16的大规模交叉开关上仅增加了12%的面积和6%的时序开销却换来了矩阵乘法核心29%的性能加速。这背后的设计思路、实现细节和踩过的坑正是我想分享的重点。2. 核心设计思路如何优雅地为AXI交叉开关注入多播能力2.1 基石选择开源AXI交叉开关作为起点我们的工作并非从零开始造轮子而是站在了巨人的肩膀上。我们选择了由Kurth等人开发的开源AXI交叉开关作为基础。这个设计本身已经非常成熟和高效它采用经典的集中式交换结构通过一个由多路复用器Mux和解复用器Demux构成的阵列将多个AXI主设备Master连接到多个AXI从设备Slave。其核心是一个地址映射表Address Map通过地址解码器将主设备的请求路由到正确的从设备端口。选择这个设计作为起点有几个关键考量开放性与可验证性开源代码允许我们深入理解其内部机制并进行定制化修改同时也便于社区复现和验证我们的成果。工业级标准基于AXIAdvanced eXtensible Interface协议这是Arm公司推出的片上总线事实标准确保了我们的扩展能与大量现有的商业和开源IP核无缝集成。清晰的架构其主从端口间通过Demux和Mux直接连接的结构相对清晰便于我们分析数据流并插入多播控制逻辑。我们的核心任务就是在这个清晰的架构上以最小的扰动增加对多播写事务的支持。2.2 灵魂一种可扩展的多地址编码方案多播面临的首要技术挑战是如何在一次事务中表达多个目标地址在网络上有诸如组播地址、目的地址列表等多种方案。但在片上系统的全局内存地址空间中我们需要一种既紧凑节省带宽和存储又能灵活表示常见数据分布模式如连续块、规则步长的编码方式。我们提出了一种基于**地址掩码Address Mask**的编码方案并将其巧妙地嵌入到AXI协议预留的aw_user信号中。这个方案的核心思想是基本原理将一个目标地址和与之等宽的掩码Mask配对。掩码中为1的位表示对应地址位是“无关位”Don‘t Care可以为0或1。通过设置n个无关位这个“地址-掩码”对就可以编码2^n个具体的地址。举例说明连续地址块假设要访问地址0x100到0x1FF这个256字节的连续空间。起始地址0x100的二进制是0001 0000 0000。这个区间的特点是低8位[7:0]从0000 0000变化到1111 1111。因此我们可以设置掩码的低8位为1即0000 0000 1111 1111。这样(addr0x100, mask0x00FF)就编码了整个区间。步长访问假设要访问所有偶数地址如0x200,0x202,0x204, ...。这些地址的二进制最低位始终为0。我们可以设置地址为0x200掩码最低位为1表示该位可0可1但结合地址固定为0实际上只匹配0同时为了匹配所有偶数地址我们需要识别出地址变化的规律。更典型的“步长”模式如访问一个二维矩阵中同一列的所有元素地址间隔固定行宽其地址集合可能无法用单个掩码完美表示但我们的编码对许多规则模式非常友好。注意这种掩码编码并非万能。它无法表示任意的、不规则的地址集合。但对于AI加速器中常见的数据广播模式如向多个核心的SPM中相同偏移地址写入相同数据或向一个连续地址区间广播这种编码非常高效。其编码大小仅与地址空间大小的对数相关与目标节点数量无关这对于拥有成百上千个核心的大规模系统至关重要。2.3 约束定义多播目标区域规则为了简化硬件设计并保证效率我们对可以配置为多播目标的地址区域即“多播规则”施加了两条约束大小必须是2的幂次。起始地址必须按大小对齐。例如一个大小为1KB2^10字节的多播区域其起始地址必须是1KB的整数倍。这两条约束使得任何区间形式的规则[start_addr, end_addr]都可以无损地转换为我们的掩码形式addr, mask。转换公式非常简单mfe.addr ife.start_addrmfe.mask ife.end_addr - ife.start_addr - 1这里的mfe代表掩码形式编码ife代表区间形式编码。mask的值其实就是区间大小减一其二进制表示中低位连续的1正好标出了“无关位”的范围。硬件上地址解码器需要被扩展以支持这种掩码匹配逻辑。3. 硬件实现详解在AXI交叉开关中嵌入多播逻辑3.1 地址解码器的扩展原有的地址解码器进行的是精确匹配或范围匹配。现在它需要处理带掩码的地址。对于每个配置的多播规则已转换为rule.addr, rule.mask和到来的请求req.addr, req.mask解码器需要判断请求的地址集合与该规则定义的区域是否有交集并计算出交集的掩码形式。判断是否有交集的逻辑可以简化为位运算合并掩码masked_bits req.mask | rule.mask。合并后为1的位表示在该位上请求和规则都允许变化或不关心。比较固定位match_bits ~(req.addr ^ rule.addr)。这行代码计算在那些合并掩码不为1的位上即“固定位”请求地址和规则地址是否相同。相同则对应位为1。判断匹配如果所有位固定位和无关位的“或”结果为1即(masked_bits | match_bits)为1则说明两者有交集。计算交集掩码和地址的逻辑也类似通过按位操作即可得到新的out.addr和out.mask。解码器的输出不再是一个单一的从设备选择信号而是一个位向量aw_select其中每一位指示对应的从设备端口是否至少包含一个目标地址。3.2 解复用器axi_demux的改造多播事务的调度与响应合并axi_demux模块负责将一个主设备的请求分发给多个从设备。支持多播后它面临两个核心挑战事务排序和响应合并。事务排序与死锁预防AXI协议允许不同ID的事务乱序完成但对同一ID的事务有顺序要求。在多播场景下一次多播写事务会衍生出多个到不同从设备的子事务共享同一个AXI ID。如果允许多播事务和单播事务乱序交错就需要复杂的缓冲区来跟踪所有子事务的状态极易导致死锁。 我们的解决方案是简化排序模型我们禁止多播事务与任何未完成的单播事务具有相同ID重叠反之亦然。这相当于为每个主设备端口维护了两个虚拟通道单播和多播同一时间只能有一个通道活跃。这虽然牺牲了一点潜在的并发性但极大地简化了控制逻辑避免了死锁。实现上axi_demux内部需要维护一个计数器来跟踪正在进行的多播事务数量并与一个可配置的最大值比较以进行流控。响应B通道合并一次多播写事务需要等待所有目标从设备都返回写响应B通道后才能向主设备返回一个最终的响应。我们使用了一个名为stream_join_dynamic的模块来实现动态的响应合并。它接收来自多个从设备的B响应并确保只在所有被寻址的从设备都返回响应后才产生一次向主设备的握手。 另一个问题是各个从设备返回的bresp响应状态如OKAY, SLVERR, DECERR可能不同。AXI协议没有规定多播场景下的响应合并语义。我们采取了保守且合理的策略如果任何一个从设备返回错误响应SLVERR或DECERR则向主设备返回SLVERR否则返回OKAY。同时我们禁止对多播事务使用独占访问Exclusive Access从而避免了EXOKAY状态的处理。合并逻辑最终简化为一个对所有从设备bresp信号的或OR归约操作。写数据W通道分发的死锁预防这是一个非常微妙但关键的坑。考虑一个经典的死锁场景两个主设备M0, M1同时发起多播目标都是两个从设备S0, S1。假设S0先接收了M0的地址相位AW后接收M1的AW而S1先接收了M1的AW后接收M0的AW。根据AXI协议一个从设备必须按照接收AW的顺序来接收对应的写数据W。那么S0会期待先收到M0的数据再收M1的数据S1则相反。如果M0和M1都试图同时发送数据它们可能会陷入等待M0等待S1准备好接收其数据而M1等待S0准备好接收其数据形成循环等待。 为了解决这个问题我们引入了aw.commit信号机制。其核心思想是一个主设备必须原子性地“获取”所有目标从设备端口的访问权限。在axi_demux中我们使用一个优先级编码器Leading-Zero Counter来确保所有相关的axi_mux模块对来自同一主设备端口的多播请求做出一致的选择。只有当所有目标axi_mux都表明其从设备端口准备好接收AW事务时axi_demux才断言aw.commit信号从而“释放”这些axi_mux允许后续的W数据传输。这打破了上述循环等待的条件有效预防了死锁。3.3 复用器axi_mux的改造多播与单播路径的仲裁axi_mux模块负责将来自多个主设备的请求合并到一个从设备端口。为了支持多播主要增加两部分逻辑路径选择从axi_demux传来的aw.is_mcast信号用于标识一个请求是多播还是单播。axi_mux内部需要为这两种事务提供独立的处理路径或共享路径但不同的控制逻辑。由于多播事务有更严格的排序限制不能与单播交错我们通常赋予多播路径更高的仲裁优先级。提交信号处理接收并处理来自axi_demux的aw.commit信号作为允许该多播事务真正占用从设备端口的最终许可。3.4 系统集成扩展到完整的加速器我们将改造后的多播交叉开关集成到了Occamy这款开源的288核AI加速器中。Occamy的片上网络是一个两级层次结构包含宽数据通路512位用于DMA和数据搬运和窄数据通路64位用于核心LSU同步和控制的交叉开关。我们替换了其中所有的交叉开关IP。更重要的是我们修改了Snitch计算簇的直接内存访问DMA引擎和加载存储单元LSU。这使得DMA引擎可以发起多播传输将数据从LLC末级缓存或其它位置一次性广播到多个计算簇的本地SPM中。LSU可以发起多播的存储操作用于高效的屏障同步或结果收集。这种软硬件协同的修改是最终能获得性能提升的基础。4. 开销与性能评估数据不说谎4.1 面积与时序开销轻量级的代价我们在GF 12LP工艺下使用1ns时钟周期约束目标1GHz对不同的NxN交叉开关进行了综合。交叉开关规模基线面积 (kGE)支持多播后面积 (kGE)面积开销频率达标情况2x227.328.6~4.9%是4x461.464.9~5.7%是8x8141.7154.9~9.3%是16x16372.3417.7~12.2%下降~6%结果分析面积开销可控面积开销随着交叉开关规模增大而上升这是因为多播控制逻辑如地址解码扩展、响应合并逻辑的复杂度与端口数N有关。但在最大的、物理上可实现的16x16规模上12.2%的额外面积开销是完全可以接受的尤其是考虑到它带来的带宽收益。时序影响微小关键路径略有增加导致16x16设计在1GHz下约有6%的时序违例需要略微降频或优化。对于更常见的8x8或更小规模都能满足1GHz目标。这表明我们的设计没有引入复杂的、深层次的关键路径。** scalability**开销的增长是亚线性的证明我们的编码和逻辑设计具有良好的可扩展性能够支撑未来规模更大的加速器互连。4.2 微基准测试多播传输的加速比我们设计了一个微基准测试让一个计算簇使用DMA将一块数据发送给其他所有计算簇。对比三种方案基线发起N次独立的单播DMA传输。软件层次化多播源簇先发给每个“组”的一个代表簇再由代表簇在组内广播。这利用了Occamy的两级网络拓扑实现了一定程度的并行。硬件多播使用我们扩展的DMA引擎和交叉开关发起一次多播传输。核心发现随着目标簇数量增加硬件多播相对于基线的加速比迅速上升在32个簇时加速比接近理想值根据阿姆达尔定律计算的等效并行比例达97%。这是因为固的顺序开销如往返延迟被平摊到了大量的传输中。即使与并行的软件层次化多播相比硬件多播仍有显著优势在32个簇的传输中平均有5.6倍的加速。这凸显了硬件原生支持在减少软件复杂度、降低延迟和提升并行度方面的价值。4.3 核心内核性能矩阵乘法的实际收益我们选择双精度矩阵乘法DGEMM作为评估内核这是AI和HPC的基石运算。在Occamy上我们将一个256x256的矩阵乘法任务大小刚好适配LLC并行到32个计算簇上。每个簇负责计算C矩阵的一个8x256的行块。关键优化点数据复用模式在分块矩阵乘法中每个簇需要重复使用自己负责的A矩阵子块8x256并与整个B矩阵256x256的各个列块依次相乘。在传统单播模式下每个簇都需要通过DMA从LLC加载整个B矩阵的当前列块。这是巨大的带宽浪费多播的应用利用硬件多播我们只需要将B矩阵的当前列块从LLC加载一次然后广播到所有32个计算簇。这立即将LLC到集群L1 SPM的读带宽需求降低了32倍。性能结果 我们在“屋顶线模型Roofline Model”中评估性能基线单播由于每个簇都独立加载B运算强度Operational Intensity, OI很低约为1.9 FLOP/Byte系统处于内存瓶颈区性能为114.4 GFLOPS。软件层次化多播通过软件模拟多播OI提升至约7 FLOP/Byte性能提升至303.0 GFLOPS。硬件多播OI大幅提升至约31.4 FLOP/Byte性能达到391.4 GFLOPS相比基线提升了3.4倍。这个实验有力地证明了即使在计算密度极高的矩阵乘法中内存子系统仍然是瓶颈。而一个轻量级、硬件实现的多播机制能够通过极小的面积代价显著优化数据移动模式将内核从内存瓶颈区推向计算瓶颈区从而释放出被束缚的算力。在我们的Occamy测试系统中最终实现了29%的端到端性能提升。5. 实操心得与避坑指南在将这套多播方案从设计文档落实到RTL代码再到集成进复杂SoC并最终验证性能的整个过程中我们积累了一些在论文中未必会详述但对实际工程至关重要的经验。5.1 协议兼容性与扩展性的权衡利用aw_user信号AXI协议的awuser/wuser/buser等信号是留给用户自定义的。用它们来传递多播掩码是保持协议向后兼容的黄金法则。这意味着任何标准的AXI主设备或监视器Monitor依然能正常处理我们的链路只是无法利用多播功能。切忌去修改AXI标准通道如awaddr,awlen的语义那会带来灾难性的兼容性问题。响应合并语义需自定义正如前文所述AXI协议没有定义多播的响应合并规则。我们选择“任一错误则整体报错”的保守策略。在实际芯片中这可能需要对软件栈驱动程序、运行时库进行相应修改使其能理解并处理这种聚合后的错误信号。更复杂的策略如返回错误位图会增加硬件开销需要根据系统需求仔细权衡。5.2 死锁预防是重中之重片上网络死锁是无声的杀手。我们遇到的W通道死锁场景非常典型。教训是在设计多播流控时不能只孤立地看一个主设备或一个从设备必须从全局事务依赖关系图的角度去思考。引入aw.commit机制来实现“原子性资源获取”是一个有效且相对简单的解决方案。在验证阶段必须构造极端并发场景的压力测试让多个主设备同时向重叠的从设备集合发起多播和单播请求以确保在各种交错情况下都不会挂死。5.3 验证策略从单元测试到系统仿真模块级验证首先对修改后的axi_demux和axi_mux进行充分的单元测试。重点测试掩码地址解码是否正确、多播/单播事务隔离机制、响应合并逻辑、以及上面提到的死锁预防逻辑。使用UVM或类似的验证方法学构造随机化测试向量非常有效。交叉开关级集成测试搭建一个包含多个主从设备的测试平台让主设备随机发起单播和多播请求检查数据是否正确路由到所有目标响应是否正确合并返回。同时要监测吞吐量和延迟。系统级性能建模与仿真在集成到Occamy这样复杂的SoC之前我们先用周期精确的性能模型如Gem5评估了多播可能带来的收益这帮助我们确定了优化的重点如DMA和LSU的修改。最终的RTL仿真使用QuestaSim虽然耗时但却是获得真实性能数据的唯一途径。建议在系统仿真中除了核心计算内核也要加入带有随机间隔的干扰流量以模拟真实的内存访问压力这样得到的性能数据更可靠。5.4 面积与性能的折衷点我们的设计选择了“禁止多播与单播事务ID交错”来简化控制逻辑。这可能会轻微影响某些混合负载下的绝对性能但换来了面积和时序上的显著优势。对于以规则数据流为主的AI加速器这是一个非常划算的 trade-off。如果你的目标应用场景有大量细粒度、交错出现的单播和多播小事务可能需要考虑更复杂的、支持交错的事务跟踪机制但这必然会带来更大的面积和功耗开销。5.5 软件生态的适配硬件提供了多播能力但需要软件来驾驭。我们修改了DMA引擎的驱动和LSU的指令但这只是第一步。要让编译器、编程模型如OpenCL、CUDA的某种变体乃至高层框架如TensorFlow/PyTorch的算子库能够自动识别出可以应用多播的数据搬运模式如矩阵乘法的权重广播并生成相应的代码才是这项技术发挥最大威力的关键。这通常需要一个支持“批量存储”或“广播存储”指令的ISA扩展以及相应的编译器优化。我们的工作为硬件层打下了基础而软件栈的优化将是一个持续的工程。6. 总结与展望这次为AXI交叉开关添加多播支持的实践再次印证了一个道理在追求极致性能的芯片设计中架构的优化往往比单纯的工艺提升或频率拉升更为有效。我们通过一个相对轻量级的硬件修改~12%面积开销在关键的矩阵乘法负载上换来了近30%的系统级性能提升投资回报率非常高。我们的方案的价值在于其实用性和可用性它基于广泛使用的AXI标准对现有设计侵入性小开销明确且可控并且已经在一个真实的多核加速器上得到了验证和性能量化。代码的开源也使得社区可以在此基础上进行进一步的探索和改进。对于未来我认为有几个方向值得深入更灵活的编码当前的掩码编码对某些非对齐或不规则的数据模式支持有限。是否可以设计一种低开销的、支持小规模地址列表的编码作为补充与一致性协议的协同虽然本文聚焦于非一致性SPM架构但在支持缓存一致性的多核CPU或异构系统中多播如何与MESI等缓存一致性协议协同工作以高效实现数据广播和失效是一个有趣的问题。自动化代码生成与优化如何从高级的算法描述或计算图中自动识别出多播机会并生成最优的多播传输指令序列是打通从硬件特性到应用性能的“最后一公里”。芯片设计尤其是在AI加速器这个快速迭代的领域永远是在面积、功耗、性能和灵活性之间走钢丝。这次多播交叉开关的设计是我们在这条钢丝上向前迈出的扎实一步。它或许不是最炫酷的技术但却是能实实在在提升能效比的工程实践。如果你正在设计大规模并行处理器的互连网络不妨仔细评估一下数据流中的广播模式也许一个简单的多播扩展就是你打破内存墙的那把钥匙。

相关新闻