Mapping Parallel Matrix Multiplication in GotoBLAS2 to the AMD Versal ACAP for Deep Learning

TL;DR

This paper presents a parallel GEMM implementation on AMD Versal ACAP, optimizing multi-level memory, vector units, and multi-AIE tile scalability for deep learning inference acceleration.

Contribution

It introduces a novel architecture-specific micro-kernel and a scalable parallel design for GEMM on Versal ACAP, leveraging multiple AI Engines for high throughput.

Findings

High parallel scalability with up to 32 AI Engines

Efficient use of Versal ACAP's memory hierarchy

Micro-kernel optimized for mixed precision arithmetic

Abstract

This paper investigates the design of parallel general matrix multiplication (GEMM) for a Versal Adaptive Compute Accelerated Platform (ACAP) equipped with a VC1902 system-on-chip and multiple Artificial Intelligence Engines (AIEs). Our efforts aim to port standard optimization techniques applied in the high-performance realization of GEMM on CPUs to the Versal ACAP. In particular, 1) we address the flexible exploitation of the Versal ACA multi-level memory hierarchy; 2) we delve into the efficient use of the vector units in the AIE tiles, proposing an architecture-specific micro-kernel for mixed precision arithmetic to address the strong demand for adaptive-precision inference in deep learning; and 3) we introduce a parallel design for GEMM that spans multiple AIE tiles, enhancing the computational throughput. We conduct experimental profiling, with up to 32 AI Engines, that demonstrates the high parallel scalability of the solution.