[Fix] Add configurable number of threads for OpenBLAS

Issue: #3082

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

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})

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)

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							*/

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",

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: $,}"

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;
-}
+}