// H2D bandwidth at the engine's expert-slab size.
// Deliberately avoids hipGetDeviceProperties: it segfaults on this ROCm build
// (hipGetDeviceCount works fine and reports 3 devices).
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <hip/hip_runtime.h>

static double one(size_t bytes, bool pinned, int iters) {
    void *dev = NULL, *host = NULL;
    if (hipMalloc(&dev, bytes) != hipSuccess) return -1;
    if (pinned) { if (hipHostMalloc(&host, bytes, hipHostMallocDefault) != hipSuccess) return -1; }
    else        { host = malloc(bytes); }
    memset(host, 0x5a, bytes);

    hipEvent_t a, b; hipEventCreate(&a); hipEventCreate(&b);
    for (int i = 0; i < 3; i++) hipMemcpy(dev, host, bytes, hipMemcpyHostToDevice);
    hipDeviceSynchronize();

    hipEventRecord(a, 0);
    for (int i = 0; i < iters; i++) hipMemcpy(dev, host, bytes, hipMemcpyHostToDevice);
    hipEventRecord(b, 0); hipEventSynchronize(b);
    float ms = 0; hipEventElapsedTime(&ms, a, b);
    double gb = (double)bytes * iters / (ms / 1000.0) / 1e9;

    if (!pinned) free(host); else hipHostFree(host);
    hipFree(dev); hipEventDestroy(a); hipEventDestroy(b);
    return gb;
}

int main(int argc, char **argv) {
    setvbuf(stdout, NULL, _IONBF, 0);
    int dev = (argc > 1) ? atoi(argv[1]) : 0;
    if (hipSetDevice(dev) != hipSuccess) { printf("setDevice(%d) failed\n", dev); return 1; }
    printf("device %d selected\n", dev);

    size_t sz[] = {2459136UL, 3UL<<20, 16UL<<20, 64UL<<20, 256UL<<20, 1UL<<30};
    printf("  %10s | %7s | %7s | %7s\n", "MiB", "pinned", "pagin.", "xfer us");
    printf("  %s\n", "-------------------------------------------");
    for (int i = 0; i < 6; i++) {
        size_t b = sz[i];
        int iters = (b >= (1UL<<30)) ? 10 : (b >= (256UL<<20)) ? 20 : 100;
        double p = one(b, true,  iters);
        double g = one(b, false, iters);
        // single-transfer latency
        void *dv, *hv;
        if (hipMalloc(&dv, b) == hipSuccess && hipHostMalloc(&hv, b, hipHostMallocDefault) == hipSuccess) {
            memset(hv, 0x5a, b);
            hipEvent_t a, c; hipEventCreate(&a); hipEventCreate(&c);
            hipMemcpy(dv, hv, b, hipMemcpyHostToDevice); hipDeviceSynchronize();
            hipEventRecord(a, 0);
            hipMemcpy(dv, hv, b, hipMemcpyHostToDevice);
            hipEventRecord(c, 0); hipEventSynchronize(c);
            float ms = 0; hipEventElapsedTime(&ms, a, c);
            printf("  %10.3f | %7.2f | %7.2f | %7.1f\n", b/1048576.0, p, g, ms*1000.0);
            hipHostFree(hv); hipFree(dv);
        } else {
            printf("  %10.3f | %7.2f | %7.2f |  (alloc failed)\n", b/1048576.0, p, g);
        }
    }
    printf("\n  2.344 MiB is the engine's exact expert slab size.\n");
    printf("  Gen5 x8 ceiling = 31.5 GB/s per direction.\n");
    return 0;
}