随缘随笔
Insights Flow
0%

论文阅读[粗读]-THE CORA TENSOR COMPILER: COMPILATION FOR RAGGED TENSORS WITH MINIMAL PADDING

发表于 更新于 阅读次数： Valine：
本文字数： 2.1k 阅读时长 ≈ 2 分钟

陈天奇在MLSys 22发的另一篇文章，讲的是如何提高输入不对齐时的表现。

已有的方法面对不对齐的数据，常常是padding到一个长度，这样会造成额外的计算。本文解决这个问题，得到了Pytorch上16倍的速度提升，TensorFlow上1.4倍的性能提升

Introduction

DL模型的输入常常不对齐，如下图所示：

最近有一些编译器开始支持不对齐的数据，但是范围、效果很受限。而经典的cuDNN，oneDNN库都采用了MASK，PADDING的方法。这种方法会带来额外的计算损失(红色部分)，而且随着batch size的增大这个损失会更加严重。

本文提出了一套compiler-based方法解决不对齐数据的问题。稀疏和密集的tensor的研究很多，但大多不能应用到rag tensor,同时本文提出CORA (Compiler for Ragged Tensors)：

  • Irregularity in generated code：可变的循环带来可变的代码，使得GPU性能下降。
    • CORA使用minimal padding方法，来平衡高效代码生成和thread remapping strategies
  • Insufficient compiler mechanisms：由于可变循环内的依赖关系复杂，Representing transformations效果很差。并且稀疏tensor的优化方法可能不能应用到rag tensor，因为后者可能更加稀疏。
    • Cora用uninterpreted functions表示循环边界，来调度operation
  • Ill-fitting computation abstractions：接口和抽象不匹配的问题。dense compilers不能表达rag op，sparse compilers没有办法生成高效的code。
    • CORA给padding，thread remapping等提供一套抽象，来解决不匹配的问题。

和已有方法对比如下：

本文主要贡献：

  • 提出CORA，第一个高效解决ragged tensor的tensor编译器
  • 为ragged tensor提供了新的API，抽象和scheduleing方法
  • 进行了详细评测。

CORA OVERVIEW

基于以下的两个灵感：

  • ragged operation中的不对齐性在计算前就可以获取到
    • 可以在真正计算前预计算不对齐性带来的辅助信息
  • ragged tensor比起稀疏tensor，可以用\(\mathcal{O}(1)\)时间访问（HASH方法虽然可以\(\mathcal{O}(1)\)稀疏tensor，但不利于GPU执行）
    • 可以辅助生成高效的代码

CORA总体执行如上图所示

术语

  • variable loops or vloops ：如果内部循环的边界是和外层变量有关的
  • vloop nest：包含多个vloop
  • variable dimensions：当存储tensor不用padding时，齐形状取决于外部tensor。这个形状称为variable dimensions，或者vdims
  • cdims：固定形状
  • ragged tensor：有最少一个vdim的tensor

CORA’S RAGGED API

CORA的API要求对于ragged tensor描述出vdim是怎么计算的，由外层循环的什么特征来定义。满足此条件后，CORA就可以自动的符合后面计算API的实现。

Scheduling Primitives

  • Loop Scheduling：CORA不允许对vloop交换顺序，因为没有发现有什么实际需求。
  • Operation Splitting：CORA可以把OP分解成多个OP，后面把拆除的部分并行计算。
  • Horizontal Fusion：CORA支持hfusion，可以把最外层循环内的几个operator并行计算
  • Loop and Storage Padding：对于minimal padding(二对齐、四对齐)，可以对齐的存储，并且不浪费很多额外的空间
  • Tensor Dimension Scheduling：CORA允许用户交换tensor的维度
  • Load Balancing：CORA允许用户去自己满足thread remapping policy，来解决这个问题。

CORA’S RAGGED API LOWERING

这一部分讨论把CORA API lower到SSA-based IR的技术

Loop and Tensor Dimension Fusion

对于左上角的二重循环，如果想要变成一层循环，就需要计算出边界以及什么时候换。CORA预计算这些信息，再执行计算部分的代码。

Bounds Inference

  • Variable Loop Fusion：
  • Named Dimensions

Storage Access Lowering

考虑上面的4维tensor，24维是vdim。

CORA提供右图左面的组织形式，可以达到快速的\(\mathcal{O}(1)\)访问时间，通过预计算。CORA需要存储额外的形状信息。

IMPLEMENTATION

CORA通过扩展TVM的代码进行实现

  • Ragged API：目前的代码只支持vdim取决于最外层的dim
  • Lowering：这部分代码基于一篇已有工作修改。

Evaluation

大概就是CORA更好

相关文章