The grant will support work on automated generation of compilers for AI accelerators.
Charith Mendis has won an NSF CAREER award that will support his ongoing work on rapid, machine-learning-driven creation of compilers for a key hardware component of many large AI systems.
AI accelerators are small machines that are assembled in large numbers to handle machine learning tasks by doing matrix multiplications. They can do so faster and with greater energy efficiency than GPUs or CPUs can, and for that reason, many companies and academic research teams have been eagerly developing more and more of them. Google’s Tensor Processing Units (TPUs) and Amazon’s Inferentia and Trainium are examples.
But there’s a catch: “Each company is developing their own compiler by hand, and there does not exist a compiler infrastructure that works across many different accelerators,” says Mendis, who is an assistant professor in Computer Science. “So how do you support hundreds of these accelerators that will come up in the future?” In particular, he asks, “How do you build compiler technologies easily, that can adapt to multiple accelerators?”
Each of the growing number of accelerators has its own nuances, meaning there can’t be a one-size-fits-all compiler. Further, accelerators get upgraded and modified over time, so their compilers must be updated to match. Unfortunately, creation and updating of compilers are typically done by hand, and are far from being quick or cheap: such efforts often require teams of dozens of programmers.
The process is especially complicated because it’s important for a compiler not merely to work, but to do so optimally. “You want to get the best out of the machine rather than underutilizing it,” notes Mendis.
So what will Mendis do to tackle this problem? In his CAREER project—a 5-year, $525K effort entitled “An Agile Compiler Framework for Spatial Dataflow Accelerators”—he will invent a way to generate compilers automatically via machine learning, and he will use formal methods, and his own newly developed techniques, to confirm the correctness of the generated compilers’ translations.
In addition to the enormous cost savings, a big advantage of his solution will be the ease of revising compilers quickly.
“Say that you change your hardware,” he says. “You can develop the new compiler very easily, because most of the compiler construction techniques are automated, not hard-coded.”
Mendis worked at Google prior to joining UIUC in 2021. While there, he built a data-driven cost model—now used in production at Google—for predicting the cost of running programs in TPUs. Cost modeling, he explains, is done when you have multiple options, and you want to know which is most profitable. Most people handle cost modeling problems with handwritten analytical models, but he instead asked the question: if we can collect a huge data set, could we simply learn a model that can rank options?
“So what we did was, we collected a huge data set of program executions on Google’s Tensor Processing Units,” Mendis says. “And using that data set, we trained a model that can properly rank these options.”
That solution is a precursor to the planned CAREER work; it is already in use inside Google’s TPU compiler to make optimization decisions, and Mendis plans to generalize the solution to work with multiple accelerator platforms.
“And the good thing is, this model is now learned through data, so it can be updated very quickly,” Mendis explains. “You don’t need a huge team to come in and change coefficients or whatever. Because you just need new data! That’s it!”
One benefit of his Google relationship is that he’s been using a large (now open-source) dataset of real-world workloads. Since he hasn’t been limited to the small or synthetic datasets usually used in academia, the outcomes of his CAREER work should be robust. He, therefore, anticipates that “there will be industrial adoption if the project goes well.”
According to NSF’s website, Faculty Early Career Development (CAREER) grants are its “most prestigious awards in support of early-career faculty who have the potential to serve as academic role models in research and education and to lead advances in the mission of their department or organization.”