// Minimal cuDNN legacy activation API check: sigmoid(x) = 1 / (1 + exp(-x)).
#include <cuda_runtime.h>
#include <cudnn.h>
#include <cmath>
#include <cstdlib>
#include <iostream>

void check_cuda(cudaError_t status) {
    if (status != cudaSuccess) {
        std::cerr << "CUDA: " << cudaGetErrorString(status) << '\n';
        std::exit(EXIT_FAILURE);
    }
}

void check_cudnn(cudnnStatus_t status) {
    if (status != CUDNN_STATUS_SUCCESS) {
        std::cerr << "cuDNN: " << cudnnGetErrorString(status) << '\n';
        std::exit(EXIT_FAILURE);
    }
}

int main() {
    constexpr int count = 10;
    float input[count], output[count];
    for (int i = 0; i < count; ++i) input[i] = static_cast<float>(i - 5);
    check_cuda(cudaSetDevice(0));
    cudaDeviceProp device;
    check_cuda(cudaGetDeviceProperties(&device, 0));
    std::cout << "GPU: " << device.name << ", cuDNN: " << cudnnGetVersion() << '\n';

    cudnnHandle_t handle;
    cudnnTensorDescriptor_t tensor;
    cudnnActivationDescriptor_t activation;
    check_cudnn(cudnnCreate(&handle));
    check_cudnn(cudnnCreateTensorDescriptor(&tensor));
    check_cudnn(cudnnSetTensor4dDescriptor(tensor, CUDNN_TENSOR_NCHW,
                                         CUDNN_DATA_FLOAT, 1, 1, 1, count));
    check_cudnn(cudnnCreateActivationDescriptor(&activation));
    check_cudnn(cudnnSetActivationDescriptor(activation, CUDNN_ACTIVATION_SIGMOID,
                                            CUDNN_PROPAGATE_NAN, 0.0));
    float *device_input, *device_output;
    check_cuda(cudaMalloc(reinterpret_cast<void **>(&device_input), sizeof(input)));
    check_cuda(cudaMalloc(reinterpret_cast<void **>(&device_output), sizeof(output)));
    check_cuda(cudaMemcpy(device_input, input, sizeof(input), cudaMemcpyHostToDevice));
    const float alpha = 1.0f, beta = 0.0f;
    check_cudnn(cudnnActivationForward(handle, activation, &alpha, tensor,
                                      device_input, &beta, tensor, device_output));
    check_cuda(cudaMemcpy(output, device_output, sizeof(output), cudaMemcpyDeviceToHost));
    for (int i = 0; i < count; ++i) {
        const double expected = 1.0 / (1.0 + std::exp(-static_cast<double>(input[i])));
        if (!std::isfinite(output[i]) || std::abs(output[i] - expected) > 1e-6) {
            std::cerr << "FAIL: sigmoid at index " << i << '\n';
            return EXIT_FAILURE;
        }
    }
    check_cuda(cudaFree(device_input));
    check_cuda(cudaFree(device_output));
    check_cudnn(cudnnDestroyActivationDescriptor(activation));
    check_cudnn(cudnnDestroyTensorDescriptor(tensor));
    check_cudnn(cudnnDestroy(handle));
    std::cout << "PASS: all 10 sigmoid values match the CPU reference\n";
    return EXIT_SUCCESS;
}