diff --git a/Cargo.lock b/Cargo.lock index 74986dcc9..ac88f2ac5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -473,7 +473,7 @@ version = "0.19.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f071cd6a7b5d51607df76aa2d426aaabc7a74bc6bdb885b8afa63a880572ad9b" dependencies = [ - "libloading", + "libloading 0.9.0", ] [[package]] @@ -867,6 +867,16 @@ version = "0.2.178" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + [[package]] name = "libloading" version = "0.9.0" @@ -927,6 +937,7 @@ version = "0.1.0" dependencies = [ "crypto", "cudarc", + "libloading 0.8.9", "math", "rand 0.8.5", "rand_chacha 0.3.1", diff --git a/bin/cli/Cargo.toml b/bin/cli/Cargo.toml index e4fcdb7fd..b744714c3 100644 --- a/bin/cli/Cargo.toml +++ b/bin/cli/Cargo.toml @@ -19,3 +19,5 @@ env_logger = "0.11" jemalloc-stats = ["dep:tikv-jemalloc-ctl"] disk-spill = ["prover/disk-spill"] instruments = ["prover/instruments", "stark/instruments"] +# GPU profiling build (Nsight): CUDA prover + instruments spans + NVTX ranges. +nvtx = ["prover/nvtx", "instruments"] diff --git a/crypto/math-cuda/Cargo.toml b/crypto/math-cuda/Cargo.toml index 7a28498da..5a214f7bd 100644 --- a/crypto/math-cuda/Cargo.toml +++ b/crypto/math-cuda/Cargo.toml @@ -30,11 +30,17 @@ cudarc = { version = "0.19", default-features = false, features = [ ] } math = { path = "../math" } rayon = "1.7" +# NVTX range emission for Nsight timelines (dlopen'd at runtime, see src/nvtx.rs). +libloading = { version = "0.8", optional = true } [features] # Test-only fault injection in FriCommitState. Production builds leave this # off so the fault check is fully elided at compile time. test-faults = [] +# NVTX ranges around every GPU pipeline entry point, for Nsight Systems +# profiling. Zero-cost when disabled; when enabled but libnvToolsExt is +# absent at runtime, every call is a cheap no-op. +nvtx = ["dep:libloading"] [dev-dependencies] crypto = { path = "../crypto" } diff --git a/crypto/math-cuda/build.rs b/crypto/math-cuda/build.rs index 7bd7c04cc..fbd70eb5b 100644 --- a/crypto/math-cuda/build.rs +++ b/crypto/math-cuda/build.rs @@ -82,6 +82,7 @@ fn compile_kernel(src: &str, out_name: &str, have_nvcc: bool) { println!("cargo:rerun-if-env-changed=CUDA_HOME"); println!("cargo:rerun-if-env-changed=CUDA_PATH"); println!("cargo:rerun-if-env-changed=CUDARC_NVCC_ARCH"); + println!("cargo:rerun-if-env-changed=LAMBDA_VM_NVCC_LINEINFO"); // When nvcc is missing from PATH, emit an empty cubin stub so the crate // still compiles. include_bytes! in src/device.rs needs the file to exist @@ -115,8 +116,15 @@ fn compile_kernel(src: &str, out_name: &str, have_nvcc: bool) { .map(|a| to_real_arch(&a)) .unwrap_or_else(|_| detect_arch()); - let status = Command::new(nvcc_path()) - .args(["--cubin", "-O3", "-std=c++17", "-arch", &arch, "-o"]) + let mut cmd = Command::new(nvcc_path()); + cmd.args(["--cubin", "-O3", "-std=c++17", "-arch", &arch]); + // SASS→source line mapping for Nsight Compute. Unlike -G this does not + // change codegen, but keep it opt-in so production cubins stay byte-stable. + if env::var("LAMBDA_VM_NVCC_LINEINFO").is_ok_and(|v| v != "0" && !v.is_empty()) { + cmd.arg("-lineinfo"); + } + let status = cmd + .arg("-o") .arg(&out_path) .arg(&src_path) .status() diff --git a/crypto/math-cuda/kernels/constraint_interp.cu b/crypto/math-cuda/kernels/constraint_interp.cu index a4fa77b72..dccd8948e 100644 --- a/crypto/math-cuda/kernels/constraint_interp.cu +++ b/crypto/math-cuda/kernels/constraint_interp.cu @@ -2,29 +2,38 @@ // // Evaluates a captured `ConstraintProgram` (lowered to the flat device blob by // `crypto/stark/src/constraint_ir/device.rs`) over every row of a -// device-resident LDE, producing the per-constraint evaluations. It is a -// transliteration of the CPU walker `eval_device_program` (same module), with -// `FieldElement` arithmetic replaced by `goldilocks.cuh` / `ext3.cuh` — the two -// are asserted bit-for-bit equal by the pre-GPU parity test, so this kernel's -// output equals the compiled prover folder. +// device-resident LDE. It is a transliteration of the CPU walker +// `eval_device_program` (same module), with `FieldElement` arithmetic replaced +// by `goldilocks.cuh` / `ext3.cuh` — the two are asserted bit-for-bit equal by +// the pre-GPU parity test, so this kernel's output equals the compiled prover +// folder. // -// Design (v1, "stripped" — mirrors OpenVM's GLOBAL=true quotient kernel): +// Design (v2, dim-split + liveness slots): // * One thread per LDE row, grid-stride over all rows (fixed launch, any size). -// * Per-thread value array in GLOBAL memory, strided by thread for coalescing: -// node `i` for this thread lives at `d_values[task_offset + i*num_threads]`. -// * ALL values are carried as ext3 `Fe3` (a base value `x` is the embedding -// `{x,0,0}`). Because embedding is a ring homomorphism and the device field -// arithmetic is bit-identical to the CPU path, doing every op in ext3 yields -// outputs bit-identical to the dim-split CPU walk — the per-node `dim` tag -// is therefore unused here and reserved for the base/ext-split optimization -// (Phase 6). `Op::Embed` is consequently the identity. +// * The lowering assigns every node a slot in one of two per-thread scratch +// classes — base (`u64`) or ext (`Fe3`) — with liveness reuse, so scratch +// is sized by the program's max-live-set, not its node count. Slots are +// strided by thread for coalescing: base slot `s` for this thread is +// `vb[s * num_threads + tid]`; ext slot `s` keeps its three components at +// `ve[(s*3 + k) * num_threads + tid]`. +// * Operands are encoded as `kind << 29 | payload` (see `OPK_*`): a slot in +// either class, or a direct reference into the tiny uniform tables +// (constants, RAP challenges, alpha powers, table offset) — uniform leaves +// never touch scratch. +// * Base-dim arithmetic runs in the base field (1 mul vs 9 for ext3), and +// mixed base×ext ops use shortcuts (`mul_base`, componentwise add/sub) +// that are bit-identical to the full ext op on the embedded operand: +// embedding is a ring homomorphism, `gl::add(x,0) == x == gl::sub(x,0)`, +// and `dot3` with zero products reduces to `gl::mul`. Where an identity +// is NOT guaranteed bitwise (negating an embedded zero limb), the full +// form is kept (`gl::sub(0, y)`, never `gl::neg(y)`). // // Output is the per-constraint eval matrix `d_evals[c*num_rows + row]` (Fe3; -// base-rooted constraints carry their value in `.a`). Fusing the -// `z*Σ(Cᵢ·βᵢ) + boundary` accumulation into this kernel (to avoid a D2H of the -// matrix — the actual data-residency win) is the pipeline-integration follow-on. +// base-rooted constraints carry their value in `.a`). The composition kernel +// below fuses the `z*Σ(Cᵢ·βᵢ) + boundary` accumulation instead, avoiding the +// matrix entirely. // -// Op tags and the `Var` packing MUST stay in sync with +// Op tags, operand kinds and the `res`/root packing MUST stay in sync with // `crypto/stark/src/constraint_ir/device.rs`. #include "goldilocks.cuh" @@ -45,12 +54,27 @@ using ext3::Fe3; #define OP_NEG 9u #define OP_EMBED 10u +// -- operand kinds (mirror device.rs OPK_*): enc = kind << 29 | payload -- +#define OPK_SHIFT 29u +#define OPK_PAYLOAD_MASK 0x1FFFFFFFu +#define OPK_BASE_SLOT 0u +#define OPK_EXT_SLOT 1u +#define OPK_BASE_CONST 2u +#define OPK_EXT_CONST 3u +#define OPK_RAP 4u +#define OPK_ALPHA 5u +#define OPK_OFFSET 6u + +// -- res / root packing: bit 31 = ext slot class, low bits = slot index -- +#define RES_EXT_BIT 0x80000000u +#define RES_SLOT_MASK 0x7FFFFFFFu + // A flat IR node. Packed into two u64 words for a pure-u64 upload (matching the // crate's device-buffer convention): -// word0 = op | (a << 32) ; word1 = b | (dim << 32) -// This mirrors the `#[repr(C)] DeviceNode { op, a, b, dim: u32 }` payload. +// word0 = op | (a << 32) ; word1 = b | (res << 32) +// This mirrors the `#[repr(C)] DeviceNode { op, a, b, res: u32 }` payload. struct Node { - uint32_t op, a, b, dim; + uint32_t op, a, b, res; }; __device__ __forceinline__ Node load_node(const uint64_t *d_nodes, uint64_t i) { @@ -60,94 +84,231 @@ __device__ __forceinline__ Node load_node(const uint64_t *d_nodes, uint64_t i) { n.op = (uint32_t)(w0 & 0xFFFFFFFFull); n.a = (uint32_t)(w0 >> 32); n.b = (uint32_t)(w1 & 0xFFFFFFFFull); - n.dim = (uint32_t)(w1 >> 32); + n.res = (uint32_t)(w1 >> 32); return n; } +// The per-proof uniform tables an operand can reference directly. +struct Uniforms { + const uint64_t *base_consts; + const Fe3 *ext_consts; + const Fe3 *rap; + const Fe3 *alpha; + Fe3 offset; +}; + +// Whether an encoded operand holds a base-field value (slot or constant). +__device__ __forceinline__ bool opk_is_base(uint32_t enc) { + uint32_t kind = enc >> OPK_SHIFT; + return kind == OPK_BASE_SLOT || kind == OPK_BASE_CONST; +} + +// Load a base-field operand (kind must be a base kind). +__device__ __forceinline__ uint64_t load_base_operand(uint32_t enc, const uint64_t *vb, + uint64_t vstride, const Uniforms &u) { + uint32_t payload = enc & OPK_PAYLOAD_MASK; + return (enc >> OPK_SHIFT) == OPK_BASE_SLOT ? vb[(uint64_t)payload * vstride] + : u.base_consts[payload]; +} + +// Load any operand as ext3, embedding base values as {x, 0, 0}. +__device__ __forceinline__ Fe3 load_ext_operand(uint32_t enc, const uint64_t *vb, const uint64_t *ve, + uint64_t vstride, const Uniforms &u) { + uint32_t kind = enc >> OPK_SHIFT; + uint32_t payload = enc & OPK_PAYLOAD_MASK; + switch (kind) { + case OPK_BASE_SLOT: + return ext3::make(vb[(uint64_t)payload * vstride], 0, 0); + case OPK_EXT_SLOT: { + const uint64_t *p = ve + (uint64_t)payload * 3 * vstride; + return ext3::make(p[0], p[vstride], p[2 * vstride]); + } + case OPK_BASE_CONST: + return ext3::make(u.base_consts[payload], 0, 0); + case OPK_EXT_CONST: + return u.ext_consts[payload]; + case OPK_RAP: + return u.rap[payload]; + case OPK_ALPHA: + return u.alpha[payload]; + default: // OPK_OFFSET + return u.offset; + } +} + +__device__ __forceinline__ void store_base_slot(uint64_t *vb, uint64_t vstride, uint32_t slot, + uint64_t v) { + vb[(uint64_t)slot * vstride] = v; +} + +__device__ __forceinline__ void store_ext_slot(uint64_t *ve, uint64_t vstride, uint32_t slot, + const Fe3 &v) { + uint64_t *p = ve + (uint64_t)slot * 3 * vstride; + p[0] = v.a; + p[vstride] = v.b; + p[2 * vstride] = v.c; +} + +// Read a root value as ext3 (base roots embed as {x, 0, 0}). +__device__ __forceinline__ Fe3 load_root(uint64_t root_enc, const uint64_t *vb, const uint64_t *ve, + uint64_t vstride) { + uint32_t enc = (uint32_t)root_enc; + uint32_t slot = enc & RES_SLOT_MASK; + if (enc & RES_EXT_BIT) { + const uint64_t *p = ve + (uint64_t)slot * 3 * vstride; + return ext3::make(p[0], p[vstride], p[2 * vstride]); + } + return ext3::make(vb[(uint64_t)slot * vstride], 0, 0); +} + // Resolve an `Op::Var` leaf against the device-resident LDE columns. // a = col (low 16 bits); b = main<<16 | offset<<8 | row (see device.rs pack_var) // Base (main) columns are column-major `d_main[col*main_stride + r]`; ext (aux) // columns store component k at `d_aux[(col*3 + k)*aux_stride + r]` (GpuLdeExt3). // The frame `offset` selects row `r = (row + offset*next_step) mod num_rows`. -__device__ __forceinline__ Fe3 read_var(uint32_t a, uint32_t b, uint64_t row, uint64_t next_step, - uint64_t num_rows, const uint64_t *d_main, - uint64_t main_stride, const uint64_t *d_aux, - uint64_t aux_stride) { - uint32_t col = a & 0xFFFFu; - bool is_main = ((b >> 16) & 1u) != 0u; +__device__ __forceinline__ uint64_t var_row(uint32_t b, uint64_t row, uint64_t next_step, + uint64_t num_rows) { uint32_t offset = (b >> 8) & 0xFFu; - uint64_t r = row + (uint64_t)offset * next_step; if (r >= num_rows) { r -= num_rows; // wrap; offset*next_step < num_rows by construction } - - if (is_main) { - uint64_t x = d_main[(uint64_t)col * main_stride + r]; - return ext3::make(x, 0, 0); - } - uint64_t base = (uint64_t)col * 3; - uint64_t a0 = d_aux[(base + 0) * aux_stride + r]; - uint64_t a1 = d_aux[(base + 1) * aux_stride + r]; - uint64_t a2 = d_aux[(base + 2) * aux_stride + r]; - return ext3::make(a0, a1, a2); + return r; } // Shared forward pass: evaluate every IR node of the program for one LDE row -// into the per-thread value scratch (node `i`'s value at `vals[i * vstride]`; -// id `i` references only nodes `< i`). The single home of the op semantics — -// both kernels below run this exact walk, so an op change stays in lockstep -// with `constraint_ir/device.rs` in one place. +// into the per-thread slot scratch. The single home of the op semantics — both +// kernels below run this exact walk, so an op change stays in lockstep with +// `constraint_ir/device.rs` in one place. +// +// Every mixed-op shortcut below must be bit-identical to the full ext3 op on +// the embedded operand; see the file header for the argument. __device__ __forceinline__ void eval_program_row( - Fe3 *vals, uint64_t vstride, const uint64_t *d_nodes, uint64_t num_nodes, - const uint64_t *d_base_consts, const Fe3 *d_ext_consts, const Fe3 *d_rap_challenges, - const Fe3 *d_alpha_powers, Fe3 table_offset, uint64_t row, uint64_t next_step, - uint64_t num_rows, const uint64_t *d_main, uint64_t main_stride, const uint64_t *d_aux, - uint64_t aux_stride) { + uint64_t *vb, uint64_t *ve, uint64_t vstride, const uint64_t *d_nodes, uint64_t num_nodes, + const Uniforms &u, uint64_t row, uint64_t next_step, uint64_t num_rows, + const uint64_t *d_main, uint64_t main_stride, const uint64_t *d_aux, uint64_t aux_stride) { for (uint64_t i = 0; i < num_nodes; i++) { Node nd = load_node(d_nodes, i); - Fe3 v; + uint32_t slot = nd.res & RES_SLOT_MASK; + bool res_ext = (nd.res & RES_EXT_BIT) != 0; switch (nd.op) { - case OP_CONST_BASE: - v = ext3::make(d_base_consts[nd.a], 0, 0); + case OP_VAR: { + uint64_t r = var_row(nd.b, row, next_step, num_rows); + uint32_t col = nd.a & 0xFFFFu; + bool is_main = ((nd.b >> 16) & 1u) != 0u; + if (is_main) { + store_base_slot(vb, vstride, slot, d_main[(uint64_t)col * main_stride + r]); + } else { + uint64_t base = (uint64_t)col * 3; + store_ext_slot(ve, vstride, slot, + ext3::make(d_aux[(base + 0) * aux_stride + r], + d_aux[(base + 1) * aux_stride + r], + d_aux[(base + 2) * aux_stride + r])); + } break; - case OP_CONST_EXT: - v = d_ext_consts[nd.a]; + } + case OP_ADD: { + if (!res_ext) { + store_base_slot(vb, vstride, slot, + goldilocks::add(load_base_operand(nd.a, vb, vstride, u), + load_base_operand(nd.b, vb, vstride, u))); + } else if (opk_is_base(nd.a)) { + // {x,0,0} + y = {add(x,y.a), y.b, y.c} (add(0,v) == v). + uint64_t x = load_base_operand(nd.a, vb, vstride, u); + Fe3 y = load_ext_operand(nd.b, vb, ve, vstride, u); + store_ext_slot(ve, vstride, slot, ext3::make(goldilocks::add(x, y.a), y.b, y.c)); + } else if (opk_is_base(nd.b)) { + Fe3 x = load_ext_operand(nd.a, vb, ve, vstride, u); + uint64_t y = load_base_operand(nd.b, vb, vstride, u); + store_ext_slot(ve, vstride, slot, ext3::make(goldilocks::add(x.a, y), x.b, x.c)); + } else { + store_ext_slot(ve, vstride, slot, + ext3::add(load_ext_operand(nd.a, vb, ve, vstride, u), + load_ext_operand(nd.b, vb, ve, vstride, u))); + } break; - case OP_VAR: - v = read_var(nd.a, nd.b, row, next_step, num_rows, d_main, main_stride, d_aux, - aux_stride); + } + case OP_SUB: { + if (!res_ext) { + store_base_slot(vb, vstride, slot, + goldilocks::sub(load_base_operand(nd.a, vb, vstride, u), + load_base_operand(nd.b, vb, vstride, u))); + } else if (opk_is_base(nd.a)) { + // {x,0,0} - y = {sub(x,y.a), sub(0,y.b), sub(0,y.c)}; sub(0,·) + // is kept literal — it is NOT bitwise `neg` on non-canonical + // limbs. + uint64_t x = load_base_operand(nd.a, vb, vstride, u); + Fe3 y = load_ext_operand(nd.b, vb, ve, vstride, u); + store_ext_slot(ve, vstride, slot, + ext3::make(goldilocks::sub(x, y.a), goldilocks::sub(0, y.b), + goldilocks::sub(0, y.c))); + } else if (opk_is_base(nd.b)) { + // x - {y,0,0} = {sub(x.a,y), x.b, x.c} (sub(v,0) == v). + Fe3 x = load_ext_operand(nd.a, vb, ve, vstride, u); + uint64_t y = load_base_operand(nd.b, vb, vstride, u); + store_ext_slot(ve, vstride, slot, ext3::make(goldilocks::sub(x.a, y), x.b, x.c)); + } else { + store_ext_slot(ve, vstride, slot, + ext3::sub(load_ext_operand(nd.a, vb, ve, vstride, u), + load_ext_operand(nd.b, vb, ve, vstride, u))); + } break; - case OP_RAP_CHALLENGE: - v = d_rap_challenges[nd.a]; + } + case OP_MUL: { + if (!res_ext) { + store_base_slot(vb, vstride, slot, + goldilocks::mul(load_base_operand(nd.a, vb, vstride, u), + load_base_operand(nd.b, vb, vstride, u))); + } else if (opk_is_base(nd.a)) { + // {x,0,0} * y = mul_base(y, x): dot3 with zero products + // reduces to gl::mul exactly. + uint64_t x = load_base_operand(nd.a, vb, vstride, u); + Fe3 y = load_ext_operand(nd.b, vb, ve, vstride, u); + store_ext_slot(ve, vstride, slot, ext3::mul_base(y, x)); + } else if (opk_is_base(nd.b)) { + Fe3 x = load_ext_operand(nd.a, vb, ve, vstride, u); + uint64_t y = load_base_operand(nd.b, vb, vstride, u); + store_ext_slot(ve, vstride, slot, ext3::mul_base(x, y)); + } else { + store_ext_slot(ve, vstride, slot, + ext3::mul(load_ext_operand(nd.a, vb, ve, vstride, u), + load_ext_operand(nd.b, vb, ve, vstride, u))); + } break; - case OP_ALPHA_POW: - v = d_alpha_powers[nd.a]; + } + case OP_NEG: { + if (!res_ext) { + store_base_slot(vb, vstride, slot, + goldilocks::neg(load_base_operand(nd.a, vb, vstride, u))); + } else { + store_ext_slot(ve, vstride, slot, + ext3::neg(load_ext_operand(nd.a, vb, ve, vstride, u))); + } break; - case OP_TABLE_OFFSET: - v = table_offset; + } + case OP_EMBED: { + store_ext_slot(ve, vstride, slot, load_ext_operand(nd.a, vb, ve, vstride, u)); break; - case OP_ADD: - v = ext3::add(vals[(uint64_t)nd.a * vstride], vals[(uint64_t)nd.b * vstride]); + } + // Uniform leaves materialize only when they are constraint roots. + case OP_CONST_BASE: + store_base_slot(vb, vstride, slot, u.base_consts[nd.a]); break; - case OP_SUB: - v = ext3::sub(vals[(uint64_t)nd.a * vstride], vals[(uint64_t)nd.b * vstride]); + case OP_CONST_EXT: + store_ext_slot(ve, vstride, slot, u.ext_consts[nd.a]); break; - case OP_MUL: - v = ext3::mul(vals[(uint64_t)nd.a * vstride], vals[(uint64_t)nd.b * vstride]); + case OP_RAP_CHALLENGE: + store_ext_slot(ve, vstride, slot, u.rap[nd.a]); break; - case OP_NEG: - v = ext3::neg(vals[(uint64_t)nd.a * vstride]); + case OP_ALPHA_POW: + store_ext_slot(ve, vstride, slot, u.alpha[nd.a]); break; - case OP_EMBED: - // All-ext representation: the base operand is already {x,0,0}. - v = vals[(uint64_t)nd.a * vstride]; + case OP_TABLE_OFFSET: + store_ext_slot(ve, vstride, slot, u.offset); break; default: - v = ext3::zero(); break; } - vals[i * vstride] = v; } } @@ -159,7 +320,7 @@ extern "C" __global__ void constraint_interp_kernel( uint64_t num_nodes, const uint64_t *__restrict__ d_base_consts, const Fe3 *__restrict__ d_ext_consts, - const uint64_t *__restrict__ d_roots, + const uint64_t *__restrict__ d_roots, // slot | ext_bit<<31, one per constraint uint64_t num_roots, // per-proof uniforms const Fe3 *__restrict__ d_rap_challenges, @@ -173,24 +334,31 @@ extern "C" __global__ void constraint_interp_kernel( uint64_t next_step, // sizing uint64_t num_rows, - // scratch: per-thread value array, [num_nodes * num_threads] - Fe3 *__restrict__ d_values) { + // scratch: per-thread slot files, [num_base_slots * num_threads] and + // [num_ext_slots * 3 * num_threads] + uint64_t *__restrict__ d_vals_base, + uint64_t *__restrict__ d_vals_ext) { uint64_t task_offset = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; uint64_t num_threads = (uint64_t)gridDim.x * blockDim.x; - Fe3 *vals = d_values + task_offset; + uint64_t *vb = d_vals_base + task_offset; + uint64_t *ve = d_vals_ext + task_offset; uint64_t vstride = num_threads; - Fe3 table_offset = *d_table_offset; + + Uniforms u; + u.base_consts = d_base_consts; + u.ext_consts = d_ext_consts; + u.rap = d_rap_challenges; + u.alpha = d_alpha_powers; + u.offset = *d_table_offset; for (uint64_t row = task_offset; row < num_rows; row += num_threads) { - eval_program_row(vals, vstride, d_nodes, num_nodes, d_base_consts, d_ext_consts, - d_rap_challenges, d_alpha_powers, table_offset, row, next_step, num_rows, - d_main, main_stride, d_aux, aux_stride); + eval_program_row(vb, ve, vstride, d_nodes, num_nodes, u, row, next_step, num_rows, d_main, + main_stride, d_aux, aux_stride); - // Emit each constraint root. + // Emit each constraint root (base roots embed as {x, 0, 0}). for (uint64_t c = 0; c < num_roots; c++) { - uint64_t root = d_roots[c]; - d_evals[c * num_rows + row] = vals[root * vstride]; + d_evals[c * num_rows + row] = load_root(d_roots[c], vb, ve, vstride); } } } @@ -204,9 +372,10 @@ extern "C" __global__ void constraint_interp_kernel( // H(row) = z_inv[row % z_len] * Σ_c beta_trans[c] * C_c(row) (transition) // + Σ_b z_b_inv[b*num_rows + row] * beta_bnd[b] * (trace_b - value_b) // -// where a base-rooted C_c is carried as the embedding {C,0,0} (all-ext, exactly -// as the interpreter), z_inv is the cyclic base transition-zerofier inverse, and -// the boundary term reads the resident trace at column `b_col[b]` (main or aux). +// where a base-rooted C_c contributes via `mul_base` (bit-identical to the +// full mul on its embedding), z_inv is the cyclic base transition-zerofier +// inverse, and the boundary term reads the resident trace at column `b_col[b]` +// (main or aux). extern "C" __global__ void constraint_composition_kernel( // output: one H(row) per LDE row Fe3 *__restrict__ d_h, @@ -239,25 +408,40 @@ extern "C" __global__ void constraint_composition_kernel( const Fe3 *__restrict__ d_b_value, // [num_boundary] const Fe3 *__restrict__ d_b_beta, // [num_boundary] const uint64_t *__restrict__ d_b_z_inv, // [num_boundary * num_rows] - // scratch value array, [num_nodes * num_threads] - Fe3 *__restrict__ d_values) { + // scratch: per-thread slot files + uint64_t *__restrict__ d_vals_base, + uint64_t *__restrict__ d_vals_ext) { uint64_t task_offset = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; uint64_t num_threads = (uint64_t)gridDim.x * blockDim.x; - Fe3 *vals = d_values + task_offset; + uint64_t *vb = d_vals_base + task_offset; + uint64_t *ve = d_vals_ext + task_offset; uint64_t vstride = num_threads; - Fe3 table_offset = *d_table_offset; + + Uniforms u; + u.base_consts = d_base_consts; + u.ext_consts = d_ext_consts; + u.rap = d_rap_challenges; + u.alpha = d_alpha_powers; + u.offset = *d_table_offset; for (uint64_t row = task_offset; row < num_rows; row += num_threads) { - eval_program_row(vals, vstride, d_nodes, num_nodes, d_base_consts, d_ext_consts, - d_rap_challenges, d_alpha_powers, table_offset, row, next_step, num_rows, - d_main, main_stride, d_aux, aux_stride); + eval_program_row(vb, ve, vstride, d_nodes, num_nodes, u, row, next_step, num_rows, d_main, + main_stride, d_aux, aux_stride); - // Transition: z_inv * Σ_c beta_c * C_c. + // Transition: z_inv * Σ_c beta_c * C_c. Base roots use mul_base — + // bit-identical to mul(beta, {v,0,0}). Fe3 sum = ext3::zero(); for (uint64_t c = 0; c < num_roots; c++) { - Fe3 cval = vals[(uint64_t)d_roots[c] * vstride]; - sum = ext3::add(sum, ext3::mul(d_beta_trans[c], cval)); + uint32_t enc = (uint32_t)d_roots[c]; + uint32_t slot = enc & RES_SLOT_MASK; + if (enc & RES_EXT_BIT) { + const uint64_t *p = ve + (uint64_t)slot * 3 * vstride; + Fe3 cval = ext3::make(p[0], p[vstride], p[2 * vstride]); + sum = ext3::add(sum, ext3::mul(d_beta_trans[c], cval)); + } else { + sum = ext3::add(sum, ext3::mul_base(d_beta_trans[c], vb[(uint64_t)slot * vstride])); + } } Fe3 h = ext3::mul_base(sum, d_z_inv[row % z_len]); diff --git a/crypto/math-cuda/kernels/keccak.cu b/crypto/math-cuda/kernels/keccak.cu index e7bb8a618..7b62789f9 100644 --- a/crypto/math-cuda/kernels/keccak.cu +++ b/crypto/math-cuda/kernels/keccak.cu @@ -474,3 +474,41 @@ extern "C" __global__ void keccak256_leaves_base_row_major_row_pair( } finalize_keccak256(st, rate_pos, hashed_leaves_out + tid * 32); } + +// Column-range variant of `keccak256_leaves_base_row_major_row_pair`: each leaf +// hashes only columns `[col_start, col_end)` of the two bit-reversed rows, +// while `m` remains the full row stride. Byte layout equals the CPU +// `commit_rows_bit_reversed_subset(data, m, col_start, col_end)` — used for +// preprocessed tables, whose precomputed and multiplicity column ranges commit +// to separate Merkle trees over the same row-major LDE. +extern "C" __global__ void keccak256_leaves_base_row_major_row_pair_range( + const uint64_t *data, + uint64_t m, + uint64_t col_start, + uint64_t col_end, + uint64_t num_rows, + uint64_t log_num_rows, + uint8_t *hashed_leaves_out) +{ + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t num_leaves = num_rows >> 1; + if (tid >= num_leaves) return; + + uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows); + uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows); + const uint64_t *row_0 = data + br_0 * m; + const uint64_t *row_1 = data + br_1 * m; + + uint64_t st[25]; + #pragma unroll + for (int i = 0; i < 25; ++i) st[i] = 0; + + uint32_t rate_pos = 0; + for (uint64_t c = col_start; c < col_end; ++c) { + absorb_lane(st, rate_pos, bswap64(goldilocks::canonical(row_0[c]))); + } + for (uint64_t c = col_start; c < col_end; ++c) { + absorb_lane(st, rate_pos, bswap64(goldilocks::canonical(row_1[c]))); + } + finalize_keccak256(st, rate_pos, hashed_leaves_out + tid * 32); +} diff --git a/crypto/math-cuda/src/barycentric.rs b/crypto/math-cuda/src/barycentric.rs index 8fdea14f1..252b451a1 100644 --- a/crypto/math-cuda/src/barycentric.rs +++ b/crypto/math-cuda/src/barycentric.rs @@ -45,9 +45,13 @@ pub fn barycentric_base( let be = backend()?; let stream = be.next_stream(); - let cols_dev = stream.clone_htod(&columns[..num_cols * col_stride])?; - let points_dev = stream.clone_htod(coset_points)?; - let inv_dev = stream.clone_htod(inv_denoms_ext3)?; + let (cols_dev, points_dev, inv_dev) = { + ( + stream.clone_htod(&columns[..num_cols * col_stride])?, + stream.clone_htod(coset_points)?, + stream.clone_htod(inv_denoms_ext3)?, + ) + }; let mut out_dev = stream.alloc_zeros::(3 * num_cols)?; let col_stride_u64 = col_stride as u64; @@ -68,8 +72,10 @@ pub fn barycentric_base( .arg(&mut out_dev) .launch(cfg)?; } - let out = stream.clone_dtoh(&out_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&out_dev)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -100,9 +106,13 @@ pub fn barycentric_ext3( let be = backend()?; let stream = be.next_stream(); - let cols_dev = stream.clone_htod(&columns[..num_cols * 3 * col_stride])?; - let points_dev = stream.clone_htod(coset_points)?; - let inv_dev = stream.clone_htod(inv_denoms_ext3)?; + let (cols_dev, points_dev, inv_dev) = { + ( + stream.clone_htod(&columns[..num_cols * 3 * col_stride])?, + stream.clone_htod(coset_points)?, + stream.clone_htod(inv_denoms_ext3)?, + ) + }; let mut out_dev = stream.alloc_zeros::(3 * num_cols)?; let col_stride_u64 = col_stride as u64; @@ -123,8 +133,10 @@ pub fn barycentric_ext3( .arg(&mut out_dev) .launch(cfg)?; } - let out = stream.clone_dtoh(&out_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&out_dev)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -149,9 +161,14 @@ pub fn barycentric_base_on_device( let be = backend()?; let stream = be.next_stream(); + main_handle.wait_ready_on(&stream)?; - let points_dev = stream.clone_htod(coset_points)?; - let inv_dev = stream.clone_htod(inv_denoms_ext3)?; + let (points_dev, inv_dev) = { + ( + stream.clone_htod(coset_points)?, + stream.clone_htod(inv_denoms_ext3)?, + ) + }; let mut out_dev = stream.alloc_zeros::(3 * num_cols)?; let col_stride_u64 = col_stride as u64; @@ -174,8 +191,10 @@ pub fn barycentric_base_on_device( .arg(&mut out_dev) .launch(cfg)?; } - let out = stream.clone_dtoh(&out_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&out_dev)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -197,6 +216,7 @@ pub fn barycentric_base_on_device_with_dev_inv_denoms( inv_offset_u64: usize, n: usize, ) -> Result> { + main_handle.wait_ready_on(stream)?; assert!(coset_points_dev.len() >= n); let inv_end = inv_offset_u64 .checked_add(3 * n) @@ -233,8 +253,10 @@ pub fn barycentric_base_on_device_with_dev_inv_denoms( .arg(&mut out_dev) .launch(cfg)?; } - let out = stream.clone_dtoh(&out_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&out_dev)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -257,9 +279,14 @@ pub fn barycentric_ext3_on_device( let be = backend()?; let stream = be.next_stream(); + aux_handle.wait_ready_on(&stream)?; - let points_dev = stream.clone_htod(coset_points)?; - let inv_dev = stream.clone_htod(inv_denoms_ext3)?; + let (points_dev, inv_dev) = { + ( + stream.clone_htod(coset_points)?, + stream.clone_htod(inv_denoms_ext3)?, + ) + }; let mut out_dev = stream.alloc_zeros::(3 * num_cols)?; let col_stride_u64 = col_stride as u64; @@ -282,8 +309,10 @@ pub fn barycentric_ext3_on_device( .arg(&mut out_dev) .launch(cfg)?; } - let out = stream.clone_dtoh(&out_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&out_dev)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -297,6 +326,7 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms( inv_offset_u64: usize, n: usize, ) -> Result> { + aux_handle.wait_ready_on(stream)?; assert!(coset_points_dev.len() >= n); let inv_end = inv_offset_u64 .checked_add(3 * n) @@ -333,8 +363,10 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms( .arg(&mut out_dev) .launch(cfg)?; } - let out = stream.clone_dtoh(&out_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&out_dev)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -347,12 +379,13 @@ pub fn gather_rows_base_on_device( rows: &[u32], stream: &Arc, ) -> Result> { + main.wait_ready_on(stream)?; let num_cols = main.m; if num_cols == 0 || rows.is_empty() { return Ok(Vec::new()); } let be = backend()?; - let rows_dev = stream.clone_htod(rows)?; + let rows_dev = { stream.clone_htod(rows)? }; let mut out = stream.alloc_zeros::(rows.len() * num_cols)?; let col_stride = main.lde_size as u64; let num_cols_u64 = num_cols as u64; @@ -373,8 +406,10 @@ pub fn gather_rows_base_on_device( .arg(&mut out) .launch(cfg)?; } - let host = stream.clone_dtoh(&out)?; - stream.synchronize()?; + let host = { stream.clone_dtoh(&out)? }; + { + stream.synchronize()?; + } Ok(host) } @@ -385,12 +420,13 @@ pub fn gather_rows_ext3_on_device( rows: &[u32], stream: &Arc, ) -> Result> { + aux.wait_ready_on(stream)?; let num_cols = aux.m; if num_cols == 0 || rows.is_empty() { return Ok(Vec::new()); } let be = backend()?; - let rows_dev = stream.clone_htod(rows)?; + let rows_dev = { stream.clone_htod(rows)? }; let mut out = stream.alloc_zeros::(rows.len() * num_cols * 3)?; let col_stride = aux.lde_size as u64; let num_cols_u64 = num_cols as u64; @@ -411,7 +447,9 @@ pub fn gather_rows_ext3_on_device( .arg(&mut out) .launch(cfg)?; } - let host = stream.clone_dtoh(&out)?; - stream.synchronize()?; + let host = { stream.clone_dtoh(&out)? }; + { + stream.synchronize()?; + } Ok(host) } diff --git a/crypto/math-cuda/src/constraint_interp.rs b/crypto/math-cuda/src/constraint_interp.rs index ba11b4f64..6d956fc65 100644 --- a/crypto/math-cuda/src/constraint_interp.rs +++ b/crypto/math-cuda/src/constraint_interp.rs @@ -6,6 +6,13 @@ //! handles, uploads the program + per-proof uniforms, launches the interpreter //! over every LDE row, and returns the per-constraint eval matrix. //! +//! The lowering dim-splits the per-thread value scratch into a base (`u64`) +//! and an ext (`3 × u64`) slot class with liveness-reused slots, so the +//! scratch here is sized by the program's max-live-set +//! (`num_base_slots`/`num_ext_slots`), not its node count. Both buffers are +//! allocated uninitialized: the topological walk writes every slot before any +//! read. +//! //! Layering note: this crate cannot see `stark`'s `DeviceProgram` type (stark //! depends on math-cuda, not the reverse), so the caller flattens the program //! into the raw `u64` slices below. The stark-side dispatch @@ -18,9 +25,9 @@ use crate::device::backend; use crate::lde::{GpuLdeBase, GpuLdeExt3}; const BLOCK_DIM: u32 = 256; -/// Cap on total threads (grid × block). Each thread owns `num_nodes` ext3 slots -/// in the global value scratch (`num_nodes × MAX_THREADS × 24 B`), so a fixed -/// cap bounds that buffer regardless of LDE size; threads grid-stride over the +/// Cap on total threads (grid × block). Each thread owns `num_base_slots` u64 +/// plus `num_ext_slots` ext3 slots of global value scratch, so a fixed cap +/// bounds those buffers regardless of LDE size; threads grid-stride over the /// remaining rows. 65536 mirrors OpenVM's quotient `TASK_SIZE`. const MAX_THREADS: u32 = 1 << 16; @@ -31,17 +38,21 @@ const MAX_THREADS: u32 = 1 << 16; /// Base-rooted constraints carry their value in component 0. /// /// Inputs (all raw limbs, matching the crate's u64 device convention): -/// - `nodes`: 2 `u64` per IR node (`op | a<<32`, then `b | dim<<32`). +/// - `nodes`: 2 `u64` per IR node (`op | a<<32`, then `b | res<<32`). +/// - `num_base_slots` / `num_ext_slots`: per-thread scratch sizes of the two +/// slot classes (from the lowering's liveness scan). /// - `base_consts`: one `u64` per base constant. /// - `ext_consts`, `rap_challenges`, `alpha_powers`: 3 `u64` per element. /// - `table_offset`: exactly 3 `u64`. -/// - `roots`: one `u64` node id per constraint. +/// - `roots`: one `u64` per constraint (`slot | ext_bit<<31`). /// - `main`/`aux`: device-resident LDE handles; `next_step` is the LDE row /// stride for a frame-offset step; `num_rows` is the number of LDE rows. #[allow(clippy::too_many_arguments)] pub fn eval_constraints_on_device( nodes: &[u64], num_nodes: usize, + num_base_slots: usize, + num_ext_slots: usize, base_consts: &[u64], ext_consts: &[u64], roots: &[u64], @@ -62,25 +73,33 @@ pub fn eval_constraints_on_device( let be = backend()?; let stream = be.next_stream(); + main.wait_ready_on(&stream)?; + aux.wait_ready_on(&stream)?; // Upload the program + uniforms (the column data never crosses PCIe — it is // already resident in `main.buf` / `aux.buf`). - let d_nodes = stream.clone_htod(nodes)?; - let d_base_consts = stream.clone_htod(base_consts)?; - let d_ext_consts = stream.clone_htod(ext_consts)?; - let d_roots = stream.clone_htod(roots)?; - let d_rap = stream.clone_htod(rap_challenges)?; - let d_alpha = stream.clone_htod(alpha_powers)?; - let d_offset = stream.clone_htod(table_offset)?; + let (d_nodes, d_base_consts, d_ext_consts, d_roots, d_rap, d_alpha, d_offset) = { + ( + stream.clone_htod(nodes)?, + stream.clone_htod(base_consts)?, + stream.clone_htod(ext_consts)?, + stream.clone_htod(roots)?, + stream.clone_htod(rap_challenges)?, + stream.clone_htod(alpha_powers)?, + stream.clone_htod(table_offset)?, + ) + }; // Fixed thread count, grid-stride over rows. let max_grid = MAX_THREADS / BLOCK_DIM; let grid = (num_rows as u32).div_ceil(BLOCK_DIM).clamp(1, max_grid); let num_threads = (grid as usize) * (BLOCK_DIM as usize); - // Per-thread value scratch ([num_nodes * num_threads] ext3) and output. - let mut d_values = stream.alloc_zeros::(num_nodes * num_threads * 3)?; - let mut d_evals = stream.alloc_zeros::(num_roots * num_rows * 3)?; + // Per-thread slot scratch, uninitialized (the walk writes before reading). + let mut d_vals_base = unsafe { stream.alloc::((num_base_slots * num_threads).max(1)) }?; + let mut d_vals_ext = unsafe { stream.alloc::((num_ext_slots * 3 * num_threads).max(1)) }?; + // Output: every (constraint, row) cell is written by the emit loop. + let mut d_evals = unsafe { stream.alloc::(num_roots * num_rows * 3) }?; let num_nodes_u64 = num_nodes as u64; let num_roots_u64 = num_roots as u64; @@ -113,11 +132,22 @@ pub fn eval_constraints_on_device( .arg(&aux_stride) .arg(&next_step_u64) .arg(&num_rows_u64) - .arg(&mut d_values) + .arg(&mut d_vals_base) + .arg(&mut d_vals_ext) .launch(cfg)?; } - let out = stream.clone_dtoh(&d_evals)?; - stream.synchronize()?; + let out = { + let pending = crate::device::async_dtoh_via( + &stream, + be.pinned_staging(), + &be.ctx, + &d_evals, + d_evals.len(), + )?; + let mut out = vec![0u64; d_evals.len()]; + pending.wait_into_u64(&mut out)?; + out + }; Ok(out) } @@ -157,6 +187,8 @@ pub struct CompositionAccum<'a> { pub fn eval_composition_on_device( nodes: &[u64], num_nodes: usize, + num_base_slots: usize, + num_ext_slots: usize, base_consts: &[u64], ext_consts: &[u64], roots: &[u64], @@ -197,35 +229,51 @@ pub fn eval_composition_on_device( let be = backend()?; let stream = be.next_stream(); + main.wait_ready_on(&stream)?; + aux.wait_ready_on(&stream)?; - let d_nodes = stream.clone_htod(nodes)?; - let d_base_consts = stream.clone_htod(base_consts)?; - let d_ext_consts = stream.clone_htod(ext_consts)?; - let d_roots = stream.clone_htod(roots)?; - let d_rap = stream.clone_htod(rap_challenges)?; - let d_alpha = stream.clone_htod(alpha_powers)?; - let d_offset = stream.clone_htod(table_offset)?; + let (d_nodes, d_base_consts, d_ext_consts, d_roots, d_rap, d_alpha, d_offset) = { + ( + stream.clone_htod(nodes)?, + stream.clone_htod(base_consts)?, + stream.clone_htod(ext_consts)?, + stream.clone_htod(roots)?, + stream.clone_htod(rap_challenges)?, + stream.clone_htod(alpha_powers)?, + stream.clone_htod(table_offset)?, + ) + }; - let d_beta_trans = stream.clone_htod(accum.beta_trans)?; - let d_z_inv = stream.clone_htod(accum.z_inv)?; - let d_b_col = stream.clone_htod(accum.b_col)?; - let d_b_is_aux = stream.clone_htod(accum.b_is_aux)?; - let d_b_value = stream.clone_htod(accum.b_value)?; - let d_b_beta = stream.clone_htod(accum.b_beta)?; + let (d_beta_trans, d_z_inv, d_b_col, d_b_is_aux, d_b_value, d_b_beta) = { + ( + stream.clone_htod(accum.beta_trans)?, + stream.clone_htod(accum.z_inv)?, + stream.clone_htod(accum.b_col)?, + stream.clone_htod(accum.b_is_aux)?, + stream.clone_htod(accum.b_value)?, + stream.clone_htod(accum.b_beta)?, + ) + }; // Per-slice upload straight from the caller's per-constraint vectors into - // the flat `b * num_rows + row` device layout — no flattened host copy. - let mut d_b_z_inv = stream.alloc_zeros::((num_boundary * num_rows).max(1))?; - for (b, slice) in accum.b_z_inv.iter().enumerate() { - let mut dst = d_b_z_inv.slice_mut(b * num_rows..(b + 1) * num_rows); - stream.memcpy_htod(*slice, &mut dst)?; + // the flat `b * num_rows + row` device layout — no flattened host copy, no + // zeroing (the slice uploads cover every element the kernel reads). + let mut d_b_z_inv = unsafe { stream.alloc::((num_boundary * num_rows).max(1)) }?; + { + for (b, slice) in accum.b_z_inv.iter().enumerate() { + let mut dst = d_b_z_inv.slice_mut(b * num_rows..(b + 1) * num_rows); + stream.memcpy_htod(*slice, &mut dst)?; + } } let max_grid = MAX_THREADS / BLOCK_DIM; let grid = (num_rows as u32).div_ceil(BLOCK_DIM).clamp(1, max_grid); let num_threads = (grid as usize) * (BLOCK_DIM as usize); - let mut d_values = stream.alloc_zeros::(num_nodes * num_threads * 3)?; - let mut d_h = stream.alloc_zeros::(num_rows * 3)?; + // Per-thread slot scratch, uninitialized (the walk writes before reading). + let mut d_vals_base = unsafe { stream.alloc::((num_base_slots * num_threads).max(1)) }?; + let mut d_vals_ext = unsafe { stream.alloc::((num_ext_slots * 3 * num_threads).max(1)) }?; + // Output: every row is written by the grid-stride loop. + let mut d_h = unsafe { stream.alloc::(num_rows * 3) }?; let num_nodes_u64 = num_nodes as u64; let num_roots_u64 = num_roots as u64; @@ -269,10 +317,16 @@ pub fn eval_composition_on_device( .arg(&d_b_value) .arg(&d_b_beta) .arg(&d_b_z_inv) - .arg(&mut d_values) + .arg(&mut d_vals_base) + .arg(&mut d_vals_ext) .launch(cfg)?; } - let out = stream.clone_dtoh(&d_h)?; - stream.synchronize()?; + let out = { + let pending = + crate::device::async_dtoh_via(&stream, be.pinned_staging(), &be.ctx, &d_h, d_h.len())?; + let mut out = vec![0u64; d_h.len()]; + pending.wait_into_u64(&mut out)?; + out + }; Ok(out) } diff --git a/crypto/math-cuda/src/deep.rs b/crypto/math-cuda/src/deep.rs index 581fbc404..6890f1a6b 100644 --- a/crypto/math-cuda/src/deep.rs +++ b/crypto/math-cuda/src/deep.rs @@ -138,6 +138,11 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms( row_stride: usize, domain_size: usize, ) -> Result> { + main_lde.wait_ready_on(stream)?; + if let Some(aux) = aux_lde { + aux.wait_ready_on(stream)?; + } + h_parts_dev.wait_ready_on(stream)?; assert_eq!(main_lde.m, num_main); assert_eq!(h_parts_dev.m, num_parts); assert_eq!(h_parts_dev.lde_size, main_lde.lde_size); @@ -175,10 +180,14 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms( let be = backend()?; // H2D only the small scalars on the caller's stream. - let h_ood_dev = stream.clone_htod(h_ood)?; - let trace_ood_dev = stream.clone_htod(trace_ood)?; - let gammas_h_dev = stream.clone_htod(gammas_h)?; - let gammas_tr_dev = stream.clone_htod(gammas_tr)?; + let (h_ood_dev, trace_ood_dev, gammas_h_dev, gammas_tr_dev) = { + ( + stream.clone_htod(h_ood)?, + stream.clone_htod(trace_ood)?, + stream.clone_htod(gammas_h)?, + stream.clone_htod(gammas_tr)?, + ) + }; // Slice the inv_denoms buffer into the H-term and trace-term views. let inv_h_view = inv_denoms_dev.slice(0..ext3_size); @@ -232,8 +241,24 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms( .launch(cfg)?; } - let out = stream.clone_dtoh(&deep_out)?; - stream.synchronize()?; + // DEEP output (domain_size * 3 u64s, ~50 MB): async D2H through the + // per-worker pinned slab instead of a blocking pageable copy. The + // labelled sync keeps its host-block measurement (now covering the + // kernels plus the DMA); the pending wait after it is instant. + let pending = { + crate::device::async_dtoh_via( + stream, + be.pinned_staging(), + &be.ctx, + &deep_out, + domain_size * 3, + )? + }; + { + stream.synchronize()?; + } + let mut out = vec![0u64; domain_size * 3]; + pending.wait_into_u64(&mut out)?; Ok(out) } @@ -257,6 +282,13 @@ fn deep_composition_ext3_impl( row_stride: usize, domain_size: usize, ) -> Result> { + main_lde.wait_ready_on(stream)?; + if let Some(aux) = aux_lde { + aux.wait_ready_on(stream)?; + } + if let Some(parts) = h_parts_dev { + parts.wait_ready_on(stream)?; + } assert_eq!(main_lde.m, num_main); if let Some(a) = aux_lde { assert_eq!(a.m, num_aux); @@ -293,12 +325,16 @@ fn deep_composition_ext3_impl( let be = backend()?; - let h_ood_dev = stream.clone_htod(h_ood)?; - let trace_ood_dev = stream.clone_htod(trace_ood)?; - let gammas_h_dev = stream.clone_htod(gammas_h)?; - let gammas_tr_dev = stream.clone_htod(gammas_tr)?; - let inv_h_dev = stream.clone_htod(inv_h)?; - let inv_t_dev = stream.clone_htod(inv_t)?; + let (h_ood_dev, trace_ood_dev, gammas_h_dev, gammas_tr_dev, inv_h_dev, inv_t_dev) = { + ( + stream.clone_htod(h_ood)?, + stream.clone_htod(trace_ood)?, + stream.clone_htod(gammas_h)?, + stream.clone_htod(gammas_tr)?, + stream.clone_htod(inv_h)?, + stream.clone_htod(inv_t)?, + ) + }; let h_lde_host_dev; let dummy_aux; @@ -358,7 +394,23 @@ fn deep_composition_ext3_impl( .launch(cfg)?; } - let out = stream.clone_dtoh(&deep_out)?; - stream.synchronize()?; + // DEEP output (domain_size * 3 u64s, ~50 MB): async D2H through the + // per-worker pinned slab instead of a blocking pageable copy. The + // labelled sync keeps its host-block measurement (now covering the + // kernels plus the DMA); the pending wait after it is instant. + let pending = { + crate::device::async_dtoh_via( + stream, + be.pinned_staging(), + &be.ctx, + &deep_out, + domain_size * 3, + )? + }; + { + stream.synchronize()?; + } + let mut out = vec![0u64; domain_size * 3]; + pending.wait_into_u64(&mut out)?; Ok(out) } diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 2bebe2cc0..9f6452689 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -25,6 +25,16 @@ use crate::ntt::{twiddles_forward, twiddles_inverse}; pub struct PinnedStaging { ptr: *mut u64, capacity_elems: usize, + /// Pool-wide first-allocation size hint (u64 elems), shared by every slot + /// of the pool this staging belongs to. Published by the prover once table + /// sizes are known (see [`Backend::set_staging_size_hints`]); zero = none. + hint_u64s: Arc, + /// Reusable completion event for [`async_dtoh_via`] copies through this + /// slot. Created once on first use and re-recorded per drain — per-call + /// cuEventCreate/Destroy measurably convoys the driver lock under load. + /// At most one drain per slot is in flight (the pending holds the slot + /// mutex), so a single event can never be aliased. + event: Option, } // SAFETY: the raw pointer aliases host memory allocated via cuMemHostAlloc. @@ -34,10 +44,12 @@ unsafe impl Send for PinnedStaging {} unsafe impl Sync for PinnedStaging {} impl PinnedStaging { - const fn empty() -> Self { + fn empty(hint_u64s: Arc) -> Self { Self { ptr: std::ptr::null_mut(), capacity_elems: 0, + hint_u64s, + event: None, } } @@ -55,7 +67,12 @@ impl PinnedStaging { self.ptr = std::ptr::null_mut(); self.capacity_elems = 0; } - let new_cap = min_elems.next_power_of_two().max(1 << 20); // at least 8 MB + // First allocation of this slot jumps straight to the prove-wide size + // hint (when the prover published one), so a slot pays cuMemHostAlloc + // once instead of climbing a realloc ladder (pinned alloc/free is + // ~50ms+ per step — see docs/gpu_baseline_ethrex_5090.md). + let target = min_elems.max(self.hint_u64s.load(Ordering::Relaxed)); + let new_cap = target.next_power_of_two().max(1 << 20); // at least 8 MB let bytes = new_cap * std::mem::size_of::(); let ptr = unsafe { cudarc::driver::result::malloc_host(bytes, 0 /* flags: non-WC */)? @@ -65,6 +82,29 @@ impl PinnedStaging { Ok(()) } + /// Record the slot's reusable event on `stream` (creating it on first use; + /// normally pre-created at backend init so no mid-prove cuEventCreate). + /// Pairs with [`PinnedStaging::sync_event`]; also re-recorded by + /// [`async_dtoh_via`] — safe because slot access is serialized by its + /// mutex, so a recorded event is always synchronized before re-recording. + pub fn record_event(&mut self, stream: &Arc) -> Result<()> { + match self.event.as_ref() { + Some(ev) => ev.record(stream), + None => { + self.event = Some(stream.record_event(None)?); + Ok(()) + } + } + } + + /// Block until the last [`PinnedStaging::record_event`] point completes. + pub fn sync_event(&self) -> Result<()> { + match self.event.as_ref() { + Some(ev) => ev.synchronize(), + None => Ok(()), + } + } + /// View of the first `len` elements. Caller must hold this `PinnedStaging` /// locked while using the slice; the slice aliases the internal pointer. /// @@ -124,7 +164,13 @@ pub struct Backend { /// `pinned_staging`; sized `num_rows * 32` bytes per slot. Lives /// alongside the LDE staging so the GPU→host D2H runs at PCIe line-rate. pinned_hashes: Vec>, + /// First-allocation size hints for the two staging pools (u64 elems), + /// shared with every slot. See [`Backend::set_staging_size_hints`]. + staging_hint_u64s: Arc, + hashes_hint_u64s: Arc, util_stream: Arc, + /// Free-list of pre-created events for [`Backend::take_event`]. + event_pool: Mutex>, next: AtomicUsize, /// VRAM budget (bytes) for table-session admission control. See /// [`detect_vram_budget_bytes`]. @@ -159,6 +205,7 @@ pub struct Backend { // keccak.cubin pub keccak256_leaves_base_row_major_row_pair: CudaFunction, + pub keccak256_leaves_base_row_major_row_pair_range: CudaFunction, pub keccak256_leaves_base_batched: CudaFunction, pub keccak256_leaves_base_row_pair_batched: CudaFunction, pub keccak256_leaves_ext3_batched: CudaFunction, @@ -319,12 +366,32 @@ impl Backend { // when no custom pool is in use. Stable across the backend's lifetime // since rayon's pool is fixed at first use. let n_slots = rayon::current_num_threads().max(1); - let pinned_staging: Vec> = (0..n_slots) - .map(|_| Mutex::new(PinnedStaging::empty())) - .collect(); - let pinned_hashes: Vec> = (0..n_slots) - .map(|_| Mutex::new(PinnedStaging::empty())) - .collect(); + let staging_hint = Arc::new(AtomicUsize::new(0)); + let hashes_hint = Arc::new(AtomicUsize::new(0)); + // Pre-create each slot's reusable event here, off the prove's critical + // path — a mid-prove cuEventCreate convoys the driver lock (~30 ms + // measured under load vs ~µs at init). + let make_pool = |hint: &Arc| -> Result>> { + let mut pool = Vec::with_capacity(n_slots); + for _ in 0..n_slots { + let mut slot = PinnedStaging::empty(Arc::clone(hint)); + slot.event = Some(ctx.new_event(None)?); + pool.push(Mutex::new(slot)); + } + Ok(pool) + }; + let pinned_staging = make_pool(&staging_hint)?; + let pinned_hashes = make_pool(&hashes_hint)?; + // Pre-create the handle-readiness event pool (see `take_event`): one + // event per device-resident handle a prove can have alive; creation + // here is ~µs each, mid-prove it convoys the driver lock. + let event_pool = { + let mut pool = Vec::with_capacity(512); + for _ in 0..512 { + pool.push(ctx.new_event(None)?); + } + Mutex::new(pool) + }; // Separate "utility" stream for twiddle uploads and other bookkeeping; // not part of the pool that callers rotate through. let util_stream = ctx.new_stream()?; @@ -361,6 +428,8 @@ impl Backend { matrix_transpose_strided: ntt.load_function("matrix_transpose_strided")?, keccak256_leaves_base_row_major_row_pair: keccak .load_function("keccak256_leaves_base_row_major_row_pair")?, + keccak256_leaves_base_row_major_row_pair_range: keccak + .load_function("keccak256_leaves_base_row_major_row_pair_range")?, keccak256_leaves_base_batched: keccak.load_function("keccak256_leaves_base_batched")?, keccak256_leaves_base_row_pair_batched: keccak .load_function("keccak256_leaves_base_row_pair_batched")?, @@ -405,6 +474,9 @@ impl Backend { streams, pinned_staging, pinned_hashes, + event_pool, + staging_hint_u64s: staging_hint, + hashes_hint_u64s: hashes_hint, util_stream, next: AtomicUsize::new(0), vram_budget_bytes, @@ -417,6 +489,19 @@ impl Backend { self.vram_budget_bytes } + /// Publish first-allocation size hints (u64 elems) for the pinned staging + /// pools, so each worker slot allocates its slab once at final size + /// instead of climbing a realloc ladder (each pinned alloc/free step costs + /// ~50ms+). Call once per prove, as soon as table sizes are known; only + /// raises (never shrinks) the current hints, and only affects slots that + /// have not allocated yet. + pub fn set_staging_size_hints(&self, lde_u64s: usize, hashes_u64s: usize) { + self.staging_hint_u64s + .fetch_max(lde_u64s, Ordering::Relaxed); + self.hashes_hint_u64s + .fetch_max(hashes_u64s, Ordering::Relaxed); + } + /// Round-robin over the stream pool. Concurrent callers get different /// streams so their kernel launches overlap on the GPU. pub fn next_stream(&self) -> Arc { @@ -533,3 +618,145 @@ pub fn backend() -> Result<&'static Backend> { let _ = BACKEND.set(b); Ok(BACKEND.get().expect("backend just initialised")) } + +// ── Asynchronous D2H through the pinned staging slabs ──────────────────────── + +/// A device→host copy enqueued into a per-worker pinned staging slab, not yet +/// awaited. Created by [`async_dtoh_via`]; consumed by one of the `wait_*` +/// methods, which block only until the copy (and everything queued before it +/// on its stream) lands. +/// +/// Holding this value keeps the staging slot's mutex locked, which is what +/// makes the whole scheme safe: no other caller (and no capacity growth) can +/// touch the slab while the DMA is in flight. +pub struct PendingD2H<'a> { + staging: std::sync::MutexGuard<'a, PinnedStaging>, + n_bytes: usize, +} + +/// Enqueue an async D2H of `n_elems` of `src` into the pinned slab of `slot`, +/// without synchronizing the stream. Unlike `stream.memcpy_dtoh` into a plain +/// (pageable) slice — which the driver services synchronously — this returns +/// as soon as the copy is queued; the returned [`PendingD2H`] is awaited at +/// the point the host actually needs the bytes. +/// +/// SAFETY contract (upheld by construction for our callers): `src` must stay +/// alive until the copy completes. Dropping a `CudaSlice` frees it +/// stream-ordered on its own stream, so a `src` allocated on `stream` may be +/// dropped after this call — the free queues behind the copy. Do NOT pass a +/// `src` owned by a *different* stream and drop it before waiting. +pub fn async_dtoh_via<'a, T: cudarc::driver::DeviceRepr>( + stream: &Arc, + slot: &'a Mutex, + ctx: &CudaContext, + src: &CudaSlice, + n_elems: usize, +) -> Result> { + use cudarc::driver::DevicePtr; + assert!(n_elems <= src.len()); + let n_bytes = n_elems * std::mem::size_of::(); + let u64_len = n_bytes.div_ceil(8); + let mut staging = slot.lock().unwrap(); + staging.ensure_capacity(u64_len, ctx)?; + ctx.bind_to_thread()?; + // SAFETY: dst is this slot's pinned allocation — stable address (only + // `ensure_capacity` moves it, and we hold the lock), pinned (registered + // via cuMemHostAlloc, so the driver DMAs directly, asynchronously). + // `device_ptr` orders the read after prior writes on `stream`. + unsafe { + let (src_ptr, _record) = src.device_ptr(stream); + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + staging.ptr as *mut core::ffi::c_void, + src_ptr, + n_bytes, + stream.cu_stream(), + ) + .result()?; + } + // Re-record the slot's reusable event (created once — per-call + // cuEventCreate/Destroy convoys the driver lock under load). + staging.record_event(stream)?; + Ok(PendingD2H { staging, n_bytes }) +} + +impl PendingD2H<'_> { + /// Number of bytes the copy deposits. + pub fn len_bytes(&self) -> usize { + self.n_bytes + } + + /// Block until the copy lands, then read the pinned bytes through `f`. + /// Consumes the pending (releasing the staging slot when `f` returns). + pub fn wait_and_read(self, f: impl FnOnce(&[u8]) -> R) -> Result { + self.staging + .event + .as_ref() + .expect("recorded by async_dtoh_via") + .synchronize()?; + // SAFETY: event completion orders the DMA before this read; the slab + // is exclusively ours while the guard lives. + let bytes = + unsafe { std::slice::from_raw_parts(self.staging.ptr as *const u8, self.n_bytes) }; + Ok(f(bytes)) + } + + /// Wait and copy the bytes out into `dst` (pageable is fine — this is a + /// plain host memcpy at RAM speed, not a DMA target). + pub fn wait_into_bytes(self, dst: &mut [u8]) -> Result<()> { + assert_eq!(dst.len(), self.n_bytes); + self.wait_and_read(|src| dst.copy_from_slice(src)) + } + + /// Wait and copy out as u64s. `dst.len() * 8` must equal the copied bytes. + pub fn wait_into_u64(self, dst: &mut [u64]) -> Result<()> { + assert_eq!(dst.len() * 8, self.n_bytes); + self.staging + .event + .as_ref() + .expect("recorded by async_dtoh_via") + .synchronize()?; + // SAFETY: as in `wait_and_read`; the slab is u64-aligned by + // construction. + let src = unsafe { std::slice::from_raw_parts(self.staging.ptr as *const u64, dst.len()) }; + dst.copy_from_slice(src); + Ok(()) + } +} + +// ── Pooled events for handle-readiness tracking ────────────────────────────── + +/// A pre-created CUDA event borrowed from the backend's free-list; returns +/// itself to the list on drop. Used as the `ready` marker on device-resident +/// handles (`GpuLdeBase`/`GpuLdeExt3`) so consumers on other streams can wait +/// device-side (`stream.wait`) instead of the producer host-blocking in a +/// final synchronize. Pooled because a mid-prove cuEventCreate convoys the +/// driver lock (see `PinnedStaging::record_event`). +pub struct PooledEvent { + event: Option, +} + +impl PooledEvent { + pub fn event(&self) -> &cudarc::driver::CudaEvent { + self.event.as_ref().expect("present until drop") + } +} + +impl Drop for PooledEvent { + fn drop(&mut self) { + if let (Some(ev), Ok(be)) = (self.event.take(), backend()) { + be.event_pool.lock().unwrap().push(ev); + } + } +} + +impl Backend { + /// Take a pre-created event from the pool (creating one only if the pool + /// ran dry, which should not happen in a normal prove). + pub fn take_event(&self) -> Result { + let ev = match self.event_pool.lock().unwrap().pop() { + Some(ev) => ev, + None => self.ctx.new_event(None)?, + }; + Ok(PooledEvent { event: Some(ev) }) + } +} diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index fb854d0a4..9cb68f3dd 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -116,7 +116,7 @@ impl FriCommitState { let tight_total_nodes = 2 * num_leaves - 1; // H2D zeta. - let zeta_dev = self.stream.clone_htod(&zeta_raw)?; + let zeta_dev = { self.stream.clone_htod(&zeta_raw)? }; let cfg = LaunchConfig { grid_dim: ((n_out as u32).div_ceil(128), 1, 1), @@ -176,7 +176,9 @@ impl FriCommitState { .launch(kcfg)?; } } - build_inner_tree_levels(self.stream.as_ref(), be, &mut nodes_dev, num_leaves)?; + { + build_inner_tree_levels(self.stream.as_ref(), be, &mut nodes_dev, num_leaves)?; + } // Update inv_twiddles for the next layer: `new[j] = old[2j]^2` for // j in 0..n_out/2. (If n_out == 1, skip; no next fold.) Writes into @@ -204,23 +206,40 @@ impl FriCommitState { } // Sync and D2H. - self.stream.synchronize()?; + { + self.stream.synchronize()?; + } - // Layer evals: 3 * n_out u64 from the output buffer. - let layer_evals: Vec = if self.a_is_input { - let view = self.evals_b.slice(0..3 * n_out); - self.stream.clone_dtoh(&view)? - } else { - let view = self.evals_a.slice(0..3 * n_out); - self.stream.clone_dtoh(&view)? + // Layer evals: 3 * n_out u64 from the output buffer, staged through + // the per-worker pinned slab (async DMA) instead of a blocking + // pageable copy. The wait is deferred past the root copy below. + let n_evals = 3 * n_out; + let pending = { + let output_evals: &CudaSlice = if self.a_is_input { + &self.evals_b + } else { + &self.evals_a + }; + crate::device::async_dtoh_via( + &self.stream, + be.pinned_staging(), + &be.ctx, + output_evals, + n_evals, + )? }; // Keep the layer tree resident on device; copy only the 32-byte root so // R4 query openings gather paths on device instead of copying the tree. + // This pageable copy drains the stream (including the evals DMA above), + // so the pending wait after it is instant — one block covers both. let mut root = [0u8; 32]; - self.stream - .memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; - self.stream.synchronize()?; + { + self.stream + .memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; + } + let mut layer_evals = vec![0u64; n_evals]; + pending.wait_into_u64(&mut layer_evals)?; self.a_is_input = !self.a_is_input; self.current_n = n_out; diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 485e005f8..27bff60d6 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -67,8 +67,15 @@ pub fn batch_inverse_ext3(a: &[u64]) -> Result> { let stream = be.next_stream(); let input_dev = stream.clone_htod(a)?; let out_dev = batch_inverse_ext3_dev(&input_dev, n, &stream)?; - let out = stream.clone_dtoh(&out_dev)?; + // Result download (3 * n u64s): async D2H through the per-worker pinned + // slab instead of a blocking pageable copy. The labelled sync keeps its + // host-block measurement (now covering the kernels plus the DMA); the + // pending wait after it is instant. + let pending = + crate::device::async_dtoh_via(&stream, be.pinned_staging(), &be.ctx, &out_dev, 3 * n)?; stream.synchronize()?; + let mut out = vec![0u64; 3 * n]; + pending.wait_into_u64(&mut out)?; Ok(out) } @@ -113,15 +120,21 @@ pub fn batch_inverse_ext3_dev( let mut prefix = unsafe { stream.alloc::(3 * n) }?; let mut suffix = unsafe { stream.alloc::(3 * n) }?; - scan_into_fwd(stream, be, input, &mut prefix, n)?; - scan_into_rev(stream, be, input, &mut suffix, n)?; + { + scan_into_fwd(stream, be, input, &mut prefix, n)?; + scan_into_rev(stream, be, input, &mut suffix, n)?; + } // total = prefix[n-1] = suffix[0]. Invert on host (one Fermat per batch). - let last_host: Vec = stream.clone_dtoh(&prefix.slice((n - 1) * 3..n * 3))?; - stream.synchronize()?; + let last_host: Vec = { stream.clone_dtoh(&prefix.slice((n - 1) * 3..n * 3))? }; + { + stream.synchronize()?; + } let inv_total = invert_ext3_host([last_host[0], last_host[1], last_host[2]])?; let mut inv_total_dev = unsafe { stream.alloc::(3) }?; - stream.memcpy_htod(&inv_total, &mut inv_total_dev)?; + { + stream.memcpy_htod(&inv_total, &mut inv_total_dev)?; + } // Combine: out[i] = prefix[i-1] * inv_total * suffix[i+1]. // SAFETY: the combine kernel writes every slot before any read. @@ -191,7 +204,7 @@ pub fn compute_and_invert_denoms_ext3_dev( )); } - let z_dev = stream.clone_htod(z_scalars_host)?; + let z_dev = { stream.clone_htod(z_scalars_host)? }; // SAFETY: the compute_denoms_ext3 kernel writes every output slot. let mut denoms = unsafe { stream.alloc::(3 * total) }?; let n_u64 = n as u64; diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 06c1a02cd..0a865ee43 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -167,25 +167,10 @@ fn d2h_bytes_via_pinned_hashes( dev_bytes: &CudaSlice, dst: &mut [u8], ) -> Result<()> { - let n_bytes = dst.len(); - let u64_len = n_bytes.div_ceil(8); - let staging_slot = be.pinned_hashes(); - let mut staging = staging_slot.lock().unwrap(); - staging.ensure_capacity(u64_len, &be.ctx)?; - let pinned = unsafe { staging.as_mut_slice(u64_len) }; - // Reinterpret the u64 pinned buffer as bytes — same allocation, just - // typed differently. SAFETY: u64 has stricter alignment than u8 and the - // byte length fits in the `u64_len` capacity (rounded up to u64). - let pinned_bytes: &mut [u8] = - unsafe { std::slice::from_raw_parts_mut(pinned.as_mut_ptr() as *mut u8, n_bytes) }; - stream.memcpy_dtoh(dev_bytes, pinned_bytes)?; - stream.synchronize()?; - - // Runs under the pinned_hashes lock, where rayon can deadlock. See - // `Backend::pinned_staging`. - dst.copy_from_slice(pinned_bytes); - drop(staging); - Ok(()) + let pending = + crate::device::async_dtoh_via(stream, be.pinned_hashes(), &be.ctx, dev_bytes, dst.len())?; + // Waits only for work queued up to the copy (event), not the whole stream. + pending.wait_into_bytes(dst) } /// Run `pointwise_mul_batched`: `buf[c*col_stride + i] *= weights[i]` for @@ -343,6 +328,46 @@ fn launch_keccak_base_row_major_row_pair( Ok(()) } +/// Column-range variant of [`launch_keccak_base_row_major_row_pair`]: leaves +/// hash only columns `[col_start, col_end)` of each bit-reversed row pair +/// (`m` stays the full row stride). Matches the CPU +/// `commit_rows_bit_reversed_subset`. +#[allow(clippy::too_many_arguments)] +fn launch_keccak_base_row_major_row_pair_range( + stream: &CudaStream, + be: &Backend, + buf: &CudaSlice, + m: u64, + col_start: u64, + col_end: u64, + num_rows: u64, + log_num_rows: u64, + leaves_out: &mut cudarc::driver::CudaViewMut<'_, u8>, +) -> Result<()> { + debug_assert!( + num_rows >= 2, + "row-major row-pair keccak requires num_rows >= 2" + ); + debug_assert!( + col_start < col_end && col_end <= m, + "column range in bounds" + ); + let cfg = keccak_launch_cfg(num_rows >> 1); + unsafe { + stream + .launch_builder(&be.keccak256_leaves_base_row_major_row_pair_range) + .arg(buf) + .arg(&m) + .arg(&col_start) + .arg(&col_end) + .arg(&num_rows) + .arg(&log_num_rows) + .arg(leaves_out) + .launch(cfg)?; + } + Ok(()) +} + /// Transpose row-major `lde_size × cols` → column-major with stride `lde_size`, /// returning the new device buffer. Used to convert the row-major LDE output to /// the column-major layout expected by downstream GPU kernels (DEEP, barycentric). @@ -386,69 +411,38 @@ enum InnerInput<'a> { Dev(&'a CudaSlice), } -/// Shared row-major LDE + Keccak + Merkle pipeline for the base and ext3 paths. -/// -/// `total_cols` is the number of base-field columns in the row-major layout: -/// `m` for base, `m * 3` for ext3. Because `Fp3 = [u64; 3]`, the three ext3 -/// components are just three adjacent base-field columns, so the same row-major -/// NTT and Keccak kernels process all of them simultaneously — no de-interleave. -/// -/// Single H2D (or D2D), row-major NTT, single D2H — no CPU-side extract or -/// transpose. Returns (merkle_nodes, column-major device buffer, row-major LDE -/// Vec, optional trace-domain column-major snapshot — `Some` iff -/// `retain_trace_col_major`). The buffer is transposed to column-major (as -/// required by the downstream GPU kernels DEEP/barycentric); callers wrap it in -/// the appropriate LDE handle. -#[allow(clippy::type_complexity)] +/// The expansion stage shared by the row-major commit pipelines: upload (or +/// D2D-copy) the row-major trace into a zero-padded `lde_size × total_cols` +/// buffer, optionally snapshot the trace-domain input column-major (for the +/// LogUp fingerprint kernel), then iNTT → coset weights → forward NTT in +/// place. Returns the row-major LDE buffer and the optional snapshot. #[allow(clippy::too_many_arguments)] -fn coset_lde_row_major_inner( +fn expand_row_major_on_stream( + stream: &Arc, + be: &Backend, input: InnerInput, n: usize, total_cols: usize, blowup_factor: usize, weights: &[u64], - what: &str, retain_trace_col_major: bool, - retain_host_lde: bool, -) -> Result<( - GpuMerkleTree, - CudaSlice, - Vec, - Option>, -)> { - let input_len = match &input { - InnerInput::Host(h) => h.len(), - InnerInput::Dev(d) => d.len(), - }; - assert_eq!(input_len, n * total_cols); - assert!(n.is_power_of_two()); - assert_eq!(weights.len(), n); - assert!(blowup_factor.is_power_of_two()); +) -> Result<(CudaSlice, Option>)> { let lde_size = n * blowup_factor; - assert_u32_domain(lde_size, what); - - // Row-pair trace commit: one Merkle leaf per bit-reversed row pair (rows 2i, - // 2i+1), matching the CPU `commit_bit_reversed(.., ROWS_PER_LEAF=2)` and the - // verifier's `verify_opening_pair`. `lde_size` is a power of two >= 2, so it - // is always even. - let num_leaves = lde_size / 2; - let nodes_bytes = KeccakCommit::FullTree.total_nodes_bytes(num_leaves); let log_n = n.trailing_zeros() as u64; let log_lde = lde_size.trailing_zeros() as u64; let n_u64 = n as u64; let lde_u64 = lde_size as u64; let cols_u64 = total_cols as u64; - let be = backend()?; - let stream = be.next_stream(); - // Fill a zeroed lde_size*total_cols buffer; only the first n*total_cols rows // carry data, the remainder are already zero (zero-padding for LDE). Host // input uploads (H2D); device input copies in place (D2D, no PCIe upload). let mut buf = stream.alloc_zeros::(lde_size * total_cols)?; - match input { - InnerInput::Host(h) => stream.memcpy_htod(h, &mut buf.slice_mut(0..n * total_cols))?, - InnerInput::Dev(d) => stream.memcpy_dtod(d, &mut buf.slice_mut(0..n * total_cols))?, + { + match input { + InnerInput::Host(h) => stream.memcpy_htod(h, &mut buf.slice_mut(0..n * total_cols))?, + InnerInput::Dev(d) => stream.memcpy_dtod(d, &mut buf.slice_mut(0..n * total_cols))?, + } } // Snapshot the trace-domain input (column-major) before the iNTT overwrites @@ -458,7 +452,7 @@ fn coset_lde_row_major_inner( // bit-reversed): dst[col*n + row] = buf[row*total_cols + col]. let trace_col_major = if retain_trace_col_major { Some(launch_row_to_col_major( - &stream, be, &buf, n, total_cols, n as u64, + stream, be, &buf, n, total_cols, n as u64, )?) } else { None @@ -495,6 +489,75 @@ fn coset_lde_row_major_inner( cols_u64, )?; + Ok((buf, trace_col_major)) +} + +/// Shared row-major LDE + Keccak + Merkle pipeline for the base and ext3 paths. +/// +/// `total_cols` is the number of base-field columns in the row-major layout: +/// `m` for base, `m * 3` for ext3. Because `Fp3 = [u64; 3]`, the three ext3 +/// components are just three adjacent base-field columns, so the same row-major +/// NTT and Keccak kernels process all of them simultaneously — no de-interleave. +/// +/// Single H2D (or D2D), row-major NTT, single D2H — no CPU-side extract or +/// transpose. Returns (merkle_nodes, column-major device buffer, row-major LDE +/// Vec, optional trace-domain column-major snapshot — `Some` iff +/// `retain_trace_col_major`). The buffer is transposed to column-major (as +/// required by the downstream GPU kernels DEEP/barycentric); callers wrap it in +/// the appropriate LDE handle. +#[allow(clippy::type_complexity)] +#[allow(clippy::too_many_arguments)] +fn coset_lde_row_major_inner( + input: InnerInput, + n: usize, + total_cols: usize, + blowup_factor: usize, + weights: &[u64], + what: &str, + retain_trace_col_major: bool, + retain_host_lde: bool, +) -> Result<( + GpuMerkleTree, + CudaSlice, + Vec, + Option>, + Arc, +)> { + let input_len = match &input { + InnerInput::Host(h) => h.len(), + InnerInput::Dev(d) => d.len(), + }; + assert_eq!(input_len, n * total_cols); + assert!(n.is_power_of_two()); + assert_eq!(weights.len(), n); + assert!(blowup_factor.is_power_of_two()); + let lde_size = n * blowup_factor; + assert_u32_domain(lde_size, what); + + // Row-pair trace commit: one Merkle leaf per bit-reversed row pair (rows 2i, + // 2i+1), matching the CPU `commit_bit_reversed(.., ROWS_PER_LEAF=2)` and the + // verifier's `verify_opening_pair`. `lde_size` is a power of two >= 2, so it + // is always even. + let num_leaves = lde_size / 2; + let nodes_bytes = KeccakCommit::FullTree.total_nodes_bytes(num_leaves); + let log_lde = lde_size.trailing_zeros() as u64; + let lde_u64 = lde_size as u64; + let cols_u64 = total_cols as u64; + + let be = backend()?; + let stream = be.next_stream(); + + let (buf, trace_col_major) = expand_row_major_on_stream( + &stream, + be, + input, + n, + total_cols, + blowup_factor, + weights, + retain_trace_col_major, + )?; + // Keccak + Merkle on-device. Each row-pair leaf reads two bit-reversed rows // of `total_cols` consecutive u64s (`lde_u64` is the bit-reverse modulus; the // kernel emits `lde_size / 2` leaves). @@ -514,45 +577,60 @@ fn coset_lde_row_major_inner( } crate::merkle::build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; - // D2H the row-major LDE first (before the handle transpose). Release the - // staging lock before the Merkle nodes transfer to minimise lock contention. - // Skipped entirely when `retain_host_lde` is false (the caller keeps the LDE - // device-only): this is the round-1 trace D2H the full-residency path - // eliminates — the big transfer/alloc win — so we return an empty host Vec. - let lde_out = if retain_host_lde { - let staging_slot = be.pinned_staging(); - let mut staging = staging_slot.lock().unwrap(); - staging.ensure_capacity(lde_size * total_cols, &be.ctx)?; - let pinned = unsafe { staging.as_mut_slice(lde_size * total_cols) }; - stream.memcpy_dtoh(&buf, pinned)?; - stream.synchronize()?; - let out = pinned[..lde_size * total_cols].to_vec(); - drop(staging); - out - } else { - Vec::new() - }; - - // Keep the Merkle tree resident on device; copy only the 32 byte root so the - // commitment is available without copying the whole tree. Query openings - // gather paths from the device tree (see merkle::gather_merkle_paths_dev). + // Copy the 32-byte root BEFORE queueing the big drain/transpose: this + // pageable copy host-blocks until everything queued so far lands, so + // keeping it early means it waits for the tree kernels only (the root is + // needed now regardless — Fiat-Shamir absorbs it before anything else). + // The Merkle tree stays resident on device; query openings gather paths + // from it (see merkle::gather_merkle_paths_dev). let mut root = [0u8; 32]; stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; + // D2H the row-major LDE (skipped when `retain_host_lde` is false — the + // full-residency path keeps the LDE device-only; that skip is the big + // transfer/alloc win, and we return an empty host Vec). + let lde_pending = if retain_host_lde { + Some(crate::device::async_dtoh_via( + &stream, + be.pinned_staging(), + &be.ctx, + &buf, + lde_size * total_cols, + )?) + } else { + None + }; + // Transpose row-major buf into column-major for the handle. Downstream // kernels (DEEP, barycentric) expect buf[c * lde_size + r] (column-major). let col_major_dev = launch_row_to_col_major(&stream, be, &buf, lde_size, total_cols, lde_u64)?; - // Synchronize before returning: the handle crosses stream boundaries. - // Downstream consumers call be.next_stream() and read handle.buf on a - // different stream, and the root copy above must have landed. - stream.synchronize()?; + // No host synchronize here: the handle carries a `ready` event instead, + // and consumers on other streams wait on it device-side + // (`wait_ready_on`). On the device-only path this makes the whole + // commit's tail (transpose) run behind the host's next work. + let ready = be.take_event()?; + ready.event().record(&stream)?; + let lde_out = match lde_pending { + Some(p) => { + let mut out = vec![0u64; lde_size * total_cols]; + p.wait_into_u64(&mut out)?; + out + } + None => Vec::new(), + }; let tree = GpuMerkleTree { nodes: Arc::new(nodes_dev), leaves_len: num_leaves, root, }; - Ok((tree, col_major_dev, lde_out, trace_col_major)) + Ok(( + tree, + col_major_dev, + lde_out, + trace_col_major, + Arc::new(ready), + )) } /// Row-major LDE + Keccak + Merkle, all on-device, keeping the Merkle tree @@ -571,7 +649,7 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( weights: &[u64], retain_host_lde: bool, ) -> Result<(GpuLdeBase, Vec)> { - let (tree, col_major_dev, lde_out, trace_col_major) = coset_lde_row_major_inner( + let (tree, col_major_dev, lde_out, trace_col_major, ready) = coset_lde_row_major_inner( InnerInput::Host(row_major), n, m, @@ -586,12 +664,125 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( m, lde_size: n * blowup_factor, tree: Some(tree), + ready: Some(ready), trace_dev: trace_col_major.map(Arc::new), trace_rows: n, }; Ok((handle, lde_out)) } +/// Row-major LDE + TWO subset Merkle trees for preprocessed tables: the +/// precomputed columns `[0, split_col)` and the multiplicity columns +/// `[split_col, m)` commit to separate trees over the same row-major LDE, +/// mirroring the CPU `commit_rows_bit_reversed_subset` pair. +/// +/// Both trees' complete node buffers are downloaded to host +/// (`(2*num_leaves - 1) * 32` bytes each, inner nodes first, root at offset 0, +/// leaves at the tail — the exact `MerkleTree::from_precomputed_nodes` +/// layout), because preprocessed-table openings walk host trees. The +/// precomputed tree is only built when `build_precomputed` is true (the +/// caller skips it on a process-cache hit). +/// +/// Returns `(precomputed_nodes, mult_nodes, handle, row_major_lde)`. The +/// handle carries the column-major LDE + trace snapshot for downstream GPU +/// rounds but NO device tree (`tree: None`) — openings for preprocessed +/// tables never gather from device. +#[allow(clippy::type_complexity)] +pub fn coset_lde_row_major_split_trees( + row_major: &[u64], + n: usize, + m: usize, + blowup_factor: usize, + weights: &[u64], + split_col: usize, + build_precomputed: bool, +) -> Result<(Option>, Vec, GpuLdeBase, Vec)> { + assert!(split_col > 0 && split_col < m, "split inside the row"); + let lde_size = n * blowup_factor; + assert_u32_domain(lde_size, "coset_lde_row_major_split lde_size"); + let num_leaves = lde_size / 2; + let nodes_bytes = KeccakCommit::FullTree.total_nodes_bytes(num_leaves); + let leaves_offset = KeccakCommit::FullTree.leaves_offset_bytes(num_leaves); + let log_lde = lde_size.trailing_zeros() as u64; + let lde_u64 = lde_size as u64; + let cols_u64 = m as u64; + + let be = backend()?; + let stream = be.next_stream(); + + let (buf, trace_col_major) = expand_row_major_on_stream( + &stream, + be, + InnerInput::Host(row_major), + n, + m, + blowup_factor, + weights, + true, + )?; + + // One subset tree per column range, built sequentially on the stream. + let build_subset_tree = |col_start: u64, col_end: u64| -> Result> { + let mut nodes_dev = unsafe { stream.alloc::(nodes_bytes) }?; + { + let mut leaves_view = + nodes_dev.slice_mut(leaves_offset..leaves_offset + num_leaves * 32); + launch_keccak_base_row_major_row_pair_range( + stream.as_ref(), + be, + &buf, + cols_u64, + col_start, + col_end, + lde_u64, + log_lde, + &mut leaves_view, + )?; + } + crate::merkle::build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; + let mut nodes_host = vec![0u8; nodes_bytes]; + { + stream.memcpy_dtoh(&nodes_dev, &mut nodes_host)?; + } + Ok(nodes_host) + }; + + let precomputed_nodes = if build_precomputed { + Some(build_subset_tree(0, split_col as u64)?) + } else { + None + }; + let mult_nodes = build_subset_tree(split_col as u64, cols_u64)?; + + // D2H the row-major LDE (preprocessed tables always keep the host copy — + // they are excluded from the device-only gate). + let lde_pending = + crate::device::async_dtoh_via(&stream, be.pinned_staging(), &be.ctx, &buf, lde_size * m)?; + + // Column-major handle for downstream GPU rounds (DEEP, barycentric, + // constraint composition). + let col_major_dev = launch_row_to_col_major(&stream, be, &buf, lde_size, m, lde_u64)?; + let ready = be.take_event()?; + ready.event().record(&stream)?; + + let lde_out = { + let mut out = vec![0u64; lde_size * m]; + lde_pending.wait_into_u64(&mut out)?; + out + }; + + let handle = GpuLdeBase { + buf: Arc::new(col_major_dev), + m, + lde_size, + tree: None, + ready: Some(Arc::new(ready)), + trace_dev: trace_col_major.map(Arc::new), + trace_rows: n, + }; + Ok((precomputed_nodes, mult_nodes, handle, lde_out)) +} + /// Row-major ext3 LDE + Keccak + Merkle, all on-device. /// /// `Fp3` is `[u64; 3]` in memory, so row-major ext3 with `m` ext3 columns is @@ -610,7 +801,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( weights: &[u64], retain_host_lde: bool, ) -> Result<(GpuLdeExt3, Vec)> { - let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( + let (tree, col_major_dev, lde_out, _, ready) = coset_lde_row_major_inner( InnerInput::Host(row_major), n, m * 3, @@ -625,6 +816,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( m, lde_size: n * blowup_factor, tree: Some(tree), + ready: Some(ready), }; Ok((handle, lde_out)) } @@ -641,7 +833,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev( weights: &[u64], retain_host_lde: bool, ) -> Result<(GpuLdeExt3, Vec)> { - let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( + let (tree, col_major_dev, lde_out, _, ready) = coset_lde_row_major_inner( InnerInput::Dev(input_dev), n, m * 3, @@ -656,6 +848,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev( m, lde_size: n * blowup_factor, tree: Some(tree), + ready: Some(ready), }; Ok((handle, lde_out)) } @@ -679,6 +872,20 @@ pub struct GpuLdeBase { pub trace_dev: Option>>, /// Row count (n) of `trace_dev`; 0 when `trace_dev` is None. pub trace_rows: usize, + /// Fires once `buf` is fully written (recorded after the producer's last + /// kernel). `None` means the producer synchronized before returning. + /// Consumers on other streams call [`GpuLdeBase::wait_ready_on`]. + pub ready: Option>, +} + +impl GpuLdeBase { + /// Make `stream` wait (device-side, no host block) until `buf` is ready. + pub fn wait_ready_on(&self, stream: &CudaStream) -> Result<()> { + match &self.ready { + Some(ev) => stream.wait(ev.event()), + None => Ok(()), + } + } } /// Handle to an ext3 LDE kept live on device, de-interleaved into 3 base @@ -692,6 +899,19 @@ pub struct GpuLdeExt3 { /// Optionally the aux or composition Merkle tree kept resident on device /// (the keep path), so R4 openings gather paths on device. None otherwise. pub tree: Option, + /// Fires once `buf` is fully written. `None` = producer synchronized. + /// Consumers on other streams call [`GpuLdeExt3::wait_ready_on`]. + pub ready: Option>, +} + +impl GpuLdeExt3 { + /// Make `stream` wait (device-side, no host block) until `buf` is ready. + pub fn wait_ready_on(&self, stream: &CudaStream) -> Result<()> { + match &self.ready { + Some(ev) => stream.wait(ev.event()), + None => Ok(()), + } + } } /// Merkle tree kept resident on device after a commit, so query openings gather @@ -744,13 +964,15 @@ pub fn coset_lde_base(evals: &[u64], blowup_factor: usize, weights: &[u64]) -> R let lde_u64 = lde_size as u64; // === 1. iNTT on first N: bit_reverse + 8-level-fused DIT body === - unsafe { - stream - .launch_builder(&be.bit_reverse_permute) - .arg(&mut buf) - .arg(&n_u64) - .arg(&log_n) - .launch(LaunchConfig::for_num_elems(n as u32))?; + { + unsafe { + stream + .launch_builder(&be.bit_reverse_permute) + .arg(&mut buf) + .arg(&n_u64) + .arg(&log_n) + .launch(LaunchConfig::for_num_elems(n as u32))?; + } } // Note: `run_ntt_body` expects a standalone CudaSlice; we pass `buf` and // the kernel walks the first `n_u64` elements via its own indexing. @@ -759,28 +981,34 @@ pub fn coset_lde_base(evals: &[u64], blowup_factor: usize, weights: &[u64]) -> R // next pointwise multiply applies both the coset shift and the 1/N factor. // === 2. Pointwise multiply first N by coset weights (includes 1/N) === - unsafe { - stream - .launch_builder(&be.pointwise_mul) - .arg(&mut buf) - .arg(&weights_dev) - .arg(&n_u64) - .launch(LaunchConfig::for_num_elems(n as u32))?; + { + unsafe { + stream + .launch_builder(&be.pointwise_mul) + .arg(&mut buf) + .arg(&weights_dev) + .arg(&n_u64) + .launch(LaunchConfig::for_num_elems(n as u32))?; + } } // === 3. Forward NTT on full buffer === - unsafe { - stream - .launch_builder(&be.bit_reverse_permute) - .arg(&mut buf) - .arg(&lde_u64) - .arg(&log_lde) - .launch(LaunchConfig::for_num_elems(lde_size as u32))?; + { + unsafe { + stream + .launch_builder(&be.bit_reverse_permute) + .arg(&mut buf) + .arg(&lde_u64) + .arg(&log_lde) + .launch(LaunchConfig::for_num_elems(lde_size as u32))?; + } } run_ntt_body(stream.as_ref(), &mut buf, fwd_tw.as_ref(), lde_u64, log_lde)?; - let out = stream.clone_dtoh(&buf)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&buf)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -824,9 +1052,9 @@ pub fn coset_lde_batch_base( let staging_slot = be.pinned_staging(); // Pinned staging. Lock and grow to max(m*n for upload, m*lde_size for - // download). Holding the guard across the whole call serialises concurrent - // batched calls that happened to hash to the same stream slot, but that's - // exactly what we want — one stream can only do one sequence at a time. + // download). The guard is held from the pack until the async uploads have + // landed (the H2D DMA reads the slab directly); the D2H drain at the end + // re-acquires the slot via `async_dtoh_via`. let mut staging = staging_slot.lock().unwrap(); staging.ensure_capacity(m * lde_size, &be.ctx)?; // SAFETY: staging is locked, the slice alias ends before we unlock. @@ -835,8 +1063,10 @@ pub fn coset_lde_batch_base( // Pack columns into the first m*n slots of the pinned buffer. Runs under // the pinned-staging lock, where rayon can deadlock. See // `Backend::pinned_staging`. - for (c, col) in columns.iter().enumerate() { - pinned[c * n..c * n + n].copy_from_slice(col); + { + for (c, col) in columns.iter().enumerate() { + pinned[c * n..c * n + n].copy_from_slice(col); + } } // Column layout: `buf[c * lde_size + r]`. Zeroed so the [n, lde_size) @@ -844,10 +1074,16 @@ pub fn coset_lde_batch_base( let mut buf = stream.alloc_zeros::(m * lde_size)?; // One memcpy per column from the pinned buffer into the strided slots. // The pinned source hits PCIe line-rate. - for c in 0..m { - let mut dst = buf.slice_mut(c * lde_size..c * lde_size + n); - stream.memcpy_htod(&pinned[c * n..c * n + n], &mut dst)?; + { + for c in 0..m { + let mut dst = buf.slice_mut(c * lde_size..c * lde_size + n); + stream.memcpy_htod(&pinned[c * n..c * n + n], &mut dst)?; + } } + // The uploads above are truly asynchronous (pinned source), so the + // staging slot must stay locked until they land; the slot's reusable event marks that + // point. It is waited just before the D2H drain re-acquires the slot. + staging.record_event(&stream)?; let inv_tw = be.inv_twiddles_for(log_n)?; let fwd_tw = be.fwd_twiddles_for(log_lde)?; @@ -913,29 +1149,44 @@ pub fn coset_lde_batch_base( m_u32, )?; + // Release the staging slot before the drain: the uploads have landed once + // the slot event fires (the NTT kernels above are queued behind them, so the + // GPU stays busy while the host waits here). + staging.sync_event()?; + drop(staging); + // Single big D2H into the reusable pinned staging buffer — pinned, one - // call to the driver, saturates PCIe. - stream.memcpy_dtoh(&buf, &mut pinned[..m * lde_size])?; - stream.synchronize()?; + // call to the driver, saturates PCIe. Enqueued without blocking; the host + // blocks once, in `wait_and_read` below. + let pending = + crate::device::async_dtoh_via(&stream, staging_slot, &be.ctx, &buf, m * lde_size)?; // Split pinned into per-column Vecs. Runs under the pinned-staging - // lock, where rayon can deadlock. See `Backend::pinned_staging`. - let out: Vec> = (0..m) - .map(|c| { - // set_len skips the O(N) zero-init that vec![0; n] would do. - // copy_from_slice below writes every slot before any reader - // sees the Vec. - #[allow(clippy::uninit_vec)] - let mut v = { - let mut v = Vec::::with_capacity(lde_size); - unsafe { v.set_len(lde_size) }; - v - }; - v.copy_from_slice(&pinned[c * lde_size..c * lde_size + lde_size]); - v - }) - .collect(); - drop(staging); + // lock (held by `pending`), where rayon can deadlock. See + // `Backend::pinned_staging`. + let out: Vec> = { + pending.wait_and_read(|bytes| { + // SAFETY: the pinned slab is u64-aligned by construction and the + // copy deposited exactly `m * lde_size` u64s. + let pinned = + unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const u64, m * lde_size) }; + (0..m) + .map(|c| { + // set_len skips the O(N) zero-init that vec![0; n] would + // do. copy_from_slice below writes every slot before any + // reader sees the Vec. + #[allow(clippy::uninit_vec)] + let mut v = { + let mut v = Vec::::with_capacity(lde_size); + unsafe { v.set_len(lde_size) }; + v + }; + v.copy_from_slice(&pinned[c * lde_size..c * lde_size + lde_size]); + v + }) + .collect() + })? + }; Ok(out) } @@ -986,15 +1237,22 @@ pub fn coset_lde_batch_base_into( staging.ensure_capacity(m * lde_size, &be.ctx)?; let pinned = unsafe { staging.as_mut_slice(m * lde_size) }; - for (c, col) in columns.iter().enumerate() { - pinned[c * n..c * n + n].copy_from_slice(col); + { + for (c, col) in columns.iter().enumerate() { + pinned[c * n..c * n + n].copy_from_slice(col); + } } let mut buf = stream.alloc_zeros::(m * lde_size)?; - for c in 0..m { - let mut dst = buf.slice_mut(c * lde_size..c * lde_size + n); - stream.memcpy_htod(&pinned[c * n..c * n + n], &mut dst)?; + { + for c in 0..m { + let mut dst = buf.slice_mut(c * lde_size..c * lde_size + n); + stream.memcpy_htod(&pinned[c * n..c * n + n], &mut dst)?; + } } + // The uploads above are truly asynchronous (pinned source); the staging + // slot stays locked until this event fires (waited before the drain). + staging.record_event(&stream)?; let inv_tw = be.inv_twiddles_for(log_n)?; let fwd_tw = be.fwd_twiddles_for(log_lde)?; @@ -1052,15 +1310,30 @@ pub fn coset_lde_batch_base_into( m_u32, )?; - stream.memcpy_dtoh(&buf, &mut pinned[..m * lde_size])?; - stream.synchronize()?; + // Release the staging slot before the drain: the uploads have landed once + // the slot event fires (the kernels above are queued behind them). + staging.sync_event()?; + drop(staging); + + // Big D2H enqueued without blocking; the host blocks once, in + // `wait_and_read` below. + let pending = + crate::device::async_dtoh_via(&stream, staging_slot, &be.ctx, &buf, m * lde_size)?; - // Copy pinned into caller outputs. Runs under the pinned-staging lock, - // where rayon can deadlock. See `Backend::pinned_staging`. - for (c, dst) in outputs.iter_mut().enumerate() { - dst.copy_from_slice(&pinned[c * lde_size..c * lde_size + lde_size]); + // Copy pinned into caller outputs. Runs under the pinned-staging lock + // (held by `pending`), where rayon can deadlock. See + // `Backend::pinned_staging`. + { + pending.wait_and_read(|bytes| { + // SAFETY: the pinned slab is u64-aligned by construction and the + // copy deposited exactly `m * lde_size` u64s. + let pinned = + unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const u64, m * lde_size) }; + for (c, dst) in outputs.iter_mut().enumerate() { + dst.copy_from_slice(&pinned[c * lde_size..c * lde_size + lde_size]); + } + })?; } - drop(staging); Ok(()) } @@ -1146,15 +1419,22 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner( // Pack columns into the pinned buffer. Runs under the pinned-staging // lock, where rayon can deadlock. See `Backend::pinned_staging`. - for (c, col) in columns.iter().enumerate() { - pinned[c * n..c * n + n].copy_from_slice(col); + { + for (c, col) in columns.iter().enumerate() { + pinned[c * n..c * n + n].copy_from_slice(col); + } } let mut buf = stream.alloc_zeros::(m * lde_size)?; - for c in 0..m { - let mut dst = buf.slice_mut(c * lde_size..c * lde_size + n); - stream.memcpy_htod(&pinned[c * n..c * n + n], &mut dst)?; + { + for c in 0..m { + let mut dst = buf.slice_mut(c * lde_size..c * lde_size + n); + stream.memcpy_htod(&pinned[c * n..c * n + n], &mut dst)?; + } } + // The uploads above are truly asynchronous (pinned source); the staging + // slot stays locked until this event fires (waited before the drain). + staging.record_event(&stream)?; let inv_tw = be.inv_twiddles_for(log_n)?; let fwd_tw = be.fwd_twiddles_for(log_lde)?; @@ -1249,16 +1529,33 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner( crate::merkle::build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; } - // D2H the LDE and the tree/leaves nodes via pinned staging. - stream.memcpy_dtoh(&buf, &mut pinned[..m * lde_size])?; + // Release the staging slot before the drain: the uploads have landed once + // the slot event fires (the kernels above are queued behind them). + staging.sync_event()?; + drop(staging); + + // D2H the LDE (async, via pinned staging, enqueued without blocking) and + // the tree/leaves nodes (via the separate pinned-hashes slot; that helper + // waits internally, and its event is recorded after the LDE copy, so the + // `wait_and_read` below is nearly instant). + let lde_pending = + crate::device::async_dtoh_via(&stream, staging_slot, &be.ctx, &buf, m * lde_size)?; d2h_bytes_via_pinned_hashes(&stream, be, &nodes_dev, nodes_out)?; - // Copy pinned into caller outputs. Runs under the pinned-staging lock, - // where rayon can deadlock. See `Backend::pinned_staging`. - for (c, dst) in outputs.iter_mut().enumerate() { - dst.copy_from_slice(&pinned[c * lde_size..c * lde_size + lde_size]); + // Copy pinned into caller outputs. Runs under the pinned-staging lock + // (held by `lde_pending`), where rayon can deadlock. See + // `Backend::pinned_staging`. + { + lde_pending.wait_and_read(|bytes| { + // SAFETY: the pinned slab is u64-aligned by construction and the + // copy deposited exactly `m * lde_size` u64s. + let pinned = + unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const u64, m * lde_size) }; + for (c, dst) in outputs.iter_mut().enumerate() { + dst.copy_from_slice(&pinned[c * lde_size..c * lde_size + lde_size]); + } + })?; } - drop(staging); if keep_device_buf { Ok(Some(GpuLdeBase { @@ -1268,6 +1565,9 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner( tree: None, trace_dev: None, trace_rows: 0, + // The pending wait above drained the stream past the last write + // to `buf`, so the handle is complete at return. + ready: None, })) } else { drop(buf); @@ -1373,10 +1673,15 @@ fn evaluate_poly_coset_batch_ext3_into_inner( pack_ext3_to_pinned_slabs(coefs, pinned, n); let mut buf = stream.alloc_zeros::(mb * lde_size)?; - for s in 0..mb { - let mut dst = buf.slice_mut(s * lde_size..s * lde_size + n); - stream.memcpy_htod(&pinned[s * n..s * n + n], &mut dst)?; + { + for s in 0..mb { + let mut dst = buf.slice_mut(s * lde_size..s * lde_size + n); + stream.memcpy_htod(&pinned[s * n..s * n + n], &mut dst)?; + } } + // The uploads above are truly asynchronous (pinned source); the staging + // slot stays locked until this event fires (waited before the drain). + staging.record_event(&stream)?; let fwd_tw = be.fwd_twiddles_for(log_lde)?; let weights_dev = stream.clone_htod(weights)?; @@ -1417,8 +1722,9 @@ fn evaluate_poly_coset_batch_ext3_into_inner( mb_u32, )?; - // Optional R2-style row-pair Merkle tree build on the LDE buffer. - if let Some(nodes_out) = merkle_nodes_out { + // Optional R2-style row-pair Merkle tree build on the LDE buffer, queued + // ahead of the drains below. + let nodes = if let Some(nodes_out) = merkle_nodes_out { let num_leaves = lde_size / 2; let tight_total_nodes = 2 * num_leaves - 1; assert_eq!(nodes_out.len(), tight_total_nodes * 32); @@ -1443,22 +1749,44 @@ fn evaluate_poly_coset_batch_ext3_into_inner( } } crate::merkle::build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; + Some((nodes_dev, nodes_out)) + } else { + None + }; - stream.memcpy_dtoh(&buf, &mut pinned[..mb * lde_size])?; + // Release the staging slot before the drain: the uploads have landed once + // the slot event fires (the kernels above are queued behind them). + staging.sync_event()?; + drop(staging); + + // LDE drain enqueued without blocking. When a tree was built, its nodes + // drain via the separate pinned-hashes slot; that helper waits internally, + // and its event is recorded after the LDE copy, so the `wait_and_read` + // below is nearly instant. + let lde_pending = + crate::device::async_dtoh_via(&stream, staging_slot, &be.ctx, &buf, mb * lde_size)?; + if let Some((nodes_dev, nodes_out)) = nodes { d2h_bytes_via_pinned_hashes(&stream, be, &nodes_dev, nodes_out)?; - } else { - stream.memcpy_dtoh(&buf, &mut pinned[..mb * lde_size])?; - stream.synchronize()?; } - unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); - drop(staging); + { + lde_pending.wait_and_read(|bytes| { + // SAFETY: the pinned slab is u64-aligned by construction and the + // copy deposited exactly `mb * lde_size` u64s. + let pinned = + unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const u64, mb * lde_size) }; + unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); + })?; + } if keep_device_buf { Ok(Some(GpuLdeExt3 { buf: std::sync::Arc::new(buf), m, lde_size, tree: None, + // The pending wait above drained the stream past the last write + // to `buf`, so the handle is complete at return. + ready: None, })) } else { drop(buf); @@ -1562,10 +1890,15 @@ pub fn coset_lde_batch_ext3_into( // Allocate + zero-pad device buffer holding 3M slabs of `lde_size`. let mut buf = stream.alloc_zeros::(mb * lde_size)?; // H2D: slab by slab into the first N slots of each `lde_size`-slab. - for s in 0..mb { - let mut dst = buf.slice_mut(s * lde_size..s * lde_size + n); - stream.memcpy_htod(&pinned[s * n..s * n + n], &mut dst)?; + { + for s in 0..mb { + let mut dst = buf.slice_mut(s * lde_size..s * lde_size + n); + stream.memcpy_htod(&pinned[s * n..s * n + n], &mut dst)?; + } } + // The uploads above are truly asynchronous (pinned source); the staging + // slot stays locked until this event fires (waited before the drain). + staging.record_event(&stream)?; let inv_tw = be.inv_twiddles_for(log_n)?; let fwd_tw = be.fwd_twiddles_for(log_lde)?; @@ -1624,13 +1957,28 @@ pub fn coset_lde_batch_ext3_into( mb_u32, )?; - stream.memcpy_dtoh(&buf, &mut pinned[..mb * lde_size])?; - stream.synchronize()?; + // Release the staging slot before the drain: the uploads have landed once + // the slot event fires (the kernels above are queued behind them). + staging.sync_event()?; + drop(staging); + + // Big D2H enqueued without blocking; the host blocks once, in + // `wait_and_read` below. + let pending = + crate::device::async_dtoh_via(&stream, staging_slot, &be.ctx, &buf, mb * lde_size)?; // Unpack: for each output column, re-interleave 3 slabs back into the - // ext3-per-element layout. - unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); - drop(staging); + // ext3-per-element layout. Runs under the pinned-staging lock (held by + // `pending`), where rayon can deadlock. See `Backend::pinned_staging`. + { + pending.wait_and_read(|bytes| { + // SAFETY: the pinned slab is u64-aligned by construction and the + // copy deposited exactly `mb * lde_size` u64s. + let pinned = + unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const u64, mb * lde_size) }; + unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); + })?; + } Ok(()) } diff --git a/crypto/math-cuda/src/lib.rs b/crypto/math-cuda/src/lib.rs index 493480991..6b58d935b 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -15,6 +15,7 @@ pub mod lde; pub mod logup; pub mod merkle; pub mod ntt; +pub mod nvtx; // Re-exported for downstream crates so they can refer to CUDA primitive // types without depending on cudarc directly. diff --git a/crypto/math-cuda/src/logup.rs b/crypto/math-cuda/src/logup.rs index e9d120ee1..68f34f378 100644 --- a/crypto/math-cuda/src/logup.rs +++ b/crypto/math-cuda/src/logup.rs @@ -141,15 +141,15 @@ pub fn logup_term_columns( let stream = be.next_stream(); let timing = std::env::var_os("LAMBDA_VM_LOGUP_TIMING").is_some(); let t0 = std::time::Instant::now(); - let main_dev = stream.clone_htod(main_cols)?; + let main_dev = { stream.clone_htod(main_cols)? }; if timing { stream.synchronize()?; } let t1 = std::time::Instant::now(); - let fp = fingerprints_into_dev(&main_dev, num_rows, d, alpha_powers, z, &stream)?; + let fp = { fingerprints_into_dev(&main_dev, num_rows, d, alpha_powers, z, &stream)? }; let n = d.num_interactions * num_rows; - let recip = batch_inverse_ext3_dev(&fp, n, &stream)?; + let recip = { batch_inverse_ext3_dev(&fp, n, &stream)? }; let total = d.num_out_cols * num_rows; let mut out = unsafe { stream.alloc::(total * 3) }?; @@ -158,12 +158,23 @@ pub fn logup_term_columns( return Ok(Vec::new()); } - let out_col_offsets = stream.clone_htod(d.out_col_offsets)?; - let out_col_interactions = stream.clone_htod(d.out_col_interactions)?; - let mult_const = stream.clone_htod(d.mult_const)?; - let mult_term_offsets = stream.clone_htod(d.mult_term_offsets)?; - let mult_term_coef = stream.clone_htod(d.mult_term_coef)?; - let mult_term_col = stream.clone_htod(d.mult_term_col)?; + let ( + out_col_offsets, + out_col_interactions, + mult_const, + mult_term_offsets, + mult_term_coef, + mult_term_col, + ) = { + ( + stream.clone_htod(d.out_col_offsets)?, + stream.clone_htod(d.out_col_interactions)?, + stream.clone_htod(d.mult_const)?, + stream.clone_htod(d.mult_term_offsets)?, + stream.clone_htod(d.mult_term_coef)?, + stream.clone_htod(d.mult_term_col)?, + ) + }; let num_rows_u32 = num_rows as u32; let num_out_u32 = d.num_out_cols as u32; unsafe { @@ -186,8 +197,17 @@ pub fn logup_term_columns( stream.synchronize()?; } let t2 = std::time::Instant::now(); - let host = stream.clone_dtoh(&out)?; - stream.synchronize()?; + // Terms download (num_out_cols * num_rows * 3 u64s): async D2H through + // the per-worker pinned slab instead of a blocking pageable copy. The + // labelled sync keeps its host-block measurement (now covering the + // kernels plus the DMA); the pending wait after it is instant. + let pending = + { crate::device::async_dtoh_via(&stream, be.pinned_staging(), &be.ctx, &out, total * 3)? }; + { + stream.synchronize()?; + } + let mut host = vec![0u64; total * 3]; + pending.wait_into_u64(&mut host)?; let t3 = std::time::Instant::now(); if timing { eprintln!( @@ -316,9 +336,11 @@ pub fn logup_aux_resident( // Resident device main = zero upload; host main = one H2D. `uploaded` owns // the staged buffer for the function scope so `main_dev` can borrow it. - let uploaded: Option> = match main { - ResidentMain::Dev(_) => None, - ResidentMain::Host(h) => Some(stream.clone_htod(h)?), + let uploaded: Option> = { + match main { + ResidentMain::Dev(_) => None, + ResidentMain::Host(h) => Some(stream.clone_htod(h)?), + } }; let main_dev: &CudaSlice = match (main, &uploaded) { (ResidentMain::Dev(d), _) => d, @@ -329,12 +351,12 @@ pub fn logup_aux_resident( sync_if(timing)?; let t_h2d = std::time::Instant::now(); - let fp = fingerprints_into_dev(main_dev, num_rows, d, alpha_powers, z, stream)?; + let fp = { fingerprints_into_dev(main_dev, num_rows, d, alpha_powers, z, stream)? }; sync_if(timing)?; let t_fp = std::time::Instant::now(); let n = d.num_interactions * num_rows; - let recip = batch_inverse_ext3_dev(&fp, n, stream)?; + let recip = { batch_inverse_ext3_dev(&fp, n, stream)? }; sync_if(timing)?; let t_inv = std::time::Instant::now(); @@ -349,12 +371,23 @@ pub fn logup_aux_resident( } let num_out = d.num_out_cols; let mut terms = unsafe { stream.alloc::(num_out * num_rows * 3) }?; - let out_col_offsets = stream.clone_htod(d.out_col_offsets)?; - let out_col_interactions = stream.clone_htod(d.out_col_interactions)?; - let mult_const = stream.clone_htod(d.mult_const)?; - let mult_term_offsets = stream.clone_htod(d.mult_term_offsets)?; - let mult_term_coef = stream.clone_htod(d.mult_term_coef)?; - let mult_term_col = stream.clone_htod(d.mult_term_col)?; + let ( + out_col_offsets, + out_col_interactions, + mult_const, + mult_term_offsets, + mult_term_coef, + mult_term_col, + ) = { + ( + stream.clone_htod(d.out_col_offsets)?, + stream.clone_htod(d.out_col_interactions)?, + stream.clone_htod(d.mult_const)?, + stream.clone_htod(d.mult_term_offsets)?, + stream.clone_htod(d.mult_term_coef)?, + stream.clone_htod(d.mult_term_col)?, + ) + }; sync_if(timing)?; let t_desc = std::time::Instant::now(); let num_rows_u32 = num_rows as u32; @@ -379,53 +412,59 @@ pub fn logup_aux_resident( let t_term = std::time::Instant::now(); // row_sum over all term columns → additive scan → accumulated column. - let mut row_sum = unsafe { stream.alloc::(num_rows * 3) }?; - unsafe { - stream - .launch_builder(&be.logup_row_sum_ext3) - .arg(&terms) - .arg(&num_out_u32) - .arg(&num_rows_u32) - .arg(&mut row_sum) - .launch(cfg(num_rows)?)?; - } - scan_add_inplace(stream, be, &mut row_sum, num_rows)?; // row_sum now holds S - let (i0, i1, i2) = (inv_n[0], inv_n[1], inv_n[2]); - let mut accumulated = unsafe { stream.alloc::(num_rows * 3) }?; - let n_u64 = num_rows as u64; - unsafe { - stream - .launch_builder(&be.logup_finalize_accum_ext3) - .arg(&row_sum) - .arg(&n_u64) - .arg(&i0) - .arg(&i1) - .arg(&i2) - .arg(&mut accumulated) - .launch(cfg(num_rows)?)?; - } - - // Assemble row-major aux buffer: committed (num_out-1) cols + accumulated. let num_committed = num_out - 1; let num_aux_cols = num_committed + 1; - let mut aux = unsafe { stream.alloc::(num_aux_cols * num_rows * 3) }?; - let num_committed_u32 = num_committed as u32; - unsafe { - stream - .launch_builder(&be.logup_assemble_aux_ext3) - .arg(&terms) - .arg(&num_committed_u32) - .arg(&accumulated) - .arg(&num_rows_u32) - .arg(&mut aux) - .launch(cfg(num_rows)?)?; + let mut row_sum; + let mut aux; + { + row_sum = unsafe { stream.alloc::(num_rows * 3) }?; + unsafe { + stream + .launch_builder(&be.logup_row_sum_ext3) + .arg(&terms) + .arg(&num_out_u32) + .arg(&num_rows_u32) + .arg(&mut row_sum) + .launch(cfg(num_rows)?)?; + } + scan_add_inplace(stream, be, &mut row_sum, num_rows)?; // row_sum now holds S + let (i0, i1, i2) = (inv_n[0], inv_n[1], inv_n[2]); + let mut accumulated = unsafe { stream.alloc::(num_rows * 3) }?; + let n_u64 = num_rows as u64; + unsafe { + stream + .launch_builder(&be.logup_finalize_accum_ext3) + .arg(&row_sum) + .arg(&n_u64) + .arg(&i0) + .arg(&i1) + .arg(&i2) + .arg(&mut accumulated) + .launch(cfg(num_rows)?)?; + } + + // Assemble row-major aux buffer: committed (num_out-1) cols + accumulated. + aux = unsafe { stream.alloc::(num_aux_cols * num_rows * 3) }?; + let num_committed_u32 = num_committed as u32; + unsafe { + stream + .launch_builder(&be.logup_assemble_aux_ext3) + .arg(&terms) + .arg(&num_committed_u32) + .arg(&accumulated) + .arg(&num_rows_u32) + .arg(&mut aux) + .launch(cfg(num_rows)?)?; + } } sync_if(timing)?; let t_accum_done = std::time::Instant::now(); // L = table_contribution = S[n-1] (sum of all term columns, all rows). - let l_host: Vec = stream.clone_dtoh(&row_sum.slice((num_rows - 1) * 3..num_rows * 3))?; - stream.synchronize()?; + let l_host: Vec = { stream.clone_dtoh(&row_sum.slice((num_rows - 1) * 3..num_rows * 3))? }; + { + stream.synchronize()?; + } if timing { let t_end = std::time::Instant::now(); let ms = |a: std::time::Instant, b: std::time::Instant| (b - a).as_secs_f64() * 1e3; diff --git a/crypto/math-cuda/src/merkle.rs b/crypto/math-cuda/src/merkle.rs index fb1125ea4..d08d50530 100644 --- a/crypto/math-cuda/src/merkle.rs +++ b/crypto/math-cuda/src/merkle.rs @@ -57,7 +57,7 @@ pub fn keccak_leaves_base( assert!(columns.len() >= total); let be = backend()?; let stream = be.next_stream(); - let cols_dev = stream.clone_htod(&columns[..total])?; + let cols_dev = { stream.clone_htod(&columns[..total])? }; let mut out_dev = stream.alloc_zeros::((num_rows / rows_per_leaf) * 32)?; let launch = if rows_per_leaf == 2 { launch_keccak_base_row_pair @@ -72,8 +72,10 @@ pub fn keccak_leaves_base( num_rows as u64, &mut out_dev.as_view_mut(), )?; - let out = stream.clone_dtoh(&out_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&out_dev)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -108,7 +110,7 @@ pub fn keccak_leaves_ext3( assert!(columns.len() >= total); let be = backend()?; let stream = be.next_stream(); - let cols_dev = stream.clone_htod(&columns[..total])?; + let cols_dev = { stream.clone_htod(&columns[..total])? }; let mut out_dev = stream.alloc_zeros::((num_rows / rows_per_leaf) * 32)?; let launch = if rows_per_leaf == 2 { launch_keccak_ext3_row_pair @@ -123,8 +125,10 @@ pub fn keccak_leaves_ext3( num_rows as u64, &mut out_dev.as_view_mut(), )?; - let out = stream.clone_dtoh(&out_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&out_dev)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -312,8 +316,10 @@ pub fn build_merkle_tree_on_device(hashed_leaves: &[u8]) -> Result> { build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, leaves_len)?; - let out = stream.clone_dtoh(&nodes_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&nodes_dev)? }; + { + stream.synchronize()?; + } Ok(out) } @@ -373,8 +379,16 @@ pub fn gather_merkle_paths_dev( .arg(&mut out) .launch(cfg)?; } - let host = stream.clone_dtoh(&out)?; - stream.synchronize()?; + // Async drain via the pinned-hashes slot (path nodes are hash output): + // enqueued without blocking, then the host waits only on the copy's event + // (which also covers the gather kernel queued before it) instead of a + // full stream sync. + let pending = + crate::device::async_dtoh_via(stream, be.pinned_hashes(), &be.ctx, &out, out.len())?; + let mut host = vec![0u8; out.len()]; + { + pending.wait_into_bytes(&mut host)?; + } Ok(host) } @@ -418,8 +432,10 @@ fn build_comp_poly_tree_nodes_dev( // below read the device `buf`, not `pinned`). Synchronize first so the // async H2D has consumed `pinned` before it is freed/reused. let mut buf = stream.alloc_zeros::(mb * lde_size)?; - stream.memcpy_htod(&pinned[..mb * lde_size], &mut buf)?; - stream.synchronize()?; + { + stream.memcpy_htod(&pinned[..mb * lde_size], &mut buf)?; + stream.synchronize()?; + } drop(staging); // Leaves into tail of a tight node buffer. diff --git a/crypto/math-cuda/src/ntt.rs b/crypto/math-cuda/src/ntt.rs index f5accc3b1..60cfe6eb3 100644 --- a/crypto/math-cuda/src/ntt.rs +++ b/crypto/math-cuda/src/ntt.rs @@ -76,23 +76,27 @@ fn ntt_inplace(input: &[u64], forward: bool) -> Result> { let be = backend()?; let stream = be.next_stream(); - let mut x_dev = stream.clone_htod(input)?; - let tw_dev = if forward { - be.fwd_twiddles_for(log_n)? - } else { - be.inv_twiddles_for(log_n)? + let mut x_dev = { stream.clone_htod(input)? }; + let tw_dev = { + if forward { + be.fwd_twiddles_for(log_n)? + } else { + be.inv_twiddles_for(log_n)? + } }; let n_u64 = n as u64; // 1. Bit-reverse: natural → bit-reversed. - unsafe { - stream - .launch_builder(&be.bit_reverse_permute) - .arg(&mut x_dev) - .arg(&n_u64) - .arg(&log_n) - .launch(LaunchConfig::for_num_elems(n as u32))?; + { + unsafe { + stream + .launch_builder(&be.bit_reverse_permute) + .arg(&mut x_dev) + .arg(&n_u64) + .arg(&log_n) + .launch(LaunchConfig::for_num_elems(n as u32))?; + } } // 2. DIT butterfly levels. For log_n >= 8 we fuse 8 levels per kernel via @@ -114,8 +118,10 @@ fn ntt_inplace(input: &[u64], forward: bool) -> Result> { } } - let out = stream.clone_dtoh(&x_dev)?; - stream.synchronize()?; + let out = { stream.clone_dtoh(&x_dev)? }; + { + stream.synchronize()?; + } Ok(out) } diff --git a/crypto/math-cuda/src/nvtx.rs b/crypto/math-cuda/src/nvtx.rs new file mode 100644 index 000000000..fa81f555c --- /dev/null +++ b/crypto/math-cuda/src/nvtx.rs @@ -0,0 +1,276 @@ +//! Minimal NVTX bindings so Nsight Systems timelines show named host-side +//! ranges (prover phases, kernel-pipeline entry points) instead of a wall of +//! anonymous CUDA API calls. +//! +//! Loading mirrors the crate's cudarc `dynamic-loading` philosophy: no +//! build-time or link-time dependency on the CUDA toolkit layout. At first use +//! we dlopen `libnvToolsExt.so` (NVTX v2, shipped with every CUDA toolkit and +//! honored by nsys/ncu); when it is absent every call is a cheap no-op, so a +//! `--features nvtx` binary runs unchanged on machines without the library. +//! Override the library path with `LAMBDA_VM_NVTX_LIB` if it lives somewhere +//! unusual. +//! +//! With the `nvtx` cargo feature *disabled* this module compiles to empty +//! inline stubs — the label closures passed to [`Range::fmt`] are never +//! evaluated and the whole thing vanishes. +//! +//! Semantics note: a [`Range`] measures the *host-side* extent of a call. For +//! the `_keep`/`_dev` entry points that enqueue async GPU work without +//! syncing, kernels execute after the range closes — that is fine: nsys +//! correlates each kernel to the range that launched it. + +pub use imp::*; + +#[cfg(feature = "nvtx")] +mod imp { + use std::ffi::{CString, c_char, c_int}; + use std::marker::PhantomData; + use std::path::PathBuf; + use std::sync::OnceLock; + + struct Api { + // Field order is drop order; the fn pointers are only valid while the + // library is loaded, and both live for the whole process anyway + // (static OnceLock). + range_push: unsafe extern "C" fn(*const c_char) -> c_int, + range_pop: unsafe extern "C" fn() -> c_int, + mark: unsafe extern "C" fn(*const c_char), + _lib: libloading::Library, + } + // SAFETY: the NVTX v2 API is thread-safe (push/pop stacks are per-thread) + // and the Library handle is only kept alive, never re-entered. + unsafe impl Send for Api {} + unsafe impl Sync for Api {} + + fn nvtx_lib_candidates() -> Vec { + let mut c = Vec::new(); + if let Some(p) = std::env::var_os("LAMBDA_VM_NVTX_LIB") { + c.push(PathBuf::from(p)); + } + c.push(PathBuf::from("libnvToolsExt.so.1")); + c.push(PathBuf::from("libnvToolsExt.so")); + let cuda_home = std::env::var_os("CUDA_HOME") + .or_else(|| std::env::var_os("CUDA_PATH")) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/usr/local/cuda")); + c.push(cuda_home.join("lib64").join("libnvToolsExt.so.1")); + c.push(cuda_home.join("lib64").join("libnvToolsExt.so")); + c + } + + fn api() -> Option<&'static Api> { + static API: OnceLock> = OnceLock::new(); + API.get_or_init(|| { + for path in nvtx_lib_candidates() { + // SAFETY: loading a shared library runs its initializers; + // libnvToolsExt is NVIDIA's stub dispatcher with no side + // effects beyond tool injection. + let Ok(lib) = (unsafe { libloading::Library::new(&path) }) else { + continue; + }; + // SAFETY: signatures match the NVTX v2 C API. + let syms = unsafe { + ( + lib.get:: c_int>( + b"nvtxRangePushA\0", + ) + .map(|s| *s), + lib.get:: c_int>(b"nvtxRangePop\0") + .map(|s| *s), + lib.get::(b"nvtxMarkA\0") + .map(|s| *s), + ) + }; + if let (Ok(range_push), Ok(range_pop), Ok(mark)) = syms { + return Some(Api { + range_push, + range_pop, + mark, + _lib: lib, + }); + } + } + None + }) + .as_ref() + } + + /// True when libnvToolsExt was found; use to skip label formatting work. + #[inline] + pub fn is_active() -> bool { + api().is_some() + } + + fn push_str(api: &Api, name: &str) { + // NVTX takes a NUL-terminated C string; a label containing NUL is a + // bug we don't care to surface here — fall back to a fixed name. + let c = CString::new(name).unwrap_or_else(|_| CString::new("invalid-label").unwrap()); + // SAFETY: `c` is a valid NUL-terminated string for the duration of the call. + unsafe { (api.range_push)(c.as_ptr()) }; + } + + /// Push a range on this thread's NVTX stack. Prefer [`Range`]; this raw + /// form exists for RAII guards that live in other crates (instruments). + #[inline] + pub fn range_push(name: &str) { + if let Some(api) = api() { + push_str(api, name); + } + } + + /// Pop this thread's innermost NVTX range. Must pair with [`range_push`]. + #[inline] + pub fn range_pop() { + if let Some(api) = api() { + // SAFETY: no arguments; unbalanced pops are handled by NVTX (no-op). + unsafe { (api.range_pop)() }; + } + } + + /// Instantaneous marker on the timeline. + #[inline] + pub fn mark(name: &str) { + if let Some(api) = api() { + let c = CString::new(name).unwrap_or_else(|_| CString::new("invalid-label").unwrap()); + // SAFETY: `c` is a valid NUL-terminated string for the duration of the call. + unsafe { (api.mark)(c.as_ptr()) }; + } + } + + /// RAII NVTX range: pushed on construction, popped on drop. `!Send` on + /// purpose — NVTX push/pop stacks are per-thread, so a guard must drop on + /// the thread that created it. + pub struct Range { + pushed: bool, + _not_send: PhantomData<*const ()>, + } + + impl Range { + #[inline] + pub fn new(name: &str) -> Range { + let pushed = api().map(|a| push_str(a, name)).is_some(); + Range { + pushed, + _not_send: PhantomData, + } + } + + /// Like [`Range::new`] but the label is only formatted when a + /// profiler-visible NVTX library is actually loaded. + #[inline] + pub fn fmt String>(label: F) -> Range { + if is_active() { + Range::new(&label()) + } else { + Range { + pushed: false, + _not_send: PhantomData, + } + } + } + } + + impl Drop for Range { + fn drop(&mut self) { + if self.pushed { + range_pop(); + } + } + } + + // --- CUDA profiler capture-range control ------------------------------- + // + // cuProfilerStart/Stop gate `nsys profile --capture-range=cudaProfilerApi`, + // letting a session capture one phase/epoch of a long prove instead of the + // whole run. Loaded from libcuda (already resident via cudarc) — separate + // dlopen so this module stays independent of cudarc's bound symbol set. + + struct ProfilerApi { + start: unsafe extern "C" fn() -> c_int, + stop: unsafe extern "C" fn() -> c_int, + _lib: libloading::Library, + } + unsafe impl Send for ProfilerApi {} + unsafe impl Sync for ProfilerApi {} + + fn profiler_api() -> Option<&'static ProfilerApi> { + static API: OnceLock> = OnceLock::new(); + API.get_or_init(|| { + for name in ["libcuda.so.1", "libcuda.so"] { + // SAFETY: libcuda is loaded by cudarc already; this bumps a refcount. + let Ok(lib) = (unsafe { libloading::Library::new(name) }) else { + continue; + }; + // SAFETY: signatures match the CUDA driver profiler API. + let syms = unsafe { + ( + lib.get:: c_int>(b"cuProfilerStart\0") + .map(|s| *s), + lib.get:: c_int>(b"cuProfilerStop\0") + .map(|s| *s), + ) + }; + if let (Ok(start), Ok(stop)) = syms { + return Some(ProfilerApi { + start, + stop, + _lib: lib, + }); + } + } + None + }) + .as_ref() + } + + /// Begin a profiler capture range (`nsys --capture-range=cudaProfilerApi`). + /// No-op without libcuda or outside a profiler session. + #[inline] + pub fn profiler_start() { + if let Some(api) = profiler_api() { + // SAFETY: no arguments; valid to call any time after libcuda loads. + unsafe { (api.start)() }; + } + } + + /// End a profiler capture range. Must pair with [`profiler_start`]. + #[inline] + pub fn profiler_stop() { + if let Some(api) = profiler_api() { + // SAFETY: no arguments; valid to call any time after libcuda loads. + unsafe { (api.stop)() }; + } + } +} + +#[cfg(not(feature = "nvtx"))] +mod imp { + /// No-op stub; see the `nvtx`-feature implementation above. + pub struct Range; + + impl Range { + #[inline(always)] + pub fn new(_name: &str) -> Range { + Range + } + #[inline(always)] + pub fn fmt String>(_label: F) -> Range { + Range + } + } + + #[inline(always)] + pub fn is_active() -> bool { + false + } + #[inline(always)] + pub fn range_push(_name: &str) {} + #[inline(always)] + pub fn range_pop() {} + #[inline(always)] + pub fn mark(_name: &str) {} + #[inline(always)] + pub fn profiler_start() {} + #[inline(always)] + pub fn profiler_stop() {} +} diff --git a/crypto/math-cuda/tests/barycentric_strided.rs b/crypto/math-cuda/tests/barycentric_strided.rs index d96f7128b..024eb77e8 100644 --- a/crypto/math-cuda/tests/barycentric_strided.rs +++ b/crypto/math-cuda/tests/barycentric_strided.rs @@ -46,6 +46,7 @@ fn run_base(log_trace: u32, blowup: usize, num_cols: usize, seed: u64) { let lde_dev = stream.clone_htod(&lde_flat).unwrap(); stream.synchronize().unwrap(); let handle = GpuLdeBase { + ready: None, buf: Arc::new(lde_dev), m: num_cols, lde_size, @@ -105,6 +106,7 @@ fn run_ext3(log_trace: u32, blowup: usize, num_cols: usize, seed: u64) { let lde_dev = stream.clone_htod(&lde_flat).unwrap(); stream.synchronize().unwrap(); let handle = GpuLdeExt3 { + ready: None, buf: Arc::new(lde_dev), m: num_cols, lde_size, diff --git a/crypto/math-cuda/tests/deep.rs b/crypto/math-cuda/tests/deep.rs index b7c027914..f7e163564 100644 --- a/crypto/math-cuda/tests/deep.rs +++ b/crypto/math-cuda/tests/deep.rs @@ -174,6 +174,7 @@ fn run_parity( stream.synchronize().unwrap(); let main_handle = GpuLdeBase { + ready: None, buf: Arc::new(main_dev), m: num_main, lde_size, @@ -183,6 +184,7 @@ fn run_parity( }; let aux_handle = if num_aux > 0 { Some(GpuLdeExt3 { + ready: None, buf: Arc::new(aux_dev), m: num_aux, lde_size, diff --git a/crypto/math-cuda/tests/gather_rows.rs b/crypto/math-cuda/tests/gather_rows.rs index 8b7f1b4a3..fe76c5898 100644 --- a/crypto/math-cuda/tests/gather_rows.rs +++ b/crypto/math-cuda/tests/gather_rows.rs @@ -23,6 +23,7 @@ fn run_base(lde_size: usize, num_cols: usize, seed: u64) { let dev = stream.clone_htod(&buf).unwrap(); stream.synchronize().unwrap(); let handle = GpuLdeBase { + ready: None, buf: Arc::new(dev), m: num_cols, lde_size, @@ -57,6 +58,7 @@ fn run_ext3(lde_size: usize, num_cols: usize, seed: u64) { let dev = stream.clone_htod(&buf).unwrap(); stream.synchronize().unwrap(); let handle = GpuLdeExt3 { + ready: None, buf: Arc::new(dev), m: num_cols, lde_size, diff --git a/crypto/stark/Cargo.toml b/crypto/stark/Cargo.toml index 78d95be67..64e650aa7 100644 --- a/crypto/stark/Cargo.toml +++ b/crypto/stark/Cargo.toml @@ -56,6 +56,10 @@ debug-checks = [] # Enables v parallel = ["dep:rayon", "crypto/parallel"] cuda = ["dep:math-cuda"] test-cuda-faults = ["cuda", "math-cuda/test-faults"] +# NVTX ranges for Nsight Systems: math-cuda pipeline entry points plus every +# instruments span (prover phases) become named timeline ranges. Pulls in +# `instruments` so the span tree exists to mirror, and `cuda` for math-cuda. +nvtx = ["cuda", "instruments", "math-cuda/nvtx"] wasm = ["dep:wasm-bindgen", "dep:serde-wasm-bindgen", "dep:web-sys"] disk-spill = ["dep:memmap2", "dep:tempfile", "dep:libc", "crypto/disk-spill"] diff --git a/crypto/stark/src/constraint_ir/device.rs b/crypto/stark/src/constraint_ir/device.rs index 6c522d103..e4e170bfc 100644 --- a/crypto/stark/src/constraint_ir/device.rs +++ b/crypto/stark/src/constraint_ir/device.rs @@ -8,24 +8,43 @@ //! field tower never reaches this module — it stays on the generic //! [`interp`](super::interp) path. //! +//! ## Slot-based value layout (dim-split) +//! +//! The kernel keeps per-thread value scratch in global memory, so its size and +//! traffic are the dominant cost of the constraint walk. The lowering therefore +//! does three things beyond serializing ops: +//! +//! - **Dim split**: a node's value lives in a *base* (`u64`) or *ext* +//! (`[u64; 3]`) slot class according to its [`Dim`] tag, and arithmetic on +//! base nodes is base-field arithmetic. Because embedding is a ring +//! homomorphism and the device field ops are bit-identical to the CPU's, +//! this is bit-for-bit equal to the all-ext evaluation it replaces, at a +//! third of the scratch traffic and ~1/9 of the multiply cost for base +//! nodes. +//! - **Liveness slot reuse**: slots are assigned by a linear scan that frees +//! an operand's slot at its last use, so the scratch working set is the +//! program's max-live-set, not its node count (root nodes are pinned: both +//! kernels read them after the walk). +//! - **Uniform propagation**: row-invariant leaves (constants, RAP +//! challenges, LogUp alpha powers, the table offset) never materialize as +//! nodes or slots; operands reference the tiny uniform tables directly. +//! They only stay as nodes in the degenerate case where one is itself a +//! constraint root. +//! //! Two things live here: //! -//! - [`DeviceProgram::lower`] — serialize a `ConstraintProgram` into flat, device-uploadable arrays: a `#[repr(C)]` -//! [`DeviceNode`] list, `u64` / `[u64; 3]` constant tables, `roots`, and -//! `num_base`. Field constants become raw limbs via `FieldElement::to_raw`, -//! which is byte-identical to how [`crate::gpu_lde`] already hands Goldilocks -//! elements to the device (a `#[repr(transparent)]` `u64` / `[u64; 3]`). +//! - [`DeviceProgram::lower`] — the lowering itself. Field constants become +//! raw limbs via `FieldElement::to_raw`, byte-identical to how +//! [`crate::gpu_lde`] hands Goldilocks elements to the device. //! //! - [`eval_device_program`] — a CPU forward pass over the *flat* node array -//! (not the [`Op`] enum), decoding leaves from the raw limb tables and -//! reproducing the exact [`interp::run`](super::interp) semantics -//! (per-node [`Dim`] drives base-vs-extension arithmetic, mixed operands -//! auto-embed). It is the model of the GPU kernel's per-thread walk, so a -//! bit-for-bit match against [`eval_program`](super::interp::eval_program) -//! pins the on-device layout and control flow *before* any CUDA exists. The -//! kernel is then a transliteration of this walk with the `FieldElement` -//! arithmetic swapped for `goldilocks.cuh` / `ext3.cuh`. +//! (not the [`Op`] enum), decoding operands exactly as the kernel does and +//! reproducing the [`interp::run`](super::interp) semantics. It is the model +//! of the GPU kernel's per-thread walk, so a bit-for-bit match against +//! [`eval_program`](super::interp::eval_program) pins the on-device layout +//! and control flow *before* any CUDA runs. The kernel is a transliteration +//! of this walk with the `FieldElement` arithmetic swapped for +//! `goldilocks.cuh` / `ext3.cuh`. use math::field::element::FieldElement; use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as GoldilocksExtension; @@ -37,48 +56,70 @@ type FpE = FieldElement; type Ext3E = FieldElement; // ------------------------------------------------------------------------- -// Wire tags — MUST match the CUDA kernel's `switch (op)` and dim checks. +// Wire tags — MUST match the CUDA kernel's `switch (op)` and operand decode. // ------------------------------------------------------------------------- -/// `a` = index into `base_consts`. +/// `a` = index into `base_consts` (only when a uniform leaf is itself a root). pub const OP_CONST_BASE: u32 = 0; -/// `a` = index into `ext_consts`. +/// `a` = index into `ext_consts` (only when a uniform leaf is itself a root). pub const OP_CONST_EXT: u32 = 1; /// Trace-cell read; `a`/`b` pack the [`Op::Var`] fields (see [`pack_var`]). pub const OP_VAR: u32 = 2; -/// `a` = index into the per-proof `rap_challenges` uniform buffer. +/// `a` = index into the per-proof `rap_challenges` uniform buffer (root-only). pub const OP_RAP_CHALLENGE: u32 = 3; -/// `a` = index into the per-proof `logup_alpha_powers` uniform buffer. +/// `a` = index into the per-proof `logup_alpha_powers` uniform buffer +/// (root-only). pub const OP_ALPHA_POW: u32 = 4; -/// The per-proof LogUp table offset uniform; no operands. +/// The per-proof LogUp table offset uniform; no operands (root-only). pub const OP_TABLE_OFFSET: u32 = 5; -/// `a`, `b` = node ids. +/// `a`, `b` = encoded operands (see `OPK_*`). pub const OP_ADD: u32 = 6; -/// `a`, `b` = node ids. +/// `a`, `b` = encoded operands. pub const OP_SUB: u32 = 7; -/// `a`, `b` = node ids. +/// `a`, `b` = encoded operands. pub const OP_MUL: u32 = 8; -/// `a` = node id. +/// `a` = encoded operand. pub const OP_NEG: u32 = 9; -/// `a` = node id (base → extension embed). +/// `a` = encoded operand (base → extension embed). pub const OP_EMBED: u32 = 10; -/// Node result is a base-field value. -pub const DIM_BASE: u32 = 0; -/// Node result is an extension-field value. -pub const DIM_EXT: u32 = 1; +// -- operand encoding: `kind << OPK_SHIFT | payload` ---------------------- + +/// Bit position of the 3-bit operand kind. +pub const OPK_SHIFT: u32 = 29; +/// Mask of the 29-bit operand payload (slot or table index). +pub const OPK_PAYLOAD_MASK: u32 = (1 << OPK_SHIFT) - 1; +/// Payload = base (`u64`) scratch-slot index. +pub const OPK_BASE_SLOT: u32 = 0; +/// Payload = ext (`[u64; 3]`) scratch-slot index. +pub const OPK_EXT_SLOT: u32 = 1; +/// Payload = `base_consts` index. +pub const OPK_BASE_CONST: u32 = 2; +/// Payload = `ext_consts` index. +pub const OPK_EXT_CONST: u32 = 3; +/// Payload = per-proof `rap_challenges` index. +pub const OPK_RAP: u32 = 4; +/// Payload = per-proof `logup_alpha_powers` index. +pub const OPK_ALPHA: u32 = 5; +/// The per-proof table offset (payload unused). +pub const OPK_OFFSET: u32 = 6; + +/// In a node's `res` word and in `roots` entries: bit 31 set = ext slot, +/// clear = base slot; low bits = the slot index. +pub const RES_EXT_BIT: u32 = 1 << 31; /// One flattened IR instruction: 16 bytes, `#[repr(C)]` for a 1:1 device -/// upload. `op` is an `OP_*` tag; the meaning of `a`/`b` depends on `op` (node -/// ids for arithmetic, table indices for constants/uniforms, packed [`Op::Var`] -/// fields for [`OP_VAR`]); `dim` is [`DIM_BASE`] or [`DIM_EXT`]. +/// upload. `op` is an `OP_*` tag; `a`/`b` are encoded operands (`OPK_*` kinds +/// for arithmetic, packed [`Op::Var`] fields for [`OP_VAR`], raw table indices +/// for root-pinned uniform leaves); `res` is the result slot with [`RES_EXT_BIT`] +/// selecting the slot class. #[repr(C)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct DeviceNode { pub op: u32, pub a: u32, pub b: u32, - pub dim: u32, + pub res: u32, } /// Pack an [`Op::Var`]'s fields into the `(a, b)` operand words: @@ -101,59 +142,204 @@ pub fn unpack_var(a: u32, b: u32) -> (bool, u8, u8, u16) { (main, offset, row, col) } -/// A [`ConstraintProgram`] lowered to flat, device-uploadable arrays. Constants -/// are canonical raw limbs (`u64` base / `[u64; 3]` extension), matching the -/// `#[repr(transparent)]` layout the GPU trace buffers already use. +/// A [`ConstraintProgram`] lowered to flat, device-uploadable arrays with +/// dim-split, liveness-reused value slots. Constants are canonical raw limbs +/// (`u64` base / `[u64; 3]` extension), matching the `#[repr(transparent)]` +/// layout the GPU trace buffers already use. #[derive(Clone, Debug)] pub struct DeviceProgram { - /// Topologically ordered instruction list (id `i` references only `< i`). + /// Topologically ordered instruction list (operands reference slots + /// already written or uniform tables). Uniform leaves and dead nodes are + /// not materialized. pub nodes: Vec, - /// Base-field constant table, indexed by [`OP_CONST_BASE`]. + /// Base-field constant table, indexed by [`OPK_BASE_CONST`] operands (and + /// [`OP_CONST_BASE`] root nodes). pub base_consts: Vec, - /// Extension-field constant table, indexed by [`OP_CONST_EXT`]. + /// Extension-field constant table, indexed by [`OPK_EXT_CONST`] operands + /// (and [`OP_CONST_EXT`] root nodes). pub ext_consts: Vec<[u64; 3]>, - /// Per-constraint root node ids, indexed by `constraint_idx`. + /// Per-constraint root slots (`slot | RES_EXT_BIT`), indexed by + /// `constraint_idx`. Root slots are pinned — never reused — so both + /// kernels can read them after the walk. pub roots: Vec, - /// Number of leading ([`DIM_BASE`]-rooted) constraints written to - /// `base_evals`; the rest go to `ext_evals`. + /// Number of leading base-rooted constraints written to `base_evals`; the + /// rest go to `ext_evals`. pub num_base: u32, + /// Size of the base (`u64`) slot class, per thread. + pub num_base_slots: u32, + /// Size of the ext (`[u64; 3]`) slot class, per thread. + pub num_ext_slots: u32, +} + +/// Whether an op is a row-invariant leaf (uniform per proof). +fn is_uniform_leaf(op: &Op) -> bool { + matches!( + op, + Op::ConstBase(_) + | Op::ConstExt(_) + | Op::RapChallenge { .. } + | Op::AlphaPow { .. } + | Op::TableOffset + ) +} + +/// The (up to two) operand node ids of an op. +fn operands(op: &Op) -> [Option; 2] { + match *op { + Op::Add(a, b) | Op::Sub(a, b) | Op::Mul(a, b) => [Some(a), Some(b)], + Op::Neg(a) | Op::Embed(a) => [Some(a), None], + _ => [None, None], + } } impl DeviceProgram { /// Lower a concrete-Goldilocks [`ConstraintProgram`] to its flat device - /// form. Pure serialization — no field arithmetic, no device access. + /// form: dim-split slot assignment with liveness reuse, uniform-leaf + /// propagation into operands, and root pinning. Pure serialization plus + /// the slot scan — no field arithmetic, no device access. pub fn lower(prog: &ConstraintProgram) -> Self { - let nodes = prog - .nodes - .iter() - .zip(prog.dims.iter()) - .map(|(op, dim)| { - let dim = match dim { - Dim::Base => DIM_BASE, - Dim::Ext => DIM_EXT, - }; - let (op, a, b) = match *op { - Op::ConstBase(idx) => (OP_CONST_BASE, idx, 0), - Op::ConstExt(idx) => (OP_CONST_EXT, idx, 0), - Op::Var { - main, - offset, - row, - col, - } => { - let (a, b) = pack_var(main, offset, row, col); - (OP_VAR, a, b) + let n = prog.nodes.len(); + assert!( + n <= OPK_PAYLOAD_MASK as usize, + "program of {n} nodes exceeds the 29-bit slot space" + ); + + // Liveness: last consumer (by node id) of every node, plus root pins. + let mut used = vec![false; n]; + let mut last_use = vec![0u32; n]; + for (i, op) in prog.nodes.iter().enumerate() { + for operand in operands(op).into_iter().flatten() { + used[operand as usize] = true; + last_use[operand as usize] = i as u32; + } + } + let mut is_root = vec![false; n]; + for &r in &prog.roots { + is_root[r as usize] = true; + used[r as usize] = true; + } + + // A node materializes (gets a slot) unless it is a propagated uniform + // leaf or dead. Uniform leaves stay only when they are roots (the + // post-walk emit reads slots). + let emitted: Vec = (0..n) + .map(|i| used[i] && (!is_uniform_leaf(&prog.nodes[i]) || is_root[i])) + .collect(); + + let enc_uniform = |op: &Op| -> u32 { + match *op { + Op::ConstBase(idx) => { + debug_assert!(idx <= OPK_PAYLOAD_MASK); + (OPK_BASE_CONST << OPK_SHIFT) | idx + } + Op::ConstExt(idx) => { + debug_assert!(idx <= OPK_PAYLOAD_MASK); + (OPK_EXT_CONST << OPK_SHIFT) | idx + } + Op::RapChallenge { idx } => (OPK_RAP << OPK_SHIFT) | idx as u32, + Op::AlphaPow { idx } => (OPK_ALPHA << OPK_SHIFT) | idx as u32, + Op::TableOffset => OPK_OFFSET << OPK_SHIFT, + _ => unreachable!("not a uniform leaf"), + } + }; + + // Linear-scan slot assignment with per-class free lists. + const UNASSIGNED: u32 = u32::MAX; + let mut slot_of = vec![UNASSIGNED; n]; + let mut free_base: Vec = Vec::new(); + let mut free_ext: Vec = Vec::new(); + let mut num_base_slots = 0u32; + let mut num_ext_slots = 0u32; + let mut nodes = Vec::with_capacity(n); + + for i in 0..n { + if !emitted[i] { + continue; + } + let op = &prog.nodes[i]; + let dim = prog.dims[i]; + + // Encode operands while their slots are still assigned. + let enc_operand = |j: u32| -> u32 { + let j = j as usize; + if !emitted[j] { + return enc_uniform(&prog.nodes[j]); + } + let slot = slot_of[j]; + debug_assert_ne!(slot, UNASSIGNED, "operand before definition"); + match prog.dims[j] { + Dim::Base => (OPK_BASE_SLOT << OPK_SHIFT) | slot, + Dim::Ext => (OPK_EXT_SLOT << OPK_SHIFT) | slot, + } + }; + + let (tag, a, b) = match *op { + Op::ConstBase(idx) => (OP_CONST_BASE, idx, 0), + Op::ConstExt(idx) => (OP_CONST_EXT, idx, 0), + Op::Var { + main, + offset, + row, + col, + } => { + let (a, b) = pack_var(main, offset, row, col); + (OP_VAR, a, b) + } + Op::RapChallenge { idx } => (OP_RAP_CHALLENGE, idx as u32, 0), + Op::AlphaPow { idx } => (OP_ALPHA_POW, idx as u32, 0), + Op::TableOffset => (OP_TABLE_OFFSET, 0, 0), + Op::Add(x, y) => (OP_ADD, enc_operand(x), enc_operand(y)), + Op::Sub(x, y) => (OP_SUB, enc_operand(x), enc_operand(y)), + Op::Mul(x, y) => (OP_MUL, enc_operand(x), enc_operand(y)), + Op::Neg(x) => (OP_NEG, enc_operand(x), 0), + Op::Embed(x) => (OP_EMBED, enc_operand(x), 0), + }; + + // Free operand slots at their last use (roots stay pinned). The + // `slot_of` reset guards the a == b double-free. + for operand in operands(op).into_iter().flatten() { + let j = operand as usize; + if emitted[j] && !is_root[j] && last_use[j] == i as u32 && slot_of[j] != UNASSIGNED + { + match prog.dims[j] { + Dim::Base => free_base.push(slot_of[j]), + Dim::Ext => free_ext.push(slot_of[j]), } - Op::RapChallenge { idx } => (OP_RAP_CHALLENGE, idx as u32, 0), - Op::AlphaPow { idx } => (OP_ALPHA_POW, idx as u32, 0), - Op::TableOffset => (OP_TABLE_OFFSET, 0, 0), - Op::Add(a, b) => (OP_ADD, a, b), - Op::Sub(a, b) => (OP_SUB, a, b), - Op::Mul(a, b) => (OP_MUL, a, b), - Op::Neg(a) => (OP_NEG, a, 0), - Op::Embed(a) => (OP_EMBED, a, 0), - }; - DeviceNode { op, a, b, dim } + slot_of[j] = UNASSIGNED; + } + } + + // Allocate the result slot (a freed operand slot may be reused — + // the kernel reads operands before writing the result). + let slot = match dim { + Dim::Base => free_base.pop().unwrap_or_else(|| { + num_base_slots += 1; + num_base_slots - 1 + }), + Dim::Ext => free_ext.pop().unwrap_or_else(|| { + num_ext_slots += 1; + num_ext_slots - 1 + }), + }; + slot_of[i] = slot; + + let res = match dim { + Dim::Base => slot, + Dim::Ext => slot | RES_EXT_BIT, + }; + nodes.push(DeviceNode { op: tag, a, b, res }); + } + + let roots = prog + .roots + .iter() + .map(|&r| { + let slot = slot_of[r as usize]; + debug_assert_ne!(slot, UNASSIGNED, "root without a slot"); + match prog.dims[r as usize] { + Dim::Base => slot, + Dim::Ext => slot | RES_EXT_BIT, + } }) .collect(); @@ -164,32 +350,10 @@ impl DeviceProgram { nodes, base_consts, ext_consts, - roots: prog.roots.clone(), + roots, num_base: prog.num_base as u32, - } - } -} - -/// A node's computed value during the walk: base or extension field element. -#[derive(Clone)] -enum Value { - Base(FpE), - Ext(Ext3E), -} - -impl Value { - /// Promote to the extension field, embedding a base value if needed. - fn to_ext(&self) -> Ext3E { - match self { - Value::Base(x) => (*x).to_extension::(), - Value::Ext(x) => *x, - } - } - - fn as_base(&self) -> &FpE { - match self { - Value::Base(x) => x, - Value::Ext(_) => panic!("expected a base value but found an extension value"), + num_base_slots, + num_ext_slots, } } } @@ -211,33 +375,12 @@ fn encode_ext(x: &Ext3E) -> [u64; 3] { [*limbs[0].value(), *limbs[1].value(), *limbs[2].value()] } -/// Apply a binary op, auto-embedding to the extension when the result dimension -/// is [`DIM_EXT`] (or either operand is already an extension value) — the exact -/// rule of [`interp::binop`](super::interp). -#[inline] -fn binop( - values: &[Value], - a: u32, - b: u32, - dim: u32, - base_op: impl Fn(FpE, FpE) -> FpE, - ext_op: impl Fn(Ext3E, Ext3E) -> Ext3E, -) -> Value { - let va = &values[a as usize]; - let vb = &values[b as usize]; - if dim == DIM_BASE - && let (Value::Base(x), Value::Base(y)) = (va, vb) - { - return Value::Base(base_op(*x, *y)); - } - Value::Ext(ext_op(va.to_ext(), vb.to_ext())) -} - /// Full prover-shaped forward pass over the *flat* device blob, in raw limbs — -/// the CPU model of the GPU kernel. Mirrors +/// the CPU model of the GPU kernel: dim-split slot files, encoded-operand +/// loads, mixed ops evaluated as full ext ops on embedded operands (the GPU's +/// mixed-op shortcuts are bit-identical to that by construction). Mirrors /// [`eval_program`](super::interp::eval_program): base-rooted constraints -/// (`c < num_base`) land in `base_evals`, the rest in `ext_evals`, with the -/// same auto-embed semantics. +/// (`c < num_base`) land in `base_evals`, the rest in `ext_evals`. /// /// `main[offset][col]` / `aux[offset][col]` are the frame's trace cells; /// `rap_challenges` / `alpha_powers` / `table_offset` are the per-proof @@ -254,71 +397,101 @@ pub fn eval_device_program( base_evals: &mut [u64], ext_evals: &mut [[u64; 3]], ) { - let mut values: Vec = Vec::with_capacity(dev.nodes.len()); + let mut base_slots = vec![FpE::zero(); dev.num_base_slots as usize]; + let mut ext_slots = vec![Ext3E::zero(); dev.num_ext_slots as usize]; + + let load_base = |enc: u32, base_slots: &[FpE]| -> FpE { + let payload = (enc & OPK_PAYLOAD_MASK) as usize; + match enc >> OPK_SHIFT { + OPK_BASE_SLOT => base_slots[payload], + OPK_BASE_CONST => FpE::from_raw(dev.base_consts[payload]), + other => panic!("base operand with non-base kind {other}"), + } + }; + let load_ext = |enc: u32, base_slots: &[FpE], ext_slots: &[Ext3E]| -> Ext3E { + let payload = (enc & OPK_PAYLOAD_MASK) as usize; + match enc >> OPK_SHIFT { + OPK_BASE_SLOT => base_slots[payload].to_extension::(), + OPK_EXT_SLOT => ext_slots[payload], + OPK_BASE_CONST => { + FpE::from_raw(dev.base_consts[payload]).to_extension::() + } + OPK_EXT_CONST => decode_ext(dev.ext_consts[payload]), + OPK_RAP => decode_ext(rap_challenges[payload]), + OPK_ALPHA => decode_ext(alpha_powers[payload]), + OPK_OFFSET => decode_ext(table_offset), + other => panic!("unknown operand kind {other}"), + } + }; for node in &dev.nodes { - let v = match node.op { - OP_CONST_BASE => Value::Base(FpE::from_raw(dev.base_consts[node.a as usize])), - OP_CONST_EXT => Value::Ext(decode_ext(dev.ext_consts[node.a as usize])), + let res_slot = (node.res & !RES_EXT_BIT) as usize; + let res_ext = node.res & RES_EXT_BIT != 0; + match node.op { + OP_CONST_BASE => base_slots[res_slot] = FpE::from_raw(dev.base_consts[node.a as usize]), + OP_CONST_EXT => ext_slots[res_slot] = decode_ext(dev.ext_consts[node.a as usize]), OP_VAR => { let (is_main, offset, _row, col) = unpack_var(node.a, node.b); if is_main { - Value::Base(FpE::from_raw(main[offset as usize][col as usize])) + base_slots[res_slot] = FpE::from_raw(main[offset as usize][col as usize]); + } else { + ext_slots[res_slot] = decode_ext(aux[offset as usize][col as usize]); + } + } + OP_RAP_CHALLENGE => ext_slots[res_slot] = decode_ext(rap_challenges[node.a as usize]), + OP_ALPHA_POW => ext_slots[res_slot] = decode_ext(alpha_powers[node.a as usize]), + OP_TABLE_OFFSET => ext_slots[res_slot] = decode_ext(table_offset), + OP_ADD => { + if res_ext { + ext_slots[res_slot] = load_ext(node.a, &base_slots, &ext_slots) + + load_ext(node.b, &base_slots, &ext_slots); + } else { + base_slots[res_slot] = + load_base(node.a, &base_slots) + load_base(node.b, &base_slots); + } + } + OP_SUB => { + if res_ext { + ext_slots[res_slot] = load_ext(node.a, &base_slots, &ext_slots) + - load_ext(node.b, &base_slots, &ext_slots); + } else { + base_slots[res_slot] = + load_base(node.a, &base_slots) - load_base(node.b, &base_slots); + } + } + OP_MUL => { + if res_ext { + ext_slots[res_slot] = load_ext(node.a, &base_slots, &ext_slots) + * load_ext(node.b, &base_slots, &ext_slots); } else { - Value::Ext(decode_ext(aux[offset as usize][col as usize])) + base_slots[res_slot] = + load_base(node.a, &base_slots) * load_base(node.b, &base_slots); } } - OP_RAP_CHALLENGE => Value::Ext(decode_ext(rap_challenges[node.a as usize])), - OP_ALPHA_POW => Value::Ext(decode_ext(alpha_powers[node.a as usize])), - OP_TABLE_OFFSET => Value::Ext(decode_ext(table_offset)), - OP_ADD => binop( - &values, - node.a, - node.b, - node.dim, - |x, y| x + y, - |x, y| x + y, - ), - OP_SUB => binop( - &values, - node.a, - node.b, - node.dim, - |x, y| x - y, - |x, y| x - y, - ), - OP_MUL => binop( - &values, - node.a, - node.b, - node.dim, - |x, y| x * y, - |x, y| x * y, - ), OP_NEG => { - let val = &values[node.a as usize]; - if node.dim == DIM_BASE { - match val { - Value::Base(x) => Value::Base(-x), - // Dim/value mismatch: keep it in the extension, as interp does. - Value::Ext(x) => Value::Ext(-*x), - } + if res_ext { + ext_slots[res_slot] = -load_ext(node.a, &base_slots, &ext_slots); } else { - Value::Ext(-val.to_ext()) + base_slots[res_slot] = -load_base(node.a, &base_slots); } } - OP_EMBED => Value::Ext(values[node.a as usize].to_ext()), + OP_EMBED => { + ext_slots[res_slot] = load_ext(node.a, &base_slots, &ext_slots); + } other => panic!("unknown device op tag {other}"), - }; - values.push(v); + } } for (c, &root) in dev.roots.iter().enumerate() { - let v = &values[root as usize]; + let slot = (root & !RES_EXT_BIT) as usize; + let is_ext = root & RES_EXT_BIT != 0; if (c as u32) < dev.num_base { - base_evals[c] = *v.as_base().value(); + assert!(!is_ext, "base-rooted constraint with an ext root slot"); + base_evals[c] = *base_slots[slot].value(); + } else if is_ext { + ext_evals[c] = encode_ext(&ext_slots[slot]); } else { - ext_evals[c] = encode_ext(&v.to_ext()); + ext_evals[c] = encode_ext(&base_slots[slot].to_extension::()); } } } @@ -474,4 +647,167 @@ mod tests { assert_eq!(ext_dev[2], encode_ext(&ext_ref[2])); } } + + /// Lowering invariants of the slot allocator: uniform leaves are + /// propagated (no nodes), the slot classes are bounded by the max-live-set + /// (strictly fewer slots than nodes for a program with dead-after-use + /// intermediates), and slot indices stay in range. + #[test] + fn lowering_reuses_slots_and_propagates_uniforms() { + let prog = all_ops_program(); + let dev = DeviceProgram::lower(&prog); + + // No uniform leaf is materialized (none is a root here). + for n in &dev.nodes { + assert!( + !matches!( + n.op, + OP_CONST_BASE + | OP_CONST_EXT + | OP_RAP_CHALLENGE + | OP_ALPHA_POW + | OP_TABLE_OFFSET + ), + "uniform leaf materialized as a node" + ); + } + // Slot classes are within bounds and smaller than the node count. + let total_slots = (dev.num_base_slots + dev.num_ext_slots) as usize; + assert!(total_slots < prog.nodes.len()); + for n in &dev.nodes { + let slot = n.res & !RES_EXT_BIT; + if n.res & RES_EXT_BIT != 0 { + assert!(slot < dev.num_ext_slots); + } else { + assert!(slot < dev.num_base_slots); + } + } + for &r in &dev.roots { + let slot = r & !RES_EXT_BIT; + if r & RES_EXT_BIT != 0 { + assert!(slot < dev.num_ext_slots); + } else { + assert!(slot < dev.num_base_slots); + } + } + } + + /// A uniform leaf that is itself a root must still materialize (the + /// post-walk emit reads a slot). + #[test] + fn uniform_root_is_materialized() { + let mut b = IrBuilder::::new(); + let c = b.const_base(7); + b.emit(0, c); + let prog = b.finish(1); + let dev = DeviceProgram::lower(&prog); + + assert!(dev.nodes.iter().any(|n| n.op == OP_CONST_BASE)); + let mut base_evals = vec![0u64; 1]; + let mut ext_evals: Vec<[u64; 3]> = vec![]; + eval_device_program( + &dev, + &[], + &[], + &[], + &[], + [0, 0, 0], + &mut base_evals, + &mut ext_evals, + ); + assert_eq!(base_evals[0], 7); + } + + /// Randomized differential: a synthetic DAG with heavy slot churn (long + /// chains whose intermediates die immediately) evaluates identically + /// through the interpreter and the slot-reusing device walk. + #[test] + fn slot_reuse_differential_random_chains() { + let mut b = IrBuilder::::new(); + let m0 = b.main(0, 0); + let m1 = b.main(0, 1); + let ch = b.challenge(0); + + // Base chain: alternating add/mul over rotating leaves. + let mut acc = m0; + for k in 0..50u64 { + let c = b.const_base(k + 2); + let t = if k % 2 == 0 { + b.add(acc, c) + } else { + b.mul(acc, m1) + }; + acc = t; + } + b.emit(0, acc); + + // Ext chain crossing dims each step. + let mut eacc = b.mul(m0, ch); + for k in 0..50u64 { + let c = b.const_base(k + 100); + let t = b.mul(eacc, c); // ext × base + let u = b.sub(t, ch); + eacc = u; + } + b.emit(1, eacc); + let prog = b.finish(1); + let dev = DeviceProgram::lower(&prog); + + // Slot reuse must keep the live-set small despite 100+ nodes. + assert!(dev.num_base_slots <= 8, "base slots {}", dev.num_base_slots); + assert!(dev.num_ext_slots <= 8, "ext slots {}", dev.num_ext_slots); + + let mut rng = SplitMix64(0xDEAD_BEEF_0BAD_F00D); + for _ in 0..500 { + let main_vals: Vec> = (0..2) + .map(|_| vec![fp(rng.next_u64()), fp(rng.next_u64())]) + .collect(); + let aux_vals: Vec> = (0..2).map(|_| vec![rng.ext()]).collect(); + let rap = vec![rng.ext()]; + let alpha = vec![rng.ext()]; + let offset = rng.ext(); + + let steps: Vec> = main_vals + .iter() + .zip(aux_vals.iter()) + .map(|(m, a)| TableView::::new(vec![m.clone()], vec![a.clone()])) + .collect(); + let frame = Frame::::new(steps); + let ctx = TransitionEvaluationContext::new_prover( + frame.as_row_frame(), + &rap, + &alpha, + &offset, + ); + let mut base_ref = vec![FpE::zero(); 1]; + let mut ext_ref = vec![Ext3E::zero(); 2]; + eval_program(&prog, &ctx, &mut base_ref, &mut ext_ref); + + let main_raw: Vec> = main_vals + .iter() + .map(|r| r.iter().map(|x| *x.value()).collect()) + .collect(); + let aux_raw: Vec> = aux_vals + .iter() + .map(|r| r.iter().map(encode_ext).collect()) + .collect(); + let rap_raw: Vec<[u64; 3]> = rap.iter().map(encode_ext).collect(); + let alpha_raw: Vec<[u64; 3]> = alpha.iter().map(encode_ext).collect(); + let mut base_dev = vec![0u64; 1]; + let mut ext_dev = vec![[0u64; 3]; 2]; + eval_device_program( + &dev, + &main_raw, + &aux_raw, + &rap_raw, + &alpha_raw, + encode_ext(&offset), + &mut base_dev, + &mut ext_dev, + ); + + assert_eq!(base_dev[0], *base_ref[0].value()); + assert_eq!(ext_dev[1], encode_ext(&ext_ref[1])); + } + } } diff --git a/crypto/stark/src/constraint_ir/gpu_interp.rs b/crypto/stark/src/constraint_ir/gpu_interp.rs index de6366d74..56bd2b11d 100644 --- a/crypto/stark/src/constraint_ir/gpu_interp.rs +++ b/crypto/stark/src/constraint_ir/gpu_interp.rs @@ -28,13 +28,13 @@ use math_cuda::lde::{GpuLdeBase, GpuLdeExt3}; use super::device::DeviceProgram; use super::ir::ConstraintProgram; -/// Pack the lowered node list into 2 `u64` per node (`op | a<<32`, `b | dim<<32`), +/// Pack the lowered node list into 2 `u64` per node (`op | a<<32`, `b | res<<32`), /// the encoding the kernel's `load_node` decodes. fn pack_nodes(dev: &DeviceProgram) -> Vec { let mut out = Vec::with_capacity(dev.nodes.len() * 2); for n in &dev.nodes { out.push(n.op as u64 | ((n.a as u64) << 32)); - out.push(n.b as u64 | ((n.dim as u64) << 32)); + out.push(n.b as u64 | ((n.res as u64) << 32)); } out } @@ -246,6 +246,8 @@ where let result = math_cuda::constraint_interp::eval_composition_on_device( &nodes, dev.nodes.len(), + dev.num_base_slots as usize, + dev.num_ext_slots as usize, &dev.base_consts, &ext_consts, &roots, @@ -300,6 +302,8 @@ where let result = math_cuda::constraint_interp::eval_constraints_on_device( &nodes, dev.nodes.len(), + dev.num_base_slots as usize, + dev.num_ext_slots as usize, &dev.base_consts, &ext_consts, &roots, diff --git a/crypto/stark/src/constraints/builder.rs b/crypto/stark/src/constraints/builder.rs index 5395c7228..4554dd4ee 100644 --- a/crypto/stark/src/constraints/builder.rs +++ b/crypto/stark/src/constraints/builder.rs @@ -292,6 +292,7 @@ pub trait ConstraintSet: Send + Sync { /// PAGE, REGISTER, the continuation GLOBAL_MEMORY / global L2G sub-tables). /// The framework still appends the LogUp constraints; this contributes nothing /// before them. +#[derive(Clone, Copy)] pub struct EmptyConstraints; impl ConstraintSet for EmptyConstraints { diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index f962dc272..78afaa2a5 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -14,6 +14,9 @@ use std::sync::atomic::{AtomicU64, Ordering}; use math_cuda::{CudaSlice, CudaStream}; +// External-profiler capture window (nsys -c cudaProfilerApi); re-exported so +// the prover crate can bracket the proving section without a math-cuda dep. + use crypto::fiat_shamir::is_transcript::IsStarkTranscript; use crypto::merkle_tree::merkle::MerkleTree; use crypto::merkle_tree::proof::Proof; @@ -644,6 +647,110 @@ where Some((tree, handle, lde_out)) } +/// Convert a GPU-built full node buffer (`(2*leaves - 1) * 32` bytes, inner +/// nodes first, root at offset 0, leaves at the tail) into a host +/// [`MerkleTree`], the exact layout `from_precomputed_nodes` expects. +fn tree_from_node_bytes(nodes: Vec) -> Option> +where + B: IsMerkleTreeBackend, +{ + debug_assert_eq!(nodes.len() % 32, 0); + let nodes: Vec<[u8; 32]> = nodes + .chunks_exact(32) + .map(|c| { + let mut n = [0u8; 32]; + n.copy_from_slice(c); + n + }) + .collect(); + MerkleTree::::from_precomputed_nodes(nodes) +} + +/// Preprocessed-table variant of [`try_expand_leaf_and_tree_row_major_keep`]: +/// one row-major GPU LDE of ALL columns plus TWO subset Merkle trees — the +/// precomputed columns `[0, split_col)` and the multiplicity columns +/// `[split_col, m)` — matching the CPU `commit_rows_bit_reversed_subset` +/// pair bit for bit. Trees come back as full HOST trees (openings for +/// preprocessed tables walk host trees); the handle keeps the column-major +/// LDE + trace snapshot device-resident for the downstream GPU rounds, with +/// no device tree. +/// +/// `build_precomputed=false` skips the precomputed tree (process-cache hit); +/// the first element is then `None`. +#[allow(clippy::type_complexity)] +pub(crate) fn try_expand_split_trees_row_major_keep( + row_major: &[FieldElement], + n: usize, + m: usize, + blowup_factor: usize, + weights: &[FieldElement], + split_col: usize, + build_precomputed: bool, +) -> Option<( + Option>, + MerkleTree, + math_cuda::lde::GpuLdeBase, + Vec>, +)> +where + F: IsField + 'static, + E: IsField + 'static, + B: IsMerkleTreeBackend, +{ + let lde_size = n.saturating_mul(blowup_factor); + if lde_size < gpu_lde_threshold() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if row_major.len() != n * m || m == 0 || n == 0 { + return None; + } + if split_col == 0 || split_col >= m { + return None; + } + + let raw: &[u64] = unsafe { from_raw_parts(row_major.as_ptr() as *const u64, n * m) }; + let weights_u64 = unsafe { weights_to_u64::(weights) }; + + GPU_LDE_CALLS.fetch_add(m as u64, Ordering::Relaxed); + GPU_LEAF_HASH_CALLS.fetch_add(1 + build_precomputed as u64, Ordering::Relaxed); + GPU_MERKLE_TREE_CALLS.fetch_add(1 + build_precomputed as u64, Ordering::Relaxed); + + let (pre_nodes, mult_nodes, handle, lde_u64) = math_cuda::lde::coset_lde_row_major_split_trees( + raw, + n, + m, + blowup_factor, + &weights_u64, + split_col, + build_precomputed, + ) + .ok()?; + + let pre_tree = match pre_nodes { + Some(nodes) => Some(tree_from_node_bytes::(nodes)?), + None => None, + }; + let mult_tree = tree_from_node_bytes::(mult_nodes)?; + + // Transmute Vec → Vec> (zero-copy, E == GoldilocksField). + let lde_out: Vec> = unsafe { + let mut v = std::mem::ManuallyDrop::new(lde_u64); + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len(), + v.capacity(), + ) + }; + + Some((pre_tree, mult_tree, handle, lde_out)) +} + /// Row-major ext3 GPU path: single H2D → row-major NTT (m*3 base-field cols) → /// row-major Keccak → Merkle → single D2H → transpose to GpuLdeExt3 handle. /// Same optimization as the base-field path: no extract_columns, no CPU transpose. @@ -1287,6 +1394,14 @@ where &weights_u64, retain_host_lde, ) + .inspect_err(|e| { + // This path has no CPU fallback (the host aux trace is empty), so the + // caller hard-aborts; surface the swallowed driver error (e.g. OOM). + eprintln!( + "[gpu] resident aux LDE failed (rows={} cols={} blowup={}): {e:?}", + ra.num_rows, ra.num_aux_cols, blowup_factor + ); + }) .ok()?; let lde_out: Vec> = unsafe { @@ -1852,7 +1967,7 @@ where let mut fri_layer_list: Vec>> = Vec::with_capacity(num_committed); - for _ in 0..num_committed { + for _layer_idx in 0..num_committed { // <<<< Receive challenge zeta_k let zeta: FieldElement = transcript.sample_field_element(); // SAFETY: E == Ext3. @@ -1990,3 +2105,86 @@ where .collect(); Some(decommits) } + +/// GPU↔CPU parity for the preprocessed split-tree commit path. Requires the +/// `cuda` feature and a visible GPU (skipped otherwise via the dispatch gate +/// returning `None` — asserted here, so a silent skip fails the test). +#[cfg(all(test, feature = "cuda"))] +mod split_tree_tests { + use super::*; + use crate::config::BatchedMerkleTreeBackend; + use crate::prover::{IsStarkProver, Prover}; + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Ext3; + + type F = GoldilocksField; + type Fp = FieldElement; + type TestProver = Prover; + + struct SplitMix64(u64); + impl SplitMix64 { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + } + + /// Both subset trees (roots, nodes via openings) must equal the CPU + /// `commit_rows_bit_reversed_subset` built over the same row-major LDE. + /// The LDE itself is parity-pinned by the existing full-row fused tests, + /// so the CPU reference consumes the GPU's returned LDE directly — this + /// isolates the tree layout/hashing under test. + #[test] + fn split_trees_match_cpu_subset_commits() { + // Above the dispatch threshold (2^19 LDE) so the GPU path must engage. + let n: usize = 1 << 18; + let blowup: usize = 2; + let m: usize = 5; + let split: usize = 2; + + let mut rng = SplitMix64(0x5EED_C0DE_5EED_C0DE); + let data: Vec = (0..n * m).map(|_| Fp::from(rng.next_u64())).collect(); + let weights: Vec = (0..n).map(|_| Fp::from(rng.next_u64())).collect(); + + let (pre_tree, mult_tree, handle, lde) = + try_expand_split_trees_row_major_keep::>( + &data, n, m, blowup, &weights, split, true, + ) + .expect("GPU split path must engage above the threshold"); + let pre_tree = pre_tree.expect("precomputed tree was requested"); + + let (cpu_pre, cpu_pre_root) = + TestProver::commit_rows_bit_reversed_subset(&lde, m, 0, split) + .expect("CPU subset commit (precomputed)"); + let (cpu_mult, cpu_mult_root) = + TestProver::commit_rows_bit_reversed_subset(&lde, m, split, m) + .expect("CPU subset commit (multiplicities)"); + + assert_eq!(pre_tree.root, cpu_pre_root, "precomputed root"); + assert_eq!(mult_tree.root, cpu_mult_root, "multiplicity root"); + + // Openings must be byte-identical at scattered positions (pins the + // full node buffers, not just the roots). + let num_leaves = n * blowup / 2; + for pos in [0usize, 1, 511, 12_345, num_leaves - 1] { + assert_eq!( + pre_tree.get_proof_by_pos(pos).unwrap().merkle_path, + cpu_pre.get_proof_by_pos(pos).unwrap().merkle_path, + "precomputed path at {pos}" + ); + assert_eq!( + mult_tree.get_proof_by_pos(pos).unwrap().merkle_path, + cpu_mult.get_proof_by_pos(pos).unwrap().merkle_path, + "multiplicity path at {pos}" + ); + } + + // The handle must carry the column-major LDE for downstream rounds: + // spot-check a few cells against the row-major host LDE. + assert_eq!(handle.m, m); + assert_eq!(handle.lde_size, n * blowup); + assert!(handle.tree.is_none(), "no device tree on the split path"); + } +} diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 96bf6ffae..06a5d2c70 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -43,9 +43,36 @@ pub struct SpanGuard { order: u32, start: Instant, start_ns: u128, + /// This span gates an nsys capture range (LAMBDA_VM_NSYS_CAPTURE_SPAN). + #[cfg(feature = "nvtx")] + capture: bool, +} + +/// Label of the span that brackets an `nsys --capture-range=cudaProfilerApi` +/// session, from `LAMBDA_VM_NSYS_CAPTURE_SPAN` (e.g. `rounds_2to4`, or +/// `proving` to capture one epoch of a continuations run). None = never. +#[cfg(feature = "nvtx")] +fn capture_span_label() -> Option<&'static str> { + static LABEL: OnceLock> = OnceLock::new(); + LABEL + .get_or_init(|| std::env::var("LAMBDA_VM_NSYS_CAPTURE_SPAN").ok()) + .as_deref() +} + +/// NVTX-only range with a runtime-formatted name (e.g. `epoch[i=3]`), for +/// callers that need per-instance identity on Nsight timelines. Instruments +/// spans require `'static` labels, so repeated phases (continuation epochs) +/// are told apart by instance order in the JSON timeline and by one of these +/// ranges in nsys. The closure only runs when a profiler-visible NVTX +/// library is loaded. +#[cfg(feature = "nvtx")] +pub fn nvtx_range_fmt String>(label: F) -> math_cuda::nvtx::Range { + math_cuda::nvtx::Range::fmt(label) } /// Open a wall-clock span; records elapsed time when the guard drops. +/// Under the `nvtx` feature the span is mirrored as an NVTX range so Nsight +/// timelines carry the same phase names as the instruments tree. pub fn span(label: &'static str) -> SpanGuard { let depth = SPAN_DEPTH.with(|d| { let v = d.get(); @@ -57,17 +84,35 @@ pub fn span(label: &'static str) -> SpanGuard { .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_nanos(); + #[cfg(feature = "nvtx")] + let capture = { + math_cuda::nvtx::range_push(label); + let capture = capture_span_label() == Some(label); + if capture { + math_cuda::nvtx::profiler_start(); + } + capture + }; SpanGuard { label, depth, order, start: Instant::now(), start_ns, + #[cfg(feature = "nvtx")] + capture, } } impl Drop for SpanGuard { fn drop(&mut self) { + #[cfg(feature = "nvtx")] + { + if self.capture { + math_cuda::nvtx::profiler_stop(); + } + math_cuda::nvtx::range_pop(); + } let wall = self.start.elapsed(); SPAN_DEPTH.with(|d| d.set(d.get().saturating_sub(1))); if let Ok(mut t) = TIMELINE.lock() { diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index 8a89ea727..c39d231d1 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -848,6 +848,39 @@ pub struct AirWithBuses< max_bus_elements: usize, } +/// Cloning an `AirWithBuses` copies its derived artifacts — the MetaBuilder-run +/// constraint metadata, the LogUp layout, and (if already forced) the captured +/// constraint IR inside the `OnceLock` — so a pre-built, pre-captured prototype +/// clones into per-shard/per-epoch instances without re-running the constraint +/// bodies. `B`/`PI` ride in `PhantomData` and need no bounds. +impl< + F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, + E: IsField + Send + Sync, + B: BoundaryConstraintBuilder, + PI, + CS: ConstraintSet + Clone, +> Clone for AirWithBuses +{ + fn clone(&self) -> Self { + Self { + context: self.context.clone(), + step_size: self.step_size, + trace_layout: self.trace_layout, + constraint_set: self.constraint_set.clone(), + logup: self.logup.clone(), + meta: self.meta.clone(), + num_base: self.num_base, + constraint_program: self.constraint_program.clone(), + auxiliary_trace_build_data: self.auxiliary_trace_build_data.clone(), + boundary_constraint_builder: PhantomData, + preprocessed_commitment: self.preprocessed_commitment, + num_precomputed_cols: self.num_precomputed_cols, + name: self.name.clone(), + max_bus_elements: self.max_bus_elements, + } + } +} + impl< F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync + 'static, E: IsField + Send + Sync + 'static, @@ -1312,6 +1345,7 @@ where /// Struct representing how each lookup air should build its auxiliary trace /// Contains a list of all lookup interactions +#[derive(Clone)] pub struct AuxiliaryTraceBuildData { pub interactions: Vec, } diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 8c44c42a9..b557be73b 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1,5 +1,5 @@ use std::marker::PhantomData; -use std::sync::{Arc, OnceLock}; +use std::sync::{Arc, Mutex, OnceLock}; #[cfg(feature = "instruments")] use std::time::{Duration, Instant}; @@ -139,18 +139,20 @@ where } } - /// Build a `TableCommit` for a preprocessed table. + /// Build a `TableCommit` for a preprocessed table. The precomputed tree + /// arrives as an `Arc` because it may be shared from the process-wide + /// cache (see [`precomputed_tree_cache_get`]). fn preprocessed( tree: BatchedMerkleTree, root: Commitment, - precomputed_tree: BatchedMerkleTree, + precomputed_tree: Arc>, precomputed_root: Commitment, num_precomputed_cols: usize, ) -> Self { Self { tree: Arc::new(tree), root, - precomputed_tree: Some(Arc::new(precomputed_tree)), + precomputed_tree: Some(precomputed_tree), precomputed_root: Some(precomputed_root), num_precomputed_cols, } @@ -172,6 +174,47 @@ where } } +/// Process-wide cache of precomputed-column Merkle trees, keyed by their +/// commitment root. The root fully determines the tree (column content, +/// domain, blowup and leaf layout all feed the hash), so a hit needs no +/// re-verification: the lookup key IS the root a rebuild would be checked +/// against. This is what makes continuation epochs stop re-committing the +/// same DECODE/BITWISE/range tables once per epoch — those trees are +/// execution-independent; only the multiplicity columns change per run. +/// Type-erased so one static serves every field instantiation. +fn precomputed_tree_cache() +-> &'static Mutex>> { + static CACHE: OnceLock< + Mutex>>, + > = OnceLock::new(); + CACHE.get_or_init(|| Mutex::new(std::collections::HashMap::new())) +} + +fn precomputed_tree_cache_get( + root: &Commitment, +) -> Option>> +where + FieldElement: AsBytes, +{ + let cache = precomputed_tree_cache().lock().unwrap(); + cache + .get(root) + .cloned() + .and_then(|any| any.downcast::>().ok()) +} + +fn precomputed_tree_cache_put( + root: Commitment, + tree: Arc>, +) where + FieldElement: AsBytes, +{ + precomputed_tree_cache() + .lock() + .unwrap() + .insert(root, tree as Arc); +} + /// A container for the results of the first round of the STARK Prove protocol. pub(crate) struct Round1 where @@ -893,6 +936,86 @@ pub trait IsStarkProver< } } + // Fused GPU split path for preprocessed tables (cuda only): one + // row-major LDE of ALL columns plus two subset Merkle trees + // (precomputed / multiplicity) built on device — leaves and levels are + // bit-identical to `commit_rows_bit_reversed_subset`, and the trees + // come back as full host trees so the preprocessed opening path and + // the process-wide precomputed-tree cache work unchanged. The handle + // keeps the LDE device-resident for the downstream GPU rounds. + #[cfg(feature = "cuda")] + if let Some((expected_precomputed_root, num_precomputed)) = precomputed { + let (trace_slice, num_cols) = trace.main_data_row_major(); + let n = if num_cols > 0 { + trace_slice.len() / num_cols + } else { + 0 + }; + #[cfg(feature = "disk-spill")] + let cache_ok = storage_mode != StorageMode::Disk; + #[cfg(not(feature = "disk-spill"))] + let cache_ok = true; + let cached_pre = cache_ok + .then(|| precomputed_tree_cache_get::(&expected_precomputed_root)) + .flatten(); + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); + if let Some((pre_tree, mult_tree, handle, main_data)) = + crate::gpu_lde::try_expand_split_trees_row_major_keep::< + Field, + Field, + BatchedMerkleTreeBackend, + >( + trace_slice, + n, + num_cols, + domain.blowup_factor, + &twiddles.coset_weights, + num_precomputed, + cached_pre.is_none(), + ) + { + #[cfg(feature = "instruments")] + crate::instruments::accum_r1_main(t_sub.elapsed(), std::time::Duration::ZERO); + let precomputed_tree = match cached_pre { + // Cache key == the root a rebuild would be verified + // against, so a hit needs no re-check. + Some(tree) => tree, + None => { + #[allow(unused_mut)] + let mut tree = pre_tree.expect("precomputed tree requested on cache miss"); + if tree.root != expected_precomputed_root { + return Err(ProvingError::PrecomputedCommitmentMismatch); + } + #[cfg(feature = "disk-spill")] + Self::spill_tree(&mut tree, storage_mode, "precomputed Merkle tree")?; + let tree = Arc::new(tree); + if cache_ok { + precomputed_tree_cache_put::( + expected_precomputed_root, + Arc::clone(&tree), + ); + } + tree + } + }; + #[allow(unused_mut)] + let mut mult_tree = mult_tree; + #[cfg(feature = "disk-spill")] + Self::spill_tree(&mut mult_tree, storage_mode, "mult Merkle tree")?; + let mult_root = mult_tree.root; + let commit = TableCommit::preprocessed( + mult_tree, + mult_root, + precomputed_tree, + expected_precomputed_root, + num_precomputed, + ); + return Ok((commit, (main_data, num_cols), Some(handle))); + } + // GPU split path declined (size threshold / tower) → CPU path below. + } + // CPU path: the trace `Table` is already row-major, so copy it directly // (one memcpy — no transpose) and expand in place with the cache-blocked // batched two-half FFT. Row-major end-to-end: no LDE-size transpose, @@ -936,15 +1059,47 @@ pub trait IsStarkProver< TableCommit::plain(tree, root) } Some((expected_precomputed_root, num_precomputed)) => { - #[allow(unused_mut)] - let (mut precomputed_tree, precomputed_root) = - Self::commit_rows_bit_reversed_subset( - &main_data, - total_cols, - 0, - num_precomputed, - ) - .ok_or(ProvingError::EmptyCommitment)?; + // Only the multiplicity columns depend on the execution; the + // precomputed-columns tree is a pure function of (content, + // domain) already pinned by `expected_precomputed_root`, so it + // is reused from the process cache when this exact commitment + // was built before — across epochs and across proves. Bypassed + // in disk-spill Disk mode, where trees are spilled (mutated). + #[cfg(feature = "disk-spill")] + let cache_ok = storage_mode != StorageMode::Disk; + #[cfg(not(feature = "disk-spill"))] + let cache_ok = true; + let precomputed_tree = match cache_ok + .then(|| precomputed_tree_cache_get::(&expected_precomputed_root)) + .flatten() + { + // Cache key == the root a rebuild would be verified + // against, so a hit needs no re-check. + Some(tree) => tree, + None => { + #[allow(unused_mut)] + let (mut tree, root) = Self::commit_rows_bit_reversed_subset( + &main_data, + total_cols, + 0, + num_precomputed, + ) + .ok_or(ProvingError::EmptyCommitment)?; + if root != expected_precomputed_root { + return Err(ProvingError::PrecomputedCommitmentMismatch); + } + #[cfg(feature = "disk-spill")] + Self::spill_tree(&mut tree, storage_mode, "precomputed Merkle tree")?; + let tree = Arc::new(tree); + if cache_ok { + precomputed_tree_cache_put::( + expected_precomputed_root, + Arc::clone(&tree), + ); + } + tree + } + }; #[allow(unused_mut)] let (mut mult_tree, mult_root) = Self::commit_rows_bit_reversed_subset( &main_data, @@ -953,23 +1108,13 @@ pub trait IsStarkProver< total_cols, ) .ok_or(ProvingError::EmptyCommitment)?; - if precomputed_root != expected_precomputed_root { - return Err(ProvingError::PrecomputedCommitmentMismatch); - } #[cfg(feature = "disk-spill")] - { - Self::spill_tree( - &mut precomputed_tree, - storage_mode, - "precomputed Merkle tree", - )?; - Self::spill_tree(&mut mult_tree, storage_mode, "mult Merkle tree")?; - } + Self::spill_tree(&mut mult_tree, storage_mode, "mult Merkle tree")?; TableCommit::preprocessed( mult_tree, mult_root, precomputed_tree, - precomputed_root, + expected_precomputed_root, num_precomputed, ) } @@ -2406,6 +2551,13 @@ pub trait IsStarkProver< #[cfg(not(feature = "cuda"))] let vram_budget = u64::MAX; + // NOTE: an earlier revision published prove-wide pinned-staging size + // hints here (`set_staging_size_hints`) so slabs allocated once at + // final size. Measured on a 5090 it BACKFIRED: every worker slot then + // pays a max-size cuMemHostAlloc (~160ms avg, 7.5s total vs 4.2s of + // ladder churn), and those allocations convoy the driver lock. Left + // out until a shared-slab design bounds the number of allocations. + // R1 main commit: only the main LDE and its Merkle scratch are resident, // so the aux columns add nothing to this phase's working set. let main_chunks = { diff --git a/crypto/stark/tests/gpu_constraint_interp.rs b/crypto/stark/tests/gpu_constraint_interp.rs index de7211dcd..0304b7e4c 100644 --- a/crypto/stark/tests/gpu_constraint_interp.rs +++ b/crypto/stark/tests/gpu_constraint_interp.rs @@ -34,7 +34,8 @@ use math_cuda::device::backend; use math_cuda::lde::{GpuLdeBase, GpuLdeExt3}; use stark::constraint_ir::device::{ - DeviceProgram, OP_ALPHA_POW, OP_RAP_CHALLENGE, OP_VAR, eval_device_program, unpack_var, + DeviceProgram, OP_ADD, OP_ALPHA_POW, OP_EMBED, OP_MUL, OP_NEG, OP_RAP_CHALLENGE, OP_SUB, + OP_VAR, OPK_ALPHA, OPK_PAYLOAD_MASK, OPK_RAP, OPK_SHIFT, eval_device_program, unpack_var, }; use stark::constraint_ir::{ConstraintProgram, IrBuilder}; @@ -119,6 +120,17 @@ fn all_ops_program() -> ConstraintProgram { /// #rap challenges, #alpha powers, and the max frame offset. fn program_footprint(dev: &DeviceProgram) -> (usize, usize, usize, usize, usize) { let (mut main_cols, mut aux_cols, mut rap_len, mut alpha_len, mut max_off) = (0, 0, 0, 0, 0); + // Uniform leaves are propagated into operand encodings, so the RAP/alpha + // footprint must be read from the operands of arithmetic nodes (the + // root-pinned leaf-node forms are kept for completeness). + let scan_operand = |enc: u32, rap_len: &mut usize, alpha_len: &mut usize| { + let payload = (enc & OPK_PAYLOAD_MASK) as usize; + match enc >> OPK_SHIFT { + OPK_RAP => *rap_len = (*rap_len).max(payload + 1), + OPK_ALPHA => *alpha_len = (*alpha_len).max(payload + 1), + _ => {} + } + }; for n in &dev.nodes { match n.op { OP_VAR => { @@ -133,6 +145,11 @@ fn program_footprint(dev: &DeviceProgram) -> (usize, usize, usize, usize, usize) } OP_RAP_CHALLENGE => rap_len = rap_len.max(n.a as usize + 1), OP_ALPHA_POW => alpha_len = alpha_len.max(n.a as usize + 1), + OP_ADD | OP_SUB | OP_MUL => { + scan_operand(n.a, &mut rap_len, &mut alpha_len); + scan_operand(n.b, &mut rap_len, &mut alpha_len); + } + OP_NEG | OP_EMBED => scan_operand(n.a, &mut rap_len, &mut alpha_len), _ => {} } } @@ -194,6 +211,7 @@ fn check_program(prog: &ConstraintProgram, label: &str, seed: u64) { stream.synchronize().expect("sync uploads"); let main = GpuLdeBase { + ready: None, buf: Arc::new(base_dev), m: main_cols, lde_size, @@ -202,6 +220,7 @@ fn check_program(prog: &ConstraintProgram, label: &str, seed: u64) { trace_rows: 0, }; let aux = GpuLdeExt3 { + ready: None, buf: Arc::new(aux_dev), m: aux_cols, lde_size, @@ -409,6 +428,7 @@ fn check_composition(prog: &ConstraintProgram, label: &str, seed: u64) let aux_dev = stream.clone_htod(&aux_flat).expect("upload aux"); stream.synchronize().expect("sync"); let main = GpuLdeBase { + ready: None, buf: Arc::new(base_dev), m: main_cols, lde_size, @@ -417,6 +437,7 @@ fn check_composition(prog: &ConstraintProgram, label: &str, seed: u64) trace_rows: 0, }; let aux = GpuLdeExt3 { + ready: None, buf: Arc::new(aux_dev), m: aux_cols, lde_size, diff --git a/prover/Cargo.toml b/prover/Cargo.toml index 15344d138..b5603e7c7 100644 --- a/prover/Cargo.toml +++ b/prover/Cargo.toml @@ -11,6 +11,8 @@ cuda = ["stark/cuda"] test-cuda-faults = ["cuda", "stark/test-cuda-faults"] debug-checks = ["stark/debug-checks"] instruments = ["stark/instruments"] +# GPU profiling build (Nsight): implies cuda — the toolkit always traces the GPU prover. +nvtx = ["cuda", "instruments", "stark/nvtx"] profile-markers = ["stark/profile-markers"] disk-spill = ["stark/disk-spill"] diff --git a/prover/src/constraints/cpu.rs b/prover/src/constraints/cpu.rs index 917445b7e..afb7fb07b 100644 --- a/prover/src/constraints/cpu.rs +++ b/prover/src/constraints/cpu.rs @@ -273,6 +273,7 @@ pub fn emit_next_pc_add_pair for CpuConstraints { diff --git a/prover/src/continuation.rs b/prover/src/continuation.rs index 169cd7278..65c0dfd50 100644 --- a/prover/src/continuation.rs +++ b/prover/src/continuation.rs @@ -48,6 +48,7 @@ //! and ELF alone (`prove_and_verify_continuation` is a thin wrapper over both). use std::collections::HashMap; +use std::sync::Arc; use crypto::fiat_shamir::default_transcript::DefaultTranscript; use executor::elf::Elf; @@ -68,7 +69,9 @@ use crate::statement::{StatementKind, absorb_continuation_global_statement, abso use crate::tables::local_to_global::{self, CellBoundary}; use crate::tables::page::{self, PageConfig}; use crate::tables::register; -use crate::tables::trace_builder::{Traces, build_init_page_data, build_initial_image_paged}; +use crate::tables::trace_builder::{ + DecodeArtifacts, Traces, build_init_page_data, build_initial_image_paged, +}; use crate::tables::types::{GoldilocksExtension, GoldilocksField}; use crate::tables::{MaxRowsConfig, global_memory}; use crate::{ @@ -139,6 +142,7 @@ fn global_transcript( /// identical trace (root-bound), so it inherits it. /// The L2G epoch-local table's single transition constraint: `MU ∈ {0,1}` /// (`MU·(1−MU) = 0`) at constraint index 0. +#[derive(Clone, Copy)] struct L2gMemoryConstraints; impl ConstraintSet for L2gMemoryConstraints { @@ -251,10 +255,10 @@ fn global_memory_air( /// and verifier iterate the identical sequence — `multi_verify` matches AIRs to sub-proofs /// positionally. Carries page bases ONLY: no cell values, so private-input bytes never /// enter the bundle (unlike the full `CellBoundary`, whose `init.value` is a private byte). -fn touched_page_bases(boundaries: &[Vec]) -> Vec { +fn touched_page_bases(boundaries: &[Arc>]) -> Vec { boundaries .iter() - .flatten() + .flat_map(|epoch| epoch.iter()) .map(|b| page::page_base_for_address(b.address)) .collect::>() .into_iter() @@ -381,6 +385,37 @@ struct EpochStart<'a> { label: u64, } +/// One epoch's proving inputs, fully derived from execution (register init, +/// traces, boundary — no dependency on any previous epoch's *proof*), so the +/// preparation of epoch i+1 can run on a producer thread while epoch i proves. +/// +/// `boundary` is shared (`Arc`): the same per-epoch boundary feeds both this +/// epoch's prove and the cross-epoch global prove, which starts as soon as the +/// producer has prepared the last epoch (see `prove_continuation`). +struct PreparedEpoch { + index: u64, + register_init: Vec, + label: u64, + traces: Traces, + boundary: Arc>, + is_final: bool, +} + +/// A collected-but-not-yet-built epoch, handed from the producer to the trace +/// builder pool. Everything sequential (execution, op collection over the +/// advancing memory image, boundary + register-fini derivation) already +/// happened on the producer; a builder turns `collected` into full trace +/// tables ([`Traces::build_from_collected`]) — pure epoch-local work — and +/// forwards the resulting [`PreparedEpoch`] to the epoch provers. +struct BuildJob { + index: u64, + register_init: Vec, + label: u64, + collected: crate::tables::trace_builder::CollectedEpoch, + boundary: Arc>, + is_final: bool, +} + /// One epoch's proof plus everything a standalone verifier needs to re-check it /// using ONLY the bundle (never the prover's in-memory traces). Each field is a /// public value the verifier re-binds: a wrong value either makes the proof's @@ -635,6 +670,7 @@ fn prove_epoch( is_final: bool, boundary: &[CellBoundary], opts: &ProofOptions, + decode_commitment: Commitment, ) -> Result { // Count this L2G table's range-check lookups into the BITWISE table so its // AreBytes/IsHalfword multiplicities balance the range-check senders. @@ -667,7 +703,10 @@ fn prove_epoch( start.register_init, ®_fini, is_final, - None, + // Computed once per prove_continuation — the DECODE commitment is a + // function of (ELF, opts) only, identical for every epoch; passing + // None here would rebuild the whole DECODE trace+LDE+tree per epoch. + Some(decode_commitment), ); let label = start.label; @@ -823,7 +862,7 @@ fn verify_epoch( /// The bus balances iff every `fini` matches the next epoch's `init` and every genesis /// matches its source (the ELF for ELF/runtime pages). fn prove_global( - boundaries: &[Vec], + boundaries: &[Arc>], elf_bytes: &[u8], init_page_data: &HashMap>, page_bases: &[u64], @@ -833,7 +872,7 @@ fn prove_global( // Each cell's final state (boundaries are in epoch order, so the last fini wins). let mut final_state: global_memory::FiniStateMap = HashMap::new(); for epoch in boundaries { - for b in epoch { + for b in epoch.iter() { final_state.insert( b.address, global_memory::FiniState { @@ -853,7 +892,7 @@ fn prove_global( let mut l2g_traces: Vec> = boundaries .iter() - .map(|epoch| local_to_global::generate_local_to_global_trace(epoch)) + .map(|epoch| local_to_global::generate_local_to_global_trace(epoch.as_slice())) .collect(); let mut gm_traces: Vec> = gm_configs .iter() @@ -1009,9 +1048,25 @@ pub fn prove_continuation( )) })?; + // Root span for the profiling toolkit (scripts/profiling): the whole + // continuation prove is one tree; per-stage spans below are recorded from + // their worker threads and told apart by label + instance order. + #[cfg(feature = "instruments")] + stark::instruments::reset_timeline(); + #[cfg(feature = "instruments")] + let __root = stark::instruments::span("prove_continuation_total"); + let elf = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; let mut executor = Executor::new(&elf, private_inputs.to_vec()) .map_err(|e| Error::Execution(format!("{e}")))?; + // The DECODE precomputed commitment depends only on (ELF, opts): compute + // it once here instead of once per epoch inside `build_epoch_airs`. + let decode_commitment = crate::tables::decode::commitment_from_elf(&elf, opts) + .map_err(|e| Error::Recursion(format!("DECODE commitment from ELF: {e}")))?; + // Same for the DECODE trace artifacts (instruction map + pristine trace): + // a pure function of the ELF, built once and shared by every epoch's trace + // build instead of re-parsed/regenerated inside the serial producer chain. + let decode_artifacts = DecodeArtifacts::from_elf(&elf)?; // The cross-epoch memory image, carried forward: epoch i+1's init is epoch i's // fini, updated in place with each epoch's touched-cell final values. @@ -1025,110 +1080,402 @@ pub fn prove_continuation( // final-state). Deliberately NOT stored in `EpochProof`/the bundle — `CellBoundary` // holds cell values (private-input bytes for private reads); only the value-free // page-base set is shipped (see `touched_page_bases`). - let mut all_boundaries: Vec> = Vec::new(); - // The previous epoch's bound final register file R_{i+1}; epoch i+1's init is - // derived from it (the cross-epoch register binding). - let mut prev_fini: Option> = None; - - let mut index: u64 = 0; - loop { - if executor.pc() == 0 { - break; - } - // The cross-epoch ordering check (IsB20 on `fini_epoch - 1 - init_epoch`) - // only spans `local_to_global::MAX_EPOCHS` epochs. Beyond that the IsB20 bus - // cannot balance, so an honest proof is impossible — fail fast with a clear - // error instead of building an unprovable trace. The verifier already - // rejects any such proof; this is a prover-side guard for a clean message. - if index >= local_to_global::MAX_EPOCHS { - return Err(Error::InvalidContinuationEpochSize(format!( - "execution needs more than {} continuation epochs (the IsB20 cross-epoch \ - ordering range); use a larger epoch size", - local_to_global::MAX_EPOCHS - ))); + // + // The producer publishes each epoch's boundary (an `Arc` share of the one it + // sends to the epoch provers) on this dedicated channel, in epoch order. The + // global-prove thread drains it until the producer hangs up (last epoch + // prepared) — the global proof depends only on these execution artifacts, + // never on an epoch *proof*, so it overlaps the epoch proves' tail instead + // of serializing after them. Proof bytes are unchanged — only the schedule. + let (boundary_tx, boundary_rx) = std::sync::mpsc::channel::>>(); + + // Three-stage epoch pipeline: a producer thread runs the + // sequential-critical work (execute + op collection over the advancing + // memory image + boundary/fini derivation), a small pool of trace builders + // turns collected epochs into trace tables, and `workers` provers prove + // them. Everything the next epoch's preparation needs is derived from + // execution, not from proofs or traces: `register_init` comes from the + // collected register end state (`register_fini`, the same value the + // generated REGISTER trace binds) and the memory image update comes from + // the boundary — so the producer chains epochs without waiting for any + // table to be built. Proof bytes are unchanged — only the schedule is. + // + // The bounded channels cap peak memory: at most one collected epoch + // queued, `builders` building, one built epoch queued, `workers` proving. + let (tx, rx) = std::sync::mpsc::sync_channel::>(1); + let (build_tx, build_rx) = std::sync::mpsc::sync_channel::>(1); + // Shared prover-pool state; declared outside the scope so scoped threads + // can borrow it. + // + // Default 3: measured on ethrex-10tx / RTX 5090 32 GB (8×8 interleaved + // runs), 3 concurrent epoch proves beat 2 by ~4% (14.9s vs 15.5s) at + // ~15 GB peak VRAM. On smaller cards (< 32 GB) set + // LAMBDA_VM_EPOCH_CONCURRENCY=2 — each concurrent 2^20-cycle epoch prove + // peaks at ~9 GB and the device-only GPU paths abort hard on VRAM + // exhaustion rather than degrade. + let workers = std::env::var("LAMBDA_VM_EPOCH_CONCURRENCY") + .ok() + .and_then(|s| s.parse::().ok()) + .filter(|&k| k >= 1) + .unwrap_or(3); + // Trace builders: each turns one collected epoch into full trace tables + // (the bulk of the old per-epoch producer latency). 2 is enough to keep + // the prove pipeline fed on the measured workloads; builds compete with + // proves for CPU, so more builders mostly reshuffle the same cores. + let builders = std::env::var("LAMBDA_VM_TRACE_BUILDERS") + .ok() + .and_then(|s| s.parse::().ok()) + .filter(|&b| b >= 1) + .unwrap_or(2); + let rx = std::sync::Mutex::new(rx); + let build_rx = std::sync::Mutex::new(build_rx); + type EpochResult = (u64, EpochProof); + let results: std::sync::Mutex> = std::sync::Mutex::new(Vec::new()); + let first_err: std::sync::Mutex> = std::sync::Mutex::new(None); + let decode_artifacts_ref = &decode_artifacts; + let prove_worker = |_worker: usize| { + loop { + // Take the next prepared epoch (lock released before proving). + let msg = { rx.lock().unwrap().recv() }; + let prepared = match msg { + Ok(Ok(p)) => p, + Ok(Err(e)) => { + first_err.lock().unwrap().get_or_insert(e); + return; + } + Err(_) => return, // channel closed: no more epochs + }; + if first_err.lock().unwrap().is_some() { + return; // another worker failed; stop consuming + } + // Per-epoch identity on Nsight timelines (dynamic NVTX name); the + // instruments span carries a static label and instances are told + // apart by order (phase_table.py reports them per instance). + #[cfg(feature = "nvtx")] + let __nvtx = + stark::instruments::nvtx_range_fmt(|| format!("epoch_prove[i={}]", prepared.index)); + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("epoch_prove"); + let start = EpochStart { + register_init: &prepared.register_init, + label: prepared.label, + }; + match prove_epoch( + &elf, + elf_bytes, + &start, + prepared.traces, + prepared.is_final, + &prepared.boundary, + opts, + decode_commitment, + ) { + Ok(epoch) => results.lock().unwrap().push((prepared.index, epoch)), + Err(e) => { + first_err.lock().unwrap().get_or_insert(e); + return; + } + } } - let register_init: Vec = if index == 0 { - register::register_init_from_entry_point(elf.entry_point) - } else { - // Epoch i+1's init is epoch i's bound fini, reused directly (same - // `register_word_address_list` order) — the cross-epoch register binding. - prev_fini.clone().ok_or_else(|| { - Error::ContinuationInvariant( - "previous epoch final registers are missing after the first epoch".to_string(), - ) - })? + }; + // Trace-builder worker: drain collected epochs, build their trace tables + // (pure epoch-local work) and forward the prepared epoch to the provers. + // Errors propagate through the prove channel, exactly like producer errors. + let build_worker = + |_lane: usize, tx: std::sync::mpsc::SyncSender>| { + loop { + let msg = { build_rx.lock().unwrap().recv() }; + let job = match msg { + Ok(Ok(j)) => j, + Ok(Err(e)) => { + let _ = tx.send(Err(e)); + return; + } + Err(_) => return, // channel closed: no more epochs + }; + if first_err.lock().unwrap().is_some() { + return; // a prover failed; stop building + } + #[cfg(feature = "nvtx")] + let __nvtx = stark::instruments::nvtx_range_fmt(|| { + format!("epoch_trace_build[i={}]", job.index) + }); + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("epoch_trace_build"); + let traces = Traces::build_from_collected( + decode_artifacts_ref, + job.collected, + // Continuation epochs use the L2G bookend: PAGE tables (the + // only image consumers in the build) are skipped. + None::<&std::collections::HashMap>, + &job.register_init, + &MaxRowsConfig::default(), + private_inputs, + job.is_final, + true, + #[cfg(feature = "disk-spill")] + stark::storage_mode::StorageMode::Ram, + ); + // Close the build span BEFORE forwarding: the send below blocks + // on prove-channel backpressure, which is waiting, not building. + #[cfg(feature = "instruments")] + drop(__sp); + #[cfg(feature = "nvtx")] + drop(__nvtx); + match traces { + Ok(traces) => { + let prepared = PreparedEpoch { + index: job.index, + register_init: job.register_init, + label: job.label, + traces, + boundary: job.boundary, + is_final: job.is_final, + }; + // A send error means the prover side hung up (its error is + // already propagating) — stop quietly. + if tx.send(Ok(prepared)).is_err() { + return; + } + } + Err(e) => { + let _ = tx.send(Err(e)); + return; + } + } + } }; - // Run one epoch; `logs` is this epoch's chunk only (the executor clears it). - let logs = match executor - .resume_with_limit(epoch_size) - .map_err(|e| Error::Execution(format!("{e}")))? - { - Some(logs) => logs.to_vec(), - None => break, - }; - let is_final = executor.pc() == 0; - - // Invariant: a non-final epoch ran the full `epoch_size` (a power of two), - // so its CPU table has no padding rows. - if !is_final && logs.len() != epoch_size { - return Err(Error::ContinuationInvariant(format!( - "intermediate epoch ran {} cycles, expected {epoch_size}", - logs.len() - ))); + // The global prove's result, produced by its own scoped thread. `None` only + // if that thread never ran to completion (a panic — surfaced by the scope). + type GlobalResult = (MultiProof, Vec, usize); + let global_result: std::sync::Mutex>> = + std::sync::Mutex::new(None); + std::thread::scope(|scope| -> Result<(), Error> { + let elf_ref = &elf; + let producer = scope.spawn(move || { + let mut prepare_all = || -> Result<(), Error> { + let mut prev_fini: Option> = None; + let mut index: u64 = 0; + loop { + if executor.pc() == 0 { + return Ok(()); + } + // The cross-epoch ordering check (IsB20 on `fini_epoch - 1 - + // init_epoch`) only spans `local_to_global::MAX_EPOCHS` epochs. + // Beyond that the IsB20 bus cannot balance, so an honest proof + // is impossible — fail fast with a clear error instead of + // building an unprovable trace. + if index >= local_to_global::MAX_EPOCHS { + return Err(Error::InvalidContinuationEpochSize(format!( + "execution needs more than {} continuation epochs (the IsB20 \ + cross-epoch ordering range); use a larger epoch size", + local_to_global::MAX_EPOCHS + ))); + } + let register_init: Vec = match (index, prev_fini.take()) { + (0, _) => register::register_init_from_entry_point(elf_ref.entry_point), + // Epoch i+1's init is epoch i's bound fini, reused directly + // (same `register_word_address_list` order) — the cross-epoch + // register binding. + (_, Some(fini)) => fini, + (_, None) => { + return Err(Error::ContinuationInvariant( + "previous epoch final registers are missing after the first epoch" + .to_string(), + )); + } + }; + + // Run one epoch; `logs` is this epoch's chunk only (the executor + // clears it). + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("epoch_execute"); + let logs = match executor + .resume_with_limit(epoch_size) + .map_err(|e| Error::Execution(format!("{e}")))? + { + Some(logs) => logs.to_vec(), + None => return Ok(()), + }; + #[cfg(feature = "instruments")] + drop(__sp); + let is_final = executor.pc() == 0; + + // Invariant: a non-final epoch ran the full `epoch_size` (a power + // of two), so its CPU table has no padding rows. + if !is_final && logs.len() != epoch_size { + return Err(Error::ContinuationInvariant(format!( + "intermediate epoch ran {} cycles, expected {epoch_size}", + logs.len() + ))); + } + + let label = local_to_global::epoch_label(index); + // Sequential-critical half only (Phases 1-2): op collection + // over the pre-epoch image. The table build (Phases 3-5) + // happens on the builder pool — nothing below needs it. + #[cfg(feature = "nvtx")] + let __nvtx = + stark::instruments::nvtx_range_fmt(|| format!("epoch_collect[i={index}]")); + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("epoch_collect"); + let collected = Traces::collect_epoch( + decode_artifacts_ref, + &image, + ®ister_init, + &logs, + is_final, + )?; + let boundary = Arc::new(local_to_global::epoch_boundary( + &mut provenance, + label, + &collected.touched_memory_cells(), + )); + // Publish this epoch's boundary for the global prove (in + // epoch order; the channel closes when the producer ends). + let _ = boundary_tx.send(Arc::clone(&boundary)); + + // R_{i+1} from the collected register end state — the exact + // value the generated REGISTER trace binds (`fini_from_trace` + // equivalence pinned by `fini_from_final_state_matches_trace`). + prev_fini = Some(collected.register_fini(®ister_init)); + + // Carry the image forward: this epoch's fini is the next + // epoch's init. + for cell in boundary.iter() { + image.set(cell.address, (cell.fini.value & 0xFF) as u8); + } + + // Close the collect span BEFORE handing off: the send below + // blocks on builder backpressure, which is waiting, not work. + #[cfg(feature = "instruments")] + drop(__sp); + #[cfg(feature = "nvtx")] + drop(__nvtx); + let job = BuildJob { + index, + register_init, + label, + collected, + boundary, + is_final, + }; + // A send error means the builder side hung up (its error is + // already propagating) — stop preparing quietly. + if build_tx.send(Ok(job)).is_err() || is_final { + return Ok(()); + } + index += 1; + } + }; + if let Err(e) = prepare_all() { + // Surface preparation errors through the builder channel (a + // builder forwards them to the provers); if the downstream side + // is already gone the error there wins. + let _ = build_tx.send(Err(e)); + } + }); + + // Trace-builder pool: collected epochs → trace tables → prove channel. + // Each builder owns a clone of the prove sender; the original is + // dropped below so the provers' channel closes once the producer and + // every builder are done. + for lane in 0..builders { + let tx = tx.clone(); + scope.spawn(move || build_worker(lane, tx)); } - - let label = local_to_global::epoch_label(index); - let traces = Traces::from_image_and_logs( - &elf, - &image, - ®ister_init, - &logs, - &MaxRowsConfig::default(), - private_inputs, - is_final, - true, - #[cfg(feature = "disk-spill")] - stark::storage_mode::StorageMode::Ram, - )?; - let boundary = - local_to_global::epoch_boundary(&mut provenance, label, &traces.touched_memory_cells); - - let start = EpochStart { - register_init: ®ister_init, - label, - }; - let epoch = prove_epoch(&elf, elf_bytes, &start, traces, is_final, &boundary, opts)?; - prev_fini = Some(epoch.reg_fini.clone()); - - // Carry the image forward: this epoch's fini is the next epoch's init. - for cell in &boundary { - image.set(cell.address, (cell.fini.value & 0xFF) as u8); + drop(tx); + + // Global prove, overlapped: drain the boundary channel until the + // producer hangs up (last epoch prepared), then prove the cross-epoch + // global memory argument WHILE the tail epochs are still proving. The + // global proof consumes only execution artifacts (boundaries, ELF, + // genesis pages) — never an epoch proof — so this is pure schedule. + let global_result_ref = &global_result; + let init_page_data_ref = &init_page_data; + scope.spawn(move || { + let mut all: Vec>> = Vec::new(); + while let Ok(b) = boundary_rx.recv() { + all.push(b); + } + let run = || -> Result { + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("prove_global"); + let num_private_input_pages = page::private_input_page_count(private_inputs); + // SINGLE source of truth: the same page-base list drives the + // committed GLOBAL_MEMORY tables and is shipped in the bundle, + // so the two can never diverge in set or order. + let touched = touched_page_bases(&all); + let global = prove_global( + &all, + elf_bytes, + init_page_data_ref, + &touched, + num_private_input_pages, + opts, + )?; + Ok((global, touched, num_private_input_pages)) + }; + *global_result_ref.lock().unwrap() = Some(run()); + }); + + // Prove epochs concurrently. Epoch proofs are mutually independent — + // each is seeded by its own label-domain-separated transcript + // (`epoch_transcript`) and nothing in an epoch's proof feeds the next + // one (the execution chain lives entirely in the producer above) — so + // `workers` provers pull prepared epochs and prove in parallel. + // Results are re-ordered by epoch index before the bundle is + // assembled, so proof bytes are identical to the sequential schedule. + // + // Sizing: each concurrent epoch prove costs its own peak VRAM/heap + // (~9 GB VRAM per 2^20-cycle epoch measured on ethrex), so the + // default is a conservative 2; tune with LAMBDA_VM_EPOCH_CONCURRENCY. + let prover_handles: Vec<_> = (0..workers) + .map(|k| scope.spawn(move || prove_worker(k))) + .collect(); + for h in prover_handles { + h.join().map_err(|_| { + Error::ContinuationInvariant("epoch prover thread panicked".to_string()) + })?; } + producer.join().map_err(|_| { + Error::ContinuationInvariant("epoch preparation thread panicked".to_string()) + })?; + Ok(()) + })?; + if let Some(e) = first_err.into_inner().unwrap() { + return Err(e); + } + let mut results = results.into_inner().unwrap(); + results.sort_by_key(|(index, _)| *index); + for (_, epoch) in results { epochs.push(epoch); - all_boundaries.push(boundary); + } + + // One global LogUp over all the (kept) local-to-global tables — proven + // concurrently by the scoped thread above; collect its result here. The + // scope guarantees the thread finished, so `None` is unreachable. + let (global, touched_page_bases, num_private_input_pages) = + global_result.into_inner().unwrap().ok_or_else(|| { + Error::ContinuationInvariant("global prove thread produced no result".to_string()) + })??; - if is_final { - break; + // Same timeline output as the monolithic path (prover/src/lib.rs): print + // the wall-clock span tree and honor LAMBDA_VM_TIMELINE_JSON. Without this, + // continuation runs record spans that are never drained (the profiling + // toolkit's phase_table.py consumes the JSON). + #[cfg(feature = "instruments")] + { + drop(__root); + let spans = stark::instruments::take_timeline(); + print!("{}", stark::instruments::format_timeline(&spans)); + if let Ok(path) = std::env::var("LAMBDA_VM_TIMELINE_JSON") { + let _ = std::fs::write(&path, stark::instruments::timeline_json(&spans)); + println!("[timeline] wrote {path}"); } - index += 1; } - // One global LogUp over all the (kept) local-to-global tables. `all_boundaries` was - // accumulated locally in the loop (never round-tripped through the bundle). - let num_private_input_pages = page::private_input_page_count(private_inputs); - // SINGLE source of truth: the same page-base list drives the committed GLOBAL_MEMORY - // tables and is shipped in the bundle, so the two can never diverge in set or order. - let touched_page_bases = touched_page_bases(&all_boundaries); - let global = prove_global( - &all_boundaries, - elf_bytes, - &init_page_data, - &touched_page_bases, - num_private_input_pages, - opts, - )?; - Ok(ContinuationProof { epochs, global, diff --git a/prover/src/tables/branch.rs b/prover/src/tables/branch.rs index b5bfe83b6..0d3c2e206 100644 --- a/prover/src/tables/branch.rs +++ b/prover/src/tables/branch.rs @@ -410,6 +410,7 @@ fn carry_1_expr>( /// - idx 2: `JALR·carry_0·(1 − carry_0)` on the register path (degree 3); /// - idx 3: `JALR·carry_1·(1 − carry_1)` on the register path (degree 3); /// - idx 4: `JALR·(1 − JALR)` (degree 2). +#[derive(Clone, Copy)] pub struct BranchConstraints; impl ConstraintSet for BranchConstraints { diff --git a/prover/src/tables/commit.rs b/prover/src/tables/commit.rs index 4660c7fb0..65e74f182 100644 --- a/prover/src/tables/commit.rs +++ b/prover/src/tables/commit.rs @@ -732,6 +732,7 @@ pub fn bus_interactions() -> Vec { /// - idx 3: `(first + end)·(1 − μ) = 0` (first/end ⇒ μ); /// - idx 4,5: `ADD` pair `address + 1 = address_incr` (unconditional); /// - idx 6,7: `ADD` pair `count_decr + 1 = count` (unconditional). +#[derive(Clone, Copy)] pub struct CommitConstraints; impl ConstraintSet for CommitConstraints { diff --git a/prover/src/tables/cpu32.rs b/prover/src/tables/cpu32.rs index 8b1bf86d6..60b83e84e 100644 --- a/prover/src/tables/cpu32.rs +++ b/prover/src/tables/cpu32.rs @@ -597,6 +597,7 @@ pub fn bus_interactions() -> Vec { /// - idx 25,26: `read_register2·imm[i]` (arg2 exclusivity); /// - idx 27-31: `(1 − μ)·flag` for `read_register1/2`, `write_register`, /// `signed`, `res_sign`. +#[derive(Clone, Copy)] pub struct Cpu32Constraints; impl ConstraintSet for Cpu32Constraints { diff --git a/prover/src/tables/dvrm.rs b/prover/src/tables/dvrm.rs index 9f979742b..c499a72bf 100644 --- a/prover/src/tables/dvrm.rs +++ b/prover/src/tables/dvrm.rs @@ -980,6 +980,7 @@ use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; /// DVRM table constraints as a single-source [`ConstraintSet`]. No column /// configuration is needed (the DVRM layout is fixed via `cols`). +#[derive(Clone, Copy)] pub struct DvrmConstraints; impl DvrmConstraints { diff --git a/prover/src/tables/ecdas.rs b/prover/src/tables/ecdas.rs index ff0c41f84..6d4b8a908 100644 --- a/prover/src/tables/ecdas.rs +++ b/prover/src/tables/ecdas.rs @@ -272,6 +272,7 @@ use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; /// ECDAS transition constraints as a single-source [`ConstraintSet`] (200 /// total). No column configuration needed (the layout is fixed via `cols`). +#[derive(Clone, Copy)] pub struct EcdasConstraints; impl EcdasConstraints { diff --git a/prover/src/tables/ecsm.rs b/prover/src/tables/ecsm.rs index 5d0a9477f..746bef91c 100644 --- a/prover/src/tables/ecsm.rs +++ b/prover/src/tables/ecsm.rs @@ -691,6 +691,7 @@ use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; /// ECSM transition constraints as a single-source [`ConstraintSet`] (413 /// total). No column configuration needed (the layout is fixed via `cols`). +#[derive(Clone, Copy)] pub struct EcsmConstraints; impl EcsmConstraints { diff --git a/prover/src/tables/eq.rs b/prover/src/tables/eq.rs index 0f20ca695..f967becf4 100644 --- a/prover/src/tables/eq.rs +++ b/prover/src/tables/eq.rs @@ -250,6 +250,7 @@ pub fn bus_interactions() -> Vec { /// - idx 0,1: `ADD` pair `b + diff = a` (unconditional); /// - idx 2: `IS_BIT(invert)` (unconditional); /// - idx 3: `res = eq XOR invert`. +#[derive(Clone, Copy)] pub struct EqConstraints; impl ConstraintSet for EqConstraints { diff --git a/prover/src/tables/keccak.rs b/prover/src/tables/keccak.rs index 9626b7e3b..7b84cbd48 100644 --- a/prover/src/tables/keccak.rs +++ b/prover/src/tables/keccak.rs @@ -461,6 +461,7 @@ pub fn bus_interactions() -> Vec { /// `state_ptr` DWordHL); /// - idx 50: `μ · carry_1 = 0` (top-lane no-overflow), where `carry_1` is the /// high carry of `addr + 192 = state_ptr[24]`. +#[derive(Clone, Copy)] pub struct KeccakConstraints; impl ConstraintSet for KeccakConstraints { diff --git a/prover/src/tables/keccak_rnd.rs b/prover/src/tables/keccak_rnd.rs index afc5dee3a..1b121a8b9 100644 --- a/prover/src/tables/keccak_rnd.rs +++ b/prover/src/tables/keccak_rnd.rs @@ -903,6 +903,7 @@ pub fn bus_interactions() -> Vec { /// The KECCAK round table's 20 transition constraints as a single /// [`ConstraintSet`]: for `x ∈ 0..5`, `hw ∈ 0..4` (idx `x·4 + hw`), the μ-gated /// `IS_BIT` on `Cxz_right[x][hw]` — `μ · Cxz_right·(1 − Cxz_right)`. +#[derive(Clone, Copy)] pub struct KeccakRndConstraints; impl ConstraintSet for KeccakRndConstraints { diff --git a/prover/src/tables/load.rs b/prover/src/tables/load.rs index 1da3cf564..a18ebddd9 100644 --- a/prover/src/tables/load.rs +++ b/prover/src/tables/load.rs @@ -485,6 +485,7 @@ use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; /// LOAD table constraints as a single-source [`ConstraintSet`]. No column /// configuration is needed (the LOAD layout is fixed via `cols`). +#[derive(Clone, Copy)] pub struct LoadConstraints; impl LoadConstraints { diff --git a/prover/src/tables/lt.rs b/prover/src/tables/lt.rs index 86e88a9f7..fb7d34267 100644 --- a/prover/src/tables/lt.rs +++ b/prover/src/tables/lt.rs @@ -356,6 +356,7 @@ use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; /// LT table constraints as a single-source [`ConstraintSet`]. No column /// configuration is needed (the LT layout is fixed via `cols`). +#[derive(Clone, Copy)] pub struct LtConstraints; impl LtConstraints { diff --git a/prover/src/tables/memw.rs b/prover/src/tables/memw.rs index 338b85467..282b0c312 100644 --- a/prover/src/tables/memw.rs +++ b/prover/src/tables/memw.rs @@ -862,6 +862,7 @@ fn w2_expr>(b: &B) -> /// - idx 4-10: `IS_BIT` on `carry[0..6]`; /// - idx 11-13: `IS_BIT` on `write2`, `write4`, `write8`; /// - idx 14: `IS_BIT` (width sum is a bit). +#[derive(Clone, Copy)] pub struct MemwConstraints; impl ConstraintSet for MemwConstraints { diff --git a/prover/src/tables/memw_aligned.rs b/prover/src/tables/memw_aligned.rs index 0853bf5ff..47678f541 100644 --- a/prover/src/tables/memw_aligned.rs +++ b/prover/src/tables/memw_aligned.rs @@ -667,6 +667,7 @@ fn w2_expr>(b: &B) -> /// - idx 2,3: `IS_BIT` on `μ_read`, `μ_write`; /// - idx 4-6: `IS_BIT` on `write2`, `write4`, `write8`; /// - idx 7: `IS_BIT` (width sum is a bit). +#[derive(Clone, Copy)] pub struct MemwAlignedConstraints; impl ConstraintSet for MemwAlignedConstraints { diff --git a/prover/src/tables/memw_register.rs b/prover/src/tables/memw_register.rs index 590a55100..99d8819f9 100644 --- a/prover/src/tables/memw_register.rs +++ b/prover/src/tables/memw_register.rs @@ -491,6 +491,7 @@ pub fn bus_interactions() -> Vec { /// The MEMW_R table's 3 transition constraints as a single [`ConstraintSet`]: /// - idx 0,1: `IS_BIT` on `μ_read`, `μ_write`; /// - idx 2: `IS_BIT<μ_sum>` with `μ_sum = μ_read + μ_write`. +#[derive(Clone, Copy)] pub struct MemwRegisterConstraints; impl ConstraintSet for MemwRegisterConstraints { diff --git a/prover/src/tables/mul.rs b/prover/src/tables/mul.rs index 181fba514..a615f74df 100644 --- a/prover/src/tables/mul.rs +++ b/prover/src/tables/mul.rs @@ -693,6 +693,7 @@ use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; /// MUL table constraints as a single-source [`ConstraintSet`]. No column /// configuration is needed (the MUL layout is fixed via `cols`). +#[derive(Clone, Copy)] pub struct MulConstraints; impl MulConstraints { diff --git a/prover/src/tables/register.rs b/prover/src/tables/register.rs index 4da1f5efe..2ec38bc0d 100644 --- a/prover/src/tables/register.rs +++ b/prover/src/tables/register.rs @@ -113,7 +113,7 @@ pub type FinalRegisterStateMap = HashMap; /// Returns the Word addresses for all register table rows. /// /// x0-x31 use addresses 0..63, x254 uses address 508, x255 uses 510..511. -fn register_word_address_list() -> [u64; NUM_REGISTER_ADDRESSES] { +pub(crate) fn register_word_address_list() -> [u64; NUM_REGISTER_ADDRESSES] { let mut addrs = [0u64; NUM_REGISTER_ADDRESSES]; // x0-x31: addresses 0..63 for (i, addr) in addrs.iter_mut().enumerate().take(64) { @@ -268,6 +268,27 @@ pub fn fini_from_trace(trace: &TraceTable) .collect() } +/// [`fini_from_trace`] without the trace: derives the same final register file +/// directly from the collected final state, mirroring exactly how +/// [`generate_register_trace`] fills the `FINI` column (accessed registers take +/// their final value; never-accessed registers keep their init). Lets the +/// continuation producer chain epochs before this epoch's REGISTER trace is +/// generated. Pinned to the trace-derived values by +/// `fini_from_final_state_matches_trace`. +pub fn fini_from_final_state(final_state: &FinalRegisterStateMap, init: &[u32]) -> Vec { + register_word_address_list() + .iter() + .take(NUM_REGISTER_ADDRESSES) + .enumerate() + .map(|(row, word_addr)| { + final_state + .get(word_addr) + .map(|state| state.value) + .unwrap_or_else(|| init.get(row).copied().unwrap_or(0)) + }) + .collect() +} + // ========================================================================= // Preprocessed commitment // ========================================================================= diff --git a/prover/src/tables/shift.rs b/prover/src/tables/shift.rs index 5ac5a393f..75d7605db 100644 --- a/prover/src/tables/shift.rs +++ b/prover/src/tables/shift.rs @@ -739,6 +739,7 @@ use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; /// SHIFT table constraints as a single-source [`ConstraintSet`]. No column /// configuration is needed (the SHIFT layout is fixed via `cols`). +#[derive(Clone, Copy)] pub struct ShiftConstraints; impl ShiftConstraints { diff --git a/prover/src/tables/store.rs b/prover/src/tables/store.rs index ac30832e1..6509dbe54 100644 --- a/prover/src/tables/store.rs +++ b/prover/src/tables/store.rs @@ -258,6 +258,7 @@ pub fn bus_interactions() -> Vec { /// - idx 0-3: `IS_BIT` on `write2`, `write4`, `write8`, `μ` (unconditional); /// - idx 4: `(Σ width)·(1 − Σ width) = 0` (width sum is a bit); /// - idx 5: `(Σ width)·(1 − μ) = 0` (width ⇒ μ). +#[derive(Clone, Copy)] pub struct StoreConstraints; impl ConstraintSet for StoreConstraints { diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 5c6b3085e..60a7f995c 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -2655,6 +2655,68 @@ fn generate_page_tables( // Trace Generation // ============================================================================= +/// Per-ELF DECODE artifacts: the parsed instruction map, the pristine DECODE +/// trace (multiplicities all zero) and its PC→row index. They are a pure +/// function of the ELF, so continuation epochs build them once +/// ([`DecodeArtifacts::from_elf`]) and share them across every epoch's trace +/// build ([`Traces::from_image_and_logs_with_decode`]) instead of re-parsing +/// the ELF and regenerating the trace per epoch. +pub struct DecodeArtifacts { + instructions: U64HashMap, + decode_trace: TraceTable, + decode_pc_to_row: decode::PcToRow, +} + +impl DecodeArtifacts { + /// Parse the ELF and generate the pristine DECODE trace. + /// + /// IMPORTANT: uses `generate_decode_trace` (same as + /// `compute_precomputed_commitment`) so the DECODE trace row ordering + /// matches the AIR's hardcoded commitment. + pub fn from_elf(elf: &Elf) -> Result { + let instructions = decode::instructions_from_elf(elf) + .map_err(|e| Error::Execution(format!("Failed to parse instructions: {e}")))?; + let (decode_trace, decode_pc_to_row) = decode::generate_decode_trace(&instructions); + Ok(Self { + instructions, + decode_trace, + decode_pc_to_row, + }) + } +} + +/// An epoch's collected operations and end state (Phases 1-2 of the trace +/// build), produced by [`Traces::collect_epoch`] and consumed by +/// [`Traces::build_from_collected`]. Collection must run in epoch order (it +/// reads the advancing memory image); everything downstream of this struct is +/// epoch-local. The cross-epoch values a continuation producer needs — the +/// touched-cell set and the final register file — are derivable here, before +/// any table is generated. +pub struct CollectedEpoch { + ops: CollectedOps, + memory_state: MemoryState, + register_state: RegisterState, +} + +impl CollectedEpoch { + /// The epoch's touched memory cells (sorted by address): the exact values + /// `build_traces` later stores in `Traces::touched_memory_cells` (both are + /// [`touched_cells_from_memory_state`] over the same immutable + /// `memory_state`), available before any table is built. + pub fn touched_memory_cells(&self) -> local_to_global::EpochTouches { + touched_cells_from_memory_state(&self.memory_state) + } + + /// The epoch's final register file (`R_{i+1}`) in + /// `register_word_address_list` order: the exact values + /// [`register::fini_from_trace`] reads off the generated REGISTER trace + /// (see [`register::fini_from_final_state`]), available before the trace + /// exists. + pub fn register_fini(&self, register_init: &[u32]) -> Vec { + register::fini_from_final_state(&self.register_state.to_final_state_map(), register_init) + } +} + /// All generated trace tables. pub struct Traces { /// CPU execution traces (split into chunks of max_rows::CPU) @@ -2976,7 +3038,7 @@ fn build_traces( memory_state: &MemoryState, register_init: &[u32], decode_trace: TraceTable, - decode_pc_to_row: decode::PcToRow, + decode_pc_to_row: &decode::PcToRow, mut register_state: RegisterState, max_rows: &super::MaxRowsConfig, #[cfg(feature = "disk-spill")] storage_mode: StorageMode, @@ -3330,7 +3392,7 @@ fn build_traces( let mut decode = decode_trace; let mut decode_lookups: Vec = cpu_ops_ref.iter().map(|op| op.decode.pc).collect(); decode_lookups.extend(std::iter::repeat_n(cpu::CPU_PADDING_PC, num_padding_rows)); - decode::update_multiplicities(&mut decode, &decode_pc_to_row, &decode_lookups); + decode::update_multiplicities(&mut decode, decode_pc_to_row, &decode_lookups); decode }; let gen_commit = || commit::generate_commit_trace(&commit_ops); @@ -4223,6 +4285,68 @@ impl Traces { l2g_memory_bookend: bool, #[cfg(feature = "disk-spill")] storage_mode: StorageMode, ) -> Result { + let artifacts = DecodeArtifacts::from_elf(elf)?; + Self::from_image_and_logs_with_decode( + &artifacts, + initial_image, + register_init, + logs, + max_rows, + private_input, + is_final, + l2g_memory_bookend, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + } + + /// [`Self::from_image_and_logs`] with the per-ELF DECODE artifacts supplied + /// by the caller. Continuation epochs of the same ELF build + /// [`DecodeArtifacts`] once and reuse them here, so the per-epoch producer + /// chain skips the ELF re-parse and the pristine DECODE trace regeneration + /// (Phase 0) — the trace is cloned (a memcpy) and its multiplicities are + /// filled per epoch. + #[allow(clippy::too_many_arguments)] + pub fn from_image_and_logs_with_decode( + artifacts: &DecodeArtifacts, + initial_image: &I, + register_init: &[u32], + logs: &[Log], + max_rows: &super::MaxRowsConfig, + private_input: &[u8], + is_final: bool, + l2g_memory_bookend: bool, + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + ) -> Result { + let collected = + Self::collect_epoch(artifacts, initial_image, register_init, logs, is_final)?; + Self::build_from_collected( + artifacts, + collected, + Some(initial_image), + register_init, + max_rows, + private_input, + is_final, + l2g_memory_bookend, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + } + + /// The sequential-critical half of an epoch's trace build: log collection + /// and op routing (Phases 1-2), which read the pre-epoch memory image and + /// produce the epoch's memory/register end state. This must run in epoch + /// order (the image advances between epochs); the table generation that + /// consumes the result ([`Self::build_from_collected`]) is epoch-local and + /// can run on another thread. + pub fn collect_epoch( + artifacts: &DecodeArtifacts, + initial_image: &I, + register_init: &[u32], + logs: &[Log], + is_final: bool, + ) -> Result { // A non-final epoch must not contain the program-terminating instruction // (next_pc == 0). Otherwise the CPU sends an ECALL bus token with no HALT // table to receive it (HALT is excluded when !is_final), producing an @@ -4231,21 +4355,10 @@ impl Traces { return Err(Error::HaltInNonFinalEpoch); } - // Phase 0: ELF → DECODE + instructions - // IMPORTANT: Use generate_decode_trace (same as compute_precomputed_commitment) - // so the DECODE trace row ordering matches the AIR's hardcoded commitment. - #[cfg(feature = "instruments")] - let __sp = stark::instruments::span("p0_decode"); - let instructions = decode::instructions_from_elf(elf) - .map_err(|e| Error::Execution(format!("Failed to parse instructions: {e}")))?; - let (decode_trace, decode_pc_to_row) = decode::generate_decode_trace(&instructions); - #[cfg(feature = "instruments")] - drop(__sp); - // Phase 1: Logs → CPU operations #[cfg(feature = "instruments")] let __sp = stark::instruments::span("p1_cpu_ops"); - let cpu_ops = collect_cpu_ops(logs, &instructions)?; + let cpu_ops = collect_cpu_ops(logs, &artifacts.instructions)?; #[cfg(feature = "instruments")] drop(__sp); @@ -4289,17 +4402,51 @@ impl Traces { #[cfg(feature = "instruments")] drop(__sp); + Ok(CollectedEpoch { + ops, + memory_state, + register_state, + }) + } + + /// The epoch-local half of an epoch's trace build: table generation + /// (Phases 3-5) over an already-collected epoch. Reads nothing sequential — + /// continuation callers run this on a builder-pool thread while the + /// producer collects the next epoch. `initial_image` is only used for PAGE + /// tables and their bitwise lookups, both skipped in continuation mode + /// (`l2g_memory_bookend`), where callers pass `None`. + #[allow(clippy::too_many_arguments)] + pub fn build_from_collected( + artifacts: &DecodeArtifacts, + collected: CollectedEpoch, + initial_image: Option<&I>, + register_init: &[u32], + max_rows: &super::MaxRowsConfig, + private_input: &[u8], + is_final: bool, + l2g_memory_bookend: bool, + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + ) -> Result { + // Phase 0 (cached): the pristine DECODE trace is cloned so + // `build_traces` can fill this epoch's multiplicities. + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("p0_decode"); + let decode_trace = artifacts.decode_trace.clone(); + let decode_pc_to_row = &artifacts.decode_pc_to_row; + #[cfg(feature = "instruments")] + drop(__sp); + // Phases 3-5 #[cfg(feature = "instruments")] let __sp = stark::instruments::span("p3to5_build_traces"); let result = build_traces( - ops, - Some(initial_image), - &memory_state, + collected.ops, + initial_image, + &collected.memory_state, register_init, decode_trace, decode_pc_to_row, - register_state, + collected.register_state, max_rows, #[cfg(feature = "disk-spill")] storage_mode, @@ -4370,7 +4517,7 @@ impl Traces { &memory_state, ®ister_init, decode_trace, - decode_pc_to_row, + &decode_pc_to_row, register_state, max_rows, #[cfg(feature = "disk-spill")] diff --git a/prover/src/test_utils.rs b/prover/src/test_utils.rs index 6dd28ce71..d7969612f 100644 --- a/prover/src/test_utils.rs +++ b/prover/src/test_utils.rs @@ -10,6 +10,7 @@ //! - Minimal trace generation for testing //! - AIR creation helpers +use std::collections::HashMap; use std::path::PathBuf; use crypto::fiat_shamir::is_transcript::IsStarkTranscript; @@ -604,7 +605,44 @@ pub fn generate_minimal_bitwise_trace(ops: &[BitwiseOperation]) -> TraceTable + 'static>( +/// Process-wide cache of pre-built, pre-captured AIR prototypes, keyed by +/// `(table name, proof options)`. Constructing an [`AirWithBuses`] runs every +/// constraint body through a MetaBuilder, and its first `constraint_program()` +/// runs them again for the IR capture — for the big tables (ECDAS/ECSM/ +/// KECCAK_RND, 16-25K IR nodes) that is by far the dominant cost of building +/// an AIR. Continuation epochs rebuild the full table set per epoch and shard +/// tables build one instance per shard, so without this cache the same walks +/// re-run dozens of times per prove. A prototype is built (and captured) once; +/// every later request clones it — the clone copies the derived metadata and +/// the already-captured IR, never re-running the bodies. +/// +/// Key correctness: for a given name the constraint set and bus interactions +/// are a pure function of the table module (PAGE embeds its page base in the +/// name), and `ProofOptions` covers everything else `AirWithBuses::new` reads. +/// The cached prototype is pristine — `with_name` / `with_preprocessed` apply +/// to the caller's clone only. +fn air_prototype_cache() +-> &'static std::sync::Mutex>> { + static CACHE: std::sync::OnceLock< + std::sync::Mutex>>, + > = std::sync::OnceLock::new(); + CACHE.get_or_init(|| std::sync::Mutex::new(HashMap::new())) +} + +type AirProtoKey = (String, u8, usize, u64, u8, u8); + +fn air_proto_key(name: &str, o: &ProofOptions) -> AirProtoKey { + ( + name.to_string(), + o.blowup_factor, + o.fri_number_of_queries, + o.coset_offset, + o.grinding_factor, + o.fri_final_poly_log_degree, + ) +} + +fn build_air + Clone + Send + Sync + 'static>( num_columns: usize, interactions: Vec, proof_options: &ProofOptions, @@ -612,14 +650,30 @@ fn build_air + 'static>( constraint_set: CS, name: &str, ) -> AirWithBuses { - AirWithBuses::new( + type Proto = AirWithBuses; + let key = air_proto_key(name, proof_options); + { + let cache = air_prototype_cache().lock().unwrap(); + if let Some(proto) = cache.get(&key).and_then(|p| p.downcast_ref::>()) { + return proto.clone(); + } + } + let air = Proto::::new( num_columns, AuxiliaryTraceBuildData { interactions }, proof_options, step_size, constraint_set, ) - .with_name(name) + .with_name(name); + // Pre-capture the constraint IR so every clone carries it (the prover's + // GPU lowering and interpreter paths force it per instance otherwise). + let _ = air.constraint_program(); + air_prototype_cache() + .lock() + .unwrap() + .insert(key, Box::new(air.clone())); + air } /// Create CPU AIR with all constraints and bus interactions. diff --git a/prover/src/tests/ir_stats_dump.rs b/prover/src/tests/ir_stats_dump.rs new file mode 100644 index 000000000..a269127d9 --- /dev/null +++ b/prover/src/tests/ir_stats_dump.rs @@ -0,0 +1,133 @@ +//! Diagnostic dump (ignored by default): per-table constraint-IR and +//! device-lowering stats — node counts by dim, op mix, and the lowered slot +//! footprint that sizes the GPU kernel's per-thread scratch. Use it to gauge +//! how a lowering change moves the scratch working set (and therefore the +//! constraint kernel's cache behavior / thread-count headroom). Run with: +//! `cargo test -p lambda-vm-prover ir_stats_dump -- --ignored --nocapture` + +use stark::constraint_ir::DeviceProgram; +use stark::constraint_ir::{ConstraintProgram, Dim, Op}; +use stark::proof::options::GoldilocksCubicProofOptions; +use stark::traits::AIR; + +use crate::tables::types::{GoldilocksExtension, GoldilocksField}; +use crate::test_utils::*; + +type Gl = GoldilocksField; +type Ext3 = GoldilocksExtension; + +fn dump(label: &str, prog: &ConstraintProgram) { + let n = prog.nodes.len(); + let mut base_nodes = 0usize; + let mut ext_nodes = 0usize; + let mut uniform = 0usize; // row-invariant: consts/challenges/alpha/offset + let mut base_mul = 0usize; + let mut ext_mul = 0usize; + let mut base_addsub = 0usize; + let mut ext_addsub = 0usize; + let mut vars_main = 0usize; + let mut vars_aux = 0usize; + let mut embeds = 0usize; + // mixed = ext-dim binop with at least one base-dim operand (implicit embed) + let mut mixed_binops = 0usize; + + for (op, dim) in prog.nodes.iter().zip(prog.dims.iter()) { + match dim { + Dim::Base => base_nodes += 1, + Dim::Ext => ext_nodes += 1, + } + match *op { + Op::ConstBase(_) + | Op::ConstExt(_) + | Op::RapChallenge { .. } + | Op::AlphaPow { .. } + | Op::TableOffset => uniform += 1, + Op::Var { main, .. } => { + if main { + vars_main += 1 + } else { + vars_aux += 1 + } + } + Op::Mul(a, b) => { + if *dim == Dim::Base { + base_mul += 1 + } else { + ext_mul += 1; + if prog.dims[a as usize] == Dim::Base || prog.dims[b as usize] == Dim::Base { + mixed_binops += 1; + } + } + } + Op::Add(a, b) | Op::Sub(a, b) => { + if *dim == Dim::Base { + base_addsub += 1 + } else { + ext_addsub += 1; + if prog.dims[a as usize] == Dim::Base || prog.dims[b as usize] == Dim::Base { + mixed_binops += 1; + } + } + } + Op::Neg(_) => {} + Op::Embed(_) => embeds += 1, + } + } + + // Lowered slot footprint: old scratch was nodes×24B per thread; new is + // base_slots×8 + ext_slots×24. + let dev = DeviceProgram::lower(prog); + let old_bytes = n * 24; + let new_bytes = dev.num_base_slots as usize * 8 + dev.num_ext_slots as usize * 24; + + println!( + "{label:12} nodes={n:6} base={base_nodes:6} ({:4.1}%) ext={ext_nodes:5} uniform={uniform:4} \ + mul(b/e)={base_mul:5}/{ext_mul:4} addsub(b/e)={base_addsub:5}/{ext_addsub:4} \ + mixed={mixed_binops:4} var(m/a)={vars_main:4}/{vars_aux:3} embed={embeds:3} \ + roots={:4} num_base={:4} | slots(b/e)={}/{} scratch/thr {}B -> {}B ({:.1}x)", + 100.0 * base_nodes as f64 / n as f64, + prog.roots.len(), + prog.num_base, + dev.num_base_slots, + dev.num_ext_slots, + old_bytes, + new_bytes, + old_bytes as f64 / new_bytes as f64, + ); +} + +fn stats(air: &dyn AIR, label: &str) { + dump(label, air.constraint_program()); +} + +#[test] +#[ignore = "analysis dump, not a test"] +fn dump_ir_stats() { + let opts = GoldilocksCubicProofOptions::with_blowup(2).expect("blowup=2 valid"); + + stats(&create_cpu_air(&opts), "CPU"); + stats(&create_bitwise_air(&opts), "BITWISE"); + stats(&create_lt_air(&opts), "LT"); + stats(&create_shift_air(&opts), "SHIFT"); + stats(&create_eq_air(&opts), "EQ"); + stats(&create_bytewise_air(&opts), "BYTEWISE"); + stats(&create_store_air(&opts), "STORE"); + stats(&create_cpu32_air(&opts), "CPU32"); + stats(&create_memw_air(&opts), "MEMW"); + stats(&create_memw_aligned_air(&opts), "MEMW_A"); + stats(&create_memw_register_air(&opts), "MEMW_R"); + stats(&create_load_air(&opts), "LOAD"); + stats(&create_decode_air(&opts), "DECODE"); + stats(&create_mul_air(&opts), "MUL"); + stats(&create_dvrm_air(&opts), "DVRM"); + stats(&create_branch_air(&opts), "BRANCH"); + stats(&create_halt_air(&opts), "HALT"); + stats(&create_commit_air(&opts), "COMMIT"); + stats(&create_page_air(&opts, 0x1000), "PAGE"); + stats(&create_register_air(&opts), "REGISTER"); + stats(&create_keccak_air(&opts), "KECCAK"); + stats(&create_keccak_rnd_air(&opts), "KECCAK_RND"); + stats(&create_keccak_rc_air(&opts), "KECCAK_RC"); + stats(&create_ecsm_air(&opts), "ECSM"); + stats(&create_ecdas_air(&opts), "ECDAS"); +} diff --git a/prover/src/tests/register_tests.rs b/prover/src/tests/register_tests.rs index 433968ab5..66dcd2662 100644 --- a/prover/src/tests/register_tests.rs +++ b/prover/src/tests/register_tests.rs @@ -133,3 +133,43 @@ fn test_precomputed_commitment_with_fini_binds_fini() { // The 3-column (with-fini) commitment differs from the 2-column monolithic one. assert_ne!(root_a, compute_precomputed_commitment(&opts, &init)); } + +/// `fini_from_final_state` must return exactly what `fini_from_trace` reads off +/// the generated REGISTER trace, for both accessed and never-accessed +/// registers — it is the continuation producer's trace-free replacement. +#[test] +fn fini_from_final_state_matches_trace() { + let entry_point = 0x8000_0123u64; + let init = register_init_from_entry_point(entry_point); + + // Touch a scattered subset of word addresses (lo/hi words, x255 PC word) + // with distinct values; leave the rest untouched. + let addr_list = register_word_address_list(); + let mut final_state = FinalRegisterStateMap::new(); + for (k, &word_addr) in addr_list.iter().take(NUM_REGISTER_ADDRESSES).enumerate() { + if k % 3 == 0 { + final_state.insert( + word_addr, + FinalRegisterWordState { + timestamp: 1000 + k as u64, + value: 0xABC0_0000 | k as u32, + }, + ); + } + } + + let trace = generate_register_trace(&final_state, &init); + assert_eq!( + fini_from_final_state(&final_state, &init), + fini_from_trace(&trace), + "trace-free FINI derivation must match the generated trace" + ); + + // Empty final state: FINI == init on every row. + let empty = FinalRegisterStateMap::new(); + let trace = generate_register_trace(&empty, &init); + assert_eq!( + fini_from_final_state(&empty, &init), + fini_from_trace(&trace) + ); +} diff --git a/prover/tests/gpu_constraint_interp_real.rs b/prover/tests/gpu_constraint_interp_real.rs index e464f2556..2cea4be1b 100644 --- a/prover/tests/gpu_constraint_interp_real.rs +++ b/prover/tests/gpu_constraint_interp_real.rs @@ -30,7 +30,8 @@ use math_cuda::device::backend; use math_cuda::lde::{GpuLdeBase, GpuLdeExt3}; use stark::constraint_ir::device::{ - DeviceProgram, OP_ALPHA_POW, OP_RAP_CHALLENGE, OP_VAR, eval_device_program, unpack_var, + DeviceProgram, OP_ADD, OP_ALPHA_POW, OP_EMBED, OP_MUL, OP_NEG, OP_RAP_CHALLENGE, OP_SUB, + OP_VAR, OPK_ALPHA, OPK_PAYLOAD_MASK, OPK_RAP, OPK_SHIFT, eval_device_program, unpack_var, }; use stark::constraint_ir::gpu_interp::try_eval_program_gpu; use stark::proof::options::GoldilocksCubicProofOptions; @@ -68,6 +69,17 @@ fn enc(x: &Fp3) -> [u64; 3] { /// program actually references. fn program_footprint(dev: &DeviceProgram) -> (usize, usize, usize, usize, usize) { let (mut main_cols, mut aux_cols, mut rap_len, mut alpha_len, mut max_off) = (0, 0, 0, 0, 0); + // Uniform leaves are propagated into operand encodings, so the RAP/alpha + // footprint must be read from the operands of arithmetic nodes (the + // root-pinned leaf-node forms are kept for completeness). + let scan_operand = |enc: u32, rap_len: &mut usize, alpha_len: &mut usize| { + let payload = (enc & OPK_PAYLOAD_MASK) as usize; + match enc >> OPK_SHIFT { + OPK_RAP => *rap_len = (*rap_len).max(payload + 1), + OPK_ALPHA => *alpha_len = (*alpha_len).max(payload + 1), + _ => {} + } + }; for n in &dev.nodes { match n.op { OP_VAR => { @@ -82,6 +94,11 @@ fn program_footprint(dev: &DeviceProgram) -> (usize, usize, usize, usize, usize) } OP_RAP_CHALLENGE => rap_len = rap_len.max(n.a as usize + 1), OP_ALPHA_POW => alpha_len = alpha_len.max(n.a as usize + 1), + OP_ADD | OP_SUB | OP_MUL => { + scan_operand(n.a, &mut rap_len, &mut alpha_len); + scan_operand(n.b, &mut rap_len, &mut alpha_len); + } + OP_NEG | OP_EMBED => scan_operand(n.a, &mut rap_len, &mut alpha_len), _ => {} } } @@ -145,6 +162,7 @@ fn check_air(air: &dyn AIR, stream.synchronize().expect("sync uploads"); let main = GpuLdeBase { + ready: None, buf: Arc::new(base_dev), m: main_cols, lde_size, @@ -153,6 +171,7 @@ fn check_air(air: &dyn AIR, trace_rows: 0, }; let aux = GpuLdeExt3 { + ready: None, buf: Arc::new(aux_dev), m: aux_cols, lde_size, diff --git a/scripts/profiling/README.md b/scripts/profiling/README.md new file mode 100644 index 000000000..276681515 --- /dev/null +++ b/scripts/profiling/README.md @@ -0,0 +1,218 @@ +# GPU profiling toolkit + +Tooling for profiling the CUDA prover on the dedicated RTX 5090 box. The full +methodology (what to measure, in what order, and how to read it) lives in +`thoughts/gpu-profiling/plan.md`; this directory is the executable part. + +## One-time machine setup + +```bash +scripts/profiling/setup_machine.sh # apt tooling, nsight, perf/eBPF perms — then REBOOT +``` + +Prereqs handled elsewhere: NVIDIA driver, and the build toolchain + guest +programs per `scripts/SERVER_SETUP.md` (`make compile-programs-asm`, etc.). +On Blackwell (RTX 5090, sm_120) CUDA ≥ 12.8 and nsys/ncu ≥ 2025.1 are hard +requirements. + +## Before every session + +```bash +sudo scripts/profiling/bench_mode.sh on # lock SM clocks, persistence, governor +make test-cuda-integration # sanity: every GPU counter fires +``` + +`bench_mode.sh off` restores defaults. Numbers taken with floating clocks are +not comparable across sessions. + +## The main entry point + +```bash +# 3 instrumented runs + phase table with GPU util per phase: +scripts/profiling/run_profile.sh executor/program_artifacts/rust/ethrex.elf \ + --private-input executor/tests/ethrex_5_transfers.bin + +# the real workload, plus an nsys-traced run and the per-phase GPU-busy report: +scripts/profiling/run_profile.sh --nsys \ + executor/program_artifacts/rust/ethrex.elf \ + --private-input executor/tests/ethrex_5_transfers.bin + +# big continuation run, nsys capture limited to one `proving` span (one epoch): +LAMBDA_VM_NSYS_CAPTURE_SPAN=epoch_prove scripts/profiling/run_profile.sh --nsys --continuations \ + executor/program_artifacts/rust/ethrex.elf \ + --private-input executor/tests/ethrex_10_transfers.bin +``` + +Each invocation produces a self-contained bundle under +`reports/__/`: + +| file | what | +|---|---| +| `env.json` | driver, clocks, sha, env — the context that makes numbers comparable | +| `phase_table.md` | warm-run phase tree: median ms, % of total, jitter, GPU util per phase | +| `phase_table_cold.md` | run 1 alone (module load, twiddle build, mempool growth) | +| `trace_perfetto.json` | span tree for ui.perfetto.dev | +| `nsys_report.nsys-rep` | open in the Nsight Systems GUI (scp to a laptop) | +| `nsys_stats.txt` | stock nsys summaries (kernels, memcpy, API, NVTX) | +| `phase_busy.md` | **the ranking input**: per-phase GPU busy %, memcpy, top kernels | + +How to read `phase_busy.md`: a phase with low busy% is host-bound — fix +pipeline/overlap/syncs, don't touch its kernels. A phase with high busy% is +kernel-bound — take its top kernels to Nsight Compute. + +## What each column means + +`phase_table.md` (from the instruments spans + the NVML sampler): + +| column | meaning | +|---|---| +| `median` | wall-clock of the span, median across warm runs; repeated spans (per-epoch, per-table) are **summed within a run** first (`n/run` = instances) | +| `% of total` | share of the root span (`prove_total` / `prove_continuation_total`) | +| `cv%` | run-to-run spread (stdev/mean) — the noise floor an optimization claim must beat | +| `gpu% / mem%` | average `nvidia-smi` utilization **sampled at 10 Hz inside the span's wall window**. Coarse: SM-occupancy-ish, includes other phases' async kernels landing in the window | +| `vram MiB` | max VRAM sampled inside the window | +| instances tables | per-instance wall, `gap→next` (host time between consecutive instances), gpu% inside the span vs inside the gap | + +`phase_busy.md` (from the nsys sqlite export; only exists with `--nsys`): + +| column | meaning | +|---|---| +| `wall ms` | union of that phase's NVTX windows (merged if overlapping) | +| `kernel-sum ms` | sum of kernel durations **attributed by correlation ID to launches made inside the phase** — can exceed wall when streams overlap | +| `gpu-busy ms / busy%` | union coverage of kernel intervals clipped to the phase windows — the honest "GPU was doing *something*" number; the one to rank phases by | +| `h2d / d2h ms/MiB` | memcpy time and volume attributed to the phase | + +The two GPU numbers answer different questions: `gpu%` (NVML) is a cheap +always-on sanity signal; `busy%` (nsys) is the precise one — trust it when +they disagree. Attribution is by *launch site*, so async kernels count toward +the phase that enqueued them even if they execute later. + +## Reference: scripts and knobs + +| script | what it does / flags | +|---|---| +| `run_profile.sh [opts] [--private-input ]` | the bundle. `--runs N` (default 3; run 1 kept separately as cold), `--nsys`, `--gpu-metrics` (needs the counters permission), `--continuations`, `--out DIR`, `--no-build`. Env: `PROFILE_FEATURES` (default `nvtx,jemalloc-stats`), `EXTRA_PROVE_ARGS` | +| `flamegraphs.sh [opts] …` | on-CPU + off-CPU SVGs. `--offcpu-secs N` overrides the capture window, `--skip-offcpu`, `--no-build`, `--continuations` | +| `bench_mode.sh on [mhz] / off` | lock/unlock SM clocks (default 90% of max), persistence, governor | +| `capture_env.sh` | env JSON to stdout — attach to anything you measure by hand | +| `phase_table.py [--util u.csv]… tl.json…` | aggregate timelines; `--instances LABEL` adds per-instance tables for deeper repeated spans, `--min-pct X` hides noise rows | +| `nsys_phase_busy.py report.sqlite [--top N]` | the GPU busy report from `nsys export --type sqlite` | +| `nvml_sampler.py -o out.csv [-i 0.1]` | standalone 10 Hz GPU util sampler (epoch-ns timestamps, aligns with span `start_ns`) | +| `timeline_to_perfetto.py tl.json > trace.json` | span tree for ui.perfetto.dev | + +Environment variables the tooling understands: + +| var | effect | +|---|---| +| `LAMBDA_VM_TIMELINE_JSON=` | prover writes the span timeline there (needs `instruments`) | +| `LAMBDA_VM_NSYS_CAPTURE_SPAN=