pub const GATHER_CU: &str = "// Stage 0: the CSR coefficient gather, plus the C matrix that snarkjs never stores.\n//\n// Twin of crates/metal/src/shaders/gather.metal. Depends on bn254_fr.cuh being\n// concatenated ahead of this file by snarkrs_gpu_kernels::unit_stages; it carries no\n// #include of its own because NVRTC has no filesystem.\n//\n// ============================================================================\n// WHY ONE THREAD PER ROW\n// ============================================================================\n//\n// The rows are short and ragged, and they get more so as the circuit grows. Measured on\n// the real proving keys, mean row length / empty rows / longest row:\n//\n// js_1x1_d8 domain 2^12 A 2.31 / 732 / 65 B 3.70 / 737 / 64\n// js_2x2_d16 domain 2^14 A 1.51 / 6224 / 67 B 2.34 / 6231 / 66\n// js_8x8_d32 domain 2^17 A 1.13 / 60696 / 79 B 1.70 / 60715 / 78\n// js_16x16_d32 domain 2^18 A 1.13 / 121848 / 95 B 1.69 / 121883 / 94\n//\n// So the mean row is one or two terms, nearly half the rows are empty (the domain rounds\n// up well past the constraint count), and the longest row in the whole key is under a\n// hundred. That rules out every cooperative mapping:\n//\n// * One block per row would put 2.31 multiplies on a 32-lane warp and run at\n// single-digit occupancy, and it would need a __shared__ reduction to combine partial\n// sums that are almost always a single term.\n// * One warp per row is the same argument with a smaller constant: 31 of 32 lanes idle\n// on the average row.\n// * A flat \"one thread per nonzero\" mapping is a scatter, which needs either atomics\n// (there is no 32-byte atomic) or a segmented reduction over the whole nonzero\n// stream. The CSR sort in snarkrs-formats exists precisely so we do not have to do that.\n//\n// One thread per row instead puts 32 consecutive rows in one warp, so a warp retires\n// roughly 74 (A) to 118 (B) Montgomery multiplies with full lane utilisation. Divergence\n// is bounded by the longest row in the warp rather than the longest row in the matrix,\n// and neighbouring constraint rows in a circom circuit have similar lengths, so the\n// ragged tail is short. An empty row costs one comparison.\n//\n// The argument carries over from Metal unchanged because a Metal SIMD group and a CUDA\n// warp are both 32 lanes with the same \"the group waits for its slowest lane\" cost model.\n//\n// Nothing here writes a location another thread reads, so there is no barrier, no atomic\n// and no ordering requirement. That is the whole payoff of the CSR sort.\n//\n// ============================================================================\n// WHY C IS COMPUTED HERE AND NOT GATHERED\n// ============================================================================\n//\n// There is no C matrix in a snarkjs zkey. snarkjs\' buildABC1 fills C by multiplying the\n// A and B evaluations pointwise, which is exact rather than an approximation: the R1CS\n// constraint *is* a*b = c, so on the constraint domain the C evaluation of row i is the\n// product of the A and B evaluations of row i. Doing it in this kernel means the product\n// reuses the two values already in registers, so C costs one Montgomery multiply per row\n// and zero extra loads. A separate pass would cost a full read of A and B (2n * 32 bytes)\n// on a stage that is already memory bound, and on a discrete card that traffic is real\n// GDDR bandwidth rather than a read out of Apple\'s unified pool.\n\n#ifndef G16_GATHER_CU\n#define G16_GATHER_CU\n\n// out_a[c] = sum over CSR row c of value_a[k] * witness[signal_a[k]]\n// out_b[c] = the same for matrix B\n// out_c[c] = out_a[c] * out_b[c]\n//\n// `witness` is Montgomery-form Fr (layout::PackedFr), as are the CSR values, so the\n// products need no conversion in either direction.\n//\n// LAUNCH: grid_dim.x = ceil(n / block_dim.x), block_dim.x free (256, the Metal twin\'s\n// preference, is the right starting point). One thread owns exactly one output row and\n// touches no other thread\'s data, so a block is just 256 independent rows.\n//\n// No __restrict__ on the read-only pointers. It would unlock the non-coherent cache path,\n// but nothing in this signature tells the compiler that `witness` and the CSR arrays come\n// from disjoint allocations, and a __restrict__ that is ever a lie is a wrong proof with\n// no diagnostic. The gather is bound by the dependent load witness[signal[k]] in any\n// case, which no cache hint fixes.\nextern \"C\" __global__ void g16_gather_abc(\n const u32* row_ptr_a,\n const u32* signal_a,\n const Fr* value_a,\n const u32* row_ptr_b,\n const u32* signal_b,\n const Fr* value_b,\n const Fr* witness,\n Fr* out_a,\n Fr* out_b,\n Fr* out_c,\n const u32 n)\n{\n u32 gid = blockIdx.x * blockDim.x + threadIdx.x;\n\n // In Metal this guard was belt and braces, because dispatch_threads sizes the final\n // threadgroup exactly. Here it is load bearing. A CUDA grid is always a whole number\n // of blocks, so for any n that is not a multiple of block_dim.x the last block runs\n // threads past the end, and without this guard they would read row_ptr[n + 1] and\n // write past out_a.\n if (gid >= n) {\n return;\n }\n\n Fr a = fr_zero();\n u32 lo = row_ptr_a[gid];\n u32 hi = row_ptr_a[gid + 1u];\n for (u32 k = lo; k < hi; k++) {\n a = fr_add(a, fr_mul(value_a[k], witness[signal_a[k]]));\n }\n\n Fr b = fr_zero();\n lo = row_ptr_b[gid];\n hi = row_ptr_b[gid + 1u];\n for (u32 k = lo; k < hi; k++) {\n b = fr_add(b, fr_mul(value_b[k], witness[signal_b[k]]));\n }\n\n out_a[gid] = a;\n out_b[gid] = b;\n out_c[gid] = fr_mul(a, b);\n}\n\n#endif // G16_GATHER_CU\n";Expand description
Stage 0: gather the A and B coefficient columns out of the CSR proving key.