X-engine correlator on a Ryzen NPU

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.

Getting peak TOPS on a Ryzen AI 7 350 NPU

I have a Framework Laptop 13 that has a Ryzen AI 7 350 CPU that includes an NPU. I have started playing with this NPU to understand how to develop software for it. While NPUs are mainly intended as accelerators for inference of ML models, they are fundamentally hardware accelerators for matrix multiplication and other similar linear algebra operations, so they are also useful for signal processing and other compute applications, which is why I am interested in them. Another reason why I am interested in this NPU is that, as I will explain below, it is very similar to the AIE-ML v2 AI engine in Versal FPGA SoCs, so this laptop is a great platform to learn how to use this AI engine.

NPUs use the concept of TOPS (tera operations per second) as a high-level marketing figure of their capabilities. An operation is generally understood as an addition or multiplication for int8 data types, since the amount of parallelization that can be achieved depends on the datatype width. The NPU on the Ryzen AI 7 350 is marketed as a 50 TOPS NPU. The main goal of this post is to understand where this number comes from, in terms of hardware execution units and capabilities, understand under which conditions it can be reached, and write a small application that reaches this TOPS value.

I think this is a good way of gaining in-depth understanding about a compute architecture. Most typical real world use cases are going to be slower than this, because the algorithms will have bottlenecks that result in hardware underutilization. By understanding how the hardware needs to be used to reach peak performance, we have a better idea of the gaps of these algorithms and also how to rewrite the algorithms to reduce the gap if possible. In a post last year about NEON kernels on the ARM Cortex-A53 I worked in a similar way, by choosing a simple kernel to accelerate and by comparing performance benchmarks with the peak performance allowed by the hardware.

Published
Categorised as Software Tagged