CoRa 张量编译器：面向最小填充的不规则张量编译

深度学习中使用的输入数据在形状和大小上常存在变化。在许多情况下，此类数据可使用具有非均匀形状的张量（即不规则张量）来表示。由于对不规则张量高效执行的支持有限且不可移植，当前的深度学习框架通常使用填充和掩码等技术使数据形状统一，然后将计算卸载到针对稠密张量代数的优化内核上。然而，此类技术会导致大量无用计算，从而造成性能损失。本文提出 CoRa，一种张量编译器，允许用户轻松生成面向各类 CPU 和 GPU 的不规则张量算子高效代码。在多种不规则张量算子以及 transformer 模型的编码器层上评估 CoRa，我们发现 CoRa（i）与手算优化的算子和 transformer 编码器实现具有竞争力；（ii）相较 PyTorch，在 Nvidia GPU 上对编码器实现 1.6 倍几何平均加速，在 ARM CPU 上对 transformer 中使用的多头注意力模块实现 1.86 倍几何平均加速。