enlarge stack size

intel · May 23, 2024 · 4683b39 · 4683b39
1 parent 4965b41
commit 4683b39
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/bestla/bestla/bestla_wrapper.h b/bestla/bestla/bestla_wrapper.h
@@ -284,11 +284,12 @@ class LauncherBase {
     template <typename ScaleT, int MTILE, class SNbits>
     static void gemv_kblock(const Param& _param, const parallel::gemm::ThreadProblemBase& _config) {
       if constexpr (support()) {
-        auto constexpr TmpSize = 3 * 1024LL;
-        auto constexpr CSize = 1 * 1024LL;
+        auto constexpr TmpSize = 16 * 1024LL;
+        auto constexpr CSize = 8 * 1024LL;
         auto StackTmp_ = alloca(TmpSize + CSize);
         auto StackTmp = utils::cpu_pointer_align<void>(StackTmp_);
         auto tmpc_ptr = reinterpret_cast<CType*>((char*)StackTmp + TmpSize);
+        static_assert(CSize >= (MTILE * GemmCore::NTILE * sizeof(float)));
         utils::GemvParamB<ScaleT> paramB = SNbits::template createB<ScaleT>(_param.paramB.packedW);
         const float* Aptr = _param.paramA.A;
         if constexpr (std::is_same_v<PrologueA,
@@ -689,8 +690,9 @@ class LauncherIntKBlock {
     template <typename ScaleT, int MTILE, class SNbits>
     static void gemv_kblock(const Param& _param, const parallel::gemm::ThreadProblemBase& _config) {
       if constexpr (support()) {
-        auto constexpr TmpSize = 3 * 1024LL;
-        auto constexpr CSize = 1 * 1024LL;
+        auto constexpr TmpSize = 16 * 1024LL;
+        auto constexpr CSize = 8 * 1024LL;
+        static_assert(CSize >= (MTILE * GemmCore::NTILE * sizeof(float)));
         auto StackTmp_ = alloca(TmpSize + CSize);
         auto StackTmp = utils::cpu_pointer_align<void>(StackTmp_);
         auto tmpc_ptr = reinterpret_cast<CType*>((char*)StackTmp + TmpSize);