// HIP adaptation of the pcie-link-capability-probe.
// Question: is the decode bottleneck the LINK, or the engine's read PATH?
// Measure H2D at the engine's actual expert slab size (2.344 MiB) and across
// a size sweep. If 2.344 MiB copies reach ~31 GB/s, the link is healthy and
// the engine is leaving bandwidth on the table. If they stall near 10 GB/s,
// the link is the wall.
//
// build: hipcc -O2 -o pcie_probe pcie_probe.hip   (scripts/pcie-probe.sh does this and guards the run)
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <hip/hip_runtime.h>

#define CK(x) do { hipError_t e = (x); if (e != hipSuccess) { \
    fprintf(stderr, "HIP %s (%s:%d)\n", hipGetErrorString(e), __FILE__, __LINE__); \
    return 1; } } while (0)

static double ms_between(hipEvent_t a, hipEvent_t b) {
    float ms = 0; hipEventElapsedTime(&ms, a, b); return (double)ms;
}

static int bench_size(size_t bytes) {
    const int iters = bytes >= (1UL<<30) ? 20
                    : bytes >= (256UL<<20) ? 30
                    : bytes >= (16UL<<20) ? 60 : 150;
    void *dev = NULL, *host = NULL, *pageable = NULL;
    hipStream_t s = NULL;
    CK(hipMalloc(&dev, bytes));
    CK(hipHostMalloc(&host, bytes, hipHostMallocDefault));   // PINNED, like the engine's slabs
    memset(host, 0x5a, bytes);
    CK(hipStreamCreate(&s));
    hipEvent_t a, b; CK(hipEventCreate(&a)); CK(hipEventCreate(&b));

    for (int i = 0; i < 3; i++) CK(hipMemcpy(dev, host, bytes, hipMemcpyHostToDevice));
    CK(hipDeviceSynchronize());

    // sync copy
    CK(hipEventRecord(a, 0));
    for (int i = 0; i < iters; i++) CK(hipMemcpy(dev, host, bytes, hipMemcpyHostToDevice));
    CK(hipEventRecord(b, 0)); CK(hipEventSynchronize(b));
    double h2d = (double)bytes * iters / (ms_between(a,b)/1000.0) / 1e9;

    // async copy on a stream (the engine's fill pipeline shape)
    CK(hipEventRecord(a, 0));
    for (int i = 0; i < iters; i++) CK(hipMemcpyAsync(dev, host, bytes, hipMemcpyHostToDevice, s));
    CK(hipEventRecord(b, 0)); CK(hipEventSynchronize(b));
    double h2da = (double)bytes * iters / (ms_between(a,b)/1000.0) / 1e9;

    // single-transfer latency at this size
    CK(hipMemcpy(dev, host, bytes, hipMemcpyHostToDevice));
    CK(hipDeviceSynchronize());
    CK(hipEventRecord(a, 0));
    CK(hipMemcpy(dev, host, bytes, hipMemcpyHostToDevice));
    CK(hipEventRecord(b, 0)); CK(hipEventSynchronize(b));
    double one_us = ms_between(a,b) * 1000.0;

    // pageable, to see whether pinning is load-bearing here
    pageable = malloc(bytes); memset(pageable, 0x5a, bytes);
    CK(hipMemcpy(dev, pageable, bytes, hipMemcpyHostToDevice));
    CK(hipEventRecord(a, 0));
    for (int i = 0; i < iters; i++) CK(hipMemcpy(dev, pageable, bytes, hipMemcpyHostToDevice));
    CK(hipEventRecord(b, 0)); CK(hipEventSynchronize(b));
    double pag = (double)bytes * iters / (ms_between(a,b)/1000.0) / 1e9;

    printf("  %8.3f MiB | %6.2f | %6.2f | %8.1f us | %6.2f\n",
           bytes/1048576.0, h2d, h2da, one_us, pag);

    free(pageable); hipHostFree(host); hipFree(dev);
    hipEventDestroy(a); hipEventDestroy(b); hipStreamDestroy(s);
    return 0;
}

int main(int argc, char **argv) {
    int dev_id = 0;
    if (argc > 1) dev_id = atoi(argv[1]);
    // No hipGetDeviceProperties here: it segfaults on this ROCm build while hipGetDeviceCount is
    // fine, and the link topology these numbers must be read against comes from sysfs, which
    // scripts/pcie-probe.sh prints. Only the device index is needed to pick what to measure.
    printf("device %d (PCI bus from sysfs, see the topology block above)\n", dev_id);
    printf("  size      | H2D   | async  | 1x xfer  | pageable   [GB/s]\n");
    printf("  %s\n", "-----------------------------------------------------------");
    size_t sz[] = {2459136UL,              /* 2.344 MiB: the engine's exact expert slab */
                   3UL<<20, 16UL<<20, 64UL<<20, 256UL<<20, 1UL<<30, 2UL<<30};
    for (int i = 0; i < 7; i++) if (bench_size(sz[i])) return 1;
    printf("\nGen5 x8 ceiling = 31.5 GB/s per direction\n");
    return 0;
}