arm64: SME GEMM kernels with packing buffers per thread (thread safety)

cpp_thread_test/CMakeLists.txt

@@ -20,6 +20,7 @@ set(CPP_THREAD_SAFETY_DGEMM_ARGS "" CACHE STRING "Arguments passed to the DGEMM
 set(CPP_THREAD_SAFETY_DGEMM_MIXED_ARGS "" CACHE STRING "Arguments passed to the mixed DGEMM thread safety test")
 set(CPP_THREAD_SAFETY_DGEMV_ARGS "" CACHE STRING "Arguments passed to the DGEMV thread safety test")
 set(CPP_THREAD_SAFETY_SHUTDOWN_ARGS "" CACHE STRING "Arguments passed to the DGEMM shutdown safety test")
+set(CPP_THREAD_SAFETY_ZGEMM_ARGS "" CACHE STRING "Arguments passed to the ZGEMM thread safety test")
 
 if (CPP_THREAD_SAFETY_TEST)
   message(STATUS "building thread safety test")
@@ -31,6 +32,10 @@ if (CPP_THREAD_SAFETY_TEST)
   target_link_libraries(dgemm_thread_safety_mixed ${CPP_THREAD_SAFETY_LIBS})
   add_test(NAME dgemm_thread_safety_mixed COMMAND ${CMAKE_CURRENT_BINARY_DIR}/dgemm_thread_safety_mixed ${CPP_THREAD_SAFETY_DGEMM_MIXED_ARGS})
 
+  add_executable(zgemm_thread_safety zgemm_thread_safety.cpp)
+  target_link_libraries(zgemm_thread_safety ${CPP_THREAD_SAFETY_LIBS})
+  add_test(NAME zgemm_thread_safety COMMAND ${CMAKE_CURRENT_BINARY_DIR}/zgemm_thread_safety ${CPP_THREAD_SAFETY_ZGEMM_ARGS})
+
   if (WIN32)
     add_executable(dgemm_thread_safety_shutdown dgemm_thread_safety_shutdown.cpp)
     target_link_libraries(dgemm_thread_safety_shutdown ${CPP_THREAD_SAFETY_LIBS})

cpp_thread_test/Makefile

@@ -1,7 +1,7 @@
 TOPDIR = ..
 include $(TOPDIR)/Makefile.system
 
-all :: dgemv_tester dgemm_tester dgemm_mixed_tester
+all :: dgemv_tester dgemm_tester dgemm_mixed_tester zgemm_tester
 
 ifeq ($(OSNAME), WINNT)
 all :: dgemm_shutdown_tester
@@ -35,6 +35,10 @@ dgemm_mixed_tester : dgemm_tester
 	$(CXX) $(COMMON_OPT) -Wall -Wextra -Wshadow -std=c++11 dgemm_thread_safety_mixed.cpp ../$(LIBNAME) $(EXTRALIB) $(FEXTRALIB) -o dgemm_mixed_tester
 	./dgemm_mixed_tester
 
+zgemm_tester : dgemm_mixed_tester
+	$(CXX) $(COMMON_OPT) -Wall -Wextra -Wshadow -std=c++11 zgemm_thread_safety.cpp ../$(LIBNAME) $(EXTRALIB) $(FEXTRALIB) -o zgemm_tester
+	./zgemm_tester
+
 dgemm_shutdown_tester : dgemm_mixed_tester
 	$(CXX) $(COMMON_OPT) -Wall -Wextra -Wshadow -std=c++11 dgemm_thread_safety_shutdown.cpp ../$(LIBNAME) $(EXTRALIB) $(FEXTRALIB) -o dgemm_shutdown_tester
 	./dgemm_shutdown_tester
@@ -45,4 +49,4 @@ dgemm_mixed_callback_tester : dgemm_mixed_tester
 endif
 
 clean ::
-	rm -f dgemv_tester dgemm_tester dgemm_mixed_tester dgemm_shutdown_tester
+	rm -f dgemv_tester dgemm_tester dgemm_mixed_tester zgemm_tester dgemm_shutdown_tester

cpp_thread_test/zgemm_thread_safety.cpp

@@ -0,0 +1,111 @@
+#include <iostream>
+#include <vector>
+#include <random>
+#include <future>
+#ifdef OPENBLAS_USE_GENERATED_CBLAS_H
+#include "generated/cblas.h"
+#else
+#include "../cblas.h"
+#endif
+#include "cpp_thread_safety_common.h"
+
+void launch_cblas_zgemm(double* A, double* B, double* C, const blasint randomMatSize){
+	const double alpha[2] = {1.0, 0.5}, beta[2] = {0.1, -0.2};
+	cblas_zgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, randomMatSize, randomMatSize, randomMatSize, alpha, A, randomMatSize, B, randomMatSize, beta, C, randomMatSize);
+}
+
+int main(int argc, char* argv[]){
+	blasint randomMatSize = 384; //dimension of the random square complex matrices used
+	uint32_t numConcurrentThreads = 96; //number of concurrent calls of the functions being tested
+	uint32_t numTestRounds = 16; //number of testing rounds before success exit
+	uint32_t maxHwThreads = GetMaxHwThreads();
+
+	if (maxHwThreads < 96)
+		numConcurrentThreads = maxHwThreads;
+
+	if (argc > 4){
+		std::cout<<"ERROR: too many arguments for thread safety tester"<<std::endl;
+		abort();
+	}
+
+	if(argc == 4){
+		std::vector<std::string> cliArgs;
+		for (int i = 1; i < argc; i++){
+			cliArgs.push_back(argv[i]);
+			std::cout<<argv[i]<<std::endl;
+		}
+		randomMatSize = std::stoul(cliArgs[0]);
+		numConcurrentThreads = std::stoul(cliArgs[1]);
+		numTestRounds = std::stoul(cliArgs[2]);
+	}
+
+	std::uniform_real_distribution<double> rngdist{-1.0, 1.0};
+	std::vector<std::vector<double>> matBlock(numConcurrentThreads*3);
+	std::vector<std::future<void>> futureBlock(numConcurrentThreads);
+
+	std::cout<<"*----------------------------*\n";
+	std::cout<<"| ZGEMM thread safety tester |\n";
+	std::cout<<"*----------------------------*\n";
+	std::cout<<"Size of random matrices(N=M=K): "<<randomMatSize<<'\n';
+	std::cout<<"Number of concurrent calls into OpenBLAS : "<<numConcurrentThreads<<'\n';
+	std::cout<<"Number of testing rounds : "<<numTestRounds<<'\n';
+	std::cout<<"This test will need "<<(static_cast<uint64_t>(randomMatSize*randomMatSize)*numConcurrentThreads*3*16)/static_cast<double>(1024*1024)<<" MiB of RAM\n"<<std::endl;
+
+	FailIfThreadsAreZero(numConcurrentThreads);
+
+	std::cout<<"Initializing random number generator..."<<std::flush;
+	std::mt19937_64 PRNG = InitPRNG();
+	std::cout<<"done\n";
+
+	std::cout<<"Preparing to test CBLAS ZGEMM thread safety\n";
+	std::cout<<"Allocating matrices..."<<std::flush;
+	for(uint32_t i=0; i<(numConcurrentThreads*3); i++){
+		matBlock[i].resize(2*randomMatSize*randomMatSize);
+	}
+	std::cout<<"done\n";
+	//pauser();
+	std::cout<<"Filling matrices with random numbers..."<<std::flush;
+	for(uint32_t i=0; i<3; i++){
+		for(uint32_t j = 0; j < static_cast<uint32_t>(2*randomMatSize*randomMatSize); j++){
+			matBlock[i][j] = rngdist(PRNG);
+		}
+	}
+	for(uint32_t i=3; i<(numConcurrentThreads*3); i++){
+		matBlock[i] = matBlock[i%3];
+	}
+	//PrintMatrices(matBlock, randomMatSize, numConcurrentThreads, 3);
+	std::cout<<"done\n";
+	std::cout<<"Testing CBLAS ZGEMM thread safety\n";
+	SetLauncherThreads(numConcurrentThreads);
+	for(uint32_t R=0; R<numTestRounds; R++){
+		std::cout<<"ZGEMM round #"<<R<<std::endl;
+		std::cout<<"Launching "<<numConcurrentThreads<<" threads simultaneously"<<LauncherName()<<"..."<<std::flush;
+		#ifdef CPP_THREAD_SAFETY_USE_OPENMP
+		#pragma omp parallel for default(none) shared(futureBlock, matBlock, randomMatSize, numConcurrentThreads)
+		#endif
+		for(uint32_t i=0; i<numConcurrentThreads; i++){
+			futureBlock[i] = std::async(std::launch::async, launch_cblas_zgemm, &matBlock[i*3][0], &matBlock[i*3+1][0], &matBlock[i*3+2][0], randomMatSize);
+			//launch_cblas_zgemm( &matBlock[i][0], &matBlock[i+1][0], &matBlock[i+2][0]);
+		}
+		std::cout<<"done\n";
+		std::cout<<"Waiting for threads to finish..."<<std::flush;
+		for(uint32_t i=0; i<numConcurrentThreads; i++){
+			futureBlock[i].get();
+		}
+		std::cout<<"done\n";
+		//PrintMatrices(matBlock, randomMatSize, numConcurrentThreads, 3);
+		std::cout<<"Comparing results from different threads..."<<std::flush;
+		for(uint32_t i=3; i<(numConcurrentThreads*3); i+=3){ //i is the index of matrix A, for a given thread
+			for(uint32_t j = 0; j < static_cast<uint32_t>(2*randomMatSize*randomMatSize); j++){
+				if (std::abs(matBlock[i+2][j] - matBlock[2][j]) > 1.0E-13){ //i+2 is the index of matrix C, for a given thread
+					std::cout<<"ERROR: one of the threads returned a different result! Index : "<<i+2<<std::endl;
+					std::cout<<"CBLAS ZGEMM thread safety test FAILED!"<<std::endl;
+					return -1;
+				}
+			}
+		}
+		std::cout<<"OK!\n"<<std::endl;
+	}
+	std::cout<<"CBLAS ZGEMM thread safety test PASSED!\n"<<std::endl;
+	return 0;
+}

kernel/arm64/sme_cgemm_kernel.c

@@ -5,6 +5,7 @@
 #include <stdalign.h>
 #include <stdbool.h>
 #include "common.h"
+#include "sme_gemm_buffer.h"
 #ifndef stdmin
 #define stdmin(a,b)   (a>b? b:a)
 #endif
@@ -174,10 +175,12 @@ static void cgemm_sme_NN(int M, int N, int K, const cfloat alpha, const cfloat *
 
 
 
-    alignas(256) /*thread_local*/ static float A_pack[MC * KC * 2];
-    alignas(256) /*thread_local*/ static float B_pack[KC * NC * 2];
 #pragma omp parallel for schedule(dynamic, 1)
     for (int j = 0; j < N; j += NC) {
+        /* per thread (and OpenMP task): the panels are never shared between threads */
+        void *pool;
+        const size_t kc = stdmin(KC, K), nb = kc * (stdmin(NC, N) + 64) * 2, na = kc * (stdmin(MC, M) + 64) * 2;
+        float *B_pack = (float *)sme_buf_get((nb + na) * sizeof(float), &pool), *A_pack = B_pack + nb;
         int current_N = stdmin(NC, N - j);
         int N_pad = (current_N + 15) & ~15;
         for (int k = 0; k < K; k += KC) {
@@ -242,6 +245,7 @@ static void cgemm_sme_NN(int M, int N, int K, const cfloat alpha, const cfloat *
                 }
             }
         }
+        sme_buf_put(pool);
     }
 }
 
@@ -263,10 +267,12 @@ static void cgemm_sme_TN(int M, int N, int K, const cfloat alpha, const cfloat *
 	}
         return;
     }
-    alignas(256) /*thread_local*/ static float A_pack[MC * KC * 2];
-    alignas(256) /*thread_local*/ static float B_pack[KC * NC * 2];
 #pragma omp parallel for schedule(dynamic, 1)
     for (int j = 0; j < N; j += NC) {
+        /* per thread (and OpenMP task): the panels are never shared between threads */
+        void *pool;
+        const size_t kc = stdmin(KC, K), nb = kc * (stdmin(NC, N) + 64) * 2, na = kc * (stdmin(MC, M) + 64) * 2;
+        float *B_pack = (float *)sme_buf_get((nb + na) * sizeof(float), &pool), *A_pack = B_pack + nb;
         int current_N = stdmin(NC, N - j);
         int N_pad = (current_N + 15) & ~15;
         for (int k = 0; k < K; k += KC) {
@@ -332,6 +338,7 @@ static void cgemm_sme_TN(int M, int N, int K, const cfloat alpha, const cfloat *
                 }
             }
         }
+        sme_buf_put(pool);
     }
 }
 
@@ -353,10 +360,12 @@ static void cgemm_sme_NT(int M, int N, int K, const cfloat alpha, const cfloat *
 	}
         return;
     }
-    alignas(256) /*thread_local*/ static float A_pack[MC * KC * 2];
-    alignas(256) /*thread_local*/ static float B_pack[KC * NC * 2];
 #pragma omp parallel for schedule(dynamic, 1)
     for (int j = 0; j < N; j += NC) {
+        /* per thread (and OpenMP task): the panels are never shared between threads */
+        void *pool;
+        const size_t kc = stdmin(KC, K), nb = kc * (stdmin(NC, N) + 64) * 2, na = kc * (stdmin(MC, M) + 64) * 2;
+        float *B_pack = (float *)sme_buf_get((nb + na) * sizeof(float), &pool), *A_pack = B_pack + nb;
         int current_N = stdmin(NC, N - j);
         int N_pad = (current_N + 15) & ~15;
         for (int k = 0; k < K; k += KC) {
@@ -422,6 +431,7 @@ static void cgemm_sme_NT(int M, int N, int K, const cfloat alpha, const cfloat *
                 }
             }
         }
+        sme_buf_put(pool);
     }
 }
 
@@ -443,10 +453,12 @@ static void cgemm_sme_TT(int M, int N, int K, const cfloat alpha, const cfloat *
 	}
         return;
     }
-    alignas(256) /*thread_local*/ static float A_pack[MC * KC * 2];
-    alignas(256) /*thread_local*/ static float B_pack[KC * NC * 2];
 #pragma omp parallel for schedule(dynamic, 1)
     for (int j = 0; j < N; j += NC) {
+        /* per thread (and OpenMP task): the panels are never shared between threads */
+        void *pool;
+        const size_t kc = stdmin(KC, K), nb = kc * (stdmin(NC, N) + 64) * 2, na = kc * (stdmin(MC, M) + 64) * 2;
+        float *B_pack = (float *)sme_buf_get((nb + na) * sizeof(float), &pool), *A_pack = B_pack + nb;
         int current_N = stdmin(NC, N - j);
         int N_pad = (current_N + 15) & ~15;
         for (int k = 0; k < K; k += KC) {
@@ -512,6 +524,7 @@ static void cgemm_sme_TT(int M, int N, int K, const cfloat alpha, const cfloat *
                 }
             }
         }
+        sme_buf_put(pool);
     }
 }
 

kernel/arm64/sme_dgemm_kernel.c

@@ -6,6 +6,7 @@
 #include <stdbool.h>
 #include <stdio.h>
 #include "common.h"
+#include "sme_gemm_buffer.h"
 #include "sme2_gemm_detect.h"
 #ifdef HAVE_SME2_GEMM
 #include "sme2_gemm_impl.h"
@@ -252,11 +253,12 @@ static void dgemm_sme_NN(int M, int N, int K, double alpha, const double *A, int
         return;
     }
 
-    alignas(256) /*thread_local*/ static double A_pack[MC * KC];
-    alignas(256) /*thread_local*/ static double B_pack[KC * NC];
-
 #pragma omp parallel for schedule(dynamic, 1)
     for (int j = 0; j < N; j += NC) {
+        /* per thread (and OpenMP task): the panels are never shared between threads */
+        void *pool;
+        const size_t kc = stdmin(KC, K), nb = kc * (stdmin(NC, N) + 64), na = kc * (stdmin(MC, M) + 64);
+        double *B_pack = (double *)sme_buf_get((nb + na) * sizeof(double), &pool), *A_pack = B_pack + nb;
 #if PROFILING
         CALI_MARK_BEGIN("loop_j");
 #endif
@@ -512,6 +514,7 @@ static void dgemm_sme_NN(int M, int N, int K, double alpha, const double *A, int
 #if PROFILING
         CALI_MARK_END("loop_j");
 #endif
+        sme_buf_put(pool);
     }
 }
 
@@ -531,11 +534,12 @@ static void dgemm_sme_TN(int M, int N, int K, double alpha, const double *A, int
         return;
     }
 
-    alignas(256) /*thread_local*/ static double A_pack[MC * KC];
-    alignas(256) /*thread_local*/ static double B_pack[KC * NC];
-
 #pragma omp parallel for schedule(dynamic, 1)
     for (int j = 0; j < N; j += NC) {
+        /* per thread (and OpenMP task): the panels are never shared between threads */
+        void *pool;
+        const size_t kc = stdmin(KC, K), nb = kc * (stdmin(NC, N) + 64), na = kc * (stdmin(MC, M) + 64);
+        double *B_pack = (double *)sme_buf_get((nb + na) * sizeof(double), &pool), *A_pack = B_pack + nb;
         int current_N = stdmin(NC, N - j);
         int N_pad = (current_N + 15) & ~15;
 
@@ -743,6 +747,7 @@ static void dgemm_sme_TN(int M, int N, int K, double alpha, const double *A, int
                 }
             }
         }
+        sme_buf_put(pool);
     }
 }
 
@@ -762,11 +767,12 @@ static void dgemm_sme_NT(int M, int N, int K, double alpha, const double *A, int
         return;
     }
 
-    alignas(256) /*thread_local*/ static double A_pack[MC * KC];
-    alignas(256) /*thread_local*/ static double B_pack[KC * NC];
-
 #pragma omp parallel for schedule(dynamic, 1)
     for (int j = 0; j < N; j += NC) {
+        /* per thread (and OpenMP task): the panels are never shared between threads */
+        void *pool;
+        const size_t kc = stdmin(KC, K), nb = kc * (stdmin(NC, N) + 64), na = kc * (stdmin(MC, M) + 64);
+        double *B_pack = (double *)sme_buf_get((nb + na) * sizeof(double), &pool), *A_pack = B_pack + nb;
         int current_N = stdmin(NC, N - j);
         int N_pad = (current_N + 15) & ~15;
 
@@ -964,6 +970,7 @@ static void dgemm_sme_NT(int M, int N, int K, double alpha, const double *A, int
                 }
             }
         }
+        sme_buf_put(pool);
     }
 }
 
@@ -983,11 +990,12 @@ static void dgemm_sme_TT(int M, int N, int K, double alpha, const double *A, int
         return;
     }
 
-    alignas(256) /*thread_local*/ static double A_pack[MC * KC];
-    alignas(256) /*thread_local*/ static double B_pack[KC * NC];
-
 #pragma omp parallel for schedule(dynamic, 1)
     for (int j = 0; j < N; j += NC) {
+        /* per thread (and OpenMP task): the panels are never shared between threads */
+        void *pool;
+        const size_t kc = stdmin(KC, K), nb = kc * (stdmin(NC, N) + 64), na = kc * (stdmin(MC, M) + 64);
+        double *B_pack = (double *)sme_buf_get((nb + na) * sizeof(double), &pool), *A_pack = B_pack + nb;
         int current_N = stdmin(NC, N - j);
         int N_pad = (current_N + 15) & ~15;
 
@@ -1179,6 +1187,7 @@ static void dgemm_sme_TT(int M, int N, int K, double alpha, const double *A, int
                 }
             }
         }
+        sme_buf_put(pool);
     }
 }