转换布局 Pass#
Author: Animesh Jain
1. 背景#
数据布局格式描述了数据在内存中的排列方式。例如，Tensorflow 框架卷积运算符的默认数据布局为 NHWC，即数据为四维，以行优先格式排列，其中 N 是第一维，C 是最后一维。数据布局在模型性能中起着重要作用，对空间和时间局部性有着显著影响。例如，TVM 中的 Intel x86 后端喜欢使用 NCHWc 布局，其中 C 维度在 2 个维度上被铺设(tiled)以有效地利用数据局部性。同样，CUDA 后端喜欢数据布局以 NCHW 格式排列。
实际上，在编译器工具链中，TVM 必须处理数据布局 - 框架解析器、Relay 布局转换和 TOPI 调度。随着我们向第三方代码生成集成的发展，可能会有他们自己的数据布局限制，在 TVM 工具链的所有级别处理布局将变得更加具有挑战性。因此，TVM 开发了新的 Relay pass—— ConvertLayout ——以减少由于布局处理而产生的一些复杂性。
如果您想直接了解 ConvertLayout Pass 的用法，请直接跳到第 4 节 —— 用法。
2. 动机和概述#
让我们看简单的场景，以了解由于不同布局而产生的复杂性——假设我们想要为 ARM 边缘设备编译 Tensorflow NHWC Graph。但是，假设我们目前在 ARM 上仅支持 NCHW 调度。因此，框架布局与 TOPI 支持的布局之间存在不匹配。处理此不匹配的一种方法是在每个卷积之前和之后插入布局转换，以便生成的卷积具有 NCHW 输入数据布局并可以使用 TOPI 调度。但是，由于存在过多的布局转换，这可能会导致性能下降。
在其他用例中也遇到了类似的问题
无法在 Nvidia GPU 上运行 TFLite 图。TOPI 仅为 GPU 提供 NCHW 调度。
AlterOpLayout中的逻辑越来越复杂，以支持不同的布局转换对。
由于额外的布局转换，TF graphs 的性能不够优化。
第三方代码生成集成（如 TensorRT）中的复杂性，它偏好数据布局为一种格式。
为了解决这些问题，引入了 ConvertLayout 传递，建立基础设施以使用最少的数据布局转换改变整个 Graph 的数据布局。在理想情况下，将只对数据进行两次布局转换，一次在开始时，一次在结束时。下面是一个示例，展示了该转换过程。
# Original graph - 2 convolutions in NHWC format.
fn (%x: Tensor[(1, 56, 56, 64), float32], %weight1: Tensor[(3, 3, 64, 32), float32], %weight2: Tensor[(3, 3, 32, 32), float32]) {
%0 = nn.conv2d(%x, %weight1, padding=[1, 1], channels=32, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO");
%1 = nn.relu(%0);
%2 = nn.conv2d(%1, %weight2, padding=[1, 1], channels=32, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO");
nn.relu(%2)
}
# After ConvertLayout - For data, there is a transform at the start and at the end.
# For weights, there are transforms to adapt to NCHW layout. These will be removed by FoldConstant pass.
fn (%x: Tensor[(1, 56, 56, 64), float32], %weight1: Tensor[(3, 3, 64, 32), float32], %weight2: Tensor[(3, 3, 32, 32), float32]) {
%0 = layout_transform(%x, src_layout="NHWC", dst_layout="NCHW") /* ty=Tensor[(1, 64, 56, 56), float32] */;
%1 = layout_transform(%weight1, src_layout="HWIO", dst_layout="OIHW") /* ty=Tensor[(32, 64, 3, 3), float32] */;
%2 = nn.conv2d(%0, %1, padding=[1, 1], channels=32, kernel_size=[3, 3]) /* ty=Tensor[(1, 32, 56, 56), float32] */;
%3 = nn.relu(%2) /* ty=Tensor[(1, 32, 56, 56), float32] */;
%4 = layout_transform(%weight2, src_layout="HWIO", dst_layout="OIHW") /* ty=Tensor[(32, 32, 3, 3), float32] */;
%5 = nn.conv2d(%3, %4, padding=[1, 1], channels=32, kernel_size=[3, 3]) /* ty=Tensor[(1, 32, 56, 56), float32] */;
%6 = nn.relu(%5) /* ty=Tensor[(1, 32, 56, 56), float32] */;
layout_transform(%6, src_layout="NCHW", dst_layout="NHWC") /* ty=Tensor[(1, 56, 56, 32), float32] */
}
3. 设计#
在深入了解 ConvertLayout 传递之前，让我们根据算子对其进行分类，分为三类，这种分类将有助于后面理解 ConvertLayout 传递的详细信息。
布局无关 (Layout agnostic) - Relu、pow 等。这些算子不受数据布局的影响，既不会影响其功能，也不会影响性能。
轻度布局敏感 (Lightly-layout sensitive) - pad、concatenate、reduce 等算子。如果在它们之前进行布局转换，这些算子的某些属性将受到功能影响。然而，就性能而言，差异并不显著。对于这些算子，只需适应前一个算子的输出数据布局即可。
重度布局敏感 (Heavily-layout sensitive) - 卷积、conv2d_transpose 等算子。这些算子在功能和性能上都受到数据布局的重大影响，同时它们的算子属性也是数据布局。通常情况下，对于这些算子，修改输入数据布局是有益的（如果输入数据布局不是有效的），而其余的“布局无关”和“轻度布局敏感”的算子则适应于这些“重度布局敏感”算子所控制的布局。
现在让我们来看一下两个相关的 Relay 算子属性。每个 Relay 算子都有属性，比如 InferType，可以由 TVM 开发人员定义。通常情况下，Relay pass 逐个遍历图中的算子，并读取这些算子的属性。例如，InferType pass 查看一个算子的 InferType 属性，确定其输出形状和类型，然后将其传递给下一个算子的 InferType 属性。类似地，在我们的上下文中，我们有两个这样的属性 - FTVMConvertLayout 和 FInferCorrectLayout。ConvertLayout pass 遍历整个图，查看这两个属性以及一个自动布局转换插入模块来处理数据布局。因此，整个过程可以分为三个步骤：
运行 FTVMConvertLayout 属性 - 这使开发人员可以将原始的 Relay expr 转换为具有新布局的新的 Relay expr，从而允许用户定义布局更改。开发人员可以使用 Python 回调函数来方便地进行操作。这仅用于对布局敏感的算子。
运行 FTVMInferCorretLayout 属性 - 我们可以将其视为布局推断。它查看原始输入布局和新输入布局，这些新布局来自于先前的算子或 FTVMConvertLayout 修改后的表达式（如果使用了该属性）。这可以被轻度布局敏感的算子用来适应新的数据布局。每个算子都进行布局推断。
自动插入布局转换 - 前面的布局推断步骤为输入表达式设置了新的布局。如果这些布局与原始布局不同，则该组件会自动插入布局转换。因此，开发人员不需要对此组件进行任何操作。
这些步骤按顺序逐个算子进行，其中 ConvertLayout 步骤不断将新布局传递给下一个算子的属性，最终逐个算子修改整个图。现在，让我们来看几个定义这两个属性的示例。
FTVMConvertLayout - 用于布局修改的 Python 回调 - 这用于 极其敏感于布局 的算子。例如，可以返回新的卷积算子，其中包括新的数据和核布局。另外两个组件将推断布局并在必要时插入布局转换。卷积算子的示例如下，将其转换为 NCHW 布局。
@reg.register_convert_op_layout("nn.conv2d")
def convert_conv2d(attrs, inputs, tinfos, desired_layouts):
"""Convert Layout pass registration for conv2d op.
Parameters
----------
attrs : tvm.attrs.Attrs
Attributes of current convolution
inputs : list of tvm.relay.Expr
The args of the Relay expr to be legalized
tinfos : list of types
List of input and output types
desired_layouts : list of layout strings
List of layouts defining our desired
layout for the data and kernel inputs respectively.
Returns
-------
result : tvm.relay.Expr
The transformed expr
"""
from tvm import relay
data, weight = inputs
new_attrs = dict(attrs)
# We expect 2 desired layouts to be specified, one for the data and one for the kernel.
assert len(desired_layouts) == 2, "A desired layout is expected for both of nn.conv2d's inputs"
# Use the first entry in desired layouts which specifies the data layout.
# The expected ordering of layouts for this operator is defined by this function.
desired_data_layout, desired_kernel_layout = map(str, desired_layouts)
assert desired_data_layout != "default", "Data layout cannot be default"
new_attrs['data_layout'] = desired_data_layout
if desired_data_layout == 'NCHW':
if desired_kernel_layout != 'default':
new_attrs['kernel_layout'] = desired_kernel_layout
else:
new_attrs['kernel_layout'] = 'OIHW'
# Actual insertion of layout transforms is taken care internally
# by ConvertLayout pass.
return relay.nn.conv2d(data, weight, **new_attrs)
raise ValueError('Layout %s is not yet supported' % desired_data_layout)
FInferCorrectLayout - 布局推断 - 目前，此属性仅在 C++ 中公开。该函数接受原始输入布局和新的输入布局（从上一个算子或来自用于布局修改的 python 回调传递），并推断出最终的数据布局。布局推断对每个算子都进行调用。使用方式可能因不同算子类别而异。对于不考虑布局的算子，我们只需在此函数中返回新的数据布局即可。对于轻度敏感和极度敏感于布局的算子，我们可以更改算子属性（如 concatenate 的 axis，pad 的 pad_width），以便我们可以适应新的数据布局，避免插入布局转换。让我们看几个例子来更好地理解这个过程。
First example is for layout agnostic operators. These operators do not have any operator attributes that are affected by data layouts, so we just adapt to new layouts.
// For operator set its attributes like following
// .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ElemwiseArbitraryLayout);
// Take arbitrary input layouts and copy to outputs.
inline Array<Array<Layout>> ElemwiseArbitraryLayout(const Attrs& attrs,
const Array<Layout>& new_in_layouts,
const Array<Layout>& old_in_layouts,
const Array<Array<IndexExpr>> &old_in_shapes) {
Layout ret;
if (new_in_layouts.defined()) {
ICHECK_GE(new_in_layouts.size(), 1);
ret = new_in_layouts[0];
} else {
for (size_t i = 0; i < old_in_layouts.size(); ++i) {
if (old_in_layouts[i].defined()) {
ret = old_in_layouts[i];
break;
}
}
}
return Array<Array<Layout>>{Array<Layout>(old_in_layouts.size(), ret), {ret}};
}
第一个例子是针对不考虑布局的算子。这些算子没有任何受数据布局影响的算子属性，因此我们只需适应新的布局。
Array<Array<Layout>> BatchNormInferCorrectLayout(const Attrs& attrs,
const Array<Layout>& new_in_layouts,
const Array<Layout>& old_in_layouts,
const Array<Array<IndexExpr>>& old_in_shapes) {
BatchNormAttrs* param = const_cast<BatchNormAttrs*>(attrs.as<BatchNormAttrs>());
size_t axis =
param->axis < 0 ? param->axis + old_in_shapes[0].size() : static_cast<size_t>(param->axis);
Layout ret = Layout::Undef();
// For example, consider old_layout = NHWC, and new_layout = NCHW, and param->axis = 3
if (new_in_layouts.defined() && old_in_layouts.defined()) {
// Get the new C axis. Extract the dim in old layout. Find the index of that dim in next layout.
// Following line gives bn_dim = C as old_layout = NHWC, axis = 3
const auto& bn_dim = old_in_layouts[0][axis];
// The new_index is 1 because new_layout = NCHW and bn_dim is C
auto new_index = new_in_layouts[0].IndexOf(bn_dim);
// We modify the layout-dependent attribute here - axis to 1.
param->axis = new_index;
// Finally, we adapt to the new layout.
ret = new_in_layouts[0];
} else if (old_in_layouts.defined()) {
ret = old_in_layouts[0];
}
// In case both new and old layouts are undefined, then there is no need of a change.
// ConvertLayout pass skips the automatic insertion of layout transforms in this case.
// Following line is not important to tutorial. But, layout inference needs to define
// the layout for all input and output data layouts. For batch norm, the other inputs
// and outputs are vector having length of C dim in the input. So, we set the other
// layouts as C. BN has 5 inputs, 3 outputs. The last 4 inputs and last 2 outputs
// have "C" layout.
Layout c_layout = Layout("C");
return Array<Array<Layout>>{{ret, c_layout, c_layout, c_layout, c_layout},
{ret, c_layout, c_layout}};
}
4. 用法#
ConvertLayout pass 的使用非常容易。该 pass 不是默认的 relay.build 管道的一部分。其预期的使用方法是在从框架到 relay 的解析器和 relay.build 模块调用之间调用它。
为了指定要转换的布局，我们创建了一个映射，将布局敏感的算子映射到该算子所需的布局列表。下面的第一个例子指定了数据布局，我们允许内核布局自动转换为 TVM 支持的内核布局（针对特定的数据布局和算子）。这是通过使用 “default” 关键字指定的。第二个示例展示了如何转换为我们选择的特定内核布局。值得注意的是，以下示例将转换为相同的布局，即 {‘nn.conv2d’: [‘NCHW’, ‘default’]} == {‘nn.conv2d’: [‘NCHW’, ‘OIHW’]}。
# TFlite framework to Relay parser - Default layout is NHWC
mod, params = relay.frontend.from_tflite(tflite_model,
shape_dict=shape_dict,
dtype_dict=dtype_dict)
# We assume our model's heavily-layout sensitive operators only consist of nn.conv2d
desired_layouts = {'nn.conv2d': ['NCHW', 'default']}
# Convert the layout to NCHW
# RemoveUnunsedFunctions is used to clean up the graph.
seq = tvm.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
relay.transform.ConvertLayout(desired_layouts)])
with tvm.transform.PassContext(opt_level=3):
mod = seq(mod)
# Call relay compilation
with relay.build_config(opt_level=3):
graph, lib, params = relay.build(mod, target, params=params)
desired_layouts = {'nn.conv2d': ['NCHW', 'OIHW']}
pass = relay.transform.ConvertLayout(desired_layouts)
布局的顺序由 register_convert_op_layout(“OPNAME”) 的实现定义，您可以参考 docstring，其中应明确说明预期的布局。在上面的例子中，它是 [data_layout, kernel_layout]。
当前实现支持几乎所有常用于图像分类模型中的算子。但是，如果在图中遇到太多的数据布局转换，很可能存在一些需要特殊处理布局的算子，如第 3 节所述。一些 pull requests 可以在这种情况下提供帮助，它们包括：
Batch Norm 的布局推断 - Batch Norm 属于轻度敏感算子的范畴。该 PR 展示了如何处理 Batch Norm 的布局推断。
卷积 的 Python 回调 - 对于高度敏感的算子，有时需要进行 Python 回调。该 PR 展示了如何为卷积算子定义 Python 回调函数。