DietCode：TVM中新的动态Shape解决方案

‍

‍近年来，随着深度学习领域动态 shape 推理场景激增，如机器翻译、语音识别等，针对动态 shape 模型的优化成为了热门研究议题。本文主要介绍了 DietCode 的整体架构设计，尝试了解 DietCode 如何针对动态 shape 进行调优。以 BERT 模型为例，DietCode 与 Ansor 相比， 搜索时间减少 5.88倍，性能提高约 70%。DietCode项目已经开始合入TVM中，但目前来看，只是在 tir 优化 Pass 中增加了 LocalPad，其他 feature 还未合入。

1 背景

在机器学习领域，提高计算密集型算子的性能是一项极具挑战的工作，目前业界主要有两种解决方案：使用供应商的算子库或使用 auto-scheduler 来完成算子调优。然而，前者需要大量的工程工作，算子开发周期长，限制多且专业性强；后者在 TVM 中只支持输入为静态形状的工作负载，以便模型编译时获得形状大小，如果输入为动态形状，将会导致极长的调优时间，这是因为整个搜索空间与形状强相关，auto-scheduler 不能为所有可能形状的工作负载构造一个确定的搜索空间。为了解决以上问题，DietCode 应运而生。

目前，TVM中 auto-scheduler（如：Ansor）的整体流程，如图1（a）所示，详细信息请参考Ansor: 为深度学习生成高性能张量程序。它的的输入为 TE&#x