AI编译器技术剖析（三）-树模型编译工具 Treelite 详解

近年来，AI应用程序已经无处不在。比如：智能家居设备由自然语言处理（NLP）和语音识别模型驱动，自动驾驶技术以计算机视觉模型为支柱。通常这些AI模型会部署在云平台、专用计算设备以及物联网传感器的内置微型芯片。

因此，我们在进行AI应用落地时，需要将AI模型从研发阶段转而部署到多种多样的生产环境，同时，也需要相当多的繁重工作。即使对于我们最熟悉的环境（例如：在 GPU 上），部署包含非标准算子的深度学习模型仍然需要大量的工程。为了解决这些繁琐的问题，AI 编译器应运而生，可以说未来十年AI编译器将迎来快速发展的黄金十年。

本系列将分享 AI 编译器的技术原理，本文为该系列第三篇。本文将针对Treelite树模型编译优化工具进行讲述。

Treelite 简介

Treelite 是一个用于将决策树集成高效部署生产环境的模型编译器。Treelite 提供了多个前端接口来与其他树库（如：XGBoost、LightGBM 和 scikit-learn）配合使用。Treelite 特别适用于大批量的数据进行推理，使用 Treelite 进行模型编译优化后的性能相比于原生的XGBoost、LightGBM模型会提升2-4倍。

Treelite 的设计

与许多树模型库进行互操作

Treelite 提供了多个前端接口来与其他树库配合使用。

首先，有一个专用接口来导入 XGBoost、LightGBM 和 scikit-learn 生成的模型。特别是，Treelite 提供了与 XGBoost 的无缝集成。

此外，Treelite 还提供模型构建器 API，让用户能够以编程方式指定其模型。如果用户使用其他包来训练模型，这非常有用。要求用户指定每个测试条件以及每个叶子节点的输出。

可扩展的模块化设计

Treelite 采用了模块化的设计，如图所示，前端（Treelite 中与其他树库交互的部分）和后端（Treelite 中生成可部署 C 文件的部分）之间有明显的分离。

模块化的设计方便以后的扩展。设计中的一个关键部分是决策树集成的通用模式。前端和后端不应以通用模式以外的方式进行通信。

该模式旨在适应随机森林和梯度提升树。

通过规则编译加快预测速度

通过将决策规则"编译"为嵌套的 if-else 条件，将给定的树集成模型转换为 C 程序。每个测试节点都转换为一对 if-else 语句，如下所示：

python 复制代码

if ( /* comparison test for the test node */ ) {
/* ... code for the left child node ... */
} else {
/* ... code for the right child node ... */
}

然后将左右子节点递归扩展为 C 代码，直到扩展到每个叶节点。

通过编译规则，我们可以实现特定于正在检查的模型的编译时优化。

以前，模型会在运行时从文件加载，并且预测逻辑不会注意到特定于该特定模型的任何信息。然而，通过规则编译，编译器可以访问正在编译的特定模型中的每一位（bit）信息，它可以使用这些信息来进一步优化生成的机器代码。作为早期演示，Treelite 提供了两种优化。

注释条件分支

我们通过计算训练数据中满足该条件的数据点的数量来预测每个条件的可能性。如果条件至少 50% 的时间为真（在训练数据上），则该条件被标记为"预期为真"。否则，它会被标记为"预期为假"。 GCC 和 clang 编译器都提供编译器内在 __builtin_expect 指定条件的可能结果。

这有助于编译器对分支的顺序做出更智能的决策，从而改进分支预测。

对条件使用整数阈值

此优化将测试节点中的所有阈值替换为整数，以便每个阈值条件执行整数比较而不是通常的浮点比较。阈值被"量化"为整数索引。

在 x86-64 等平台上，将浮点比较替换为整数比较，通过减少可执行代码大小和改善数据局部性可提高性能。

Treelite 生态

Treelite可以与如下工具一起搭配使用：

TL2cgen：TL2cgen 是决策树模型的模型编译器。您可以将任何决策树模型（随机森林、梯度提升模型）转换为 C 代码，并将其作为本机二进制文件发布。
Triton Inference Server FIL Backend：Triton 是一款机器学习推理服务框架，可轻松且高度优化地部署在几乎所有主要框架中训练的模型。该后端特别有助于在 Triton 中使用树模型（包括使用 XGBoost、LightGBM、Scikit-Learn 和 cuML 训练的模型）。
RAPIDS Forest Inference Library：Forest Inference Library 提供了一个轻量级、灵活的 API，用于从 GPU 上基于集成树模型预测结果。集成树可以是在 XGBoost、cuML、scikit-learn 或 LightGBM 中训练的梯度提升决策树 (GBDT) 或随机森林 (RF) 模型。

Treelite 和 TL2cgen 有什么关系？

从4.0版本开始，Treelite不再支持将树模型编译成C代码。 Treelite 的这一部分已迁移到 TL2cgen。

Treelite（从 4.0 开始）现在是一个小型库，使其他 C++ 应用程序能够在磁盘和网络上交换和存储决策树。通过使用 Treelite，应用程序编写者可以以最少的代码重复和高水平的正确地支持多种树模型。树使用有效的二进制格式存储。

TL2cgen 是一个模型编译器，可将树模型转换为 C 代码。 TL2cgen 与 Treelite 无缝集成。 Treelite 支持的任何树模型都可以使用 TL2cgen 进行转换。

TL2cgen

简介

TL2cgen 是决策树模型的模型编译器。您可以将任何决策树模型（随机森林、梯度提升模型）转换为 C 代码，并将其作为本机二进制文件发布。

TL2cgen 与 Treelite 无缝集成。 Treelite 支持的任何树模型都可以使用 TL2cgen 进行转换。

预测函数优化

TL2cgen 提供系统级优化以提高预测性能。但是模型信息被完整保留；优化仅影响预测的执行方式。

1. 条件分支注释

此优化分析并注释测试节点中的每个阈值条件以提高性能。

1.1 使用方法

第一步是为您的集成模型生成分支注释记录。首先需确保训练数据准备好。

python 复制代码

# model = your ensemble model (object of type treelite.Model)
# dmat = training data (object of type tl2cgen.DMatrix)

# Annotate branches by iterating over the training data
tl2cgen.annotate_branch(model, dmat, path="mymodel-annotation.json")

# Save the branch annotation record as a JSON file
annotator.save(path="mymodel-annotation.json")

要使用分支注释记载，请在导出模型时提供编译器参数annotate_in：

ini 复制代码

# Export a source directory
tl2cgen.generate_c_code(model, dirpath="./mymodel",
                        params={"annotate_in": "mymodel-annotation.json"})

# Export a source directory, packaged in a zip archive
tl2cgen.export_srcpkg(model, toolchain="gcc", pkgpath="./mymodel.zip",
                      libname="mymodel.so",
                      params={"annotate_in": "mymodel-annotation.json"})

# Export a shared library
tl2cgen.export_lib(model, toolchain="gcc", libpath="./mymodel.so",
                   params={"annotate_in": "mymodel-annotation.json"})

1.2 技术细节

原理：

现代 CPU 严重依赖一种称为分支预测的技术，在该技术中，它们提前"猜测"每个if-else分支中条件表达式的结果。给定一个程序：

scss 复制代码

if ( [conditional expression] ) {
  foo();
} else {
  bar();
}

如果给定条件可能为真，CPU 将预取函数foo()的指令。如果条件可能为假，CPU 将预取函数bar()的指令。可以说，正确预测条件分支对性能有很大影响。每次CPU正确预测分支时，它都可以保留之前预取的指令。每次CPU错误预测时，它必须丢弃预取的指令并重新获取另一组指令。

决策树集成的预测函数是非常困难的，因为分支预测是必须很好猜测的条件分支：

scss 复制代码

/* A slice of prediction function */
float predict_margin(const float* data) {
  float sum = 0.0f;
  if (!(data[0].missing != -1) || data[0].fvalue <= 9.5) {
    if (!(data[0].missing != -1) || data[0].fvalue <= 3.5) {
      if (!(data[10].missing != -1) || data[10].fvalue <= 0.74185) {
        if (!(data[0].missing != -1) || data[0].fvalue <= 1.5) {
          if (!(data[2].missing != -1) || data[2].fvalue <= 2.08671) {
            if ( (data[4].missing != -1) && data[4].fvalue <= 2.02632) {
              if (!(data[3].missing != -1) || data[3].fvalue <= 0.763339) {
                sum += (float)0.00758165;
              } else {
                sum += (float)0.0060202;
              }
            } else {
              if ( (data[1].missing != -1) && data[1].fvalue <= 0.0397456) {
                sum += (float)0.00415399;
              } else {
                sum += (float)0.00821985;
              }
            }
/* and so forth... */

事实上，检测节点中的每个阈值条件都需要预测。虽然 CPU 缺乏足够的信息来对这些条件做出正确的猜测，但我们可以通过提供一些信息来帮助做出正确的猜测。

为 C 编译器提供分支信息的机制：

我们通过计算训练数据中满足该条件的数据点的数量来预测每个条件的可能性。请参见下图的说明。

如果条件至少 50% 的时间为true（在训练数据上），则该条件被标记为"预期为真"：

sql 复制代码

/* expected to be true */
if ( __builtin_expect( [condition], 1 ) ) {
  ...
} else {
  ...
}

如果某个条件至少有 50% 的时间为false，则该条件将被标记为"预期为假"：

sql 复制代码

/* expected to be false */
if ( __builtin_expect( [condition], 0 ) ) {
  ...
} else {
  ...
}

关于表达式 __builtin_expect 的说明：

__builtin_expect 表达式是编译器固有的，用于为 C 编译器提供分支预测信息。 gcc 和 clang 都支持它。不幸的是，Microsoft Visual C++ 没有。要利用分支注释，请确保在目标计算机上使用 gcc 或 clang。

2. 对条件使用整数阈值

此优化将检测节点中的所有阈值替换为整数，以便每个阈值条件执行整数比较而不是通常的浮点比较。即将阈值量化为整数索引。

量化前:

css 复制代码

if (data[3].fvalue < 1.5) {  /* floating-point comparison */
  ...
}

量化后:

css 复制代码

if (data[3].qvalue < 3) {     /* integer comparison */
  ...
}

2.1 使用方法

只需在导出模型时添加编译器参数 quantize=1 即可：

ini 复制代码

# Export a source directory
tl2cgen.generate_c_code(model, dirpath="./mymodel",
                        params={"quantize": 1})

# Export a source directory, packaged in a zip archive
tl2cgen.export_srcpkg(model, toolchain="gcc", pkgpath="./mymodel.zip",
                      libname="mymodel.so",
                      params={"quantize": 1})

# Export a shared library
tl2cgen.export_lib(model, toolchain="gcc", libpath="./mymodel.so",
                   params={"quantize": 1})

2.2 技术细节

原理：

在某些平台（例如：x86-64）上，用整数替换浮点阈值有助于通过减少可执行代码大小 和提高数据局部性来提高性能。之所以如此，是因为在这些平台上，整数常量可以作为比较指令的一部分嵌入，而浮点常量则不能。

在看看 x86-64 平台整数比较如下：

css 复制代码

a <= 4

生成一条汇编指令：

x86asm 复制代码

cmpl    $4, 8(%rsp)       ;    8(%rsp) contains the variable a

由于整数常量 4 嵌入到比较指令 cmpl 中，因此我们只需从内存中获取变量 a 即可。

另一方面，浮点比较如下：

css 复制代码

b < 1.2f

生成两个汇编指令：

perl 复制代码

movss   250(%rip), %xmm0  ;  250(%rip) contains the constant 1.2f
ucomiss  12(%rsp), %xmm0  ;   12(%rsp) contains the variable b

浮点常量 1.2f 没有嵌入到比较指令 ucomiss 中。在进行比较之前，必须将常量使用 movss 提取到寄存器 xmm0 中。

总而言之，浮点比较需要的指令数量是整数比较的两倍，从而增加了可执行代码的大小；浮点比较涉及额外的提取指令 (movss)，可能会导致缓存未命中。

但是使用整数阈值会增加预测时的开销成本（详情见下面的特征映射的机制），因此应该确保整数比较的好处超过开销成本。

特征映射的机制：

当启用quantize选项时，TL2cgen 将收集集成树模型中出现的所有阈值。对于每个特征，将生成一个列表，其中按升序列出阈值：

less 复制代码

/* example of how per-feature threshold list may look like */

Feature 0:  [1.5, 6.5, 12.5]
Feature 3:  [0.15, 0.35, 1.5]
Feature 6:  [7, 9, 10, 135]

使用这些列表，我们可以通过简单的查找将任何数据点转换为整数索引。对于上例中的特征 0，值将映射到整数索引，如下所示：

ini 复制代码

Let x be the value of feature 0.

Assign -1 if          x  <  1.5
Assign  0 if          x ==  1.5
Assign  1 if   1.5  < x  <  6.5
Assign  2 if          x ==  6.5
Assign  3 if   6.5  < x  < 12.5
Assign  4 if          x == 12.5
Assign  5 if          x  > 12.5

一个浮点向量如何转换为整数索引向量的具体示例：

r 复制代码

feature id   0     1        2      3      4        5      6
            [7, missing, missing, 0.2, missing, missing, 20 ]
         => [3, missing, missing,   1, missing, missing,  5 ]

由于预测函数仍然需要接受浮点特征，因此在实际预测之前将对特征进行内部转换。

如果没有quantize选项预测函数如下所示：

arduino 复制代码

float predict_margin(const Entry* data) {
  /* ... Run through the trees to compute the leaf output score ... */

  return score;
}

现在，它有一个额外的步骤，将传入的数据向量映射为整数：

arduino 复制代码

float predict_margin(const Entry* data) {
  /* ... 将数据中的特征值量化为整数索引   ... */

  /* ... Run through the trees to compute the leaf output score ... */
  return score;
}

编译器参数

编译器参数影响从集成树模型生成C预测函数的方式。

annotate_in：模型注释文件的名称。
quantize：是否量化阈值点（0：否，>0：是）
parallel_comp：启用并行编译的选项；如果设置为非零，树将均匀分布到 [parallel_comp] 个文件中。设置此选项可以缩短编译时间并减少编译期间的内存消耗。
verbose：是否产生额外的消息。如果>0，则产生额外的消息。
native_lib_name：本机库名称（不带扩展名）
code_folding_req：用于折叠很少访问的子树的参数（无 if/else 块）；所有数据计数比决策树根节点低[code_folding_req]的节点将被折叠。要禁用折叠，请设置为 +inf。如果 hessian 总和可用，它们将用作数据计数的代理。
dump_array_as_elf：仅当编译器设置为故障安全时适用。如果设置为正值，故障安全编译器将不会向 C 代码发出大型常量数组。相反，数组将作为 ELF 二进制文件发出（仅限 Linux）。对于大型数组，直接转储 ELF 二进制文件比将它们传递给 C 编译器要快得多。

Treelite 实践

下面以随机森林为例，演示使用 Treelite 进行模型部署优化。

ini 复制代码

import treelite
import numpy as np
import time
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier

# 使用分类生成器生成10000个样本，1000个特征，分类标签数为2。
X, y = make_classification(n_samples=10000, n_features=1000, n_classes=2)

# 初始化树集成模型随机森林并进行模型训练
rf = RandomForestClassifier(n_estimators=10).fit(X, y)

# 从 scikit-learn 模型对象加载树集成模型
model = treelite.sklearn.import_model(rf)

output_path = "/Users/liguodong/output/model-inference/chapter02"

toolchain = 'gcc'

# 生成预测代码并转换为动态共享库，同时，根据指定文件保存生成的动态共享库。
model.export_lib(toolchain=toolchain, libpath=output_path + '/rfmodel.so', verbose=True)

总结

本文介绍了一个用于将决策树集成模型高效部署生产环境的模型编译器 Treelite 。

码字不易，如果觉得有帮助，欢迎点赞收藏加关注。

参考文档：

论文：Treelite: toolbox for decision tree deployment
Github：github.com/dmlc/treeli...
文档：treelite.readthedocs.io/en/latest/
Treelite：树模型部署加速工具（支持XGBoost、LightGBM和Sklearn）
Treelite模型加速