fix apple math library single prec

block-hczhai · Apr 29, 2023 · 6574017 · 6574017
1 parent 0016748
commit 6574017
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 10 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -14,6 +14,7 @@ OPTION(USE_COMPLEX "ComplexNumber" OFF)
 OPTION(USE_SG "SpinGeneral" OFF)
 OPTION(USE_SINGLE_PREC "SinglePrecision" OFF)
 OPTION(ARCH_ARM64 "MacOS arch arm64" OFF)
+OPTION(APPLE_ACC_SINGLE_PREC "Fix Apple Accelerate single prec" ON)
 OPTION(BUILD_LIB "Build python block2.so" OFF)
 OPTION(BUILD_CLIB "Build C++ block2.so" OFF)
 
@@ -491,6 +492,10 @@ IF(${F77UNDERSCORE})
     SET(MKL_FLAG ${MKL_FLAG} "-D_F77UNDERSCORE")
 ENDIF()
 
+IF (APPLE AND ${USE_SINGLE_PREC} AND (NOT ${USE_MKL_ANY}) AND ${APPLE_ACC_SINGLE_PREC})
+    SET(MKL_FLAG ${MKL_FLAG} "-D_APPLE_ACC_SINGLE_PREC")
+ENDIF()
+
 IF (${BUILD_LIB})
     ADD_LIBRARY(${PROJECT_NAME} MODULE ${PYBIND_SRCS} src/pybind.cpp ${SRCS})
     SET_TARGET_PROPERTIES(${PROJECT_NAME} PROPERTIES SUFFIX "${PYLIB_SUFFIX}" PREFIX "")

diff --git a/src/core/matrix_functions.hpp b/src/core/matrix_functions.hpp
@@ -54,11 +54,11 @@ extern void FNAME(saxpy)(const MKL_INT *n, const float *sa, const float *sx,
                          const MKL_INT *incy) noexcept;
 
 // vector dot product
-extern float FNAME(sdot)(const MKL_INT *n, const float *dx, const MKL_INT *incx,
+extern FRETT FNAME(sdot)(const MKL_INT *n, const float *dx, const MKL_INT *incx,
                          const float *dy, const MKL_INT *incy) noexcept;
 
 // Euclidean norm of a vector
-extern float FNAME(snrm2)(const MKL_INT *n, const float *x,
+extern FRETT FNAME(snrm2)(const MKL_INT *n, const float *x,
                           const MKL_INT *incx) noexcept;
 
 // matrix multiplication
@@ -346,7 +346,7 @@ inline double xnrm2<double>(const MKL_INT *n, const double *x,
 template <>
 inline float xnrm2<float>(const MKL_INT *n, const float *x,
                           const MKL_INT *incx) noexcept {
-    return FNAME(snrm2)(n, x, incx);
+    return (float)FNAME(snrm2)(n, x, incx);
 }
 
 template <typename FL>
@@ -380,7 +380,7 @@ inline double xdot<double>(const MKL_INT *n, const double *dx,
 template <>
 inline float xdot<float>(const MKL_INT *n, const float *dx, const MKL_INT *incx,
                          const float *dy, const MKL_INT *incy) noexcept {
-    return FNAME(sdot)(n, dx, incx, dy, incy);
+    return (float)FNAME(sdot)(n, dx, incx, dy, incy);
 }
 
 template <typename FL>
@@ -848,9 +848,9 @@ struct GMatrixFunctions<
     // c(.T) = bra.T * a(.T) * ket
     // return nflop. (.T) is always transpose conjugate
     static size_t left_partial_rotate(const GMatrix<FL> &a, bool conj_a,
-                         const GMatrix<FL> &c, bool conj_c,
-                         const GMatrix<FL> &bra, const GMatrix<FL> &ket,
-                         FL scale) {
+                                      const GMatrix<FL> &c, bool conj_c,
+                                      const GMatrix<FL> &bra,
+                                      const GMatrix<FL> &ket, FL scale) {
         shared_ptr<VectorAllocator<FL>> d_alloc =
             make_shared<VectorAllocator<FL>>();
         GMatrix<FL> work(nullptr, conj_a ? a.n : a.m, ket.n);
@@ -866,9 +866,9 @@ struct GMatrixFunctions<
     // c(.T) = bra.c * a(.T) * ket.t = (a(~.T) * bra.t).T * ket.t
     // return nflop. (.T) is always transpose conjugate
     static size_t right_partial_rotate(const GMatrix<FL> &a, bool conj_a,
-                         const GMatrix<FL> &c, bool conj_c,
-                         const GMatrix<FL> &bra, const GMatrix<FL> &ket,
-                         FL scale) {
+                                       const GMatrix<FL> &c, bool conj_c,
+                                       const GMatrix<FL> &bra,
+                                       const GMatrix<FL> &ket, FL scale) {
         shared_ptr<VectorAllocator<FL>> d_alloc =
             make_shared<VectorAllocator<FL>>();
         GMatrix<FL> work(nullptr, conj_a ? a.m : a.n, bra.m);

diff --git a/src/core/utils.hpp b/src/core/utils.hpp
@@ -58,6 +58,12 @@ namespace block2 {
 #define FNAME(X) X
 #endif
 
+#ifdef _APPLE_ACC_SINGLE_PREC
+#define FRETT double
+#else
+#define FRETT float
+#endif
+
 #ifdef _WIN32
 
 struct timeval {