Dynamism is common in AI computation, e.g., the dynamic tensor shapes and the dynamic control flows in models. Due to the long compilation time, existing runtime compilation damages the model efficiency, while the offline compilers either suffer from the long compilation time and device memory footprint to cover all the possible execution instances of a dynamic model, or sacrifice optimization opportunities for usability. In this paper, we rethink the feasibility of runtime compilation for dynamic models and identify that the key for it to work is to speed up the compilation or hide the compilation overhead. To do this, we propose a real-time compiler, DVM. In DVM, we design a runtime operator compiler based on a bytecode virtual machine to perform effective and efficient compilation for each dynamic operator instance given its input. Specifically, instead of compiling programs into machine code, we encode the operator program into bytecode on the CPU and decode the bytecode into virtual instructions for direct execution on the NPU. Based on the runtime operator compiler, we further propose an operator fuser, which performs symbol-deduction-based fusion on static graphs and runtime fusion on dynamic graphs. Both pattern- and stacking-based fusion are supported to increase fusion opportunities. Evaluation on operators, subgraphs, and models shows that, compared with TorchInductor, PyTorch-eager and MindSpore-graph-O0, we are up to 11.77$\times$ better in terms of the operator/model efficiency and up to 5 orders of magnitude faster in terms of the maximum compilation time.
翻译:动态性是AI计算中的常见现象,例如模型中的动态张量形状和动态控制流。由于编译时间较长,现有的运行时编译会损害模型效率,而离线编译器要么为覆盖动态模型所有可能的执行实例而承受较长的编译时间和设备内存开销,要么为了可用性而牺牲优化机会。本文重新思考了动态模型运行时编译的可行性,并指出其成功的关键在于加速编译或隐藏编译开销。为此,我们提出一种实时编译器DVM。在DVM中,我们设计了一种基于字节码虚拟机的运行时算子编译器,能够针对每个动态算子实例的输入执行高效编译。具体而言,我们不将程序编译为机器码,而是在CPU上对算子程序进行字节码编码,并在NPU上将字节码解码为虚拟指令直接执行。基于该运行时算子编译器,我们进一步提出了一种算子融合器,该融合器在静态图上执行基于符号推导的融合,在动态图上执行运行时融合。同时支持基于模式融合和基于堆叠融合,以增加融合机会。在算子、子图和模型上的评估表明,与TorchInductor、PyTorch-eager和MindSpore-graph-O0相比,我们的算子/模型效率最高提升11.77倍,最大编译时间最快提升5个数量级。