Efficient Reduce and AllReduce communication collectives are a critical cornerstone of high-performance computing (HPC) applications. We present the first systematic investigation of Reduce and AllReduce on the Cerebras Wafer-Scale Engine (WSE). This architecture has been shown to achieve unprecedented performance both for machine learning workloads and other computational problems like FFT. We introduce a performance model to estimate the execution time of algorithms on the WSE and validate our predictions experimentally for a wide range of input sizes. In addition to existing implementations, we design and implement several new algorithms specifically tailored to the architecture. Moreover, we establish a lower bound for the runtime of a Reduce operation on the WSE. Based on our model, we automatically generate code that achieves near-optimal performance across the whole range of input sizes. Experiments demonstrate that our new Reduce and AllReduce algorithms outperform the current vendor solution by up to 3.27x. Additionally, our model predicts performance with less than 4% error. The proposed communication collectives increase the range of HPC applications that can benefit from the high throughput of the WSE. Our model-driven methodology demonstrates a disciplined approach that can lead the way to further algorithmic advancements on wafer-scale architectures.

翻译：高效的Reduce和AllReduce通信集合操作是高性能计算（HPC）应用的关键基石。本文首次系统性地研究了Cerebras晶圆级引擎（WSE）上的Reduce和AllReduce操作。该架构已在机器学习工作负载及FFT等其他计算问题上展现出前所未有的性能。我们引入了一个性能模型来估计WSE上算法的执行时间，并针对广泛输入规模范围通过实验验证了预测结果。除现有实现外，我们还设计并实现了多个专门适配该架构的新算法。此外，我们建立了WSE上Reduce操作运行时间的一个下界。基于此模型，我们自动生成了在完整输入规模范围内实现近乎最优性能的代码。实验表明，我们提出的新Reduce和AllReduce算法性能相较于现有供应商解决方案最高提升达3.27倍。同时，该模型预测性能的误差低于4%。所提出的通信集合操作扩大了能从WSE高吞吐量中受益的HPC应用范围。我们的模型驱动方法论展现了一种严谨的研究路径，可引领晶圆级架构上实现进一步的算法突破。