已合并
修复quant mm代码中的整型乘法溢出int32_t代码风险 #5374
oscillated创建于 5月28日
修复quant mm代码中的整型乘法溢出int32_t代码风险 #5374
已合并
共 2 个文件变更+15-13
| @@ -354,19 +354,19 @@ template <typename xType, typename wType, typename scaleType, typename yType, Qu | |||
| 354 | __aicore__ inline void QuantBatchMatmulV4Msd<xType, wType, scaleType, yType, quantType, bTrans, weightNz >::MMCompute(uint32_t mIdx, uint32_t nIdx, uint64_t workSpaceOffset) | 354 | __aicore__ inline void QuantBatchMatmulV4Msd<xType, wType, scaleType, yType, quantType, bTrans, weightNz >::MMCompute(uint32_t mIdx, uint32_t nIdx, uint64_t workSpaceOffset) |
| 355 | { | 355 | { |
| 356 | if ASCEND_IS_AIC { | 356 | if ASCEND_IS_AIC { |
| 357 | - uint64_t x1Offset = mIdx * baseM_ * kSize_; | 357 | + uint64_t x1Offset = static_cast<uint64_t>(mIdx) * baseM_ * kSize_; |
| 358 | uint64_t x2Offset = 0; | 358 | uint64_t x2Offset = 0; |
| 359 | if constexpr (bTrans == true) { | 359 | if constexpr (bTrans == true) { |
| 360 | if constexpr (weightNz == true) { | 360 | if constexpr (weightNz == true) { |
| 361 | - x2Offset = nIdx * baseN_ * 64; | 361 | + x2Offset = static_cast<uint64_t>(nIdx) * baseN_ * 64; |
| 362 | } else if constexpr (weightNz == false) { | 362 | } else if constexpr (weightNz == false) { |
| 363 | - x2Offset = nIdx * kSize_ * baseN_; | 363 | + x2Offset = static_cast<uint64_t>(nIdx) * kSize_ * baseN_; |
| 364 | } | 364 | } |
| 365 | } else if constexpr (bTrans == false) { | 365 | } else if constexpr (bTrans == false) { |
| 366 | if constexpr (weightNz == true) { | 366 | if constexpr (weightNz == true) { |
| 367 | - x2Offset = nIdx * baseN_ * ops::Aligned(kSize_, uint32_t(16)); | 367 | + x2Offset = static_cast<uint64_t>(nIdx) * baseN_ * ops::Aligned(kSize_, uint32_t(16)); |
| 368 | } else if constexpr (weightNz == false) { | 368 | } else if constexpr (weightNz == false) { |
| 369 | - x2Offset = nIdx * baseN_; | 369 | + x2Offset = static_cast<uint64_t>(nIdx) * baseN_; |
| 370 | } | 370 | } |
| 371 | } | 371 | } |
| 372 | uint32_t curSingleN = baseN_; | 372 | uint32_t curSingleN = baseN_; |
| @@ -397,13 +397,13 @@ __aicore__ inline void QuantBatchMatmulV4Msd<xType, wType, scaleType, yType, qua | |||
| 397 | mmObj_.SetSingleShape(curSingleM, curSingleN, groupSize_); | 397 | mmObj_.SetSingleShape(curSingleM, curSingleN, groupSize_); |
| 398 | GlobalTensor<wType> weightSlice; | 398 | GlobalTensor<wType> weightSlice; |
| 399 | for (uint32_t loopK = 0; loopK < groupNum_; loopK++) { | 399 | for (uint32_t loopK = 0; loopK < groupNum_; loopK++) { |
| 400 | - mmObj_.SetTensorA(x1Global_[x1Offset + loopK * groupSize_]); | 400 | + mmObj_.SetTensorA(x1Global_[x1Offset + static_cast<uint64_t>(loopK) * groupSize_]); |
| 401 | - auto weightSlice = x2Global_[x2Offset + loopK * groupSize_ * nSize_]; | 401 | + auto weightSlice = x2Global_[x2Offset + static_cast<uint64_t>(loopK) * groupSize_ * nSize_]; |
| 402 | if (blockDimM_ == 1) { | 402 | if (blockDimM_ == 1) { |
| 403 | weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); | 403 | weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); |
| 404 | } | 404 | } |
| 405 | mmObj_.SetTensorB(weightSlice); | 405 | mmObj_.SetTensorB(weightSlice); |
| 406 | - mmObj_.SetQuantVector(x2ScaleGlobal_[loopK * nSize_ + x2Offset]); | 406 | + mmObj_.SetQuantVector(x2ScaleGlobal_[static_cast<uint64_t>(loopK) * nSize_ + x2Offset]); |
| 407 | mmObj_.Iterate(); | 407 | mmObj_.Iterate(); |
| 408 | mmObj_.GetTensorC(mmOutGlobal_[workSpaceOffset], loopK == 0 ? 0 : 1, true); | 408 | mmObj_.GetTensorC(mmOutGlobal_[workSpaceOffset], loopK == 0 ? 0 : 1, true); |
| 409 | } | 409 | } |
| @@ -421,7 +421,7 @@ __aicore__ inline void QuantBatchMatmulV4Msd<xType, wType, scaleType, yType, qua | |||
| 421 | curCubeSingleN = nSize_ - nIdx * baseN_; | 421 | curCubeSingleN = nSize_ - nIdx * baseN_; |
| 422 | } | 422 | } |
| 423 | uint32_t curCubeSingleM = baseM_ / 2; // 2: 2 lines int4 to 1 line int8 | 423 | uint32_t curCubeSingleM = baseM_ / 2; // 2: 2 lines int4 to 1 line int8 |
| 424 | - uint64_t outOffset = mIdx * curCubeSingleM * tilingData_->nSize + nIdx * baseN_; | 424 | + uint64_t outOffset = static_cast<uint64_t>(mIdx) * curCubeSingleM * tilingData_->nSize + static_cast<uint64_t>(nIdx) * baseN_; |
| 425 | if (mIdx == blockDimM_ - 1) { | 425 | if (mIdx == blockDimM_ - 1) { |
| 426 | curCubeSingleM = mSize_ - mIdx * curCubeSingleM; | 426 | curCubeSingleM = mSize_ - mIdx * curCubeSingleM; |
| 427 | } | 427 | } |
| @@ -348,10 +348,12 @@ __aicore__ inline void WeightQuantBatchMatmulV2Common< | |||
| 348 | GM_ADDR x, GM_ADDR weight, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, GM_ADDR quantScale, | 348 | GM_ADDR x, GM_ADDR weight, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, GM_ADDR quantScale, |
| 349 | GM_ADDR quantOffset, GM_ADDR bias, GM_ADDR y) | 349 | GM_ADDR quantOffset, GM_ADDR bias, GM_ADDR y) |
| 350 | { | 350 | { |
| 351 | - xGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ xType*>(x), tiling_->matmulTiling.M * tiling_->matmulTiling.Ka); | 351 | + xGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ xType*>(x), |
| 352 | - wGlobal_.SetGlobalBuffer( | 352 | + static_cast<uint64_t>(tiling_->matmulTiling.M) * tiling_->matmulTiling.Ka); |
| 353 | - reinterpret_cast<__gm__ wType*>(weight), tiling_->matmulTiling.Kb * tiling_->matmulTiling.N); | 353 | + wGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ wType*>(weight), |
| 354 | - yGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ yType*>(y), tiling_->matmulTiling.M * tiling_->matmulTiling.N); | 354 | + static_cast<uint64_t>(tiling_->matmulTiling.Kb) * tiling_->matmulTiling.N); |
| 355 | + yGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ yType*>(y), | ||
| 356 | + static_cast<uint64_t>(tiling_->matmulTiling.M) * tiling_->matmulTiling.N); | ||
| 355 | biasGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ biasType*>(bias), tiling_->matmulTiling.N); | 357 | biasGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ biasType*>(bias), tiling_->matmulTiling.N); |
| 356 | biasFlag_ = static_cast<bool>(tiling_->matmulTiling.isBias); | 358 | biasFlag_ = static_cast<bool>(tiling_->matmulTiling.isBias); |
| 357 | offsetGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ xType*>(antiquantOffset), tiling_->matmulTiling.N); | 359 | offsetGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ xType*>(antiquantOffset), tiling_->matmulTiling.N); |