Back in May I wrote a post explaining the architecture of AMD NPUs and demonstrating a simple kernel that reaches close to the peak TOPS allowed by the hardware resources. That program was very simple and there was no data movement between the NPU compute tiles and the rest of the system. After writing that post, I wanted to write a more complex program to understand better the NPU architecture and the software stack, and what the challenges are when working on real-world projects.
I decided to implement an X-engine correlator, such as the ones commonly used in radio astronomy. This is a problem that I’m personally interested in, and it is basically a big matrix multiplication, with some small caveats that I will explain below. As the main goal of NPUs is to be used for matrix multiplication in machine learning inference, I thought “how hard can this be?”. After a few months of working on and off on this project, I now have something ready to show.
The design has actually been quite challenging. There are many nuances I will get into below, and many dead-end paths I have taken until I discovered that when using some particular approaches performance would be killed for one reason or another, so I had to backtrack multiple times. I should also mention that while it is usually very clear what is happening at low level (which was one of the main messages of my previous post), the lack of flexibility of the hardware is often quite limiting, in particular regarding the programmability of the DMA engines, so it is quite challenging to design approaches that actually work well in the hardware.
This post is a full presentation of my X-engine implementation.