VTensor: Using Virtual Tensors to Build a Layout-Oblivious AI Programming Framework

Abstract

Tensor是开发AI算法中一个常见的编程接口。layout涉及了tensor数据在内存中的顺序并且会通过影响数据局部性而影响性能，因此AI库中layout是有默认的使用规范。由于AI应用中可以使用任意的layout，现有的AI系统并不会提供编程抽象来保护AI库约定的layout，开发人员需要编写很多于layout相关的代码，这减低了开发新的库或者开发新算子的效率。此外，开发者要指定layout转换的操作给内部算子来解决不确定的输入layout，从而失去了layout优化的机会。基于多态性的思想，我们提出一个与layout无关的虚拟tensor编程接口，称为VTensor框架，这允许开发人员在写一个新算子的时候不需要考虑tensor的物理layout。另外，VTensor框架在运行时会执行全局layout推断，以一种简单、清晰的方式解决虚拟tensor中对layout的需求，并且执行运行时layout无关优化来全局减少layout转换操作的数量。实验证明，开发者可以避免写layout相关的代码。和TF相比可以平均减少写一个新算子时47.82%的代码行数，平均提升18.65%的性能。

2. 动机和背景

2.1中介绍基于数据流图的layout-aware编程模型
之后通过示例来介绍layout-aware编程模型中的挑战
2.2专注于可维护性
2.3专注于layout优化

2.1 TF的设计

Tensor贯穿TF模型，并且与layout强相关(这一点做过AI部署工作的同学应该非常熟悉，万恶的NHWC，以及ONNX的NCHW等等)。
VTensor中将tensor的语义从计算图和库中抽离，将他们重新表示为虚拟tensor和物理tensor。

更直白地讲，常量视为物理Tensor，数据流视为虚拟Tensor

2.2 TF等维护性差(VTensor优势1)

在实现算子库的时候，通常需要将所有输入layout的情况都考虑到，否则一旦用户的layout比较独特，那么这个算子执行大概率会失败。使用Vtensor时,只需要专注于算子的计算逻辑即可,访问tensor的信息不需要知道物理layout。很多逻辑无关的代码都是自动生成的，当需要分配物理tensor时，直接调用相应的接口，就可以通过VTensor来声明。

2.3 未经优化的layout转换

将layout的转换从算子中提出，创建单独的layout转换节点

3. VTensor框架预览

提供四个编程接口来定义算子、描述lib、说明如何内部调用库

• VTensor API：可以通过编写Compute函数访问VTensor来实现算子的计算逻辑，Dispatch API作为lib和设备的调用入口。
• PTensor API：开发者可以声明和VTensor相对应的PTensor。PTensor通常是在库调用函数中使用，其中包含有具体的layout，可以为确定状态NCHW，也可以为非确定状态ANY，同时提供require函数来指定layout。
• library描述文件：促进多library支持。每个library都需要一个这样的文件来描述PTensor的物理layout和library之间layout名称的映射。同时在描述文件中会指定layout转换程序，并从多个备选方案中选择layout的指导方针(人话就是指定library名称来选择对应的转换pass) <硬件不同时可以根据这个描述文件来自动做layout的转换>
• 框架API：向VTensor框架中注册处理程序。运行时layout优化期间会调用这些处理程序。每个libraryInvoker函数都必须被注册为调用的入口函数，这个函数中会包含如何从抽象参数列表中创建实际的参数。同时会插入一个tag，用于标记eval和infer时执行不同长度的操作，eval只会执行tag之前，infer全部
  同时提出DLR(dynamic layout resolver)在运行时动态优化数据布局。通过分析访问模式，该框架可以自动调整内存布局以提高性能。