For many years, optimisation of compilers went through a predictable evolution consisting of studying hardware, coming up with optimisation guidelines, implementing them in compilers, and improving those guidelines over the years.AI workloads are changing this pattern.

Modern neural networks use more and more sophisticated CPUs, GPUs, TPUs and other hardware accelerators. The process of optimisation depends on many criteria, such as tensor shape, memory arrangement, kernel fusion, cache efficiency, register pressure, as well as hardware-specific instructions. Learning the rule for every combination is becoming more and more challenging.

The industry is responding to this challenge via a quiet race: the implementation of machine learning into the compiler.

Traditional Compiler Heuristics Are Losing Their Effectiveness

Despite the power of conventional compilers, they still rely on expertly designed heuristics to make certain decisions.

A compiler may have to decide:

  • to inline a function,

  • to combine operations,

  • to modify layouts in memory,

  • to select a particular kernel,

  • to determine the sizes of tiles,

  • to reorder loops,

  • to change the way of allocating registers, or

  • to choose an optimisation method.

Traditionally, it has been up to engineers to define rules for making such decisions.

The problem is that modern AI workloads have introduced huge optimisation spaces; thus, the configuration that works for one model or accelerator could fail badly in the case of another one.

This led scientists to suggest that machine learning might find a way of dealing with the situation by establishing correlations between program parameters and performance results.

Thus, machine learning can learn the best paths to take, rather than being instructed by engineers about every possible situation.

The Compiler Develops into an Intelligent System

Primarily, the principle is quite elementary to understand.

To begin with, the compiler should analyse the source code of the program and get some information regarding the flow of control, loops, operations, memory access and dimensions of tensors.

The next thing that needs to be done is to apply an ML model in order to predict optimisation choices.

Then, the compiler compiles those options and sees how they work. The feedback received will be used to make better decisions in the future.

The whole process is:

Program → Optimization → Measurement → Learning Process → Better Optimization

This is in contrast to a compiler that operates according to the same principles.

MLGO – Google’s machine learning organisation that helps LLVM

A good example is Machine Learning Guided Optimisation, or MLGO, from Google.

The MLGO project organises machine learning techniques for use with a real operating compiler. First, the team experimented with function call inlining and word allocation, both problems that were previously attempted by applying traditional algorithms.

Teaching LLVM When to Inline

Inlining means performing a function call by inserting the function’s own code directly in the program.

This allows eliminating the overhead associated with a function call and getting more chances to apply additional optimisations. However, too much inlining increases the size of the program.

The question remains:

When should a compiler call an inlining function?

MLGO does not depend solely on a set of predetermined rules, but it allows for new policies to be learned by the compiler. According to Google, its ML-based system for inlining for size resulted in a code-size decrease of 7% in comparison to LLVM’s -Oz optimiser. In addition, they found that trained models can generalise to various targets and changing software environments.

This is a significant advancement as the compiler, instead of just applying pre-defined rules, is able to learn a policy.

Professional Perspective: The MLGO Research Team

According to the researchers at Google who developed MLGO, it is always harder to maintain complex handcrafted compiler heuristics. Machine learning can find optimisation opportunities that cannot be easily expressed manually.

The importance of MLGO is that this method was developed for the industrial compiler environment and wasn’t just a laboratory project.

The next important area is autotuning.

Unlike predictive analytics, which stays focused on one solution for the compiler, an autotune expert tests out different available solutions and compares their efficiency.

This is especially applicable to a number of AI fields.

Take, for instance, the case of matrix multiplication. For the compiler to make decisions, it might have to check on the following parameters:

  • Tile dimensions

  • Thread configuration

  • Memory layout

  • Loop ordering

  • Shared-memory usage

  • Vectorization

  • Kernel fusion

There is rarely one perfect solution. The best choice might depend on the type of accelerator being used, tensor dimensions and the type of workload being processed.

Machine learning can help find the answer in a big search space.

TVM and Ansor Turn Optimization Into Search

Apache TVM has developed into a major research platform for the machine learning compilation arena. The ecosystem includes automated tuning and scheduling technologies such as AutoTVM and Ansor.

Ansor produces tensor programs and provides for the search of the optimisation space. It performs searches in the area, using techniques such as evolutionary search alongside a learned cost model to find suitable candidates.

The researchers have shown improvements of up to 3.8× on Intel CPUs, 2.6× on ARM CPUs, and 1.7× on NVIDIA GPUs for tasks tested relative to existing methods at that time.

The important lesson is that when optimisation space is too large to be explored by humans, machine-guided search can perform effectively instead.

XLA from Google Demonstrates Efficiency in Production

XLA refers to Accelerated Linear Algebra. It serves as just one of the notable examples of an important technology.

XLA serves as a tool for compiling and optimising ML workloads for TPUs, CPUs, and GPUs.

Optimisation comes with its problems since optimisation tasks often interfere with each other.

Therefore, many optimisations need to be performed at once: optimising the layout of a tensor, performing operator fusion and optimising code generation parameters at the same time.

Research conducted by Google scientists has investigated autotuning methods that can perform several tasks of a compiler pipeline in various modes. One of the studies in which 150 ML training and inference models were examined showed about 5% improvement in execution time, with certain workloads producing up to 2.4 times better results than an optimised XLA compiler.

Creating the Infrastructure with MLIR

In addition to this, another important aspect of this transformation is MLIR (Multiple Intermediate Language Representation).

MLIR is a part of the LLVM ecosystem and is responsible for creating tools necessary for constructing different tools which are capable of working with different levels of abstraction and hardware. The importance of this lies in the fact that we are seeing the splitting of AI hardware types.

We have processors, CPUs, GPUs, TPUs, NPUs, and tailored accelerators which come with different features.

The simplified AI compilation process can be described as follows:

PyTorch/JAX model

↓

Graph form of the model

↓

MLIR

↓

Optimization processes

↓

Machine architecture representation

↓

Machine code

Different stages of this process are considered to be the points at which machine learning can help find the best strategy of optimisation.

Expert Perspective: Chris Lattner and the MLIR Concept

Chris Lattner and other MLIR researchers have pointed to the necessity of having a compiler infrastructure that is reusable across various levels of abstraction and hardware configurations.

This need grows further as AI companies design their own unique hardware. Without a shared compiler infrastructure, every hardware accelerator may require its own dedicated optimisation.

PyTorch Integrates Compilation into AI Development

Moreover, this development is observable in PyTorch too.

The module torch. compile utilises TorchDynamo so as to capture the Python programs and also TorchInductor to develop optimised code that is suitable for the underlying hardware.

This implies that programmers may enjoy even higher-level programming of models as the compiler does most of the heavy lifting.

The developer does not have to create every optimised kernel by himself.

The goal of the compiler stack is to automatically find an optimal implementation.

AOTInductor enhances this process even more by compiling the PyTorch model into production artefacts that may operate independently from the Python program.

The overall tendency is obvious:

The AI frameworks are evolving into compiler platforms.

LLMs are Getting Into the Compiler Optimisation Cycle

In future developments, larger neural models may be utilised in helping compiler optimisation.

There have already been efforts to assess the potential of transformer-based models for optimising LLVM.

For instance, one of the projects utilised a 7 billion-parameter transformer model to optimise LLVM assembly code with the goal of shrinking its size. Researchers noted that their innovation was able to decrease instruction count in comparison with the compiler baseline by 3% and provide compilable code 91% of the time.

This does not suggest that large language models will definitely eliminate the need for traditional compilers.

An even more plausible outcome of the ongoing attempts of researchers in optimising compilers with the use of neural models is that the latter will be utilised as an instrument among other tools for compiler systems, while the effectiveness of traditional compiler tools will still be relied on.

The Hidden Arms Race Is Really About Hardware Efficiency

The secret arms race is about the efficiency of hardware. The competition regarding AI is typically regarded in terms of the characteristics of models and their accelerators. However, there is another pressing question: with what efficiency does the software work on the hardware? Suppose you have two accelerators that have nearly the same performance on paper. If the compiler for one of them can utilise the hardware better than for the other one, this can lead to the following benefits:

  • Lower costs of inference

  • More throughput

  • Lower latency

  • Better energy efficiency

  • Fewer servers

  • Better economics

Thus, the role of compilers in this competition is especially important.

The Most Pressing Issue: Trust in Machine Learning

There is a critical issue with machine-learning-oriented compilers.

Predictive errors are made by machine-learning engines.

An optimiser can use the beneficial optimisation process without interfering with the behaviour of the program.

As a result, the next generation of compilers will likely introduce these elements into their design:

ML-based forecasts + established compilation rules + testing evaluations + verification of correctness of the documentation

Machine learning can suggest a transformation, but the system should still have the necessary tools to do a verification of the transformation’s legality and safety.

Making the Optimiser More Efficient

There is a different issue. It appears that the process of optimisation can also require a considerable amount of resources for its completion.

If a compiler tests multiple potential kernels, the execution of the search process can become cost-inefficient.This results in another optimisation issue. The task is about making the optimisation process as efficient as possible.

Use of estimates can benefit the system as it can predict which candidates can be tested on the physical hardware. Thus, instead of testing every candidate, the system will limit its search to the most promising ones.

What Is Going To Follow?

The forthcoming version of compiler-AI systems will most likely incorporate a variety of techniques, including:

  • Reinforcement learning

  • Cost estimation models

  • Evolutionary search

  • Bayesian optimization

  • Graphically-based programs

  • Code transformation with LLM

  • Compiler automation

  • Optimisation of hardware

  • Permanent performance feedback

The trend is already visible in research based on reinforcement learning techniques that aid in making decisions in compiler passes.

The end goal should not be the total automation of the job of compiler engineers.

The Real Race Is Below the AI Model

Competition within the AI space is hot among the many different AI models and accelerators out there; but below all of these tools lies something else – the compiler.

For years, compilers would turn handwritten code into machine code using a series of highly engineered algorithms.

But now, even that process is seeing the influence of machine learning.

Examples of learned policies for compilation are provided by MLGO. Learned search is shown through projects like TVM and Ansor. Production ML workloads can be improved with autotuning through XLA. Increasingly complex compiler pipelines are made possible via MLIR, and PyTorch makes compilation a key part of developing AI.

The result is a whole new competitive frontier. It’s not just about who builds the most powerful AI model or accelerator anymore. It’s increasingly about:

Who can build the most intelligent system for making use of this hardware?

This ML arms race happens inside the compiler – and those who win it will be able to get substantially more performance out of their AI hardware.