Matrix Multiplication - 2024.1 English

AI Engine-ML Kernel and Graph Programming Guide (UG1603)

Document ID
Release Date
2024.1 English

The AI Engine API provides a aie::mmul class template for a vector-based matrix multiplication. Multiple intermediate matrix multiplication results are accumulated to give the final result. For more details on the supported matrix multiplication shapes (M*K*N) and data types, see Matrix Multiplication in the AI Engine API User Guide (UG1529).

The aie::mmul operations mul and mac accept row based data for the vector-based matrix multiplication. Then for the mac operation, arrange the data by M*K or K*N for better performance. This data shuffling can be done either in PL or AI Engine.

This section gives an example of (64 * 64) x (64 * 64) dimensions matrix multiplication. The data type is int8 x int8. The matrix multiplication shape 4*16*8 is chosen.

The input data is assumed to be row based. Prior to the matrix multiplication kernel, the input data is shuffled in AI Engine-ML memory tile to form 4*16 matrix and 16*8 matrix.

The picture below shows the graph.

Figure 1. Matrix Multiplication Kernel Graph