TVM计算图分割--Collage

最新推荐文章于 2026-03-25 11:38:15 发布

原创

最新推荐文章于 2026-03-25 11:38:15 发布 · 1.1k 阅读

标签

#TVM #计算图分割

1 背景

为满足高效部署的需要，整合大量优化的tensor代数库和运行时做为后端成为必要之举。现在的深度学习后端可以分为两类：1）算子库(operator kernel libraries)，为每个DL算子单独提供高效地低阶kernel实现。这些库一般也支持算子融合，可以将多个算子按照融合规则组合为一个kernel实现。2）图推理库(graph inference libraries)，将整个DL模型作为输入并生成高效地运行时代码。除了包含优化的算子库，图推理库还可以处理跨kernel的优化如内存优化。

无缝整合后端面临两个挑战，一是要为不同编程模型整合多种可用后端同时保证性能问题；二是优化后端的分配策略权衡每个后端的性能优势。但是现有的框架都是基于规则的方式，按照固定的优先策略将每个算子降阶到不同后端，如上图所示。比如在Pytorch中卷积操作优先选择cuDNN，矩阵乘优先选择cuBLAS。手动地选择后端，这会降低系统潜在性能，减缓不同层之间的连续性，并且开发者需要快速改变编码规则适配新版本的代码库。并且即便对于同一类型的算子，最优后端也是依赖于硬件和算子配置的。

为解决这些问题，Collage提供了一个无缝的整合方案，通过提供有表达力的后端接口允许用户精确地指定多种后端，并且Collage可以为给定模型自动搜寻最优的后端安置策略。

2 结构

Collage包含两个关键部分，一是提供了具有表达力的后端注册接口来指定后端的能力，该能力基于支持的算子类型、配置信息和融合规则确定。该接口只需要理解提供的匹配语言即可。二是后端分配策略。Collage采用了两阶段的优化策略来处理operator kernel library 和graph inference library两种类型后端的特征。该系统会将全部后端和硬件考虑在内，自动地匹配计算图和后端算子pattern，找到最优的分配方案。

下图列举了Collage的整体结构，以DNN模型和可用的后端作为输入，然后为底层硬件优化后端分配策略。这里考虑的是对于给定目标硬件选择不同组后端，通过度量组件M来评估性能特点。

2.1 Backend pattern abstraction

准确的指定执行平台对于发挥后端全部能力至关重要。后端模板（backend pattern）定义了一组算子和可以直接部署在各个后端的算子组合。然后Collage提供了一个两阶段的抽象，对于简单的pattern，直接列举出全部支持的算子pattern。对于复杂情况，用户可以指定pattern创建规则，如指定简单算子类型和复杂算子融合规则。当融合规则给定，整个计算图上的满足条件的模板都会被保存在后端模板注册表中。

Collage的pattern是Relay Pattern Language的拓展，为所有支持的算子直接指定pattern。对于复杂情况难以直接指定的，可以通过pattern rule指定一组算子的融合规则。该规则可以通过python语法特性来表示。

对于每一个后端，用户可以通过直接指定pattern和pattern rule来提供编程接口。如上图所示，如果后端支持的算子有限，可以将所有pattern枚举出来如3-12行代码；对于复杂pattern，Collage的pattern生成器可以指定复杂规则如14-53行所示。生成器内部可以指定检查器判断约束条件（20-30行），可以以python代码的形式指定融合规则（32-49行）。

在确定好pattern后，Collage会搜索所有满足条件的规则，并将这些规则实例添加到backend pattern registry中。上图展示了搜索逻辑，pattern generator会按照pattern rule贪心地搜索。对于每个算子首先判断是否可在后端执行，然后进一步向外扩展一步判断是否满足融合规则，一旦条件满足则将对应pattern添加至pattern注册表，然后进一步扩大一步匹配范围。

2.2 Backend placement optimizer

当可用的模板注册好后，Collage会采用一个两阶段的优化方法为当前执行环境找到一个最优的后端安置策略。首先，算子库级的分配优化器会采用动态规划算法为每个算子寻找最优的后端平台，不需要考虑跨内核优化问题。其次，图级分配优化器通过演化搜索算法找最优后端。该部分会考虑跨内核的优化，并作为算子级优化的补充。

2.2.1 问题定义

对于多个后端，Collage的目的是最大化利用可用后端发挥性能最大化。对于后端选择问题，可以将计算图记作G，后端pattern记作B， $b=(p,d)\in B$ 是一个算子和pattern对儿，p表示具体pattern，d表示后端的标识。p和d都不是唯一对应的，一个p可以在多个d中被满足，一个d可以包含多个p，这样p和d的组合可以唯一确定一个后端和pattern。

对于M个匹配的子图 $g_i$ 和后端patern $b_i$ ,集合 $P(G) = (g_i, b_i) | b_i \in B, g_i \in G, g_i \cap g_j = \varnothing forall i,j \in 1,2,\cdots M where i \neq j$ 可以表示计算图G上的后端分配策略即将子图 $g_i$ 放置在 $b_i$ 上，其中Cost(P(G))表示函数P(G)的执行时间。Collage的目的是通过最小化Cost(P(G))找到一个分配策略 $P_{opt}$ 。