一句话：提出动态形状张量编译器 MikPoly，用运行时形状驱动的微内核动态聚合与轻量代价模型实现高效优化。
动态形状神经网络中算子的形状在运行时才确定，导致静态形状张量编译器需反复编译，编译开销高昂且难以实用。本文提出 MikPoly——一种基于微内核聚合（micro-kernel polymerization）的动态形状张量编译器。它采用两阶段优化：先静态生成多个微内核，再在运行时根据已知的张量算子形状，通过轻量级代价模型动态地将这些微内核组合成适配当前形状的内核，从而兼顾优化效果与编译效率。作者在 GPU Tensor Core 和昇腾 NPU 两类代表性加速器上，使用常见的动态形状算子与神经网络进行评估，结果显示 MikPoly 能有效优化动态形状负载，相比当前最先进的厂商库平均取得 1.49 倍的性能提升。
In recent times, dynamic-shape neural networks have gained widespread usage in intelligent applications to address complex tasks, introducing challenges in optimizing tensor programs due to their dynamic nature. As the operators' shapes are determined at runtime in dynamic scenarios, the compilation process becomes expensive, limiting the practicality of existing static-shape tensor compilers. To address the need for effective and efficient optimization of dynamic-shape neural networks, this paper introduces MikPoly, a novel dynamic-shape tensor compiler based on micro-kernel polymerization. MikPoly employs a two-stage optimization approach, dynamically combining multiple statically generated micro-kernels using a lightweight cost model based on the shape of a tensor operator known at runtime. We evaluate the effectiveness of MikPoly by employing popular dynamic-shape operators and neural networks on two representative accelerators, namely GPU Tensor Cores and Ascend NPUs. Our experimental results demonstrate that MikPoly effectively optimizes dynamic-shape workloads, yielding an average performance improvement of 1.49× over state-of-the-art vendor libraries.
@inproceedings{YuLZCFX24,
title = {Optimizing Dynamic-Shape Neural Networks on Accelerators via On-the-Fly Micro-Kernel Polymerization},
author = {Feng Yu and Guangli Li and Jiacheng Zhao and Huimin Cui and Xiaobing Feng and Jingling Xue},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620665.3640390}
}