面向Ascend AI 加速器的并行扫描

我们设计并实现了使用Ascend AI 加速器的并行前缀和（scan）算法。Ascend 加速器具有专用的计算单元：用于高效矩阵乘法的立方体单元，以及针对向量运算优化的向量单元。提出的 scan算法的一个关键特点是广泛利用由立方体单元支持的矩阵乘法和累加。为展示这些算法的有效性，我们还实现并评估了几种在 AI 工作负载中常用的 scan-based 操作符，包括排序、张量掩码和 top-k / top-p 抽样。我们的单核结果显示，与仅使用向量实现相比，针对足够大的输入长度，速度提升范围为 $5\times$ 至 $9.6\times$。此外，我们 presented 一个多核 scan算法，充分利用Ascend的立方体单元和向量单元，达到了内存复制所实现内存带宽的74.9%。进一步，我们的基数排序实现利用矩阵乘法进行并行分割，展示了矩阵引擎提升复杂操作的潜力，相对于仅使用向量的基线实现实现了最高 $3.3\times$ 的速度提升。