Convert Layout Pass
作者:Animesh Jain
1. 背景
数据布局格式(Data layout format)描述了数据在内存中的布局方式。例如,卷积算子的 TensorFlow 框架默认数据布局是 NHWC,即数据是四维的,并以行优先格式布局,N 为第一维,C 为最后一维。
数据布局在模型性能中起主要作用,对空间和时间局部性有重大影响。例如,TVM 中的 Intel x86 后端往往采用 NCHWc 布局,其中 C 维度以二维形式平铺,从而有效利用数据局部性(data locality)。同样,CUDA 后端的数据布局往往采用 NCHW 格式。
本质上,TVM 必须处理整个编译器工具链中的数据布局——框架解析器、Relay 布局转换和 TOPI schedule。随着我们转向第三方 codegen 集成(它可能有自己的数据布局限制),处理 TVM 工具链中所有级别的布局变得更具挑战性。因此,我们开发了一个新的 Relay pass——ConvertLayout——减少布局处理引发的问题。
若要了解 ConvertLayout Pass 的用法,请直接跳到本节第 4 部分 - 用法。