Bring Your Own Codegen to TVM

《Bring Your Own Codegen to TVM》文档总结如下： ### 核心内容概述 本文档介绍了如何将自定义代码生成（Codegen）集成到 Apache TVM 中，特别是在运行时实现子图的分割与分发执行。该方法适用于设计深度学习芯片的场景，帮助芯片厂商更好地支持复杂模型，同时提升性能。 ### 问题背景 - 深度学习芯片通常针对常用操作器（如 Conv2D、Dense、ReLU 等）进行优化，但难以支持新操作器（如 Non Maximum Suppression, NMS）。 - TVM 支持部分芯片不支持的操作器，如何充分利用 TVM 的能力是关键。 ### 解决方案概述 - **实现构建逻辑**：自定义代码生成工具，生成目标设备的二进制文件或库。 - **实现运行时分发器**：在 TVM 运行时分发子图到自定义代码生成的执行引擎中。 - **运行时路径**：在 `src/runtime/contrib/` 目录下实现外部函数调用逻辑。 ### 工作流程 1. **标注子图**：通过实现操作器级或图级标注器，标记支持的操作器或子图。 - **操作器级标注**：在 `python/tvm/relay/op/contrib/` 下实现模板，判断操作器是否支持。 - **图级标注**：直接标记支持的子图。 2. **划分子图**：系统根据标注结果划分 Relay IR 图。 3. **编译子图**：自定义代码生成工具生成可执行文件。 ### 案例展示 - 使用 Intel MKL-DNN（DNNL）库作为示例，展示如何将 MobileNet 模型划分并通过自定义代码生成工具构建。 - 代码示例： ```python mod = relay.build_extern(mod, "dnnl") exe = relay.create_executor("vm", mod=mod, ctx=tvm.cpu(0)) ``` ### 实现选项 1. **操作器级标注**：通过布尔函数判断操作器支持。 2. **图级标注**：直接标记子图。 ### 系统概览 - **Relay IR 图**：通过标注器划分为支持和不支持的子图。 - **运行时调度**：通过自定义分发器将子图分发到目标设备或代码生成工具。 - **执行流程**：启动 Relay 运行时（如 VM 或图运行时），调用外部函数执行子图。 ### 总结 - 提供了一种灵活的方法，允许硬件厂商利用 TVM 的强大功能，同时支持自定义优化。 - 通过分离构建和运行时逻辑，提升模型性能和硬件利用率。 ### 相关资源 - 系统原型：[GitHub Pull Request](https://github.com/apache/incubator-tvm/pull/4258) - 请求评论（RFC）：[TVM 讨论区](https://discuss.tvm.ai/t/bring-your-own-codegen-to-tvm/4501)