Using KML in the Eigen Library

Introduction to the Eigen Library

The Eigen library is a C++ template library for linear algebra operations. It provides high-performance matrix operations and linear algebra operations, and is widely used in scientific computing, machine learning, and computer graphics. For details, visit the Eigen official website. Using BLAS/LAPACK from Eigen describes a method of enabling open-source BLAS and LAPACK libraries. The BLAS and LAPACK libraries in KML can also be used. External BLAS interfaces provide matrix computing interfaces with higher performance for Eigen. In addition to standard BLAS and LAPACK interfaces, KML provides additional BLAS interfaces for Eigen. For details, see Table 1.

KML-based Optimization and Adaptation for Eigen

Based on the EIGEN_USE_BLAS mechanism, KML further extends and optimizes multiple types of Eigen interfaces to improve the overall performance in application scenarios. You can obtain the optimized Eigen version from kunpeng-eigen.

For basic operators such as matrix addition, usage remains identical to the native Eigen interface. To enable the optimized path, simply replace the compilation macro -DEIGEN_USE_BLAS with -DKUNPENG_USE in the compilation phase. The usage and verification processes of other optimized interfaces will be described in detail in subsequent sections.

The following table lists the Eigen functions that can be optimized.

Enabling KML Through Eigen TensorContraction

Enable KML acceleration for 2D float tensor contractions within Eigen TensorContraction module by replacing underlying operators while retaining the Eigen ThreadPool for multi-thread scheduling.

1. Download the optimized Eigen library at kunpeng-eigen
2. Go to the eigen/eigen_blas directory, which acts as the Eigen-KML adaptation layer. Conduct an independent compilation and connect to Eigen in binary mode. The compile command is as follows:

   clang -shared -fPIC -o libeigen_blas.so eigen_blas.c -L${PATH_TO_KBLAS} -lkblas

   Replace PATH_TO_KBLAS with the path of the KML dynamic library. After the compilation is complete, the libeigen_blas.so dynamic library is generated.

3. In the root directory of Eigen, set OMP_NUM_THREADS=1, KUNPENG_OPT_CONTRACT=1, and KUNPENG_CONTRACT_NUM_THREADS=number_of_target_threads, and then run the compile command.

   Sample code of matmul_test.cpp:

   #include <iostream>
   #include <chrono>
   #include <Eigen/Core>
   #include <unsupported/Eigen/CXX11/Tensor>
   #include <unsupported/Eigen/CXX11/ThreadPool>
   using namespace Eigen;
   int main() {
       int M = 5000, N = 1024, K = 1024;
       Tensor<float, 2, RowMajor> t_a(M, K);
       Tensor<float, 2, RowMajor> t_b(K, N);
       Tensor<float, 2, RowMajor> t_c(M, N);
       t_a.setRandom();
       t_b.setRandom();
       Eigen::array<Eigen::IndexPair<int>, 1> product_dims = { Eigen::IndexPair<int>(1, 0) };
       int num_threads = 4;
       Eigen::ThreadPool pool(num_threads);
       Eigen::ThreadPoolDevice device(&pool, num_threads);
       unsetenv("KUNPENG_OPT_CONTRACT"); 
     
       t_c.device(device) = t_a.contract(t_b, product_dims);
       auto start_ori = std::chrono::high_resolution_clock::now();
       t_c.device(device) = t_a.contract(t_b, product_dims);
       auto end_ori = std::chrono::high_resolution_clock::now();
       auto dur_ori = std::chrono::duration_cast<std::chrono::microseconds>(end_ori - start_ori).count();
       setenv("KUNPENG_OPT_CONTRACT", "1", 1);
       setenv("KUNPENG_CONTRACT_NUM_THREADS", "4", 1);
       t_c.device(device) = t_a.contract(t_b, product_dims);
       auto start_opt = std::chrono::high_resolution_clock::now();
       t_c.device(device) = t_a.contract(t_b, product_dims);
       auto end_opt = std::chrono::high_resolution_clock::now();
       
       auto dur_opt = std::chrono::duration_cast<std::chrono::microseconds>(end_opt - start_opt).count();
       std::cout << "--- MatMul (Contraction) KML Optimization Test ---" << std::endl;
       std::cout << "Scale          : " << M << "x" << N << "x" << K << std::endl;
       std::cout << "Original Time  : " << dur_ori << " us" << std::endl;
       std::cout << "Optimized Time : " << dur_opt << " us" << std::endl;
       std::cout << "Speedup        : " << (float)dur_ori / dur_opt << "x" << std::endl;
       return 0;
   }

   Compile command:

   clang++ matmul_test.cpp -I ./eigen -I ./eigen/eigen_blas -o matmul_test -march=armv9-a+sve -DEIGEN_USE_THREADS -L ./eigen/eigen_blas -leigen_blas

   Command output:

   --- MatMul (Contraction) KML Optimization Test ---
   Scale          : 1024x1024x1024
   Original Time  : 2131122 us
   Optimized Time : 5400 us
   Speedup        : 394.652x

Optimizing Eigen tanh

On the Arm architecture, the Eigen tanh function defaults to NEON instruction set for vectorization. However, the Scalable Vector Extension (SVE) supports variable-length vectors ranging from 128 bits to 2048 bits, allowing for more tanh computations to be processed at once, thereby improving performance. To further enhance efficiency, SVE intrinsic optimizations during tanh computation have been implemented specifically for the float data type.

1. Download the optimized Eigen library at kunpeng-eigen
2. Go to the root directory of Eigen and compile the following sample test code.
   tanh_test.cpp

   #include <iostream>
   #include <chrono>
   #include <vector>
   #include <Eigen/Core>
   #include <unsupported/Eigen/CXX11/Tensor>
   #include <unsupported/Eigen/CXX11/ThreadPool>
   using namespace Eigen;
   int main() {
       int M = 500, N = 500;
       int size = M * N;
       std::vector<float> input_raw(size);
       std::vector<float> output_raw(size);
       for(int i = 0; i < size; ++i) input_raw[i] = static_cast<float>(rand());
       TensorMap<Tensor<float, 2, RowMajor>> input_map(input_raw.data(), M, N);
       TensorMap<Tensor<float, 2, RowMajor>> output_map(output_raw.data(), M, N);
       Eigen::ThreadPool pool(1);
       Eigen::ThreadPoolDevice device(&pool, 1);
       auto start_ori = std::chrono::high_resolution_clock::now();
       output_map.device(device) = (1.0f * input_map).tanh();
       auto end_ori = std::chrono::high_resolution_clock::now();
       auto dur_ori = std::chrono::duration_cast<std::chrono::microseconds>(end_ori - start_ori).count();
       auto start_opt = std::chrono::high_resolution_clock::now();
       output_map.device(device) = input_map.tanh();
       auto end_opt = std::chrono::high_resolution_clock::now();
       auto dur_opt = std::chrono::duration_cast<std::chrono::microseconds>(end_opt - start_opt).count();
       std::cout << "--- Tanh SVE Optimization ---" << std::endl;
       std::cout << "Original Time  : " << dur_ori << " us" << std::endl;
       std::cout << "Optimized Time : " << dur_opt << " us" << std::endl;
       std::cout << "Speedup        : " << (float)dur_ori / dur_opt << "x" << std::endl;
       return 0;
   }

   Compile command:

   clang++ test_tanh.cpp -I ./eigen -o test_tanh -march=armv9-a+sve -I ./eigen/eigen_blas -DEIGEN_USE_THREADS

   Command output:

   --- Tanh SVE Optimization ---
   Original Time  : 43847 us
   Optimized Time : 3987 us
   Speedup        : 10.9975x