Co-exploration of Hardware Acceleration and Tensor Programs via Unified Compilation Interface

Domain-specific hardware accelerators hold great promise for achieving significant enhancements in both speed and energy efficiency compared to general-purpose processors. Nevertheless, the design and development process of these accelerators often demands a substantial amount of manual effort, particularly in the realms of hardware design and the development of high-performance tensor programs. Extensive research has been dedicated to optimizing these aspects individually, which has resulted in a proliferation of possibilities, posing considerable challenges. To tackle this challenge, we propose BACO, an end-to-end one-loop search framework. The aim of BACO is to provide a unified compilation interface that facilitates the co-exploration of hardware and software. The key element of BACO is its multi-level abstraction, which not only precisely defines the behavior of custom hardware instructions and configuration state, but also enables fine-grained optimization of computation and data movement through the concept of a block. Building upon this abstraction, we develop a programming language and compiler that incorporates differentiable performance models to automatically explore various high-performance design points. Additionally, BACO employs an efficient strategy for estimating importance, allowing for simultaneous optimization of multiple layers to enhance end-to-end performance. This approach facilitates rapid co-exploration while reducing development efforts. Experimental results demonstrate that BACO surpasses existing state-of-the-art search solutions and manually tuned versions by a substantial margin. Moreover, we plan to open source the whole codebase.