/**
 * 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 dynamic_quant_tiling_arch35.cpp
 * \brief
 */
#include "dynamic_quant_tiling_arch35.h"

#include "register/op_impl_registry.h"
#include "log/log.h"
#include "util/math_util.h"
#include "tiling/platform/platform_ascendc.h"
#include "platform/platform_infos_def.h"
#include "error_util.h"
#include "quant/dynamic_quant/op_kernel/arch35/dynamic_quant_struct.h"
#include "quant/dynamic_quant/op_kernel/arch35/dynamic_quant_arch35_tilingdata.h"

using namespace ge;
using namespace AscendC;
using namespace DynamicQuantOp;

namespace optiling {
constexpr uint32_t OUTPUT_NUM_DYNAMIC_QUANT_V2 = 3;
constexpr uint32_t X_INDEX = 0;
constexpr uint32_t SMOOTH_INDEX = 1;
constexpr uint32_t GROUP_INDEX = 2;
constexpr uint32_t DST_TYPE_ATTR_INDEX = 0;
constexpr uint32_t IS_SYMMETRICAL_ATTR_INDEX = 1;
constexpr uint32_t QUANT_MODE_ATTR_INDEX = 2;
constexpr uint32_t DST_TYPE_MAX_ATTR_INDEX = 3;
constexpr uint32_t Y_INDEX = 0;
constexpr uint32_t SCALE_INDEX = 1;
constexpr uint32_t OFFSET_INDEX = 2;
constexpr uint32_t ONE = 1;
// 16 * 1024 * 1024
constexpr uint32_t SYS_WORKSPACE_SIZE = 16777216;
constexpr uint32_t COMPARE_INT = 255;
constexpr uint32_t RESERVED_LENGTH = 1024;
constexpr uint32_t MAX_EXPERT_NUM = 1024;
constexpr uint32_t MOE_SMOOTH_NUM = 2;
constexpr int64_t TILING_KEY_EMPTY_TENSOR = 999;
constexpr int64_t EVEN_FACTOR = 2;
constexpr uint32_t FLOAT_NUM_ONE_RPT = 128;
constexpr uint32_t ASYMMETRICAL = 2;
constexpr uint32_t USE_BUFFER_NUM = 2;
constexpr uint32_t SUM_INPUT_OUTPUT_NUM = 2;
// optional B16 smooth is 2 bytes
constexpr uint64_t SMOOTH_BYTES_SIZE = 2;
// B16 x is 2 bytes, B8/B4 y is 1 bytes, B32 scale is 4 bytes
constexpr uint64_t REQUIRED_BYTES_SIZE = 7;
// offset B32 offset is 4 bytes
constexpr uint64_t OFFSET_BYTES_SIZE = 4;
// double buffer offset B32 offset is 4 bytes
constexpr uint64_t DB_OFFSET_BYTES_SIZE = 8;
// double buffer optional B16 smooth is 4 bytes
constexpr uint64_t DB_SMOOTH_BYTES_SIZE = 4;
// double buffer B16 x, B8/B4 y, B32 scale are 14 bytes
constexpr uint64_t DB_REQUIRED_BYTES_SIZE = 14;
static map<const ge::DataType, const uint32_t> g_dTypeLen = {{ge::DT_INT32, 4}, {ge::DT_INT64, 8}};

// perchannel constants
// each element is 2 byte, so n base size is 128 elements for one register.
constexpr uint32_t SPLIT_M_SCHEDULE_MODE = 1;
constexpr uint32_t PER_CHANNEL_EXCLUDE_DIM = 2;
constexpr uint32_t PER_CHANNEL_N_BASE_SIZE = 64;

constexpr float HIFLOAT8_MAX_VALUE = 32768.0;
constexpr float FLT_EPSILON = 1e-6f;

template <uint32_t base, typename T = uint32_t>
auto AlignUp(T a) -> T
{
    return (a + base - 1) / base * base;
}

void DynamicQuantRegbaseTiling::SetTilingKey(gert::TilingContext* context)const
{
    int64_t tilingKey = GET_TPL_TILING_KEY(
        static_cast<uint32_t>(useDb), quantMode_, static_cast<uint32_t>(hasSmooth),
        static_cast<uint32_t>(isSymmetrical_));
    OP_LOGD(context, "regbase tilingKey is %ld", tilingKey);
    context->SetTilingKey(tilingKey);
}
ge::graphStatus DynamicQuantRegbaseTiling::CheckOpDim(
    const gert::StorageShape* shape1, const gert::StorageShape* shape2, uint32_t shape1Dim, uint32_t shape2Dim) const
{
    // 检查两个Tensor的shape是否一样
    if (shape1Dim != shape2Dim) {
        return ge::GRAPH_FAILED;
    }

    for (uint32_t i = 0; i < shape1Dim; i++) {
        if (shape1->GetStorageShape().GetDim(i) != shape2->GetStorageShape().GetDim(i)) {
            return ge::GRAPH_FAILED;
        }
    }
    return ge::GRAPH_SUCCESS;
}

// 检查输入的shape是否符合要求，设置groupNum
ge::graphStatus DynamicQuantRegbaseTiling::CheckOpInputShape(gert::TilingContext* context)
{
    auto xShape = context->GetInputShape(X_INDEX);
    OP_CHECK_NULL_WITH_CONTEXT(context, xShape);
    size_t xDimNum = xShape->GetStorageShape().GetDimNum();
    size_t xSizeNum = xShape->GetStorageShape().GetShapeSize();
    isEmptyTensor = (xSizeNum == 0UL);

    if (xDimNum <= 1UL) {
        OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(context->GetNodeName(), "x", std::to_string(xDimNum),
                                              "The shape dim of x must be greater than or equal to 2");
        return ge::GRAPH_FAILED;
    }
    int64_t xDimLast = xShape->GetStorageShape().GetDim(xDimNum - 1);
    int64_t xDimSecendToLast = xShape->GetStorageShape().GetDim(xDimNum - PER_CHANNEL_EXCLUDE_DIM);

    if (yDtype == ge::DT_INT4) {
        OP_CHECK_IF(
            (xDimLast % EVEN_FACTOR),
OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context->GetNodeName(), "x", Ops::Base::ToString(xShape->GetStorageShape()),
                                               "When the yDtype is DT_INT4, the tail axis of x must be an even number"),
            return ge::GRAPH_FAILED);
    }

    auto smoothShape = context->GetOptionalInputShape(SMOOTH_INDEX);
    if (smoothShape != nullptr) {
        auto groupShape = context->GetOptionalInputShape(GROUP_INDEX);
        if (groupShape != nullptr) {
            groupNum = groupShape->GetStorageShape().GetDim(groupShape->GetStorageShape().GetDimNum() - 1);
            OP_CHECK_IF(
                (groupNum <= 0 || groupNum > MAX_EXPERT_NUM),
                OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "group_index", std::to_string(groupNum),
                                               "The value of group_index must be within the range [1, 1024]"),
                return ge::GRAPH_FAILED);
            OP_CHECK_IF(
                (groupNum > 0 && isPerChannel_),
                OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "groupNum", std::to_string(groupNum),
                                               "The vaule of groupNum must be less than or equal to 0 when the quantization mode is perchannel"),
                return ge::GRAPH_FAILED);
        }

        size_t smoothDimNum = smoothShape->GetStorageShape().GetDimNum();
        // 针对moe场景下的校验
        if (groupNum >= 1UL) {
            OP_CHECK_IF(
                (smoothDimNum != MOE_SMOOTH_NUM),
                OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(context->GetNodeName(), "smooth_scales", std::to_string(smoothDimNum),
                                              "When groupNum is more than 0, the shape dim of smooth_scales must be equal to 2"),
                return ge::GRAPH_FAILED);
            int64_t smoothDimFirst = smoothShape->GetStorageShape().GetDim(0);
            if (groupNum != static_cast<size_t>(smoothDimFirst)) {
                OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(context->GetNodeName(), "smooth_scales, group_index",
                                                std::to_string(smoothDimFirst) + ", " + std::to_string(groupNum),
                                                "The value of smooth_scales first dim and group_num must be the same");
                return ge::GRAPH_FAILED;
            }
        } else {
            if (smoothDimNum != 1UL) {
                OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(context->GetNodeName(), "smooth_scales", std::to_string(smoothDimNum),
                                              "The shape dim of smooth_scales must be 1");
                return ge::GRAPH_FAILED;
            }
        }
        int64_t smoothDimLast = smoothShape->GetStorageShape().GetDim(smoothDimNum - 1);
        if (!isPerChannel_ && xDimLast != smoothDimLast) {
            OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "x, smooth_scales",
                                                std::to_string(xDimLast) + ", " + std::to_string(smoothDimLast),
                                                "The shapes of x and smooth_scales on the last dim must be the same");
            return ge::GRAPH_FAILED;
        }
        if (isPerChannel_ && xDimSecendToLast != smoothDimLast) {
            OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "x, smooth_scales",
                                                std::to_string(xDimSecendToLast) + ", " + std::to_string(smoothDimLast),
                                                "[-2] dim of x must be equal to [-1] dim of smooth_scales when quant_mode is perchannel");
            return ge::GRAPH_FAILED;
        }
        hasSmooth = true;
    }
    return ge::GRAPH_SUCCESS;
}

// 检查输出shape
ge::graphStatus DynamicQuantRegbaseTiling::CheckOpOutputShape(gert::TilingContext* context) const
{
    auto xShape = context->GetInputShape(X_INDEX);
    size_t xDimNum = xShape->GetStorageShape().GetDimNum();
    auto yShape = context->GetOutputShape(Y_INDEX);
    OP_CHECK_NULL_WITH_CONTEXT(context, yShape);
    size_t yDimNum = yShape->GetStorageShape().GetDimNum();
    // 检查x和y
    OP_CHECK_IF(
        (CheckOpDim(xShape, yShape, xDimNum, yDimNum) != ge::GRAPH_SUCCESS),
        OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "x, y",
                                                Ops::Base::ToString(xShape->GetStorageShape()) + ", " + Ops::Base::ToString(yShape->GetStorageShape()),
                                                "The shapes of x and y must be the same"),
        return ge::GRAPH_FAILED);

    // 检查x和scale
    auto scaleShape = context->GetOutputShape(SCALE_INDEX);
    OP_CHECK_NULL_WITH_CONTEXT(context, scaleShape);
    size_t scaleDimNum = scaleShape->GetStorageShape().GetDimNum();

    if (quantMode_ == TPL_PER_TENSOR_FULL_LOAD) {
        OP_CHECK_IF(
            scaleDimNum != 1 || scaleShape->GetStorageShape().GetDim(0) != 1,
            OP_LOGE_FOR_INVALID_SHAPE(context->GetNodeName(), "scale", Ops::Base::ToString(scaleShape->GetStorageShape()), "[1]"),
            return ge::GRAPH_FAILED);
    } else if (isPerChannel_){
        // check batch dims
        OP_CHECK_IF(
            (CheckOpDim(xShape, scaleShape, xDimNum - PER_CHANNEL_EXCLUDE_DIM, scaleDimNum - 1) != ge::GRAPH_SUCCESS),
            OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "scale, x",
                                                Ops::Base::ToString(scaleShape->GetStorageShape()) + ", " + Ops::Base::ToString(xShape->GetStorageShape()),
                                                "The shapes of scale excluding the last dim and the shape of x excluding the last two dims must be the same"),
            return ge::GRAPH_FAILED);
        // check last dim
        OP_CHECK_IF(
            (xShape->GetStorageShape().GetDim(xDimNum - 1) != scaleShape->GetStorageShape().GetDim(scaleDimNum - 1)),
            OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "scale, x",
                                                std::to_string(scaleShape->GetStorageShape().GetDim(scaleDimNum - 1)) + ", " + std::to_string(xShape->GetStorageShape().GetDim(xDimNum - 1)),
                                                "Shape[-1] of scale must be equal to Shape[-1] of x"),
            return ge::GRAPH_FAILED);        
    } else {
        OP_CHECK_IF(
            (CheckOpDim(xShape, scaleShape, xDimNum - 1, scaleDimNum) != ge::GRAPH_SUCCESS),
            OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "scale, x",
                                                Ops::Base::ToString(scaleShape->GetStorageShape()) + ", " + Ops::Base::ToString(xShape->GetStorageShape()),
                                                "The shapes of scale and x excluding the last dim must be the same"),
            return ge::GRAPH_FAILED);
    }

    // 检查scale和offset
    if (context->GetComputeNodeOutputNum() == OUTPUT_NUM_DYNAMIC_QUANT_V2) {
        auto offsetShape = context->GetOutputShape(OFFSET_INDEX);
        OP_CHECK_NULL_WITH_CONTEXT(context, offsetShape);
        OPS_ERR_IF(
            scaleShape->GetStorageShape() != offsetShape->GetStorageShape(),
OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "scale, offset",
                                                Shape2String(scaleShape->GetStorageShape()) + ", " + Shape2String(offsetShape->GetStorageShape()),
                                                "The shapes of scale and offset must be the same"),
            return ge::GRAPH_FAILED);
    }

    return ge::GRAPH_SUCCESS;
}

// 检查算子全部变量的shape
ge::graphStatus DynamicQuantRegbaseTiling::CheckOpShape(gert::TilingContext* context)
{
    OP_CHECK_IF(
        (CheckOpInputShape(context) != ge::GRAPH_SUCCESS),
        OP_LOGE(context, "input shape check failed!"), return ge::GRAPH_FAILED);

    OP_CHECK_IF(
        (CheckOpOutputShape(context) != ge::GRAPH_SUCCESS),
        OP_LOGE(context, "output shape check failed!"), return ge::GRAPH_FAILED);
    return ge::GRAPH_SUCCESS;
}

// 检查输出的数据类型
ge::graphStatus DynamicQuantRegbaseTiling::CheckOutputDtype(gert::TilingContext* context)
{
    auto yDesc = context->GetOutputDesc(Y_INDEX);
    OP_CHECK_NULL_WITH_CONTEXT(context, yDesc);
    yDtype = yDesc->GetDataType();
    std::vector<ge::DataType> ySupportDtype = {
        ge::DT_INT8, ge::DT_INT4, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_HIFLOAT8};
    if (std::find(ySupportDtype.begin(), ySupportDtype.end(), yDtype) == ySupportDtype.end()) {
        OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "y", ge::TypeUtils::DataTypeToSerialString(static_cast<ge::DataType>(yDtype)),
                                  "DT_INT8, DT_INT4, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN, DT_HIFLOAT8");
        return ge::GRAPH_FAILED;
    }

    auto scaleDesc = context->GetOutputDesc(SCALE_INDEX);
    OP_CHECK_NULL_WITH_CONTEXT(context, scaleDesc);
    auto scaleDtype = scaleDesc->GetDataType();
    if (scaleDtype != ge::DataType::DT_FLOAT) {
        OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "scale", ge::TypeUtils::DataTypeToSerialString(scaleDtype), "DT_FLOAT");
        return ge::GRAPH_FAILED;
    }

    if (context->GetComputeNodeOutputNum() == OUTPUT_NUM_DYNAMIC_QUANT_V2) {
        auto offsetDesc = context->GetOutputDesc(OFFSET_INDEX);
        OP_CHECK_NULL_WITH_CONTEXT(context, offsetDesc);
        auto offsetDtype = scaleDesc->GetDataType();
        if (offsetDtype != ge::DataType::DT_FLOAT) {
            OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "offset", ge::TypeUtils::DataTypeToSerialString(offsetDtype), "DT_FLOAT");
            return ge::GRAPH_FAILED;
        }
    }
    return ge::GRAPH_SUCCESS;
}

// 检查输入的数据类型
ge::graphStatus DynamicQuantRegbaseTiling::CheckInputDtype(gert::TilingContext* context)
{
    auto xDesc = context->GetInputDesc(X_INDEX);
    OP_CHECK_NULL_WITH_CONTEXT(context, xDesc);
    auto xDtype = xDesc->GetDataType();
    if (xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16) {
        OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "x", ge::TypeUtils::DataTypeToSerialString(xDtype), "DT_FLOAT16, DT_BF16");
        return ge::GRAPH_FAILED;
    }

    auto smoothDesc = context->GetOptionalInputDesc(SMOOTH_INDEX);
    if (smoothDesc != nullptr) {
        auto smoothDtype = smoothDesc->GetDataType();
        if (smoothDtype != ge::DataType::DT_FLOAT16 && smoothDtype != ge::DataType::DT_BF16) {
OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "smooth_scale", ge::TypeUtils::DataTypeToSerialString(smoothDtype), "DT_FLOAT16, DT_BF16");
            return ge::GRAPH_FAILED;
        }
        if (xDtype != smoothDtype) {
            OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(context->GetNodeName(), "x, smooth_scales",
                                                ge::TypeUtils::DataTypeToSerialString(xDtype) + ", " + ge::TypeUtils::DataTypeToSerialString(smoothDtype),
                                                "The dtype of x must be the same as the dtype of smooth_scales");
            return ge::GRAPH_FAILED;
        }
        auto groupDesc = context->GetOptionalInputDesc(GROUP_INDEX);
        if (groupDesc != nullptr) {
            auto groupDtype = groupDesc->GetDataType();
            if (groupDtype != ge::DataType::DT_INT32) {
OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "group_index", ge::TypeUtils::DataTypeToSerialString(groupDtype), "DT_INT32");
                return ge::GRAPH_FAILED;
            }
            groupDtypeSize = g_dTypeLen[groupDtype];
        }
    }
    return ge::GRAPH_SUCCESS;
}

// 检查attr是否符合要求
ge::graphStatus DynamicQuantRegbaseTiling::CheckAttrs(gert::TilingContext* context)
{
    auto* attrs = context->GetAttrs();
    OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
    const int32_t* dstTypePtr = attrs->GetAttrPointer<int32_t>(DST_TYPE_ATTR_INDEX);
    OP_CHECK_NULL_WITH_CONTEXT(context, dstTypePtr);
    int32_t dstType = *dstTypePtr;
    if (dstType != yDtype) {
	    OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(context->GetNodeName(), "dst_type, y",
                                                ge::TypeUtils::DataTypeToSerialString(static_cast<ge::DataType>(dstType)) + ", " + ge::TypeUtils::DataTypeToSerialString(static_cast<ge::DataType>(yDtype)),
                                                "The dtypes of dst_type and y must be the same");
        return ge::GRAPH_FAILED;
    }

    // check dynamicquantv2 attr
    if (context->GetComputeNodeOutputNum() == OUTPUT_NUM_DYNAMIC_QUANT_V2) {
        // get and check quant mode
        const char* quantModeAttr = attrs->GetAttrPointer<char>(QUANT_MODE_ATTR_INDEX);
        OP_CHECK_NULL_WITH_CONTEXT(context, quantModeAttr);

        // 获取dstTypeMax的值
        const float* dstTypeMaxAttr = attrs->GetAttrPointer<float>(DST_TYPE_MAX_ATTR_INDEX);
        dstTypeMax = (dstTypeMaxAttr != nullptr) ? *dstTypeMaxAttr : 0.0f;

        if (dstTypeMax > HIFLOAT8_MAX_VALUE || dstTypeMax < 0) {
            OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "dst_type_max", std::to_string(dstTypeMax),
                                               "The value of dst_type_max must be within the range [0, 32768]");
            return ge::GRAPH_FAILED;
        }

        if (std::fabs(dstTypeMax) < FLT_EPSILON) {
            dstTypeMax = HIFLOAT8_MAX_VALUE;
        }

        std::string quantModeStr = quantModeAttr;
        if (quantModeStr == "pertoken") {
            quantMode_ = TPL_COMMON_FULL_LOAD;
        } else if (quantModeStr == "pertensor") {
            quantMode_ = TPL_PER_TENSOR_FULL_LOAD;
        } else if (quantModeStr == "perchannel") {
            // quantMode_ will be computed later
            isPerChannel_ = true;
        } else {
            OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "quant_mode", quantModeStr,
                                               "The value of quant_mode must be pertoken, pertensor or perchannel");
            return ge::GRAPH_FAILED;
        }

        // get and check is_symmetrical
        const bool* isSymmetricalAttr = attrs->GetAttrPointer<bool>(IS_SYMMETRICAL_ATTR_INDEX);
        OP_CHECK_NULL_WITH_CONTEXT(context, isSymmetricalAttr);
        isSymmetrical_ = isSymmetricalAttr == nullptr ? false : *isSymmetricalAttr;
    }

    return ge::GRAPH_SUCCESS;
}

// 检查所有变量是否符合要求
ge::graphStatus DynamicQuantRegbaseTiling::CheckOpParams(gert::TilingContext* context)
{
    OP_CHECK_IF(
        (CheckInputDtype(context) != ge::GRAPH_SUCCESS),
        OP_LOGE(context, "x or smooth_scales dtype is invalid."),
        return ge::GRAPH_FAILED);

    OP_CHECK_IF(
        (CheckOutputDtype(context) != ge::GRAPH_SUCCESS),
        OP_LOGE(context, "op output dtype is invalid."),
        return ge::GRAPH_FAILED);

    OP_CHECK_IF(
        (CheckAttrs(context) != ge::GRAPH_SUCCESS),
        OP_LOGE(context, "op attrs is invalid."), return ge::GRAPH_FAILED);

    OP_CHECK_IF(
        (CheckOpShape(context) != ge::GRAPH_SUCCESS),
        OP_LOGE(context, "input or output shape is invalid."),
        return ge::GRAPH_FAILED);
    return ge::GRAPH_SUCCESS;
}

// 重置大shape变量
void DynamicQuantRegbaseTiling::ResetLargeTilingParams()
{
    useDb = false;
    isEmptyTensor = false;
    innerLoopEle = 0UL;
    innerLoopTimes = 0UL;
    innerLoopTail = 0UL;
    groupNum = 0UL;
    groupDtypeSize = 0UL;
    hasSmooth = false;
    quantMode_ = 0UL;
    dstTypeMax = HIFLOAT8_MAX_VALUE;
    isPerChannel_ = false;

    mLen = 0;
    nLen = 0;
    totalBatchLen = 0;
    nBlockSize = 1;
    nBlockNum = 1;
    mBlockSize = 0;
    mBlockNum = 1;
    totalBlockNum = 0;
    blockPerHead = 0;
    blockPerTail = 0;
    nBaseLoopNum = 1;
    batchBlockSize = 1;
    batchTailBlockSize = 1;
    batchBlockNum = 1;
}

// 打印tiling参数
void DynamicQuantRegbaseTiling::PrintTilingData(gert::TilingContext* context) const
{
    DynamicQuantTilingDataArch35* tilingDataPtr = context->GetTilingData<DynamicQuantTilingDataArch35>();
    OP_LOGD(
        context,
        "tilingData is coreNum:%u, rowLen:%u, headCoreNum:%u, rowPerHeadCore:%u, "
        "rowPerTailCore:%u, multiRowNumHeadCore:%u, multiRowNumTailCore:%u, "
        "innerLoopEle:%u, innerLoopTimes:%u, innerLoopTail:%u, groupNum:%u, hasSmooth:%u, dstTypeMax:%f",
        tilingDataPtr->coreNum, tilingDataPtr->rowLen, tilingDataPtr->headCoreNum, tilingDataPtr->rowPerHeadCore,
        tilingDataPtr->rowPerTailCore, tilingDataPtr->multiRowNumHeadCore, tilingDataPtr->multiRowNumTailCore,
        tilingDataPtr->innerLoopEle, tilingDataPtr->innerLoopTimes, tilingDataPtr->innerLoopTail,
        tilingDataPtr->groupNum, tilingDataPtr->hasSmooth, tilingDataPtr->dstTypeMax);
}

// 赋值tiling参数
void DynamicQuantRegbaseTiling::SetTilingData(gert::TilingContext* context) const
{
    SetTilingKey(context);
    DynamicQuantTilingDataArch35* tilingDataPtr = context->GetTilingData<DynamicQuantTilingDataArch35>();
    tilingDataPtr->coreNum = coreNum;
    tilingDataPtr->rowLen = rowLen;
    tilingDataPtr->headCoreNum = headCoreNum;
    tilingDataPtr->rowPerHeadCore = rowPerHeadCore;
    tilingDataPtr->rowPerTailCore = rowPerTailCore;
    tilingDataPtr->multiRowNumHeadCore = multiRowNumHeadCore;
    tilingDataPtr->multiRowNumTailCore = multiRowNumTailCore;
    tilingDataPtr->innerLoopEle = innerLoopEle;
    tilingDataPtr->innerLoopTimes = innerLoopTimes;
    tilingDataPtr->innerLoopTail = innerLoopTail;
    tilingDataPtr->groupNum = groupNum;
    tilingDataPtr->hasSmooth = hasSmooth ? 1 : 0;
    tilingDataPtr->dstTypeMax = dstTypeMax;
}

// 处理空Tensor的tiling
ge::graphStatus DynamicQuantRegbaseTiling::DoEmptyTensorTiling(gert::TilingContext* context) const
{
    int64_t tilingKey = GET_TPL_TILING_KEY(0U, TPL_EMPTY_TENSOR, 0U, 0U);
    OP_LOGD(context, "DoEmptyTensorTiling, the last dim must be 0, tilingKey is %ld", tilingKey);
    context->SetTilingKey(tilingKey);

    size_t* workSpaces = context->GetWorkspaceSizes(1);
    OP_CHECK_NULL_WITH_CONTEXT(context, workSpaces);
    workSpaces[0] = SYS_WORKSPACE_SIZE;
    context->SetBlockDim(1);
    return ge::GRAPH_SUCCESS;
}

// 设置rowLen为x的最后一维
// 设置headCoreNum，headCore的数量
// 设置rowPerHeadCore、rowPerTailCore，每个核处理的row数量，如果有tailCore的话，headCore处理的row比tailCore多一个
// 设置coreNum,最红使用的核数
void DynamicQuantRegbaseTiling::CalculateCoreNum(const gert::TilingContext* context)
{
    const gert::StorageShape* xShape = context->GetInputShape(X_INDEX);
    rowLen = xShape->GetStorageShape().GetDim(xShape->GetStorageShape().GetDimNum() - 1);
    // 获取输入x的维度x-1
    size_t dimNum = xShape->GetStorageShape().GetDimNum() - 1;
    uint64_t tempHeadCoreNum = 1;
    for (size_t i = 0; i < dimNum; i++) {
        tempHeadCoreNum *= xShape->GetStorageShape().GetDim(i);
    }
    // 设置rowNum为x除了最后一维以外所有维度的乘积
    rowNum = tempHeadCoreNum;
    headCoreNum = rowNum % vectorCoreNum;
    rowPerHeadCore = (rowNum + vectorCoreNum - 1U) / vectorCoreNum;
    rowPerTailCore = rowNum / vectorCoreNum;
    coreNum = std::max(std::min(vectorCoreNum, rowNum), ONE);
}

// 设置useDb，打开DB的情况
// 设置quantMode_
// 对于大shape设置innerLoopEle、innerLoopTimes、innerLoopTail、multiRowNumHeadCore、multiRowNumTailCore
// 对于全载场景，设置multiRowNumHeadCore、multiRowNumTailCore，每次UB搬运处理的row数量
void DynamicQuantRegbaseTiling::CalculateTilingData()
{
    // 每个row的元素数量对齐16
    uint32_t alignedRowLen = AlignUp<16>(rowLen);
    uint64_t maxUseUbSize = ubSize - RESERVED_LENGTH;
    uint32_t smoothBuffer = (hasSmooth ? 1UL : 0UL);
    uint32_t offsetBuffer = 1UL;

    uint64_t calcSize = static_cast<uint64_t>(alignedRowLen) *
                        (smoothBuffer * SMOOTH_BYTES_SIZE + REQUIRED_BYTES_SIZE + offsetBuffer * OFFSET_BYTES_SIZE);
    // 如果开double buffer的话，buffer一次处理的数量会有变化
    uint64_t calcDbSize =
        static_cast<uint64_t>(alignedRowLen) *
        (smoothBuffer * DB_SMOOTH_BYTES_SIZE + DB_REQUIRED_BYTES_SIZE + offsetBuffer * DB_OFFSET_BYTES_SIZE);

    // 一个UB装不下一个row的情况
    // 一个UB可以装下一个row，但是无法开启double buffer
    // 一个UB可以装下一个row，但是可以开启double buffer
    if (calcSize > maxUseUbSize) {
        uint64_t calcLargeBuf =
            2UL * (static_cast<uint64_t>(smoothBuffer) * 2UL + 4UL); // for inQueue,outQueue&smoothQueue DB
        maxUseUbSize -= RESERVED_LENGTH;                             // for scaleQueue DB
        useDb = true;
        innerLoopEle = static_cast<uint32_t>(maxUseUbSize) / static_cast<uint32_t>(calcLargeBuf) / FLOAT_NUM_ONE_RPT *
                       FLOAT_NUM_ONE_RPT;
        innerLoopTimes = rowLen / innerLoopEle;
        innerLoopTail = rowLen % innerLoopEle;
        multiRowNumHeadCore = std::min({COMPARE_INT, ONE, rowPerHeadCore});
        multiRowNumTailCore = std::min({COMPARE_INT, ONE, rowPerTailCore});

        if (quantMode_ != TPL_PER_TENSOR_FULL_LOAD) {
            quantMode_ = groupNum > 0 ? TPL_MOE_LARGE_SHAPE : TPL_COMMON_LARGE_SHAPE;
        } else {
            quantMode_ = groupNum > 0 ? TPL_MOE_PER_TENSOR_LARGE_SHAPE : TPL_PER_TENSOR_LARGE_SHAPE;
        }
    } else if (calcDbSize > maxUseUbSize) {
        useDb = false;
        // 取UB可以处理的row数量
        uint32_t ubAvail = static_cast<uint32_t>(maxUseUbSize) / static_cast<uint32_t>(calcSize);
        multiRowNumHeadCore = std::min({ubAvail, COMPARE_INT, rowPerHeadCore});
        multiRowNumTailCore = std::min({ubAvail, COMPARE_INT, rowPerTailCore});
        if (quantMode_ != TPL_PER_TENSOR_FULL_LOAD) {
            quantMode_ = groupNum > 0 ? TPL_MOE_FULL_LOAD : TPL_COMMON_FULL_LOAD;
        } else {
            quantMode_ = groupNum > 0 ? TPL_MOE_PER_TENSOR_FULL_LOAD : TPL_PER_TENSOR_FULL_LOAD;
        }
    } else {
        useDb = true;
        uint32_t ubAvail = static_cast<uint32_t>(maxUseUbSize) / static_cast<uint32_t>(calcDbSize);
        multiRowNumHeadCore = std::min({ubAvail, COMPARE_INT, rowPerHeadCore});
        multiRowNumTailCore = std::min({ubAvail, COMPARE_INT, rowPerTailCore});
        if (quantMode_ != TPL_PER_TENSOR_FULL_LOAD) {
            quantMode_ = groupNum > 0 ? TPL_MOE_FULL_LOAD : TPL_COMMON_FULL_LOAD;
        } else {
            quantMode_ = groupNum > 0 ? TPL_MOE_PER_TENSOR_FULL_LOAD : TPL_PER_TENSOR_FULL_LOAD;
        }
    }
}

void DynamicQuantRegbaseTiling::IsCapableForFullLoad(const gert::TilingContext* context)
{
    bool iscapable = false;
    int64_t tempNBlockSize = 1;
    int64_t tempMBlockSize = 1;
    int64_t tempBatchBlockSize = 1;
    uint64_t maxUseUbSize = ubSize - RESERVED_LENGTH * 8;
    for (int32_t i = 1; i <= nMaxLoopNum; i++) {
        // maximum m blocksize in ub, buffers in ub are: x, y, scale, offset, smoothscale
        // db * (mMaxBlockSize * nBlockSize * (sizeof(half) + sizeof(int8)) + 4 * nBlockSize * sizeof(float) + mLen * sizeof(half)) < maxUbSize
        tempNBlockSize = i * PER_CHANNEL_N_BASE_SIZE;
        tempMBlockSize = (static_cast<int64_t>(maxUseUbSize) - SUM_INPUT_OUTPUT_NUM * USE_BUFFER_NUM * sizeof(float) * tempNBlockSize - USE_BUFFER_NUM * sizeof(int16_t) * mLen) / 
                         (USE_BUFFER_NUM * (sizeof(int16_t) + sizeof(int8_t)) * tempNBlockSize);
        if (tempMBlockSize < mLen) {
            break;
        }
        iscapable = true;
        nBaseLoopNum = i;
        // try to load multiple batch each time
        tempBatchBlockSize = (static_cast<int64_t>(maxUseUbSize) - USE_BUFFER_NUM * sizeof(int16_t) * mLen) / 
                             (tempNBlockSize * (mLen * USE_BUFFER_NUM * (sizeof(int16_t) + sizeof(int8_t)) + SUM_INPUT_OUTPUT_NUM * USE_BUFFER_NUM * sizeof(float)));
        batchBlockSize = tempBatchBlockSize > 0 ? tempBatchBlockSize : 1;
        batchBlockNum = Ops::Base::CeilDiv(static_cast<int64_t>(totalBatchLen), batchBlockSize);
        batchTailBlockSize = totalBatchLen - (batchBlockNum - 1) * batchBlockSize;
        if (tempNBlockSize >= nLen) {
            break;
        }
    }
    if (iscapable) {
        quantMode_ = TPL_PER_CHANNEL_FULL_LOAD;
        OP_LOGD(context->GetNodeName(), "Entering perchannel full load template");
    }
    nBlockSize = PER_CHANNEL_N_BASE_SIZE * nBaseLoopNum;
    nBlockNum = Ops::Base::CeilDiv(nLen, nBlockSize);
    totalBlockNum = batchBlockNum * nBlockNum;
}

void DynamicQuantRegbaseTiling::IsCapableForRecompute(gert::TilingContext* context)
{
    bool iscapable = false;
    int64_t newSize;
    int64_t tempBlockNumForN;
    int64_t tempTailBlockNum;
    uint64_t maxUseUbSize = ubSize - RESERVED_LENGTH * 8;
    // recompute
    if (quantMode_ != TPL_PER_CHANNEL_FULL_LOAD && nBlockNum * totalBatchLen >= vectorCoreNum) {
        iscapable = true;
        tempBlockNumForN = Ops::Base::CeilDiv(nLen, nBlockSize);
        tempTailBlockNum = nLen - (tempBlockNumForN - 1) * nBlockSize;
        float ratio = tempTailBlockNum * 1.0 / nBlockSize;
        for (int i = 2; i <= nMaxLoopNum; i++) {
            newSize = PER_CHANNEL_N_BASE_SIZE * i;
            tempBlockNumForN = Ops::Base::CeilDiv(nLen, newSize);
            if (tempBlockNumForN * totalBatchLen < vectorCoreNum) {
                break;
            }
            tempTailBlockNum = nLen - (tempBlockNumForN - 1) * newSize;
            float newRatio = tempTailBlockNum * 1.0 / newSize;
            if (newRatio >= ratio) {
                nBlockSize = newSize; // 当尾块填充比更优时，更新N轴切块大小
                nBlockNum = tempBlockNumForN;
                nBaseLoopNum = i;
                ratio = newRatio;
            }
        }
        // maximum m blocksize in ub, buffers in ub are: x, smoothscale, y, col_max, col_min, scale, offset
        // db * (mMaxBlockSize * nBlockSize * (sizeof(half) + sizeof(int8)) + 4 * nBlockSize * sizeof(float) + mMaxBlockSize * sizeof(half)) < maxUbSize
        mBlockSize = (static_cast<int64_t>(maxUseUbSize) - 4 * USE_BUFFER_NUM * sizeof(float) * nBlockSize) / (nBlockSize * USE_BUFFER_NUM * (sizeof(int16_t) + sizeof(int8_t)) + USE_BUFFER_NUM * sizeof(int16_t));
        OPS_ERR_IF(
            mBlockSize <= 0, OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "mBlockSize", std::to_string(mBlockSize),
                                               "The value of mBlockSize must be greater than 0"),
            return);
        mBlockNum = Ops::Base::CeilDiv(mLen, mBlockSize);
        totalBlockNum = totalBatchLen * nBlockNum;
    }
    if (iscapable) {
        quantMode_ = TPL_PER_CHANNEL_RECOMPUTE;
        OP_LOGD(context->GetNodeName(), "Entering perchannel recompute template");
    }
}

void DynamicQuantRegbaseTiling::IsCapableForSplitM(gert::TilingContext* context)
{
    int64_t mMaxBlockSize;
    uint64_t maxUseUbSize = ubSize - RESERVED_LENGTH * 8;
    if (quantMode_ == TPL_PER_CHANNEL_SPLIT_M) {
        context->SetScheduleMode(SPLIT_M_SCHEDULE_MODE);
        // db * (mMaxBlockSize * nBlockSize * (sizeof(half) + sizeof(int8)) + 4 * nBlockSize * sizeof(float) + mMaxBlockSize * sizeof(half)) < maxUbSize
        mMaxBlockSize = (static_cast<int64_t>(maxUseUbSize) - 4 * USE_BUFFER_NUM * sizeof(float) * nBlockSize) / (nBlockSize * USE_BUFFER_NUM * (sizeof(int16_t) + sizeof(int8_t)) + USE_BUFFER_NUM * sizeof(int16_t));
        mBlockSize = mMaxBlockSize;
        mBlockNum = Ops::Base::CeilDiv(mLen, mMaxBlockSize);
        batchBlockSize = 1;
        batchBlockNum = totalBatchLen;
        batchTailBlockSize = 1;
        totalBlockNum = batchBlockNum * mBlockNum * nBlockNum;
        OP_LOGD(context->GetNodeName(), "Entering perchannel split m template");
    }
}

ge::graphStatus DynamicQuantRegbaseTiling::CalculateTilingDataForPerChannel(gert::TilingContext* context)
{
    const gert::StorageShape* xShape = context->GetInputShape(X_INDEX);
    nLen = xShape->GetStorageShape().GetDim(xShape->GetStorageShape().GetDimNum() - 1);
    OPS_ERR_IF(
            nLen <= 0, OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "nLen", std::to_string(nLen),
                                               "The value of nLen must be greater than 0"),
            return ge::GRAPH_FAILED);
    mLen = xShape->GetStorageShape().GetDim(xShape->GetStorageShape().GetDimNum() - PER_CHANNEL_EXCLUDE_DIM);
    OPS_ERR_IF(
            mLen <= 0, OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "mLen", std::to_string(mLen),
                                               "The value of mLen must be greater than 0"),
            return ge::GRAPH_FAILED);
    // 获取输入x的维度x-2
    size_t dimNum = xShape->GetStorageShape().GetDimNum() - PER_CHANNEL_EXCLUDE_DIM;
    uint64_t tempBatchLen = 1;
    for (size_t i = 0; i < dimNum; i++) {
        tempBatchLen *= xShape->GetStorageShape().GetDim(i);
    }
    // 设置totalBatchLen为x除了最后两维以外所有维度的乘积
    totalBatchLen = tempBatchLen;
    batchBlockSize = tempBatchLen;
    // default template for perchannel is splitm
    quantMode_ = TPL_PER_CHANNEL_SPLIT_M;
    // check full load template
    IsCapableForFullLoad(context);
    // check recompute template
    IsCapableForRecompute(context);
    // check split m template
    IsCapableForSplitM(context);
    headCoreNum = totalBlockNum % vectorCoreNum;
    headCoreNum = headCoreNum > 0 ? headCoreNum : vectorCoreNum;
    blockPerHead = Ops::Base::CeilDiv(totalBlockNum, static_cast<int64_t>(vectorCoreNum));
    blockPerTail = blockPerHead - 1;    
    coreNum = std::max(vectorCoreNum, ONE);
    // 默认开db
    useDb = true;
    SetTilingDataForPerChannel(context);
    PrintTilingDataForPerChannel(context);
    return ge::GRAPH_SUCCESS;
}

void DynamicQuantRegbaseTiling::SetTilingDataForPerChannel(gert::TilingContext* context) const {
    DynamicQuantTilingDataArch35* tilingDataPtr = context->GetTilingData<DynamicQuantTilingDataArch35>();
    tilingDataPtr->totalBatchLen = totalBatchLen;
    tilingDataPtr->mLen = mLen;
    tilingDataPtr->mBlockSize = mBlockSize;
    tilingDataPtr->mTailBlockSize = mLen - mBlockSize * (mBlockNum - 1);
    tilingDataPtr->mBlockNum = mBlockNum;
    tilingDataPtr->nLen = nLen;
    tilingDataPtr->nBlockSize = nBlockSize;
    tilingDataPtr->nTailBlockSize = nLen - nBlockSize * (nBlockNum - 1);
    tilingDataPtr->nBlockNum = nBlockNum;
    tilingDataPtr->nBaseSize = PER_CHANNEL_N_BASE_SIZE;
    tilingDataPtr->nBaseLoopNum = nBaseLoopNum;
    tilingDataPtr->blockPerHead = blockPerHead;
    tilingDataPtr->blockPerTail = blockPerTail;
    tilingDataPtr->totalBlockNum = totalBlockNum;
    tilingDataPtr->batchBlockSize = batchBlockSize;
    tilingDataPtr->batchTailBlockSize = batchTailBlockSize;
    tilingDataPtr->batchBlockNum = batchBlockNum;
    tilingDataPtr->dstTypeMax = dstTypeMax;
}

// 打印perchannel模板相关tiling参数
void DynamicQuantRegbaseTiling::PrintTilingDataForPerChannel(gert::TilingContext* context) const
{
    DynamicQuantTilingDataArch35* tilingDataPtr = context->GetTilingData<DynamicQuantTilingDataArch35>();
    OP_LOGD(
        context,
        "tilingData is totalBatchLen:%u, mLen:%u, mBlockSize:%u, mTailBlockSize:%u, mBlockNum:%u, "
        "nLen:%u, nBlockSize:%u, nTailBlockSize:%u, nBlockNum:%u, "
        "nBaseSize:%u, nBaseLoopNum:%u, blockPerHead:%u, blockPerTail:%u, totalBlockNum:%u"
        "batchBlockSize:%u, batchTailBlockSize:%u, batchBlockNum:%u",
        tilingDataPtr->totalBatchLen, tilingDataPtr->mLen, tilingDataPtr->mBlockSize, tilingDataPtr->mTailBlockSize,
        tilingDataPtr->mBlockNum, tilingDataPtr->nLen, tilingDataPtr->nBlockSize, tilingDataPtr->nTailBlockSize,
        tilingDataPtr->nBlockNum, tilingDataPtr->nBaseSize, tilingDataPtr->nBaseLoopNum, tilingDataPtr->blockPerHead,
        tilingDataPtr->blockPerTail, tilingDataPtr->totalBlockNum, tilingDataPtr->batchBlockSize,
        tilingDataPtr->batchTailBlockSize, tilingDataPtr->batchBlockNum);
}

// 获取核数vectorCoreNum和UB大小ubSize
ge::graphStatus DynamicQuantRegbaseTiling::GetCompileInfo(gert::TilingContext* context)
{
    auto compileInfo = context->GetCompileInfo<DynamicQuantCompileInfo>();
    OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
    vectorCoreNum = compileInfo->vectorCoreNum;
    ubSize = compileInfo->ubSize;
    OP_CHECK_IF(
        (vectorCoreNum <= 0 || ubSize <= 0),
        OP_LOGE(context, "RunFusionKernelTiling GetCompileInfo Failed, coreNum:%u, ubSize:%lu.",
            vectorCoreNum, ubSize),
        return ge::GRAPH_FAILED);
    return ge::GRAPH_SUCCESS;
}

ge::graphStatus DynamicQuantRegbaseTiling::RunFusionKernelTiling(gert::TilingContext* context)
{
    ResetLargeTilingParams();
    isSymmetrical_ = (context->GetComputeNodeOutputNum() != OUTPUT_NUM_DYNAMIC_QUANT_V2);

    OP_CHECK_IF(
        (GetCompileInfo(context) != ge::GRAPH_SUCCESS),
        OP_LOGE(context, "RunFusionKernelTiling GetCompileInfo failed."),
        return ge::GRAPH_FAILED);
    OP_CHECK_IF(
        (CheckOpParams(context) != ge::GRAPH_SUCCESS),
        OP_LOGE(context, "RunFusionKernelTiling CheckOpParams failed."),
        return ge::GRAPH_FAILED);

    if (isEmptyTensor) {
        return DoEmptyTensorTiling(context);
    }
    if (!isPerChannel_) {
        CalculateCoreNum(context);
        CalculateTilingData();
    } else {
        CalculateTilingDataForPerChannel(context);
    }
    SetTilingData(context);
    PrintTilingData(context);

    size_t* workSpaces = context->GetWorkspaceSizes(1);
    OP_CHECK_NULL_WITH_CONTEXT(context, workSpaces);

    size_t normalWorkSpace = coreNum * sizeof(float);
    size_t perChannelWorkSpace = totalBatchLen * nLen * sizeof(float);
    workSpaces[0] = SYS_WORKSPACE_SIZE;
    if (quantMode_ == TPL_PER_CHANNEL_SPLIT_M) {
        workSpaces[0] += isSymmetrical_ ? perChannelWorkSpace : perChannelWorkSpace * ASYMMETRICAL;
    } else {
        workSpaces[0] += isSymmetrical_ ? normalWorkSpace : normalWorkSpace * ASYMMETRICAL;
    }

    if (quantMode_ == TPL_PER_TENSOR_FULL_LOAD || quantMode_ == TPL_PER_TENSOR_LARGE_SHAPE ||
        quantMode_ == TPL_MOE_PER_TENSOR_FULL_LOAD || quantMode_ == TPL_MOE_PER_TENSOR_LARGE_SHAPE){
        context->SetScheduleMode(1); // 设置为batch mode模式，所有核同时启动
    }
    context->SetBlockDim(coreNum);
    return ge::GRAPH_SUCCESS;
}

} // namespace optiling