use Eigen as a BLAS alternative (LeelaChessZero#858)

* import Eigen 3.3.7

Author: borg323

meson.build

@@ -195,7 +195,11 @@ if get_option('build_backends')
   endif
 
   if get_option('blas')
-    if get_option('mkl') and mkl_lib.found()
+    if get_option('eigen')
+      add_project_arguments('-DUSE_EIGEN', language : 'cpp')
+      has_blas = true
+
+    elif get_option('mkl') and mkl_lib.found()
       add_project_arguments('-DUSE_MKL', language : 'cpp')
       includes += include_directories(get_option('mkl_include'))
       deps += [ mkl_lib ]

meson_options.txt

@@ -103,6 +103,11 @@ option('mkl',
        value: true,
        description: 'Enable MKL BLAS support')
 
+option('eigen',
+       type: 'boolean',
+       value: false,
+       description: 'Use EIGEN as a BLAS alternative')
+
 option('accelerate',
        type: 'boolean',
        value: true,

src/neural/blas/convolution1.cc

@@ -19,7 +19,19 @@
 #include "neural/blas/convolution1.h"
 #include "neural/blas/blas.h"
 
+#ifdef USE_EIGEN
+#include <Eigen/Dense>
+#endif
+
 namespace lczero {
+#ifdef USE_EIGEN
+template <typename T>
+using EigenMatrixMap =
+    Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
+template <typename T>
+using ConstEigenMatrixMap =
+    Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
+#endif
 
 void Convolution1::Forward(const size_t batch_size, const size_t input_channels,
                            const size_t output_channels, const float* input,
@@ -44,7 +56,7 @@ void Convolution1::Forward(const size_t batch_size, const size_t input_channels,
 
     const float* batch_input = input + i * kSquares * input_channels;
     float* batch_output = output + i * kSquares * output_channels;
-
+#ifndef USE_EIGEN
     cblas_sgemm(CblasRowMajor,         // Row major formar
                 CblasNoTrans,          // A not transposed
                 CblasNoTrans,          // B not transposed
@@ -59,6 +71,12 @@ void Convolution1::Forward(const size_t batch_size, const size_t input_channels,
                 0.0f,                  // beta
                 batch_output,          // C
                 kSquares);             // ldc, leading rank of B
+#else
+    auto C_mat = EigenMatrixMap<float>(batch_output, kSquares, output_channels);
+    C_mat.noalias() =
+        ConstEigenMatrixMap<float>(batch_input, kSquares, input_channels) *
+        ConstEigenMatrixMap<float>(weights, input_channels, output_channels);
+#endif
   }
 }
 

src/neural/blas/fully_connected_layer.cc

@@ -23,7 +23,24 @@
 #include <cassert>
 #include <cmath>
 
+#ifdef USE_EIGEN
+#include <Eigen/Dense>
+#endif
+
 namespace lczero {
+#ifdef USE_EIGEN
+template <typename T>
+using EigenVectorMap = Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 1>>;
+template <typename T>
+using ConstEigenVectorMap =
+    Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>>;
+template <typename T>
+using EigenMatrixMap =
+    Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
+template <typename T>
+using ConstEigenMatrixMap =
+    Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
+#endif
 
 void FullyConnectedLayer::Forward1D(size_t batch_size, const size_t input_size,
                                     const size_t output_size,
@@ -41,11 +58,17 @@ void FullyConnectedLayer::Forward1D(size_t batch_size, const size_t input_size,
     //
     //   rows  output_size      output_size          input_size
     //
-
+#ifndef USE_EIGEN
     cblas_sgemv(CblasRowMajor, CblasNoTrans,
                 // M     K
                 (int)output_size, (int)input_size, 1.0f, weights,
                 (int)input_size, inputs, 1, 0.0f, outputs, 1);
+#else
+    EigenVectorMap<float> y(outputs, output_size);
+    y.noalias() = ConstEigenMatrixMap<float>(weights, input_size, output_size)
+                      .transpose() *
+                  ConstEigenVectorMap<float>(inputs, input_size);
+#endif
   } else {
     // more columns, matrix-matrix multiplication
     //
@@ -66,7 +89,7 @@ void FullyConnectedLayer::Forward1D(size_t batch_size, const size_t input_size,
     // passing a matrix A[m][n], the value should be m.
     //    cblas_sgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
     //                ldb, beta, C, N);
-
+#ifndef USE_EIGEN
     cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans,
                 (int)output_size,   // M
                 (int)batch_size,    // N
@@ -79,6 +102,13 @@ void FullyConnectedLayer::Forward1D(size_t batch_size, const size_t input_size,
                 0.0f,               // beta
                 outputs,            // C
                 (int)output_size);  // ldc, leading rank of C
+#else
+    auto C_mat = EigenMatrixMap<float>(outputs, output_size, batch_size);
+    C_mat.noalias() =
+        ConstEigenMatrixMap<float>(weights, input_size, output_size)
+            .transpose() *
+        ConstEigenMatrixMap<float>(inputs, input_size, batch_size);
+#endif
   }
   if (apply_relu) {
     for (size_t i = 0; i < batch_size; i++) {
@@ -104,7 +134,12 @@ float FullyConnectedLayer::Forward0D(const size_t size, const float* x,
   // float cblas_sdot(const int N, const float *X, const int incX, const float
   // *Y,
   // const int incY);
+#ifndef USE_EIGEN
   return cblas_sdot((int)size, x, 1, y, 1);
+#else
+  return ConstEigenVectorMap<float>(x, size)
+      .dot(ConstEigenVectorMap<float>(y, size));
+#endif
 }
 
 }  // namespace lczero

src/neural/blas/network_blas.cc

@@ -33,6 +33,10 @@
 #include <cmath>
 #include <iostream>
 
+#ifdef USE_EIGEN
+#include <Eigen/Core>
+#endif
+
 namespace lczero {
 namespace {
 
@@ -339,7 +343,9 @@ void BlasComputation::EncodePlanes(const InputPlanes& sample, float* buffer) {
 
 BlasNetwork::BlasNetwork(const WeightsFile& file, const OptionsDict& options)
     : weights_(file.weights()) {
+#ifndef USE_EIGEN
   int blas_cores = options.GetOrDefault<int>("blas_cores", 1);
+#endif
   max_batch_size_ =
       static_cast<size_t>(options.GetOrDefault<int>("batch_size", 256));
 
@@ -378,6 +384,11 @@ BlasNetwork::BlasNetwork(const WeightsFile& file, const OptionsDict& options)
                                                        pol_channels, channels);
   }
 
+#ifdef USE_EIGEN
+  CERR << "Using Eigen version " << EIGEN_WORLD_VERSION << "."
+       << EIGEN_MAJOR_VERSION << "." << EIGEN_MINOR_VERSION;
+#endif
+
 #ifdef USE_OPENBLAS
   int num_procs = openblas_get_num_procs();
   blas_cores = std::min(num_procs, blas_cores);

src/neural/blas/winograd_convolution3.cc

@@ -29,7 +29,19 @@
 #include "winograd_transform_ispc.h"
 #endif
 
+#ifdef USE_EIGEN
+#include <Eigen/Dense>
+#endif
+
 namespace lczero {
+#ifdef USE_EIGEN
+template <typename T>
+using EigenMatrixMap =
+    Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
+template <typename T>
+using ConstEigenMatrixMap =
+    Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
+#endif
 
 WinogradConvolution3::WinogradConvolution3(const size_t max_batch_size,
                                            const size_t max_input_layers,
@@ -206,7 +218,7 @@ void WinogradConvolution3::Sgemm(const size_t batch_size, const float* weights,
 
     auto offset_v = b * batch_size * input_channels * kTiles;
     auto offset_m = b * batch_size * output_channels * kTiles;
-
+#ifndef USE_EIGEN
     cblas_sgemm(CblasColMajor,               // Row major format
                 CblasNoTrans,                // A no trans
                 CblasNoTrans,                // B no trans
@@ -220,6 +232,14 @@ void WinogradConvolution3::Sgemm(const size_t batch_size, const float* weights,
                 (int)input_channels, 0.0f,   // ldV
                 &M_[offset_m],               // M
                 (int)output_channels);       // ldM
+#else
+    auto C_mat = EigenMatrixMap<float>(&M_[offset_m], output_channels,
+                                       batch_size * kTiles);
+    C_mat.noalias() = ConstEigenMatrixMap<float>(
+                          &weights[offset_u], output_channels, input_channels) *
+                      ConstEigenMatrixMap<float>(&V_[offset_v], input_channels,
+                                                 batch_size * kTiles);
+#endif
   }
 
 #endif

third_party/Eigen/CMakeLists.txt

@@ -0,0 +1,19 @@
+include(RegexUtils)
+test_escape_string_as_regex()
+
+file(GLOB Eigen_directory_files "*")
+
+escape_string_as_regex(ESCAPED_CMAKE_CURRENT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
+
+foreach(f ${Eigen_directory_files})
+  if(NOT f MATCHES "\\.txt" AND NOT f MATCHES "${ESCAPED_CMAKE_CURRENT_SOURCE_DIR}/[.].+" AND NOT f MATCHES "${ESCAPED_CMAKE_CURRENT_SOURCE_DIR}/src")
+    list(APPEND Eigen_directory_files_to_install ${f})
+  endif()
+endforeach(f ${Eigen_directory_files})
+
+install(FILES
+  ${Eigen_directory_files_to_install}
+  DESTINATION ${INCLUDE_INSTALL_DIR}/Eigen COMPONENT Devel
+  )
+
+install(DIRECTORY src DESTINATION ${INCLUDE_INSTALL_DIR}/Eigen COMPONENT Devel FILES_MATCHING PATTERN "*.h")

third_party/Eigen/Cholesky

@@ -0,0 +1,46 @@
+// This file is part of Eigen, a lightweight C++ template library
+// for linear algebra.
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#ifndef EIGEN_CHOLESKY_MODULE_H
+#define EIGEN_CHOLESKY_MODULE_H
+
+#include "Core"
+#include "Jacobi"
+
+#include "src/Core/util/DisableStupidWarnings.h"
+
+/** \defgroup Cholesky_Module Cholesky module
+  *
+  *
+  *
+  * This module provides two variants of the Cholesky decomposition for selfadjoint (hermitian) matrices.
+  * Those decompositions are also accessible via the following methods:
+  *  - MatrixBase::llt()
+  *  - MatrixBase::ldlt()
+  *  - SelfAdjointView::llt()
+  *  - SelfAdjointView::ldlt()
+  *
+  * \code
+  * #include <Eigen/Cholesky>
+  * \endcode
+  */
+
+#include "src/Cholesky/LLT.h"
+#include "src/Cholesky/LDLT.h"
+#ifdef EIGEN_USE_LAPACKE
+#ifdef EIGEN_USE_MKL
+#include "mkl_lapacke.h"
+#else
+#include "src/misc/lapacke.h"
+#endif
+#include "src/Cholesky/LLT_LAPACKE.h"
+#endif
+
+#include "src/Core/util/ReenableStupidWarnings.h"
+
+#endif // EIGEN_CHOLESKY_MODULE_H
+/* vim: set filetype=cpp et sw=2 ts=2 ai: */

third_party/Eigen/CholmodSupport

@@ -0,0 +1,48 @@
+// This file is part of Eigen, a lightweight C++ template library
+// for linear algebra.
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#ifndef EIGEN_CHOLMODSUPPORT_MODULE_H
+#define EIGEN_CHOLMODSUPPORT_MODULE_H
+
+#include "SparseCore"
+
+#include "src/Core/util/DisableStupidWarnings.h"
+
+extern "C" {
+  #include <cholmod.h>
+}
+
+/** \ingroup Support_modules
+  * \defgroup CholmodSupport_Module CholmodSupport module
+  *
+  * This module provides an interface to the Cholmod library which is part of the <a href="http://www.suitesparse.com">suitesparse</a> package.
+  * It provides the two following main factorization classes:
+  * - class CholmodSupernodalLLT: a supernodal LLT Cholesky factorization.
+  * - class CholmodDecomposiiton: a general L(D)LT Cholesky factorization with automatic or explicit runtime selection of the underlying factorization method (supernodal or simplicial).
+  *
+  * For the sake of completeness, this module also propose the two following classes:
+  * - class CholmodSimplicialLLT
+  * - class CholmodSimplicialLDLT
+  * Note that these classes does not bring any particular advantage compared to the built-in
+  * SimplicialLLT and SimplicialLDLT factorization classes.
+  *
+  * \code
+  * #include <Eigen/CholmodSupport>
+  * \endcode
+  *
+  * In order to use this module, the cholmod headers must be accessible from the include paths, and your binary must be linked to the cholmod library and its dependencies.
+  * The dependencies depend on how cholmod has been compiled.
+  * For a cmake based project, you can use our FindCholmod.cmake module to help you in this task.
+  *
+  */
+
+#include "src/CholmodSupport/CholmodSupport.h"
+
+#include "src/Core/util/ReenableStupidWarnings.h"
+
+#endif // EIGEN_CHOLMODSUPPORT_MODULE_H
+