EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
csrc/moe/hc_pre_inv_rms/CMakeLists.txt Normal file
View File

@@ -0,0 +1,19 @@
# -----------------------------------------------------------------------------------------------------------
# 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(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
if(NOT ENABLE_TEST AND NOT BENCHMARK)
list(REMOVE_ITEM CURRENT_DIRS tests)
endif()
foreach(SUB_DIR ${CURRENT_DIRS})
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
add_subdirectory(${SUB_DIR})
endif()
endforeach()

62
csrc/moe/hc_pre_inv_rms/op_host/CMakeLists.txt Normal file
View File

@@ -0,0 +1,62 @@
# This program is free software, you can redistribute it and/or modify it.
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under 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.
# ======================================================================================================================
# add_ops_compile_options(
# OP_NAME HcPreInvRms
# OPTIONS --cce-auto-sync=off
# -Wno-deprecated-declarations
# -Werror
# -mllvm -cce-aicore-hoist-movemask=false
# --op_relocatable_kernel_binary=true
# )
# set(hc_pre_inv_rms_depends nn/norm/hc_pre_inv_rms PARENT_SCOPE)
# target_sources(op_host_aclnn PRIVATE
# op_host/hc_pre_inv_rms_def.cpp
# )
# target_sources(optiling PRIVATE
# op_host/hc_pre_inv_rms_tiling.cpp
# )
# if (NOT BUILD_OPEN_PROJECT)
# target_sources(opmaster_ct PRIVATE
# op_host/hc_pre_inv_rms_tiling.cpp
# )
# endif ()
# target_include_directories(optiling PRIVATE
# ${CMAKE_CURRENT_SOURCE_DIR}/op_host
# )
# target_sources(opsproto PRIVATE
# op_host/hc_pre_inv_rms_proto.cpp
# )
add_op_to_compiled_list()
if (BUILD_OPEN_PROJECT)
target_sources(op_host_aclnn PRIVATE
hc_pre_inv_rms_def.cpp
)
endif()
add_ops_compile_options(
OP_NAME HcPreInvRms
OPTIONS --cce-auto-sync=off
-Wno-deprecated-declarations
-mllvm -cce-aicore-hoist-movemask=false
--op_relocatable_kernel_binary=true
)
if (NOT BUILD_OPS_RTY_KERNEL)
add_modules_sources(OPTYPE hc_pre_inv_rms ACLNNTYPE aclnn)
endif()

51
csrc/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_def.cpp Normal file
View File

@@ -0,0 +1,51 @@
/**
* 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.
*/
/*!
* \file hc_pre_inv_rms.cpp
* \brief HcPreInvRms op host config
*/
#include <cstdint>
#include "register/op_def_registry.h"
namespace ops {
class HcPreInvRms : public OpDef {
public:
explicit HcPreInvRms(const char *name) : OpDef(name)
{
this->Input("x")
.ParamType(REQUIRED)
.DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
this->Output("y")
.ParamType(REQUIRED)
.DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
this->Attr("epsilon")
.AttrType(OPTIONAL)
.Float(1e-6f);
this->AICore().AddConfig("ascend910b");
this->AICore().AddConfig("ascend910_93");
OpAICoreConfig regbaseCfg;
regbaseCfg.DynamicCompileStaticFlag(true)
.DynamicRankSupportFlag(true)
.DynamicShapeSupportFlag(true)
.ExtendCfgInfo("opFile.value", "hc_pre_inv_rms");
this->AICore().AddConfig("ascend950", regbaseCfg);
}
};
OP_ADD(HcPreInvRms);
} // namespace ops

67
csrc/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_proto.cpp Normal file
View File

@@ -0,0 +1,67 @@
/**
* 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.
*/
/*!
* \file hc_pre_inv_rms_proto.cpp
* \brief
*/
#include <graph/utils/type_utils.h>
#include <register/op_impl_registry.h>
#include "error/ops_error.h"
using namespace ge;
namespace ops {
const int32_t INPUT_IDX_X = 0;
const int32_t INDEX_OUTPUT_Y = 0;
const static int64_t DIM_0 = 0;
const static int64_t DIM_1 = 1;
const static int64_t DIM_2 = 2;
const static int64_t DIM_3 = 3;
const static int64_t BS_INPUT_DIM_NUM = 4;
const static int64_t TND_INPUT_DIM_NUM = 3;
static ge::graphStatus InferShape4HcPreInvRms(gert::InferShapeContext* context)
{
OPS_LOG_I(context->GetNodeName(), "Begin to do InferShape4HcPreInvRms.");
const gert::Shape* xShape = context->GetInputShape(INPUT_IDX_X);
OPS_LOG_E_IF_NULL(context, xShape, return ge::GRAPH_FAILED);
auto xDimNum = xShape->GetDimNum();
auto yShape = context->GetOutputShape(INDEX_OUTPUT_Y);
// The first one or two dimensions of y match those of x, and the last dimension of y is 1.
// x: (b, s, hc, d) --> y: (b, s, 1) or x: (b * s, hc, d) --> y: (b * s, 1)
yShape->SetDimNum(xDimNum);
if (xDimNum == BS_INPUT_DIM_NUM) {
yShape->SetDim(DIM_0, xShape->GetDim(DIM_0));
yShape->SetDim(DIM_1, xShape->GetDim(DIM_1));
yShape->SetDim(DIM_2, 1);
} else if (xDimNum == TND_INPUT_DIM_NUM) {
yShape->SetDim(DIM_0, xShape->GetDim(DIM_0));
yShape->SetDim(DIM_1, 1);
}
OPS_LOG_I(context->GetNodeName(), "End to do InferShape4HcPreInvRms");
return ge::GRAPH_SUCCESS;
}
static ge::graphStatus InferDtype4HcPreInvRms(gert::InferDataTypeContext* context)
{
OPS_LOG_I(context->GetNodeName(), "InferDtype4HcHost enter");
context->SetOutputDataType(INDEX_OUTPUT_Y, ge::DT_FLOAT);
OPS_LOG_I(context->GetNodeName(), "InferDtype4HcPreInvRms end");
return GRAPH_SUCCESS;
}
IMPL_OP_INFERSHAPE(HcPreInvRms)
.InferShape(InferShape4HcPreInvRms)
.InferDataType(InferDtype4HcPreInvRms);
} // namespace ops

359
csrc/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling.cpp Normal file
View File

@@ -0,0 +1,359 @@
/**
* 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.
*/
/* !
* \file hc_pre_inv_rms_tiling.cpp
* \brief
*/
#include <graph/utils/type_utils.h>
#include "hc_pre_inv_rms_tiling.h"
#include "hc_pre_inv_rms_tiling_arch35.h"
#include "hc_pre_inv_rms_tiling_large_d.h"
namespace optiling {
const static int64_t DEFAULT_WORKSPACE_SIZE = 16777216; // 预留16M空间
const static int64_t X_INPUT_INDEX = 0;
const static int64_t Y_OUTPUT_INDEX = 0;
const static int64_t EPS_ATTR_INDEX = 0;
const static size_t X_INPUT_BS_FUSED_DIMS = 3;
const static size_t X_INPUT_DIMS = 4;
const static int64_t UB_BLOCK_SIZE = 32;
const static uint64_t TILING_KEY_FULL_LOAD = 1000;
const static int64_t DIM_0 = 0;
const static int64_t DIM_1 = 1;
const static int64_t DIM_2 = 2;
const static int64_t DIM_3 = 3;
const static int64_t B16_TYPE_BYTE_SIZE = 2;
const static int64_t B32_TYPE_BYTE_SIZE = 4;
template <typename T>
static inline T CeilDiv(T num, T rnd)
{
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)));
}
template <typename T>
static inline T CeilAlign(T num, T rnd)
{
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd));
}
class HcPreInvRmsTilingBase {
public:
explicit HcPreInvRmsTilingBase(gert::TilingContext *context) : context_(context)
{
Reset();
}
~HcPreInvRmsTilingBase() = default;
bool IsCapable()
{
return true;
}
// 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小
ge::graphStatus GetPlatformInfo();
// 2、获取INPUT/OUTPUT/ATTR信息
ge::graphStatus GetShapeAttrsInfo();
// 3、计算数据切分TilingData
ge::graphStatus DoOpTiling();
// 4、计算高阶API的TilingData
ge::graphStatus DoLibApiTiling();
// 5、计算TilingKey
uint64_t GetTilingKey() const;
// 6、计算Workspace 大小
ge::graphStatus GetWorkspaceSize();
// 7、保存Tiling数据
ge::graphStatus PostTiling();
void Reset();
private:
ge::graphStatus CheckInputShape();
ge::graphStatus CheckAttr();
ge::graphStatus CheckOutShape();
void SplitA();
void CalUbFactorA();
const gert::Shape *xShape_ = nullptr;
const gert::Shape *yShape_ = nullptr;
float eps_ = 1e-6f;
int64_t A_ = 0;
int64_t R_ = 0;
int64_t inputDtypeSize_;
int64_t outputDtypeSize_;
const char *opName_ = "";
HcPreInvRmsFullLoadTilingData invRmsTilingData_;
gert::TilingContext *context_ = nullptr;
uint64_t workspaceSize_ = 0;
uint64_t coreNum_ = 0;
int64_t ubSize_ = 0;
int64_t ubBlockSize_ = 0;
};
ge::graphStatus HcPreInvRmsTilingBase::CheckInputShape()
{
size_t xDimNum = xShape_->GetDimNum();
OPS_ERR_IF(xDimNum != X_INPUT_DIMS && xDimNum != X_INPUT_BS_FUSED_DIMS,
OPS_LOG_E(context_, "The dim number of x is: %zu, but it should be %zu or %zu(bs fused)."
, xDimNum, X_INPUT_DIMS, X_INPUT_BS_FUSED_DIMS),
return ge::GRAPH_FAILED);
if (xDimNum == X_INPUT_DIMS) {
A_ = xShape_->GetDim(DIM_0) * xShape_->GetDim(DIM_1);
R_ = xShape_->GetDim(DIM_2) * xShape_->GetDim(DIM_3);
} else if (xDimNum == X_INPUT_BS_FUSED_DIMS) {
A_ = xShape_->GetDim(DIM_0);
R_ = xShape_->GetDim(DIM_1) * xShape_->GetDim(DIM_2);
}
invRmsTilingData_.set_A(A_);
invRmsTilingData_.set_R(R_);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingBase::CheckAttr()
{
OPS_ERR_IF(eps_ <= 0, OPS_LOG_E(context_, "epsilon is: %ld, but it should not be less than 0.", eps_), return ge::GRAPH_FAILED);
invRmsTilingData_.set_epsilon(eps_);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingBase::GetShapeAttrsInfo()
{
opName_ = context_->GetNodeName();
// 获取输入shape信息
auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED);
xShape_ = &xShapePtr->GetStorageShape();
// 获取输出shape
auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED);
yShape_ = &yShapePtr->GetStorageShape();
// 获取输入dtype
auto xDesc = context_->GetInputDesc(X_INPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, xDesc, return ge::GRAPH_FAILED);
auto xDtype = xDesc->GetDataType();
OPS_ERR_IF(
(xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16),
OPS_LOG_E(context_, "x dtype %s error, only supports float32, float16 and bfloat16. please check.",
ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()),
return ge::GRAPH_FAILED);
// 获取输出dtype
auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED);
auto yDtype = yDesc->GetDataType();
OPS_ERR_IF((yDtype != ge::DataType::DT_FLOAT),
OPS_LOG_E(context_, "y out dtype %s error, only support float32, please check",
ge::TypeUtils::DataTypeToSerialString(yDtype).c_str()),
return ge::GRAPH_FAILED);
// 获取属性
auto attrs = context_->GetAttrs();
OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED);
const float *epsPtr = attrs->GetAttrPointer<float>(EPS_ATTR_INDEX);
if (epsPtr != nullptr) {
eps_ = *epsPtr;
}
OPS_LOG_I(context_, "Attr eps is: %f ", eps_);
inputDtypeSize_ = static_cast<int64_t>(ge::GetSizeByDataType(context_->GetInputDesc(X_INPUT_INDEX)->GetDataType()));
outputDtypeSize_ = static_cast<int64_t>(ge::GetSizeByDataType(context_->GetOutputDesc(Y_OUTPUT_INDEX)->GetDataType()));
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingBase::GetPlatformInfo()
{
auto platformInfo = context_->GetPlatformInfo();
OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
coreNum_ = ascendcPlatform.GetCoreNumAiv();
OPS_ERR_IF(
coreNum_ <= 0, OPS_LOG_E(context_->GetNodeName(), "coreNum must be greater than 0."),
return ge::GRAPH_FAILED);
// 获取UB大小
uint64_t ubSizePlatForm;
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
ubSize_ = static_cast<int64_t>(ubSizePlatForm);
OPS_ERR_IF(
ubSize_ <= 0, OPS_LOG_E(context_->GetNodeName(), "ubSize must be greater than 0."),
return ge::GRAPH_FAILED);
ubBlockSize_ = UB_BLOCK_SIZE; // 32: ub block size
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingBase::CheckOutShape()
{
OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)),
OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0),
xShape_->GetDim(0)),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
void HcPreInvRmsTilingBase::SplitA()
{
int64_t blockFactorA = CeilDiv(A_, static_cast<int64_t>(coreNum_));
int64_t blockNumA = CeilDiv(A_, blockFactorA);
int64_t blockTailFactorA = A_ % blockFactorA == 0 ? blockFactorA : A_ % blockFactorA;
invRmsTilingData_.set_blockNumA(blockNumA);
invRmsTilingData_.set_blockFactorA(blockFactorA);
invRmsTilingData_.set_blockTailFactorA(blockTailFactorA);
int64_t ubFactorA = invRmsTilingData_.get_ubFactorA();
if (ubFactorA > blockFactorA) {
invRmsTilingData_.set_ubFactorA(blockFactorA);
}
}
void HcPreInvRmsTilingBase::CalUbFactorA()
{
int64_t rAlignSize = CeilAlign(R_ * inputDtypeSize_, UB_BLOCK_SIZE);
int64_t ubFactorA = 1;
if (inputDtypeSize_ == B16_TYPE_BYTE_SIZE) {
ubFactorA = ubSize_ / (4 * rAlignSize + 2 * outputDtypeSize_ + R_ / 16);
} else if (inputDtypeSize_ == B32_TYPE_BYTE_SIZE) {
ubFactorA = ubSize_ / (2 * rAlignSize + 2 * outputDtypeSize_ + R_ / 16);
}
invRmsTilingData_.set_ubFactorA(ubFactorA);
}
ge::graphStatus HcPreInvRmsTilingBase::DoOpTiling()
{
auto ret = GetPlatformInfo();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = GetShapeAttrsInfo();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = CheckInputShape();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = CheckOutShape();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = CheckAttr();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
CalUbFactorA();
SplitA();
ret = PostTiling();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingBase::DoLibApiTiling()
{
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingBase::GetWorkspaceSize()
{
// 计算workspace大小
workspaceSize_ = DEFAULT_WORKSPACE_SIZE;
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingBase::PostTiling()
{
context_->SetTilingKey(GetTilingKey());
context_->SetBlockDim(invRmsTilingData_.get_blockNumA());
size_t *currentWorkspace = context_->GetWorkspaceSizes(1);
currentWorkspace[0] = workspaceSize_;
invRmsTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(),
context_->GetRawTilingData()->GetCapacity());
context_->GetRawTilingData()->SetDataSize(invRmsTilingData_.GetDataSize());
return ge::GRAPH_SUCCESS;
}
uint64_t HcPreInvRmsTilingBase::GetTilingKey() const
{
return TILING_KEY_FULL_LOAD;
}
void HcPreInvRmsTilingBase::Reset()
{
opName_ = nullptr;
return;
}
ge::graphStatus TilingForHcPreInvRms(gert::TilingContext *context)
{
OPS_LOG_I(context, "TilingForHcPreInvRms start");
OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForHcPreInvRms", "Tiling context is null"),
return ge::GRAPH_FAILED);
auto platformInfo = context->GetPlatformInfo();
OPS_ERR_IF(platformInfo == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForHcPreInvRms", "Tiling platformInfo is null"),
return ge::GRAPH_FAILED);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
auto socVersion = ascendcPlatform.GetSocVersion();
if (socVersion == platform_ascendc::SocVersion::ASCEND950) {
OPS_LOG_I(context, "Using arch35 tiling for ASCEND950");
HcPreInvRmsRegbase::HcPreInvRmsTilingRegbase hcPreInvRmsTilingRegbase(context);
return hcPreInvRmsTilingRegbase.DoOpTiling();
}
auto xShapePtr = context->GetInputShape(0);
if (xShapePtr == nullptr) {
HcPreInvRmsTilingBase invRmsTilingBase(context);
return invRmsTilingBase.DoOpTiling();
}
auto &xShape = xShapePtr->GetStorageShape();
size_t xDimNum = xShape.GetDimNum();
int64_t R = 0;
if (xDimNum == X_INPUT_DIMS) {
R = xShape.GetDim(DIM_2) * xShape.GetDim(DIM_3);
} else if (xDimNum == X_INPUT_BS_FUSED_DIMS) {
R = xShape.GetDim(DIM_1) * xShape.GetDim(DIM_2);
}
if (R == 28672) {
OPS_LOG_I(context, "Using large_d tiling for R=28672");
HcPreInvRmsLargeD::HcPreInvRmsTilingLargeD invRmsTilingLargeD(context);
return invRmsTilingLargeD.DoOpTiling();
}
HcPreInvRmsTilingBase invRmsTilingBase(context);
return invRmsTilingBase.DoOpTiling();
}
static ge::graphStatus TilingPrepareForHcPreInvRms(gert::TilingParseContext *context)
{
(void)context;
return ge::GRAPH_SUCCESS;
}
IMPL_OP_OPTILING(HcPreInvRms)
.Tiling(TilingForHcPreInvRms)
.TilingParse<HcPreInvRmsCompileInfo>(TilingPrepareForHcPreInvRms);
} // namespace optiling

52
csrc/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling.h Normal file
View File

@@ -0,0 +1,52 @@
/**
* 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.
*/
/*!
* \file hc_pre_inv_rms_tiling.h
* \brief
*/
#ifndef HC_PRE_INV_RMS_TILING_H_
#define HC_PRE_INV_RMS_TILING_H_
#include <cmath>
#include <cstdint>
#include <vector>
#include <algorithm>
#include <iostream>
#include "register/op_impl_registry.h"
#include "platform/platform_infos_def.h"
#include "exe_graph/runtime/tiling_context.h"
#include "tiling/platform/platform_ascendc.h"
#include "register/op_def_registry.h"
#include "register/tilingdata_base.h"
#include "tiling/tiling_api.h"
#include "error/ops_error.h"
#include "platform/platform_info.h"
namespace optiling {
BEGIN_TILING_DATA_DEF(HcPreInvRmsFullLoadTilingData)
TILING_DATA_FIELD_DEF(int64_t, A); // A轴大小
TILING_DATA_FIELD_DEF(int64_t, R); // R轴大小
TILING_DATA_FIELD_DEF(int64_t, blockNumA); // 使用核数
TILING_DATA_FIELD_DEF(int64_t, blockFactorA); // 每个核处理的A个数
TILING_DATA_FIELD_DEF(int64_t, blockTailFactorA); // 尾核处理的A个数
TILING_DATA_FIELD_DEF(int64_t, ubFactorA); // 每次UB循环处理的A个数
TILING_DATA_FIELD_DEF(float, epsilon);
END_TILING_DATA_DEF;
REGISTER_TILING_DATA_CLASS(HcPreInvRms, HcPreInvRmsFullLoadTilingData)
struct HcPreInvRmsCompileInfo {};
} // namespace optiling
#endif // HC_PRE_INV_RMS_TILING_H_

310
csrc/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling_arch35.h Normal file
View File

@@ -0,0 +1,310 @@
/**
* 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.
*/
/* !
* \file hc_pre_inv_rms_tiling_arch35.cpp
* \brief
*/
#include <graph/utils/type_utils.h>
#include "hc_pre_inv_rms_tiling.h"
namespace optiling {
namespace HcPreInvRmsRegbase{
const static int64_t DEFAULT_WORKSPACE_SIZE = 16777216; // 预留16M空间
const static int64_t X_INPUT_INDEX = 0;
const static int64_t Y_OUTPUT_INDEX = 0;
const static int64_t EPS_ATTR_INDEX = 0;
const static size_t X_INPUT_BS_FUSED_DIMS = 3;
const static size_t X_INPUT_DIMS = 4;
const static int64_t UB_BLOCK_SIZE = 32;
const static uint64_t FULL_LOAD_REGBASE_TILING_KEY = 2000;
const static int64_t DIM_0 = 0;
const static int64_t DIM_1 = 1;
const static int64_t DIM_2 = 2;
const static int64_t DIM_3 = 3;
template <typename T>
static inline T CeilDiv(T num, T rnd)
{
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)));
}
template <typename T>
static inline T CeilAlign(T num, T rnd)
{
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd));
}
class HcPreInvRmsTilingRegbase {
public:
explicit HcPreInvRmsTilingRegbase(gert::TilingContext *context) : context_(context)
{
Reset();
}
~HcPreInvRmsTilingRegbase() = default;
bool IsCapable()
{
if (socVersion_ != platform_ascendc::SocVersion::ASCEND950) {
return false;
}
return true;
}
// 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小
ge::graphStatus GetPlatformInfo();
// 2、获取INPUT/OUTPUT/ATTR信息
ge::graphStatus GetShapeAttrsInfo();
// 3、计算数据切分TilingData
ge::graphStatus DoOpTiling();
// 4、计算高阶API的TilingData
ge::graphStatus DoLibApiTiling();
// 5、计算TilingKey
uint64_t GetTilingKey() const;
// 6、计算Workspace 大小
ge::graphStatus GetWorkspaceSize();
// 7、保存Tiling数据
ge::graphStatus PostTiling();
void Reset();
private:
ge::graphStatus CheckInputShape();
ge::graphStatus CheckAttr();
ge::graphStatus CheckOutShape();
void SplitA();
void CalUbFactorA();
const gert::Shape *xShape_ = nullptr;
const gert::Shape *yShape_ = nullptr;
float eps_ = 1e-6f;
int64_t A_ = 0;
int64_t R_ = 0;
int64_t inputDtypeSize_;
int64_t outputDtypeSize_;
const char *opName_ = "";
HcPreInvRmsFullLoadTilingData invRmsTilingData_;
gert::TilingContext *context_ = nullptr;
uint64_t workspaceSize_ = 0;
uint64_t coreNum_ = 0;
int64_t ubSize_ = 0;
int64_t ubBlockSize_ = 0;
platform_ascendc::SocVersion socVersion_;
};
ge::graphStatus HcPreInvRmsTilingRegbase::CheckInputShape()
{
size_t xDimNum = xShape_->GetDimNum();
OPS_ERR_IF(xDimNum != X_INPUT_DIMS && xDimNum != X_INPUT_BS_FUSED_DIMS,
OPS_LOG_E(context_, "The dim number of x is: %zu, but it should be %zu or %zu(bs fused)."
, xDimNum, X_INPUT_DIMS, X_INPUT_BS_FUSED_DIMS),
return ge::GRAPH_FAILED);
if (xDimNum == X_INPUT_DIMS) {
A_ = xShape_->GetDim(DIM_0) * xShape_->GetDim(DIM_1);
R_ = xShape_->GetDim(DIM_2) * xShape_->GetDim(DIM_3);
} else if (xDimNum == X_INPUT_BS_FUSED_DIMS) {
A_ = xShape_->GetDim(DIM_0);
R_ = xShape_->GetDim(DIM_1) * xShape_->GetDim(DIM_2);
}
invRmsTilingData_.set_A(A_);
invRmsTilingData_.set_R(R_);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingRegbase::CheckAttr()
{
OPS_ERR_IF(eps_ <= 0, OPS_LOG_E(context_, "epsilon is: %ld, but it should not be less than 0.", eps_), return ge::GRAPH_FAILED);
invRmsTilingData_.set_epsilon(eps_);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingRegbase::GetShapeAttrsInfo()
{
opName_ = context_->GetNodeName();
// 获取输入shape信息
auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED);
xShape_ = &xShapePtr->GetStorageShape();
// 获取输出shape
auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED);
yShape_ = &yShapePtr->GetStorageShape();
// 获取输入dtype
auto xDesc = context_->GetInputDesc(X_INPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, xDesc, return ge::GRAPH_FAILED);
auto xDtype = xDesc->GetDataType();
OPS_ERR_IF(
(xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16),
OPS_LOG_E(context_, "x dtype %s error, only supports float32, float16 and bfloat16. please check.",
ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()),
return ge::GRAPH_FAILED);
// 获取输出dtype
auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED);
auto yDtype = yDesc->GetDataType();
OPS_ERR_IF((yDtype != ge::DataType::DT_FLOAT),
OPS_LOG_E(context_, "y out dtype %s error, only support float32, please check",
ge::TypeUtils::DataTypeToSerialString(yDtype).c_str()),
return ge::GRAPH_FAILED);
// 获取属性
auto attrs = context_->GetAttrs();
OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED);
const float *epsPtr = attrs->GetAttrPointer<float>(EPS_ATTR_INDEX);
if (epsPtr != nullptr) {
eps_ = *epsPtr;
}
OPS_LOG_I(context_, "Attr eps is: %f ", eps_);
inputDtypeSize_ = static_cast<int64_t>(ge::GetSizeByDataType(context_->GetInputDesc(X_INPUT_INDEX)->GetDataType()));
outputDtypeSize_ = static_cast<int64_t>(ge::GetSizeByDataType(context_->GetOutputDesc(Y_OUTPUT_INDEX)->GetDataType()));
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingRegbase::GetPlatformInfo()
{
auto platformInfo = context_->GetPlatformInfo();
OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
coreNum_ = ascendcPlatform.GetCoreNumAiv();
socVersion_ = ascendcPlatform.GetSocVersion();
OPS_ERR_IF(
coreNum_ <= 0, OPS_LOG_E(context_->GetNodeName(), "coreNum must be greater than 0."),
return ge::GRAPH_FAILED);
// 获取UB大小
uint64_t ubSizePlatForm;
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
ubSize_ = static_cast<int64_t>(ubSizePlatForm);
OPS_ERR_IF(
ubSize_ <= 0, OPS_LOG_E(context_->GetNodeName(), "ubSize must be greater than 0."),
return ge::GRAPH_FAILED);
ubBlockSize_ = UB_BLOCK_SIZE; // 32: ub block size
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingRegbase::CheckOutShape()
{
OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)),
OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0),
xShape_->GetDim(0)),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
void HcPreInvRmsTilingRegbase::SplitA()
{
int64_t blockFactorA = CeilDiv(A_, static_cast<int64_t>(coreNum_));
int64_t blockNumA = CeilDiv(A_, blockFactorA);
int64_t blockTailFactorA = A_ % blockFactorA == 0 ? blockFactorA : A_ % blockFactorA;
invRmsTilingData_.set_blockNumA(blockNumA);
invRmsTilingData_.set_blockFactorA(blockFactorA);
invRmsTilingData_.set_blockTailFactorA(blockTailFactorA);
int64_t ubFactorA = invRmsTilingData_.get_ubFactorA();
if (ubFactorA > blockFactorA) {
invRmsTilingData_.set_ubFactorA(blockFactorA);
}
}
void HcPreInvRmsTilingRegbase::CalUbFactorA()
{
// ubFactorA * R_AlignSize * 2 + ubFactorA * outputDtypeSize_ + 32 = ubSize_
int64_t rAlignSize = CeilAlign(R_ * inputDtypeSize_, UB_BLOCK_SIZE);
int64_t ubFactorA = (ubSize_ - UB_BLOCK_SIZE) / (2 * rAlignSize + outputDtypeSize_);
invRmsTilingData_.set_ubFactorA(ubFactorA);
}
ge::graphStatus HcPreInvRmsTilingRegbase::DoOpTiling()
{
auto ret = GetPlatformInfo();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = GetShapeAttrsInfo();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = CheckInputShape();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = CheckOutShape();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = CheckAttr();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
CalUbFactorA();
SplitA();
ret = PostTiling();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingRegbase::DoLibApiTiling()
{
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingRegbase::GetWorkspaceSize()
{
// 计算workspace大小
workspaceSize_ = DEFAULT_WORKSPACE_SIZE;
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingRegbase::PostTiling()
{
context_->SetTilingKey(GetTilingKey());
context_->SetBlockDim(invRmsTilingData_.get_blockNumA());
size_t *currentWorkspace = context_->GetWorkspaceSizes(1);
currentWorkspace[0] = workspaceSize_;
invRmsTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(),
context_->GetRawTilingData()->GetCapacity());
context_->GetRawTilingData()->SetDataSize(invRmsTilingData_.GetDataSize());
return ge::GRAPH_SUCCESS;
}
uint64_t HcPreInvRmsTilingRegbase::GetTilingKey() const
{
// 全载场景
return FULL_LOAD_REGBASE_TILING_KEY;
}
void HcPreInvRmsTilingRegbase::Reset()
{
opName_ = nullptr;
return;
}
} // namespace HcPreInvRmsRegbase
} // namespace optiling

291
csrc/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling_large_d.h Normal file
View File

@@ -0,0 +1,291 @@
/**
* 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.
*/
/* !
* \file hc_pre_inv_rms_tiling_large_d.cpp
* \brief tiling for large d (R=28672, d=7168)
*/
#include <graph/utils/type_utils.h>
#include "hc_pre_inv_rms_tiling.h"
namespace optiling {
namespace HcPreInvRmsLargeD{
const static int64_t DEFAULT_WORKSPACE_SIZE = 16777216;
const static int64_t X_INPUT_INDEX = 0;
const static int64_t Y_OUTPUT_INDEX = 0;
const static int64_t EPS_ATTR_INDEX = 0;
const static size_t X_INPUT_BS_FUSED_DIMS = 3;
const static size_t X_INPUT_DIMS = 4;
const static int64_t UB_BLOCK_SIZE = 32;
const static uint64_t FULL_LOAD_LARGE_D_TILING_KEY = 1001;
const static int64_t R_LARGE_D = 28672;
const static int64_t DIM_0 = 0;
const static int64_t DIM_1 = 1;
const static int64_t DIM_2 = 2;
const static int64_t DIM_3 = 3;
template <typename T>
static inline T CeilDiv(T num, T rnd)
{
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)));
}
template <typename T>
static inline T CeilAlign(T num, T rnd)
{
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd));
}
class HcPreInvRmsTilingLargeD {
public:
explicit HcPreInvRmsTilingLargeD(gert::TilingContext *context) : context_(context)
{
Reset();
}
~HcPreInvRmsTilingLargeD() = default;
bool IsCapable()
{
return true;
}
ge::graphStatus GetPlatformInfo();
ge::graphStatus GetShapeAttrsInfo();
ge::graphStatus DoOpTiling();
ge::graphStatus DoLibApiTiling();
uint64_t GetTilingKey() const;
ge::graphStatus GetWorkspaceSize();
ge::graphStatus PostTiling();
void Reset();
private:
ge::graphStatus CheckInputShape();
ge::graphStatus CheckAttr();
ge::graphStatus CheckOutShape();
void SplitA();
void CalUbFactorA();
const gert::Shape *xShape_ = nullptr;
const gert::Shape *yShape_ = nullptr;
float eps_ = 1e-6f;
int64_t A_ = 0;
int64_t R_ = 0;
int64_t inputDtypeSize_;
int64_t outputDtypeSize_;
const char *opName_ = "";
HcPreInvRmsFullLoadTilingData invRmsTilingData_;
gert::TilingContext *context_ = nullptr;
uint64_t workspaceSize_ = 0;
uint64_t coreNum_ = 0;
int64_t ubSize_ = 0;
int64_t ubBlockSize_ = 0;
};
ge::graphStatus HcPreInvRmsTilingLargeD::CheckInputShape()
{
size_t xDimNum = xShape_->GetDimNum();
OPS_ERR_IF(xDimNum != X_INPUT_DIMS && xDimNum != X_INPUT_BS_FUSED_DIMS,
OPS_LOG_E(context_, "The dim number of x is: %zu, but it should be %zu or %zu(bs fused)."
, xDimNum, X_INPUT_DIMS, X_INPUT_BS_FUSED_DIMS),
return ge::GRAPH_FAILED);
if (xDimNum == X_INPUT_DIMS) {
A_ = xShape_->GetDim(DIM_0) * xShape_->GetDim(DIM_1);
R_ = xShape_->GetDim(DIM_2) * xShape_->GetDim(DIM_3);
} else if (xDimNum == X_INPUT_BS_FUSED_DIMS) {
A_ = xShape_->GetDim(DIM_0);
R_ = xShape_->GetDim(DIM_1) * xShape_->GetDim(DIM_2);
}
OPS_ERR_IF(R_ != R_LARGE_D,
OPS_LOG_E(context_, "R is: %ld, but large_d tiling only supports R=%ld.", R_, R_LARGE_D),
return ge::GRAPH_FAILED);
invRmsTilingData_.set_A(A_);
invRmsTilingData_.set_R(R_);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingLargeD::CheckAttr()
{
OPS_ERR_IF(eps_ <= 0, OPS_LOG_E(context_, "epsilon is: %ld, but it should not be less than 0.", eps_), return ge::GRAPH_FAILED);
invRmsTilingData_.set_epsilon(eps_);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingLargeD::GetShapeAttrsInfo()
{
opName_ = context_->GetNodeName();
auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED);
xShape_ = &xShapePtr->GetStorageShape();
auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED);
yShape_ = &yShapePtr->GetStorageShape();
auto xDesc = context_->GetInputDesc(X_INPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, xDesc, return ge::GRAPH_FAILED);
auto xDtype = xDesc->GetDataType();
OPS_ERR_IF(
(xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16),
OPS_LOG_E(context_, "x dtype %s error, only supports float32, float16 and bfloat16. please check.",
ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()),
return ge::GRAPH_FAILED);
auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX);
OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED);
auto yDtype = yDesc->GetDataType();
OPS_ERR_IF((yDtype != ge::DataType::DT_FLOAT),
OPS_LOG_E(context_, "y out dtype %s error, only support float32, please check",
ge::TypeUtils::DataTypeToSerialString(yDtype).c_str()),
return ge::GRAPH_FAILED);
auto attrs = context_->GetAttrs();
OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED);
const float *epsPtr = attrs->GetAttrPointer<float>(EPS_ATTR_INDEX);
if (epsPtr != nullptr) {
eps_ = *epsPtr;
}
OPS_LOG_I(context_, "Attr eps is: %f ", eps_);
inputDtypeSize_ = static_cast<int64_t>(ge::GetSizeByDataType(context_->GetInputDesc(X_INPUT_INDEX)->GetDataType()));
outputDtypeSize_ = static_cast<int64_t>(ge::GetSizeByDataType(context_->GetOutputDesc(Y_OUTPUT_INDEX)->GetDataType()));
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingLargeD::GetPlatformInfo()
{
auto platformInfo = context_->GetPlatformInfo();
OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
coreNum_ = ascendcPlatform.GetCoreNumAiv();
OPS_ERR_IF(
coreNum_ <= 0, OPS_LOG_E(context_->GetNodeName(), "coreNum must be greater than 0."),
return ge::GRAPH_FAILED);
uint64_t ubSizePlatForm;
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
ubSize_ = static_cast<int64_t>(ubSizePlatForm);
OPS_ERR_IF(
ubSize_ <= 0, OPS_LOG_E(context_->GetNodeName(), "ubSize must be greater than 0."),
return ge::GRAPH_FAILED);
ubBlockSize_ = UB_BLOCK_SIZE;
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingLargeD::CheckOutShape()
{
OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)),
OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0),
xShape_->GetDim(0)),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
void HcPreInvRmsTilingLargeD::SplitA()
{
int64_t blockFactorA = CeilDiv(A_, static_cast<int64_t>(coreNum_));
int64_t blockNumA = CeilDiv(A_, blockFactorA);
int64_t blockTailFactorA = A_ % blockFactorA == 0 ? blockFactorA : A_ % blockFactorA;
invRmsTilingData_.set_blockNumA(blockNumA);
invRmsTilingData_.set_blockFactorA(blockFactorA);
invRmsTilingData_.set_blockTailFactorA(blockTailFactorA);
int64_t ubFactorA = invRmsTilingData_.get_ubFactorA();
if (ubFactorA > blockFactorA) {
invRmsTilingData_.set_ubFactorA(blockFactorA);
}
}
void HcPreInvRmsTilingLargeD::CalUbFactorA()
{
invRmsTilingData_.set_ubFactorA(1);
}
ge::graphStatus HcPreInvRmsTilingLargeD::DoOpTiling()
{
auto ret = GetPlatformInfo();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = GetShapeAttrsInfo();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = CheckInputShape();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = CheckOutShape();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
ret = CheckAttr();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
CalUbFactorA();
SplitA();
ret = PostTiling();
if (ret != ge::GRAPH_SUCCESS) {
return ret;
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingLargeD::DoLibApiTiling()
{
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingLargeD::GetWorkspaceSize()
{
workspaceSize_ = DEFAULT_WORKSPACE_SIZE;
return ge::GRAPH_SUCCESS;
}
ge::graphStatus HcPreInvRmsTilingLargeD::PostTiling()
{
context_->SetTilingKey(GetTilingKey());
context_->SetBlockDim(invRmsTilingData_.get_blockNumA());
size_t *currentWorkspace = context_->GetWorkspaceSizes(1);
currentWorkspace[0] = workspaceSize_;
invRmsTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(),
context_->GetRawTilingData()->GetCapacity());
context_->GetRawTilingData()->SetDataSize(invRmsTilingData_.GetDataSize());
return ge::GRAPH_SUCCESS;
}
uint64_t HcPreInvRmsTilingLargeD::GetTilingKey() const
{
return FULL_LOAD_LARGE_D_TILING_KEY;
}
void HcPreInvRmsTilingLargeD::Reset()
{
opName_ = nullptr;
return;
}
} // namespace HcPreInvRmsLargeD
} // namespace optiling

50
csrc/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms.cpp Normal file
View File

@@ -0,0 +1,50 @@
/**
* 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.
*/
/*!
* \file hc_pre_inv_rms_apt.cpp
* \brief hc_pre_inv_rms kernel file
*/
#include "hc_pre_inv_rms_full_load.h"
#include "hc_pre_inv_rms_full_load_large_d.h"
#if defined(__DAV_C310__)
#include "hc_pre_inv_rms_full_load_regbase.h"
using namespace HcPreInvRmsRegbase;
#endif
#include "kernel_operator.h"
using namespace AscendC;
using namespace HcPreInvRms;
using namespace HcPreInvRmsLargeD;
#define FULL_LOAD_TILING_KEY 1000
#define FULL_LOAD_LARGE_D_TILING_KEY 1001
#define REGBASE_FULL_LOAD_TILING_KEY 2000
extern "C" __global__ __aicore__ void hc_pre_inv_rms(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
TPipe pipe;
GET_TILING_DATA(tilingData, tiling);
if (TILING_KEY_IS(FULL_LOAD_TILING_KEY)) {
HcPreInvRmsFullLoad<DTYPE_X> op;
op.Init(x, y, workspace, &tilingData, &pipe);
op.Process();
} else if (TILING_KEY_IS(FULL_LOAD_LARGE_D_TILING_KEY)) {
HcPreInvRmsFullLoadLargeD<DTYPE_X> op;
op.Init(x, y, workspace, &tilingData, &pipe);
op.Process();
}
#if defined(__DAV_C310__)
else if (TILING_KEY_IS(REGBASE_FULL_LOAD_TILING_KEY)) {
HcPreInvRmsFullLoadRegbase<DTYPE_X> op;
op.Init(x, y, workspace, &tilingData, &pipe);
op.Process();
}
#endif
}

248
csrc/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load.h Normal file
View File

@@ -0,0 +1,248 @@
/**
* 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.
*/
/*!
* \file hc_pre_inv_rms.h
* \brief inv rms file
*/
#ifndef ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_H_
#define ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_H_
#include "kernel_operator.h"
namespace HcPreInvRms {
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t FLOAT_BTYPE_SIZE = 4;
constexpr uint32_t PER_REPEAT_LEN_B32 = 64;
constexpr uint32_t UB_BLOCK_SIZE = 32;
constexpr int32_t B16_TYPE_BYTE_SIZE = 2;
constexpr int32_t B32_TYPE_BYTE_SIZE = 4;
constexpr int32_t ONE_COUNT = 1;
constexpr int32_t FOUR_FOLD = 4;
constexpr int32_t DST_REP_STRIDE = 1;
constexpr int32_t SRC_BLK_STRIDE = 1;
constexpr int32_t SRC_REP_STRIDE = 8;
template <typename T>
class HcPreInvRmsFullLoad {
public:
__aicore__ inline HcPreInvRmsFullLoad() {};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe);
__aicore__ inline void Process();
__aicore__ inline void CopyIn(uint64_t idx, uint64_t curUbFactorA);
__aicore__ inline void Compute(uint64_t curUbFactorA);
__aicore__ inline void ComputeB16(uint64_t curUbFactorA);
__aicore__ inline void ComputeB32(uint64_t curUbFactorA);
__aicore__ inline void CopyOut(uint64_t idx, uint64_t curUbFactorA);
private:
TPipe* pipe_;
TQue<QuePosition::VECIN, 1> inQueueX;
TQue<QuePosition::VECOUT, 1> outQueueY;
TBuf<TPosition::VECCALC> castBuf;
TBuf<TPosition::VECCALC> reduceBuf;
GlobalTensor<T> xGm;
GlobalTensor<float> yGm;
int64_t A; // 输入数据 A 轴大小
int64_t R; // 输入数据 R 轴大小
int64_t blockNumA; // 使用的核数
int64_t blockFactorA; // 每个核处理的A个数
int64_t blockTailFactorA; // 尾核处理的A个数
int64_t ubFactorA; // 每次ub循环处理的A个数
int32_t blockIdx_;
float epsilon; // 算子参数
uint32_t curBlockFactorA; // 当前核处理的A个数
uint32_t rAlign;
uint32_t rAlignB32;
uint32_t reduceBufNum;
};
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe)
{
A = tiling->A;
R = tiling->R;
blockNumA = tiling->blockNumA;
blockFactorA = tiling->blockFactorA;
blockTailFactorA = tiling->blockTailFactorA;
ubFactorA = tiling->ubFactorA;
epsilon = tiling->epsilon;
rAlign = ((R * sizeof(T) + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / sizeof(T));
rAlignB32 = ((R * FLOAT_BTYPE_SIZE + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / FLOAT_BTYPE_SIZE);
pipe_ = pipe;
blockIdx_ = GetBlockIdx();
if (blockIdx_ < blockNumA - 1) {
this->curBlockFactorA = this->blockFactorA;
} else if (blockIdx_ == blockNumA - 1) {
this->curBlockFactorA = this->blockTailFactorA;
} else {
return;
}
xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactorA * R, curBlockFactorA * R);
yGm.SetGlobalBuffer((__gm__ float*)y + blockIdx_ * blockFactorA, curBlockFactorA);
// pipe alloc memory to queue, the unit is Bytes
pipe_->InitBuffer(inQueueX, BUFFER_NUM, ubFactorA * rAlign * sizeof(T));
pipe_->InitBuffer(outQueueY, BUFFER_NUM, ubFactorA * FLOAT_BTYPE_SIZE);
reduceBufNum = (rAlignB32 + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32;
pipe_->InitBuffer(reduceBuf, ubFactorA * reduceBufNum * FLOAT_BTYPE_SIZE);
if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) {
pipe_->InitBuffer(castBuf, ubFactorA * rAlignB32 * FLOAT_BTYPE_SIZE);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::Process()
{
if (blockIdx_ >= blockNumA) {
return;
}
uint64_t aUbLoopCount = (curBlockFactorA + ubFactorA - 1) / ubFactorA; // Ub循环次数
uint64_t tailUbFactorA = curBlockFactorA - (aUbLoopCount - 1) * ubFactorA; // 最后一次Ub循环的A轴大小
uint64_t curUbFactorA = ubFactorA;
for (uint64_t idx = 0; idx < aUbLoopCount; idx++) {
if (idx == aUbLoopCount - 1) {
curUbFactorA = tailUbFactorA;
}
CopyIn(idx, curUbFactorA);
Compute(curUbFactorA);
CopyOut(idx, curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::CopyIn(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
DataCopyPadExtParams<T> dataCopyPadParams{false, 0, 0, 0};
int64_t xGmStartAddr = idx * R * ubFactorA;
DataCopyExtParams dataCopyParams{
static_cast<uint16_t>(curUbFactorA), static_cast<uint32_t>(R * sizeof(T)), 0, 0, 0};
DataCopyPad(xLocal, xGm[xGmStartAddr], dataCopyParams, dataCopyPadParams);
inQueueX.EnQue<T>(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::Compute(uint64_t curUbFactorA)
{
if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) {
ComputeB16(curUbFactorA);
} else if constexpr (sizeof(T) == B32_TYPE_BYTE_SIZE) {
ComputeB32(curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::ComputeB16(uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
LocalTensor<float> castLocal = castBuf.Get<float>();
LocalTensor<float> reduceLocal = reduceBuf.Get<float>();
int32_t perFoldElems = (rAlignB32 + FOUR_FOLD - 1) / FOUR_FOLD; // 4096
int32_t perFoldRepTime = (perFoldElems + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32; // 64
AscendC::Cast(castLocal, xLocal, AscendC::RoundMode::CAST_NONE, R);
PipeBarrier<PIPE_V>();
AscendC::Mul(castLocal, castLocal, castLocal, R);
for (int idx = 0; idx < curUbFactorA; idx++) {
for (int j = 0; j < FOUR_FOLD; j++) {
PipeBarrier<PIPE_V>();
WholeReduceSum(reduceLocal[idx * reduceBufNum + j * perFoldRepTime], castLocal[idx * rAlignB32 + j * perFoldElems], PER_REPEAT_LEN_B32, perFoldRepTime,
DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
}
PipeBarrier<PIPE_V>();
WholeReduceSum(reduceLocal, reduceLocal, PER_REPEAT_LEN_B32, FOUR_FOLD, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
PipeBarrier<PIPE_V>();
WholeReduceSum(yLocal[idx], reduceLocal, FOUR_FOLD, 1, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
}
float meanCof = 1.0f / R;
PipeBarrier<PIPE_V>();
AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Sqrt(yLocal, yLocal, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA);
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::ComputeB32(uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
int32_t perFoldElems = (rAlignB32 + FOUR_FOLD - 1) / FOUR_FOLD; // 4096
int32_t perFoldRepTime = (perFoldElems + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32; // 64
LocalTensor<float> reduceLocal = reduceBuf.Get<float>();
PipeBarrier<PIPE_V>();
AscendC::Mul(xLocal, xLocal, xLocal, R);
for (int idx = 0; idx < curUbFactorA; idx++) {
for (int j = 0; j < FOUR_FOLD; j++) {
PipeBarrier<PIPE_V>();
WholeReduceSum(reduceLocal[idx * reduceBufNum + j * perFoldRepTime], xLocal[idx * rAlignB32 + j * perFoldElems], PER_REPEAT_LEN_B32, perFoldRepTime,
DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
}
PipeBarrier<PIPE_V>();
WholeReduceSum(reduceLocal, reduceLocal, PER_REPEAT_LEN_B32, FOUR_FOLD, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
PipeBarrier<PIPE_V>();
WholeReduceSum(yLocal[idx], reduceLocal, FOUR_FOLD, 1, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
}
float meanCof = 1.0f / R;
PipeBarrier<PIPE_V>();
AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Sqrt(yLocal, yLocal, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA);
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::CopyOut(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<float> yLocal = outQueueY.DeQue<float>();
AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(curUbFactorA * sizeof(float)), 0, 0, 0};
DataCopyPad(yGm[idx * ubFactorA], yLocal, copyParams);
outQueueY.FreeTensor(yLocal);
}
} // namespace HcPreInvRms
#endif // ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_H_

262
csrc/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_large_d.h Normal file
View File

@@ -0,0 +1,262 @@
/**
* 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.
*/
/*!
* \file hc_pre_inv_rms_full_load_large_d.h
* \brief inv rms for large d (R=28672, d=7168)
*/
#ifndef ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_LARGE_D_H_
#define ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_LARGE_D_H_
#include "kernel_operator.h"
namespace HcPreInvRmsLargeD {
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 1;
constexpr int32_t FLOAT_BTYPE_SIZE = 4;
constexpr uint32_t PER_REPEAT_LEN_B32 = 64;
constexpr uint32_t UB_BLOCK_SIZE = 32;
constexpr int32_t B16_TYPE_BYTE_SIZE = 2;
constexpr int32_t B32_TYPE_BYTE_SIZE = 4;
constexpr int32_t HALf_INTERVAL = 2;
constexpr int32_t INDEX_TWO = 2;
constexpr int32_t INDEX_FOUR = 4;
constexpr int32_t INDEX_EIGHT = 8;
constexpr int32_t INDEX_SIXTEEN = 16;
__aicore__ inline int32_t FindPowerTwo(int32_t n)
{
n |= n >> 1;
n |= n >> INDEX_TWO;
n |= n >> INDEX_FOUR;
n |= n >> INDEX_EIGHT;
n |= n >> INDEX_SIXTEEN;
return (n + 1) >> 1;
}
__aicore__ inline void ReduceSumHalfInterval(
const LocalTensor<float>& dst_local, const LocalTensor<float>& src_local, int32_t count)
{
if (likely(count > PER_REPEAT_LEN_B32)) {
int32_t bodyCount = FindPowerTwo(count);
int32_t tailCount = count - bodyCount;
if (tailCount > 0) {
Add(src_local, src_local, src_local[bodyCount], tailCount);
PipeBarrier<PIPE_V>();
}
while (bodyCount > PER_REPEAT_LEN_B32) {
bodyCount = bodyCount / HALf_INTERVAL;
Add(src_local, src_local, src_local[bodyCount], bodyCount);
PipeBarrier<PIPE_V>();
}
AscendCUtils::SetMask<float>(PER_REPEAT_LEN_B32);
} else {
AscendCUtils::SetMask<float>(count);
}
#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220
if (g_coreType == AIV) {
WholeReduceSum<float, false>(dst_local, src_local, MASK_PLACEHOLDER, 1, 0, 1, 0);
}
#else
WholeReduceSum<float, false>(dst_local, src_local, MASK_PLACEHOLDER, 1, 1, 1, DEFAULT_REPEAT_STRIDE);
#endif
PipeBarrier<PIPE_V>();
}
template <typename T>
class HcPreInvRmsFullLoadLargeD {
public:
__aicore__ inline HcPreInvRmsFullLoadLargeD() {};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe);
__aicore__ inline void Process();
__aicore__ inline void CopyIn(uint64_t idx, uint64_t curUbFactorA);
__aicore__ inline void Compute(uint64_t curUbFactorA);
__aicore__ inline void ComputeB16(uint64_t curUbFactorA);
__aicore__ inline void ComputeB32(uint64_t curUbFactorA);
__aicore__ inline void CopyOut(uint64_t idx, uint64_t curUbFactorA);
private:
TPipe* pipe_;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
TBuf<TPosition::VECCALC> castBuf;
TBuf<TPosition::VECCALC> reduceBuf;
GlobalTensor<T> xGm;
GlobalTensor<float> yGm;
int64_t A;
int64_t R;
int64_t blockNumA;
int64_t blockFactorA;
int64_t blockTailFactorA;
int64_t ubFactorA;
int32_t blockIdx_;
float epsilon;
uint32_t curBlockFactorA;
uint32_t rAlign;
uint32_t rAlignB32;
uint32_t reduceBufNum;
};
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe)
{
A = tiling->A;
R = tiling->R;
blockNumA = tiling->blockNumA;
blockFactorA = tiling->blockFactorA;
blockTailFactorA = tiling->blockTailFactorA;
ubFactorA = tiling->ubFactorA;
epsilon = tiling->epsilon;
rAlign = ((R * sizeof(T) + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / sizeof(T));
rAlignB32 = ((R * FLOAT_BTYPE_SIZE + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / FLOAT_BTYPE_SIZE);
pipe_ = pipe;
blockIdx_ = GetBlockIdx();
if (blockIdx_ < blockNumA - 1) {
this->curBlockFactorA = this->blockFactorA;
} else if (blockIdx_ == blockNumA - 1) {
this->curBlockFactorA = this->blockTailFactorA;
} else {
return;
}
xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactorA * R, curBlockFactorA * R);
yGm.SetGlobalBuffer((__gm__ float*)y + blockIdx_ * blockFactorA, curBlockFactorA);
pipe_->InitBuffer(inQueueX, BUFFER_NUM, ubFactorA * rAlign * sizeof(T));
pipe_->InitBuffer(outQueueY, BUFFER_NUM, ubFactorA * FLOAT_BTYPE_SIZE);
reduceBufNum = (rAlignB32 + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32;
pipe_->InitBuffer(reduceBuf, ubFactorA * reduceBufNum * FLOAT_BTYPE_SIZE);
if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) {
pipe_->InitBuffer(castBuf, ubFactorA * rAlignB32 * FLOAT_BTYPE_SIZE);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::Process()
{
if (blockIdx_ >= blockNumA) {
return;
}
uint64_t aUbLoopCount = (curBlockFactorA + ubFactorA - 1) / ubFactorA;
uint64_t tailUbFactorA = curBlockFactorA - (aUbLoopCount - 1) * ubFactorA;
uint64_t curUbFactorA = ubFactorA;
for (uint64_t idx = 0; idx < aUbLoopCount; idx++) {
if (idx == aUbLoopCount - 1) {
curUbFactorA = tailUbFactorA;
}
CopyIn(idx, curUbFactorA);
Compute(curUbFactorA);
CopyOut(idx, curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::CopyIn(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
DataCopyPadExtParams<T> dataCopyPadParams{false, 0, 0, 0};
int64_t xGmStartAddr = idx * R * ubFactorA;
DataCopyExtParams dataCopyParams{
static_cast<uint16_t>(curUbFactorA), static_cast<uint32_t>(R * sizeof(T)), 0, 0, 0};
DataCopyPad(xLocal, xGm[xGmStartAddr], dataCopyParams, dataCopyPadParams);
inQueueX.EnQue<T>(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::Compute(uint64_t curUbFactorA)
{
if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) {
ComputeB16(curUbFactorA);
} else if constexpr (sizeof(T) == B32_TYPE_BYTE_SIZE) {
ComputeB32(curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::ComputeB16(uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
LocalTensor<float> castLocal = castBuf.Get<float>();
LocalTensor<float> reduceLocal = reduceBuf.Get<float>();
AscendC::Cast(castLocal, xLocal, AscendC::RoundMode::CAST_NONE, R);
PipeBarrier<PIPE_V>();
AscendC::Mul(castLocal, castLocal, castLocal, R);
for (int idx = 0; idx < curUbFactorA; idx++) {
ReduceSumHalfInterval(yLocal[idx], castLocal[idx * rAlignB32], R);
}
float meanCof = 1.0f / R;
PipeBarrier<PIPE_V>();
AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Sqrt(yLocal, yLocal, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA);
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::ComputeB32(uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
LocalTensor<float> reduceLocal = reduceBuf.Get<float>();
PipeBarrier<PIPE_V>();
AscendC::Mul(xLocal, xLocal, xLocal, R);
for (int idx = 0; idx < curUbFactorA; idx++) {
ReduceSumHalfInterval(yLocal[idx], xLocal[idx * rAlignB32], R);
}
float meanCof = 1.0f / R;
PipeBarrier<PIPE_V>();
AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Sqrt(yLocal, yLocal, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA);
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::CopyOut(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<float> yLocal = outQueueY.DeQue<float>();
AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(curUbFactorA * sizeof(float)), 0, 0, 0};
DataCopyPad(yGm[idx * ubFactorA], yLocal, copyParams);
outQueueY.FreeTensor(yLocal);
}
} // namespace HcPreInvRmsLargeD
#endif // ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_LARGE_D_H_

303
csrc/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_regbase.h Normal file
View File

@@ -0,0 +1,303 @@
/**
* 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.
*/
/*!
* \file hc_pre_inv_rms.h
* \brief inv rms file
*/
#ifndef ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_REGBASE_H_
#define ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_REGBASE_H_
#include "kernel_operator.h"
namespace HcPreInvRmsRegbase {
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t FLOAT_BTYPE_SIZE = 4;
constexpr uint32_t VF_LEN_B32 = 64;
constexpr uint32_t UB_BLOCK_SIZE = 32;
constexpr uint32_t FOLD_FOUR = 4;
constexpr AscendC::MicroAPI::CastTrait castTraitB162B32Even = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::UNKNOWN,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::UNKNOWN,
};
template <typename T>
__aicore__ inline void LoadInputData(AscendC::MicroAPI::RegTensor<float>& dst, __local_mem__ T* src, AscendC::MicroAPI::MaskReg pregLoop, uint32_t srcOffset)
{
if constexpr (IsSameType<T, float>::value) {
DataCopy(dst, src + srcOffset);
} else if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
AscendC::MicroAPI::RegTensor<T> tmp;
DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(tmp, src + srcOffset);
Cast<float, T, castTraitB162B32Even>(dst, tmp, pregLoop);
}
}
template <typename T>
class HcPreInvRmsFullLoadRegbase {
public:
__aicore__ inline HcPreInvRmsFullLoadRegbase() {};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe);
__aicore__ inline void Process();
__aicore__ inline void CopyIn(uint64_t idx, uint64_t curUbFactorA);
__aicore__ inline void Compute(uint64_t idx, uint64_t curUbFactorA);
__aicore__ inline void ComputeFullLoadVF(LocalTensor<float>& yLocal, LocalTensor<T>& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA);
__aicore__ inline void ComputeFullLoadVfPerf(LocalTensor<float>& yLocal, LocalTensor<T>& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA);
__aicore__ inline void CopyOut(uint64_t idx, uint64_t curUbFactorA);
private:
TPipe* pipe_;
TQue<QuePosition::VECIN, 1> inQueueX;
TQue<QuePosition::VECOUT, 1> outQueueY;
GlobalTensor<T> xGm;
GlobalTensor<float> yGm;
int64_t A; // 输入数据 A 轴大小
int64_t R; // 输入数据 R 轴大小
int64_t blockNumA; // 使用的核数
int64_t blockFactorA; // 每个核处理的A个数
int64_t blockTailFactorA; // 尾核处理的A个数
int64_t ubFactorA; // 每次ub循环处理的A个数
int32_t blockIdx_;
float epsilon; // 算子参数
uint32_t curBlockFactorA; // 当前核处理的A个数
uint32_t rAlign;
};
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe)
{
A = tiling->A;
R = tiling->R;
blockNumA = tiling->blockNumA;
blockFactorA = tiling->blockFactorA;
blockTailFactorA = tiling->blockTailFactorA;
ubFactorA = tiling->ubFactorA;
epsilon = tiling->epsilon;
rAlign = ((R * sizeof(T) + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / sizeof(T));
pipe_ = pipe;
blockIdx_ = GetBlockIdx();
if (blockIdx_ < blockNumA - 1) {
this->curBlockFactorA = this->blockFactorA;
} else if (blockIdx_ == blockNumA - 1) {
this->curBlockFactorA = this->blockTailFactorA;
} else {
return;
}
xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactorA * R, curBlockFactorA * R);
yGm.SetGlobalBuffer((__gm__ float*)y + blockIdx_ * blockFactorA, curBlockFactorA);
// pipe alloc memory to queue, the unit is Bytes
pipe_->InitBuffer(inQueueX, BUFFER_NUM, ubFactorA * rAlign * sizeof(T));
pipe_->InitBuffer(outQueueY, BUFFER_NUM, ubFactorA * FLOAT_BTYPE_SIZE);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::Process()
{
if (blockIdx_ >= blockNumA) {
return;
}
uint64_t aUbLoopCount = (curBlockFactorA + ubFactorA - 1) / ubFactorA; // Ub循环次数
uint64_t tailUbFactorA = curBlockFactorA - (aUbLoopCount - 1) * ubFactorA; // 最后一次Ub循环的A轴大小
uint64_t curUbFactorA = ubFactorA;
for (uint64_t idx = 0; idx < aUbLoopCount; idx++) {
if (idx == aUbLoopCount - 1) {
curUbFactorA = tailUbFactorA;
}
CopyIn(idx, curUbFactorA);
Compute(idx, curUbFactorA);
CopyOut(idx, curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::CopyIn(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
DataCopyPadExtParams<T> dataCopyPadParams{false, 0, 0, 0};
int64_t xGmStartAddr = idx * R * ubFactorA;
DataCopyExtParams dataCopyParams{
static_cast<uint16_t>(curUbFactorA), static_cast<uint32_t>(R * sizeof(T)), 0, 0, 0};
DataCopyPad(xLocal, xGm[xGmStartAddr], dataCopyParams, dataCopyPadParams);
inQueueX.EnQue<T>(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::Compute(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
if (R % 256 == 0) {
ComputeFullLoadVfPerf(yLocal, xLocal, rAlign, R, curUbFactorA);
} else {
ComputeFullLoadVF(yLocal, xLocal, rAlign, R, curUbFactorA);
}
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::ComputeFullLoadVF(LocalTensor<float>& yLocal, LocalTensor<T>& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA)
{
__ubuf__ T* xAddr = (__ubuf__ T*)xLocal.GetPhyAddr();
__ubuf__ float* yAddr = (__ubuf__ float*)yLocal.GetPhyAddr();
uint32_t vfLen = VF_LEN_B32;
uint16_t iLoopNum = curUbFactorA;
uint16_t needLoopNum = (rAlign + vfLen - 1) / vfLen; // 需要VF循环次数
uint16_t fourLoopNum = (needLoopNum + FOLD_FOUR - 1) / FOLD_FOUR; // 需要四循环次数
uint16_t formerFourLoopNum = fourLoopNum - 1;
uint16_t tailFourLoop = needLoopNum - formerFourLoopNum * FOLD_FOUR; // 最后一次四循环需处理VF数
uint16_t formerFourLoopElems = formerFourLoopNum * FOLD_FOUR * vfLen; // 除尾块外，处理元素个数
uint32_t tailFourLoopElems = rNum - formerFourLoopElems; // 最后一次四循环需处理实际元素个数
__VEC_SCOPE__
{
AscendC::MicroAPI::RegTensor<float> vregX1;
AscendC::MicroAPI::RegTensor<float> vregX2;
AscendC::MicroAPI::RegTensor<float> vregX3;
AscendC::MicroAPI::RegTensor<float> vregX4;
AscendC::MicroAPI::RegTensor<float> vregX;
AscendC::MicroAPI::RegTensor<float> vregSum;
AscendC::MicroAPI::RegTensor<float> vregR;
AscendC::MicroAPI::RegTensor<float> vregOne;
AscendC::MicroAPI::MaskReg preg;
AscendC::MicroAPI::MaskReg pregAll = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
AscendC::MicroAPI::MaskReg pregOne = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::VL1>();
AscendC::MicroAPI::Duplicate(vregOne, 1.0f);
for (uint16_t i = 0; i < iLoopNum; i++) {
AscendC::MicroAPI::Duplicate(vregSum, 0.0f); // 用于累加的vreg
for (uint16_t j = 0; j < formerFourLoopNum; j++) {
uint32_t srcOffset1 = i * rAlign + 4 * j * vfLen;
LoadInputData<T>(vregX1, xAddr, pregAll, srcOffset1);
AscendC::MicroAPI::Mul(vregX1, vregX1, vregX1, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX1, pregAll);
uint32_t srcOffset2 = i * rAlign + (4 * j + 1) * vfLen;
LoadInputData<T>(vregX2, xAddr, pregAll, srcOffset2);
AscendC::MicroAPI::Mul(vregX2, vregX2, vregX2, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX2, pregAll);
uint32_t srcOffset3 = i * rAlign + (4 * j + 2) * vfLen;
LoadInputData<T>(vregX3, xAddr, pregAll, srcOffset3);
AscendC::MicroAPI::Mul(vregX3, vregX3, vregX3, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX3, pregAll);
uint32_t srcOffset4 = i * rAlign + (4 * j + 3) * vfLen;
LoadInputData<T>(vregX4, xAddr, pregAll, srcOffset4);
AscendC::MicroAPI::Mul(vregX4, vregX4, vregX4, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX4, pregAll);
}
tailFourLoopElems = rNum - formerFourLoopElems;
for (uint16_t j = 0; j < tailFourLoop; j++) {
preg = AscendC::MicroAPI::UpdateMask<float>(tailFourLoopElems);
uint32_t srcOffset = i * rAlign + formerFourLoopElems + j * vfLen;
LoadInputData<T>(vregX, xAddr, preg, srcOffset);
AscendC::MicroAPI::Mul(vregX, vregX, vregX, preg);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX, pregAll);
}
Reduce(vregSum, vregSum, pregAll);
AscendC::MicroAPI::Duplicate(vregR, (float)rNum);
AscendC::MicroAPI::Div(vregSum, vregSum, vregR, pregOne);
AscendC::MicroAPI::Adds(vregSum, vregSum, epsilon, pregOne);
AscendC::MicroAPI::Sqrt(vregSum, vregSum, pregOne);
AscendC::MicroAPI::Div(vregSum, vregOne, vregSum, pregOne);
AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(yAddr + i, vregSum, pregOne);
}
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::ComputeFullLoadVfPerf(LocalTensor<float>& yLocal, LocalTensor<T>& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA)
{
__ubuf__ T* xAddr = (__ubuf__ T*)xLocal.GetPhyAddr();
__ubuf__ float* yAddr = (__ubuf__ float*)yLocal.GetPhyAddr();
uint32_t vfLen = VF_LEN_B32;
uint16_t iLoopNum = curUbFactorA;
uint16_t jLoopNum = (rAlign / vfLen) / FOLD_FOUR;
__VEC_SCOPE__
{
AscendC::MicroAPI::RegTensor<float> vregX1;
AscendC::MicroAPI::RegTensor<float> vregX2;
AscendC::MicroAPI::RegTensor<float> vregX3;
AscendC::MicroAPI::RegTensor<float> vregX4;
AscendC::MicroAPI::RegTensor<float> vregSum;
AscendC::MicroAPI::RegTensor<float> vregR;
AscendC::MicroAPI::RegTensor<float> vregOne;
AscendC::MicroAPI::MaskReg pregAll = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
AscendC::MicroAPI::MaskReg pregOne = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::VL1>();
AscendC::MicroAPI::Duplicate(vregOne, 1.0f);
for (uint16_t i = 0; i < iLoopNum; i++) {
AscendC::MicroAPI::Duplicate(vregSum, 0.0f); // 用于累加的vreg
for (uint16_t j = 0; j < jLoopNum; j++) {
uint32_t srcOffset1 = i * rAlign + 4 * j * vfLen;
LoadInputData<T>(vregX1, xAddr, pregAll, srcOffset1);
AscendC::MicroAPI::Mul(vregX1, vregX1, vregX1, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX1, pregAll);
uint32_t srcOffset2 = i * rAlign + (4 * j + 1) * vfLen;
LoadInputData<T>(vregX2, xAddr, pregAll, srcOffset2);
AscendC::MicroAPI::Mul(vregX2, vregX2, vregX2, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX2, pregAll);
uint32_t srcOffset3 = i * rAlign + (4 * j + 2) * vfLen;
LoadInputData<T>(vregX3, xAddr, pregAll, srcOffset3);
AscendC::MicroAPI::Mul(vregX3, vregX3, vregX3, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX3, pregAll);
uint32_t srcOffset4 = i * rAlign + (4 * j + 3) * vfLen;
LoadInputData<T>(vregX4, xAddr, pregAll, srcOffset4);
AscendC::MicroAPI::Mul(vregX4, vregX4, vregX4, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX4, pregAll);
}
Reduce(vregSum, vregSum, pregAll);
AscendC::MicroAPI::Duplicate(vregR, (float)rNum);
AscendC::MicroAPI::Div(vregSum, vregSum, vregR, pregOne);
AscendC::MicroAPI::Adds(vregSum, vregSum, epsilon, pregOne);
AscendC::MicroAPI::Sqrt(vregSum, vregSum, pregOne);
AscendC::MicroAPI::Div(vregSum, vregOne, vregSum, pregOne);
AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(yAddr + i, vregSum, pregOne);
}
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::CopyOut(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<float> yLocal = outQueueY.DeQue<float>();
AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(curUbFactorA * sizeof(float)), 0, 0, 0};
DataCopyPad(yGm[idx * ubFactorA], yLocal, copyParams);
outQueueY.FreeTensor(yLocal);
}
} // namespace HcPreInvRmsRegbase
#endif // ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_REGBASE_H_