/**
 * 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.
 */

/* Generated By CANNBot */


/*!
 * \file inv.h
 * \brief Inv 算子 Kernel 类定义（arch35 / Ascend950）
 *
 * Inv(x) = 1 / x
 *
 * 精度策略：
 *   - float32：直接 Div(ones, x)（不使用 Reciprocal，精度要求）
 *   - float16 / bfloat16：Cast to fp32 -> Div(ones, x) -> Cast back
 *   - int32：截断向零整数倒数，等价三值映射 {+1->+1, -1->-1, 其他->0}。
 *            整型 Compare(==1)/Compare(==-1) + 嵌套 Select 实现，禁止对 INT_MIN 取负、不经 fp32 中转。
 *
 * 模板分发（由 TilingKey 对应）：
 *   key 0 -> Inv<float>
 *   key 1 -> Inv<half>
 *   key 2 -> Inv<bfloat16_t>
 *   key 3 -> Inv<int32_t>
 *
 * Buffer 布局：
 *   inputQueue (BUFFER_NUM=2)  : ubFactor * sizeof(T)
 *   outputQueue(BUFFER_NUM=2)  : ubFactor * sizeof(T)
 *   tmpBuf1_                   : ubFactor * sizeof(float)  -- 仅 FP16/BF16 路径使用（xFloat32）
 *   tmpBuf2_                   : ubFactor * sizeof(float)  -- 仅浮点路径使用（ones 向量，Init 中一次性填充）
 *   int32 路径不分配 tmpBuf1_/tmpBuf2_（纯整型 Compare/Select，无 fp32 中转，落实评审 MED-1/MED-2）
 */
#ifndef OPS_MATH_INV_KERNEL_ARCH35_INV_H_
#define OPS_MATH_INV_KERNEL_ARCH35_INV_H_

#include "kernel_operator.h"
#include "kernel_tiling/kernel_tiling.h"
#include "op_kernel/platform_util.h"
#include "inv_tiling_data.h"
#include "inv_tiling_key.h"

namespace NsInv {

using AscendC::TPipe;
using AscendC::TQue;
using AscendC::TBuf;
using AscendC::QuePosition;
using AscendC::GlobalTensor;
using AscendC::LocalTensor;
using AscendC::DataCopyExtParams;
using AscendC::DataCopyPad;
using AscendC::DataCopyPadExtParams;
using AscendC::RoundMode;
using AscendC::GetBlockIdx;
using AscendC::Cast;
using AscendC::Duplicate;
using AscendC::Div;

static constexpr int32_t BUFFER_NUM = 2;

template <typename T>
class Inv {
public:
    __aicore__ inline Inv() {}

    __aicore__ inline void Init(GM_ADDR self, GM_ADDR out, const InvTilingData* tilingData);
    __aicore__ inline void Process();

private:
    __aicore__ inline void CopyIn(int64_t gmOffset, int64_t currentNum);
    __aicore__ inline void Compute(int64_t currentNum);
    __aicore__ inline void CopyOut(int64_t gmOffset, int64_t currentNum);

    __aicore__ inline void ComputeFloat32(LocalTensor<float>& xLocal,
                                          LocalTensor<float>& yLocal,
                                          int64_t alignedNum);
    template <typename SrcT>
    __aicore__ inline void ComputeWithCast(LocalTensor<SrcT>& xLocal,
                                           LocalTensor<SrcT>& yLocal,
                                           int64_t currentNum,
                                           int64_t alignedNum);
    // int32 截断向零整数倒数：三值映射 Select(x==-1, -1, Select(x==1, 1, 0))
    // 纯整型 Compare/Select，禁止对 INT_MIN 取负、不经 fp32 中转
    __aicore__ inline void ComputeInt32(LocalTensor<int32_t>& xLocal,
                                        LocalTensor<int32_t>& yLocal,
                                        int64_t currentNum);

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inputQueue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outputQueue;
    TBuf<QuePosition::VECCALC> tmpBuf1_;
    TBuf<QuePosition::VECCALC> tmpBuf2_;

    GlobalTensor<T> selfGM_;
    GlobalTensor<T> outGM_;

    int64_t blockOffset_ = 0;
    int64_t blockLen_ = 0;
    int64_t ubFactor_ = 0;
};

template <typename T>
__aicore__ inline void Inv<T>::Init(GM_ADDR self, GM_ADDR out, const InvTilingData* tilingData)
{
    ubFactor_ = tilingData->ubFactor;

    if (tilingData->totalElements == 0 || tilingData->blockFactor == 0) {
        blockOffset_ = 0;
        blockLen_ = 0;
        return;
    }

    blockOffset_ = tilingData->blockFactor * static_cast<int64_t>(GetBlockIdx());
    int64_t remaining = tilingData->totalElements - blockOffset_;
    if (remaining <= 0) {
        blockLen_ = 0;
        return;
    }
    blockLen_ = (remaining > tilingData->blockFactor) ? tilingData->blockFactor : remaining;

    selfGM_.SetGlobalBuffer((__gm__ T*)self + blockOffset_, blockLen_);
    outGM_.SetGlobalBuffer((__gm__ T*)out + blockOffset_, blockLen_);

    pipe.InitBuffer(inputQueue,  BUFFER_NUM, ubFactor_ * sizeof(T));
    pipe.InitBuffer(outputQueue, BUFFER_NUM, ubFactor_ * sizeof(T));
    // MED-1：fp32 专用工作缓冲（tmpBuf1_ xFloat / tmpBuf2_ ones）仅浮点路径分配。
    // int32 路径为纯整型 Compare/Select，不经 fp32 中转，二者均不分配、不 Duplicate(1.0f)，
    // 与 Host bytesPerElem(int32) 严格一一对应（落实评审 MED-2）。
    if constexpr (std::is_floating_point_v<T> ||
                  std::is_same_v<T, half> ||
                  std::is_same_v<T, bfloat16_t>) {
        if constexpr (!std::is_same_v<T, float>) {
            pipe.InitBuffer(tmpBuf1_, ubFactor_ * sizeof(float));
        }
        pipe.InitBuffer(tmpBuf2_, ubFactor_ * sizeof(float));

        // ones 向量是常量，一次性填充，避免每次 Compute 重复执行
        LocalTensor<float> ones = tmpBuf2_.template Get<float>();
        Duplicate(ones, 1.0f, static_cast<int32_t>(ubFactor_));
    }
}

template <typename T>
__aicore__ inline void Inv<T>::CopyIn(int64_t gmOffset, int64_t currentNum)
{
    LocalTensor<T> xLocal = inputQueue.template AllocTensor<T>();
    DataCopyExtParams copyParams;
    copyParams.blockCount = 1;
    copyParams.blockLen = static_cast<uint32_t>(currentNum * sizeof(T));
    copyParams.srcStride = 0;
    copyParams.dstStride = 0;
    DataCopyPad(xLocal, selfGM_[gmOffset], copyParams, {false, 0, 0, 0});
    inputQueue.EnQue(xLocal);
}

template <typename T>
__aicore__ inline void Inv<T>::CopyOut(int64_t gmOffset, int64_t currentNum)
{
    LocalTensor<T> yLocal = outputQueue.template DeQue<T>();
    DataCopyExtParams copyParams;
    copyParams.blockCount = 1;
    copyParams.blockLen = static_cast<uint32_t>(currentNum * sizeof(T));
    copyParams.srcStride = 0;
    copyParams.dstStride = 0;
    DataCopyPad(outGM_[gmOffset], yLocal, copyParams);
    outputQueue.FreeTensor(yLocal);
}

template <typename T>
__aicore__ inline void Inv<T>::ComputeFloat32(LocalTensor<float>& xLocal,
                                              LocalTensor<float>& yLocal,
                                              int64_t alignedNum)
{
    LocalTensor<float> ones = tmpBuf2_.template Get<float>();
    Div(yLocal, ones, xLocal, static_cast<int32_t>(alignedNum));
}

template <typename T>
template <typename SrcT>
__aicore__ inline void Inv<T>::ComputeWithCast(LocalTensor<SrcT>& xLocal,
                                               LocalTensor<SrcT>& yLocal,
                                               int64_t currentNum,
                                               int64_t alignedNum)
{
    LocalTensor<float> xFloat = tmpBuf1_.template Get<float>();
    LocalTensor<float> ones = tmpBuf2_.template Get<float>();

    Cast(xFloat, xLocal, RoundMode::CAST_NONE, static_cast<uint32_t>(alignedNum));
    Div(xFloat, ones, xFloat, static_cast<int32_t>(alignedNum));
    Cast(yLocal, xFloat, RoundMode::CAST_RINT, static_cast<uint32_t>(alignedNum));
}

// int32 截断向零整数倒数：三值映射 y = Select(x==-1, -1, Select(x==1, 1, 0))。
// 纯整型 Compare/Select（MicroAPI，arch35 实证：adjacent_difference / floor_mod），
// 不经 fp32 中转，不对 INT_MIN 取负 —— x=0/|x|>=2/INT_MIN/INT_MAX 均落入 else->0。
//
// 迭代二 A1-Main 边界/对齐落实（结构对齐同仓已上库 adjacent_difference 的 MicroAPI 尾块模式）：
//   - 多 rank（1D~8D）：Host 已将任意 rank 扁平化为 totalElements，本分支只处理一维 currentNum，
//                       对 rank 无依赖。
//   - 256B 对齐 / 尾块：vLength = VRegSize/sizeof(int32)（arch35=64），loopTimes 向上取整覆盖尾块；
//                       remain 经 UpdateMask<int32> 按引用逐段自递减，最后一段 mask 仅覆盖有效元素。
//                       Store 始终带 mask，DataCopyPad 仅搬入 currentNum 个元素，尾部 UB 脏数据不会写出，
//                       [7]/[65] 等非 vLength 整数倍长度下无越界、无脏写。
//   - 空 Tensor：currentNum<=0 直接返回（Process() 亦已用 blockLen_<=0 早返回，此处为双重防御）。
template <typename T>
__aicore__ inline void Inv<T>::ComputeInt32(LocalTensor<int32_t>& xLocal,
                                            LocalTensor<int32_t>& yLocal,
                                            int64_t currentNum)
{
    if (currentNum <= 0) {
        return;
    }

    namespace MA = AscendC::MicroAPI;
    using AscendC::CMPMODE;
    __local_mem__ int32_t* srcAddr = (__ubuf__ int32_t*)xLocal.GetPhyAddr();
    __local_mem__ int32_t* dstAddr = (__ubuf__ int32_t*)yLocal.GetPhyAddr();

    constexpr uint32_t vLength = Ops::Base::GetVRegSize() / sizeof(int32_t);
    // currentNum <= ubFactor，loopTimes 远小于 uint16_t 上界(65535)，强转不溢出
    uint16_t loopTimes = static_cast<uint16_t>((currentNum + vLength - 1) / vLength);

    __VEC_SCOPE__
    {
        MA::RegTensor<int32_t> regZero;
        MA::RegTensor<int32_t> regOne;
        MA::RegTensor<int32_t> regNegOne;
        MA::MaskReg fullMask = MA::CreateMask<int32_t>();
        MA::Duplicate(regZero, 0, fullMask);
        MA::Duplicate(regOne, 1, fullMask);
        MA::Duplicate(regNegOne, -1, fullMask);

        uint32_t remain = static_cast<uint32_t>(currentNum);
        for (uint16_t i = 0; i < loopTimes; i++) {
            MA::MaskReg mask = MA::UpdateMask<int32_t>(remain);
            MA::RegTensor<int32_t> regX;
            MA::RegTensor<int32_t> regTmp;
            MA::RegTensor<int32_t> regY;
            MA::MaskReg maskEqP;
            MA::MaskReg maskEqN;

            MA::DataCopy<int32_t, MA::LoadDist::DIST_NORM>(regX, srcAddr + i * vLength);
            // in/out 同为 int32 等位宽，Compare/Select 无需 MaskPack；若后续扩展到混合位宽需复核
            // x == 1 ? 1 : 0
            MA::Compare<int32_t, CMPMODE::EQ>(maskEqP, regX, regOne, mask);
            MA::Select(regTmp, regOne, regZero, maskEqP);
            // x == -1 ? -1 : tmp（x 同时为 1 和 -1 不可能，互斥安全）
            MA::Compare<int32_t, CMPMODE::EQ>(maskEqN, regX, regNegOne, mask);
            MA::Select(regY, regNegOne, regTmp, maskEqN);
            MA::DataCopy<int32_t, MA::StoreDist::DIST_NORM>(dstAddr + i * vLength, regY, mask);
        }
    }
}

template <typename T>
__aicore__ inline void Inv<T>::Compute(int64_t currentNum)
{
    LocalTensor<T> xLocal = inputQueue.template DeQue<T>();
    LocalTensor<T> yLocal = outputQueue.template AllocTensor<T>();

    constexpr int64_t floatBlock = 32 / sizeof(float);
    constexpr int64_t typeBlock  = 32 / sizeof(T);
    constexpr int64_t alignBlock = (floatBlock > typeBlock) ? floatBlock : typeBlock;
    int64_t alignedNum = ((currentNum + alignBlock - 1) / alignBlock) * alignBlock;

    if constexpr (std::is_same_v<T, float>) {
        ComputeFloat32(xLocal, yLocal, alignedNum);
    } else if constexpr (std::is_integral_v<T>) {
        ComputeInt32(xLocal, yLocal, currentNum);
    } else {
        ComputeWithCast(xLocal, yLocal, currentNum, alignedNum);
    }

    outputQueue.template EnQue<T>(yLocal);
    inputQueue.FreeTensor(xLocal);
}

template <typename T>
__aicore__ inline void Inv<T>::Process()
{
    if (blockLen_ <= 0) {
        return;
    }
    int64_t loopCount = (blockLen_ + ubFactor_ - 1) / ubFactor_;
    for (int64_t i = 0; i < loopCount; i++) {
        int64_t gmOffset = i * ubFactor_;
        int64_t currentNum = (i == (loopCount - 1)) ? (blockLen_ - gmOffset) : ubFactor_;
        CopyIn(gmOffset, currentNum);
        Compute(currentNum);
        CopyOut(gmOffset, currentNum);
    }
}

} // namespace NsInv
#endif // OPS_MATH_INV_KERNEL_ARCH35_INV_H_