// Copyright (c) 2026 Huawei Technologies Co., Ltd
// All rights reserved.
//
// Licensed under the BSD 3-Clause License  (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <torch/extension.h>
#include "torch_npu/csrc/core/npu/NPUStream.h"
#include "kernel_operator.h"

constexpr uint32_t BUFFER_NUM = 2;  //tensor num for each queue
class KernelTrig {
public:
    __aicore__ inline KernelTrig() {}

    // Initialize the global memory and buffer queues
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR out_sin, GM_ADDR out_cos, GM_ADDR out_tan, uint32_t totalLength)
    {
        this->blockLength = totalLength / AscendC::GetBlockNum();
        this->tileNum = 8;
        this->tileLength = this->blockLength / this->tileNum / BUFFER_NUM;
        xGm.SetGlobalBuffer((__gm__ float *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        sinGm.SetGlobalBuffer((__gm__ float *)out_sin + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        cosGm.SetGlobalBuffer((__gm__ float *)out_cos + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        tanGm.SetGlobalBuffer((__gm__ float *)out_tan + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(float));
        pipe.InitBuffer(outQueueSin, BUFFER_NUM, this->tileLength * sizeof(float));
        pipe.InitBuffer(outQueueCos, BUFFER_NUM, this->tileLength * sizeof(float));
        pipe.InitBuffer(outQueueTan, BUFFER_NUM, this->tileLength * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        AscendC::LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
        AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    {
        AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        AscendC::LocalTensor<float> sinLocal = outQueueSin.AllocTensor<float>();
        AscendC::LocalTensor<float> cosLocal = outQueueCos.AllocTensor<float>();
        AscendC::LocalTensor<float> tanLocal = outQueueTan.AllocTensor<float>();

        AscendC::Sin(sinLocal, xLocal, this->tileLength);
        AscendC::Cos(cosLocal, xLocal, this->tileLength);
        AscendC::Tan(tanLocal, xLocal, this->tileLength);

        outQueueSin.EnQue<float>(sinLocal);
        outQueueCos.EnQue<float>(cosLocal);
        outQueueTan.EnQue<float>(tanLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        // Copy the sin, cos, and tan values from local memory to global memory (inplace modification)
        AscendC::LocalTensor<float> sinLocal = outQueueSin.DeQue<float>();
        AscendC::LocalTensor<float> cosLocal = outQueueCos.DeQue<float>();
        AscendC::LocalTensor<float> tanLocal = outQueueTan.DeQue<float>();
        AscendC::DataCopy(sinGm[progress * this->tileLength], sinLocal, this->tileLength);
        AscendC::DataCopy(cosGm[progress * this->tileLength], cosLocal, this->tileLength);
        AscendC::DataCopy(tanGm[progress * this->tileLength], tanLocal, this->tileLength);
        outQueueSin.FreeTensor(sinLocal);
        outQueueCos.FreeTensor(cosLocal);
        outQueueTan.FreeTensor(tanLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueSin, outQueueCos, outQueueTan;
    AscendC::GlobalTensor<float> xGm;
    AscendC::GlobalTensor<float> sinGm;
    AscendC::GlobalTensor<float> cosGm;
    AscendC::GlobalTensor<float> tanGm;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
};

__global__ __vector__ void trig_inplace_custom(GM_ADDR x, GM_ADDR out_sin, GM_ADDR out_cos, GM_ADDR out_tan,
                                               uint32_t totalLength)
{
    KernelTrig op;
    op.Init(x, out_sin, out_cos, out_tan, totalLength);
    op.Process();
}

namespace cpp_extension_acs {
at::Tensor ascendc_trig(const at::Tensor &x, const at::Tensor &out_sin, const at::Tensor &out_cos)
{
    auto acl_stream = c10_npu::getCurrentNPUStream().stream(true);
    at::Tensor out_tan = at::empty_like(x);
    uint32_t blockDim = 8;
    uint32_t totalLength = 1;
    for (uint32_t size : x.sizes()) {
        totalLength *= size;
    }
    // Launch the custom kernel using <<<>>>
    trig_inplace_custom<<<blockDim, nullptr, acl_stream>>>(
        (uint8_t *)(x.mutable_data_ptr()), (uint8_t *)(out_sin.mutable_data_ptr()),
        (uint8_t *)(out_cos.mutable_data_ptr()), (uint8_t *)(out_tan.mutable_data_ptr()), totalLength);
    return out_tan;
}

}  // namespace cpp_extension_acs

at::Tensor trig_impl_meta(const at::Tensor& x, const at::Tensor& out_sin, const at::Tensor& out_cos)
{
    return at::empty_like(x);
}

// Define a new operator
TORCH_LIBRARY_FRAGMENT(cpp_extension_acs, m)
{
    m.def("ascendc_trig(Tensor x, Tensor(a!) out_sin, Tensor(b!) out_cos) -> Tensor");
}

// Register implementation for the "PrivateUse1" backend
TORCH_LIBRARY_IMPL(cpp_extension_acs, PrivateUse1, m)
{
    m.impl("ascendc_trig", TORCH_FN(cpp_extension_acs::ascendc_trig));
}

// Define a simple model using the custom operation
TORCH_LIBRARY_IMPL(cpp_extension_acs, Meta, m)
{
    m.impl("ascendc_trig", &trig_impl_meta);
}