Researchers at the Massachusetts Institute of Technology's Computer Science and Artificial Intelligence Laboratory have created a new programming language called Exo 2.
Researchers at the Massachusetts Institute of Technology's Computer Science and Artificial Intelligence Laboratory have created a new programming language called Exo 2.
Exo 2 belongs to a new category of programming languages that MIT professor Jonathan Ragan-Kelly calls “user-schedulable languages” (USLs). Instead of hoping that an opaque compiler will automatically generate the fastest code, USLs put programmers in the driver’s seat, allowing them to write “schedules” that explicitly control how the compiler generates the code. This allows performance engineers to turn simple programs that specify what they want to compute into complex programs that do the same thing as the original specification, but much, much faster, writes Tech Xplore.
One of the limitations of existing USLs (such as the original Exo) is their relatively fixed set of scheduling operations, which makes it difficult to reuse scheduling code across different "cores" (individual components of a high-performance library). Exo 2, on the other hand, allows users to define new scheduling operations outside the compiler, which makes it easier to create reusable scheduling libraries.
According to the researchers, Exo 2 can reduce overall scheduling code by a factor of 100 and deliver performance that rivals state-of-the-art implementations on many different platforms, including the basic linear algebra routines (BLAS) that power many machine learning applications. This makes it an attractive option for HPC engineers focused on optimizing cores for different operations, data types, and target architectures.
“This is a bottom-up approach to automation, rather than ML/AI-powered search in high-performance code,” says researcher Yuka Ikarashi. “This means that performance engineers and hardware developers can write their own scheduling library, which is a set of optimization methods for their hardware to achieve maximum performance.”
One of the main advantages of Exo 2 is that it reduces the amount of coding effort required at any given time by reusing planning code for different applications and hardware targets.
The researchers implemented a scheduling library of approximately 2000 lines of code in Exo 2 that encapsulates multiple optimizations specific to linear algebra and specific targets (AVX512, AVX2, Neon, and Gemmini hardware accelerators). This library unifies the scheduling efforts across over 80 high-performance cores, each containing up to a dozen lines of code, providing performance comparable to or better than MKL, OpenBLAS, BLIS, and Halide.
Exo 2 includes a new mechanism called Cursors, which provide what the researchers call a “stable reference” to point to object code throughout the scheduling process. Ikarashi says that a stable reference is important for users to encapsulate schedules in a library function because it makes the scheduling code independent of object code transformations.
“We believe that USLs should be designed to be extensible, rather than having a fixed set of operations,” says Ikarashi. “This way, the language can evolve and support large projects by implementing libraries that meet a variety of optimization requirements and application areas.”
The design of Exo 2 allows performance engineers to focus on high-level optimization strategies while ensuring that the underlying object code remains functionally equivalent through the use of safe primitives. In the future, the team hopes to extend Exo 2's support for different types of hardware accelerators, such as GPUs. Several ongoing projects are aimed at improving the compiler's analysis itself in terms of correctness, compile time, and expressiveness.
