UniCoMo: A Unified Learning-Based Cost Model for Tensorized Program Tuning
Zihan Wang, Lei Gong, Wenqi Lou, Qianyu Cheng, Xianglan Chen, Chao Wang, Xuehai Zhou · 2024
Tensorized programs use hardware intrinsics on accelerators to significantly improve tensor computation performance. The trend of hardware customization has led to the emergence of massive hardware accelerators and intrinsics, which poses a significant engineering challenge for manually coding efficient tensorized programs. Tensorized program tuning in deep learning compilers (DLCs) is considered an effective approach to address the challenge. At the core of program tuning relies the design of the cost model, but currently there is still a lack of cost models specifically designed for tensorized programs, which hampers the co-optimization of DLCs and hardware accelerators. In this paper, we propose UniCoMo, a unified cost model for tensorized program tuning on various hardware platforms and hardware intrinsics. We first analyze the design challenges introduced by tensorized programs for cost models in terms of feature representation and transfer prediction. And then, for feature representation, we propose a unified feature representation for tensorized programs by using program behavior as a template, mining program features with the schedule attention matrix, and incorporating hardware intrinsic abstraction. For transfer prediction, we propose a unified transfer prediction strategy for tensorized program cost model based on lifelong learning and transfer learning. To meet training and testing requirements, we constructed a dataset dedicated to tensorized program tuning. Results show that UniCoMo maintains the state-of-the-art accuracy while significantly improving adaptability to diverse execution environments and enabling flexible transfer prediction. It can speed up search time by 9.8× and improve inference speed by 1.9× open-sourced at https://github.com/ZhW-loop/UniCoMo.