Bring Your Own Codegen to TVM

• 语言: 英语
• 格式: .pdf
• 来源: github.com/tvmai

摘要

文档介绍了如何将自定义代码生成器（codegen）集成到TVM中，以支持特定的操作或子图。主要内容包括：1）标记支持的操作或子图，可通过操作级注解或图级注解实现；2）实现自定义代码生成器的构建逻辑，包括生成二进制文件或库；3）在运行时通过分发生成的可执行引擎来处理子图。文档还展示了如何通过示例（如Intel MKL-DNN库）实现这些功能，并提供了系统概述和实现步骤。

AI总结

# 《Bring Your Own Codegen to TVM》

## 核心观点与关键信息

1. **背景与挑战**
- 设计并制造深度学习芯片时，芯片可能仅支持常用操作（如 conv2d、dense、ReLU 等），但对于复杂模型（如 YOLO），可能需要支持不常见的操作（如 Non Maximum Suppression，NMS）。
- TVM 提供了灵活的解决方案，支持通过自定义代码生成器（BYOC，Bring Your Own Codegen）来扩展功能。

2. **自定义代码生成器的实现**
- **实现构建逻辑**：
- 编写生成二进制或库文件的逻辑。
- 将生成的可执行引擎在运行时分发到目标设备。
- **实现运行时分发器**：
- 编写 TVM 运行时模块，将子图分发到自定义引擎。
- 运行时路径：`src/runtime/contrib//.{h, cc}`。

3. **标记支持的操作或子图**
- **操作级注解**：
- 通过 Python 模板判断操作是否支持（返回布尔值）。
- 示例：`def conv2d(attrs, args): return is_float32(args)`。
- **图级注解**：
- 实现一个图注解器，标记整个图中可以使用自定义代码生成器的部分。

4. **案例展示**
- **Intel MKL-DNN 示例**：

     import numpy as np
     from tvm import relay
     mod, params = relay.testing.mobilenet.get_workload(batch_size=1)
     mod = relay.build_extern(mod, "dnnl")
     exe = relay.create_executor("vm", mod=mod, ctx=tvm.cpu(0))
     data = np.random.uniform(size=(1, 3, 224, 224)).astype("float32")
     out = exe.evaluate()(data, **params)

5. **总结**
- 通过自定义代码生成器，用户可以扩展 TVM 的功能，支持特定的操作或子图，从而优化深度学习芯片的性能。
- TVM 提供了灵活的接口和清晰的实现路径，帮助用户快速集成自定义代码生成器。