/**
 * HIP Vector Addition Test — AMD BC-250 ROCm Validation
 * 
 * Uses hipMallocManaged (unified memory) — REQUIRED for BC-250 which is an
 * APU-like device with shared system memory (no dedicated VRAM).
 * Standard hipMalloc + hipMemcpy will crash the system!
 *
 * Verifies GPU compute works end-to-end:
 * 1. HIP runtime initializes & device query
 * 2. Managed memory allocation (unified address space)
 * 3. GPU kernel execution
 * 4. Results are numerically correct
 *
 * Compile: hipcc --offload-arch=gfx1010 -o hip_vector_add hip_vector_add.cpp
 * Run:     HSA_OVERRIDE_GFX_VERSION=10.1.0 ./hip_vector_add
 */
#include <hip/hip_runtime.h>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <unistd.h>  // _exit() — bypasses C++ destructors

#define HIP_CHECK(call) do { \
    hipError_t err = call; \
    if (err != hipSuccess) { \
        fprintf(stderr, "HIP Error: %s at %s:%d\n", \
                hipGetErrorString(err), __FILE__, __LINE__); \
        fflush(stderr); \
        exit(1); \
    } \
} while(0)

__global__ void vectorAdd(const float* A, const float* B, float* C, int N) {
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < N) {
        C[i] = A[i] + B[i];
    }
}

int main() {
    printf("=== HIP Vector Addition Test (Managed Memory) ===\n");
    printf("    Target: AMD BC-250 (gfx1013, APU shared memory)\n\n");
    fflush(stdout);

    // Step 1: Query device
    printf("[1/4] Querying HIP device... ");
    fflush(stdout);
    int deviceCount = 0;
    HIP_CHECK(hipGetDeviceCount(&deviceCount));

    if (deviceCount == 0) {
        printf("FAIL: No HIP devices found!\n");
        return 1;
    }

    hipDeviceProp_t props;
    HIP_CHECK(hipGetDeviceProperties(&props, 0));
    printf("OK\n");
    printf("       Device:          %s\n", props.name);
    printf("       GCN Arch:        %s\n", props.gcnArchName);
    printf("       Compute Units:   %d\n", props.multiProcessorCount);
    printf("       Total Memory:    %zu MB (shared system RAM)\n", props.totalGlobalMem / (1024*1024));
    printf("       Integrated:      %s\n", props.integrated ? "YES (APU)" : "NO");
    printf("       Managed Memory:  %s\n", props.managedMemory ? "YES" : "NO");
    fflush(stdout);

    // Step 2: Allocate MANAGED memory (unified — safe for APU/shared memory)
    const int N = 1 << 16; // 65536 elements
    size_t bytes = N * sizeof(float);

    printf("[2/4] Allocating managed memory (%.1f KB x3)... ", bytes/1024.0);
    fflush(stdout);

    float *A = nullptr, *B = nullptr, *C = nullptr;
    HIP_CHECK(hipMallocManaged(&A, bytes));
    HIP_CHECK(hipMallocManaged(&B, bytes));
    HIP_CHECK(hipMallocManaged(&C, bytes));
    printf("OK\n");
    fflush(stdout);

    // Initialize on host (managed memory is accessible from both CPU and GPU)
    for (int i = 0; i < N; i++) {
        A[i] = sinf(i) * sinf(i);
        B[i] = cosf(i) * cosf(i);
        C[i] = 0.0f;
    }

    // Step 3: Launch kernel
    int threadsPerBlock = 256;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;

    printf("[3/4] Launching kernel (%d blocks x %d threads)... ", blocksPerGrid, threadsPerBlock);
    fflush(stdout);

    hipEvent_t start, stop;
    HIP_CHECK(hipEventCreate(&start));
    HIP_CHECK(hipEventCreate(&stop));

    HIP_CHECK(hipEventRecord(start));
    hipLaunchKernelGGL(vectorAdd, dim3(blocksPerGrid), dim3(threadsPerBlock), 0, 0, A, B, C, N);
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipEventRecord(stop));
    HIP_CHECK(hipEventSynchronize(stop));

    float ms = 0;
    HIP_CHECK(hipEventElapsedTime(&ms, start, stop));
    printf("OK (%.3f ms)\n", ms);
    fflush(stdout);

    // Step 4: Verify on host (managed memory — no memcpy needed!)
    printf("[4/4] Verifying results... ");
    fflush(stdout);

    // Ensure GPU is done
    HIP_CHECK(hipDeviceSynchronize());

    int errors = 0;
    for (int i = 0; i < N; i++) {
        float expected = A[i] + B[i]; // sin²(x) + cos²(x) = 1.0
        if (fabsf(C[i] - expected) > 1e-5) {
            if (errors < 5) {
                fprintf(stderr, "\n  Mismatch at [%d]: got %f, expected %f", 
                        i, C[i], expected);
            }
            errors++;
        }
    }

    if (errors == 0) {
        printf("PASSED — all %d elements correct (sin²+cos²=1.0)\n", N);
    } else {
        printf("FAILED — %d/%d mismatches\n", errors, N);
    }
    fflush(stdout);

    // CRITICAL BC-250 WORKAROUND:
    // The HIP runtime's static destructors trigger KIQ queue teardown on the
    // BC-250's shared-memory architecture, causing "timeout waiting for kiq fence"
    // and a full system hang (no GPU reset possible on shared RAM APU).
    //
    // We MUST use _exit() to terminate immediately, bypassing:
    //   - C++ static destructors (HIP runtime cleanup)
    //   - atexit() handlers
    //   - HIP's internal queue teardown via KIQ
    //
    // Memory is reclaimed by the OS. The GPU queues are released by KFD
    // when the process file descriptors are closed, which is safer than
    // the HIP runtime's explicit teardown path.

    printf("\n=== ROCm HIP Compute: %s ===\n", errors == 0 ? "FULLY OPERATIONAL" : "FAILED");
    fflush(stdout);
    fflush(stderr);
    _exit(errors > 0 ? 1 : 0);  // HARD EXIT — bypasses HIP destructors
}
