/**
 * 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 sqrt.asc
 * \brief Pointwise Sqrt operator migrated from cann-samples
 */

#include <ATen/Operators.h>
#include <torch/all.h>
#include <torch/library.h>
#include "torch_npu/csrc/core/npu/NPUStream.h"
#include "torch_npu/csrc/framework/OpCommand.h"
#include "kernel_operator.h"
#include "platform/platform_ascendc.h"
#include <type_traits>

namespace ascend_ops {
namespace Sqrt {

// Register the operator's schema
TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m)
{
    m.def("sqrt(Tensor x) -> Tensor");
}

// Meta function implementation of Sqrt
torch::Tensor sqrt_meta(const torch::Tensor &x)
{
    auto z = torch::empty_like(x);
    return z;
}

// Register the Meta implementation
TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, Meta, m)
{
    m.impl("sqrt", sqrt_meta);
}

std::tuple<int64_t, int64_t, int64_t> calc_tiling_params(int64_t totalLength)
{
    constexpr static int64_t MIN_ELEMS_PER_CORE = 1024;
    constexpr static int64_t PIPELINE_DEPTH = 2;
    constexpr static int64_t BUFFER_NUM = 2; // 1 input + 1 output
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
    uint64_t ubSize;
    ascendcPlatform->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    int64_t coreNum = ascendcPlatform->GetCoreNumAiv();
    TORCH_CHECK(coreNum > 0, "coreNum must be positive.");
    int64_t numBlocks = std::min(coreNum, (totalLength + MIN_ELEMS_PER_CORE - 1) / MIN_ELEMS_PER_CORE);
    numBlocks = std::max(numBlocks, static_cast<int64_t>(1));
    int64_t blockLength = (totalLength + numBlocks - 1) / numBlocks;
    int64_t tileSize = ubSize / PIPELINE_DEPTH / BUFFER_NUM;
    return std::make_tuple(numBlocks, blockLength, tileSize);
}

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

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR z,
        int64_t totalLength, int64_t blockLength, uint32_t tileSize)
    {
        xGm_.SetGlobalBuffer((__gm__ T *)x + blockLength * AscendC::GetBlockIdx());
        zGm_.SetGlobalBuffer((__gm__ T *)z + blockLength * AscendC::GetBlockIdx());

        pipe_.InitBuffer(inQueueX_, PIPELINE_DEPTH, tileSize);
        pipe_.InitBuffer(outQueueZ_, PIPELINE_DEPTH, tileSize);

        int64_t currentBlockLength = totalLength - AscendC::GetBlockIdx() * blockLength;
        if (currentBlockLength > blockLength) {
            currentBlockLength = blockLength;
        }
        elementNumPerTile_ = tileSize / sizeof(T);
        tileNum_ = currentBlockLength / elementNumPerTile_;
        tailTileElementNum_ = currentBlockLength - tileNum_ * elementNumPerTile_;
    }

    __aicore__ inline void Process()
    {
        for (int64_t i = 0; i < tileNum_; ++i) {
            int64_t offset = i * elementNumPerTile_;
            CopyIn(offset, elementNumPerTile_);
            Compute(elementNumPerTile_);
            CopyOut(offset, elementNumPerTile_);
        }
        if (tailTileElementNum_ > 0) {
            int64_t offset = tileNum_ * elementNumPerTile_;
            CopyIn(offset, tailTileElementNum_);
            Compute(tailTileElementNum_);
            CopyOut(offset, tailTileElementNum_);
        }
    }

private:
    __aicore__ inline void CopyIn(int64_t offset, int64_t count)
    {
        AscendC::DataCopyExtParams copyParams;
        copyParams.blockCount = 1;
        copyParams.blockLen = count * sizeof(T);
        copyParams.srcStride = 0;
        copyParams.dstStride = 0;
        AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        AscendC::LocalTensor<T> xLocal = inQueueX_.AllocTensor<T>();
        AscendC::DataCopyPad(xLocal, xGm_[offset], copyParams, padParams);
        inQueueX_.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int64_t count)
    {
        AscendC::LocalTensor<T> xLocal = inQueueX_.DeQue<T>();
        AscendC::LocalTensor<T> zLocal = outQueueZ_.AllocTensor<T>();
        AscendC::Sqrt(zLocal, xLocal, count);
        outQueueZ_.EnQue(zLocal);
        inQueueX_.FreeTensor(xLocal);
    }

    __aicore__ inline void CopyOut(int64_t offset, int64_t count)
    {
        AscendC::LocalTensor<T> zLocal = outQueueZ_.DeQue<T>();
        AscendC::DataCopyExtParams copyParams;
        copyParams.blockCount = 1;
        copyParams.blockLen = count * sizeof(T);
        copyParams.srcStride = 0;
        copyParams.dstStride = 0;
        AscendC::DataCopyPad(zGm_[offset], zLocal, copyParams);
        outQueueZ_.FreeTensor(zLocal);
    }

    static constexpr int64_t PIPELINE_DEPTH = 2;
    AscendC::TPipe pipe_;
    AscendC::GlobalTensor<T> xGm_, zGm_;
    AscendC::TQue<AscendC::TPosition::VECIN, PIPELINE_DEPTH> inQueueX_;
    AscendC::TQue<AscendC::TPosition::VECOUT, PIPELINE_DEPTH> outQueueZ_;
    int64_t elementNumPerTile_ = 0;
    int64_t tileNum_ = 0;
    int64_t tailTileElementNum_ = 0;
};

template <typename T>
__global__ __aicore__ __vector__ void sqrt_kernel(
    GM_ADDR x, GM_ADDR z, int64_t totalLength, int64_t blockLength, uint32_t tileSize)
{
    KernelSqrt<T> op;
    op.Init(x, z, totalLength, blockLength, tileSize);
    op.Process();
}

torch::Tensor sqrt_npu(const torch::Tensor &x)
{
    const c10::OptionalDeviceGuard guard(x.device());
    auto z = sqrt_meta(x);
    auto stream = c10_npu::getCurrentNPUStream().stream(false);
    int64_t totalLength = x.numel();
    int64_t numBlocks, blockLength, tileSize;
    std::tie(numBlocks, blockLength, tileSize) = calc_tiling_params(totalLength);
    auto x_ptr = (GM_ADDR)x.data_ptr();
    auto z_ptr = (GM_ADDR)z.data_ptr();
    auto acl_call = [=]() -> int {
        AT_DISPATCH_SWITCH(
            x.scalar_type(), "sqrt_npu",
            AT_DISPATCH_CASE(torch::kFloat32, [&] {
                using scalar_t = float;
                sqrt_kernel<scalar_t><<<numBlocks, nullptr, stream>>>(x_ptr, z_ptr, totalLength, blockLength, tileSize);
            })
        );
        return 0;
    };
    at_npu::native::OpCommand::RunOpApi("Sqrt", acl_call);
    return z;
}

// Register the NPU implementation
TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, PrivateUse1, m)
{
    m.impl("sqrt", sqrt_npu);
}

} // namespace Sqrt
} // namespace ascend_ops