Matrix-Multiply Assist (MMA) exploitation in OpenBLAS on IBM AIX

The Matrix-Multiply Assist (MMA) in an IBM Power10 processor-based system provides the computational strength and data bandwidth to handle the demanding AI inferencing and machine learning (ML) workloads. This facility is a natural match for implementing numerical linear algebra computations. OpenBLAS is a widely used open source BLAS (Basic Linear Algebra Subprograms) library to speed up linear algebra computations with low-level routines that operate on vectors and matrices with platform-specific optimizations. This library has been optimized to use the MMA features on Power10 platform with the IBM AIX 7.3 operating system. This tutorial provides an overview of the MMA features, and the steps to use the MMA-optimized OpenBLAS library on AIX along with the performance metrics for the OpenBLAS benchmarks.

IBM Power10 Matrix-Multiply Assist

The high-throughput MMA facility introduced in IBM Power10 provides a new set of instructions to implement numerical linear algebra operations on small matrices. This help to accelerate computation-intensive kernels, such as matrix multiplication, convolution, and discrete Fourier transform. Also, performance per core of a Power10 processor is four times better than the previous generation, IBM Power9 processor. These make inferencing on host cpus lucrative than using separate inferencing accelerators.

The MMA instructions are fully integrated into the IBM Power Instruction Set Architecture (ISA) and are fetched, decoded and dispatched like any other Power instruction. They use the 128-bit vector-scalar registers for input and a new set of eight registers called accumulators for output. Each accumulator has 512 bits and can hold a 4 x 4 matrix of 32-bit elements (or 4 x 2 matrix of 64-bit elements). MMA facility is implemented without adding an architected state, that is, each 512-bit accumulator register is backed up by four 128-bit Vector Scalar eXtension (VSX) registers.

MMA architecture supports four floating-point data types which includes FP32 (IEEE single-precision), FP64 (IEEE double-precision), FP16 (IEEE half-precision) and bfloat16. In addition, it supports integer operations on INT16, INT8, and INT4 types. These data types support the diverse needs of machine learning. The use of compiler built-in datatypes and functions is the recommended way of using the MMA facility.

Figure 1: Power10 SMT8 Core and MMA Facility

To summarize, Power10 comes with an AI-infused powerful core for inference acceleration with the following highlights:

• Power10 MMA instructions, including:

  • Eight 512 bit architected accumulator registers
  • Four parallel units per SMT8 core
  • A new class of instructions to perform a BLAS3-class rank-k update

• New AI data types, including:

  • Double-precision and single-precision for high-performance computing (HPC) use cases
  • Lower-precision integer and half-precision floating point for deep learning (DL) inference use cases
  • Lane masking for small / odd matrix sizes

OpenBLAS

In scientific computing, OpenBLAS is a widely used open source BLAS library to speed up linear algebra computations with low-level routines that operate on vectors and matrices with platform-specific optimizations.

OpenBLAS routines are divided into three groups:

• BLAS Level 1 routines perform vector-only operations, such as vector-vector addition.
• BLAS Level 2 routines perform matrix-vector operations, such as matrix-vector multiplication.
• BLAS Level 3 contains matrix-matrix operations, such as matrix-matrix multiplication.

Extensively used BLAS operation in scientific computing and AI workloads is Matrix Multiplication (Level-3 BLAS).

MMA - Exploitation in OpenBLAS on AIX

Different level BLAS routines have been optimized to exploit MMA when running on Power10 processor-based systems with AIX (big endian) and Linux (little endian). OpenBLAS APIs have been internally modified to use MMA ‘C’ data types, functions, and instructions while running on Power10. To leverage the MMA benefit while running AI workloads, users have to just download the MMA-optimized OpenBLAS package from the AIX toolbox on a Power10 processor-based system and run their AI programs without any modifications.

Perform the following steps to download and compile a test program using the OpenBLAS API:

1. Install OpenBLAS on Power10 with AIX 7.3 or later version from AIX Toolbox along with gcc 10.3.0 compiler.

   dnf install gcc
   dnf install openblas openblas-devel
   

2. Create a following sample OpenBLAS ‘C’ code (test_cblas_dgemm.c) for double precision floating point Matrix Multiplication (DGEMM).

   #include <cblas.h>
   #include <stdio.h>
   void main()
   {
   int i=0;
   double A[6] = {1.0,2.0,1.0,-3.0,4.0,-1.0};
   double B[6] = {1.0,2.0,1.0,-3.0,4.0,-1.0};
   double C[9] = {.5,.5,.5,.5,.5,.5,.5,.5,.5};
   cblas_dgemm(CblasColMajor, CblasNoTrans, CblasTrans,3,3,2,1,A, 3, B, 3,2,C,3);
   
   for(i=0; i<9; i++)
   printf("%lf ", C[i]);
   printf("\n");
   }
   

3. Compile the OpenBLAS program using the installed MMA-optimized OpenBLAS library.

   export PATH=/opt/freeware/bin:$PATH
   gcc -maix64 -o test_cblas_open test_cblas_dgemm.c -I /opt/freeware/include/OpenBLAS_POWER10/ -L /opt/freeware/lib/ -lopenblas -lpthread -lgfortran
   

4. Run the OpenBLAS program.

   ./test_cblas_open
   11.000000 -9.000000 5.000000 -9.000000 21.000000 -1.000000 5.000000 -1.000000 3.000000
   

   Performance of MMA-optimized OpenBLAS on Power10

The GEMM benchmark available in the OpenBLAS repository was used to measure the performance of matrix multiplication for different matrix sizes and different precision data types. The graphs in the following figures show the performance comparison of matrix multiplication of different datatypes on Power9 and Power10 with and without MMA exploitation.

Figure 2: Double-precision complex matrix multiplication - ZGEMM

Figure 3: Single-precision complex matrix multiplication - CGEMM

Figure 4: Double-precision floating point matrix multiplication - DGEMM

Conclusion

With MMA-optimized OpenBLAS, we see around four to five times performance improvement on Power10 when compared to Power9. Thus, users can leverage this facility to accelerate AI workloads running on AIX. These optimizations are part of continuous and committed efforts from IBM to improve AI inferencing on AIX.

References

• OpenBLAS documentation
• Matrix-Multiply Assist Best Practices Guide
• A matrix math facility for Power ISA processor
• AIX Toolbox