TVM 部署指南：使用 mrvl 将模型编译并运行于 Marvell MLIP（Octeon DPU / 模拟器）

TVM 部署指南使用 mrvl 将模型编译并运行于 Marvell MLIPOcteon DPU / 模拟器【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址: https://gitcode.com/gh_mirrors/tvm7/tvm本文基于 TVM 官方文档 mrvl.rst 编写面向希望在 MarvellMarvell(R)Octeon 系列 DPU 上部署深度学习推理模型的开发者。文章完整覆盖启用 mrvl 的 TVM 构建 → TVMC 命令行编译 → 模拟器运行 → Python API 端到端推理的完整链路并结合当前仓库源码深入讲解 mrvl 分区、代码生成与运行时的工作原理帮助读者真正跑通一次 Marvell MLIP 的模型编译与推理。1. 背景Marvell MLIP 与 TVM 的 mrvl 支持Marvell 提供一系列高性能数据处理单元DPUData Processing Unit集成了通用计算、高速 I/O 与工作负载加速器其中包含 Marvell 的机器学习推理处理器MLIPMachine Learning Inference Processor——一个高度优化的集成式推理引擎。TVM 通过mrvl库接入 MLIP在编译期将模型划分为由 mrvl 支持、可卸载到 MLIP 加速执行的区域以及由 LLVM 生成通用计算代码的区域。在运行时mrvl 库支持两种执行方式MLIP 硬件上的原生执行面向 Octeon 系列带 ML 加速器的处理器Marvell ML 模拟器mrvl-mlsim用于在无硬件环境下完成功能验证。需要说明的是当前仓库状态下面向硬件加速的运行时支持仍处于 WIPWork In Progress阶段仓库文档明确标注硬件加速的运行时支持将在未来的 PR 中补齐见 mrvl.rst 第 3.2 节。因此本文以模拟器路径为主展开这是当前可以实际跑通验证的路径。2. 构建启用 mrvl 支持的 TVM2.1 获取 TVM 源码首先克隆 TVM 仓库含子模块git clone --recursive https://github.com/apache/tvm tvm cd tvm2.2 使用 docker 构建并启动 mrvl 容器仓库在 docker/Dockerfile.demo_mrvl 中提供了现成的 mrvl 演示镜像其内部会依次完成克隆 TVM 主仓库 → 通过 tests/scripts/task_config_build_mrvl.sh 生成开启 mrvl 的构建配置并cmake make编译 → 安装 Python 包 → 克隆并配置 Marvell 工具链MarvellMLTools将其bin目录加入$PATH。构建并进入容器./docker/build.sh demo_mrvl bash # 构建 docker 容器 ./docker/bash.sh tvm.demo_mrvl # 加载 docker 镜像若自行从源码构建关键编译开关为USE_MRVL。参考 tests/scripts/task_config_build_mrvl.sh在build/config.cmake中追加set(USE_MRVL ON)即可开启 mrvl 目标支持该脚本同时开启了USE_LLVM、USE_LIBBACKTRACE等配套选项。容器内已把 Marvell 工具链加入PATH这是后续编译的前提——mrvl 代码生成阶段会调用 Marvell 后端编译器mrvl-tmlc与模拟器mrvl-mlsim。3. 使用 TVMC 命令行编译模型TVMC 是 TVM 面向性能优化的命令行驱动模型可针对 mrvl 目标完成编译与运行。除了 TVMC 的通用选项外mrvl 还提供若干专属选项以--target-mrvl-option形式传入。3.1 TVMC 编译命令语法python3 -m tvm.driver.tvmc compile --targetmrvl, llvm --target-llvm-options --target-mrvl-options --tvm-generic-options model_file.onnx--targetmrvl, llvm是一个复合目标composite targetTVMC 会在编译前自动调用 mrvl 的分区流水线。在 python/tvm/driver/tvmc/composite_target.py 中可以看到mrvl已被注册到REGISTERED_CODEGENmrvl: { config_key: relay.ext.mrvl.options, pass_default: True, pass_pipeline: partition_for_mrvl, },即 TVMC 编译时默认会执行partition_for_mrvl完成图分区并将命令行中的 mrvl 选项写入 PassContext 配置relay.ext.mrvl.options。3.2 完整示例ARMv9 核 集成 MLIP 的 cn10ka 处理器下面的命令针对ARMv9 core 与集成 MLIP 的 cn10ka 处理器且仅使用块内 4 个 tilepython3 -m tvm.driver.tvmc compile --targetmrvl, llvm \ --target-llvm-mtripleaarch64-linux-gnu --target-llvm-mcpuneoverse-n2 \ --target-mrvl-num_tiles4 \ --cross-compiler aarch64-linux-gnu-gcc \ --output model.tar \ mnist-12.onnx参数解析参数含义--targetmrvl, llvm复合目标mrvl 负责受支持算子的加速编译LLVM 负责其余通用计算--target-llvm-mtripleaarch64-linux-gnuLLVM 目标三元组ARM 64 位 Linux--target-llvm-mcpuneoverse-n2LLVM 微架构ARMv9 Neoverse-N2--target-mrvl-num_tiles4MLIP 使用的最大 tile 数--cross-compiler aarch64-linux-gnu-gcc交叉编译器用于生成 ARM 侧的可执行/运行环境--output model.tar输出打包的模型库3.3 TVMC 的 mrvl 专属命令行选项--target-mrvl-mcpu --target-mrvl-num_tiles --target-mrvl-mattr各选项的完整说明如下与 compiler_attr.cc 中MrvlCompilerConfigNode的属性定义一一对应mcpuMarvell ML 推理处理器的 CPU 类别可选值{cn10ka, cnf10kb}默认cn10ka。num_tiles允许使用的最大 tile 数量可选值{1, 2, 4, 8}默认8。mattrmrvl 的属性集合用于指定数据类型、代码生成选项与优化策略可选值{quantize, wb_pin_ocm}。mattr 支持的属性明细1. quantize数据类型可选值{fp16, int8}默认fp16int8仍处于 WIP完整支持将在未来 PR 中加入。2. wb_pin_ocm权重偏置驻留片上内存通过将模型的权重weights与偏置bias预加载到片上内存on chip memory来优化运行时可选值{0, 1}默认0不预加载。从 python/tvm/relay/op/contrib/mrvl.py 的add_attributes实现可以看到如果用户在命令行没有显式给出某些选项代码会自动补齐合理的默认值最终拼成传给 Marvell 后端的编译选项字符串未指定arch时补-archmlip未指定quantize时补-quantizefp16未指定wb_pin_ocm时补-wb_pin_ocm0未指定num_tiles时补-num_tiles8。这些属性随后会被写入分区后的函数compiler_opts_string、working_dir等属性由 codegen 阶段读取并最终嵌入运行时。4. 面向模拟器 LLVM / x86_64编译 ONNX 模型在 TVMC 的 mrvl 流程中模型被划分为Marvell 区域与LLVM 区域。对每个 Marvell 子图编译会生成序列化产物nodes.json与const.jsonnodes.json是适合 Marvell 编译器mrvl-tmlc消费的模型图表示Marvell 编译器将模型图编译为包含 MLIP 指令的模型二进制model binary。4.1 面向模拟器的模型编译python3 -m tvm.driver.tvmc compile --targetmrvl, llvm \ --target-mrvl-num_tiles4 --output model.tar model.onnx这里不指定--cross-compiler与 ARM 三元组属于 x86_64 主机上的默认路径生成的模型二进制交由 MLIP 模拟器执行。4.2 在 x86_64 主机上使用 MLIP 模拟器运行模型编译产物由 Marvell 的 MLIP 模拟器mrvl-mlsim仿真运行python3 -m tvm.driver.tvmc run --inputs infer.npz --outputs predict.npz model.tar --number0--inputs infer.npz/--outputs predict.npz指定输入与输出文件numpy npz 格式--number0运行全部批次的推理。从运行时源码 mrvl_runtime.cc 可以看到模拟器运行时模块的类型键为mrvl_simtype_key() mrvl_sim它保存子图的符号名、nodes_json 与编译生成的二进制bin_code执行时调用RunMarvellSimulator定义于 mrvl_sw_runtime_lib.h并在 mrvl.py 的tvm.mrvl.RunSim中于临时目录内运行模拟器命令。5. 使用 Python API 编译并运行模型MNIST 端到端除 TVMC 外也可以直接用 TVM Python API 完成编译与推理。下面以 MNIST 为例完整走一遍。5.1 下载 MNIST ONNX 模型cd $HOME wget https://github.com/onnx/models/raw/main/validated/vision/classification/mnist/model/mnist-12.onnx5.2 导入依赖模块import tvm, onnx import numpy as np import tvm.relay as relay from tvm.contrib import graph_executor from tvm.relay.op.contrib.mrvl import partition_for_mrvl from tvm.relay.build_module import build from keras.datasets import mnist5.3 加载模型并构建 Relay 图onnx_model onnx.load(mnist-12.onnx) shape_dict {Input3 : (1,1,28,28)} mod, params relay.frontend.from_onnx(onnx_model, shape_dict)5.4 定义选项字典并完成 mrvl 分区调用partition_for_mrvl对图进行注解annotate与分区partition所有被 mrvl 支持的算子都会被标记并卸载到 Marvell 硬件加速器其余算子则走常规的 LLVM 编译与 ARM 代码生成。tvm_target llvm option_dict {num_tiles: 4} mod partition_for_mrvl(mod, params, **option_dict)从 mrvl.py 的实现可以看到partition_for_mrvl内部会执行一整套预处理流水线InferType、去除 dropout/copy、FoldConstant、SimplifyExpr、ConvertLayout到 NHWC、MergeComposite按 mrvl 模式表融合、AnnotateTarget(mrvl)、MergeCompilerRegions、PartitionGraph随后再对回退到 LLVM 的非计算密集/不支持子图进行反分区repartition_mrvl_subgraphs最后把用户选项以属性形式附加到各个 mrvl 子图函数上。5.5 构建 Relay 图使用partition_for_mrvl返回的新模块进行构建并同时把选项字典放入 PassContext 配置与 TVMC 注册的relay.ext.mrvl.options一致with tvm.transform.PassContext(opt_level3, config{relay.ext.mrvl.options : option_dict}): model_lib relay.build(mod, tvm_target, paramsparams)5.6 生成运行时图dev tvm.cpu() model_rt_graph graph_executor.GraphModule(model_libdefault)5.7 准备测试数据并设置输入(train_X, train_y), (test_X, test_y) mnist.load_data() image tvm.nd.array(test_X[0].reshape(1, 1, 28, 28).astype(float32) / 255) inputs_dict {} inputs_dict[Input3] image model_rt_graph.set_input(**inputs_dict)5.8 运行推理并打印输出model_rt_graph.run() output_tensor model_rt_graph.get_output(0).numpy() print (output_tensor)6. 源码级原理mrvl 的代码生成与运行时6.1 支持算子与模式表mrvl 后端支持的模式注册在 mrvl.py 的mrvl_pattern_table中主要包括mrvl.conv2d_nhwc2nhwc卷积支持 pad、bias_add/add、batch_norm、relu 的组合以及 depthwise 卷积mrvl.fc_ni2no全连接nn.dense bias activation可带 flatten/reshape/layout_transformmrvl.maxpool2d_nhwc2nhwc/mrvl.avgpool2d_nhwc2nhwc/mrvl.globalavgpool2d_nhwc2nhwc各类池化mrvl.sum逐元素加法mrvl.concat拼接仅支持 4 维 NHWC、batch size 为 1、axis3mrvl.layout_transform_nchw2nhwc、mrvl.reshape、mrvl.batch_flatten布局变换与形状变换。同时每个算子都注册了target.mrvl属性函数如conv2d_nhwc2nhwc、fc_ni2no用于在注解阶段校验算子是否真的可被 mrvl 支持——例如卷积要求 NHWC 布局、float32 输入输出、batch size ≤ 8、grouped conv 仅支持groups channels的 depthwise 情形。这些校验逻辑集中体现在 mrvl.py。6.2 JSON 序列化与 mrvl-tmlc 调用编译期由 codegen.cc 中的MrvlJSONSerializer将 Relay 复合函数序列化为 JSON 图含常量、shape、dtype、min/max、base64 权重数据并针对每种复合模式Conv2D、FC、Pool、Sum、Concat、Reshape 等生成对应 JSON 节点。随后 python/tvm/contrib/mrvl.py 的tvm.mrvl.CompileModel会把 nodes.json 与 consts.json 写入工作目录并调用mrvl-tmlc编译器mrvl-tmlc -mn symbol -f1 nodes.json -f2 consts.json compiler_opts -b batch_size编译成功后读取bin_symbol/symbol.bin并做 base64 编码嵌入生成的 TVM 模块若mrvl-tmlc不在$PATH中会抛出明确错误。环境变量MRVL_SAVE_MODEL_BIN可用于保留编译产生的中间二进制便于调试。6.3 运行时模块与模拟器模拟器运行时模块mrvl_runtime.cc实现了二进制可序列化与可运行属性通过runtime.mrvl_runtime_create注册并以runtime.module.loadbinary_mrvl_sim支持二进制加载。子图的输入/输出数量从 nodes.json 中的num_subgraph_inputs/num_subgraph_outputs读取由 mrvl.py 写入。6.4 测试用例参考仓库在 tests/python/contrib/test_mrvl/test_mrvl.py 提供了完整的验证用例均以requires_mrvl标记需要 Marvell 工具链环境test_mrvl_fuse验证 conv2dbias/batch_norm/relu 融合分区、sum 模式分区以及 MobileNet 整网分区数量test_conv2d验证mrvl.conv2d_nhwc2nhwc的代码生成与模拟器运行option_dict {num_tiles: 1}test_dense验证mrvl.fc_ni2no的代码生成与模拟器运行。这些测试是理解哪些图结构可被分区、分区后 JSON 长什么样的最佳参考样例。7. 注意事项与当前限制硬件运行时为 WIPmrvl 对 MLIP 硬件的原生运行时支持尚未合入当前代码实际可验证路径为 Marvell ML 模拟器mrvl-mlsimint8 量化为 WIPquantizeint8的完整支持待未来 PR 提供当前默认且推荐使用fp16batch size 限制mrvl 后端要求 batch size ≤ 8mrvl.py 中mrvl-tmlc编译时会做硬性检查且多个算子模式仅支持 batch size 为 1工具链依赖编译与模拟运行都要求 Marvell 工具mrvl-tmlc、mrvl-mlsim位于$PATH推荐直接使用仓库提供的 Dockerfile.demo_mrvl 镜像环境布局要求mrvl 区域内部使用 NHWC / OHWI 布局分区流水线会自动完成 NCHW→NHWC 的转换并在主图回切 NCHW无需用户手工干预。8. 总结通过 mrvl 目标TVM 将 Marvell MLIP 深度集成进统一的编译与部署流程--targetmrvl, llvm复合目标让受支持算子自动卸载到 Marvell 加速器、其余算子回退 LLVMmcpu、num_tiles、mattrquantize、wb_pin_ocm四个编译选项覆盖了处理器型号、算力规模与精度/内存优化策略Python API 侧则以partition_for_mrvlrelay.buildgraph_executor的标准三件套完成端到端推理。对于当前无法接触 MLIP 硬件的开发者模拟器路径mrvl-mlsim提供了完整的功能验证手段配合 test_mrvl.py 可快速确认自己的模型是否能够被 mrvl 后端正确分区与编译。【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址: https://gitcode.com/gh_mirrors/tvm7/tvm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考