Antkillerfarm Hacking V7.0

TVM实战, Bazel

2021-11-10

TVM实战

问题的由来

最近客户反馈我们的backend导入Pytorch模型会出错，而TFLite模型是OK的。

打印模型的IR后，我们发现：

这是Pytorch模型的IR片段：

  %0 = qnn.quantize(%input, 0.0186579f, 114, out_dtype="uint8", axis=1);
  %1 = nn.pad(%0, 114f, pad_width=[[0, 0], [0, 0], [1, 1], [1, 1]]);
  %2 = qnn.quantize(%features.0.0_weight, 0.00288958f, 0, out_dtype="int8", axis=0);
  %3 = qnn.conv2d(%1, %2, 114, 0, 0.0186579f, 0.00288958f, strides=[2, 2], padding=[0, 0, 0, 0], channels=32, kernel_size=[3, 3], out_dtype="int32");
  %4 = qnn.quantize(%features.0.0_bias, 5.39136e-05f, 0, out_dtype="int32", axis=0);
  %5 = nn.bias_add(%3, %4);
  %6 = qnn.requantize(%5, 5.39136e-05f, 0, 0.0150183f, 0, axis=1, out_dtype="int32");
  %7 = clip(%6, a_min=0f, a_max=255f);
  %8 = cast(%7, dtype="uint8");

这是TFLite模型的IR片段：

  %0 = qnn.quantize(%input, 0.0186579f /* ty=float32 */, 114 /* ty=int32 */, out_dtype="uint8", axis=1) /* ty=Tensor[(1, 3, 224, 224), uint8] */;
  %1 = nn.pad(%0, 114f /* ty=float32 */, pad_width=[[0, 0], [0, 0], [1, 1], [1, 1]]) /* ty=Tensor[(1, 3, 226, 226), uint8] */;
  %2 = qnn.conv2d(%1, meta[relay.Constant][0] /* ty=Tensor[(32, 3, 3, 3), int8] */, 114 /* ty=int32 */, 0 /* ty=int32 */, 0.0186579f /* ty=float32 */, 0.00288958f /* ty=float32 */, strides=[2, 2], padding=[0, 0, 0, 0], channels=32, kernel_size=[3, 3], out_dtype="int32") /* ty=Tensor[(1, 32, 112, 112), int32] */;
  %3 = nn.bias_add(%2, meta[relay.Constant][1] /* ty=Tensor[(32), int32] */) /* ty=Tensor[(1, 32, 112, 112), int32] */;
  %4 = qnn.requantize(%3, 5.39136e-05f /* ty=float32 */, 0 /* ty=int32 */, 0.0150183f /* ty=float32 */, 0 /* ty=int32 */, axis=1, out_dtype="uint8") /* ty=Tensor[(1, 32, 112, 112), uint8] */;

可以看出同一个QnnConv2D两者的展开形式存在一定的差异，但是语义上却基本是一致的。

Relay IR

在继续主题之前，我们首先来看看这个展开形式的差异是如何造成的。

TFLite/Pytorch模型导入之后，有一个转换成Relay IR的过程，也就是所谓的frontend。

相关代码在：

TFLite: python/tvm/relay/frontend/tflite.py

Pytorch: python/tvm/relay/frontend/pytorch.py和python/tvm/relay/frontend/qnn_torch.py

可以看出Pytorch的量化模型的weight是浮点数，而TFLite的量化模型的weight是整数。

有frontend自然就有backend：

python/tvm/relay/op/contrib/ethosn.py

和切割计算图有关的代码主要是：

seq = tvm.transform.Sequential(
    [
        ......
        transform.MergeComposite(pattern_table()),
        transform.AnnotateTarget("ethos-n"),
        transform.MergeCompilerRegions(),
        transform.PartitionGraph(),
    ]
)
seq(mod)

这里的Pass基本看名字就知道功能了，唯一需要关注的是MergeComposite。

这是pattern_table的其中一个表项：

def qnn_conv_pattern():
    pattern = is_op("nn.pad")(wildcard(), wildcard()) | wildcard()
    pattern = is_op("qnn.conv2d")(
        pattern, is_constant(), is_constant(), is_constant(), is_constant(), is_constant()
    )
    pattern = is_op("nn.bias_add")(pattern, is_constant())
    pattern = is_op("qnn.requantize")(
        pattern, is_constant(), is_constant(), is_constant(), is_constant()
    )
    return pattern

通常IR会定的比较细粒度，而AI硬件更喜欢粗粒度的op。所以往往若干个IR op组合起来，才能得到一个AI硬件op。这时就需要进行模板匹配。

一般来说，数据可以分为变量和常量（is_constant()），如果既可能是变量，也可能是常量的话，就用wildcard()匹配。

MergeComposite会将这个模板打包成一个函数，其中的变量会写为该函数的参数，而其中的常量则被放到全局的meta data里。

添加Pass

下面我们来看看如何添加Pass进行这样的转换。

class QnnQuantizeConstFold : public DFPatternRewrite {
 public:
  QnnQuantizeConstFold() {
    data_pat_ = IsConstant();
    pattern_ = IsOp("qnn.quantize")({data_pat_, IsConstant(), IsConstant()});
  }

  Expr Callback(const Expr& pre, const Expr& post,
                const Map<DFPattern, Array<Expr>>& node_map) const override {
    ......

    if (output->dtype == DataType::Int(8)) {
      return QuantizeData<int8_t>(......);
    } else if (output->dtype == DataType::Int(32)) {
      return QuantizeData<int32_t>(......);
    }
    return post;
  }

 protected:
  DFPattern data_pat_;
};

Rewrite_是Pass最重要的函数。

Pass有很多种，其中最简单的当属DFPatternRewrite。它的Rewrite_已经写好了，我们仅仅需要实现Callback即可。

Callback的参数中，pre表示原始Expr，post表示替换后的Expr，返回值也是替换后的Expr。（方便使用链式调用？）

如果Pass什么都不做的话，直接返回post就可以了。

PS：虽然pre和post在运行之初是相同的，但是一旦替换开始，两者就有差异了，所以如果是写入的内容，一定要引用post里的那份。

Python版的Pass

Pass不仅能用C++写，也可用python写。

class QnnQuantizeConstFold(tvm.relay.dataflow_pattern.DFPatternCallback):
    def __init__(self, require_type=False):
        super().__init__(require_type)
        self.pattern = is_op("qnn.quantize")(
            is_constant(), is_constant(), is_constant()
        )

    def callback(self, pre, post, node_map):
        ......
        if (dtype == "int8"):
             return tvm.relay.Constant(tvm.nd.array(data.astype(np.int8)))
        if (dtype == "int32"):
            return tvm.relay.Constant(tvm.nd.array(data.astype(np.int32)))
        return post

可以看出，写法也是大同小异，只是更便于集成，也不用写FFI接口了。

使用方法：

func = mod["main"]
func = tvm.relay.Function(func.params, func.body, None, func.type_params, func.attrs)
func = tvm.relay.dataflow_pattern.rewrite(RemoveClipAfterRequantize(), func)
func = tvm.relay.dataflow_pattern.rewrite(QnnQuantizeConstFold(), func)
mod = tvm.IRModule.from_expr(func)

这里展示了DFPatternCallback的使用方法，还有IRModule和Expr的相互转换方法。

参考：

https://tvm.apache.org/docs/reference/langref/relay_pattern.html

Pattern Matching in Relay

tvmc

tvmc是tvm提供的一个命令行接口。

export TARGET="bnns, llvm -device=arm_cpu -mtriple=aarch64-linux-gnu"
python3 -m tvm.driver.tvmc compile ./mobilenet_v1_0.25_224_quant.tflite --target "$TARGET" -o tvmc.tar --cross-compiler "$CC" --cross-compiler-options "$CC_OPTIONS"

要点如下：

  1. 注意TARGET的写法，这展示了如何将一个包含空格的字符串作为一个bash变量传递给命令行参数的做法。定义变量时的双引号和使用时的双引号都是必不可少的。

  2. 有些backend需要partation，用来将可执行的op分配到该设备上，因此该TARGET包含了两部分（用逗号分隔），其中第一部分用于partation。

相关代码在：

python/tvm/driver/tvmc/composite_target.py

Bazel

安装：

echo "deb [arch=amd64] http://storage.googleapis.com/bazel-apt stable jdk1.8" | sudo tee /etc/apt/sources.list.d/bazel.list
curl https://bazel.build/bazel-release.pub.gpg | sudo apt-key add -
sudo apt update && sudo apt install bazel

官网：

https://bazel.build/

文档：

https://docs.bazel.build/versions/master/bazel-overview.html

下载：

https://github.com/bazelbuild/bazel/releases

安装：

chmod +x bazel-<version>-installer-linux-x86_64.sh

./bazel-<version>-installer-linux-x86_64.sh --user

需要注意的是，即使这种安装也只是部分安装，安装后仍然需要联网下载依赖，并不能达到离线安装的效果。这也是Bazel的安装文件体积越来越小的根本原因。只有早期100M+的安装包，才是能离线安装的。

bazel使用Starlark语言编写扩展，后者的语法主要源自python，但并不是通用语言，也没有python那么强的功能。

Starlark官网：

https://docs.bazel.build/versions/master/skylark/language.html

使用示例：

https://github.com/antkillerfarm/antkillerfarm_crazy/tree/master/cpp/bazel

  • 每个项目的根目录必须有一个WORKSPACE文件。

  • 其他源代码目录下有一个BUILD文件。

编译：

bazel build <target>

这里的target在上例中就是cc_binary中的name = "test"：

bazel build test

bazel还可以从网上下载依赖文件：http_archive

清理：

bazel clean --expunge

可以通过在工作区中运行bazel info output_base来查找该工作区的输出库。但请注意，有一个包含最后一个命令输出的command.log文件，而bazel info本身是一个命令，因此这将覆盖command.log。

参考：

http://www.cnblogs.com/Jack47/p/bazel-faq.html

Google软件构建工具Bazel FAQ

https://zhuanlan.zhihu.com/p/47397799

bazel项目添加automake/autoconf项目解决办法

https://www.cnblogs.com/puyangsky/p/7596282.html

bazel的使用

https://zhuanlan.zhihu.com/p/421489117

创建宏与规则

引用第三方库

git_repository(
    name = "XXX",
    remote = "https://github.com/XXX/XXX.git",
    tag = "v1.1",
    patches = ["xxx.patch"],
    patch_args = ["-p1"],
    verbose = True,
)

https://brentley.dev/patching-bazel-external-dependencies/

Patching Bazel external dependencies

动态链接库

dlopen系列的so原则上代码不能互相依赖。一个解耦的办法就是同一份代码编译之后，同时链接到两个so中。示例：

cc_shared_library(
    name = "npu",
    roots = [":npu_lib"],
    static_deps = [
        "//:__subpackages__",
        "@//:__subpackages__",
        "@tsl//:__subpackages__",
    ],
)

上文中的static_deps就是来干这个事情的。

Bazelisk

这是基于Go语言编写的Bazel启动器，它会为你的工作区下载最适合的Bazel，并且透明的将命令转发给该Bazel。

由于Bazellisk提供了和Bazel一样的接口，因此通常直接将其命名为bazel：

sudo wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.16.0/bazelisk-linux-amd64
sudo chmod +x /usr/local/bin/bazel

Antkillerfarm
antkillerfarm@sohu.com

github.com/antkillerfarm

Fork me on GitHub