当前位置：首页 > article >正文

TVM计算图分割--分割方式

article 2024/11/17 6:26:55

文章目录

TVM中的计算图分割方式
- - 1. Partition Pass
  - 2. dataflow_pattern
  - 3. 内置图分割接口
  - 4. Pipeline Executor
  - 5. BYOC框架
  - 6. UMA

深度学习模型通常是用计算图来表示的。计算图是一种有向无环图，其中节点代表算子，表示一个操作，节点之间的边表示算子之间的数据依赖。计算图分割是指将计算图按照所属执行平台划分为若干个子图，每个子图分别执行在对应硬件上，以最大化值执行效率。

TVM中的计算图分割方式

1. Partition Pass

TVM中的所有涉及计算图分割的方式底层实现都是通过Partition Pass来实现的。因为在TVM中，计算图是通过Relay表示的，封装为IRModule对象。计算图分割操作是以IR上的优化Pass实现的。经过Partition Pass之后的子图会包装为复合函数并在原始IRModule中以函数调用的方式出现。
在这里插入图片描述

2. dataflow_pattern

TVM的前端接口还提供了一组算子匹配函数封装在dataflow_pattern中，可以按照制定的pattern匹配子图，并且封装了分割函数可以把匹配的子图划分出来。
在这里插入图片描述

dataflow_pattern还可以进一步做子图rewrite。具体参考：
https://tvm.apache.org/docs/reference/langref/relay_pattern.html
https://tvm.apache.org/docs/reference/api/python/relay/dataflow_pattern.html

3. 内置图分割接口

这些内置的接口都对应着一个明确的后端硬件，其实现过程还是基于BYOC框架完成的。

# 这种内置的接口实质是基于BYOC框架实现的
# ACL
from tvm.relay.op.contrib.arm_compute_lib import partition_for_arm_compute_lib
module = partition_for_arm_compute_lib(module)

# TensorRT
from tvm.relay.op.contrib.tensorrt import partition_for_tensorrt
mod, config = partition_for_tensorrt(mod, params)

#Vitis
from tvm.relay.op.contrib.vitis_ai import partition_for_vitis_ai
mod = partition_for_vitis_ai(mod, params, dpu=dpu_target)

4. Pipeline Executor

使用graph_split的分割方式是把relay表示的计算图(一个IRModule)分割为两个子计算图即两个IRModule。与普通Relay表示的计算图一样，只不过是编译需要pipeline_executor_build对象。该分割方式不涉及后端执行平台的信息。

# 自定义逻辑拆分
from test_pipeline_executor import graph_split

split_config = [{
   "op_name": "nn.relu", "op_index": 0}]
subgraphs = graph_split(net["main"], split_config, params)
mod0, mod1 = subgraphs