[Fix] Add configurable number of threads for OpenBLAS

Issue: #3082

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3416702c8ad17666c57029dbadd730f86572115c..99fbc8f2d8936916607ffe0857cef69cfbe5a31f 100644 (file)
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -643,6 +643,27 @@ ProcessPackage(LIBZ LIBRARY z INCLUDE zlib.h INCLUDE_SUFFIXES include/zlib
 ProcessPackage(SODIUM LIBRARY sodium INCLUDE sodium.h
                INCLUDE_SUFFIXES include/libsodium include/sodium
                ROOT ${LIBSODIUM_ROOT_DIR} MODULES libsodium>=1.0.0)
+ProcessPackage(BLAS OPTIONAL LIBRARY openblas blas
+               INCLUDE cblas.h INCLUDE_SUFFIXES include/openblas
+               include/blas
+               ROOT ${BLAS_ROOT_DIR}
+               LIB_OUTPUT BLAS_REQUIRED_LIBRARIES)
+IF(WITH_BLAS)
+       MESSAGE(STATUS "Use openblas to accelerate kann")
+       IF(NOT BLAS_INCLUDE)
+               FIND_FILE(HAVE_CBLAS_H HINTS "${RSPAMD_SEARCH_PATH}"
+                               NAMES cblas.h
+                               DOC "Path to cblas.h header")
+               IF(NOT HAVE_CBLAS_H)
+                       MESSAGE(STATUS "Blas header cblas.h has not been found, use internal workaround")
+               ELSE()
+                       ADD_DEFINITIONS(-DHAVE_CBLAS_H)
+               ENDIF()
+       ELSE()
+               ADD_DEFINITIONS(-DHAVE_CBLAS_H)
+       ENDIF()
+       ADD_DEFINITIONS(-DHAVE_CBLAS)
+ENDIF(WITH_BLAS)
 
 IF(ENABLE_HYPERSCAN MATCHES "ON")
        ProcessPackage(HYPERSCAN LIBRARY hs INCLUDE hs.h INCLUDE_SUFFIXES

diff --git a/contrib/kann/CMakeLists.txt b/contrib/kann/CMakeLists.txt
index ba38e14c8f9e8b5ab957575139ac202b73fe330d..5b3ea024b70de8c04ec66abf7dc01d79452df27c 100644 (file)
--- a/contrib/kann/CMakeLists.txt
+++ b/contrib/kann/CMakeLists.txt
@@ -15,21 +15,7 @@ ProcessPackage(BLAS OPTIONAL LIBRARY openblas blas
         LIB_OUTPUT BLAS_REQUIRED_LIBRARIES)
 IF(WITH_BLAS)
     MESSAGE(STATUS "Use openblas to accelerate kann")
-    IF(NOT BLAS_INCLUDE)
-        FIND_FILE(HAVE_CBLAS_H HINTS "${RSPAMD_SEARCH_PATH}"
-            NAMES cblas.h
-            DOC "Path to cblas.h header")
-        IF(NOT HAVE_CBLAS_H)
-            MESSAGE(STATUS "Blas header cblas.h has not been found, use internal workaround")
-        ELSE()
-            ADD_DEFINITIONS(-DHAVE_CBLAS_H)
-        ENDIF()
-    ELSE()
-        ADD_DEFINITIONS(-DHAVE_CBLAS_H)
-    ENDIF()
-
-    TARGET_LINK_LIBRARIES(rspamd-kann  ${BLAS_REQUIRED_LIBRARIES})
-    ADD_DEFINITIONS(-DHAVE_CBLAS)
+    TARGET_LINK_LIBRARIES(rspamd-kann ${BLAS_REQUIRED_LIBRARIES})
 ENDIF(WITH_BLAS)
 
 INSTALL(TARGETS rspamd-kann LIBRARY DESTINATION ${RSPAMD_LIBDIR})
\ No newline at end of file

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 6411601fa2ab34bb6530ef2ccd63f6bd3bdd3b23..797ae2074c1b8766e84df1131f08f064f3ce3260 100644 (file)
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -206,6 +206,10 @@ IF (ENABLE_HYPERSCAN MATCHES "ON")
        TARGET_LINK_LIBRARIES(rspamd-server hs)
 ENDIF()
 
+IF (WITH_BLAS)
+       TARGET_LINK_LIBRARIES(rspamd-server ${BLAS_REQUIRED_LIBRARIES})
+ENDIF()
+
 TARGET_LINK_LIBRARIES(rspamd-server ${RSPAMD_REQUIRED_LIBRARIES})
 
 ADD_EXECUTABLE(rspamd ${RSPAMDSRC} ${CMAKE_CURRENT_BINARY_DIR}/workers.c)

diff --git a/src/libserver/cfg_file.h b/src/libserver/cfg_file.h
index 4eea4db161c4871ee2709d7a6d76833c4aaffa39..3ee832d5ac7e72e9db283480972cb3cd4983bcb5 100644 (file)
--- a/src/libserver/cfg_file.h
+++ b/src/libserver/cfg_file.h
@@ -463,6 +463,7 @@ struct rspamd_config {
        guint lua_gc_pause;                                /**< lua gc pause                                                                            */
        guint full_gc_iters;                            /**< iterations between full gc cycle                                   */
        guint max_lua_urls;                             /**< maximum number of urls to be passed to Lua                 */
+       guint max_blas_threads;                         /**< maximum threads for openblas when learning ANN             */
 
        GList *classify_headers;                        /**< list of headers using for statistics                               */
        struct module_s **compiled_modules;                /**< list of compiled C modules                                                      */

diff --git a/src/libserver/cfg_rcl.c b/src/libserver/cfg_rcl.c
index 5ef913f36f425d36c63cfef1375b365193cfa114..7936c872546b4dc56639703b4eefbdf98b1a20bd 100644 (file)
--- a/src/libserver/cfg_rcl.c
+++ b/src/libserver/cfg_rcl.c
@@ -2200,7 +2200,13 @@ rspamd_rcl_config_init (struct rspamd_config *cfg, GHashTable *skip_sections)
                                rspamd_rcl_parse_struct_integer,
                                G_STRUCT_OFFSET (struct rspamd_config, max_lua_urls),
                                RSPAMD_CL_FLAG_INT_32,
-                               "Maximum count of URLs to pass to Lua to avoid DoS");
+                               "Maximum count of URLs to pass to Lua to avoid DoS (default: 1024)");
+               rspamd_rcl_add_default_handler (sub,
+                               "max_blas_threads",
+                               rspamd_rcl_parse_struct_integer,
+                               G_STRUCT_OFFSET (struct rspamd_config, max_blas_threads),
+                               RSPAMD_CL_FLAG_INT_32,
+                               "Maximum number of Blas threads for learning neural networks (default: 1)");
 
                /* Neighbours configuration */
                rspamd_rcl_add_section_doc (&sub->subsections, "neighbours", "name",

diff --git a/src/libserver/cfg_utils.c b/src/libserver/cfg_utils.c
index 5851e250f6f0a295f0f3cdb3013f4774176a3f7a..98a50bea3d418c717da47b88395bd3791ac9cead 100644 (file)
--- a/src/libserver/cfg_utils.c
+++ b/src/libserver/cfg_utils.c
@@ -196,6 +196,7 @@ rspamd_config_new (enum rspamd_config_init_flags flags)
        cfg->log_error_elt_maxlen = 1000;
        cfg->cache_reload_time = 30.0;
        cfg->max_lua_urls = 1024;
+       cfg->max_blas_threads = 1;
 
        /* Default log line */
        cfg->log_format_str = "id: <$mid>,$if_qid{ qid: <$>,}$if_ip{ ip: $,}"

diff --git a/src/libutil/util.c b/src/libutil/util.c
index ecdd7b2b9e6b26447788c0e31e718068f894621e..5baffebd95e6a21e60b23f942335945e28386d78 100644 (file)
--- a/src/libutil/util.c
+++ b/src/libutil/util.c
@@ -2401,6 +2401,14 @@ rspamd_free_zstd_dictionary (struct zstd_dictionary *dict)
        }
 }
 
+#ifdef HAVE_CBLAS
+#ifdef HAVE_CBLAS_H
+#include "cblas.h"
+#else
+extern void openblas_set_num_threads(int num_threads);
+#endif
+#endif
+
 void
 rspamd_config_libs (struct rspamd_external_libs_ctx *ctx,
                struct rspamd_config *cfg)
@@ -2494,6 +2502,9 @@ rspamd_config_libs (struct rspamd_external_libs_ctx *ctx,
                        ZSTD_freeCStream (ctx->out_zstream);
                        ctx->out_zstream = NULL;
                }
+#ifdef HAVE_CBLAS
+               openblas_set_num_threads (cfg->max_blas_threads);
+#endif
        }
 }
 
@@ -3267,4 +3278,4 @@ rspamd_set_counter_ema (struct rspamd_counter_data *cd,
        cd->number ++;
 
        return cd->mean;
-}
+}
\ No newline at end of file