diff --git a/cpp_thread_test/CMakeLists.txt b/cpp_thread_test/CMakeLists.txt index c81aacd836..cef0a51eb9 100644 --- a/cpp_thread_test/CMakeLists.txt +++ b/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}) diff --git a/cpp_thread_test/Makefile b/cpp_thread_test/Makefile index 56c4759d0a..139811a420 100644 --- a/cpp_thread_test/Makefile +++ b/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 diff --git a/cpp_thread_test/zgemm_thread_safety.cpp b/cpp_thread_test/zgemm_thread_safety.cpp new file mode 100644 index 0000000000..910f6cd67a --- /dev/null +++ b/cpp_thread_test/zgemm_thread_safety.cpp @@ -0,0 +1,111 @@ +#include +#include +#include +#include +#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"< cliArgs; + for (int i = 1; i < argc; i++){ + cliArgs.push_back(argv[i]); + std::cout< rngdist{-1.0, 1.0}; + std::vector> matBlock(numConcurrentThreads*3); + std::vector> 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*randomMatSize)*numConcurrentThreads*3*16)/static_cast(1024*1024)<<" MiB of RAM\n"<(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(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 : "< #include #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); } } diff --git a/kernel/arm64/sme_dgemm_kernel.c b/kernel/arm64/sme_dgemm_kernel.c index 6edd82fde0..294e1feb29 100644 --- a/kernel/arm64/sme_dgemm_kernel.c +++ b/kernel/arm64/sme_dgemm_kernel.c @@ -6,6 +6,7 @@ #include #include #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); } } diff --git a/kernel/arm64/sme_gemm_buffer.h b/kernel/arm64/sme_gemm_buffer.h new file mode 100644 index 0000000000..fae8687fcf --- /dev/null +++ b/kernel/arm64/sme_gemm_buffer.h @@ -0,0 +1,74 @@ +/*************************************************************************** +Copyright (c) 2026, The OpenBLAS Project +All rights reserved. +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: +1. Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in +the documentation and/or other materials provided with the +distribution. +3. Neither the name of the OpenBLAS project nor the names of +its contributors may be used to endorse or promote products +derived from this software without specific prior written permission. +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE OPENBLAS PROJECT OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF +THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*****************************************************************************/ + +/* Packing buffers of the SME GEMM kernels sme_[sdcz]gemm_kernel.c, private to the calling thread. */ + +#ifndef SME_GEMM_BUFFER_H +#define SME_GEMM_BUFFER_H + +#include + +#define SME_BUF_CACHE (4UL << 20) /* largest buffer kept per thread; larger ones come from blas_memory_alloc */ + +#if !defined(_WIN32) +#include +static pthread_key_t sme_buf_key; +static pthread_once_t sme_buf_once = PTHREAD_ONCE_INIT; +static void sme_buf_free(void *p) { free(p); } +static void sme_buf_key_init(void) { pthread_key_create(&sme_buf_key, sme_buf_free); } + +static void *sme_buf_cached(size_t bytes) { + pthread_once(&sme_buf_once, sme_buf_key_init); + size_t *p = (size_t *)pthread_getspecific(sme_buf_key); + if (p == NULL || p[0] < bytes) { + free(p); + void *q = NULL; + if (posix_memalign(&q, 256, SME_BUF_CACHE + 256)) q = NULL; + p = (size_t *)q; + if (p) p[0] = SME_BUF_CACHE; + pthread_setspecific(sme_buf_key, p); + if (p == NULL) return NULL; + } + return (char *)p + 256; +} +#else +static void *sme_buf_cached(size_t bytes) { (void)bytes; return NULL; } +#endif + +/* At least bytes (at most BUFFER_SIZE) for the calling thread; *pool is the buffer to return with sme_buf_put. */ +static void *sme_buf_get(size_t bytes, void **pool) { + void *p = bytes <= SME_BUF_CACHE ? sme_buf_cached(bytes) : NULL; + *pool = NULL; + if (p == NULL) p = *pool = blas_memory_alloc(0); + return p; +} + +static void sme_buf_put(void *pool) { + if (pool) blas_memory_free(pool); +} + +#endif diff --git a/kernel/arm64/sme_sgemm_kernel.c b/kernel/arm64/sme_sgemm_kernel.c index 94e174b627..fb21261948 100644 --- a/kernel/arm64/sme_sgemm_kernel.c +++ b/kernel/arm64/sme_sgemm_kernel.c @@ -5,6 +5,7 @@ #include #include #include "common.h" +#include "sme_gemm_buffer.h" #include "sme2_gemm_detect.h" #ifdef HAVE_SME2_GEMM #include "sme2_gemm_impl.h" @@ -157,11 +158,12 @@ static void sgemm_sme_NN(blasint M, blasint N, blasint K, float alpha, const flo return; } - alignas(256) /*thread_local*/ static float A_pack[MC * KC]; - alignas(256) /*thread_local*/ static float B_pack[KC * NC]; - #pragma omp parallel for schedule(dynamic, 1) for (blasint 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); + float *B_pack = (float *)sme_buf_get((nb + na) * sizeof(float), &pool), *A_pack = B_pack + nb; blasint current_N = stdmin(NC, N - j); blasint N_pad = (current_N + 31) & ~31; for (blasint k = 0; k < K; k += KC) { @@ -260,6 +262,7 @@ static void sgemm_sme_NN(blasint M, blasint N, blasint K, float alpha, const flo } } } + sme_buf_put(pool); } } @@ -276,11 +279,12 @@ static void sgemm_sme_TN(int M, int N, int K, float alpha, const float *A, int l return; } - alignas(256) /*thread_local*/ static float A_pack[MC * KC]; - alignas(256) /*thread_local*/ static float 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); + 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 + 31) & ~31; for (int k = 0; k < K; k += KC) { @@ -381,6 +385,7 @@ static void sgemm_sme_TN(int M, int N, int K, float alpha, const float *A, int l } } } + sme_buf_put(pool); } } @@ -397,11 +402,12 @@ static void sgemm_sme_NT(int M, int N, int K, float alpha, const float *A, int l return; } - alignas(256) /*thread_local*/ static float A_pack[MC * KC]; - alignas(256) /*thread_local*/ static float 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); + 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 + 31) & ~31; for (int k = 0; k < K; k += KC) { @@ -497,6 +503,7 @@ static void sgemm_sme_NT(int M, int N, int K, float alpha, const float *A, int l } } } + sme_buf_put(pool); } } @@ -513,11 +520,12 @@ static void sgemm_sme_TT(int M, int N, int K, float alpha, const float *A, int l return; } - alignas(256) /*thread_local*/ static float A_pack[MC * KC]; - alignas(256) /*thread_local*/ static float 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); + 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 + 31) & ~31; for (int k = 0; k < K; k += KC) { @@ -615,6 +623,7 @@ static void sgemm_sme_TT(int M, int N, int K, float alpha, const float *A, int l } } } + sme_buf_put(pool); } } diff --git a/kernel/arm64/sme_zgemm_kernel.c b/kernel/arm64/sme_zgemm_kernel.c index 6de23b9755..3540d6813b 100644 --- a/kernel/arm64/sme_zgemm_kernel.c +++ b/kernel/arm64/sme_zgemm_kernel.c @@ -3,6 +3,7 @@ #include #include #include "common.h" +#include "sme_gemm_buffer.h" #ifndef stdmin #define stdmin(a,b) (a>b? b:a) #endif @@ -163,10 +164,12 @@ static void zgemm_sme_NN(int M, int N, int K, const zdouble alpha, const zdouble return; } - alignas(256) /*thread_local*/ static double A_pack[MC * KC * 2]; - alignas(256) /*thread_local*/ static double 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; + 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 + 7) & ~7; for (int k = 0; k < K; k += KC) { @@ -232,6 +235,7 @@ static void zgemm_sme_NN(int M, int N, int K, const zdouble alpha, const zdouble } } } + sme_buf_put(pool); } } @@ -248,10 +252,12 @@ static void zgemm_sme_TN(int M, int N, int K, const zdouble alpha, const zdouble return; } - alignas(256) /*thread_local*/ static double A_pack[MC * KC * 2]; - alignas(256) /*thread_local*/ static double 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; + 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 + 7) & ~7; for (int k = 0; k < K; k += KC) { @@ -317,6 +323,7 @@ static void zgemm_sme_TN(int M, int N, int K, const zdouble alpha, const zdouble } } } + sme_buf_put(pool); } } @@ -352,10 +359,12 @@ static void zgemm_sme_NT(blasint M, blasint N, blasint K, const zdouble alpha, c } - alignas(256) /*thread_local*/ static double A_pack[MC * KC * 2]; - alignas(256) /*thread_local*/ static double B_pack[KC * NC * 2 *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; + 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 + 7) & ~7; for (int k = 0; k < K; k += KC) { @@ -421,6 +430,7 @@ static void zgemm_sme_NT(blasint M, blasint N, blasint K, const zdouble alpha, c } } } + sme_buf_put(pool); } } @@ -437,10 +447,12 @@ static void zgemm_sme_TT(int M, int N, int K, const zdouble alpha, const zdouble return; } - alignas(256) /*thread_local*/ static double A_pack[MC * KC * 2]; - alignas(256) /*thread_local*/ static double 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; + 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 + 7) & ~7; for (int k = 0; k < K; k += KC) { @@ -506,6 +518,7 @@ static void zgemm_sme_TT(int M, int N, int K, const zdouble alpha, const zdouble } } } + sme_buf_put(pool); } }