欢迎来到尧图网

客户服务 关于我们

您的位置:首页 > 科技 > 能源 > TVM计算图分割--分割方式

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

2024/11/30 11:48:06 来源:https://blog.csdn.net/sb985/article/details/143616119  浏览:    关键词:TVM计算图分割--分割方式

文章目录

  • 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[

版权声明:

本网仅为发布的内容提供存储空间,不对发表、转载的内容提供任何形式的保证。凡本网注明“来源:XXX网络”的作品,均转载自其它媒体,著作权归作者所有,商业转载请联系作者获得授权,非商业转载请注明出处。

我们尊重并感谢每一位作者,均已注明文章来源和作者。如因作品内容、版权或其它问题,请及时与我们联系,联系邮箱:809451989@qq.com,投稿邮箱:809451989@qq.com