#include "../../../elementwise/nvidia/elementwise_nvidia.cuh"

#include "../cuda/kernel.cuh"
#include "hardswish_nvidia.cuh"

#include <cuda_runtime.h>

namespace op::hardswish::nvidia {
namespace {

inline bool can_use_contiguous_fast_path(const op::elementwise::ElementwiseInfo &info) {
    return info.isOutputContiguous() && info.getInputSize() == 1 && info.getInputContiguous()[0] && !info.getInputBroadcasted()[0];
}

template <typename T>
__global__ void hardswish_contiguous_kernel(size_t numel, T *out, const T *in) {
    const auto op = op::hardswish::cuda::HardSwishOp{};
    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    while (idx < numel) {
        out[idx] = op(in[idx]);
        idx += blockDim.x * gridDim.x;
    }
}

template <typename T>
infiniStatus_t launch_fast_path(size_t numel,
                                void *output,
                                const std::vector<const void *> &inputs,
                                void *stream) {
    if (numel == 0) {
        return INFINI_STATUS_SUCCESS;
    }

    constexpr int BLOCK_SIZE = 256;
    int grid = static_cast<int>((numel + BLOCK_SIZE - 1) / BLOCK_SIZE);
    grid = std::min(grid, 65535);

    auto *out_ptr = reinterpret_cast<T *>(output);
    auto *in_ptr = reinterpret_cast<const T *>(inputs[0]);
    auto cuda_stream = reinterpret_cast<cudaStream_t>(stream);

    hardswish_contiguous_kernel<<<grid, BLOCK_SIZE, 0, cuda_stream>>>(numel, out_ptr, in_ptr);
    cudaError_t err = cudaGetLastError();
    return err == cudaSuccess ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR;
}

} // namespace

Descriptor::~Descriptor() = default;

infiniStatus_t Descriptor::create(
    infiniopHandle_t handle_,
    Descriptor **desc_ptr,
    infiniopTensorDescriptor_t out_desc,
    std::vector<infiniopTensorDescriptor_t> input_desc_vec) {

    auto handle = reinterpret_cast<device::nvidia::Handle *>(handle_);
    auto dtype = out_desc->dtype();

    const auto &input_desc = input_desc_vec.at(0);
    const auto &output_shape = out_desc->shape();
    const auto &input_shape = input_desc->shape();

    CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64);

    CHECK_SAME_SHAPE(output_shape, input_shape);

    CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec)

    return INFINI_STATUS_SUCCESS;
}

infiniStatus_t Descriptor::calculate(
    void *workspace,
    size_t workspace_size,
    void *output,
    std::vector<const void *> inputs,
    void *stream) const {

    const bool fast_path = can_use_contiguous_fast_path(_info);
    if (fast_path) {
        switch (_dtype) {
        case INFINI_DTYPE_BF16:
            return launch_fast_path<cuda_bfloat16>(_info.getOutputSize(), output, inputs, stream);
        case INFINI_DTYPE_F16:
            return launch_fast_path<half>(_info.getOutputSize(), output, inputs, stream);
        case INFINI_DTYPE_F32:
            return launch_fast_path<float>(_info.getOutputSize(), output, inputs, stream);
        case INFINI_DTYPE_F64:
            return launch_fast_path<double>(_info.getOutputSize(), output, inputs, stream);
        default:
            break;
        }
    }

    if (workspace_size < _workspace_size) {
        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
    }

    switch (_dtype) {
    case INFINI_DTYPE_BF16:
        return _device_info->calculate<256, cuda::HardSwishOp, cuda_bfloat16>(_info, workspace, output, inputs, stream);
    case INFINI_DTYPE_F16:
        return _device_info->calculate<256, cuda::HardSwishOp, half>(_info, workspace, output, inputs, stream);
    case INFINI_DTYPE_F32:
        return _device_info->calculate<256, cuda::HardSwishOp, float>(_info, workspace, output, inputs, stream);
    case INFINI_DTYPE_F64:
        return _device_info->calculate<256, cuda::HardSwishOp, double>(_info, workspace, output, inputs, stream);
    default:
        return INFINI_STATUS_BAD_TENSOR_DTYPE;
    }

    return INFINI_STATUS_SUCCESS;
}
} // namespace op::hardswish::nvidia
