Hierarchical Mixture of Experts: Generalizable Learning for High-Level Synthesis

TL;DR

The paper tackles domain generalization in ML-driven HLS pragma design for FPGA performance prediction. It introduces a hierarchical Mixture of Experts with three low-level MoEs operating on node, block, and graph granularities and a high-level MoE to aggregate them, together with a two-stage training strategy to prevent expert polarization. The approach yields improved predictive generalization and practical FPGA speedups on the HLSyn benchmark, with extensive ablations validating the contribution of each component. Code for the method is publicly released to enable reproducibility and further research.

Abstract

High-level synthesis (HLS) is a widely used tool in designing Field Programmable Gate Array (FPGA). HLS enables FPGA design with software programming languages by compiling the source code into an FPGA circuit. The source code includes a program (called "kernel") and several pragmas that instruct hardware synthesis, such as parallelization, pipeline, etc. While it is relatively easy for software developers to design the program, it heavily relies on hardware knowledge to design the pragmas, posing a big challenge for software developers. Recently, different machine learning algorithms, such as GNNs, have been proposed to automate the pragma design via performance prediction. However, when applying the trained model on new kernels, the significant domain shift often leads to unsatisfactory performance. We propose a more domain-generalizable model structure: a two-level hierarchical Mixture of Experts (MoE), that can be flexibly adapted to any GNN model. Different expert networks can learn to deal with different regions in the representation space, and they can utilize similar patterns between the old kernels and new kernels. In the low-level MoE, we apply MoE on three natural granularities of a program: node, basic block, and graph. The high-level MoE learns to aggregate the three granularities for the final decision. To train the hierarchical MoE stably, we further propose a two-stage training method to avoid expert polarization. Extensive experiments verify the effectiveness of the proposed hierarchical MoE. We publicized our codes at https://github.com/weikai-li/HierarchicalMoE.

Figures (6)

• Figure 1: Motivation of utilizing MoE. The two kernels share a similar nested loop, while the other parts are different. Similar parts can utilize similar experts via similar gating, while different parts can utilize different experts. This is just an illustration. Actually, we do not apply MoE on different granularities in the same model in this way. We aggregate them like Figure \ref{['fig:model']}.
• Figure 2: Illustration of hierarchical MoE. The three low-level models are aggregated via the high-level gating network.
• Figure 3: Source code of the "syr2k" kernel. The source code consists of the program that describes the FPGA's function, and several pragmas inserted in the program that instruct the hardware synthesis process, such as parallelization, pipeline, etc. "auto{pragma index}" are pragma's placeholders. We utilize the machine learning model to find the best pragmas to insert in these placeholders.
• Figure 4: Visualization of the graph embeddings. Each point is a pragma design. Points belonging to the same kernel have the same color. Points from the same kernel are clustered together, while points from different kernels are located in different places. Due to the limited number of colors, different kernels might use the same color. Due to the limited space, we only show the names of the target kernels.
• Figure 5: Code snippet of the "syrk" kernel. The line inside this nested loop is the only difference between the "syr2k" kernel and the "syrk" kernel. However, their best pragma designs are very different.