GEMM Tiling Configuration Search Algorithm based on Deep Reinforcement Learning
Zhenshan Bao, Yiming Li, Wenbo Zhang · 2024
Generalized matrix multiplication (GEMM) has become the fundamental operator in deep learning. Operators of deep learning need to be transformed and optimized into executable code for specific hardware. However, existing deep learning frameworks rely on handwritten libraries that are not scalable for various hardware target. In this paper we introduce a compiler-level tuning algorithm based on Neighborhood Actor Advantage Critic (N-A2C) method for GEMM tiling optimization. Experiment results shows that our method achieves 8% and 20% computational time savings of GEMM compared Xgboost and RNN methods when exploring 1% of the search space.