From 1ff0a0409568f4dbdac1d5b99f2c7020b80614ed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=99=88=E6=B0=91=E5=AE=89?= <chenminan@hisilicon.com>
Date: Fri, 2 Aug 2024 14:50:47 +0800
Subject: [PATCH] fix int4 size

---
 impl/matmul/matmul_impl.h | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/impl/matmul/matmul_impl.h b/impl/matmul/matmul_impl.h
index e61f9a91..b9758a64 100644
--- a/impl/matmul/matmul_impl.h
+++ b/impl/matmul/matmul_impl.h
@@ -1022,6 +1022,11 @@ __aicore__ inline void MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_
 
     int aMatrixByteSize = var.baseMK_ * sizeof(SrcT);
     int bMatrixByteSize = var.baseKN_ * sizeof(SrcT);
+    // The size of the int4b_t type obtained by sizeof is 1, but in reality, its actual size is 0.5
+    if constexpr (IsSameType<SrcT, int4b_t>::value) {
+        aMatrixByteSize = aMatrixByteSize / 2;
+        bMatrixByteSize = bMatrixByteSize / 2;
+    }
     if constexpr (!PhyPosIsL1(A_TYPE::pos)) {
         if (var.tiling_->depthA1 > DB_FACTOR) {
             if (var.tiling_->depthA1 < var.kIter_ * var.tiling_->stepM) {
@@ -1285,6 +1290,11 @@ __aicore__ inline void MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_
 
     int aMatrixByteSize = var.baseMK_ * sizeof(SrcT);
     int bMatrixByteSize = var.baseKN_ * sizeof(SrcT);
+    // The size of the int4b_t type obtained by sizeof is 1, but in reality, its actual size is 0.5
+    if constexpr (IsSameType<SrcT, int4b_t>::value) {
+        aMatrixByteSize = aMatrixByteSize / 2;
+        bMatrixByteSize = bMatrixByteSize / 2;
+    }
 
     if constexpr (!PhyPosIsL1(A_TYPE::pos)) {
         uint32_t cacheA1Size = var.tiling_->stepM * var.tiling_->stepKa;
@@ -1449,6 +1459,11 @@ __aicore__ inline void MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_
 
     int aMatrixByteSize = var.baseMK_ * sizeof(SrcT);
     int bMatrixByteSize = var.baseKN_ * sizeof(SrcT);
+    // The size of the int4b_t type obtained by sizeof is 1, but in reality, its actual size is 0.5
+    if constexpr (IsSameType<SrcT, int4b_t>::value) {
+        aMatrixByteSize = aMatrixByteSize / 2;
+        bMatrixByteSize = bMatrixByteSize / 2;
+    }
 
     if constexpr (A_TYPE::ibShare) {
         ASCENDC_ASSERT((B_TYPE::ibShare == false), {
-- 
Gitee