/**
 * This file is part of the OpenBOAT project at Harbin Institute of Technology (HIT)
 * and is contributed to the CANN Open Software.
 *
 * Copyright (c) 2025 AISS Group, Harbin Institute of Technology (HIT).
 * All Rights Reserved.
 *
 * Authors (accounts):
 * - Li Zhi <@hitLeechi>
 * - Su Tonghua <@sutonghua>
 *
 * This program is free software: you can redistribute it and/or modify it.
 * Licensed under the CANN Open Software License Agreement Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * See the LICENSE file at the root of the repository for the full text of the License.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTIES OF ANY KIND, EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 */

/*!
 * \file range.cpp
 * \brief
*/

#ifndef RANGE_H
#define RANGE_H

#include "kernel_operator.h"
#include "kernel_tiling/kernel_tiling.h"
#include "range_tiling_data.h"
#include "range_tiling_key.h"
using namespace std;
constexpr uint32_t OUTPUT_DATA_TYPE_SIZE = 4;                           // 输出float的字节数
constexpr uint32_t BLOCK_SIZE = 32;                                     // 32B内存对齐基准
constexpr uint32_t TILE_ELEM_NUM = BLOCK_SIZE / OUTPUT_DATA_TYPE_SIZE;  // float的32B对齐元素数
constexpr uint32_t BUFFER_NUM = 2;                                      // 双缓冲
namespace NsRange {
using namespace AscendC;
template <typename TStart, typename TEnd, typename TStep> 
class Range{
    public:
    __aicore__ inline Range() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR s, GM_ADDR z, uint32_t _totalLength,
                                uint32_t _blockLength, uint32_t _tileNum,
                                uint32_t _lastTileLength, uint32_t _formerLength,
                                uint32_t _formerTileNum, uint32_t _formerLastTileLength,
                                uint32_t _formerNum, uint32_t _tailLength,
                                uint32_t _tailTileNum, uint32_t _tailLastTileLength,
                                uint32_t _isEvenCore)
    {
        // 不论输入是什么类型，kernel侧始终使用float进行计算
        this->start = *reinterpret_cast<const __gm__ float*>(x);
        this->end = *reinterpret_cast<const __gm__ float*>(y);
        this->step = *reinterpret_cast<const __gm__ float*>(s);
        this->origSize = _totalLength;

        uint32_t coreId = AscendC::GetBlockIdx();
        uint32_t currentBlockLength = 0;

        if (_isEvenCore)
        {
            this->blockLength = _blockLength;
            this->tileNum = _tileNum;
            this->lastTileLength = _lastTileLength;
            this->blockOffset = blockLength * coreId;
            currentBlockLength = blockLength;
        }
        else
        {
            if (coreId < _formerNum)
            {
                this->blockLength = _formerLength;
                this->tileNum = _formerTileNum;
                this->lastTileLength = _formerLastTileLength;
                this->blockOffset = _formerLength * coreId;
                currentBlockLength = _formerLength;
            }
            else
            {
                this->blockLength = _tailLength;
                this->tileNum = _tailTileNum;
                this->lastTileLength = _tailLastTileLength;
                this->blockOffset = _formerLength * _formerNum + _tailLength * (coreId - _formerNum);
                currentBlockLength = _tailLength;
            }
        }
        zGm.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(z) + this->blockOffset, currentBlockLength);
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, TILE_ELEM_NUM * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        if (origSize == 1)
        {
            AscendC::LocalTensor<float> tileOutput = outQueueZ.AllocTensor<float>();
            for (int32_t i = 0; i < TILE_ELEM_NUM; i++) tileOutput.SetValue(i, (i == 0) ? start : 0.0f);
            outQueueZ.EnQue(tileOutput);
            CopyOut(0);
            return;
        }
        for (int32_t tileIdx = 0; tileIdx < tileNum; tileIdx++)
        {
            Compute(tileIdx);
            CopyOut(tileIdx);
        }
    }

private:
    __aicore__ inline void CopyIn() {}
    
    __aicore__ inline void Compute(int32_t tileIdx)
    {
        AscendC::LocalTensor<float> tileOutput = outQueueZ.AllocTensor<float>();
        int32_t tileElemNum = TILE_ELEM_NUM;  // 8
        int32_t tileGlobalStart = blockOffset + tileIdx * tileElemNum;

        for (int32_t i = 0; i < tileElemNum; i++)
        {
            int32_t globalIdx = tileGlobalStart + i;
            if (globalIdx < origSize)
                tileOutput.SetValue(i, start + step * static_cast<float>(globalIdx));
            else
                tileOutput.SetValue(i, 0.0f);
        }
        outQueueZ.EnQue(tileOutput);
       
    }

    __aicore__ inline void CopyOut(int32_t tileIdx)
    {
        AscendC::LocalTensor<float> zLocal = outQueueZ.DeQue<float>();
        int32_t copyElemNum;
        if (tileIdx == tileNum - 1) copyElemNum = lastTileLength;
        else copyElemNum = TILE_ELEM_NUM;
        
        int32_t copyOffset = tileIdx * TILE_ELEM_NUM;
        if (copyOffset + copyElemNum > blockLength) copyElemNum = blockLength - copyOffset;
        
        AscendC::DataCopy(zGm[copyOffset], zLocal, copyElemNum);
        outQueueZ.FreeTensor(zLocal);

    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::GlobalTensor<float> zGm;

    uint32_t blockLength;    // 当前核心处理的float总元素数
    uint32_t tileNum;        // 当前核心的Tile总数（每个Tile 8个float元素）
    uint32_t lastTileLength; // 最后一个Tile的float元素数
    uint32_t blockOffset;    // 当前核心的全局float元素偏移
    uint32_t origSize;       // 原始有效元素数
    float start;
    float end;
    float step;
};

} //namespace NsRange    
#endif // RANGE_H