pub const NTT_CU: &str = "// Stages 1-3: the six NTTs, the bit-reverse permutation, the iNTT normalisation and the\n// coset shift, all folded into two kernels.\n//\n// Twin of crates/metal/src/shaders/ntt.metal. Depends on bn254_fr.cuh being\n// concatenated ahead of this file. No #include: NVRTC has no filesystem.\n//\n// ============================================================================\n// THE ALGORITHM IS THE CPU BACKEND\'S, EXACTLY\n// ============================================================================\n//\n// snarkrs_ntt::CpuNtt is decimation in time: bit-reverse the input, then log(n) passes with\n// half = 1, 2, 4, ..., n/2, where butterfly j of a block reads twiddles[j * n/(2*half)]\n// from a natural-order table of the root\'s powers. Everything below computes that same\n// sequence of butterflies in that same order with that same twiddle table, so the CPU,\n// the Metal and the CUDA backends agree bit for bit and not merely up to a permutation.\n//\n// Two of the CPU backend\'s separate steps are folded into the load epilogue, which is\n// exact rather than approximate:\n//\n// * The iNTT\'s 1/n normalisation. The CPU multiplies every output by size_inv after the\n// transform; we multiply every input by it before. The transform is Fr-linear, so\n// these are the same map, and doing it on load costs nothing because the element is\n// already in a register.\n// * The stage 2 coset shift, x[j] *= shift^j. Applied to the input of the forward\n// transform, again on load. The powers are NOT the twiddles: the twiddles are powers\n// of the domain\'s own root of unity, the shift is a primitive 2n-th root (see\n// CpuCircuit::new for why that specific element and not Fr::GENERATOR). snarkrs_field\'s\n// Domain does not build this table and does not cache the twiddles either, so the\n// host builds all three once at prepare time and leaves them resident.\n//\n// ============================================================================\n// BATCHING: WHY TWO LAUNCHES PER TRANSFORM AND NOT log(n)\n// ============================================================================\n//\n// One launch per pass means log(n) round trips through device memory. At 2^18 that is\n// 18 passes x 16 MB of read+write traffic. On Apple silicon the measured streaming rate\n// for this field at that size was 30 GB/s against a bus that does 300, i.e. memory bound\n// with the ALU idle; on a discrete card the same shape holds and the DRAM round trip is\n// worse, not better. So passes are batched: a block pulls a slice into __shared__, runs\n// several passes there with only a __syncthreads() between them, and writes back once.\n//\n// Which passes can share a slice is fixed by the index algebra. Pass t (0-based) flips\n// bit t of an element index. So the passes in [s0, s0+k) touch only bits [s0, s0+k), and\n// a group is any set of indices agreeing on every other bit: fix `low` = bits [0, s0) and\n// `high` = bits [s0+k, log n), vary the k middle bits. That is a contiguous run of 2^k\n// elements when s0 = 0, and a run strided by 2^s0 otherwise. This is the standard six-step\n// decomposition, written as one kernel that takes (s0, k) rather than two kernels.\n//\n// ----------------------------------------------------------------------------\n// THE SHARED MEMORY BUDGET, WORKED OUT FOR CUDA RATHER THAN COPIED FROM METAL\n// ----------------------------------------------------------------------------\n//\n// The slice size is capped by __shared__ memory per block, which is NOT the same number\n// as Metal\'s maxThreadgroupMemoryLength and must not be assumed to be:\n//\n// * Every architecture from Volta up gives a block 48 KB without an opt in. Going above\n// that needs cudaFuncSetAttribute(MaxDynamicSharedMemorySize), which raises the\n// ceiling to 64 KB on Turing (sm_75), 99 KB on the consumer Ampere and Ada parts\n// (sm_86, sm_89) and 163 KB on A100 (sm_80). That opt in is deliberately not taken:\n// it buys one extra pass at most (see the arithmetic below) and costs a per-function\n// driver call plus an architecture-dependent limit the host would have to query and\n// get right on every card.\n// * An Fr is 32 bytes, so 48 KB holds 1536 of them. A slice must be a power of two\n// because it is 2^k elements, and the largest power of two at or below 1536 is\n// 1024 = 2^10. So k <= 10.\n// * Raising the ceiling to 64 KB would give 2048 elements and k <= 11, one extra pass.\n// At 2^18 that turns an 18 = 9 + 9 split into 18 = 9 + 9 again (two batches either\n// way), so it changes nothing at the sizes this prover actually runs.\n//\n// k <= 10 is the same cap the M2 Max reaches from its 32768-byte threadgroup limit\n// (32768 / 32 = 1024 exactly), by arithmetic coincidence rather than by copying. That is\n// convenient: both backends run the identical pass split at every artifact size, so a\n// Metal-versus-CUDA timing compares two implementations of one schedule.\n//\n// The host splits log(n) into ceil(log n / 10) batches of near-equal size rather than\n// greedily filling the first: an 18 = 10 + 8 split wastes half the second kernel\'s\n// parallelism and, at other sizes, produces the degenerate two-element tail kernel that\n// profiling of bellperson caught at 2^26. 18 = 9 + 9 instead, which is 512 elements and\n// 16 KB per block. That also matters for occupancy here in a way it did not on Metal: a\n// full 1024-element slice is 32 KB, and a Turing SM has 64 KB of shared memory in total,\n// so a maximal batch would cap residency at two blocks per SM. 16 KB allows four, as far\n// as shared memory is concerned; registers may still bind first, which is a reason to\n// measure rather than to assume.\n//\n// At the 2^18 domain of the largest artifact that is TWO launches per transform, twelve\n// for the six transforms of a proof, and the whole of stages 0 to 4 is fourteen launches\n// on one stream with one synchronize at the end. A CUDA kernel launch on an already-warm\n// stream is a few microseconds of host time and is asynchronous, so the fourteen queue up\n// without the host blocking; the Metal backend batches into one command buffer for the\n// same reason and reaches the same place.\n//\n// ============================================================================\n// WHAT IS *NOT* DONE HERE\n// ============================================================================\n//\n// No twiddle caching in __shared__. The slice is already 16 KB of the 48 KB budget and\n// the twiddles a batch touches are 2^(k-1) distinct values, another 8 KB at k = 9, which\n// would push a block to 24 KB and halve how many fit on an SM. The twiddle reads hit L2\n// hard (every block in a launch reads the same table) so the cache already does most of\n// this. Left on the table deliberately.\n\n#ifndef G16_NTT_CU\n#define G16_NTT_CU\n\n// Defined in pointwise.cu, which snarkrs_gpu_kernels::unit_stages concatenates AFTER this file.\n//\n// Note the ordering differs from the Metal backend on purpose. There, pointwise.metal is\n// placed ahead of ntt.metal so the definition is in scope. Here the unit order is fixed\n// as gather / ntt / pointwise by snarkrs_gpu_kernels::unit_stages, so the fused store\n// epilogue below reaches its helper through this forward declaration instead. Same\n// translation unit, so __forceinline__ still applies; only the textual order changed.\n__device__ __forceinline__ void g16_store_h(Fr* h_mont, Fr* h_std, u32 i, Fr v);\n\n// Load-time scaling and the store-time epilogue. Both are uniform across the whole\n// launch, so each costs a scalar compare and neither ever diverges a warp.\n//\n// Enumerators rather than the `__constant__` the general MSL translation rule would give\n// for `constant uint`, for two reasons. That rule is about tables the kernel reads at run\n// time, like FR_N in bn254_fr.cuh; these are compile-time tags and want to fold into\n// immediate operands, not into a load from constant memory. And two of the five (NONE and\n// PLAIN) are never compared against, since the branches test only for the non-default\n// modes, so as variables they draw nvcc warning #177 \"declared but never referenced\".\n// Deleting them would leave the mode space half documented, and the host-side Rust mirror\n// does use both values. An enumerator is never an unused variable.\nenum G16ScaleMode : u32 {\n G16_SCALE_NONE = 0u, // forward transform of an already-shifted vector\n G16_SCALE_CONST = 1u, // iNTT: multiply by size_inv\n G16_SCALE_TABLE = 2u, // stage 2: multiply by shift^j, j the natural index\n};\n\nenum G16StoreMode : u32 {\n G16_STORE_PLAIN = 0u,\n G16_STORE_JOIN = 1u, // stage 4 fused in: write H = A*B - this, not this\n};\n\n// Mirrors snarkrs_cuda::stages::NttParams. 52 bytes, all 4-byte aligned, no padding, which is\n// what makes it safe to hand across as a by-value kernel parameter: CUDA copies the\n// struct into the parameter space verbatim and the Rust side is #[repr(C)] with the same\n// five u32 followed by a PackedFr.\n//\n// Passed BY VALUE, not by pointer, everywhere. Mixing the two is the kind of mismatch\n// that reads a modulus out of the wrong offset and produces a wrong proof rather than a\n// fault, so it is stated here as a contract and not left to each call site.\nstruct NttParams {\n u32 log_n; // log2 of the domain size\n u32 s0; // index of the first pass in this batch\n u32 k; // number of passes in this batch; slice is 2^k elements\n u32 scale_mode;\n u32 store_mode;\n Fr kscale; // the G16_SCALE_CONST multiplier, Montgomery form\n};\n\nstatic_assert(sizeof(NttParams) == 52, \"NttParams must match stages::NttParams\");\n\n// Reverse all 32 bits of x. MSL has `reverse_bits` as a language builtin; CUDA has the\n// __brev intrinsic, which lowers to the single PTX instruction brev.b32.\n//\n// Written as a named wrapper rather than inlined at the call site because it is the one\n// place where a builtin is assumed to exist under NVRTC (which compiles without any CUDA\n// header). If a future NVRTC drops it, the fix is the six-line shift-and-mask reversal\n// here and nothing else changes.\n__device__ __forceinline__ u32 g16_brev(u32 x) {\n return __brev(x);\n}\n\n// The k passes of one batch, over a slice already resident in __shared__.\n//\n// `low` is bits [0, s0) of every index in the slice, which is what makes the twiddle\n// index depend on more than the local position: at global pass s0 + t the butterfly\'s\n// index within its block is low + (m mod 2^t) * 2^s0, where m is the local index, and the\n// twiddle is that times 2^(log_n - s0 - t - 1). For the s0 = 0 head this collapses to the\n// familiar (m mod 2^t) << (log_n - t - 1).\n//\n// `sh` is a pointer into __shared__ and `tw` a pointer into global memory. CUDA infers\n// both address spaces from the pointer\'s provenance, which is why the MSL `threadgroup` /\n// `device` qualifiers simply disappear in the port rather than turning into anything.\n__device__ __forceinline__ void g16_ntt_batch(Fr* sh,\n const Fr* tw,\n u32 log_n, u32 s0, u32 k, u32 low,\n u32 tid, u32 tgsz)\n{\n // A zero-width batch has no butterflies and `1u << (k - 1u)` would shift by\n // 0xFFFFFFFF, which is undefined and lets ptxas hand back any `halves` it likes,\n // including one that walks the slice off the end. k comes from NttParams, so the\n // branch is uniform across the block and none of the __syncthreads() below is\n // reached by only some of its threads.\n if (k == 0u) {\n return;\n }\n\n u32 halves = 1u << (k - 1u); // butterflies per pass\n for (u32 t = 0u; t < k; t++) {\n u32 hl = 1u << t;\n u32 twshift = log_n - s0 - t - 1u;\n for (u32 bfy = tid; bfy < halves; bfy += tgsz) {\n u32 jm = bfy & (hl - 1u);\n u32 lo = ((bfy >> t) << (t + 1u)) | jm;\n u32 hi = lo + hl;\n Fr u = sh[lo];\n Fr v = fr_mul(sh[hi], tw[(low + (jm << s0)) << twshift]);\n sh[lo] = fr_add(u, v);\n sh[hi] = fr_sub(u, v);\n }\n // Outside the strided loop, so every thread in the block reaches it the same\n // number of times even when halves < blockDim.x and some threads do no work at\n // all. A __syncthreads() that only some threads reach is undefined behaviour on\n // CUDA, not merely slow.\n __syncthreads();\n }\n}\n\n// Batch 0. Out of place, src -> dst, and it absorbs the bit-reverse permutation into the\n// load: after the permutation position i holds src[reverse(i)], so the thread that owns\n// output position i simply reads src at the reversed index. The CPU backend does the same\n// permutation as an in-place pass of swaps; skipping it as a launch saves a full read\n// and write of the domain per transform, six times per proof.\n//\n// LAUNCH: grid_dim.x = max(1, n >> k) blocks, block_dim.x free (256), and\n// shared_mem_bytes = (1 << k) * 32, the slice and nothing else. One BLOCK owns the\n// contiguous run of 2^k elements starting at blockIdx.x << k and runs all k passes over\n// it; one THREAD strides over that run by blockDim.x, so it handles 2^k / blockDim.x\n// elements on load and store and the same count of butterflies per pass. blockDim.x need\n// not divide 2^k and need not be at most 2^k; both loops are `for (m = tid; m < blk;\n// m += tgsz)` precisely so any block width is legal.\nextern \"C\" __global__ void g16_ntt_head(\n const Fr* src,\n Fr* dst,\n const Fr* tw,\n const Fr* ptab, // shift^j table, G16_SCALE_TABLE only\n const Fr* join_a, // G16_STORE_JOIN only\n const Fr* join_b,\n Fr* h_mont,\n Fr* h_std,\n const NttParams p)\n{\n // Dynamic __shared__, sized by the launch\'s shared_mem_bytes. It has to be dynamic\n // rather than a fixed `__shared__ Fr sh[1024]`: k varies with the domain and with the\n // pass split, and a static 32 KB array would charge every launch the maximum and cut\n // occupancy to one block per SM even for a 2^12 domain.\n extern __shared__ Fr sh[];\n\n u32 tid = threadIdx.x;\n u32 tgsz = blockDim.x;\n u32 blk = 1u << p.k;\n u32 base = blockIdx.x << p.k;\n // A domain of size 1 has log_n = 0, and `>> 32` is undefined, so the degenerate case\n // is a uniform branch rather than a shift the compiler is free to fold to anything.\n u32 revshift = 32u - p.log_n;\n\n for (u32 m = tid; m < blk; m += tgsz) {\n u32 i = base + m;\n u32 s = (p.log_n == 0u) ? 0u : (g16_brev(i) >> revshift);\n Fr x = src[s];\n if (p.scale_mode == G16_SCALE_CONST) {\n x = fr_mul(x, p.kscale);\n } else if (p.scale_mode == G16_SCALE_TABLE) {\n // Indexed by the *source* index, which is the natural-order position the\n // shift power belongs to. Indexing by `i` here would apply shift^bitrev(j),\n // which is a different and wrong vector.\n x = fr_mul(x, ptab[s]);\n }\n sh[m] = x;\n }\n __syncthreads();\n\n g16_ntt_batch(sh, tw, p.log_n, 0u, p.k, 0u, tid, tgsz);\n\n for (u32 m = tid; m < blk; m += tgsz) {\n u32 i = base + m;\n if (p.store_mode == G16_STORE_JOIN) {\n g16_store_h(h_mont, h_std, i, fr_sub(fr_mul(join_a[i], join_b[i]), sh[m]));\n } else {\n dst[i] = sh[m];\n }\n }\n}\n\n// Every batch after the first. In place on `a`, slice strided by 2^s0.\n//\n// Block g owns the slice with low = g mod 2^s0 and high = g >> s0, i.e. base index\n// low + high * 2^(s0+k) and elements base + m * 2^s0 for m in [0, 2^k). There are\n// 2^(log_n - k) blocks and the two halves of g cover exactly the log_n - k bits outside\n// the batch, so the blocks partition the domain with no overlap and no gap.\n//\n// LAUNCH: identical geometry to the head, grid_dim.x = max(1, n >> k) blocks,\n// shared_mem_bytes = (1 << k) * 32. This kernel is not launched with k = 0 today, because\n// a zero-width batch only ever appears as the single batch of a size-1 domain and that\n// one is served by the head; g16_ntt_batch returns early on k = 0 regardless.\nextern \"C\" __global__ void g16_ntt_tail(\n Fr* a,\n const Fr* tw,\n const Fr* join_a,\n const Fr* join_b,\n Fr* h_mont,\n Fr* h_std,\n const NttParams p)\n{\n extern __shared__ Fr sh[];\n\n u32 tid = threadIdx.x;\n u32 tgsz = blockDim.x;\n u32 blk = 1u << p.k;\n u32 low = blockIdx.x & ((1u << p.s0) - 1u);\n u32 high = blockIdx.x >> p.s0;\n u32 base = low + (high << (p.s0 + p.k));\n\n for (u32 m = tid; m < blk; m += tgsz) {\n sh[m] = a[base + (m << p.s0)];\n }\n __syncthreads();\n\n g16_ntt_batch(sh, tw, p.log_n, p.s0, p.k, low, tid, tgsz);\n\n for (u32 m = tid; m < blk; m += tgsz) {\n u32 i = base + (m << p.s0);\n if (p.store_mode == G16_STORE_JOIN) {\n g16_store_h(h_mont, h_std, i, fr_sub(fr_mul(join_a[i], join_b[i]), sh[m]));\n } else {\n a[i] = sh[m];\n }\n }\n}\n\n#endif // G16_NTT_CU\n";Expand description
Stages 1 to 3: the six transforms, iNTT then coset shift then forward NTT.