[论文解读] The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding

Pratik Fegade, Tianqi Chen+2 more|arXiv (Cornell University)|Oct 19, 2021

Parallel Computing and Optimization Techniques被引用 8

一句话总结

CoRa 是一种张量编译器，通过最小化填充来编译不规则张量操作，从而在 CPU 和 GPU 上实现高效执行，通过为不规则数据形状生成优化代码。在 NVIDIA V100 GPU 上，CoRa 相较于 PyTorch 实现了 1.6× 的几何平均加速；在 64 核 ARM CPU 上，相较于 TensorFlow 实现了 1.37× 的加速，用于 Transformer 多头注意力。

ABSTRACT

There is often variation in the shape and size of input data used for deep learning. In many cases, such data can be represented using tensors with non-uniform shapes, or ragged tensors. Due to limited and non-portable support for efficient execution on ragged tensors, current deep learning frameworks generally use techniques such as padding and masking to make the data shapes uniform and then offload the computations to optimized kernels for dense tensor algebra. Such techniques can, however, lead to a lot of wasted computation and therefore, a loss in performance. This paper presents CoRa, a tensor compiler that allows users to easily generate efficient code for ragged tensor operators targeting a wide range of CPUs and GPUs. Evaluating CoRa on a variety of operators on ragged tensors as well as on an encoder layer of the transformer model, we find that CoRa (i)performs competitively with hand-optimized implementations of the operators and the transformer encoder and (ii) achieves, over PyTorch, a 1.6X geomean speedup for the encoder on an Nvidia GPU and a 1.86X geomean speedup for the multi-head attention module used in transformers on an ARM CPU.

研究动机与目标

• 解决深度学习框架在处理可变形状数据时因填充和掩码导致的性能瓶颈。
• 实现在 CPU 和 GPU 等多样化硬件平台上的不规则张量操作高效、可移植的编译。
• 通过自动化高性能代码生成来减少对手动优化内核的依赖，适用于不规则张量操作。
• 克服现有密集张量和稀疏张量编译器在表达或优化不规则张量操作方面的局限性。
• 提供统一、可扩展的编译栈，支持 CPU 和 GPU 后端，因形状不规则导致的性能损失最小化。

提出的方法

• 设计一种领域特定语言（DSL），用于表达具有可变大小张量切片和不规则循环边界的不规则张量操作。
• 引入一种代码生成流水线，将高层操作描述转换为优化的 LLVM 或 CUDA 代码，通过形状感知调度避免填充。
• 利用循环变换技术（如分块、融合和并行化）处理可变循环边界，同时保持数据局部性和指令级并行性。
• 集成 Z3 进行约束求解，集成 LLVM 进行低层代码生成，实现内存访问模式和内核启动配置的自动优化。
• 通过生成针对目标平台的代码并结合硬件感知优化（如寄存器分块和 GPU 上的 warp 级原语），支持 CPU 和 GPU 后端。
• 采用模块化架构，将操作定义与代码生成解耦，便于扩展至新操作和硬件平台。

实验结果

研究问题

• RQ1 张量编译器能否在不依赖填充或掩码的情况下为不规则张量操作生成高性能代码？
• RQ2 CoRa 自动生成的代码在多种硬件平台上的性能与手工优化实现相比如何？
• RQ3 与基于填充的方法相比，CoRa 在 Transformer 模型中在多大程度上减少了计算浪费？
• RQ4 CoRa 是否能在保持可移植性和低实现成本的同时，实现 CPU 和 GPU 工作负载的竞争力性能？
• RQ5 CoRa 在优化复杂真实场景操作（如多头注意力和 Transformer 编码器层）方面的有效性如何？

主要发现

• 在 NVIDIA V100 GPU 上，CoRa 相较于 PyTorch 在 Transformer 编码器层实现了 1.6× 的几何平均加速，表现出与手工优化内核相当的竞争力。
• 在 64 核 ARM Neoverse N1 CPU 上，CoRa 相较于 TensorFlow 在多头注意力模块实现了 1.37× 的几何平均加速，显著优于基线。
• CoRa 生成内核的性能与手工优化实现相当，表明自动代码生成可达到专家级优化水平。
• CoRa 通过避免填充减少了计算浪费，尤其在大批量场景中，填充开销显著，如 Transformer 编码器的 FLOP 分析所示。
• 该编译器成功处理了不规则张量操作中的复杂不规则循环结构，即使在 GPU 等高度并行架构上也未牺牲性能。
• 在八个 NLP 数据集上的评估表明，CoRa 在不同序列长度和数据分布下性能稳定，验证了其泛化能力。