EPCC PhD research: Compiling HPC codes for novel hardware architectures

In recent years a number of custom hardware accelerators have been produced for accelerating particular workloads, such as deep learning. Their architectures differ dramatically from traditional CPUs and GPUs, and show promise for accelerating the wider class of traditional HPC workloads as well. However, programming them often requires developers to become experts in the architectures and redesign all their codes for each, leaving their expertise of existing programming models like OpenMP behind. 

Tenstorrent Blackhole

To those who know these accelerators well this sounds like a dream. For example, if we have a look at the Tenstorrent Blackhole, we see something that looks very different to standard CPU or GPU clusters today. The Blackhole features a mesh of "Tensix cores", which each contain their own SRAM memory, have vector and matrix accelerators, and are interconnected through a network on the chip. Each Tensix core itself comprises five RISC-V cores, which each run their own code and can bring data in from the network, compute on that data, or send data out to the network.

The Tensix cores can communicate by sending data directly into each other's SRAM, and the idea is that programmers should aim to use this to keep the data that is being operated on near to the hardware that operates on it, rather than resort to using slow, chip-shared DRAM memory. This approach of keeping data near to its compute units is generally referred to as a spatial computing architecture.

^{Above: diagram of Tenstorrent Blackhole.}

Programming the accelerator is intricate, and attaining high performance more so. Each core is capable of running three kernels simultaneously, one to bring data into the core, one to operate on it, and one to send data out of the core. Each requires the programmer to work with bare-metal constructs directly such as memory addresses and registers. Computations must be done on 32x32 element tiles, and synchronisation must be inserted at many levels to ensure correct parallel programs. 

Whilst this complexity enables flexibility for the HPC programmer, it places a great burden on them to design their code with the architecture in mind.

Running HPC codes on advanced architectures

The question my work attempts to answer is whether there exists a way to take existing HPC codes and just run them on these advanced architectures? If so, can it achieve or come close to the performance of hand-written codes? 

So far I have explored these questions using Multi-Level Intermediate Representation (MLIR) compilers. These compilers allow code to be transformed progressively, in small steps from a very high-level abstraction down to hardware instructions. Each transformation adds a little more information to the representation and brings it closer to a runnable program on the target architecture. 

My hypothesis is that the right abstractions can enable performant code across architectures by shifting the burden of managing the hardware intricacies into the compiler. This research thus far has considered a subset of this problem and focused specifically on running directive-based loop programs (written with OpenMP or OpenACC) on the Tenstorrent accelerators. 

Performance on Tenstorrent

We now have a working compiler that can compile these programs for the Tenstorrent accelerators, but do we see any reasonable performance?

So far, yes! By building on existing work at EPCC, we are able to lift pragma-annotated codes written in Fortran/C/C++ to more abstract tensor programs, and then progressively specialise these tensor programs for the Tenstorrent architecture. We do this through a multi-level compiler stack built with the MLIR and xDSL frameworks, with representations and code transformations specific to the Tenstorrent architecture. For programs already written in a compatible form, the HPC programmer need not change the code, and instead change only the compiler.

Through this work we have demonstrated correct programs across various HPC codes, including N-body simulations, spectral-element methods, and stencils. We have compared these against CPU (AMD EPYC 9825) and GPU (Nvidia A100) baselines and seen reasonable performance, and strong performance for matrix-multiply focused workloads which suit the architecture well. 

Not only can we run these complicated codes, but when we compare against the existing programming approaches for Tenstorrent accelerators, we find that our compiler-based solution achieves performance within a small factor of vendor-written and hand-written HPC foundational kernels, at a fraction of the lines of code. 

Our approach successfully enables developers to continue using their knowledge of existing programming models and target a novel architecture.

^{Above: Comparing the runtime performance and code complexity of our compiler-based approach against vendor-handwritten code (TT-NN) and codes written in the vendor-provided programming language (TT-Lang).}

Next steps

With this foundational work in place, my next steps are to try to share some of this infrastructure across some other accelerators, such as the RISC-V-based ET-SoC-1, to determine what compiler representations are sharable, and what remains a challenge in accelerator-abstract compilation. 

If any of this work interests you, please get in touch or follow my work:

My website

LinkedIn

GitHub

Author

Jake Davies, EPCC PhD student

Read more

PhD study at EPCC