/* cgbn-batch-worker.cu — bend form B: batched 256-bit modular arithmetic
 *                       via NVlabs CGBN (Cooperative Groups Big Numbers).
 *
 * Plan: examples/cuda-fanout/plans/form-B-bignum-cgbn.md
 * Catalog: CATALOG.md form B (cuda-bignum-cgbn).
 *
 * Op coverage: 0x01 mod-add, 0x02 mod-sub, 0x03 mod-mul, 0x04 mod-sqr,
 *              0x05 mod-inv, 0x06 mod-exp, 0x07 mod-reduce,
 *              0x08 add-no-mod, 0x09 mul-no-mod,
 *              0x0A mod-inv-batch (Montgomery's trick: 1 inv + 3(N-1) muls).
 * Bitwidth: 256 only (day-1 freeze).
 *
 * Wire (BCGB input):
 *   "BCGB" (4 B magic)
 *   u32 op_id
 *   u32 bitwidth         (256 for now)
 *   u32 n_instances
 *   u8[bitwidth/8] modulus            (OMITTED for 0x08, 0x09)
 *   u8[n_instances * bitwidth/8] operand_a
 *   u8[n_instances * bitwidth/8] operand_b   (OMITTED for unary 0x04, 0x05, 0x07)
 *
 * Wire (BCGR output):
 *   "BCGR"               (4 B magic)
 *   u32 status           (0 = ok)
 *   u32 n_instances
 *   u8[n_instances * width_out/8] result
 *     width_out = bitwidth          for 0x01..0x08, 0x0A
 *     width_out = 2 * bitwidth      for 0x09 mul-no-mod
 *
 * Op 0x0A wire shape (mod-inv-batch — same layout as 0x05 mod-inv):
 *   16 B header + 32 B modulus + n*32 B xs  →  12 B header + n*32 B inverses
 *   Algorithm: Montgomery's batched-inversion trick
 *     phase 1: t[0] = x[0]; t[i] = t[i-1] * x[i] mod m       (N-1 muls)
 *     phase 2: u = t[n-1]^-1 mod m                           (1 inv)
 *     phase 3: for i in n-1..1:                              (2(N-1) muls)
 *                result[i] = u * t[i-1] mod m
 *                u = u * x[i] mod m
 *              result[0] = u
 *   Cost: 1 modular inverse + 3(N-1) modular multiplications.
 *
 * Build:
 *   make cgbn-batch-worker CGBN_INC=/path/to/CGBN/include
 *
 * License: AGPLv3. Wraps NVlabs CGBN headers (BSD-3-Clause); CGBN
 * source is not redistributed by this binary, only linked in by header.
 */

#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <time.h>

#include <gmp.h>
#include <cgbn/cgbn.h>

#define CUDA_CHECK(stmt) do { \
    cudaError_t err = (stmt); \
    if (err != cudaSuccess) { \
        fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \
                #stmt, __FILE__, __LINE__, cudaGetErrorString(err)); \
        return -1; \
    } \
} while (0)

/* CGBN environment for 256-bit instances, TPI=8 (warp owns 4 instances).
 * Per plan §1: this gives 128 in-flight instances per SM × 82 SMs on
 * a 3090 = 10,496 concurrent 256-bit ops per kernel wave.
 */
#define TPI 8
#define BITS_256 256

typedef cgbn_context_t<TPI>            context_t;
typedef cgbn_env_t<context_t, BITS_256> env_256_t;
typedef cgbn_mem_t<BITS_256>            mem_256_t;
typedef cgbn_mem_t<2 * BITS_256>        mem_512_t;

/* Op codes — must match plan §2 dispatcher table. */
enum {
    OP_MOD_ADD     = 0x01,
    OP_MOD_SUB     = 0x02,
    OP_MOD_MUL     = 0x03,
    OP_MOD_SQR     = 0x04,
    OP_MOD_INV     = 0x05,
    OP_MOD_EXP     = 0x06,
    OP_MOD_REDUCE  = 0x07,
    OP_ADD_NO_MOD  = 0x08,
    OP_MUL_NO_MOD  = 0x09,
    OP_MOD_INV_BATCH = 0x0A,
};

/* ── Kernels: one __global__ per op, 256-bit, batched ──────────────── */

__global__ void mod_add_256_kernel(
    mem_256_t *result,
    mem_256_t *a,
    mem_256_t *b,
    mem_256_t *modulus,
    uint32_t n_instances)
{
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t instance = tid / TPI;
    if (instance >= n_instances) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_a, bn_b, bn_m, bn_r;

    cgbn_load(env, bn_a, &a[instance]);
    cgbn_load(env, bn_b, &b[instance]);
    cgbn_load(env, bn_m, modulus);

    /* r = a + b; if carry-out OR r >= m, r -= m. With operands in [0, m)
     * and m at most 2^256-1, sum fits in 257 bits. cgbn_add returns the
     * 257th carry bit; we subtract m when carry==1 OR r >= m.            */
    int32_t carry = cgbn_add(env, bn_r, bn_a, bn_b);
    if (carry != 0 || cgbn_compare(env, bn_r, bn_m) >= 0) {
        cgbn_sub(env, bn_r, bn_r, bn_m);
    }
    cgbn_store(env, &result[instance], bn_r);
}

__global__ void mod_sub_256_kernel(
    mem_256_t *result,
    mem_256_t *a,
    mem_256_t *b,
    mem_256_t *modulus,
    uint32_t n_instances)
{
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t instance = tid / TPI;
    if (instance >= n_instances) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_a, bn_b, bn_m, bn_r;

    cgbn_load(env, bn_a, &a[instance]);
    cgbn_load(env, bn_b, &b[instance]);
    cgbn_load(env, bn_m, modulus);

    /* r = a - b; if borrow (a < b), r += m. cgbn_sub returns nonzero
     * (typically -1) when underflow occurs.                              */
    int32_t borrow = cgbn_sub(env, bn_r, bn_a, bn_b);
    if (borrow != 0) {
        cgbn_add(env, bn_r, bn_r, bn_m);
    }
    cgbn_store(env, &result[instance], bn_r);
}

__global__ void mod_mul_256_kernel(
    mem_256_t *result,
    mem_256_t *a,
    mem_256_t *b,
    mem_256_t *modulus,
    uint32_t n_instances)
{
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t instance = tid / TPI;
    if (instance >= n_instances) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_a, bn_b, bn_m, bn_r;

    cgbn_load(env, bn_a, &a[instance]);
    cgbn_load(env, bn_b, &b[instance]);
    cgbn_load(env, bn_m, modulus);

    /* a * b mod m via wide-mul + reduce. */
    typename env_256_t::cgbn_wide_t bn_wide;
    cgbn_mul_wide(env, bn_wide, bn_a, bn_b);
    cgbn_rem_wide(env, bn_r, bn_wide, bn_m);

    cgbn_store(env, &result[instance], bn_r);
}

__global__ void mod_sqr_256_kernel(
    mem_256_t *result,
    mem_256_t *a,
    mem_256_t *modulus,
    uint32_t n_instances)
{
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t instance = tid / TPI;
    if (instance >= n_instances) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_a, bn_m, bn_r;

    cgbn_load(env, bn_a, &a[instance]);
    cgbn_load(env, bn_m, modulus);

    /* a^2 mod m via sqr_wide + rem_wide (avoids the second cgbn_t load). */
    typename env_256_t::cgbn_wide_t bn_wide;
    cgbn_sqr_wide(env, bn_wide, bn_a);
    cgbn_rem_wide(env, bn_r, bn_wide, bn_m);

    cgbn_store(env, &result[instance], bn_r);
}

__global__ void mod_inv_256_kernel(
    mem_256_t *result,
    mem_256_t *a,
    mem_256_t *modulus,
    uint32_t n_instances)
{
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t instance = tid / TPI;
    if (instance >= n_instances) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_a, bn_m, bn_r;

    cgbn_load(env, bn_a, &a[instance]);
    cgbn_load(env, bn_m, modulus);

    /* cgbn_modular_inverse returns true on success; if gcd(a,m) != 1 we
     * write zero — caller must avoid feeding non-invertible inputs.     */
    bool ok = cgbn_modular_inverse(env, bn_r, bn_a, bn_m);
    if (!ok) {
        cgbn_set_ui32(env, bn_r, 0);
    }
    cgbn_store(env, &result[instance], bn_r);
}

/* mod-inv-batch (op 0x0A): Montgomery's batched-inversion trick.
 *
 * Day-1 sequential shape: launch with TPI=8 threads (1 CGBN instance), all
 * 8 threads cooperate as a single CGBN cgbn_t inside cooperative-group ops.
 * The kernel itself iterates the n inputs in a sequential loop.
 *
 * Win comes not from intra-phase parallelism (Phase 1/3 are sequential) but
 * from replacing N independent modular inverses (each ~hundreds of binary
 * GCD steps under cgbn_modular_inverse) with N-1 simple modular multiplies
 * plus 1 inverse. At n=100k a single thread doing 100k modmuls at CGBN's
 * kernel rate (~ns each) still finishes well under the time N parallel
 * mod-invs take at 42 Mops/s (~2.4 ms).
 *
 * Buffers:
 *   t[]      scratch prefix-products workspace (n * 32 B)
 *   result[] final inverses
 */
__global__ void mod_inv_batch_256_kernel(
    mem_256_t *result,
    mem_256_t *t,
    mem_256_t *xs,
    mem_256_t *modulus,
    uint32_t n_instances)
{
    /* Single cooperating CGBN instance — TPI threads in lockstep. */
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= TPI) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_x, bn_m, bn_acc, bn_u, bn_prev_t;
    typename env_256_t::cgbn_wide_t bn_wide;

    cgbn_load(env, bn_m, modulus);

    /* Phase 1: prefix products.
     *   t[0]   = x[0]
     *   t[i]   = t[i-1] * x[i] mod m,  for i = 1 .. n-1
     */
    cgbn_load(env, bn_acc, &xs[0]);
    cgbn_store(env, &t[0], bn_acc);
    for (uint32_t i = 1; i < n_instances; i++) {
        cgbn_load(env, bn_x, &xs[i]);
        cgbn_mul_wide(env, bn_wide, bn_acc, bn_x);
        cgbn_rem_wide(env, bn_acc, bn_wide, bn_m);
        cgbn_store(env, &t[i], bn_acc);
    }

    /* Phase 2: single modular inverse on t[n-1]. */
    bool ok = cgbn_modular_inverse(env, bn_u, bn_acc, bn_m);
    if (!ok) {
        cgbn_set_ui32(env, bn_u, 0);
    }

    /* Phase 3: reverse sweep.
     *   for i = n-1 .. 1:
     *     result[i] = u * t[i-1] mod m
     *     u         = u * x[i]   mod m
     *   result[0] = u
     */
    for (int64_t i = (int64_t)n_instances - 1; i >= 1; i--) {
        cgbn_load(env, bn_prev_t, &t[i - 1]);
        cgbn_mul_wide(env, bn_wide, bn_u, bn_prev_t);
        cgbn_rem_wide(env, bn_acc, bn_wide, bn_m);
        cgbn_store(env, &result[i], bn_acc);

        cgbn_load(env, bn_x, &xs[i]);
        cgbn_mul_wide(env, bn_wide, bn_u, bn_x);
        cgbn_rem_wide(env, bn_u, bn_wide, bn_m);
    }
    cgbn_store(env, &result[0], bn_u);
}

__global__ void mod_exp_256_kernel(
    mem_256_t *result,
    mem_256_t *a,
    mem_256_t *b,
    mem_256_t *modulus,
    uint32_t n_instances)
{
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t instance = tid / TPI;
    if (instance >= n_instances) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_a, bn_b, bn_m, bn_r;

    cgbn_load(env, bn_a, &a[instance]);
    cgbn_load(env, bn_b, &b[instance]);
    cgbn_load(env, bn_m, modulus);

    /* a^b mod m via binary ladder (CGBN's built-in). */
    cgbn_modular_power(env, bn_r, bn_a, bn_b, bn_m);

    cgbn_store(env, &result[instance], bn_r);
}

__global__ void mod_reduce_256_kernel(
    mem_256_t *result,
    mem_256_t *a,
    mem_256_t *modulus,
    uint32_t n_instances)
{
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t instance = tid / TPI;
    if (instance >= n_instances) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_a, bn_m, bn_r;

    cgbn_load(env, bn_a, &a[instance]);
    cgbn_load(env, bn_m, modulus);

    /* a is already <= 256 bits, so cgbn_rem suffices (no wide path). */
    cgbn_rem(env, bn_r, bn_a, bn_m);

    cgbn_store(env, &result[instance], bn_r);
}

__global__ void add_no_mod_256_kernel(
    mem_256_t *result,
    mem_256_t *a,
    mem_256_t *b,
    uint32_t n_instances)
{
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t instance = tid / TPI;
    if (instance >= n_instances) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_a, bn_b, bn_r;

    cgbn_load(env, bn_a, &a[instance]);
    cgbn_load(env, bn_b, &b[instance]);

    /* Straight 256-bit add, truncated (carry-out discarded). */
    cgbn_add(env, bn_r, bn_a, bn_b);

    cgbn_store(env, &result[instance], bn_r);
}

__global__ void mul_no_mod_256_kernel(
    mem_512_t *result,
    mem_256_t *a,
    mem_256_t *b,
    uint32_t n_instances)
{
    uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t instance = tid / TPI;
    if (instance >= n_instances) return;

    context_t  ctx(cgbn_no_checks);
    env_256_t  env(ctx);
    env_256_t::cgbn_t bn_a, bn_b;

    cgbn_load(env, bn_a, &a[instance]);
    cgbn_load(env, bn_b, &b[instance]);

    /* Full 512-bit product. Store as low half then high half so the byte
     * stream is little-endian across the entire 64-byte product.        */
    typename env_256_t::cgbn_wide_t bn_wide;
    cgbn_mul_wide(env, bn_wide, bn_a, bn_b);

    mem_256_t *out_pair = (mem_256_t *)&result[instance];
    cgbn_store(env, &out_pair[0], bn_wide._low);
    cgbn_store(env, &out_pair[1], bn_wide._high);
}

/* ── Host: process_one_bin ─────────────────────────────────────────── */

static int read_file_to_buffer(const char *path, uint8_t **out_buf, size_t *out_len) {
    FILE *f = fopen(path, "rb");
    if (!f) { fprintf(stderr, "open %s failed\n", path); return -1; }
    fseek(f, 0, SEEK_END);
    long sz = ftell(f);
    fseek(f, 0, SEEK_SET);
    uint8_t *buf = (uint8_t*)malloc(sz);
    if (!buf) { fclose(f); return -1; }
    if (fread(buf, 1, sz, f) != (size_t)sz) {
        free(buf); fclose(f); return -1;
    }
    fclose(f);
    *out_buf = buf;
    *out_len = sz;
    return 0;
}

static int write_buffer_to_file(const char *path, const uint8_t *buf, size_t len) {
    FILE *f = fopen(path, "wb");
    if (!f) { fprintf(stderr, "create %s failed\n", path); return -1; }
    if (fwrite(buf, 1, len, f) != len) { fclose(f); return -1; }
    fclose(f);
    return 0;
}

static uint32_t rd32_le(const uint8_t *p) {
    return (uint32_t)p[0] | ((uint32_t)p[1] << 8)
         | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

static void wr32_le(uint8_t *p, uint32_t v) {
    p[0] = v & 0xff; p[1] = (v >> 8) & 0xff;
    p[2] = (v >> 16) & 0xff; p[3] = (v >> 24) & 0xff;
}

/* Op-shape classifier. Three families:
 *   FAMILY_BIN_MOD   binary with modulus  : 0x01, 0x02, 0x03, 0x06
 *                    input  = 16 + bpi + 2*n*bpi
 *                    output = 12 + n*bpi
 *   FAMILY_UN_MOD    unary  with modulus  : 0x04, 0x05, 0x07
 *                    input  = 16 + bpi + n*bpi
 *                    output = 12 + n*bpi
 *   FAMILY_NO_MOD    binary no modulus    : 0x08, 0x09
 *                    input  = 16 + 2*n*bpi
 *                    output = 12 + n*bpi (0x08) or 12 + 2*n*bpi (0x09)
 */
enum op_family {
    FAMILY_BIN_MOD,
    FAMILY_UN_MOD,
    FAMILY_NO_MOD,
    FAMILY_UNKNOWN,
};

static enum op_family classify_op(uint32_t op_id) {
    switch (op_id) {
        case OP_MOD_ADD:
        case OP_MOD_SUB:
        case OP_MOD_MUL:
        case OP_MOD_EXP:
            return FAMILY_BIN_MOD;
        case OP_MOD_SQR:
        case OP_MOD_INV:
        case OP_MOD_REDUCE:
        case OP_MOD_INV_BATCH:
            return FAMILY_UN_MOD;
        case OP_ADD_NO_MOD:
        case OP_MUL_NO_MOD:
            return FAMILY_NO_MOD;
        default:
            return FAMILY_UNKNOWN;
    }
}

static int process_one_bin(const char *in_path, const char *out_path) {
    uint8_t *in_buf = NULL;
    size_t in_len = 0;
    if (read_file_to_buffer(in_path, &in_buf, &in_len) != 0) return -1;

    if (in_len < 16 || memcmp(in_buf, "BCGB", 4) != 0) {
        fprintf(stderr, "bad BCGB magic / short input (%zu bytes)\n", in_len);
        free(in_buf); return -1;
    }
    uint32_t op_id      = rd32_le(in_buf + 4);
    uint32_t bitwidth   = rd32_le(in_buf + 8);
    uint32_t n          = rd32_le(in_buf + 12);

    if (bitwidth != BITS_256) {
        fprintf(stderr, "unsupported bitwidth %u (worker ships 256 only)\n",
                bitwidth);
        free(in_buf); return -1;
    }

    enum op_family fam = classify_op(op_id);
    if (fam == FAMILY_UNKNOWN) {
        fprintf(stderr, "unsupported op_id 0x%02x\n", op_id);
        free(in_buf); return -1;
    }

    size_t bpi = bitwidth / 8;
    size_t expected = 0;
    switch (fam) {
        case FAMILY_BIN_MOD: expected = 16 + bpi + 2 * (size_t)n * bpi; break;
        case FAMILY_UN_MOD:  expected = 16 + bpi +     (size_t)n * bpi; break;
        case FAMILY_NO_MOD:  expected = 16 +           2 * (size_t)n * bpi; break;
        default: break;
    }
    if (in_len != expected) {
        fprintf(stderr, "BCGB payload size %zu != expected %zu (op=0x%02x n=%u)\n",
                in_len, expected, op_id, n);
        free(in_buf); return -1;
    }

    /* Per-family pointer carving. */
    const uint8_t *modulus_bytes = NULL;
    const uint8_t *a_bytes       = NULL;
    const uint8_t *b_bytes       = NULL;
    switch (fam) {
        case FAMILY_BIN_MOD:
            modulus_bytes = in_buf + 16;
            a_bytes       = modulus_bytes + bpi;
            b_bytes       = a_bytes + (size_t)n * bpi;
            break;
        case FAMILY_UN_MOD:
            modulus_bytes = in_buf + 16;
            a_bytes       = modulus_bytes + bpi;
            break;
        case FAMILY_NO_MOD:
            a_bytes       = in_buf + 16;
            b_bytes       = a_bytes + (size_t)n * bpi;
            break;
        default: break;
    }

    /* Allocate device buffers. */
    mem_256_t *d_modulus = NULL, *d_a = NULL, *d_b = NULL, *d_r = NULL;
    mem_256_t *d_t = NULL;       /* scratch for OP_MOD_INV_BATCH phase 1. */
    mem_512_t *d_r_wide = NULL;

    if (modulus_bytes) {
        CUDA_CHECK(cudaMalloc(&d_modulus, sizeof(mem_256_t)));
        CUDA_CHECK(cudaMemcpy(d_modulus, modulus_bytes, bpi,
                              cudaMemcpyHostToDevice));
    }
    CUDA_CHECK(cudaMalloc(&d_a, (size_t)n * sizeof(mem_256_t)));
    CUDA_CHECK(cudaMemcpy(d_a, a_bytes, (size_t)n * bpi,
                          cudaMemcpyHostToDevice));
    if (b_bytes) {
        CUDA_CHECK(cudaMalloc(&d_b, (size_t)n * sizeof(mem_256_t)));
        CUDA_CHECK(cudaMemcpy(d_b, b_bytes, (size_t)n * bpi,
                              cudaMemcpyHostToDevice));
    }
    if (op_id == OP_MUL_NO_MOD) {
        CUDA_CHECK(cudaMalloc(&d_r_wide, (size_t)n * sizeof(mem_512_t)));
    } else {
        CUDA_CHECK(cudaMalloc(&d_r, (size_t)n * sizeof(mem_256_t)));
    }
    if (op_id == OP_MOD_INV_BATCH) {
        CUDA_CHECK(cudaMalloc(&d_t, (size_t)n * sizeof(mem_256_t)));
    }

    /* Launch geometry: TPI threads per instance, 128 instances per block. */
    uint32_t threads_per_block = 128 * TPI;
    uint32_t total_threads = n * TPI;
    uint32_t blocks = (total_threads + threads_per_block - 1) / threads_per_block;

    /* OP_MOD_INV_BATCH runs as a single cooperating CGBN instance
     * (TPI threads, sequential phases inside the kernel).               */
    uint32_t batch_inv_threads = TPI;
    uint32_t batch_inv_blocks  = 1;

    struct timespec t0, t1;
    clock_gettime(CLOCK_MONOTONIC, &t0);
    switch (op_id) {
        case OP_MOD_ADD:
            mod_add_256_kernel<<<blocks, threads_per_block>>>(d_r, d_a, d_b, d_modulus, n);
            break;
        case OP_MOD_SUB:
            mod_sub_256_kernel<<<blocks, threads_per_block>>>(d_r, d_a, d_b, d_modulus, n);
            break;
        case OP_MOD_MUL:
            mod_mul_256_kernel<<<blocks, threads_per_block>>>(d_r, d_a, d_b, d_modulus, n);
            break;
        case OP_MOD_SQR:
            mod_sqr_256_kernel<<<blocks, threads_per_block>>>(d_r, d_a, d_modulus, n);
            break;
        case OP_MOD_INV:
            mod_inv_256_kernel<<<blocks, threads_per_block>>>(d_r, d_a, d_modulus, n);
            break;
        case OP_MOD_EXP:
            mod_exp_256_kernel<<<blocks, threads_per_block>>>(d_r, d_a, d_b, d_modulus, n);
            break;
        case OP_MOD_REDUCE:
            mod_reduce_256_kernel<<<blocks, threads_per_block>>>(d_r, d_a, d_modulus, n);
            break;
        case OP_ADD_NO_MOD:
            add_no_mod_256_kernel<<<blocks, threads_per_block>>>(d_r, d_a, d_b, n);
            break;
        case OP_MUL_NO_MOD:
            mul_no_mod_256_kernel<<<blocks, threads_per_block>>>(d_r_wide, d_a, d_b, n);
            break;
        case OP_MOD_INV_BATCH:
            mod_inv_batch_256_kernel<<<batch_inv_blocks, batch_inv_threads>>>(
                d_r, d_t, d_a, d_modulus, n);
            break;
    }
    CUDA_CHECK(cudaDeviceSynchronize());
    clock_gettime(CLOCK_MONOTONIC, &t1);
    double kernel_ms = (t1.tv_sec - t0.tv_sec) * 1e3
                     + (t1.tv_nsec - t0.tv_nsec) / 1e6;

    /* Output sizing. */
    size_t out_bytes_per_inst = (op_id == OP_MUL_NO_MOD) ? 2 * bpi : bpi;
    size_t out_size = 12 + (size_t)n * out_bytes_per_inst;
    uint8_t *out_buf = (uint8_t*)malloc(out_size);
    if (!out_buf) {
        if (d_modulus) cudaFree(d_modulus);
        if (d_a) cudaFree(d_a);
        if (d_b) cudaFree(d_b);
        if (d_r) cudaFree(d_r);
        if (d_t) cudaFree(d_t);
        if (d_r_wide) cudaFree(d_r_wide);
        free(in_buf); return -1;
    }
    memcpy(out_buf, "BCGR", 4);
    wr32_le(out_buf + 4, 0);   /* status = ok */
    wr32_le(out_buf + 8, n);
    if (op_id == OP_MUL_NO_MOD) {
        CUDA_CHECK(cudaMemcpy(out_buf + 12, d_r_wide, (size_t)n * out_bytes_per_inst,
                              cudaMemcpyDeviceToHost));
    } else {
        CUDA_CHECK(cudaMemcpy(out_buf + 12, d_r, (size_t)n * out_bytes_per_inst,
                              cudaMemcpyDeviceToHost));
    }

    if (d_modulus) cudaFree(d_modulus);
    if (d_a) cudaFree(d_a);
    if (d_b) cudaFree(d_b);
    if (d_r) cudaFree(d_r);
    if (d_t) cudaFree(d_t);
    if (d_r_wide) cudaFree(d_r_wide);
    free(in_buf);

    int wr = write_buffer_to_file(out_path, out_buf, out_size);
    free(out_buf);
    fprintf(stderr, "cgbn op=0x%02x bitwidth=%u n=%u kernel=%.2fms\n",
            op_id, bitwidth, n, kernel_ms);
    return wr;
}

/* ── main ─────────────────────────────────────────────────────────── */

int main(int argc, char **argv) {
    if (argc >= 2 && !strcmp(argv[1], "--daemon")) {
        int dev_count = 0;
        cudaGetDeviceCount(&dev_count);
        if (dev_count == 0) {
            fprintf(stderr, "no CUDA devices\n");
            return 1;
        }
        cudaFree(0);
        fprintf(stdout, "ready\n"); fflush(stdout);

        char line[4096];
        while (fgets(line, sizeof(line), stdin)) {
            char *nl = strchr(line, '\n');
            if (nl) *nl = 0;
            if (!strcmp(line, "quit")) {
                fprintf(stdout, "bye\n"); fflush(stdout);
                break;
            }
            if (!strncmp(line, "process-bin ", 12)) {
                char *args = line + 12;
                char *sep = strchr(args, ' ');
                if (!sep) {
                    fprintf(stdout, "error bad-args\n"); fflush(stdout);
                    continue;
                }
                *sep = 0;
                char *in_path  = args;
                char *out_path = sep + 1;
                int r = process_one_bin(in_path, out_path);
                fprintf(stdout, r == 0 ? "done %s\n" : "error process-failed\n",
                        out_path);
                fflush(stdout);
            } else if (line[0]) {
                fprintf(stdout, "error unknown-command\n"); fflush(stdout);
            }
        }
        return 0;
    }

    if (argc >= 4 && !strcmp(argv[1], "--binary")) {
        return process_one_bin(argv[2], argv[3]);
    }

    fprintf(stderr,
        "usage: %s --binary <input.bin> <output.bin>\n"
        "       %s --daemon  (read commands on stdin)\n",
        argv[0], argv[0]);
    return 1;
}
