文章目录
- TVM中的计算图分割方式
- 1. Partition Pass
- 2. dataflow_pattern
- 3. 内置图分割接口
- 4. Pipeline Executor
- 5. BYOC框架
- 6. Collage
- 7. 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_splitsplit_config = [{"op_name": "nn.relu", "op_index": 0}]
subgraphs = graph_split(net["main"], split_config, params)
mod0, mod1 = subgraphs[0], subgraphs[