Abstract: ML libraries, often written in architecture-specific programming languages (ASPLs) that target domain-specific architectures, are key to efficient ML systems. However, writing these high-performance ML libraries is challenging because it requires expert knowledge of both ML algorithms and the ASPL. Large language models (LLMs), on the other hand, have shown general coding capabilities. However, challenges remain when using LLMs for generating ML libraries using ASPLs because 1) this task is complicated even for human experts and 2) there are limited code examples due to the esoteric and evolving nature of ASPLs. We present an adaptive self-improvement agentic system that enables LLMs to perform such complex reasoning under limited data by iteratively improving their capability through self-generated experience. In order to evaluate the effectiveness of our system, we construct a benchmark of a typical ML library and generate ASPL code with both open and closed-source LLMs on this benchmark. Our results show improvements of up to $3.9\times$ over a baseline single LLM.

Lay Summary:

This paper presents a new system that helps artificial intelligence (AI) models improve themselves over time, specifically for writing the code needed to run machine learning (ML) programs on specialized hardware. Creating this code is usually a difficult and time-consuming task, even for experts, because it requires deep knowledge of both ML techniques and the hardware-specific programming languages. The authors show how large language models can be turned into a team of AI agents that learn from their own past successes and mistakes. This self-improvement process allows them to write better code without needing thousands of examples. They tested their system on a new, cutting-edge programming language and showed it could solve nearly all the critical tasks and be up to four times better than using one AI model alone. This approach could make it much easier and faster to build high-performance ML systems in the future.