一句话：提出两阶段算子到算子编译流程的统一 MLIR 前端，兼容多框架模型并提升编译与执行效率。
MLIR 前端（如 IREE-TF、Torch-MLIR、ONNX-MLIR）负责连接 ML 框架与底层硬件，但存在高层 IR 转换不灵活、高层优化机会受限、兼容性与效率下降等问题，造成软件碎片化。UFront 提出统一的 MLIR 前端，采用两阶段的算子到算子编译工作流：第一阶段瞬时追踪模型，将追踪到的计算节点委托为标准 DNN 算子，在不依赖 MLIR pass 的情况下把不同框架的模型转换为统一高层 IR，并执行常量折叠、算子融合等图优化；第二阶段利用所提出的 lowering pattern 将高层 IR 直接转换为标准 TOSA IR，消除变换冗余并保持与现有编译后端的兼容。实验表明，UFront 在多种框架的流行 DNN 模型上具有更高兼容性、更快的端到端编译速度，并能生成执行更高效的二进制。
Automatic code generation for ML systems has gained popularity with the advent of compiler techniques like Multi-Level Intermediate Representation (Multi-Level IR, or MLIR). State-of-the-art MLIR frontends, including IREE-TF, Torch-MLIR, and ONNX-MLIR, aim to bridge the gap between ML frameworks and low-level hardware architectures through MLIR's progressive lowering pipeline. However, existing MLIR frontends encounter challenges such as inflexible high-level IR conversion, limited higher-level optimization opportunities, and reduced compatibility and efficiency, leading to software fragmentation and restricting their practical applications within the MLIR ecosystem. To address these challenges, we introduce UFront, a unified MLIR frontend employing a two-stage operator-to-operator compilation workflow. Unlike traditional frontends that compile model source code into binaries step by step with different MLIR transform passes, UFront decouples the process into two distinct stages. It first performs instantaneous model tracing, delegates traced computing nodes as standard Deep Neural Network (DNN) operators and transforms models written in different frameworks into unified high-level IR without relying on MLIR passes, enhancing conversion flexibility. Meanwhile, it performs high-level graph optimizations such as constant folding and operator fusion to produce more efficient high-level IR. In the second stage, UFront directly converts high-level IR into standard TOSA IR using proposed lowering patterns, eliminating transform redundancies and ensuring lower-level compatibility with existing ML compiler backends. This two-stage compilation approach enables consistent end-to-end code generation and optimization of various DNN models written in different formats within a single workflow. Extensive experiments on popular DNN models written in various frameworks demonstrate that UFront exhibits higher compatibility, faster end-to-end compilation, and is capable of producing more efficient binary execution compared to SOTA works.
DOI 原文 ·
@inproceedings{paperbot8345,
title = {UFront: Toward A Unified MLIR Frontend for Deep Learning},
author = {Guoqing Bao and Heng Shi and Chengyi Cui and Yalin Zhang and Jianguo Yao},
booktitle = {ASE},
year = {2024},
doi = {10.1145/3691620.3695002}
}