* Copyright (c) 2026 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.
*/
* Scale Kernel — ScaleKernel<T, RANK>
* RegBase Broadcast 范式,基于 DESIGN.md §3.5
* 目标架构: DAV_3510 (arch35)
*/
#pragma once
#include "kernel_operator.h"
#include "scale_tiling_struct.h"
#include "scale_struct.h"
template <typename T>
__simd_vf__ inline void ScaleMulAddVF(__ubuf__ T* dstAddr, __ubuf__ T* src0Addr, __ubuf__ T* src1Addr, uint32_t count,
uint32_t oneRepeatSize, uint16_t repeatTimes);
__aicore__ inline void GetCoreRange(int64_t core_id, int64_t tiles_main, int64_t cores_tail, int64_t& start,
int64_t& end)
{
if (core_id < cores_tail) {
start = core_id * (tiles_main + 1);
end = start + tiles_main + 1;
} else {
start = cores_tail * (tiles_main + 1) + (core_id - cores_tail) * tiles_main;
end = start + tiles_main;
}
}
__aicore__ inline int64_t GetUBSplitRange(int64_t a_o_off, int64_t a_o, int64_t a_i, int64_t a_i_tail)
{
return (a_o_off == a_o - 1) ? a_i_tail : a_i;
}
__aicore__ inline bool FlatToEffectiveCoord(int64_t flat, const int64_t* max_bro_shape, int64_t rank,
int64_t split_axis, int64_t a_i, int64_t a_o, int64_t* eff_coord)
{
for (int64_t d = 0; d < rank; d++)
eff_coord[d] = 0;
if (a_o <= 0)
return false;
int64_t a_o_off = flat % a_o;
int64_t outer = flat / a_o;
for (int64_t d = split_axis - 1; d >= 0; d--) {
eff_coord[d] = outer % max_bro_shape[d];
outer /= max_bro_shape[d];
}
eff_coord[split_axis] = a_o_off * a_i;
return true;
}
__aicore__ inline int64_t CalcInputOffset(const int64_t* eff_coord, const int64_t* strides, int64_t rank)
{
int64_t offset = 0;
for (int64_t d = 0; d < rank; d++)
offset += eff_coord[d] * strides[d];
return offset;
}
__aicore__ inline int64_t CalcOutputOffset(const int64_t* eff_coord, const int64_t* strides, int64_t rank)
{
int64_t offset = 0;
for (int64_t d = 0; d < rank; d++)
offset += eff_coord[d] * strides[d];
return offset;
}
template <typename T, int64_t RANK>
class ScaleKernel {
static constexpr int64_t ND = (RANK <= 5) ? RANK : 5;
static constexpr bool NEED_CAST = !std::is_same_v<T, float>;
static constexpr uint32_t VL_F32 = AscendC::GetVecLen() / sizeof(float);
AscendC::TPipe pipe_;
const ScaleTilingData<RANK>* td_;
AscendC::GlobalTensor<T> gmIn_[kMaxInputSlots];
AscendC::GlobalTensor<T> gmOut_[kMaxOutputSlots];
AscendC::TBuf<AscendC::TPosition::VECCALC> buf_[kPhysNodes];
AscendC::MultiCopyParams<T, ND> nddmaParams_[kMaxInputSlots];
int64_t nddmaDims_[kMaxInputSlots];
int64_t nddmaOuterBase_[kMaxInputSlots];
int64_t splitRunNd_[kMaxInputSlots];
int64_t splitRunStatic_[kMaxInputSlots];
public:
__aicore__ inline void Init(GM_ADDR inputs[kMaxInputSlots], GM_ADDR outputs[kMaxOutputSlots],
const ScaleTilingData<RANK>* td)
{
td_ = td;
for (int i = 0; i < kMaxInputSlots; i++)
gmIn_[i].SetGlobalBuffer((__gm__ T*)inputs[i]);
for (int i = 0; i < kMaxOutputSlots; i++)
gmOut_[i].SetGlobalBuffer((__gm__ T*)outputs[i]);
for (int i = 0; i < kPhysNodes; i++)
pipe_.InitBuffer(buf_[i], td_->per_buf_bytes);
const int64_t* dstShape = td_->max_bro_shape;
int64_t k = td_->split.axis;
for (int inp = 0; inp < kMaxInputSlots; inp++) {
int64_t runTop[RANK];
int64_t d = RANK - 1;
while (d >= k) {
int64_t top = d;
while (top > k && td_->input_strides[inp][top - 1] == dstShape[top] * td_->input_strides[inp][top]) {
top--;
}
for (int64_t t = top; t <= d; t++) {
runTop[t] = top;
}
d = top - 1;
}
int64_t tops[RANK];
int64_t topCnt = 0;
for (int64_t t = k; t < RANK; t++) {
if (runTop[t] == t) {
tops[topCnt++] = t;
}
}
int64_t keptRuns = (topCnt < ND) ? topCnt : ND;
int64_t keptLo = (topCnt > ND) ? tops[topCnt - ND] : k;
int64_t nd = 0;
int64_t inner = 1;
splitRunNd_[inp] = -1;
splitRunStatic_[inp] = 1;
for (int64_t ri = topCnt - 1; ri >= topCnt - keptRuns; ri--) {
int64_t top = tops[ri];
int64_t bot = (ri + 1 < topCnt) ? (tops[ri + 1] - 1) : (RANK - 1);
int64_t sz = (top == k) ? 1 : dstShape[top];
for (int64_t t = top + 1; t <= bot; t++) {
sz *= dstShape[t];
}
nddmaParams_[inp].loopInfo.loopSize[nd] = sz;
nddmaParams_[inp].loopInfo.loopSrcStride[nd] = td_->input_strides[inp][bot];
nddmaParams_[inp].loopInfo.loopDstStride[nd] = inner;
nddmaParams_[inp].loopInfo.loopLpSize[nd] = 0;
nddmaParams_[inp].loopInfo.loopRpSize[nd] = 0;
if (top == k) {
splitRunNd_[inp] = nd;
splitRunStatic_[inp] = sz;
}
inner *= sz;
nd++;
}
for (; nd < ND; nd++) {
nddmaParams_[inp].loopInfo.loopSize[nd] = 1;
nddmaParams_[inp].loopInfo.loopSrcStride[nd] = 0;
nddmaParams_[inp].loopInfo.loopDstStride[nd] = inner;
nddmaParams_[inp].loopInfo.loopLpSize[nd] = 0;
nddmaParams_[inp].loopInfo.loopRpSize[nd] = 0;
}
nddmaDims_[inp] = RANK - keptLo;
nddmaOuterBase_[inp] = 1;
for (int64_t t = k + 1; t < keptLo; t++) {
nddmaOuterBase_[inp] *= dstShape[t];
}
}
}
__aicore__ inline void Process()
{
if constexpr (NEED_CAST) {
ProcessWithCast();
} else {
ProcessFP32();
}
}
private:
struct PipeEvents {
int32_t mte2toV;
int32_t vtoMte2;
int32_t vtoMte3;
int32_t mte3toMte2;
};
__aicore__ inline PipeEvents FetchEvents()
{
return {
static_cast<int32_t>(GetTPipePtr()->FetchEventID(AscendC::HardEvent::MTE2_V)),
static_cast<int32_t>(GetTPipePtr()->FetchEventID(AscendC::HardEvent::V_MTE2)),
static_cast<int32_t>(GetTPipePtr()->FetchEventID(AscendC::HardEvent::V_MTE3)),
static_cast<int32_t>(GetTPipePtr()->FetchEventID(AscendC::HardEvent::MTE3_MTE2)),
};
}
__aicore__ inline int64_t ComputeInnerCount()
{
int64_t inner_count = 1;
for (int64_t d = td_->split.axis + 1; d < RANK; d++)
inner_count *= td_->max_bro_shape[d];
return inner_count;
}
__aicore__ inline void ProcessFP32()
{
PipeEvents ev = FetchEvents();
int64_t start, end;
GetCoreRange(AscendC::GetBlockIdx(), td_->multicore.tiles_main, td_->multicore.cores_tail, start, end);
int64_t inner_count = ComputeInnerCount();
int64_t coord[8] = {};
for (int64_t flat = start; flat < end; flat++) {
int64_t a_i_seg = GetUBSplitRange(flat % td_->split.a_o, td_->split.a_o, td_->split.a_i,
td_->split.a_i_tail);
int64_t count = a_i_seg * inner_count;
FlatToEffectiveCoord(flat, td_->max_bro_shape, RANK, td_->split.axis, td_->split.a_i, td_->split.a_o,
coord);
if (flat != start)
AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(ev.mte3toMte2);
TileFP32(coord, count, a_i_seg, ev, flat != end - 1);
}
}
__aicore__ inline void TileFP32(const int64_t* coord, int64_t count, int64_t a_i_seg, const PipeEvents& ev,
bool not_last)
{
constexpr int B0 = 0, B1 = 1, B2 = 2, B3 = 3;
constexpr int IN_X = 0, IN_SCALE = 1, IN_BIAS = 2, OUT_Y = 0;
CopyInBrc(coord, IN_X, B0, a_i_seg);
AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
CopyInBrc(coord, IN_SCALE, B1, a_i_seg);
AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
int outBuf = B3;
if (td_->has_bias) {
CopyInBrc(coord, IN_BIAS, B2, a_i_seg);
AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
uint16_t rep = AscendC::CeilDivision(count, VL_F32);
asc_vf_call<ScaleMulAddVF<float>>((__ubuf__ float*)buf_[B2].Get<float>().GetPhyAddr(),
(__ubuf__ float*)buf_[B0].Get<float>().GetPhyAddr(),
(__ubuf__ float*)buf_[B1].Get<float>().GetPhyAddr(), count, VL_F32, rep);
outBuf = B2;
} else {
AscendC::Mul(buf_[B3].Get<float>(), buf_[B0].Get<float>(), buf_[B1].Get<float>(), count);
}
AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(ev.vtoMte3);
AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(ev.vtoMte3);
CopyOutOne(coord, OUT_Y, outBuf, a_i_seg);
if (not_last)
AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(ev.mte3toMte2);
}
__aicore__ inline void ProcessWithCast()
{
PipeEvents ev = FetchEvents();
int64_t start, end;
GetCoreRange(AscendC::GetBlockIdx(), td_->multicore.tiles_main, td_->multicore.cores_tail, start, end);
int64_t inner_count = ComputeInnerCount();
int64_t coord[8] = {};
for (int64_t flat = start; flat < end; flat++) {
int64_t a_i_seg = GetUBSplitRange(flat % td_->split.a_o, td_->split.a_o, td_->split.a_i,
td_->split.a_i_tail);
int64_t count = a_i_seg * inner_count;
FlatToEffectiveCoord(flat, td_->max_bro_shape, RANK, td_->split.axis, td_->split.a_i, td_->split.a_o,
coord);
if (flat != start)
AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(ev.mte3toMte2);
TileCast(coord, count, a_i_seg, ev, flat != end - 1);
}
}
__aicore__ inline void TileCast(const int64_t* coord, int64_t count, int64_t a_i_seg, const PipeEvents& ev,
bool not_last)
{
constexpr int B0 = 0, B1 = 1, B2 = 2, B3 = 3;
constexpr int IN_X = 0, IN_SCALE = 1, IN_BIAS = 2, OUT_Y = 0;
CopyInBrc(coord, IN_X, B0, a_i_seg);
AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
AscendC::Cast(buf_[B1].template Get<float>(), buf_[B0].template Get<T>(), AscendC::RoundMode::CAST_NONE, count);
AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(ev.vtoMte2);
AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(ev.vtoMte2);
CopyInBrc(coord, IN_SCALE, B0, a_i_seg);
AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
AscendC::Cast(buf_[B2].template Get<float>(), buf_[B0].template Get<T>(), AscendC::RoundMode::CAST_NONE, count);
if (td_->has_bias) {
AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(ev.vtoMte2);
AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(ev.vtoMte2);
CopyInBrc(coord, IN_BIAS, B0, a_i_seg);
AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(ev.mte2toV);
AscendC::Cast(buf_[B3].template Get<float>(), buf_[B0].template Get<T>(), AscendC::RoundMode::CAST_NONE,
count);
uint16_t rep = AscendC::CeilDivision(count, VL_F32);
asc_vf_call<ScaleMulAddVF<float>>((__ubuf__ float*)buf_[B3].template Get<float>().GetPhyAddr(),
(__ubuf__ float*)buf_[B1].template Get<float>().GetPhyAddr(),
(__ubuf__ float*)buf_[B2].template Get<float>().GetPhyAddr(), count,
VL_F32, rep);
} else {
AscendC::Mul(buf_[B3].template Get<float>(), buf_[B1].template Get<float>(), buf_[B2].template Get<float>(),
count);
}
AscendC::Cast(buf_[B0].template Get<T>(), buf_[B3].template Get<float>(), AscendC::RoundMode::CAST_RINT, count);
AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(ev.vtoMte3);
AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(ev.vtoMte3);
CopyOutOne(coord, OUT_Y, B0, a_i_seg);
if (not_last)
AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(ev.mte3toMte2);
}
__aicore__ inline void CopyInBrc(const int64_t* coord, int inputIdx, int slot, int64_t a_i_seg)
{
int64_t off = CalcInputOffset(coord, td_->input_strides[inputIdx], RANK);
auto params = nddmaParams_[inputIdx];
if (splitRunNd_[inputIdx] >= 0) {
params.loopInfo.loopSize[splitRunNd_[inputIdx]] = splitRunStatic_[inputIdx] * a_i_seg;
}
static constexpr AscendC::NdDmaConfig cfg = {false, AscendC::NdDmaConfig::unsetPad,
AscendC::NdDmaConfig::unsetPad, false};
if (nddmaDims_[inputIdx] >= RANK - td_->split.axis) {
AscendC::DataCopy<T, ND, cfg>(buf_[slot].template Get<T>(), gmIn_[inputIdx][off], params);
} else {
AscendC::LocalTensor<T> buf = buf_[slot].template Get<T>();
int64_t inner = 1;
for (int64_t nd = 0; nd < ND; nd++) {
inner *= params.loopInfo.loopSize[nd];
}
const int64_t* dstShape = td_->max_bro_shape;
int64_t outerCnt = nddmaOuterBase_[inputIdx] * a_i_seg;
int64_t elem_base = off;
for (int64_t oi = 0; oi < outerCnt; oi++) {
int64_t elem_adj = 0, tmp = oi;
for (int64_t d = RANK - nddmaDims_[inputIdx] - 1; d >= td_->split.axis; d--) {
int64_t sz = (d == td_->split.axis) ? a_i_seg : dstShape[d];
elem_adj += (tmp % sz) * td_->input_strides[inputIdx][d];
tmp /= sz;
}
AscendC::DataCopy<T, ND, cfg>(buf[oi * inner], gmIn_[inputIdx][elem_base + elem_adj], params);
}
}
}
__aicore__ inline void CopyOutOne(const int64_t* coord, int outputIdx, int slot, int64_t a_i_seg)
{
int64_t off = CalcOutputOffset(coord, td_->output_strides[outputIdx], RANK);
int64_t split_elems = a_i_seg;
int64_t inner_elems = 1;
for (int64_t d = td_->split.axis + 1; d < RANK; d++)
inner_elems *= td_->output_shapes[outputIdx][d];
int64_t cnt = split_elems * inner_elems;
AscendC::DataCopyExtParams extParams;
extParams.blockCount = 1;
extParams.blockLen = cnt * sizeof(T);
extParams.srcStride = 0;
extParams.dstStride = 0;
AscendC::DataCopyPad(gmOut_[outputIdx][off], buf_[slot].template Get<T>(), extParams);
}
};
template <typename T>
__simd_vf__ inline void ScaleMulAddVF(__ubuf__ T* dstAddr, __ubuf__ T* src0Addr, __ubuf__ T* src1Addr, uint32_t count,
uint32_t oneRepeatSize, uint16_t repeatTimes)
{
AscendC::Reg::RegTensor<T> srcReg0, srcReg1, dstReg;
AscendC::Reg::MaskReg mask;
AscendC::Reg::AddrReg aReg;
for (uint16_t i = 0; i < repeatTimes; ++i) {
aReg = AscendC::Reg::CreateAddrReg<T>(i, oneRepeatSize);
uint32_t remain = count - (uint32_t)i * oneRepeatSize;
mask = AscendC::Reg::UpdateMask<T>(remain);
AscendC::Reg::LoadAlign(srcReg0, src0Addr, aReg);
AscendC::Reg::LoadAlign(srcReg1, src1Addr, aReg);
AscendC::Reg::LoadAlign(dstReg, dstAddr, aReg);
AscendC::Reg::MulAddDst(dstReg, srcReg0, srcReg1, mask);
AscendC::Reg::StoreAlign(dstAddr, dstReg, aReg, mask);
}
}