#!/bin/bash
export HSA_ENABLE_SDMA=0

cat > /tmp/run_test2.hip << 'HIPEOF'
#include <hip/hip_runtime.h>
#include <cstdio>

__global__ void add_kernel(float *a, float *b, float *c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) c[i] = a[i] + b[i];
}

int main() {
    hipDeviceProp_t prop;
    hipGetDeviceProperties(&prop, 0);
    printf("Device: %s\n", prop.name);
    printf("GCN Arch: %s\n", prop.gcnArchName);

    const int N = 256;
    float h_a[N], h_b[N], h_c[N];
    for (int i = 0; i < N; i++) { h_a[i] = (float)i; h_b[i] = (float)(i * 2); }

    float *d_a, *d_b, *d_c;
    hipMallocManaged(&d_a, N * sizeof(float));
    hipMallocManaged(&d_b, N * sizeof(float));
    hipMallocManaged(&d_c, N * sizeof(float));
    hipMemcpy(d_a, h_a, N * sizeof(float), hipMemcpyHostToDevice);
    hipMemcpy(d_b, h_b, N * sizeof(float), hipMemcpyHostToDevice);

    add_kernel<<<1, N>>>(d_a, d_b, d_c, N);
    hipDeviceSynchronize();

    hipError_t err = hipGetLastError();
    if (err != hipSuccess) {
        printf("KERNEL FAIL: %s\n", hipGetErrorString(err));
        hipFree(d_a); hipFree(d_b); hipFree(d_c);
        return 1;
    }

    hipMemcpy(h_c, d_c, N * sizeof(float), hipMemcpyDeviceToHost);

    int ok = 1;
    for (int i = 0; i < N; i++) {
        if (h_c[i] != h_a[i] + h_b[i]) { ok = 0; printf("Mismatch at %d: %f vs %f\n", i, h_c[i], h_a[i]+h_b[i]); break; }
    }
    printf("Compute: %s\n", ok ? "PASS" : "FAIL");

    hipFree(d_a); hipFree(d_b); hipFree(d_c);
    return ok ? 0 : 1;
}
HIPEOF

for arch in gfx1013 gfx1010 gfx10-1-generic; do
    echo "=== Runtime test: $arch ==="
    /opt/rocm/bin/hipcc --offload-arch=$arch /tmp/run_test2.hip -o /tmp/run_test2_${arch} 2>&1
    if [ $? -eq 0 ]; then
        echo "Compiled OK. Running..."
        timeout 15 /tmp/run_test2_${arch} 2>&1
        echo "Exit: $?"
    else
        echo "Compile FAILED"
    fi
    echo ""
done