/**
 * gpu.mjs — the model behind the visualization.
 *
 * Two kinds of numbers live in this file, and they are kept strictly apart:
 *
 *   1. SPECIFICATIONS (`GPUS`) — copied from the vendor's own datasheet /
 *      architecture whitepaper. Each entry carries a `source`. Nothing here is
 *      measured, estimated, or rounded to taste.
 *
 *   2. DERIVED QUANTITIES — everything else. Computed from (1) by the functions
 *      below, which are covered by `gpu.test.mjs`. `derivedFp32Tflops()` exists
 *      specifically so the test suite can check the spec table against itself:
 *      SMs x cores x 2 x clock must reproduce the published TFLOPS.
 *
 * Where a quantity is genuinely an estimate (how many FLOPs a GELU costs), it is
 * a named exported constant with the assumption written next to it, so it can be
 * seen and changed rather than buried inside a formula.
 *
 * Runs unmodified in Node and in the browser. No dependencies.
 */

export const WARP_SIZE = 32;

/** Bytes per element, by tensor dtype. */
export const DTYPE_BYTES = { fp32: 4, tf32: 4, bf16: 2, fp16: 2, fp8: 1, int8: 1 };

/**
 * Allocation granularity on compute capability 7.0+ (Volta, Turing, Ampere, Ada,
 * Hopper): registers are handed to a warp in units of 256, shared memory to a
 * block in units of 128 bytes.
 */
export const REG_ALLOC_UNIT_PER_WARP = 256;
export const SMEM_ALLOC_UNIT_BYTES = 128;

/* ------------------------------------------------------------------ *
 * 1. Published specifications
 * ------------------------------------------------------------------ */

export const GPUS = {
  a100: {
    id: 'a100',
    name: 'A100 SXM4 80GB',
    arch: 'Ampere (GA100)',
    cc: '8.0',
    source: 'NVIDIA A100 Tensor Core GPU datasheet + NVIDIA Ampere GA100 whitepaper',
    sms: 108,
    fp32CoresPerSm: 64,
    tensorCoresPerSm: 4,
    tensorCoreGen: '3rd gen',
    clockBoostMhz: 1410,
    // Peak throughput, dense (no structural sparsity), in TFLOPS.
    fp32Tflops: 19.5,
    tf32TensorTflops: 156,
    fp16TensorTflops: 312,
    tensorAccumulateNote: 'FP16/BF16 multiply, FP32 accumulate',
    memoryType: 'HBM2e',
    memoryGb: 80,
    bandwidthGbs: 2039,
    l2Mb: 40,
    // Unified L1 / shared memory per SM, and the most one block may claim.
    l1SmemPerSmKb: 192,
    maxSmemPerBlockKb: 164,
    registersPerSm: 65536,
    maxThreadsPerSm: 2048,
    maxWarpsPerSm: 64,
    maxBlocksPerSm: 32,
    maxThreadsPerBlock: 1024,
  },
  h100: {
    id: 'h100',
    name: 'H100 SXM5 80GB',
    arch: 'Hopper (GH100)',
    cc: '9.0',
    source: 'NVIDIA H100 Tensor Core GPU datasheet + NVIDIA Hopper GH100 whitepaper',
    sms: 132,
    fp32CoresPerSm: 128,
    tensorCoresPerSm: 4,
    tensorCoreGen: '4th gen',
    clockBoostMhz: 1980,
    fp32Tflops: 67,
    tf32TensorTflops: 494.5,
    fp16TensorTflops: 989.4,
    fp8TensorTflops: 1978.9,
    tensorAccumulateNote: 'FP16/BF16 multiply, FP32 accumulate',
    memoryType: 'HBM3',
    memoryGb: 80,
    bandwidthGbs: 3350,
    l2Mb: 50,
    l1SmemPerSmKb: 256,
    maxSmemPerBlockKb: 227,
    registersPerSm: 65536,
    maxThreadsPerSm: 2048,
    maxWarpsPerSm: 64,
    maxBlocksPerSm: 32,
    maxThreadsPerBlock: 1024,
  },
  rtx4090: {
    id: 'rtx4090',
    name: 'GeForce RTX 4090',
    arch: 'Ada Lovelace (AD102)',
    cc: '8.9',
    source: 'NVIDIA RTX 4090 specifications + NVIDIA Ada GPU architecture whitepaper',
    sms: 128,
    fp32CoresPerSm: 128,
    tensorCoresPerSm: 4,
    tensorCoreGen: '4th gen',
    clockBoostMhz: 2520,
    fp32Tflops: 82.6,
    // Consumer Ada runs tensor ops at half rate when accumulating in FP32.
    fp16TensorTflops: 165.2,
    tensorAccumulateNote:
      'FP16 multiply, FP32 accumulate (consumer Ada runs FP32 accumulate at half rate)',
    memoryType: 'GDDR6X',
    memoryGb: 24,
    bandwidthGbs: 1008,
    l2Mb: 72,
    l1SmemPerSmKb: 128,
    maxSmemPerBlockKb: 100,
    registersPerSm: 65536,
    maxThreadsPerSm: 1536,
    maxWarpsPerSm: 48,
    maxBlocksPerSm: 24,
    maxThreadsPerBlock: 1024,
  },
};

export const GPU_IDS = Object.keys(GPUS);

/**
 * Approximate load-to-use latencies, in SM clock cycles. These are NOT vendor
 * specifications — no vendor publishes them. They come from published
 * microbenchmark studies (Jia et al., "Dissecting the NVIDIA Volta / Turing /
 * Ampere GPU architecture via microbenchmarking") and vary with architecture,
 * access pattern and contention, so they are given as ranges and should be read
 * as orders of magnitude rather than as measurements of any particular chip.
 */
export const LATENCY_CYCLES = [
  { level: 'Register', min: 1, max: 6, note: 'operand collector / dependent-issue latency' },
  { level: 'Shared memory (L1)', min: 20, max: 40, note: 'conflict-free access' },
  { level: 'L2 cache', min: 190, max: 300, note: 'on-die, shared by every SM' },
  { level: 'Global (HBM)', min: 400, max: 800, note: 'off-chip DRAM, on an L2 miss' },
];

export const LATENCY_SOURCE =
  'Jia et al., microbenchmark studies of Volta/Turing/Ampere. Ranges, not vendor specs.';

/* ------------------------------------------------------------------ *
 * 2. Derived quantities
 * ------------------------------------------------------------------ */

/**
 * Peak FP32 throughput implied by the structural specs:
 *   SMs x FP32 lanes/SM x 2 FLOP (one fused multiply-add) x clock.
 * The test suite asserts this reproduces each published `fp32Tflops`, which is
 * what keeps the spec table honest.
 */
export function derivedFp32Tflops(gpu) {
  return (gpu.sms * gpu.fp32CoresPerSm * 2 * gpu.clockBoostMhz * 1e6) / 1e12;
}

export function totalFp32Cores(gpu) {
  return gpu.sms * gpu.fp32CoresPerSm;
}

/** Peak tensor-core throughput used for roofline purposes, in FLOP/s. */
export function peakTensorFlops(gpu) {
  return (gpu.fp16TensorTflops ?? gpu.fp32Tflops) * 1e12;
}

export function peakBandwidthBytesPerSec(gpu) {
  return gpu.bandwidthGbs * 1e9;
}

/**
 * Arithmetic intensity: FLOPs performed per byte moved to or from off-chip
 * memory. The single most useful number in GPU performance work.
 */
export function arithmeticIntensity(flops, bytes) {
  if (bytes <= 0) return Infinity;
  return flops / bytes;
}

/**
 * The roofline: you cannot exceed the hardware's peak FLOP/s, and you cannot
 * exceed bandwidth x arithmetic intensity.
 */
export function attainableFlops(ai, peakFlops, peakBandwidth) {
  return Math.min(peakFlops, ai * peakBandwidth);
}

/**
 * The ridge point — the arithmetic intensity at which a kernel stops being
 * limited by memory and starts being limited by arithmetic. Below it, no amount
 * of faster math helps.
 */
export function ridgePoint(peakFlops, peakBandwidth) {
  return peakFlops / peakBandwidth;
}

export function isMemoryBound(ai, peakFlops, peakBandwidth) {
  return ai < ridgePoint(peakFlops, peakBandwidth);
}

/* ------------------------------ GEMM ------------------------------ */

/** C[MxN] = A[MxK] . B[KxN]. One multiply and one add per inner-product term. */
export function gemmFlops(m, n, k) {
  return 2 * m * n * k;
}

/** Minimum off-chip traffic: read A, read B, write C, each exactly once. */
export function gemmBytes(m, n, k, bytesPerElement) {
  return bytesPerElement * (m * k + k * n + m * n);
}

export function gemmIntensity(m, n, k, bytesPerElement) {
  return arithmeticIntensity(gemmFlops(m, n, k), gemmBytes(m, n, k, bytesPerElement));
}

/* --------------------------- Tiled matmul --------------------------- */

/**
 * Off-chip traffic for an NxN matmul, counted in elements.
 *
 * Naive — every thread computes one C element by streaming a full row of A and a
 * full column of B from global memory: 2N reads per element over N^2 elements.
 *
 * Tiled — each block cooperatively stages a `tile` x `tile` patch of A and of B
 * in shared memory and reuses every staged value `tile` times. A block walks
 * N/tile tile-pairs, reading 2 * (N/tile) * tile^2 elements; there are (N/tile)^2
 * blocks, so 2N^3/tile in total.
 *
 * The reduction is exactly a factor of `tile`. That factor is the entire reason
 * shared memory exists.
 */
export function matmulTraffic(n, tile, bytesPerElement = 4) {
  const naiveReads = 2 * n ** 3;
  const tiledReads = tile > 0 ? (2 * n ** 3) / tile : Infinity;
  const writes = n ** 2;
  return {
    naiveReads,
    tiledReads,
    writes,
    naiveTotal: naiveReads + writes,
    tiledTotal: tiledReads + writes,
    naiveBytes: (naiveReads + writes) * bytesPerElement,
    tiledBytes: (tiledReads + writes) * bytesPerElement,
    reuseFactor: tile,
    reduction: naiveReads / tiledReads,
    tilesPerDim: Math.ceil(n / tile),
    blocks: Math.ceil(n / tile) ** 2,
    // Both the A tile and the B tile must be resident at once.
    smemBytesPerBlock: 2 * tile * tile * bytesPerElement,
  };
}

/* ---------------------------- Attention ---------------------------- */

/**
 * Off-chip traffic for one attention head at sequence length `n`, head dim `d`.
 *
 * Unfused — the n x n score matrix is materialized in HBM, read back by the
 * softmax, written again, then read once more by the PV matmul. Four n^2 trips
 * through memory that compute nothing.
 *
 * Fused (FlashAttention-style) — the scores never leave the SM. Only Q, K, V come
 * in and O goes out, so traffic is linear in n while the FLOPs stay quadratic,
 * and arithmetic intensity therefore grows with sequence length.
 */
export function attentionTraffic({ seqLen: n, headDim: d, bytesPerElement: b = 2, fused = false }) {
  const flops = 4 * n * n * d; // 2n^2*d for QK^T, 2n^2*d for PV
  const bytes = fused
    ? b * 4 * n * d // read Q, K, V + write O
    : b * (4 * n * d + 4 * n * n); // + write S, read S, write P, read P
  return { flops, bytes, ai: arithmeticIntensity(flops, bytes) };
}

/* ------------------- Elementwise / normalization ------------------- */

/**
 * FLOP-per-element estimates for the non-GEMM ops. These are op counts for the
 * usual formulation, not measurements — a different GELU approximation shifts
 * them. It barely matters: these kernels sit two orders of magnitude below any
 * ridge point, so their position is set by bytes, not by FLOPs.
 */
export const FLOPS_PER_ELEMENT = {
  gelu: 8, // tanh approximation
  softmax: 5, // max, subtract, exp, sum, divide
  layernorm: 5, // mean, variance, normalize, scale, shift
  residual: 1, // a single add
};

/** An elementwise kernel reading `reads` tensors and writing `writes` tensors. */
export function elementwiseKernel({
  elements,
  flopsPerElement,
  reads = 1,
  writes = 1,
  bytesPerElement = 2,
}) {
  const flops = elements * flopsPerElement;
  const bytes = elements * (reads + writes) * bytesPerElement;
  return { flops, bytes, ai: arithmeticIntensity(flops, bytes) };
}

/* ----------------- Transformer layer kernel catalog ----------------- */

/**
 * The kernels of one transformer layer, with arithmetic intensities computed
 * from the shapes. `tokens` is batch x sequence — the number of rows flowing
 * through the layer, which is exactly what separates training/prefill (many
 * rows, GEMM-shaped) from autoregressive decode (one row, matvec-shaped).
 */
export function transformerLayerKernels({
  tokens = 4096,
  dModel = 4096,
  heads = 32,
  seqLen = 4096,
  bytesPerElement = 2,
} = {}) {
  const b = bytesPerElement;
  const headDim = dModel / heads;
  const ffn = 4 * dModel;
  const kernels = [];

  const gemm = (name, m, n, k, note) =>
    kernels.push({
      name,
      kind: 'gemm',
      shape: `[${m} x ${k}] . [${k} x ${n}]`,
      flops: gemmFlops(m, n, k),
      bytes: gemmBytes(m, n, k, b),
      ai: gemmIntensity(m, n, k, b),
      note,
    });

  gemm('QKV projection', tokens, 3 * dModel, dModel, 'One GEMM produces all three projections.');

  const unfused = attentionTraffic({ seqLen, headDim, bytesPerElement: b, fused: false });
  kernels.push({
    name: 'Attention (unfused)',
    kind: 'attention',
    shape: `n = ${seqLen}, d_head = ${headDim}`,
    ...unfused,
    note: 'Materializes the n x n score matrix in HBM — four extra n^2 round trips.',
  });

  const fused = attentionTraffic({ seqLen, headDim, bytesPerElement: b, fused: true });
  kernels.push({
    name: 'Attention (fused / FlashAttention)',
    kind: 'attention',
    shape: `n = ${seqLen}, d_head = ${headDim}`,
    ...fused,
    note: 'Scores stay in SRAM. Traffic is linear in n while FLOPs stay quadratic, so AI = n/2.',
  });

  gemm('Attention output proj', tokens, dModel, dModel);
  gemm('FFN up', tokens, ffn, dModel, `Expands to 4 x d_model = ${ffn}.`);
  gemm('FFN down', tokens, dModel, ffn);

  const ew = (name, elements, flopsPerElement, reads, writes, note) => {
    const r = elementwiseKernel({ elements, flopsPerElement, reads, writes, bytesPerElement: b });
    kernels.push({
      name,
      kind: 'elementwise',
      shape: `${elements.toLocaleString('en-US')} elements`,
      ...r,
      note,
    });
  };

  ew('GELU', tokens * ffn, FLOPS_PER_ELEMENT.gelu, 1, 1, 'Pure bandwidth. The classic fusion candidate.');
  ew('LayerNorm', tokens * dModel, FLOPS_PER_ELEMENT.layernorm, 2, 1, 'Two reads: one for the statistics, one to normalize.');
  ew('Residual add', tokens * dModel, FLOPS_PER_ELEMENT.residual, 2, 1, 'One add per element. Very nearly pure memory traffic.');

  return kernels;
}

/** A single decode step: one token, so every GEMM collapses to a matrix-vector product. */
export function decodeStepKernel({ dModel = 4096, bytesPerElement = 2 } = {}) {
  const flops = gemmFlops(1, dModel, dModel);
  const bytes = gemmBytes(1, dModel, dModel, bytesPerElement);
  return {
    name: 'Decode matvec (batch = 1)',
    kind: 'gemm',
    shape: `[1 x ${dModel}] . [${dModel} x ${dModel}]`,
    flops,
    bytes,
    ai: arithmeticIntensity(flops, bytes),
    note: 'Every weight is fetched to be used exactly once. AI ~ 1 — hopelessly memory bound, which is why LLM decoding is a bandwidth problem and why batching helps so much.',
  };
}

/* ----------------------------- Occupancy ----------------------------- */

export function ceilTo(value, unit) {
  return Math.ceil(value / unit) * unit;
}

/**
 * How many warps an SM can keep resident, given a block's resource appetite.
 * Occupancy is not a goal in itself — it is the SM's supply of other work to
 * issue while some warp waits several hundred cycles on HBM.
 */
export function occupancy({ threadsPerBlock, registersPerThread, sharedMemPerBlockBytes = 0 }, gpu) {
  const warpsPerBlock = Math.ceil(threadsPerBlock / WARP_SIZE);

  const byWarps = Math.floor(gpu.maxWarpsPerSm / warpsPerBlock);

  const regsPerWarp = ceilTo(registersPerThread * WARP_SIZE, REG_ALLOC_UNIT_PER_WARP);
  const byRegisters = Math.floor(gpu.registersPerSm / (regsPerWarp * warpsPerBlock));

  const smemPerSmBytes = gpu.l1SmemPerSmKb * 1024;
  const smemPerBlock = ceilTo(sharedMemPerBlockBytes, SMEM_ALLOC_UNIT_BYTES);
  const bySharedMem = smemPerBlock > 0 ? Math.floor(smemPerSmBytes / smemPerBlock) : Infinity;

  const limits = [
    { name: 'Warp slots', blocks: byWarps },
    { name: 'Registers', blocks: byRegisters },
    { name: 'Shared memory', blocks: bySharedMem },
    { name: 'Blocks/SM cap', blocks: gpu.maxBlocksPerSm },
  ];

  const blocksPerSm = Math.max(0, Math.min(...limits.map((l) => l.blocks)));
  const activeWarps = blocksPerSm * warpsPerBlock;
  const limiter = limits.filter((l) => l.blocks === blocksPerSm).map((l) => l.name);

  return {
    warpsPerBlock,
    blocksPerSm,
    activeWarps,
    activeThreads: activeWarps * WARP_SIZE,
    occupancy: activeWarps / gpu.maxWarpsPerSm,
    regsPerWarp,
    smemPerBlock,
    limits,
    limiter,
  };
}

/* -------------------------- Memory coalescing -------------------------- */

/**
 * The memory system does not serve threads, it serves 32-byte sectors. A warp's
 * 32 addresses are satisfied by however many distinct sectors they touch — 4 if
 * the lanes read contiguous 4-byte values, up to 32 if they are scattered.
 */
export function coalescedTransactions(byteAddresses, sectorBytes = 32, elementBytes = 4) {
  const sectors = new Set(byteAddresses.map((a) => Math.floor(a / sectorBytes)));
  const bytesRequested = byteAddresses.length * elementBytes;
  const bytesFetched = sectors.size * sectorBytes;
  return {
    sectors: sectors.size,
    sectorIds: [...sectors].sort((x, y) => x - y),
    bytesRequested,
    bytesFetched,
    efficiency: bytesRequested / bytesFetched,
    wasteFactor: bytesFetched / bytesRequested,
  };
}

/** Warp lane addresses for a strided access `base + lane * stride` (in elements). */
export function stridedWarpAddresses(stride, elementBytes = 4, base = 0, lanes = WARP_SIZE) {
  return Array.from({ length: lanes }, (_, lane) => base + lane * stride * elementBytes);
}

/* --------------------------- Warp divergence --------------------------- */

/**
 * A warp has one instruction pointer. When its lanes disagree about a branch the
 * hardware executes both sides with the non-participating lanes masked off, so
 * the cost is the sum of the two paths rather than the longer of them — and the
 * masked lanes did no useful work while the other side ran.
 */
export function divergenceCost(laneTakesBranch, cyclesIfTrue, cyclesIfFalse) {
  const taken = laneTakesBranch.filter(Boolean).length;
  const notTaken = laneTakesBranch.length - taken;
  const divergent = taken > 0 && notTaken > 0;

  const cycles = divergent ? cyclesIfTrue + cyclesIfFalse : taken > 0 ? cyclesIfTrue : cyclesIfFalse;

  // Lane-cycles that did useful work, over the lane-cycles the warp occupied.
  const usefulLaneCycles = taken * cyclesIfTrue + notTaken * cyclesIfFalse;
  const totalLaneCycles = laneTakesBranch.length * cycles;

  return {
    taken,
    notTaken,
    divergent,
    pathsExecuted: divergent ? 2 : 1,
    cycles,
    idealCycles: Math.max(taken > 0 ? cyclesIfTrue : 0, notTaken > 0 ? cyclesIfFalse : 0),
    laneEfficiency: totalLaneCycles === 0 ? 1 : usefulLaneCycles / totalLaneCycles,
  };
}

/* ------------------------------ Formatting ------------------------------ */

export function formatBytes(n) {
  const units = ['B', 'KB', 'MB', 'GB', 'TB', 'PB'];
  let i = 0;
  let v = n;
  while (v >= 1024 && i < units.length - 1) {
    v /= 1024;
    i += 1;
  }
  return `${Number.isInteger(v) || v >= 100 ? Math.round(v) : v.toFixed(1)} ${units[i]}`;
}

export function formatFlops(n) {
  const units = ['FLOP', 'KFLOP', 'MFLOP', 'GFLOP', 'TFLOP', 'PFLOP'];
  let i = 0;
  let v = n;
  while (v >= 1000 && i < units.length - 1) {
    v /= 1000;
    i += 1;
  }
  return `${Number.isInteger(v) || v >= 100 ? Math.round(v) : v.toFixed(1)} ${units[i]}`;
}

export function formatCount(n) {
  if (n >= 1e12) return `${(n / 1e12).toFixed(1)}T`;
  if (n >= 1e9) return `${(n / 1e9).toFixed(1)}B`;
  if (n >= 1e6) return `${(n / 1e6).toFixed(1)}M`;
  if (n >= 1e3) return `${(n / 1e3).toFixed(1)}K`;
  return String(Math.round(n));
}