A Heterogeneous CNN Compilation Framework for RISC-V CPU and NPU Integration Based on ONNX-MLIR

With the continuous advancement of convolutional neural networks (CNNs), many neural network processing units (NPU) have emerged in recent years. NPUs offer advantages such as improved energy efficiency and low latency compared to traditional processors. However, NPUs often struggle to keep pace with the growing complexity of algorithmic models, which limits the implementation of certain AI applications. Central processing units (CPU) are known for their versatility, but are computationally inefficient. Combining the strengths of both architectures could present a viable solution to these challenges. Despite this potential, there is currently insufficient academic focus on CPU/NPU co-computing, and few architectures or toolchains have been developed to support such integration. In this paper, we propose a RISC-V CPU/NPU co-computing-based compilation framework to solve the problem of mapping from algorithms to heterogeneous architectures. To bridge the computational gap between heterogeneous platforms, within ONNX-MLIR, we developed three phases to support operator packing, algorithm adaption, and data coordination. To integrate heterogeneous compilation platforms, we introduce four additional functions to handle task scheduling, weight quantization, weight rearrangement, and memory management. In particular: (1) to optimize the transition overhead, an operator packing technique is introduced by extending the ONNX-MLIR framework with custom operators, which enhances NPU efficiency by reducing the number of transitions between operations. (2) At the aspect of task scheduling, a three-mode task scheduler is developed to allow users to customize task distribution for correctness verification, performance optimization, or detailed analysis on heterogeneous platforms, which demonstrates the versatility of our framework. (3) In terms of memory management, an optimized scheme for CPU/NPU architectures is proposed to ensure accurate data access. This memory management employs a dual-end approach with self-checks and a release-after-use mechanism to further improve efficiency. Experimental results demonstrate that the framework effectively generates instructions for various CNNs, significantly enhancing the efficiency and performance of heterogeneous architectures. It achieves up to a 7.06× reduction in code density and up to a 5.58× improvement in memory usage, narrowing the computational gap between CPUs and NPUs.