From 4ee494695967d287087124fb45d0eb8f9075b6c3 Mon Sep 17 00:00:00 2001 From: luoyu-intel Date: Mon, 20 Nov 2023 19:37:25 +0800 Subject: [PATCH] fix bug of block_size=128 --- .../core/mlas/lib/x86_64/jblas/jblas/jit_blas_wrapper.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_wrapper.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_wrapper.h index fd1255fcf99ed..2858405f3bfa2 100644 --- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_wrapper.h +++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_wrapper.h @@ -230,10 +230,11 @@ class LauncherKBlock { int blk_msize, int blk_nsize, AType* tmpA, BType* tmpB, CType* tmpBlk, AccType* tmpC, void* tmpcache) { int n_padded = utils::padto(blk_nsize, GemmCore::NTILE); + assert(_param.K % _param.KBlock == 0); for (int iterk = 0; iterk < _param.K; iterk += _param.KBlock) { memset(tmpBlk, 0, sizeof(CType) * blk_msize * _config.block[1]); for (int iblkk = 0; iblkk < _param.KBlock; iblkk += _config.block[2]) { - int k_remain = utils::remainsize(iterk + iblkk, _param.K, _config.block[2]); + int k_remain = utils::remainsize(iterk + iblkk, iterk + _param.KBlock, _config.block[2]); int k_padded = utils::padto(k_remain, GemmCore::KTILE); int k_paddedle = utils::padto_le(k_remain, GemmCore::KTILE); auto bptr_cache = tmpB;