Tensor language model enables generative scheduling for efficient tensor compilation.

Area of Science:

• Artificial Intelligence
• Compiler Design
• High-Performance Computing

Background:

• Deep learning tasks and heterogeneous computing demand compilers with low compile times and high performance.
• Existing tensor compilers rely on slow exhaustive search or heuristic methods that compromise generality and optimization quality.

Purpose of the Study:

• To introduce the Tensor Language Model (TLM), a novel generative framework for optimizing tensor programs.
• To redefine tensor program optimization as a language modeling problem, leveraging a GPT-2 architecture.

Main Methods:

• TLM utilizes a GPT-2 architecture pre-trained on millions of tensor programs represented as compact tensor code.
• The model treats optimization as a sequence-to-sequence problem, generating optimized tensor schedules directly.
• It avoids runtime search or reinforcement learning, enabling faster compilation.

Main Results:

• TLM achieves compilation speeds up to 61 times faster than search-based compilers (e.g., Ansor, MetaSchedule).
• It outperforms heuristic models (e.g., Roller) by up to 2.25 times in speed, maintaining similar runtime efficiency.
• Experiments on ResNet-50, BERT, GPT-2, and LLAMA-7B demonstrate TLM's effectiveness.

Conclusions:

• TLM presents a scalable, hardware-agnostic, and reproducible generative paradigm for next-generation deep learning compilers.
• It offers a compelling trade-off between compilation time and execution performance.