* Copyright (c) 2025 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/
* \file ascend_dequant_c310_impl.h
* \brief
*/
#if !defined(__ASCENDC_INCLUDE_INTERNAL_HEADERS__)
#pragma message("impl/adv_api/detail/quantization/dequant/ascend_dequant_c310_impl.h is an internal header file and must not be used directly. Functions or variables defined in this file may be removed in the future. Please use \"#include \"adv_api/quantization/ascend_dequant.h\"\" and use public functions or variables defined in interface headers files.")
#define __ASCENDC_INCLUDE_INTERNAL_HEADERS__
#define __UNDEF_ASCENDC_INCLUDE_INTERNAL_HEADERS_QUANTIZATION_DEQUANT_ASCEND_DEQUANT_C310_IMPL_H__
#endif
#ifndef LIB_ASCEND_DEQUANT_ASCEND_DEQUANT_C310_IMPL_H
#define LIB_ASCEND_DEQUANT_ASCEND_DEQUANT_C310_IMPL_H
#include "kernel_basic_intf.h"
#include "kernel_tensor.h"
#include "kernel_tiling/kernel_tiling.h"
#include "ascend_dequant_common.h"
#include "../../quantization/quant/ascend_quant_c310_impl.h"
#include "../../common/check.h"
namespace AscendC {
constexpr uint32_t ASCENDC_DEQUANT_B32_VF_LEN = GetVecLen() / sizeof(uint32_t);
template <typename dstT, typename scaleT, DeQuantMode mode>
__simd_vf__ inline void DequantPerchannelVFImpl(__ubuf__ half* dstUb, __ubuf__ int32_t* srcUb,
__ubuf__ float* scaleUb, DequantParams params)
{
uint32_t rowNum = params.m;
uint32_t N = params.n;
uint32_t calCount = params.calCount;
uint32_t oneBlockNum = ONE_BLK_SIZE / sizeof(dstT);
uint32_t dstInner = CeilDivision(N, oneBlockNum) * oneBlockNum;
Reg::MaskReg preg;
Reg::RegTensor<int32_t> s32vreg;
Reg::RegTensor<float> f32vreg;
Reg::RegTensor<half> b16vreg;
Reg::RegTensor<float> scaleB32Vreg0;
Reg::RegTensor<float> scaleB32Vreg1;
uint32_t sregLower = ASCENDC_DEQUANT_B32_VF_LEN;
uint16_t repeat = static_cast<uint16_t>(CeilDivision(calCount, sregLower));
for (uint16_t i = 0; i < static_cast<uint16_t>(rowNum); ++i) {
uint32_t sreg = calCount;
for (uint16_t j = 0; j < repeat; ++j) {
preg = Reg::UpdateMask<uint32_t>(sreg);
Reg::LoadAlign<int32_t, Reg::LoadDist::DIST_NORM>(s32vreg, srcUb + i * N + j * sregLower);
Reg::Cast<float, int32_t, MrgZRndA>(f32vreg, s32vreg, preg);
Reg::LoadAlign<float, Reg::LoadDist::DIST_DINTLV_B32>(scaleB32Vreg0, scaleB32Vreg1,
scaleUb + 2 * j * sregLower);
Reg::Mul(f32vreg, f32vreg, scaleB32Vreg0, preg);
Reg::Cast<dstT, float, LayoutZMrgZRndRSatS>(b16vreg, f32vreg, preg);
Reg::StoreAlign<dstT, Reg::StoreDist::DIST_PACK_B32>(dstUb + i * dstInner + j * sregLower, b16vreg,
preg);
}
}
}
template <typename dstT, typename scaleT, DeQuantMode mode>
__aicore__ inline void DequantPerchannelImpl(const LocalTensor<half>& dstTensor, const LocalTensor<int32_t>& srcTensor,
const LocalTensor<uint64_t>& deqScale, DequantParams& params)
{
CheckTensorPosition(dstTensor, "dstTensor", "VECIN, VECOUT, VECCALC");
CheckTensorPosition(srcTensor, "srcTensor", "VECIN, VECOUT, VECCALC");
CheckTensorPosition(deqScale, "deqScale", "VECIN, VECOUT, VECCALC");
__ubuf__ half* dstUb = (__ubuf__ half*)dstTensor.GetPhyAddr();
__ubuf__ int32_t* srcUb = (__ubuf__ int32_t*)srcTensor.GetPhyAddr();
__ubuf__ float* scaleUb = reinterpret_cast<__ubuf__ float*>(deqScale.GetPhyAddr());
DequantPerchannelVFImpl<dstT, scaleT, mode>(dstUb, srcUb, scaleUb, params);
}
template <typename dstT, typename scaleT, DeQuantMode mode>
__simd_vf__ inline void DequantPerchannelVFImpl(__ubuf__ dstT* dstUb, __ubuf__ int32_t* srcUb,
__ubuf__ scaleT* scaleUb, DequantParams params)
{
uint32_t rowNum = params.m;
uint32_t N = params.n;
uint32_t calCount = params.calCount;
uint32_t oneBlockNum = ONE_BLK_SIZE / sizeof(dstT);
uint32_t dstInner = CeilDivision(N, oneBlockNum) * oneBlockNum;
Reg::MaskReg preg;
Reg::RegTensor<int32_t> s32vreg;
Reg::RegTensor<float> f32vreg;
Reg::RegTensor<dstT> b16vreg;
Reg::RegTensor<scaleT> scaleVreg;
Reg::RegTensor<float> scaleB32Vreg;
uint32_t sregLower = ASCENDC_DEQUANT_B32_VF_LEN;
uint16_t repeat = static_cast<uint16_t>(CeilDivision(calCount, sregLower));
for (uint16_t i = 0; i < static_cast<uint16_t>(rowNum); ++i) {
uint32_t sreg = calCount;
for (uint16_t j = 0; j < repeat; ++j) {
preg = Reg::UpdateMask<uint32_t>(sreg);
Reg::LoadAlign<int32_t, Reg::LoadDist::DIST_NORM>(s32vreg, srcUb + i * N + j * sregLower);
Reg::Cast<float, int32_t, MrgZRndA>(f32vreg, s32vreg, preg);
if constexpr (SupportType<scaleT, bfloat16_t>()) {
Reg::LoadAlign<bfloat16_t, Reg::LoadDist::DIST_UNPACK_B16>(scaleVreg,
scaleUb + j * sregLower);
Reg::Cast<float, bfloat16_t, layoutZMrgZ>(scaleB32Vreg, scaleVreg, preg);
} else {
Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(scaleB32Vreg, scaleUb + j * sregLower);
}
Reg::Mul(f32vreg, f32vreg, scaleB32Vreg, preg);
if constexpr (SupportType<dstT, bfloat16_t, half>()) {
Reg::Cast<dstT, float, LayoutZMrgZRndRSatS>(b16vreg, f32vreg, preg);
Reg::StoreAlign<dstT, Reg::StoreDist::DIST_PACK_B32>(dstUb + i * dstInner + j * sregLower, b16vreg,
preg);
} else {
Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM_B32>(dstUb + i * dstInner + j * sregLower,
f32vreg, preg);
}
}
}
}
template <typename dstT, typename scaleT, DeQuantMode mode>
__aicore__ inline void DequantPerchannelImpl(const LocalTensor<dstT>& dstTensor, const LocalTensor<int32_t>& srcTensor,
const LocalTensor<scaleT>& deqScale, DequantParams& params)
{
CheckTensorPosition(dstTensor, "dstTensor", "VECIN, VECOUT, VECCALC");
CheckTensorPosition(srcTensor, "srcTensor", "VECIN, VECOUT, VECCALC");
CheckTensorPosition(deqScale, "deqScale", "VECIN, VECOUT, VECCALC");
__ubuf__ dstT* dstUb = (__ubuf__ dstT*)dstTensor.GetPhyAddr();
__ubuf__ int32_t* srcUb = (__ubuf__ int32_t*)srcTensor.GetPhyAddr();
__ubuf__ scaleT* scaleUb = (__ubuf__ scaleT*)deqScale.GetPhyAddr();
DequantPerchannelVFImpl<dstT, scaleT, mode>(dstUb, srcUb, scaleUb,params);
}
template <typename dstT, typename scaleT, DeQuantMode mode>
__simd_vf__ inline void DequantPertensorVFImpl(__ubuf__ dstT* dstUb, __ubuf__ int32_t* srcUb,
const scaleT deqScale, DequantParams params)
{
uint32_t rowNum = params.m;
uint32_t N = params.n;
uint32_t calCount = params.calCount;
uint32_t oneBlockNum = ONE_BLK_SIZE / sizeof(dstT);
uint32_t dstInner = CeilDivision(N, oneBlockNum) * oneBlockNum;
Reg::MaskReg preg;
Reg::RegTensor<int32_t> s32vreg;
Reg::RegTensor<float> f32vreg;
Reg::RegTensor<bfloat16_t> b16vreg;
uint32_t sregLower = ASCENDC_DEQUANT_B32_VF_LEN;
uint16_t repeat = static_cast<uint16_t>(CeilDivision(calCount, sregLower));
for (uint16_t i = 0; i < static_cast<uint16_t>(rowNum); ++i) {
uint32_t sreg = calCount;
for (uint16_t j = 0; j < repeat; ++j) {
preg = Reg::UpdateMask<uint32_t>(sreg);
Reg::LoadAlign<int32_t, Reg::LoadDist::DIST_NORM>(s32vreg, srcUb + i * N + j * sregLower);
Reg::Cast<float, int32_t, MrgZRndA>(f32vreg, s32vreg, preg);
if constexpr (SupportType<scaleT, bfloat16_t>()) {
Reg::Muls(f32vreg, f32vreg, ToFloat(deqScale), preg);
} else {
Reg::Muls(f32vreg, f32vreg, deqScale, preg);
}
if constexpr (SupportType<dstT, bfloat16_t>()) {
Reg::Cast<bfloat16_t, float, LayoutZMrgZRndRSatS>(b16vreg, f32vreg, preg);
Reg::StoreAlign<bfloat16_t, Reg::StoreDist::DIST_PACK_B32>(dstUb + i * dstInner + j * sregLower,
b16vreg, preg);
} else {
Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM_B32>(dstUb + i * dstInner + j * sregLower,
f32vreg, preg);
}
}
}
}
template <typename dstT, typename scaleT, DeQuantMode mode>
__aicore__ inline void DequantPertensorImpl(const LocalTensor<dstT>& dstTensor, const LocalTensor<int32_t>& srcTensor,
const scaleT deqScale, DequantParams& params)
{
CheckTensorPosition(dstTensor, "dstTensor", "VECIN, VECOUT, VECCALC");
CheckTensorPosition(srcTensor, "srcTensor", "VECIN, VECOUT, VECCALC");
__ubuf__ dstT* dstUb = (__ubuf__ dstT*)dstTensor.GetPhyAddr();
__ubuf__ int32_t* srcUb = (__ubuf__ int32_t*)srcTensor.GetPhyAddr();
DequantPertensorVFImpl<dstT, scaleT, mode>(dstUb, srcUb,deqScale, params);
}
template <typename scaleT>
__simd_callee__ inline void LoadPerTokenScale(__ubuf__ scaleT* addr, Reg::RegTensor<scaleT>& vreg)
{
if constexpr (SupportType<scaleT, half, bfloat16_t>()) {
Reg::LoadAlign<scaleT, Reg::LoadDist::DIST_BRC_B16>(vreg, addr);
} else {
Reg::LoadAlign<scaleT, Reg::LoadDist::DIST_BRC_B32>(vreg, addr);
}
}
template <typename dstT>
__simd_callee__ inline void StoreRes(__ubuf__ dstT* dstAddr, Reg::RegTensor<float>& vreg,
Reg::MaskReg& preg)
{
if constexpr (SupportType<dstT, half, bfloat16_t>()) {
Reg::RegTensor<dstT> tempVreg;
Reg::Cast<dstT, float, LayoutZMrgZRndRSatS>(tempVreg, vreg, preg);
Reg::StoreAlign<dstT, Reg::StoreDist::DIST_PACK_B32>(dstAddr, tempVreg, preg);
} else {
Reg::StoreAlign<dstT, Reg::StoreDist::DIST_NORM_B32>(dstAddr, vreg, preg);
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__simd_vf__ inline void DeQuantPerTokenForS32VF(__ubuf__ dstT* dstUb, __ubuf__ srcT* srcUb,
__ubuf__ scaleT* scaleUb, const AscendDeQuantParam para)
{
uint16_t rowNum = para.calCount / para.n;
uint32_t vecLen = ASCENDC_QUANT_B32_VF_LEN;
uint16_t repeat = static_cast<uint16_t>(CeilDivision(para.n, vecLen));
Reg::MaskReg preg;
Reg::RegTensor<int32_t> srcVreg;
Reg::RegTensor<float> f32Vreg;
Reg::RegTensor<scaleT> scaleVreg;
Reg::RegTensor<float> scaleF32Vreg;
for (uint16_t i = 0; i < static_cast<uint16_t>(rowNum); ++i) {
LoadPerTokenScale<scaleT>(scaleUb + i, scaleVreg);
uint32_t sreg = para.n;
for (uint16_t j = 0; j < repeat; ++j) {
preg = Reg::UpdateMask<uint32_t>(sreg);
Reg::LoadAlign<int32_t, Reg::LoadDist::DIST_NORM>(srcVreg, srcUb + i * para.n + j * vecLen);
Reg::Cast<float, int32_t, MrgZRndA>(f32Vreg, srcVreg, preg);
if constexpr (SupportType<scaleT, half, bfloat16_t>()) {
Reg::Cast<float, scaleT, layoutZMrgZ>(scaleF32Vreg, scaleVreg, preg);
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, f32Vreg, scaleF32Vreg, preg);
} else {
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, f32Vreg, scaleVreg, preg);
}
StoreRes<dstT>(dstUb + i * para.n + j * vecLen, f32Vreg, preg);
}
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__aicore__ inline void DeQuantPerTokenForS32(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor,
const LocalTensor<scaleT>& scaleTensor, const LocalTensor<scaleT>& offsetTensor,
const AscendDeQuantParam& para)
{
__ubuf__ dstT* dstUb = (__ubuf__ dstT*)dstTensor.GetPhyAddr();
__ubuf__ srcT* srcUb = (__ubuf__ srcT*)srcTensor.GetPhyAddr();
__ubuf__ scaleT* scaleUb = (__ubuf__ scaleT*)scaleTensor.GetPhyAddr();
DeQuantPerTokenForS32VF<dstT, srcT, scaleT, config>(dstUb, srcUb, scaleUb, para);
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__simd_vf__ inline void DeQuantPerTokenForF32VF(__ubuf__ dstT* dstUb, __ubuf__ srcT* srcUb,
__ubuf__ scaleT* scaleUb, const AscendDeQuantParam para)
{
uint16_t rowNum = para.calCount / para.n;
uint32_t vecLen = ASCENDC_QUANT_B32_VF_LEN;
uint16_t repeat = static_cast<uint16_t>(CeilDivision(para.n, vecLen));
Reg::MaskReg preg;
Reg::RegTensor<float> srcVreg;
Reg::RegTensor<float> f32Vreg;
Reg::RegTensor<scaleT> scaleVreg;
Reg::RegTensor<float> scaleF32Vreg;
for (uint16_t i = 0; i < static_cast<uint16_t>(rowNum); ++i) {
LoadPerTokenScale<scaleT>(scaleUb + i, scaleVreg);
uint32_t sreg = para.n;
for (uint16_t j = 0; j < repeat; ++j) {
preg = Reg::UpdateMask<uint32_t>(sreg);
Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(srcVreg, srcUb + i * para.n + j * vecLen);
if constexpr (SupportType<scaleT, half, bfloat16_t>()) {
Reg::Cast<float, scaleT, layoutZMrgZ>(scaleF32Vreg, scaleVreg, preg);
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, srcVreg, scaleF32Vreg, preg);
} else {
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, srcVreg, scaleVreg, preg);
}
StoreRes<dstT>(dstUb + i * para.n + j * vecLen, f32Vreg, preg);
}
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__aicore__ inline void DeQuantPerTokenForF32(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor,
const LocalTensor<scaleT>& scaleTensor, const LocalTensor<scaleT>& offsetTensor,
const AscendDeQuantParam& para)
{
__ubuf__ dstT* dstUb = (__ubuf__ dstT*)dstTensor.GetPhyAddr();
__ubuf__ srcT* srcUb = (__ubuf__ srcT*)srcTensor.GetPhyAddr();
__ubuf__ scaleT* scaleUb = (__ubuf__ scaleT*)scaleTensor.GetPhyAddr();
DeQuantPerTokenForF32VF<dstT, srcT, scaleT, config>(dstUb, srcUb, scaleUb, para);
}
template <typename T>
__simd_callee__ inline void GetPerGroupScale(__ubuf__ T* scaleUb, const int32_t start, const AscendDeQuantParam& para,
const AscendDeQuantConfig& config, Reg::RegTensor<T>& scaleReg)
{
uint32_t groupSize = para.groupSize;
if constexpr (SupportType<T, half, bfloat16_t>()) {
Reg::MaskReg preg = Reg::CreateMask<uint16_t, Reg::MaskPattern::ALL>();
Reg::RegTensor<int16_t> vci_vreg;
Reg::RegTensor<uint16_t> index_vreg;
Reg::RegTensor<uint16_t> gsize_vreg;
Reg::Duplicate(gsize_vreg, static_cast<uint16_t>(groupSize));
Reg::Arange(vci_vreg, static_cast<int16_t>(start));
Reg::Div(index_vreg, (Reg::RegTensor<uint16_t> &)vci_vreg, gsize_vreg, preg);
Reg::Gather(scaleReg, scaleUb, index_vreg, preg);
} else {
Reg::MaskReg preg = Reg::CreateMask<uint32_t, Reg::MaskPattern::ALL>();
Reg::RegTensor<int32_t> vci_vreg;
Reg::RegTensor<uint32_t> index_vreg;
Reg::RegTensor<uint32_t> gsize_vreg;
Reg::Duplicate(gsize_vreg, static_cast<uint32_t>(groupSize));
Reg::Arange(vci_vreg, static_cast<int32_t>(start));
Reg::Div(index_vreg, (Reg::RegTensor<uint32_t> &)vci_vreg, gsize_vreg, preg);
Reg::Gather(scaleReg, scaleUb, index_vreg, preg);
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__simd_vf__ inline void DeQuantPerGroupForColS32VF(__ubuf__ dstT* dstUb, __ubuf__ srcT* srcUb,
__ubuf__ scaleT* scaleUb, const AscendDeQuantParam para)
{
uint16_t rowNum = para.calCount / para.n;
uint32_t vecLen = ASCENDC_QUANT_B32_VF_LEN;
uint16_t repeat = static_cast<uint16_t>(CeilDivision(para.n, vecLen));
uint32_t sreg = para.n;
uint16_t scaleK = static_cast<uint16_t>(CeilDivision(para.n, para.groupSize));
Reg::MaskReg preg;
Reg::RegTensor<int32_t> srcVreg;
Reg::RegTensor<float> f32Vreg;
Reg::RegTensor<scaleT> oriScaleVreg;
Reg::RegTensor<scaleT> tempVreg;
Reg::RegTensor<int32_t> offsetVreg;
Reg::RegTensor<scaleT> scaleVreg;
Reg::RegTensor<float> scaleF32Vreg;
Reg::RegTensor<scaleT> zeroVreg;
if constexpr (SupportType<scaleT, half, bfloat16_t>()) {
Reg::MaskReg b16FullPreg = Reg::CreateMask<uint16_t, Reg::MaskPattern::ALL>();
Reg::Duplicate(zeroVreg, static_cast<scaleT>(0), b16FullPreg);
} else {
Reg::MaskReg b32FullPreg = Reg::CreateMask<uint32_t, Reg::MaskPattern::ALL>();
Reg::Duplicate(zeroVreg, static_cast<scaleT>(0), b32FullPreg);
}
for (uint16_t i = 0; i < static_cast<uint16_t>(rowNum); ++i) {
sreg = para.n;
for (uint16_t j = 0; j < repeat; ++j) {
preg = Reg::UpdateMask<uint32_t>(sreg);
if constexpr (SupportType<scaleT, half, bfloat16_t>()) {
GetPerGroupScale<scaleT>(scaleUb + i * scaleK, j * vecLen, para, config, oriScaleVreg);
Reg::Interleave(scaleVreg, tempVreg, oriScaleVreg, zeroVreg);
Reg::Cast<float, scaleT, layoutZMrgZ>(scaleF32Vreg, scaleVreg, preg);
} else {
GetPerGroupScale<scaleT>(scaleUb + i * scaleK, j * vecLen, para, config, scaleF32Vreg);
}
Reg::LoadAlign<int32_t, Reg::LoadDist::DIST_NORM>(srcVreg, srcUb + i * para.n + j * vecLen);
Reg::Cast<float, int32_t, MrgZRndA>(f32Vreg, srcVreg, preg);
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, f32Vreg, scaleF32Vreg, preg);
StoreRes<dstT>(dstUb + i * para.n + j * vecLen, f32Vreg, preg);
}
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__aicore__ inline void DeQuantPerGroupForColS32(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor,
const LocalTensor<scaleT>& scaleTensor, const LocalTensor<scaleT>& offsetTensor,
const AscendDeQuantParam& para)
{
__ubuf__ dstT* dstUb = (__ubuf__ dstT*)dstTensor.GetPhyAddr();
__ubuf__ srcT* srcUb = (__ubuf__ srcT*)srcTensor.GetPhyAddr();
__ubuf__ scaleT* scaleUb = (__ubuf__ scaleT*)scaleTensor.GetPhyAddr();
DeQuantPerGroupForColS32VF<dstT, srcT, scaleT, config>(dstUb, srcUb, scaleUb, para);
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__simd_vf__ inline void DeQuantPerGroupForColF32VF(__ubuf__ dstT* dstUb, __ubuf__ srcT* srcUb,
__ubuf__ scaleT* scaleUb, const AscendDeQuantParam para)
{
uint16_t rowNum = para.calCount / para.n;
uint32_t vecLen = ASCENDC_QUANT_B32_VF_LEN;
uint16_t repeat = static_cast<uint16_t>(CeilDivision(para.n, vecLen));
uint16_t scaleK = static_cast<uint16_t>(CeilDivision(para.n, para.groupSize));
Reg::MaskReg preg;
Reg::RegTensor<float> srcVreg;
Reg::RegTensor<float> f32Vreg;
Reg::RegTensor<scaleT> oriScaleVreg;
Reg::RegTensor<scaleT> tempVreg;
Reg::RegTensor<float> offsetVreg;
Reg::RegTensor<scaleT> scaleVreg;
Reg::RegTensor<float> scaleF32Vreg;
Reg::RegTensor<dstT> dstVreg;
Reg::RegTensor<scaleT> zeroVreg;
if constexpr (SupportType<scaleT, half, bfloat16_t>()) {
Reg::MaskReg b16FullPreg = Reg::CreateMask<uint16_t, Reg::MaskPattern::ALL>();
Reg::Duplicate(zeroVreg, static_cast<scaleT>(0), b16FullPreg);
} else {
Reg::MaskReg b32FullPreg = Reg::CreateMask<uint32_t, Reg::MaskPattern::ALL>();
Reg::Duplicate(zeroVreg, static_cast<scaleT>(0), b32FullPreg);
}
for (uint16_t i = 0; i < static_cast<uint16_t>(rowNum); ++i) {
uint32_t sreg = para.n;
for (uint16_t j = 0; j < repeat; ++j) {
preg = Reg::UpdateMask<uint32_t>(sreg);
if constexpr (SupportType<scaleT, half, bfloat16_t>()) {
GetPerGroupScale<scaleT>(scaleUb + i * scaleK, j * vecLen, para, config, oriScaleVreg);
Reg::Interleave(scaleVreg, tempVreg, oriScaleVreg, zeroVreg);
Reg::Cast<float, scaleT, layoutZMrgZ>(scaleF32Vreg, scaleVreg, preg);
} else {
GetPerGroupScale<scaleT>(scaleUb + i * scaleK, j * vecLen, para, config, scaleF32Vreg);
}
Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(srcVreg, srcUb + i * para.n + j * vecLen);
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, srcVreg, scaleF32Vreg, preg);
StoreRes<dstT>(dstUb + i * para.n + j * vecLen, f32Vreg, preg);
}
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__aicore__ inline void DeQuantPerGroupForColF32(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor,
const LocalTensor<scaleT>& scaleTensor, const LocalTensor<scaleT>& offsetTensor,
const AscendDeQuantParam& para)
{
__ubuf__ dstT* dstUb = (__ubuf__ dstT*)dstTensor.GetPhyAddr();
__ubuf__ srcT* srcUb = (__ubuf__ srcT*)srcTensor.GetPhyAddr();
__ubuf__ scaleT* scaleUb = (__ubuf__ scaleT*)scaleTensor.GetPhyAddr();
DeQuantPerGroupForColF32VF<dstT, srcT, scaleT, config>(dstUb, srcUb, scaleUb, para);
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__simd_callee__ inline void DeQuantPerGroupForRowTailBlock(__ubuf__ dstT* dstUb, __ubuf__ srcT* srcUb,
__ubuf__ scaleT* scaleUb, uint16_t repeat,
uint16_t tailRow, uint32_t n, uint32_t vecLen)
{
Reg::MaskReg preg;
Reg::RegTensor<scaleT> scaleVreg;
Reg::RegTensor<float> f32ScaleVreg;
Reg::RegTensor<srcT> srcVreg;
Reg::RegTensor<float> f32Vreg;
Reg::MaskReg b32FullPreg = Reg::CreateMask<uint32_t, Reg::MaskPattern::ALL>();
for (uint16_t i = 0; i < tailRow; ++i) {
uint32_t sreg = n;
for (uint16_t j = 0; j < repeat; ++j) {
if constexpr (SupportType<scaleT, half, bfloat16_t>()) {
Reg::LoadAlign<scaleT, Reg::LoadDist::DIST_UNPACK_B16>(scaleVreg, scaleUb + j * vecLen);
Reg::Cast<float, scaleT, layoutZMrgZ>(f32ScaleVreg, scaleVreg, b32FullPreg);
} else {
Reg::LoadAlign<scaleT, Reg::LoadDist::DIST_NORM>(f32ScaleVreg, scaleUb + j * vecLen);
}
preg = Reg::UpdateMask<uint32_t>(sreg);
Reg::LoadAlign<srcT, Reg::LoadDist::DIST_NORM>(srcVreg, srcUb + i * n + j * vecLen);
if constexpr (SupportType<srcT, int32_t>()) {
Reg::Cast<float, int32_t, MrgZRndA>(f32Vreg, srcVreg, preg);
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, f32Vreg, f32ScaleVreg, preg);
} else {
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, srcVreg, f32ScaleVreg, preg);
}
StoreRes<dstT>(dstUb + i * n + j * vecLen, f32Vreg, preg);
}
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__simd_vf__ inline void DeQuantPerGroupForRowVF(__ubuf__ dstT* dstUb, __ubuf__ srcT* srcUb,
__ubuf__ scaleT* scaleUb, const AscendDeQuantParam para, uint16_t rowNum, uint16_t tailRow)
{
uint16_t mainRowGroup = rowNum / para.groupSize;
uint32_t vecLen = ASCENDC_QUANT_B32_VF_LEN;
uint16_t repeat = static_cast<uint16_t>(CeilDivision(para.n, vecLen));
Reg::MaskReg preg;
Reg::RegTensor<scaleT> scaleVreg;
Reg::RegTensor<float> f32ScaleVreg;
Reg::RegTensor<srcT> srcVreg;
Reg::RegTensor<float> f32Vreg;
Reg::MaskReg b32FullPreg = Reg::CreateMask<uint32_t, Reg::MaskPattern::ALL>();
for (uint16_t i = 0; i < mainRowGroup; ++i) {
for (uint16_t j = 0; j < static_cast<uint16_t>(para.groupSize); ++j) {
uint32_t sreg = para.n;
for (uint16_t k = 0; k < repeat; ++k) {
if constexpr (SupportType<scaleT, half, bfloat16_t>()) {
Reg::LoadAlign<scaleT, Reg::LoadDist::DIST_UNPACK_B16>(scaleVreg, scaleUb + i * para.n + k * vecLen);
Reg::Cast<float, scaleT, layoutZMrgZ>(f32ScaleVreg, scaleVreg, b32FullPreg);
} else {
Reg::LoadAlign<scaleT, Reg::LoadDist::DIST_NORM>(
f32ScaleVreg, scaleUb + i * para.n + k * vecLen);
}
preg = Reg::UpdateMask<uint32_t>(sreg);
Reg::LoadAlign<srcT, Reg::LoadDist::DIST_NORM>(
srcVreg, srcUb + (i * para.groupSize + j) * para.n + k * vecLen);
if constexpr (SupportType<srcT, int32_t>()) {
Reg::Cast<float, int32_t, MrgZRndA>(f32Vreg, srcVreg, preg);
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, f32Vreg, f32ScaleVreg, preg);
} else {
Reg::Mul<float, Reg::MaskMergeMode::ZEROING>(f32Vreg, srcVreg, f32ScaleVreg, preg);
}
StoreRes<dstT>(dstUb + (i * para.groupSize + j) * para.n + k * vecLen, f32Vreg, preg);
}
}
}
DeQuantPerGroupForRowTailBlock<dstT, srcT, scaleT, config>(
dstUb + mainRowGroup * para.groupSize * para.n, srcUb + mainRowGroup * para.groupSize * para.n,
scaleUb + mainRowGroup * para.n, repeat, tailRow, para.n, vecLen);
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__aicore__ inline void DeQuantPerGroupForRow(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor,
const LocalTensor<scaleT>& scaleTensor, const LocalTensor<scaleT>& offsetTensor,
const AscendDeQuantParam& para)
{
__ubuf__ dstT* dstUb = (__ubuf__ dstT*)dstTensor.GetPhyAddr();
__ubuf__ srcT* srcUb = (__ubuf__ srcT*)srcTensor.GetPhyAddr();
__ubuf__ scaleT* scaleUb = (__ubuf__ scaleT*)scaleTensor.GetPhyAddr();
uint16_t rowNum = para.calCount / para.n;
uint16_t tailRow = rowNum % para.groupSize;
DeQuantPerGroupForRowVF<dstT, srcT, scaleT, config>(dstUb, srcUb, scaleUb, para, rowNum, tailRow);
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__aicore__ inline void AscendDeQuantPerToken(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor,
const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<scaleT>& scaleTensor,
const LocalTensor<scaleT>& offsetTensor, const AscendDeQuantParam& para)
{
if constexpr (SupportType<srcT, int32_t>()) {
DeQuantPerTokenForS32<dstT, srcT, scaleT, config>(dstTensor, srcTensor, scaleTensor, offsetTensor, para);
} else if constexpr (SupportType<srcT, float>()) {
DeQuantPerTokenForF32<dstT, srcT, scaleT, config>(dstTensor, srcTensor, scaleTensor, offsetTensor, para);
} else {
ASCENDC_ASSERT((false), { KERNEL_LOG(KERNEL_ERROR, "unsupport srcT for AscendDeQuant!"); });
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__aicore__ inline void AscendDeQuantPerGroupForCol(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor,
const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<scaleT>& scaleTensor,
const LocalTensor<scaleT>& offsetTensor, const AscendDeQuantParam& para)
{
if constexpr (SupportType<srcT, int32_t>()) {
DeQuantPerGroupForColS32<dstT, srcT, scaleT, config>(dstTensor, srcTensor, scaleTensor, offsetTensor, para);
} else if constexpr (SupportType<srcT, float>()) {
DeQuantPerGroupForColF32<dstT, srcT, scaleT, config>(dstTensor, srcTensor, scaleTensor, offsetTensor, para);
} else {
ASCENDC_ASSERT((false), { KERNEL_LOG(KERNEL_ERROR, "unsupport srcT for AscendDeQuant!"); });
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config>
__aicore__ inline void AscendDeQuantPerGroupForRow(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor,
const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<scaleT>& scaleTensor,
const LocalTensor<scaleT>& offsetTensor, const AscendDeQuantParam& para)
{
if constexpr (SupportType<srcT, int32_t, float>()) {
DeQuantPerGroupForRow<dstT, srcT, scaleT, config>(dstTensor, srcTensor, scaleTensor, offsetTensor, para);
} else {
ASCENDC_ASSERT((false), { KERNEL_LOG(KERNEL_ERROR, "unsupport srcT for AscendDeQuant!"); });
}
}
template <typename dstT, typename srcT, typename scaleT, const AscendDeQuantConfig& config, const AscendDeQuantPolicy& policy>
__aicore__ inline void AscendDequantImpl(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor,
const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<scaleT>& scaleTensor,
const LocalTensor<scaleT>& offsetTensor, const AscendDeQuantParam& para)
{
if ASCEND_IS_AIC {
return;
}
CheckTensorPosition(dstTensor, "dstTensor", "VECIN, VECOUT, VECCALC");
CheckTensorPosition(srcTensor, "srcTensor", "VECIN, VECOUT, VECCALC");
CheckTensorPosition(scaleTensor, "scaleTensor", "VECIN, VECOUT, VECCALC");
CheckTensorPosition(offsetTensor, "offsetTensor", "VECIN, VECOUT, VECCALC");
static_assert(SupportType<srcT, int32_t, float>(),
"AscendDequant only support int32_t/float input dtype");
static_assert(SupportType<dstT, bfloat16_t, half, float>(),
"AscendDequant only support bfloat16_t/half/float output dtype");
static_assert(SupportType<scaleT, bfloat16_t, half, float>(),
"AscendDequant only support bfloat16_t/half/float scaleT dtype");
static_assert(((policy == AscendDeQuantPolicy::PER_TOKEN) || (policy == AscendDeQuantPolicy::PER_GROUP)),
"unsupported policy for AscendDequant in current device!");
ASCENDC_ASSERT((para.calCount <= srcTensor.GetSize() && para.calCount <= dstTensor.GetSize() && para.calCount >= 0), {
KERNEL_LOG(KERNEL_ERROR, "calCount is %u, which should be in [0, min(%u, %u)]",
para.calCount, srcTensor.GetSize(), dstTensor.GetSize());
});
if constexpr (policy == AscendDeQuantPolicy::PER_TOKEN) {
AscendDeQuantPerToken<dstT, srcT, scaleT, config>(dstTensor, srcTensor, sharedTmpBuffer, scaleTensor, offsetTensor, para);
} else if constexpr (policy == AscendDeQuantPolicy::PER_GROUP) {
static_assert(
((config.kDim == 0) || (config.kDim == 1)), "AscendDequant PerGroup only support kDim is axis 0/1!");
ASCENDC_ASSERT((para.groupSize > 0 && para.groupSize % 32 == 0),
{ KERNEL_LOG(KERNEL_ERROR, "groupSize must be an integer multiple of 32 and greater than 0 !"); });
if constexpr (config.kDim == 1) {
AscendDeQuantPerGroupForCol<dstT, srcT, scaleT, config>(dstTensor, srcTensor, sharedTmpBuffer, scaleTensor, offsetTensor, para);
} else {
AscendDeQuantPerGroupForRow<dstT, srcT, scaleT, config>(dstTensor, srcTensor, sharedTmpBuffer, scaleTensor, offsetTensor, para);
}
}
}
}
#endif
#if defined(__UNDEF_ASCENDC_INCLUDE_INTERNAL_HEADERS_QUANTIZATION_DEQUANT_ASCEND_DEQUANT_C310_IMPL_H__)
#undef __ASCENDC_INCLUDE_INTERNAL_HEADERS__
#undef __UNDEF_ASCENDC_INCLUDE_INTERNAL_HEADERS_QUANTIZATION_DEQUANT_ASCEND_DEQUANT_C310_IMPL_H__
#endif
