/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
 * MindIE is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * 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 FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

/*!
 * \file dropmask.h
 * \brief
 */

#ifndef DROPMASK_H
#define DROPMASK_H

#include "util.h"

using AscendC::DROPOUT_MODE_BIT_MISALIGN;
using AscendC::DropOutShapeInfo;
using AscendC::DropOut;

struct DropMaskInfo {
    // for compute dropout mask offset
    // 参数按B N G S1 S2全部切分设置进行偏移计算，没有切分的轴对应的参数设置为合适的0或者原始值
    int64_t n2G; // n2 * g
    int64_t gSize; // g
    int64_t s1Size; // s1
    int64_t s2Size; // s2
    int64_t gOutIdx; // g out index
    int64_t bSSOffset; // boidx * s1 * s2 ===bSSOffset
    int64_t n2OutIdx; // n out index
    int64_t s1OutIdx; // s1 out index   ===s1oIdx
    int64_t s1InnerIdx; // s1 inner index, 配比 ===loopIdx
    int64_t s1BaseSize; // S1基本块大小
    int64_t splitS1BaseSize; // s1 split size ===vec1S1BaseSize
    int64_t s2StartIdx; // s2 start index
    int64_t s2Idx; // s2 index =====s2LoopCount
    int64_t s2BaseNratioSize; // s2的配比长度: s2BaseSize(S2基本块大小) * nRatio

    // for copy in dropout mask
    uint32_t s1CopySize;
    uint32_t s2CopySize;
    int64_t s2TotalSize;

    // for compute dropout mask
    uint32_t firstAxis;
    uint32_t lstAxis;
    uint32_t maskLstAxis;
    int64_t vecCoreOffset = 0;
    float keepProb;

    bool boolMode;
};

template <bool hasDrop> __aicore__ inline int64_t ComputeDropOffset(DropMaskInfo &dropMaskInfo) {
    if constexpr (hasDrop == true) {
        // boidx * n2 * g* s1 * s2
        int64_t bOffset = dropMaskInfo.bSSOffset * dropMaskInfo.n2G;
        // n2oIdx * g * s1 *s2
        int64_t n2Offset = dropMaskInfo.n2OutIdx * dropMaskInfo.gSize * dropMaskInfo.s1Size * dropMaskInfo.s2Size;
        // goIdx * s1 * s2
        int64_t gOffset = dropMaskInfo.gOutIdx * dropMaskInfo.s1Size * dropMaskInfo.s2Size;
        // s1oIdx * s1BaseSize * s2Size + s1innerindex * vec1S1BaseSize * s2Size
        int64_t s1Offset = (dropMaskInfo.s1OutIdx * dropMaskInfo.s1BaseSize + dropMaskInfo.vecCoreOffset +
                               dropMaskInfo.s1InnerIdx * dropMaskInfo.splitS1BaseSize) *
            dropMaskInfo.s2Size;
        // s2StartIdx + s2index * s2BaseNratioSize
        int64_t s2Offset = dropMaskInfo.s2StartIdx + dropMaskInfo.s2Idx * dropMaskInfo.s2BaseNratioSize;
        return bOffset + n2Offset + gOffset + s1Offset + s2Offset;
    } else {
        return 0;
    }
}

template <bool hasDrop>
__aicore__ inline void CopyInDropMask(LocalTensor<uint8_t> &dstTensor, GlobalTensor<uint8_t> &srcBoolTensor,
    GlobalTensor<uint8_t> &srcByteTensor, DropMaskInfo &dropMaskInfo, int64_t alignedSize = blockBytes) {
    if constexpr (hasDrop == true) {
        int64_t dropMaskOffset = ComputeDropOffset<hasDrop>(dropMaskInfo);
        if (unlikely(dropMaskInfo.boolMode)) {
            BoolCopyIn(dstTensor, srcBoolTensor, dropMaskOffset, dropMaskInfo.s1CopySize, dropMaskInfo.s2CopySize,
                dropMaskInfo.s2TotalSize, alignedSize);
        } else {
            Bit2Int8CopyIn(dstTensor, srcByteTensor, dropMaskOffset, 1, dropMaskInfo.s1CopySize,
                dropMaskInfo.s2CopySize, dropMaskInfo.s2TotalSize, alignedSize);
        }
        return;
    }
}

template <typename T, bool hasDrop>
__aicore__ inline void ComputeDropMask(LocalTensor<T> &dstTensor, LocalTensor<T> &srcTensor,
    LocalTensor<uint8_t> &dropoutBuffer, LocalTensor<uint8_t> &tmpDropBuffer, DropMaskInfo &dropMaskInfo) {
    if constexpr (hasDrop == true) {
        DropOutShapeInfo dropOutShapeInfo;
        dropOutShapeInfo.firstAxis = dropMaskInfo.firstAxis;
        dropOutShapeInfo.srcLastAxis = dropMaskInfo.lstAxis;

        if (unlikely(dropMaskInfo.boolMode)) {
            dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis, blockBytes) * blockBytes;
            DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo);
        } else {
            dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis / byteBitRatio, blockBytes) * blockBytes;
            if (likely(dropMaskInfo.lstAxis / byteBitRatio % blockBytes == 0)) {
                DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo);
            } else {
                DropOut<T, false, DROPOUT_MODE_BIT_MISALIGN>(
                    dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo);
            }
        }
        return;
    }
}

#endif // DROPMASK_H