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


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

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[0], subgraphs[
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

小源er

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值