TVM项目中对PyTorch ExportedProgram的支持与实现

背景介绍

TVM作为一个深度学习编译器，需要支持多种前端框架的模型导入。PyTorch作为当前最流行的深度学习框架之一，其模型导入功能一直是TVM开发的重点。随着PyTorch 2.0引入了新的导出机制torch.export.ExportedProgram，TVM项目也相应地进行了适配和升级。

技术挑战

PyTorch的torch.export.ExportedProgram是一种新的中间表示(IR)，与之前的FX Graph相比有显著不同。TVM需要解决以下几个技术挑战：

1. IR适配：需要理解并适配新的IR结构和操作语义
2. 操作符支持：需要实现对新IR中各种操作符的转换支持
3. 基础设施升级：需要升级PyTorch版本和相关CI环境
4. 代码重构：需要优化现有代码结构，提高可维护性

实现方案

TVM团队采取了分阶段的实现策略：

第一阶段：代码清理与重构

首先对现有的TorchFX导入器进行了全面的代码清理和重构工作，包括：

• 统一一元操作符转换器
• 重构二元操作符转换器
• 优化神经网络相关操作符转换
• 整理统计和搜索操作符转换
• 改进张量操作和创建相关转换器
• 完善数据类型转换处理

这一阶段将原有近2000行的代码进行了模块化和标准化，为后续工作奠定了基础。

第二阶段：基础设施升级

为了支持新的ExportedProgram，需要升级PyTorch版本和相关构建环境：

• 升级PyTorch到支持新特性的版本
• 更新CI测试镜像以确保兼容性
• 验证新版本在现有功能上的稳定性

第三阶段：新导入器实现

在完成基础工作后，实现了专门的ExportedProgramImporter，包括：

• 设计新的导入器架构
• 逐步扩展支持的操作符范围：
  • 一元操作符
  • 二元操作符
  • 神经网络操作符
  • 统计操作符
  • 搜索操作符
  • 张量操作相关功能
  • 张量创建相关功能

第四阶段：文档完善

最后完成了相关技术文档，包括：

• 新特性的使用说明
• 实现细节的技术文档
• 示例代码和最佳实践

技术细节

新的导入器实现采用了更模块化的设计，将通用逻辑与特定IR处理逻辑分离。统计显示，重构后的代码虽然总行数有所增加（从1947行增至2342行），但代码结构更加清晰，注释也更加完善（从48条增至113条）。

总结

TVM对PyTorch ExportedProgram的支持是一个系统工程，涉及从底层操作符到高层架构的多方面改进。通过分阶段的实施策略，团队成功地在保持现有功能稳定的同时，增加了对新特性的支持。这一工作不仅提升了TVM与PyTorch的兼容性，也为未来支持更多新特性奠定了良好的基础。