Use Threadpool to schedule the work

.github/workflows/build.yml

@@ -326,7 +326,7 @@ jobs:
 #          sudo apt-get install cmake
 #
 #      - name: Configure
-#        run: cmake . -DCMAKE_BUILD_TYPE=Debug -DLLAMA_SANITIZE_${{ matrix.sanitizer }}=ON
+#        run: cmake . -DCMAKE_BUILD_TYPE=RelWithDebInfo -DLLAMA_SANITIZE_${{ matrix.sanitizer }}=ON
 #
 #      - name: Build
 #        run: |

CMakeLists.txt

@@ -64,6 +64,113 @@ option(LLAMA_OPENBLAS               "llama: use OpenBLAS"
 option(LLAMA_BUILD_TESTS            "llama: build tests"    ${LLAMA_STANDALONE})
 option(LLAMA_BUILD_EXAMPLES         "llama: build examples" ${LLAMA_STANDALONE})
 
+INCLUDE(CheckCSourceRuns)
+
+SET(AVX_CODE "
+  #include <immintrin.h>
+  int main()
+  {
+    __m256 a;
+    a = _mm256_set1_ps(0);
+    return 0;
+  }
+")
+
+SET(AVX512_CODE "
+  #include <immintrin.h>
+  int main()
+  {
+    __m512i a = _mm512_set_epi8(0, 0, 0, 0, 0, 0, 0, 0,
+                                0, 0, 0, 0, 0, 0, 0, 0,
+                                0, 0, 0, 0, 0, 0, 0, 0,
+                                0, 0, 0, 0, 0, 0, 0, 0,
+                                0, 0, 0, 0, 0, 0, 0, 0,
+                                0, 0, 0, 0, 0, 0, 0, 0,
+                                0, 0, 0, 0, 0, 0, 0, 0,
+                                0, 0, 0, 0, 0, 0, 0, 0);
+    __m512i b = a;
+    __mmask64 equality_mask = _mm512_cmp_epi8_mask(a, b, _MM_CMPINT_EQ);
+    return 0;
+  }
+")
+
+SET(AVX2_CODE "
+  #include <immintrin.h>
+  int main()
+  {
+    __m256i a = {0};
+    a = _mm256_abs_epi16(a);
+    __m256i x;
+    _mm256_extract_epi64(x, 0); // we rely on this in our AVX2 code
+    return 0;
+  }
+")
+
+SET(FMA_CODE "
+  #include <immintrin.h>
+  int main()
+  {
+    __m256 acc = _mm256_setzero_ps();
+    const __m256 d = _mm256_setzero_ps();
+    const __m256 p = _mm256_setzero_ps();
+    acc = _mm256_fmadd_ps( d, p, acc );
+    return 0;
+  }
+")
+
+MACRO(CHECK_SSE type flags)
+  SET(__FLAG_I 1)
+  SET(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS})
+  FOREACH(__FLAG ${flags})
+    IF(NOT ${type}_FOUND)
+      SET(CMAKE_REQUIRED_FLAGS ${__FLAG})
+      CHECK_C_SOURCE_RUNS("${${type}_CODE}" HAS_${type}_${__FLAG_I})
+      IF(HAS_${type}_${__FLAG_I})
+        SET(${type}_FOUND TRUE CACHE BOOL "${type} support")
+        SET(${type}_FLAGS "${__FLAG}" CACHE STRING "${type} flags")
+      ENDIF()
+      MATH(EXPR __FLAG_I "${__FLAG_I}+1")
+    ENDIF()
+  ENDFOREACH()
+  SET(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE})
+
+  IF(NOT ${type}_FOUND)
+    SET(${type}_FOUND FALSE CACHE BOOL "${type} support")
+    SET(${type}_FLAGS "" CACHE STRING "${type} flags")
+  ENDIF()
+
+  MARK_AS_ADVANCED(${type}_FOUND ${type}_FLAGS)
+
+ENDMACRO()
+
+IF(${LLAMA_AVX})
+    CHECK_SSE("AVX" " ;-mavx;/arch:AVX")
+    IF(NOT ${AVX_FOUND})
+        set(LLAMA_AVX OFF)
+    ENDIF()
+ENDIF()
+
+IF(${LLAMA_AVX2})
+    CHECK_SSE("AVX2" " ;-mavx2 -mfma;/arch:AVX2")
+    IF(NOT ${AVX2_FOUND})
+        set(LLAMA_AVX2 OFF)
+    ENDIF()
+ENDIF()
+
+IF(${LLAMA_AVX512})
+    CHECK_SSE("AVX512" " ;-mavx512f -mavx512dq -mavx512vl -mavx512bw -mfma;/arch:AVX512")
+    IF(NOT ${AVX512_FOUND})
+        set(LLAMA_AVX512 OFF)
+    ENDIF()
+ENDIF()
+
+IF(${LLAMA_FMA})
+    CHECK_SSE("FMA" " ;-mfma;")
+    IF (NOT ${FMA_FOUND})
+        set(LLAMA_FMA OFF)
+    ENDIF()
+ENDIF()
+
 #
 # Compile flags
 #
@@ -229,7 +336,10 @@ endif()
 
 add_library(ggml OBJECT
             ggml.c
-            ggml.h)
+            ggml.h
+            thpool.c
+            thpool.h
+            )
 
 target_include_directories(ggml PUBLIC .)
 target_compile_features(ggml PUBLIC c_std_11) # don't bump

Makefile

@@ -139,6 +139,9 @@ default: main quantize perplexity embedding
 # Build library
 #
 
+thpool.o: thpool.c thpool.h
+	$(CC)  $(CFLAGS) -c thpool.c -o thpool.o
+
 ggml.o: ggml.c ggml.h
 	$(CC)  $(CFLAGS)   -c ggml.c -o ggml.o
 
@@ -151,20 +154,20 @@ common.o: examples/common.cpp examples/common.h
 clean:
 	rm -vf *.o main quantize perplexity embedding
 
-main: examples/main/main.cpp ggml.o llama.o common.o
-	$(CXX) $(CXXFLAGS) examples/main/main.cpp ggml.o llama.o common.o -o main $(LDFLAGS)
+main: examples/main/main.cpp thpool.o ggml.o llama.o common.o
+	$(CXX) $(CXXFLAGS) examples/main/main.cpp thpool.o ggml.o llama.o common.o -o main $(LDFLAGS)
 	@echo
 	@echo '====  Run ./main -h for help.  ===='
 	@echo
 
-quantize: examples/quantize/quantize.cpp ggml.o llama.o
-	$(CXX) $(CXXFLAGS) examples/quantize/quantize.cpp ggml.o llama.o -o quantize $(LDFLAGS)
+quantize: examples/quantize/quantize.cpp thpool.o ggml.o llama.o
+	$(CXX) $(CXXFLAGS) examples/quantize/quantize.cpp thpool.o ggml.o llama.o -o quantize $(LDFLAGS)
 
 perplexity: examples/perplexity/perplexity.cpp ggml.o llama.o common.o
-	$(CXX) $(CXXFLAGS) examples/perplexity/perplexity.cpp ggml.o llama.o common.o -o perplexity $(LDFLAGS)
+	$(CXX) $(CXXFLAGS) examples/perplexity/perplexity.cpp thpool.o ggml.o llama.o common.o -o perplexity $(LDFLAGS)
 
-embedding: examples/embedding/embedding.cpp ggml.o llama.o common.o
-	$(CXX) $(CXXFLAGS) examples/embedding/embedding.cpp ggml.o llama.o common.o -o embedding $(LDFLAGS)
+embedding: examples/embedding/embedding.cpp thpool.o ggml.o llama.o common.o
+	$(CXX) $(CXXFLAGS) examples/embedding/embedding.cpp thpool.o ggml.o llama.o common.o -o embedding $(LDFLAGS)
 
 #
 # Tests