From c683c69ccb0499fc4528ece39f292f53b8c71609 Mon Sep 17 00:00:00 2001 From: Arthur Paulino Date: Fri, 4 Sep 2026 02:22:30 +0000 Subject: [PATCH] Balance large CUDA proofs across CPU and GPU Large mixed-height proofs no longer make an all-GPU or all-CPU placement choice. Partition complete height groups from measured device capacity and host parallelism, overlap Rayon LDE construction with CUDA, and merge canonical Goldilocks/BLAKE3 digest frontiers into one byte-compatible Merkle commitment. Carry hybrid commitments through lookup, quotient, openings, and FRI. Keep reusable matrices resident, defer host copies until they are needed, spill by phase-specific byte budgets, evaluate host-backed quotient inputs through bounded staging buffers, split exceptionally large lookup rows across CPU and GPU, and batch resident FRI reductions without changing transcript order. Reuse pinned host allocations to keep repeated spill traffic bounded. Methodology: build Ix cb23e50 in release mode against this checkout, compile Init once, split its manifest into 16 deterministic shards, prove shards 7 and 8, independently reprove their recursive wrap slots 12 and 13 with --jobs 1 --max-ram 240, then independently reprove structural join slot 14. Measure external wall time and maximum RSS with /usr/bin/time -v on an RTX PRO 6000 Blackwell (97,887 MiB), CUDA 13.3/driver 595.84, a 32-thread Xeon Platinum 8559C host, and 249 GiB RAM. The configuration uses 64 FRI queries. The CPU baseline was 102.27s and 117.25s for the shard proofs, 90.81s and 93.11s for their lifts, and about 74.10s for the join (the join is isolated by subtracting the two lifts from a 258.02s recursive run): 477.54s total. CUDA measured 57.87s, 65.84s, 41.03s, 41.59s, and 33.26s respectively: 239.59s total, a 1.99x end-to-end speedup. The recursive portion falls from 258.02s to 115.88s (2.23x), and peak external RSS falls from 209.03 GiB to 143.27 GiB. Excluding Aiur execution and witness generation, CUDA STARK proving took 24.28s for shard 7 and 34.09s for shard 8. Validate the normal CUDA-independent build with 35 release tests. Validate the opt-in CUDA path with 53 release tests, clippy with warnings denied, and the complete cuda/smoke.sh suite, including field arithmetic, DFT/coset-LDE, BLAKE3/MMCS, lookup/quotient, interpolation/FRI, forced spill/reopen, CPU verification, and byte-for-byte CPU/CUDA proof comparison. --- Cargo.lock | 1 + Cargo.toml | 3 +- cuda/kernels.cu | 1109 ++++++++++++++++++++++++++++++++++-- flake.nix | 2 +- rust-toolchain.toml | 2 +- src/cuda/mmcs.rs | 1077 ++++++++++++++++++++++++++++++++++- src/cuda/mod.rs | 932 ++++++++++++++++++++++++++++-- src/cuda/pcs.rs | 1329 ++++++++++++++++++++++++++++++++++++++----- src/lookup.rs | 53 ++ src/types.rs | 563 +++++++++++++++--- 10 files changed, 4752 insertions(+), 319 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 68b4682..7c13591 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -464,6 +464,7 @@ dependencies = [ "p3-symmetric", "p3-util", "rand", + "rayon", "serde", "tracing", "tracing-subscriber", diff --git a/Cargo.toml b/Cargo.toml index 52315db..e3aac95 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ rust-version = "1.98" [dependencies] tracing = "0.1" itertools = { version = "0.14", optional = true } +rayon = { version = "1", optional = true } serde = { version = "1", features = ["derive"] } bincode = { version = "2", features = ["serde"] } p3-air = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" } @@ -47,7 +48,7 @@ parallel = ["p3-maybe-rayon/parallel"] # Use the first-party CUDA Goldilocks DFT/LDE backend. Enabling this feature # requires a CUDA toolkit at build time and an NVIDIA GPU at runtime; the # default CPU build never invokes nvcc or links the CUDA runtime. -cuda = ["dep:itertools"] +cuda = ["dep:itertools", "dep:rayon"] # Similar to `release`, but preserves debug info [profile.dev-ci] diff --git a/cuda/kernels.cu b/cuda/kernels.cu index a9d6680..9541102 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -9,6 +9,12 @@ #include #include #include +#include +#include +#include +#include +#include +#include #include namespace { @@ -367,6 +373,109 @@ class HostRegistration { bool registered_ = false; }; +// Mixed quotient evaluation double-buffers host rows through pinned memory. +// Allocating and releasing those large buffers for every circuit is expensive, +// so retain one complete double-buffer working set across sequential jobs. +// Leases make reuse safe when multiple provers call into CUDA concurrently: +// callers beyond the two cached slots receive an uncached temporary buffer. +struct PinnedHostPoolEntry { + uint64_t* pointer = nullptr; + size_t capacity = 0; + bool in_use = false; +}; + +constexpr size_t PINNED_HOST_POOL_SIZE = 2; +PinnedHostPoolEntry pinned_host_pool[PINNED_HOST_POOL_SIZE]; +pthread_mutex_t pinned_host_pool_mutex = PTHREAD_MUTEX_INITIALIZER; + +class PinnedHostBuffer { + public: + PinnedHostBuffer() = default; + PinnedHostBuffer(const PinnedHostBuffer&) = delete; + PinnedHostBuffer& operator=(const PinnedHostBuffer&) = delete; + + ~PinnedHostBuffer() { + if (pool_index_ < PINNED_HOST_POOL_SIZE) { + pthread_mutex_lock(&pinned_host_pool_mutex); + pinned_host_pool[pool_index_].in_use = false; + pthread_mutex_unlock(&pinned_host_pool_mutex); + } else if (pointer_ != nullptr) { + cudaFreeHost(pointer_); + } + } + + cudaError_t allocate(size_t bytes) { + if (bytes == 0 || pointer_ != nullptr) { + return cudaErrorInvalidValue; + } + + pthread_mutex_lock(&pinned_host_pool_mutex); + size_t reusable = PINNED_HOST_POOL_SIZE; + for (size_t i = 0; i < PINNED_HOST_POOL_SIZE; ++i) { + const auto& entry = pinned_host_pool[i]; + if (!entry.in_use && entry.pointer != nullptr && entry.capacity >= bytes && + (reusable == PINNED_HOST_POOL_SIZE || + entry.capacity < pinned_host_pool[reusable].capacity)) { + reusable = i; + } + } + if (reusable < PINNED_HOST_POOL_SIZE) { + auto& entry = pinned_host_pool[reusable]; + entry.in_use = true; + pointer_ = entry.pointer; + pool_index_ = reusable; + pthread_mutex_unlock(&pinned_host_pool_mutex); + return cudaSuccess; + } + + size_t available = PINNED_HOST_POOL_SIZE; + for (size_t i = 0; i < PINNED_HOST_POOL_SIZE; ++i) { + if (!pinned_host_pool[i].in_use && + (available == PINNED_HOST_POOL_SIZE || + pinned_host_pool[i].capacity < pinned_host_pool[available].capacity)) { + available = i; + } + } + if (available == PINNED_HOST_POOL_SIZE) { + pthread_mutex_unlock(&pinned_host_pool_mutex); + return cudaMallocHost(reinterpret_cast(&pointer_), bytes); + } + + auto& entry = pinned_host_pool[available]; + entry.in_use = true; + uint64_t* old_pointer = entry.pointer; + const size_t old_capacity = entry.capacity; + entry.pointer = nullptr; + entry.capacity = 0; + pool_index_ = available; + pthread_mutex_unlock(&pinned_host_pool_mutex); + + cudaError_t status = old_pointer == nullptr ? cudaSuccess : cudaFreeHost(old_pointer); + const bool old_pointer_freed = status == cudaSuccess; + if (status == cudaSuccess) { + status = cudaMallocHost(reinterpret_cast(&pointer_), bytes); + } + + pthread_mutex_lock(&pinned_host_pool_mutex); + entry.pointer = status == cudaSuccess ? pointer_ + : (old_pointer_freed ? nullptr : old_pointer); + entry.capacity = status == cudaSuccess ? bytes + : (old_pointer_freed ? 0 : old_capacity); + if (status != cudaSuccess) { + entry.in_use = false; + pool_index_ = PINNED_HOST_POOL_SIZE; + } + pthread_mutex_unlock(&pinned_host_pool_mutex); + return status; + } + + uint64_t* get() { return pointer_; } + + private: + uint64_t* pointer_ = nullptr; + size_t pool_index_ = PINNED_HOST_POOL_SIZE; +}; + struct ResidentMerkleTree { uint8_t* rows = nullptr; uint8_t* digests = nullptr; @@ -459,6 +568,35 @@ struct ConstraintLookup { struct Ext2 { uint64_t c0; uint64_t c1; }; +struct PendingLookupLde { + ResidentLde* lde = nullptr; + Ext2* deltas = nullptr; + uint64_t* multiplicities = nullptr; + uint64_t* args = nullptr; + uint64_t* norms = nullptr; + uint64_t* norm_inverses = nullptr; + size_t* arg_offsets = nullptr; + Ext2* conjugates = nullptr; + size_t height = 0; + size_t slots = 0; + size_t extended_height = 0; + size_t num_lookups = 0; + size_t args_width = 0; + size_t group_size = 0; + size_t scratch_rows = 0; + + ~PendingLookupLde() { + if (norm_inverses != nullptr) cudaFree(norm_inverses); + if (norms != nullptr) cudaFree(norms); + if (conjugates != nullptr) cudaFree(conjugates); + if (arg_offsets != nullptr) cudaFree(arg_offsets); + if (args != nullptr) cudaFree(args); + if (multiplicities != nullptr) cudaFree(multiplicities); + if (deltas != nullptr) cudaFree(deltas); + if (lde != nullptr) destroy_resident_lde(lde); + } +}; + // The resident PCS canonicalizes every committed LDE, and interpolation // tables are serialized with `as_canonical_u64`. Restrict the faster // canonical-input arithmetic to this boundary instead of weakening the @@ -1045,21 +1183,79 @@ __global__ void evaluate_constraint_graph( } } +struct QuotientMatrix { + const uint64_t* values; + size_t width; + bool staged; + bool staged_next; +}; + +struct QuotientHostSource { + const uint64_t* values; + size_t width; + size_t offset; + bool needs_next; +}; + +struct QuotientPackTask { + uint64_t* destination; + const QuotientHostSource* sources; + size_t source_count; + size_t storage_start; + size_t begin; + size_t end; + size_t quotient_size; + size_t next_step; + unsigned int log_size; +}; + +static inline size_t reverse_low_bits_host(size_t value, unsigned int bits) { + uint64_t x=static_cast(value); + x=((x>>1)&0x5555555555555555ULL)|((x&0x5555555555555555ULL)<<1); + x=((x>>2)&0x3333333333333333ULL)|((x&0x3333333333333333ULL)<<2); + x=((x>>4)&0x0f0f0f0f0f0f0f0fULL)|((x&0x0f0f0f0f0f0f0f0fULL)<<4); + x=((x>>8)&0x00ff00ff00ff00ffULL)|((x&0x00ff00ff00ff00ffULL)<<8); + x=((x>>16)&0x0000ffff0000ffffULL)|((x&0x0000ffff0000ffffULL)<<16); + x=(x>>32)|(x<<32); + return static_cast(x>>(64U-bits)); +} + +static void* pack_quotient_rows(void* opaque) { + const auto& task=*static_cast(opaque); + for(size_t local=task.begin;local(blockIdx.x)*slot_count*tile : shared_values; const unsigned int log_size = static_cast(__ffsll(quotient_size) - 1); - for (size_t row = static_cast(blockIdx.x) * tile + lane; - row < quotient_size; row += static_cast(gridDim.x) * tile) { + for (size_t work_row = static_cast(blockIdx.x) * tile + lane; + work_row < work_count; work_row += static_cast(gridDim.x) * tile) { + const size_t work_index = work_start + work_row; + const size_t row = storage_order ? reverse_low_bits(work_index, log_size) : work_index; const size_t sr = reverse_low_bits(row, log_size); const size_t nr = reverse_low_bits((row + next_step) & (quotient_size - 1), log_size); for (size_t i = 0; i < node_count; ++i) { @@ -1069,8 +1265,11 @@ __global__ void evaluate_quotient( uint64_t v = 0; switch (n.op) { case 0: v = n.value; break; - case 1: { const ResidentLde* m = n.aux < 2 ? preprocessed : n.aux < 4 ? main : stage2; - const size_t r = (n.aux & 1U) ? nr : sr; v = m->values[r * m->width + n.a]; break; } + case 1: { const QuotientMatrix m = n.aux < 2 ? preprocessed : n.aux < 4 ? main : stage2; + const bool next = (n.aux & 1U) != 0; + const size_t r = m.staged ? work_row * (m.staged_next ? 2 : 1) + static_cast(next) + : (next ? nr : sr); + v = m.values[r * m.width + n.a]; break; } case 2: v = publics[n.a]; break; case 3: v = selectors[row]; break; case 4: v = selectors[quotient_size + row]; break; @@ -1097,9 +1296,11 @@ __global__ void evaluate_quotient( for (size_t g = 0; g < groups; ++g) { Ext2 constraints[8]; size_t count = 0; if (lookup_count == 0) { + const size_t s2_next = stage2.staged ? 2 * work_row + 1 : nr; + const size_t s2_cur = stage2.staged ? 2 * work_row : sr; constraints[count++] = ext2_add(ext2_sub( - {stage2->values[nr * stage2->width], stage2->values[nr * stage2->width + 1]}, - {stage2->values[sr * stage2->width], stage2->values[sr * stage2->width + 1]}), injection); + {stage2.values[s2_next * stage2.width], stage2.values[s2_next * stage2.width + 1]}, + {stage2.values[s2_cur * stage2.width], stage2.values[s2_cur * stage2.width + 1]}), injection); } else { const size_t begin = g * group_size; const size_t end = begin + group_size < lookup_count ? begin + group_size : lookup_count; @@ -1107,9 +1308,11 @@ __global__ void evaluate_quotient( for (size_t j=begin;j0;--k) { f=ext2_mul(f,gamma,ext_w); f.c0=goldilocks_add(f.c0,values[lookup_args[l.arg_start+k-1]*tile+lane]); } messages[j-begin]=ext2_add(f,beta); product=ext2_mul(product,messages[j-begin],ext_w); } - const Ext2 source{stage2->values[sr*stage2->width+2*g],stage2->values[sr*stage2->width+2*g+1]}; - Ext2 target = g+1values[sr*stage2->width+2*g+2],stage2->values[sr*stage2->width+2*g+3]} - : ext2_add({stage2->values[nr*stage2->width],stage2->values[nr*stage2->width+1]},injection); + const size_t s2_cur = stage2.staged ? 2 * work_row : sr; + const size_t s2_next = stage2.staged ? 2 * work_row + 1 : nr; + const Ext2 source{stage2.values[s2_cur*stage2.width+2*g],stage2.values[s2_cur*stage2.width+2*g+1]}; + Ext2 target = g+1(handles[index]); @@ -1724,6 +1927,141 @@ cudaError_t hash_resident_lde_group(uint8_t* digests, return status; } +cudaError_t hash_partitioned_lde_group( + uint8_t* digests, const void* const* handles, + const uint64_t* const* host_values, const size_t* widths, + const size_t* heights, size_t matrix_count, size_t height) { + bool has_host = false; + bool has_resident = false; + size_t total_width = 0; + for (size_t index = 0; index < matrix_count; ++index) { + if (heights[index] != height) { + continue; + } + if (widths[index] > SIZE_MAX - total_width) { + return cudaErrorInvalidValue; + } + total_width += widths[index]; + has_host = has_host || host_values[index] != nullptr; + has_resident = has_resident || handles[index] != nullptr; + } + if (!has_host) { + return hash_resident_lde_group(digests, handles, matrix_count, height); + } + if (!has_resident || total_width == 0 || + total_width > (32 * 1024) / sizeof(uint64_t) || + !product_fits(height, total_width)) { + return cudaErrorInvalidValue; + } + + const uint64_t** host_columns = + new (std::nothrow) const uint64_t*[total_width]; + size_t* host_strides = new (std::nothrow) size_t[total_width]; + void** registered = new (std::nothrow) void*[matrix_count](); + if (host_columns == nullptr || host_strides == nullptr || + registered == nullptr) { + delete[] host_columns; + delete[] host_strides; + delete[] registered; + return cudaErrorMemoryAllocation; + } + + cudaError_t status = cudaSuccess; + size_t column_offset = 0; + for (size_t index = 0; status == cudaSuccess && index < matrix_count; + ++index) { + if (heights[index] != height) { + continue; + } + const uint64_t* values = nullptr; + if (const ResidentLde* lde = + static_cast(handles[index])) { + values = lde->values; + } else { + const size_t elements = heights[index] * widths[index]; + void* host = const_cast(host_values[index]); + status = cudaHostRegister(host, elements * sizeof(uint64_t), + cudaHostRegisterMapped | + cudaHostRegisterPortable); + if (status != cudaSuccess) { + break; + } + registered[index] = host; + void* mapped = nullptr; + status = cudaHostGetDevicePointer(&mapped, host, 0); + values = static_cast(mapped); + } + for (size_t column = 0; status == cudaSuccess && column < widths[index]; + ++column) { + host_columns[column_offset] = values + column; + host_strides[column_offset++] = widths[index]; + } + } + + DeviceBuffer device_columns; + DeviceBuffer device_strides; + DeviceBuffer combined_rows; + constexpr size_t ROW_STAGING_BYTES = size_t(32) << 20; + const size_t row_bytes = total_width * sizeof(uint64_t); + const size_t rows_per_chunk = + (ROW_STAGING_BYTES / row_bytes) > 0 ? (ROW_STAGING_BYTES / row_bytes) : 1; + if (status == cudaSuccess) { + status = device_columns.allocate(total_width); + } + if (status == cudaSuccess) { + status = device_strides.allocate(total_width); + } + if (status == cudaSuccess) { + status = combined_rows.allocate( + (height < rows_per_chunk ? height : rows_per_chunk) * total_width); + } + if (status == cudaSuccess) { + status = cudaMemcpy(device_columns.get(), host_columns, + total_width * sizeof(uint64_t*), + cudaMemcpyHostToDevice); + } + if (status == cudaSuccess) { + status = cudaMemcpy(device_strides.get(), host_strides, + total_width * sizeof(size_t), cudaMemcpyHostToDevice); + } + for (size_t row_start = 0; status == cudaSuccess && row_start < height; + row_start += rows_per_chunk) { + const size_t rows = + (height - row_start < rows_per_chunk) ? height - row_start + : rows_per_chunk; + const size_t count = rows * total_width; + gather_resident_lde_group<<>>( + combined_rows.get(), + reinterpret_cast(device_columns.get()), + reinterpret_cast(device_strides.get()), row_start, + rows, total_width); + status = cudaGetLastError(); + if (status == cudaSuccess) { + status = launch_blake3_rows( + digests + row_start * 32, + reinterpret_cast(combined_rows.get()), row_bytes, + rows); + } + } + if (status == cudaSuccess) { + status = cudaStreamSynchronize(cudaStreamPerThread); + } + for (size_t index = 0; index < matrix_count; ++index) { + if (registered[index] != nullptr) { + const cudaError_t unregister_status = + cudaHostUnregister(registered[index]); + if (status == cudaSuccess) { + status = unregister_status; + } + } + } + delete[] host_columns; + delete[] host_strides; + delete[] registered; + return status; +} + cudaError_t copy_to_device(DeviceBuffer& destination, const uint64_t* source, size_t elements) { cudaError_t status = destination.allocate(elements); @@ -2096,12 +2434,13 @@ extern "C" int multi_stark_cuda_lde_copy_rows( extern "C" int multi_stark_cuda_mixed_lde_open_row( int device_id, uint64_t* output, const void* const* handles, - size_t handle_count, size_t index) { + size_t handle_count, size_t max_height, size_t index) { if (output == nullptr || handles == nullptr || handle_count == 0) return static_cast(cudaErrorInvalidValue); - cudaError_t status=cudaSetDevice(device_id); size_t max_height=0,total=0; + cudaError_t status=cudaSetDevice(device_id); size_t handles_max_height=0,total=0; for(size_t i=0;i(handles[i]); if(l==nullptr || l->height==0 || (l->height&(l->height-1))!=0 || l->width>SIZE_MAX-total) return static_cast(cudaErrorInvalidValue); - if(l->height>max_height)max_height=l->height; total+=l->width;} + if(l->height>handles_max_height)handles_max_height=l->height; total+=l->width;} + if (!is_power_of_two(max_height) || max_height < handles_max_height || index >= max_height) return static_cast(cudaErrorInvalidValue); const uint64_t** hv=new(std::nothrow) const uint64_t*[handle_count]; size_t* hw=new(std::nothrow) size_t[3*handle_count]; if(hv==nullptr||hw==nullptr){delete[] hv;delete[] hw;return static_cast(cudaErrorMemoryAllocation);} size_t* hr=hw+handle_count;size_t* ho=hr+handle_count; size_t off=0;for(size_t i=0;i(handles[i]);hv[i]=l->values;hw[i]=l->width; @@ -2119,17 +2458,19 @@ extern "C" int multi_stark_cuda_mixed_lde_open_row( extern "C" int multi_stark_cuda_mixed_lde_open_rows( int device_id, uint64_t* output, const void* const* handles, - size_t handle_count, const uint64_t* indices, size_t query_count) { + size_t handle_count, size_t max_height, const uint64_t* indices, + size_t query_count) { if (output == nullptr || handles == nullptr || handle_count == 0 || indices == nullptr || query_count == 0) return static_cast(cudaErrorInvalidValue); - cudaError_t status=cudaSetDevice(device_id);size_t max_height=0,total=0; + cudaError_t status=cudaSetDevice(device_id);size_t handles_max_height=0,total=0; const uint64_t** hv=new(std::nothrow) const uint64_t*[handle_count]; size_t* hm=new(std::nothrow) size_t[3*handle_count]; if(hv==nullptr||hm==nullptr){delete[] hv;delete[] hm;return static_cast(cudaErrorMemoryAllocation);} size_t* hw=hm,*hh=hm+handle_count,*ho=hh+handle_count;size_t off=0; for(size_t i=0;i(handles[i]); if(l==nullptr||l->height==0||(l->height&(l->height-1))!=0||l->width>SIZE_MAX-total){delete[] hm;delete[] hv;return static_cast(cudaErrorInvalidValue);} - hv[i]=l->values;hw[i]=l->width;hh[i]=l->height;ho[i]=off;off+=l->width;total+=l->width;if(l->height>max_height)max_height=l->height;} + hv[i]=l->values;hw[i]=l->width;hh[i]=l->height;ho[i]=off;off+=l->width;total+=l->width;if(l->height>handles_max_height)handles_max_height=l->height;} + if(!is_power_of_two(max_height)||max_height(cudaErrorInvalidValue);} for(size_t i=0;i=max_height){delete[] hm;delete[] hv;return static_cast(cudaErrorInvalidValue);} const uint64_t** dv=nullptr;size_t* dm=nullptr;uint64_t* di=nullptr;uint64_t* dout=nullptr; @@ -2263,10 +2604,15 @@ extern "C" int multi_stark_cuda_quotient_values( if(status==cudaSuccess&&dynamic_shared>48*1024) status=configure_quotient_shared_memory(device_id,dynamic_shared); if(status==cudaSuccess) { + const auto* prep = static_cast(preprocessed_handle); + const auto* main = static_cast(main_handle); + const auto* stage2 = static_cast(stage2_handle); evaluate_quotient<<(blocks),static_cast(tile),dynamic_shared>>>( dout,dn,node_count,slot_count,dr,root_count,dl,lookup_count,da,group_size, - static_cast(preprocessed_handle),static_cast(main_handle), - static_cast(stage2_handle),dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch); + {prep ? prep->values : nullptr, prep ? prep->width : 0, false, false}, + {main->values, main->width, false, false}, + {stage2->values, stage2->width, false, false}, + dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch,0,quotient_size,false); status=cudaGetLastError(); } if(status==cudaSuccess) status=cudaMemcpy(output,dout,2*quotient_size*sizeof(uint64_t),cudaMemcpyDeviceToHost); @@ -2348,10 +2694,15 @@ extern "C" int multi_stark_cuda_quotient_lde( if(status==cudaSuccess&&dynamic_shared>48*1024) status=configure_quotient_shared_memory(device_id,dynamic_shared); if(status==cudaSuccess) { + const auto* prep = static_cast(preprocessed_handle); + const auto* main = static_cast(main_handle); + const auto* stage2 = static_cast(stage2_handle); evaluate_quotient<<(blocks),static_cast(tile),dynamic_shared>>>( quotient,dn,node_count,slot_count,dr,root_count,dl,lookup_count,da,group_size, - static_cast(preprocessed_handle),static_cast(main_handle), - static_cast(stage2_handle),dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch); + {prep ? prep->values : nullptr, prep ? prep->width : 0, false, false}, + {main->values, main->width, false, false}, + {stage2->values, stage2->width, false, false}, + dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch,0,quotient_size,false); status=cudaGetLastError(); } if(status==cudaSuccess)status=launch_dif(quotient,quotient_size,2,device_quotient_twiddles); @@ -2382,6 +2733,229 @@ extern "C" int multi_stark_cuda_quotient_lde( return static_cast(status); } +// Evaluate a quotient when one or more committed LDEs have been spilled to +// host memory. The quotient rows are independent until the DFT, so bounded +// pairs of current/next trace rows are staged through two buffers while the +// previous buffer executes. Iterating in bit-reversed storage order keeps the +// current-row read contiguous; host workers gather the corresponding next +// rows in parallel. Resident inputs remain in place. +extern "C" int multi_stark_cuda_quotient_lde_mixed( + int device_id, void** output_handle, const void* nodes, size_t node_count, + size_t slot_count, const uint32_t* roots, size_t root_count, + const void* lookups, size_t lookup_count, const uint32_t* lookup_args, + size_t lookup_arg_count, size_t group_size, + const void* preprocessed_handle, const uint64_t* preprocessed_host, + size_t preprocessed_height, size_t preprocessed_width, + const void* main_handle, const uint64_t* main_host, + size_t main_height, size_t main_width, + const void* stage2_handle, const uint64_t* stage2_host, + size_t stage2_height, size_t stage2_width, + const uint64_t* publics, size_t public_count, + uint64_t coset_shift, uint64_t coset_generator, uint64_t trace_last, + uint64_t vanishing_start, uint64_t vanishing_step, const uint64_t* alpha, + size_t constraint_count, const uint64_t* delta, uint64_t ext_w, + size_t quotient_size, size_t next_step, size_t quotient_degree, + size_t log_blowup, const uint64_t* quotient_twiddles, + const uint64_t* lde_twiddles, const uint64_t* slice_weights) { + const auto one_source = [](const void* handle, const uint64_t* host) { + return (handle != nullptr) != (host != nullptr); + }; + if (output_handle == nullptr || nodes == nullptr || roots == nullptr || + !one_source(main_handle, main_host) || + !one_source(stage2_handle, stage2_host) || publics == nullptr || + alpha == nullptr || delta == nullptr || quotient_twiddles == nullptr || + lde_twiddles == nullptr || slice_weights == nullptr || node_count == 0 || + slot_count == 0 || group_size == 0 || !is_power_of_two(quotient_size) || + !is_power_of_two(quotient_degree) || quotient_degree > quotient_size || + quotient_size % quotient_degree != 0 || !is_power_of_two(next_step) || + next_step > quotient_size || log_blowup >= sizeof(size_t) * 8 || + main_width == 0 || stage2_width == 0 || main_height < quotient_size || + stage2_height < quotient_size || + (preprocessed_handle != nullptr && preprocessed_host != nullptr) || + (preprocessed_host != nullptr && + (preprocessed_width == 0 || preprocessed_height < quotient_size)) || + (lookup_count != 0 && (lookups == nullptr || lookup_args == nullptr))) { + return static_cast(cudaErrorInvalidValue); + } + *output_handle = nullptr; + const size_t trace_height = quotient_size / quotient_degree; + if (trace_height > (SIZE_MAX >> log_blowup)) { + return static_cast(cudaErrorInvalidValue); + } + const size_t lde_height = trace_height << log_blowup; + const size_t width = 2 * quotient_degree; + if (!product_fits(lde_height, width)) { + return static_cast(cudaErrorInvalidValue); + } + + cudaError_t status = cudaSetDevice(device_id); + auto align8=[](size_t n){return (n+7)&~size_t(7);}; size_t bytes=0; + auto reserve=[&](size_t n){size_t at=bytes;bytes+=align8(n);return at;}; + const size_t on=reserve(node_count*sizeof(ConstraintNode)), oroot=reserve(root_count*sizeof(uint32_t)); + const size_t ol=reserve(lookup_count*sizeof(ConstraintLookup)), oa=reserve(lookup_arg_count*sizeof(uint32_t)); + const size_t op=reserve(public_count*sizeof(uint64_t)), os=reserve(4*quotient_size*sizeof(uint64_t)); + const size_t oalpha=reserve(2*constraint_count*sizeof(uint64_t)), od=reserve(2*sizeof(uint64_t)); + const size_t oo=reserve(2*quotient_size*sizeof(uint64_t)); uint8_t* allocation=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&allocation),bytes); + auto cp=[&](size_t at,const void* src,size_t n){if(status==cudaSuccess&&n!=0)status=cudaMemcpy(allocation+at,src,n,cudaMemcpyHostToDevice);}; + cp(on,nodes,node_count*sizeof(ConstraintNode));cp(oroot,roots,root_count*sizeof(uint32_t)); + cp(ol,lookups,lookup_count*sizeof(ConstraintLookup));cp(oa,lookup_args,lookup_arg_count*sizeof(uint32_t)); + cp(op,publics,public_count*sizeof(uint64_t)); + cp(oalpha,alpha,2*constraint_count*sizeof(uint64_t));cp(od,delta,2*sizeof(uint64_t)); + auto* dn=reinterpret_cast(allocation+on);auto* dr=reinterpret_cast(allocation+oroot); + auto* dl=reinterpret_cast(allocation+ol);auto* da=reinterpret_cast(allocation+oa); + auto* dp=reinterpret_cast(allocation+op);auto* ds=reinterpret_cast(allocation+os); + auto* dal=reinterpret_cast(allocation+oalpha);auto* dd=reinterpret_cast(allocation+od); + auto* quotient=reinterpret_cast(allocation+oo); + if(status==cudaSuccess)status=generate_coset_selectors(ds,quotient_size,next_step, + coset_shift,coset_generator,trace_last,vanishing_start,vanishing_step); + + const uint64_t *device_quotient_twiddles=nullptr,*device_lde_twiddles=nullptr,*device_weights=nullptr; + if(status==cudaSuccess)status=cached_device_constants(device_id,quotient_twiddles,quotient_size/2,2,0,0,&device_quotient_twiddles); + if(status==cudaSuccess)status=cached_device_constants(device_id,lde_twiddles,lde_height/2,2,0,0,&device_lde_twiddles); + if(status==cudaSuccess)status=cached_device_constants(device_id,slice_weights,quotient_degree,4,slice_weights[0],quotient_degree>1?slice_weights[1]:0,&device_weights); + + const auto* prep_resident=static_cast(preprocessed_handle); + const auto* main_resident=static_cast(main_handle); + const auto* stage2_resident=static_cast(stage2_handle); + if (status==cudaSuccess && + ((prep_resident && prep_resident->height < quotient_size) || + (main_resident && main_resident->height < quotient_size) || + (stage2_resident && stage2_resident->height < quotient_size))) { + status=cudaErrorInvalidValue; + } + + QuotientHostSource host_sources[3]; + size_t host_source_count=0; + size_t staged_row_width=0; + bool needs_next[3]={false,false,true}; + const auto* host_nodes=static_cast(nodes); + for(size_t i=0;iquotient_size)chunk_rows=quotient_size; + if(chunk_rows>=128)chunk_rows=(chunk_rows/128)*128; + } + const size_t staging_elements=chunk_rows*staged_row_width; + const size_t staging_bytes=staging_elements*sizeof(uint64_t); + size_t staging_offset=0; + for(size_t i=0;i(&device_staging[i]),staging_bytes); + if(status==cudaSuccess)status=cudaStreamCreateWithFlags(&streams[i],cudaStreamNonBlocking); + } + + size_t budget=0;if(status==cudaSuccess)status=quotient_shared_memory_budget(device_id,&budget); + size_t tile=budget/(slot_count*sizeof(uint64_t));bool global=tile<28; + if(global)tile=128;else if(tile>32)tile=32; + const size_t block_cap=global?256:1024; + uint64_t* scratch[2]={nullptr,nullptr}; + if(status==cudaSuccess&&global){ + const size_t scratch_bytes=block_cap*slot_count*tile*sizeof(uint64_t); + status=persistent_malloc(reinterpret_cast(&scratch[0]),scratch_bytes); + if(status==cudaSuccess)status=persistent_malloc(reinterpret_cast(&scratch[1]),scratch_bytes); + } + const size_t dynamic_shared=global?0:slot_count*tile*sizeof(uint64_t); + if(status==cudaSuccess&&dynamic_shared>48*1024) + status=configure_quotient_shared_memory(device_id,dynamic_shared); + + const unsigned int log_size=static_cast(__builtin_ctzll(quotient_size)); + const auto pack=[&](uint64_t* destination,size_t storage_start,size_t count){ + const long online=sysconf(_SC_NPROCESSORS_ONLN); + size_t workers=online>0?static_cast(online):1; + if(workers>64)workers=64; + const size_t useful=(count+4095)/4096; + if(workers>useful)workers=useful; + if(workers==0)workers=1; + pthread_t threads[64]; + QuotientPackTask tasks[64]; + bool launched[64]={}; + for(size_t worker=0;workervalues:nullptr,resident?resident->width:0,false,false}; + size_t offset=0;bool source_needs_next=false; + for(size_t i=0;iblock_cap)blocks=block_cap; + evaluate_quotient<<(blocks),static_cast(tile),dynamic_shared,streams[buffer]>>>( + quotient,dn,node_count,slot_count,dr,root_count,dl,lookup_count,da,group_size, + prep,main,stage2,dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch[buffer], + storage_start,count,true); + status=cudaGetLastError();stream_busy[buffer]=status==cudaSuccess; + } + for(size_t i=0;i<2;++i)if(status==cudaSuccess&&stream_busy[i])status=cudaStreamSynchronize(streams[i]); + + if(status==cudaSuccess)status=launch_dif(quotient,quotient_size,2,device_quotient_twiddles); + ResidentLde* lde=nullptr; + if(status==cudaSuccess)status=create_resident_lde(&lde); + if(status==cudaSuccess){lde->height=lde_height;lde->width=width;status=cudaMalloc(reinterpret_cast(&lde->values),lde_height*width*sizeof(uint64_t));} + if(status==cudaSuccess)status=cudaMemset(lde->values,0,lde_height*width*sizeof(uint64_t)); + if(status==cudaSuccess){ + gather_shifted_quotient_slices<<>>( + lde->values,quotient,device_weights,quotient_size,trace_height,quotient_degree,2); + status=cudaGetLastError(); + } + if(status==cudaSuccess)status=launch_dif(lde->values,lde_height,width,device_lde_twiddles); + if(status==cudaSuccess)status=cudaStreamSynchronize(0); + if(status==cudaSuccess)*output_handle=lde;else if(lde)destroy_resident_lde(lde); + for(size_t i=0;i<2;++i){ + if(streams[i])cudaStreamDestroy(streams[i]); + if(device_staging[i])persistent_free(device_staging[i]); + if(scratch[i])persistent_free(scratch[i]); + } + cudaFree(allocation); + return static_cast(status); +} + extern "C" int multi_stark_cuda_lde_interpolate(int device_id,uint64_t* output,const void* handle, size_t height,const uint64_t* inv_denoms,const uint64_t* coset,const uint64_t* scale,uint64_t ext_w){ if(!output||!handle||!inv_denoms||!coset||!scale)return static_cast(cudaErrorInvalidValue); @@ -2530,6 +3104,25 @@ extern "C" int multi_stark_cuda_reduced_create(int device_id,void** handle,size_ if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&r->values),height*sizeof(Ext2));if(status==cudaSuccess)status=cudaMemset(r->values,0,height*sizeof(Ext2)); if(status!=cudaSuccess){delete r;return static_cast(status);}*handle=r;return static_cast(cudaSuccess); } +extern "C" int multi_stark_cuda_reduced_add_host(int device_id,void* handle, + const uint64_t* values,size_t height){ + if(!handle||!values)return static_cast(cudaErrorInvalidValue); + auto* r=static_cast(handle); + if(height!=r->height)return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);const size_t bytes=height*sizeof(Ext2); + if(status==cudaSuccess&&r->scratch_bytesscratch)status=cudaFree(r->scratch);r->scratch=nullptr;r->scratch_bytes=0; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&r->scratch),bytes); + if(status==cudaSuccess)r->scratch_bytes=bytes; + } + if(status==cudaSuccess)status=cudaMemcpy(r->scratch,values,bytes,cudaMemcpyHostToDevice); + if(status==cudaSuccess){ + add_scaled_ext2<<>>( + r->values,reinterpret_cast(r->scratch),height,{1,0},7); + status=cudaGetLastError(); + } + return static_cast(status); +} extern "C" int multi_stark_cuda_lookup_trace(int device_id,uint64_t* output,uint64_t* total, const uint64_t* multiplicities,const uint64_t* args,const size_t* arg_offsets, size_t height,size_t num_lookups,size_t args_width,size_t group_size, @@ -2653,11 +3246,19 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui *output_handle=nullptr;const size_t slots=(num_lookups+group_size-1)/group_size; const size_t width=2*slots,count=height*slots,extended_height=height<(cudaErrorInvalidValue); - cudaError_t status=cudaSetDevice(device_id);const size_t message_count=height*num_lookups; + const bool profile=getenv("MULTI_STARK_CUDA_MEMORY_LOG")!=nullptr&& + extended_height*width*sizeof(uint64_t)>=(size_t(1)<<30); + const auto now=[](){timespec ts{};clock_gettime(CLOCK_MONOTONIC,&ts); + return static_cast(ts.tv_sec)+static_cast(ts.tv_nsec)*1e-9;}; + const double started=now();double allocated=started,rows_done=started,scan_done=started; + constexpr size_t LOOKUP_ROWS_PER_CHUNK=size_t(1)<<16; + const size_t chunk_rows=height(&dm),height*num_lookups*sizeof(uint64_t)); - if(status==cudaSuccess&&args_width)status=cudaMalloc(reinterpret_cast(&da),height*args_width*sizeof(uint64_t)); + cudaError_t status=cudaSetDevice(device_id); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&dm),chunk_rows*num_lookups*sizeof(uint64_t)); + if(status==cudaSuccess&&args_width)status=cudaMalloc(reinterpret_cast(&da),chunk_rows*args_width*sizeof(uint64_t)); if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&offsets),(num_lookups+1)*sizeof(size_t)); if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&deltas),count*sizeof(Ext2)); if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&conjugates),message_count*sizeof(Ext2)); @@ -2666,15 +3267,23 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui if(status==cudaSuccess)status=create_resident_lde(&lde); if(status==cudaSuccess){lde->height=extended_height;lde->width=width;status=cudaMalloc(reinterpret_cast(&lde->values),extended_height*width*sizeof(uint64_t));} if(status==cudaSuccess)status=cudaMemset(lde->values,0,extended_height*width*sizeof(uint64_t)); - if(status==cudaSuccess)status=cudaMemcpy(dm,multiplicities,height*num_lookups*sizeof(uint64_t),cudaMemcpyHostToDevice); - if(status==cudaSuccess&&args_width)status=cudaMemcpy(da,args,height*args_width*sizeof(uint64_t),cudaMemcpyHostToDevice); if(status==cudaSuccess)status=cudaMemcpy(offsets,arg_offsets,(num_lookups+1)*sizeof(size_t),cudaMemcpyHostToDevice); - if(status==cudaSuccess){lookup_messages<<>>(conjugates,norms,da,offsets,height,num_lookups,args_width,{beta[0],beta[1]},{gamma[0],gamma[1]},ext_w);status=cudaGetLastError();} - if(status==cudaSuccess)status=batch_inverse_norms(norm_inverses,norms,message_count); - if(status==cudaSuccess){lookup_group_deltas_batched<<>>(deltas,dm,conjugates,norm_inverses,height,num_lookups,group_size,ext_w);status=cudaGetLastError();} + allocated=now(); + for(size_t row_start=0;status==cudaSuccess&&row_start>>(conjugates,norms,da,offsets,rows,num_lookups,args_width,{beta[0],beta[1]},{gamma[0],gamma[1]},ext_w);status=cudaGetLastError();} + if(status==cudaSuccess)status=batch_inverse_norms(norm_inverses,norms,messages); + if(status==cudaSuccess){lookup_group_deltas_batched<<>>(deltas+row_start*slots,dm,conjugates,norm_inverses,rows,num_lookups,group_size,ext_w);status=cudaGetLastError();} + } + if(profile&&status==cudaSuccess)status=cudaStreamSynchronize(cudaStreamPerThread); + rows_done=now(); if(status==cudaSuccess)status=exclusive_scan_ext2(reinterpret_cast(lde->values),deltas,count); if(status==cudaSuccess)status=cudaMemcpy(total,lde->values+2*(count-1),sizeof(Ext2),cudaMemcpyDeviceToHost); if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); + scan_done=now(); const uint64_t *dit=nullptr,*dshift=nullptr,*dft=nullptr; if(status==cudaSuccess)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); @@ -2684,10 +3293,224 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui if(status==cudaSuccess)status=launch_dif(lde->values,extended_height,width,dft); if(status==cudaSuccess){canonicalize_goldilocks<<>>(lde->values,extended_height*width);status=cudaGetLastError();} if(status==cudaSuccess)status=cudaStreamSynchronize(0); + if(profile){const double finished=now();fprintf(stderr, + "[multi-stark/cuda] lookup phases: height=%zu lookups=%zu slots=%zu args_width=%zu allocate=%.3fs rows=%.3fs scan=%.3fs dft=%.3fs\n", + height,num_lookups,slots,args_width,allocated-started,rows_done-allocated, + scan_done-rows_done,finished-scan_done);} if(status==cudaSuccess)*output_handle=lde;else destroy_resident_lde(lde); cudaFree(norm_inverses);cudaFree(norms);cudaFree(conjugates);cudaFree(deltas); cudaFree(offsets);cudaFree(da);cudaFree(dm);return static_cast(status); } + +extern "C" int multi_stark_cuda_lookup_lde_begin_partitioned( + int device_id, void** pending_handle, const size_t* arg_offsets, + size_t height, size_t num_lookups, size_t args_width, + size_t group_size, size_t added_bits) { + if (!pending_handle || !arg_offsets || !height || !num_lookups || + !group_size || group_size > 8 || !is_power_of_two(height) || + num_lookups > SIZE_MAX - (group_size - 1) || + added_bits >= sizeof(size_t) * 8 || height > (SIZE_MAX >> added_bits)) { + return static_cast(cudaErrorInvalidValue); + } + *pending_handle = nullptr; + const size_t slots = (num_lookups + group_size - 1) / group_size; + const size_t width = 2 * slots; + const size_t extended_height = height << added_bits; + constexpr size_t LOOKUP_ROWS_PER_CHUNK = size_t(1) << 16; + const size_t scratch_rows = std::min(height, LOOKUP_ROWS_PER_CHUNK); + if (!product_fits(extended_height, width) || + !product_fits(height, slots) || + !product_fits(scratch_rows, num_lookups) || + !product_fits(scratch_rows, args_width)) { + return static_cast(cudaErrorInvalidValue); + } + + auto* pending = new (std::nothrow) PendingLookupLde; + if (!pending) return static_cast(cudaErrorMemoryAllocation); + pending->height = height; + pending->slots = slots; + pending->extended_height = extended_height; + pending->num_lookups = num_lookups; + pending->args_width = args_width; + pending->group_size = group_size; + pending->scratch_rows = scratch_rows; + + const size_t message_count = scratch_rows * num_lookups; + cudaError_t status = cudaSetDevice(device_id); + if (status == cudaSuccess) + status = cudaMalloc(reinterpret_cast(&pending->multiplicities), + message_count * sizeof(uint64_t)); + if (status == cudaSuccess && args_width) + status = cudaMalloc(reinterpret_cast(&pending->args), + scratch_rows * args_width * sizeof(uint64_t)); + if (status == cudaSuccess) + status = cudaMalloc(reinterpret_cast(&pending->arg_offsets), + (num_lookups + 1) * sizeof(size_t)); + if (status == cudaSuccess) + status = cudaMalloc(reinterpret_cast(&pending->deltas), + height * slots * sizeof(Ext2)); + if (status == cudaSuccess) + status = cudaMalloc(reinterpret_cast(&pending->conjugates), + message_count * sizeof(Ext2)); + if (status == cudaSuccess) + status = cudaMalloc(reinterpret_cast(&pending->norms), + message_count * sizeof(uint64_t)); + if (status == cudaSuccess) + status = cudaMalloc(reinterpret_cast(&pending->norm_inverses), + message_count * sizeof(uint64_t)); + if (status == cudaSuccess) status = create_resident_lde(&pending->lde); + if (status == cudaSuccess) { + pending->lde->height = extended_height; + pending->lde->width = width; + status = cudaMalloc(reinterpret_cast(&pending->lde->values), + extended_height * width * sizeof(uint64_t)); + } + if (status == cudaSuccess) + status = cudaMemset(pending->lde->values, 0, + extended_height * width * sizeof(uint64_t)); + if (status == cudaSuccess) + status = cudaMemcpy(pending->arg_offsets, arg_offsets, + (num_lookups + 1) * sizeof(size_t), + cudaMemcpyHostToDevice); + if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); + if (status == cudaSuccess) { + *pending_handle = pending; + } else { + delete pending; + } + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lookup_lde_gpu_rows_partitioned( + int device_id, void* pending_handle, const uint64_t* multiplicities, + const uint64_t* args, size_t row_start, size_t rows, + const uint64_t* beta, const uint64_t* gamma, uint64_t ext_w) { + auto* pending = static_cast(pending_handle); + if (!pending || !multiplicities || !rows || !beta || !gamma || + (pending->args_width && !args) || rows > pending->scratch_rows || + row_start > pending->height || rows > pending->height - row_start) { + return static_cast(cudaErrorInvalidValue); + } + const size_t messages = rows * pending->num_lookups; + cudaError_t status = cudaSetDevice(device_id); + if (status == cudaSuccess) + status = cudaMemcpy(pending->multiplicities, + multiplicities + row_start * pending->num_lookups, + messages * sizeof(uint64_t), cudaMemcpyHostToDevice); + if (status == cudaSuccess && pending->args_width) + status = cudaMemcpy(pending->args, + args + row_start * pending->args_width, + rows * pending->args_width * sizeof(uint64_t), + cudaMemcpyHostToDevice); + if (status == cudaSuccess) { + lookup_messages<<>>( + pending->conjugates, pending->norms, pending->args, + pending->arg_offsets, rows, pending->num_lookups, + pending->args_width, {beta[0], beta[1]}, {gamma[0], gamma[1]}, + ext_w); + status = cudaGetLastError(); + } + if (status == cudaSuccess) + status = batch_inverse_norms(pending->norm_inverses, pending->norms, + messages); + if (status == cudaSuccess) { + lookup_group_deltas_batched<<slots), THREADS>>>( + pending->deltas + row_start * pending->slots, + pending->multiplicities, pending->conjugates, + pending->norm_inverses, rows, pending->num_lookups, + pending->group_size, ext_w); + status = cudaGetLastError(); + } + if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lookup_lde_cpu_rows_partitioned( + int device_id, void* pending_handle, const uint64_t* deltas, + size_t row_start, size_t rows) { + auto* pending = static_cast(pending_handle); + if (!pending || !deltas || !rows || row_start > pending->height || + rows > pending->height - row_start) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status == cudaSuccess) + status = cudaMemcpy(pending->deltas + row_start * pending->slots, + deltas, rows * pending->slots * sizeof(Ext2), + cudaMemcpyHostToDevice); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( + int device_id, void* pending_handle, void** output_handle, uint64_t* total, + const uint64_t* inverse_twiddles, const uint64_t* shift_powers, + const uint64_t* forward_twiddles, uint64_t height_inverse) { + auto* pending = static_cast(pending_handle); + if (!pending || !output_handle || !total || !inverse_twiddles || + !shift_powers || !forward_twiddles) { + return static_cast(cudaErrorInvalidValue); + } + *output_handle = nullptr; + cudaError_t status = cudaSetDevice(device_id); + const size_t count = pending->height * pending->slots; + if (status == cudaSuccess) + status = exclusive_scan_ext2(reinterpret_cast(pending->lde->values), + pending->deltas, count); + if (status == cudaSuccess) + status = cudaMemcpy(total, pending->lde->values + 2 * (count - 1), + sizeof(Ext2), cudaMemcpyDeviceToHost); + if (status == cudaSuccess) + status = cudaMemcpy(total + 2, pending->deltas + count - 1, + sizeof(Ext2), cudaMemcpyDeviceToHost); + + const uint64_t* inverse = nullptr; + const uint64_t* shifts = nullptr; + const uint64_t* forward = nullptr; + if (status == cudaSuccess) + status = cached_device_constants(device_id, inverse_twiddles, + pending->height / 2, 1, 0, 0, &inverse); + if (status == cudaSuccess) + status = cached_device_constants(device_id, shift_powers, + pending->height, 3, shift_powers[0], + pending->height > 1 ? shift_powers[1] : 0, + &shifts); + if (status == cudaSuccess) + status = cached_device_constants(device_id, forward_twiddles, + pending->extended_height / 2, 2, 0, 0, + &forward); + const size_t width = 2 * pending->slots; + if (status == cudaSuccess) + status = launch_dif(pending->lde->values, pending->height, width, inverse); + if (status == cudaSuccess) { + bit_reverse_scale_and_shift<<height * width), THREADS>>>( + pending->lde->values, pending->height, width, + strict_log2(pending->height), height_inverse, shifts); + status = cudaGetLastError(); + } + if (status == cudaSuccess) + status = launch_dif(pending->lde->values, pending->extended_height, + width, forward); + if (status == cudaSuccess) { + canonicalize_goldilocks<<extended_height * width), THREADS>>>( + pending->lde->values, pending->extended_height * width); + status = cudaGetLastError(); + } + if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); + if (status == cudaSuccess) { + *output_handle = pending->lde; + pending->lde = nullptr; + } + delete pending; + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lookup_lde_cancel_partitioned( + int device_id, void* pending_handle) { + const cudaError_t status = cudaSetDevice(device_id); + if (status == cudaSuccess) delete static_cast(pending_handle); + return static_cast(status); +} + extern "C" int multi_stark_cuda_reduced_add(int device_id,void* reduced,const void* lde_handle,size_t height, const uint64_t* inv_denoms,const uint64_t* alpha_powers,const uint64_t* reduced_y,const uint64_t* alpha_offset,uint64_t ext_w){ if(!reduced||!lde_handle||!inv_denoms||!alpha_powers||!reduced_y||!alpha_offset)return static_cast(cudaErrorInvalidValue); @@ -2719,6 +3542,40 @@ extern "C" int multi_stark_cuda_lde_release_trace(int device_id,void* handle){ return static_cast(status); } +extern "C" int multi_stark_cuda_lde_release_values(int device_id, void* handle) { + if (!handle) return static_cast(cudaSuccess); + cudaError_t status = cudaSetDevice(device_id); + auto* lde = static_cast(handle); + if (status == cudaSuccess && lde->values) { + // This is an admission-control eviction, not a short-lived scratch + // release. Make the memory globally available before returning: the + // LDE may have been created on a Rayon worker's per-thread stream and + // its replacement upload may run on a different stream. An async free + // would make the Rust-side free-byte projection optimistic and can + // still produce cudaErrorMemoryAllocation on that upload. + status = persistent_free(lde->values); + lde->values = nullptr; + } + if (status == cudaSuccess && lde->interpolation_scratch) { + status = persistent_free(lde->interpolation_scratch); + lde->interpolation_scratch = nullptr; + lde->interpolation_scratch_bytes = 0; + } + if (status == cudaSuccess && lde->trace_values) { + status = persistent_free(lde->trace_values); + lde->trace_values = nullptr; + } + if (status == cudaSuccess && lde->host_trace_registered) { + status = cudaHostUnregister(const_cast(lde->host_trace_values)); + lde->host_trace_registered = false; + } + if (status == cudaSuccess) { + lde->host_trace_values = nullptr; + lde->trace_height = 0; + } + return static_cast(status); +} + extern "C" int multi_stark_cuda_lde_attach_trace( int device_id, void* handle, const uint64_t* trace, size_t height, size_t width) { @@ -3176,22 +4033,56 @@ extern "C" int multi_stark_cuda_mixed_merkle_destroy(int device_id, return static_cast(cudaGetLastError()); } -extern "C" int multi_stark_cuda_mixed_merkle_create_from_ldes( +extern "C" int multi_stark_cuda_mixed_merkle_create_hybrid( int device_id, void** handle, uint8_t* root, - const void* const* lde_handles, size_t lde_count) { - if (handle == nullptr || root == nullptr || lde_handles == nullptr || - lde_count == 0) { + const void* const* lde_handles, const uint64_t* const* host_values, + const size_t* widths, const size_t* heights, size_t matrix_count, + const uint8_t* const* host_digest_groups, + const size_t* host_digest_heights, size_t host_digest_group_count) { + if (handle == nullptr || root == nullptr || + matrix_count == 0 || lde_handles == nullptr || host_values == nullptr || + widths == nullptr || heights == nullptr || + (host_digest_group_count != 0 && + (host_digest_groups == nullptr || host_digest_heights == nullptr))) { return static_cast(cudaErrorInvalidValue); } size_t max_height = 0; - for (size_t index = 0; index < lde_count; ++index) { + for (size_t index = 0; index < matrix_count; ++index) { const ResidentLde* lde = static_cast(lde_handles[index]); - if (lde == nullptr || !is_power_of_two(lde->height) || lde->width == 0) { + bool prehashed = false; + for (size_t group = 0; group < host_digest_group_count; ++group) { + prehashed = prehashed || host_digest_heights[group] == heights[index]; + } + if ((lde != nullptr && host_values[index] != nullptr) || + (lde == nullptr && host_values[index] == nullptr && !prehashed) || + !is_power_of_two(heights[index]) || widths[index] == 0 || + !product_fits(heights[index], widths[index]) || + (lde != nullptr && + (lde->height != heights[index] || lde->width != widths[index]))) { + return static_cast(cudaErrorInvalidValue); + } + if (heights[index] > max_height) { + max_height = heights[index]; + } + } + for (size_t group = 0; group < host_digest_group_count; ++group) { + if (host_digest_groups[group] == nullptr || + !is_power_of_two(host_digest_heights[group])) { return static_cast(cudaErrorInvalidValue); } - if (lde->height > max_height) { - max_height = lde->height; + for (size_t previous = 0; previous < group; ++previous) { + if (host_digest_heights[previous] == host_digest_heights[group]) { + return static_cast(cudaErrorInvalidValue); + } + } + bool found_height = false; + for (size_t index = 0; index < matrix_count; ++index) { + found_height = found_height || + heights[index] == host_digest_heights[group]; + } + if (!found_height) { + return static_cast(cudaErrorInvalidValue); } } if (max_height > SIZE_MAX / 64) { @@ -3211,20 +4102,33 @@ extern "C" int multi_stark_cuda_mixed_merkle_create_from_ldes( status = cudaMalloc(reinterpret_cast(&tree->digests), 64 * max_height); size_t max_injected_height = 0; - for (size_t index = 0; index < lde_count; ++index) { - const ResidentLde* lde = - static_cast(lde_handles[index]); - if (lde->height < max_height && lde->height > max_injected_height) { - max_injected_height = lde->height; + for (size_t index = 0; index < matrix_count; ++index) { + if (heights[index] < max_height && heights[index] > max_injected_height) { + max_injected_height = heights[index]; + } + } + for (size_t group = 0; group < host_digest_group_count; ++group) { + const size_t height = host_digest_heights[group]; + if (height < max_height && height > max_injected_height) { + max_injected_height = height; } } DeviceBuffer injected_digests; if (status == cudaSuccess && max_injected_height != 0) { status = injected_digests.allocate(4 * max_injected_height); } + const auto hash_group = [&](uint8_t* output, size_t height) { + for (size_t group = 0; group < host_digest_group_count; ++group) { + if (host_digest_heights[group] == height) { + return cudaMemcpy(output, host_digest_groups[group], height * 32, + cudaMemcpyHostToDevice); + } + } + return hash_partitioned_lde_group(output, lde_handles, host_values, + widths, heights, matrix_count, height); + }; if (status == cudaSuccess) { - status = hash_resident_lde_group(tree->digests, lde_handles, - lde_count, max_height); + status = hash_group(tree->digests, max_height); } size_t count = max_height; @@ -3238,15 +4142,12 @@ extern "C" int multi_stark_cuda_mixed_merkle_create_from_ldes( true); bool inject = false; - for (size_t index = 0; index < lde_count; ++index) { - const ResidentLde* lde = - static_cast(lde_handles[index]); - inject = inject || lde->height == count; + for (size_t index = 0; index < matrix_count; ++index) { + inject = inject || heights[index] == count; } if (status == cudaSuccess && inject) { - status = hash_resident_lde_group( - reinterpret_cast(injected_digests.get()), - lde_handles, lde_count, count); + status = hash_group( + reinterpret_cast(injected_digests.get()), count); } if (status == cudaSuccess && inject) { status = launch_blake3_digest_pairs( @@ -3267,6 +4168,85 @@ extern "C" int multi_stark_cuda_mixed_merkle_create_from_ldes( return static_cast(cudaSuccess); } +extern "C" int multi_stark_cuda_hash_hybrid_height_group( + int device_id, uint8_t* host_digests, const void* const* lde_handles, + const uint64_t* const* host_values, const size_t* widths, + const size_t* heights, size_t matrix_count, size_t height) { + if (device_id < 0 || host_digests == nullptr || lde_handles == nullptr || + host_values == nullptr || widths == nullptr || heights == nullptr || + matrix_count == 0 || !is_power_of_two(height) || height > SIZE_MAX / 4) { + return static_cast(cudaErrorInvalidValue); + } + bool has_resident = false; + for (size_t index = 0; index < matrix_count; ++index) { + const ResidentLde* lde = + static_cast(lde_handles[index]); + if ((lde == nullptr) == (host_values[index] == nullptr) || + widths[index] == 0 || heights[index] != height || + (lde != nullptr && + (lde->width != widths[index] || lde->height != height))) { + return static_cast(cudaErrorInvalidValue); + } + has_resident = has_resident || lde != nullptr; + } + if (!has_resident) { + return static_cast(cudaErrorInvalidValue); + } + + cudaError_t status = cudaSetDevice(device_id); + DeviceBuffer device_digests; + if (status == cudaSuccess) { + status = device_digests.allocate(height * 4); + } + if (status == cudaSuccess) { + status = hash_partitioned_lde_group( + reinterpret_cast(device_digests.get()), lde_handles, + host_values, widths, heights, matrix_count, height); + } + if (status == cudaSuccess) { + status = cudaMemcpy(host_digests, device_digests.get(), height * 32, + cudaMemcpyDeviceToHost); + } + return static_cast(status); +} + +extern "C" int multi_stark_cuda_mixed_merkle_create_from_ldes( + int device_id, void** handle, uint8_t* root, + const void* const* lde_handles, size_t lde_count) { + if (lde_handles == nullptr || lde_count == 0) { + return static_cast(cudaErrorInvalidValue); + } + const uint64_t** host_values = + new (std::nothrow) const uint64_t*[lde_count](); + size_t* widths = new (std::nothrow) size_t[lde_count]; + size_t* heights = new (std::nothrow) size_t[lde_count]; + if (host_values == nullptr || widths == nullptr || heights == nullptr) { + delete[] host_values; + delete[] widths; + delete[] heights; + return static_cast(cudaErrorMemoryAllocation); + } + for (size_t index = 0; index < lde_count; ++index) { + const ResidentLde* lde = + static_cast(lde_handles[index]); + if (lde == nullptr) { + delete[] host_values; + delete[] widths; + delete[] heights; + return static_cast(cudaErrorInvalidValue); + } + widths[index] = lde->width; + heights[index] = lde->height; + } + const int status = multi_stark_cuda_mixed_merkle_create_hybrid( + device_id, handle, root, lde_handles, host_values, widths, heights, + lde_count, nullptr, nullptr, 0); + delete[] host_values; + delete[] widths; + delete[] heights; + return status; +} + extern "C" int multi_stark_cuda_fri_merkle_create( int device_id,void** handle,uint8_t* root,const void* codeword_handle,size_t arity){ if(!codeword_handle||!is_power_of_two(arity))return static_cast(cudaErrorInvalidValue); @@ -3349,5 +4329,26 @@ extern "C" int multi_stark_cuda_memory_info(int device_id, size_t* free_bytes, } cudaError_t status = cudaSetDevice(device_id); if (status == cudaSuccess) status = cudaMemGetInfo(free_bytes, total_bytes); + // cudaMemGetInfo excludes pages retained by cudaMallocAsync's default + // pool, even though subsequent stream allocations can reuse them. Treat + // the unused part of that pool as available for admission decisions; using + // raw driver-free bytes alone causes needless spills after a large stage. + if (status == cudaSuccess) { + cudaMemPool_t pool = nullptr; + uint64_t reserved = 0; + uint64_t used = 0; + const cudaError_t pool_status = cudaDeviceGetDefaultMemPool(&pool, device_id); + if (pool_status == cudaSuccess && + cudaMemPoolGetAttribute(pool, cudaMemPoolAttrReservedMemCurrent, + &reserved) == cudaSuccess && + cudaMemPoolGetAttribute(pool, cudaMemPoolAttrUsedMemCurrent, &used) == + cudaSuccess && + reserved > used) { + const uint64_t reusable = reserved - used; + *free_bytes = reusable > SIZE_MAX - *free_bytes + ? SIZE_MAX + : *free_bytes + static_cast(reusable); + } + } return static_cast(status); } diff --git a/flake.nix b/flake.nix index b4848e1..b8fbe8f 100644 --- a/flake.nix +++ b/flake.nix @@ -40,7 +40,7 @@ # Pins the Rust toolchain rustToolchain = fenix.packages.${system}.fromToolchainFile { file = ./rust-toolchain.toml; - sha256 = "sha256-P30Tm3O7vQAE725YtDCDHGjNrSsfZO4us11UwJGZSJo="; + sha256 = "sha256-p8h3Sl/YRByZfZTAKXdsvF6xEenXKrXSVvpphmZENH4="; }; craneLib = (crane.mkLib pkgs).overrideToolchain rustToolchain; diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 23de054..06581a9 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,4 +1,4 @@ [toolchain] # The default profile includes rustc, rust-std, cargo, rust-docs, rustfmt and clippy. profile = "default" -channel = "1.98" +channel = "1.98.1" diff --git a/src/cuda/mmcs.rs b/src/cuda/mmcs.rs index 8f19966..7b419b6 100644 --- a/src/cuda/mmcs.rs +++ b/src/cuda/mmcs.rs @@ -1,10 +1,15 @@ use p3_commit::{BatchOpening, BatchOpeningRef, Mmcs}; +use p3_field::PrimeField64; use p3_goldilocks::Goldilocks; use p3_matrix::{Dimensions, Matrix, dense::RowMajorMatrix}; +use p3_maybe_rayon::prelude::*; use p3_merkle_tree::{MerkleTreeError, MerkleTreeMmcs}; -use p3_symmetric::MerkleCap; +use p3_symmetric::{CryptographicHasher, MerkleCap}; -use super::{CudaLde, CudaMixedMerkleTree, mixed_lde_open_row, mixed_lde_open_rows}; +use super::{ + CudaLde, CudaMixedMerkleTree, mixed_lde_open_row, mixed_lde_open_row_at_height, + mixed_lde_open_rows, mixed_lde_open_rows_at_height, +}; use crate::types::Blake3CompressionFunction; use p3_blake3::Blake3; use p3_symmetric::SerializingHasher; @@ -12,23 +17,136 @@ use p3_symmetric::SerializingHasher; type CpuMmcs = MerkleTreeMmcs, Blake3CompressionFunction, 2, 32>; type CpuData = >::ProverData; +pub(crate) type DeferredMatrix = Option<(Dimensions, std::thread::JoinHandle)>; +type DeferredWorker = Option>>>; + +#[derive(Clone, Copy)] +pub enum CudaMatrixSource<'a> { + Resident(&'a CudaLde), + Host(&'a RowMajorMatrix), +} pub enum CudaMmcsData { Cpu(CpuData), + Hybrid { + // Drop resident LDEs before their retained host traces. + resident: Vec>, + resident_active: Vec, + materialize: Box M + Send + Sync>, + committed_matrices: Vec>, + deferred_matrices: Vec>, + dimensions: Vec, + retained_traces: Vec>, + tree: CudaMixedMerkleTree, + }, Cuda { resident: std::sync::Arc>, materialize: Option Vec + Send + Sync>>, // `materialize` can own the last Arc to `resident`; drop it before // retained host matrices so pinned trace pointers cannot outlive them. - matrices: std::sync::OnceLock>, + committed_matrices: std::sync::OnceLock>, + retained_traces: std::sync::OnceLock>, tree: CudaMixedMerkleTree, }, } +fn hybrid_matrix<'a, M>( + matrices: &'a [std::sync::OnceLock], + deferred: &[DeferredWorker], + index: usize, +) -> &'a M { + matrices[index].get_or_init(|| { + let started = std::time::Instant::now(); + let matrix = deferred[index] + .as_ref() + .expect("hybrid matrix has neither storage nor deferred computation") + .lock() + .expect("deferred matrix lock poisoned") + .take() + .expect("deferred matrix worker already consumed") + .join() + .expect("deferred matrix worker panicked"); + if super::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] waited for deferred matrix {index}: {:.3}s", + started.elapsed().as_secs_f64() + ); + } + matrix + }) +} + +fn hybrid_resident_candidates( + data: &CudaMmcsData, + protected_index: Option, +) -> Vec<(usize, usize, bool)> { + let CudaMmcsData::Hybrid { + resident, + resident_active, + committed_matrices, + .. + } = data + else { + return Vec::new(); + }; + resident + .iter() + .enumerate() + .filter_map(|(index, lde)| { + if protected_index == Some(index) { + return None; + } + let lde = lde.as_ref()?; + resident_active[index] + .load(std::sync::atomic::Ordering::Acquire) + .then_some(( + index, + lde.height().saturating_mul(lde.width()), + committed_matrices[index].get().is_some(), + )) + }) + .collect() +} + +fn evict_hybrid_resident(data: &CudaMmcsData, index: usize) { + let CudaMmcsData::Hybrid { + resident, + resident_active, + materialize, + committed_matrices, + .. + } = data + else { + panic!("attempted to evict a non-hybrid CUDA commitment"); + }; + let lde = resident[index] + .as_ref() + .expect("hybrid resident matrix is missing its LDE"); + committed_matrices[index].get_or_init(|| materialize(lde)); + // SAFETY: admission transitions run between proving stages, when no CUDA + // operation can access this LDE or its trace. + unsafe { lde.release_values() }; + resident_active[index].store(false, std::sync::atomic::Ordering::Release); +} + impl CudaMmcsData { pub(crate) fn resident(&self, index: usize) -> Option<&CudaLde> { match self { Self::Cuda { resident, .. } => resident.get(index), + Self::Hybrid { + resident, + resident_active, + .. + } => { + if resident_active + .get(index)? + .load(std::sync::atomic::Ordering::Acquire) + { + resident.get(index)?.as_ref() + } else { + None + } + } Self::Cpu(_) => None, } } @@ -38,10 +156,12 @@ impl CudaMmcsData> { pub(crate) fn resident_with_trace(&self, index: usize) -> Option<&CudaLde> { match self { Self::Cuda { - resident, matrices, .. + resident, + retained_traces, + .. } => { let lde = resident.get(index)?; - if let Some(retained) = matrices.get() { + if let Some(retained) = retained_traces.get() { // SAFETY: the matrix is retained in this prover data and // drops after every Arc owner of `resident`. This proving // path serializes attachment changes with CUDA use. @@ -49,11 +169,219 @@ impl CudaMmcsData> { } Some(lde) } + Self::Hybrid { + resident, + resident_active, + retained_traces, + .. + } => { + if !resident_active + .get(index)? + .load(std::sync::atomic::Ordering::Acquire) + { + return None; + } + let lde = resident.get(index)?.as_ref()?; + if let Some(trace) = retained_traces.get(index)?.as_ref() { + // SAFETY: the trace is owned by this prover data and drops + // after the resident LDE which holds the registered pointer. + unsafe { lde.attach_trace(trace) }; + } + Some(lde) + } Self::Cpu(_) => None, } } } +#[cfg(test)] +pub(crate) fn hash_cpu_height_groups( + matrices: &[Option>], +) -> Vec<(usize, Vec<[u8; 32]>)> { + let mut groups = std::collections::BTreeMap::>>::new(); + for matrix in matrices.iter().flatten() { + groups.entry(matrix.height()).or_default().push(matrix); + } + groups + .into_iter() + .map(|(height, matrices)| { + let hasher = SerializingHasher::new(Blake3); + let digests = (0..height) + .into_par_iter() + .map(|row| { + hasher.hash_iter(matrices.iter().flat_map(|matrix| matrix.row(row).unwrap())) + }) + .collect(); + (height, digests) + }) + .collect() +} + +pub(crate) fn hash_host_only_height_groups( + matrices: &[Option>], + resident: &[Option], + deferred_dimensions: &[Option], + prehashed_heights: &std::collections::BTreeSet, +) -> Vec<(usize, Vec<[u8; 32]>)> { + assert_eq!(matrices.len(), resident.len()); + assert_eq!(matrices.len(), deferred_dimensions.len()); + let mut groups = std::collections::BTreeMap::>::new(); + for (index, ((matrix, lde), deferred)) in matrices + .iter() + .zip(resident) + .zip(deferred_dimensions) + .enumerate() + { + assert!( + usize::from(matrix.is_some()) + + usize::from(lde.is_some()) + + usize::from(deferred.is_some()) + == 1, + "every partitioned matrix must have exactly one storage backend" + ); + let height = matrix.as_ref().map_or_else( + || { + lde.as_ref() + .map_or_else(|| deferred.unwrap().height, CudaLde::height) + }, + Matrix::height, + ); + assert!( + deferred.is_none() || prehashed_heights.contains(&height), + "a deferred matrix requires a pre-hashed height group" + ); + groups.entry(height).or_default().push(index); + } + + let hasher = SerializingHasher::new(Blake3); + groups + .into_iter() + .filter_map(|(height, indices)| { + if prehashed_heights.contains(&height) + || indices.iter().any(|&index| resident[index].is_some()) + { + // CUDA hashes fully resident groups directly and assembles mixed + // groups from resident and mapped host columns in bounded chunks. + return None; + } + let sources = indices + .iter() + .map(|&index| matrices[index].as_ref().unwrap()) + .collect::>(); + let digests = (0..height) + .into_par_iter() + .map(|row| { + hasher.hash_iter(sources.iter().flat_map(|matrix| matrix.row(row).unwrap())) + }) + .collect(); + Some((height, digests)) + }) + .collect() +} + +fn hybrid_max_height(dimensions: &[Dimensions]) -> usize { + dimensions + .iter() + .map(|dimensions| dimensions.height) + .max() + .unwrap() +} + +fn hybrid_open_row>( + resident: &[Option], + resident_active: &[std::sync::atomic::AtomicBool], + matrices: &[std::sync::OnceLock], + deferred: &[DeferredWorker], + dimensions: &[Dimensions], + index: usize, +) -> Vec> { + let max_height = hybrid_max_height(dimensions); + assert!(index < max_height, "MMCS opening index out of bounds"); + let active_ldes = resident + .iter() + .zip(resident_active) + .filter_map(|(resident, active)| { + active + .load(std::sync::atomic::Ordering::Acquire) + .then_some(resident.as_ref()) + .flatten() + }) + .collect::>(); + let mut gpu_rows = if !active_ldes.is_empty() { + mixed_lde_open_row_at_height(active_ldes, max_height, index).into_iter() + } else { + Vec::new().into_iter() + }; + resident + .iter() + .zip(resident_active) + .enumerate() + .map(|(matrix_index, (resident, active))| { + if resident.is_some() && active.load(std::sync::atomic::Ordering::Acquire) { + gpu_rows.next().unwrap() + } else { + let matrix = hybrid_matrix(matrices, deferred, matrix_index); + let row = index >> (max_height.trailing_zeros() - matrix.height().trailing_zeros()); + matrix.row(row).unwrap().into_iter().collect() + } + }) + .collect() +} + +fn hybrid_open_rows>( + resident: &[Option], + resident_active: &[std::sync::atomic::AtomicBool], + matrices: &[std::sync::OnceLock], + deferred: &[DeferredWorker], + dimensions: &[Dimensions], + indices: &[usize], +) -> Vec>> { + if indices.is_empty() { + return Vec::new(); + } + let max_height = hybrid_max_height(dimensions); + assert!(indices.iter().all(|&index| index < max_height)); + let active_ldes = resident + .iter() + .zip(resident_active) + .filter_map(|(resident, active)| { + active + .load(std::sync::atomic::Ordering::Acquire) + .then_some(resident.as_ref()) + .flatten() + }) + .collect::>(); + let mut gpu_queries = if !active_ldes.is_empty() { + mixed_lde_open_rows_at_height(active_ldes, max_height, indices).into_iter() + } else { + (0..indices.len()) + .map(|_| Vec::new()) + .collect::>() + .into_iter() + }; + indices + .iter() + .map(|&index| { + let mut gpu_rows = gpu_queries.next().unwrap().into_iter(); + resident + .iter() + .zip(resident_active) + .enumerate() + .map(|(matrix_index, (resident, active))| { + if resident.is_some() && active.load(std::sync::atomic::Ordering::Acquire) { + gpu_rows.next().unwrap() + } else { + let matrix = hybrid_matrix(matrices, deferred, matrix_index); + let row = index + >> (max_height.trailing_zeros() - matrix.height().trailing_zeros()); + matrix.row(row).unwrap().into_iter().collect() + } + }) + .collect() + }) + .collect() +} + #[derive(Clone, Debug)] pub struct CudaMmcs { cpu: CpuMmcs, @@ -72,6 +400,67 @@ impl CudaMmcs { pub trait CudaCommitMmcs: Mmcs { fn cuda_device_id(&self) -> i32; + fn hash_cuda_hybrid_height_group( + &self, + resident: &[Option<&CudaLde>], + host: &[Option<&RowMajorMatrix>], + ) -> Vec<[u8; 32]>; + + fn is_cuda_resident>(&self, data: &Self::ProverData) -> bool; + + fn is_matrix_cuda_resident>( + &self, + data: &Self::ProverData, + index: usize, + ) -> bool; + + /// Reports whether a matrix can be accessed without waiting for deferred + /// host materialization. + fn is_matrix_source_ready>( + &self, + data: &Self::ProverData, + index: usize, + ) -> bool; + + /// Materializes and releases the largest hybrid-resident matrices until + /// the device has at least `target_free_bytes` available. Returns the + /// measured free bytes after eviction. + fn ensure_device_headroom>( + &self, + data: &Self::ProverData, + target_free_bytes: usize, + protected_index: Option, + ) -> usize; + + /// Selects spill candidates across all supplied commitments instead of + /// letting iteration order decide which commitment gives up residency. + fn ensure_device_headroom_batch( + &self, + data: &[&Self::ProverData>], + target_free_bytes: usize, + ) -> usize; + + fn matrix_dimensions(&self, data: &Self::ProverData>) -> Vec; + + fn cpu_matrices<'a>( + &self, + data: &'a Self::ProverData>, + ) -> Vec>>; + + fn with_resident_matrix( + &self, + data: &Self::ProverData>, + index: usize, + f: impl FnOnce(&CudaLde) -> R, + ) -> R; + + fn with_matrix_source( + &self, + data: &Self::ProverData>, + index: usize, + f: impl FnOnce(CudaMatrixSource<'_>) -> R, + ) -> R; + fn resident_or_upload>( &self, data: &Self::ProverData, @@ -81,6 +470,20 @@ pub trait CudaCommitMmcs: Mmcs { ldes: Vec, ) -> (Self::Commitment, Self::ProverData>); + fn commit_cuda_spillable( + &self, + ldes: Vec, + ) -> (Self::Commitment, Self::ProverData>); + + fn commit_cuda_hybrid( + &self, + resident: Vec>, + host_matrices: Vec>>, + deferred_matrices: Vec>>, + retained_traces: Vec>>, + host_digest_groups: Vec<(usize, Vec<[u8; 32]>)>, + ) -> (Self::Commitment, Self::ProverData>); + fn retain_matrices( &self, data: &mut Self::ProverData>, @@ -91,6 +494,11 @@ pub trait CudaCommitMmcs: Mmcs { &self, ldes: Vec>, ) -> (Self::Commitment, Self::ProverData>); + + fn commit_cpu_storage( + &self, + ldes: Vec>, + ) -> (Self::Commitment, Self::ProverData>); } impl CudaCommitMmcs for CudaMmcs { @@ -98,15 +506,398 @@ impl CudaCommitMmcs for CudaMmcs { self.device_id } + fn hash_cuda_hybrid_height_group( + &self, + resident: &[Option<&CudaLde>], + host: &[Option<&RowMajorMatrix>], + ) -> Vec<[u8; 32]> { + CudaMixedMerkleTree::hash_hybrid_height_group(self.device_id, resident, host) + } + + fn is_cuda_resident>(&self, data: &Self::ProverData) -> bool { + matches!(data, CudaMmcsData::Cuda { .. }) + } + + fn is_matrix_cuda_resident>( + &self, + data: &Self::ProverData, + index: usize, + ) -> bool { + match data { + CudaMmcsData::Cpu(_) => false, + CudaMmcsData::Hybrid { + resident, + resident_active, + .. + } => { + resident + .get(index) + .expect("hybrid matrix index out of bounds"); + resident_active[index].load(std::sync::atomic::Ordering::Acquire) + } + CudaMmcsData::Cuda { resident, .. } => { + assert!(index < resident.len(), "CUDA matrix index out of bounds"); + true + } + } + } + + fn is_matrix_source_ready>( + &self, + data: &Self::ProverData, + index: usize, + ) -> bool { + match data { + CudaMmcsData::Cpu(data) => { + assert!( + index < self.cpu.get_matrix_heights(data).len(), + "CPU matrix index out of bounds" + ); + true + } + CudaMmcsData::Cuda { resident, .. } => { + assert!(index < resident.len(), "CUDA matrix index out of bounds"); + true + } + CudaMmcsData::Hybrid { + resident_active, + committed_matrices, + deferred_matrices, + .. + } => { + let active = resident_active + .get(index) + .expect("hybrid matrix index out of bounds") + .load(std::sync::atomic::Ordering::Acquire); + active + || committed_matrices[index].get().is_some() + || deferred_matrices[index].as_ref().is_some_and(|worker| { + worker + .lock() + .expect("deferred matrix lock poisoned") + .as_ref() + .is_some_and(std::thread::JoinHandle::is_finished) + }) + } + } + } + + fn ensure_device_headroom>( + &self, + data: &Self::ProverData, + target_free_bytes: usize, + protected_index: Option, + ) -> usize { + let (mut free_bytes, _) = super::device_memory_info(self.device_id); + if super::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] admission start: target={} free={}", + target_free_bytes, free_bytes + ); + } + if free_bytes >= target_free_bytes { + return free_bytes; + } + let mut candidates = hybrid_resident_candidates(data, protected_index); + while free_bytes < target_free_bytes && !candidates.is_empty() { + let deficit_cells = target_free_bytes + .saturating_sub(free_bytes) + .div_ceil(size_of::()); + let candidate = candidates + .iter() + .enumerate() + .filter(|(_, (_, cells, _))| *cells >= deficit_cells) + .min_by_key(|(_, (_, cells, _))| *cells) + .or_else(|| { + candidates + .iter() + .enumerate() + .max_by_key(|(_, (_, cells, _))| *cells) + }) + .map(|(position, _)| position) + .unwrap(); + let (index, cells, _) = candidates.swap_remove(candidate); + evict_hybrid_resident(data, index); + // Eviction is a synchronous device release. Re-read the allocator + // instead of inferring availability from logical matrix sizes: + // CUDA allocation granularity and per-thread memory pools make a + // byte-sum projection unreliable across Rayon worker streams. + free_bytes = super::device_memory_info(self.device_id).0; + if super::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] evicted matrix {index}: cells={cells} free={free_bytes}" + ); + } + } + free_bytes + } + + fn ensure_device_headroom_batch( + &self, + data: &[&Self::ProverData>], + target_free_bytes: usize, + ) -> usize { + let (mut measured_free_bytes, _) = super::device_memory_info(self.device_id); + if super::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] batch admission start: target={} free={}", + target_free_bytes, measured_free_bytes + ); + } + if measured_free_bytes >= target_free_bytes { + return measured_free_bytes; + } + let initial_free_bytes = measured_free_bytes; + let mut released_bytes = 0usize; + let mut candidates = data + .iter() + .enumerate() + .flat_map(|(data_index, data)| { + hybrid_resident_candidates(data, None).into_iter().map( + move |(matrix_index, cells, host_ready)| { + (data_index, matrix_index, cells, host_ready) + }, + ) + }) + .collect::>(); + while initial_free_bytes + .saturating_add(released_bytes) + .max(measured_free_bytes) + < target_free_bytes + && !candidates.is_empty() + { + let available_bytes = initial_free_bytes + .saturating_add(released_bytes) + .max(measured_free_bytes); + let deficit_cells = target_free_bytes + .saturating_sub(available_bytes) + .div_ceil(size_of::()); + let ready = candidates.iter().any(|candidate| candidate.3); + let eligible = |candidate: &&(usize, usize, usize, bool)| !ready || candidate.3; + let candidate = candidates + .iter() + .enumerate() + .filter(|(_, candidate)| eligible(candidate)) + .filter(|(_, (_, _, cells, _))| *cells <= deficit_cells) + .max_by_key(|(_, (_, _, cells, _))| *cells) + .or_else(|| { + candidates + .iter() + .enumerate() + .filter(|(_, candidate)| eligible(candidate)) + .min_by_key(|(_, (_, _, cells, _))| *cells) + }) + .map(|(position, _)| position) + .unwrap(); + let (data_index, matrix_index, cells, host_ready) = candidates.swap_remove(candidate); + evict_hybrid_resident(data[data_index], matrix_index); + // `release_values` synchronously frees this exact payload. CUDA's + // free-memory counters advance at allocator granularity, however, + // so requiring `cudaMemGetInfo` to reflect every small release can + // evict hundreds of matrices to cover a sub-granule reporting gap. + // The released payload is a conservative lower bound (the LDE may + // also release interpolation scratch), while the fresh measurement + // catches unrelated memory becoming available. + released_bytes = + released_bytes.saturating_add(cells.saturating_mul(size_of::())); + measured_free_bytes = super::device_memory_info(self.device_id).0; + if super::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] batch evicted commitment {data_index} matrix {matrix_index}: cells={cells} host_ready={host_ready} measured_free={measured_free_bytes} projected_free={}", + initial_free_bytes.saturating_add(released_bytes) + ); + } + } + initial_free_bytes + .saturating_add(released_bytes) + .max(measured_free_bytes) + } + + fn matrix_dimensions( + &self, + data: &Self::ProverData>, + ) -> Vec { + match data { + CudaMmcsData::Cpu(data) => self + .cpu + .get_matrices(data) + .into_iter() + .map(Matrix::dimensions) + .collect(), + CudaMmcsData::Hybrid { dimensions, .. } => dimensions.clone(), + CudaMmcsData::Cuda { resident, .. } => resident + .iter() + .map(|lde| Dimensions { + width: lde.width(), + height: lde.height(), + }) + .collect(), + } + } + + fn cpu_matrices<'a>( + &self, + data: &'a Self::ProverData>, + ) -> Vec>> { + match data { + CudaMmcsData::Cpu(data) => self.cpu.get_matrices(data).into_iter().map(Some).collect(), + CudaMmcsData::Hybrid { + resident, + resident_active, + committed_matrices, + deferred_matrices, + .. + } => resident + .iter() + .zip(resident_active) + .enumerate() + .map(|(index, (resident, active))| { + (resident.is_none() || !active.load(std::sync::atomic::Ordering::Acquire)) + .then(|| hybrid_matrix(committed_matrices, deferred_matrices, index)) + }) + .collect(), + CudaMmcsData::Cuda { resident, .. } => (0..resident.len()).map(|_| None).collect(), + } + } + + fn with_resident_matrix( + &self, + data: &Self::ProverData>, + index: usize, + f: impl FnOnce(&CudaLde) -> R, + ) -> R { + match data { + CudaMmcsData::Cuda { resident, .. } => f(resident + .get(index) + .expect("CUDA matrix index out of bounds")), + CudaMmcsData::Hybrid { + resident, + resident_active, + committed_matrices, + deferred_matrices, + .. + } => { + let resident = resident + .get(index) + .expect("hybrid matrix index out of bounds"); + if resident_active[index].load(std::sync::atomic::Ordering::Acquire) { + let lde = resident.as_ref().expect("active hybrid matrix has no LDE"); + f(lde) + } else { + let matrix = hybrid_matrix(committed_matrices, deferred_matrices, index); + if super::memory_diagnostics_enabled() { + let (free_bytes, _) = super::device_memory_info(self.device_id); + eprintln!( + "[multi-stark/cuda] upload hybrid matrix {index}: bytes={} free={free_bytes}", + matrix + .height() + .saturating_mul(matrix.width()) + .saturating_mul(size_of::()) + ); + } + let uploaded = CudaLde::from_row_major_matrix(self.device_id, matrix); + f(&uploaded) + } + } + CudaMmcsData::Cpu(data) => { + let matrices = self.cpu.get_matrices(data); + let matrix = matrices.get(index).expect("CPU matrix index out of bounds"); + if super::memory_diagnostics_enabled() { + let (free_bytes, _) = super::device_memory_info(self.device_id); + eprintln!( + "[multi-stark/cuda] upload CPU matrix {index}: bytes={} free={free_bytes}", + matrix + .height() + .saturating_mul(matrix.width()) + .saturating_mul(size_of::()) + ); + } + let uploaded = CudaLde::from_row_major_matrix(self.device_id, matrix); + f(&uploaded) + } + } + } + + fn with_matrix_source( + &self, + data: &Self::ProverData>, + index: usize, + f: impl FnOnce(CudaMatrixSource<'_>) -> R, + ) -> R { + match data { + CudaMmcsData::Cuda { resident, .. } => f(CudaMatrixSource::Resident( + resident + .get(index) + .expect("CUDA matrix index out of bounds"), + )), + CudaMmcsData::Hybrid { + resident, + resident_active, + committed_matrices, + deferred_matrices, + .. + } => { + if resident_active + .get(index) + .expect("hybrid matrix index out of bounds") + .load(std::sync::atomic::Ordering::Acquire) + { + f(CudaMatrixSource::Resident( + resident[index] + .as_ref() + .expect("active hybrid matrix has no LDE"), + )) + } else { + f(CudaMatrixSource::Host(hybrid_matrix( + committed_matrices, + deferred_matrices, + index, + ))) + } + } + CudaMmcsData::Cpu(data) => { + let matrices = self.cpu.get_matrices(data); + f(CudaMatrixSource::Host( + matrices.get(index).expect("CPU matrix index out of bounds"), + )) + } + } + } + + fn commit_cpu_storage( + &self, + ldes: Vec>, + ) -> ( + Self::Commitment, + Self::ProverData>, + ) { + let (commitment, data) = self.cpu.commit(ldes); + (commitment, CudaMmcsData::Cpu(data)) + } + fn retain_matrices( &self, data: &mut Self::ProverData>, retained: Vec>, ) { - if let CudaMmcsData::Cuda { matrices, .. } = data { - matrices + match data { + CudaMmcsData::Cuda { + retained_traces, .. + } => retained_traces .set(retained) - .expect("fresh CUDA trace matrix cell"); + .expect("fresh CUDA trace matrix cell"), + CudaMmcsData::Hybrid { + retained_traces, .. + } => { + assert_eq!(retained_traces.len(), retained.len()); + assert!(retained_traces.iter().all(Option::is_none)); + for (slot, matrix) in retained_traces.iter_mut().zip(retained) { + *slot = Some(matrix); + } + } + CudaMmcsData::Cpu(_) => { + panic!("CPU commitment cannot retain CUDA trace matrices") + } } } @@ -116,6 +907,9 @@ impl CudaCommitMmcs for CudaMmcs { ) -> std::sync::Arc> { match data { CudaMmcsData::Cuda { resident, .. } => std::sync::Arc::clone(resident), + CudaMmcsData::Hybrid { .. } => { + panic!("hybrid commitments cannot be converted to all-resident storage") + } CudaMmcsData::Cpu(cpu) => std::sync::Arc::new( self.cpu .get_matrices(cpu) @@ -154,7 +948,114 @@ impl CudaCommitMmcs for CudaMmcs { .map(CudaLde::to_row_major_matrix) .collect() })), - matrices: std::sync::OnceLock::new(), + committed_matrices: std::sync::OnceLock::new(), + retained_traces: std::sync::OnceLock::new(), + tree, + }, + ) + } + + fn commit_cuda_spillable( + &self, + ldes: Vec, + ) -> ( + Self::Commitment, + Self::ProverData>, + ) { + let matrix_count = ldes.len(); + self.commit_cuda_hybrid( + ldes.into_iter().map(Some).collect(), + (0..matrix_count).map(|_| None).collect(), + (0..matrix_count).map(|_| None).collect(), + (0..matrix_count).map(|_| None).collect(), + Vec::new(), + ) + } + + fn commit_cuda_hybrid( + &self, + resident: Vec>, + host_matrices: Vec>>, + deferred_matrices: Vec>>, + retained_traces: Vec>>, + host_digest_groups: Vec<(usize, Vec<[u8; 32]>)>, + ) -> ( + Self::Commitment, + Self::ProverData>, + ) { + assert_eq!(resident.len(), host_matrices.len()); + assert_eq!(resident.len(), deferred_matrices.len()); + assert_eq!(resident.len(), retained_traces.len()); + assert!( + resident + .iter() + .zip(&host_matrices) + .zip(&deferred_matrices) + .all(|((gpu, cpu), deferred)| usize::from(gpu.is_some()) + + usize::from(cpu.is_some()) + + usize::from(deferred.is_some()) + == 1), + "every hybrid matrix must have exactly one storage backend" + ); + let deferred_dimensions = deferred_matrices + .iter() + .map(|deferred| deferred.as_ref().map(|(dimensions, _)| *dimensions)) + .collect::>(); + let dimensions = resident + .iter() + .zip(&host_matrices) + .zip(&deferred_dimensions) + .map(|((resident, host), deferred)| { + resident.as_ref().map_or_else( + || { + host.as_ref() + .map_or_else(|| deferred.unwrap(), Matrix::dimensions) + }, + |lde| Dimensions { + width: lde.width(), + height: lde.height(), + }, + ) + }) + .collect::>(); + let resident_refs = resident.iter().map(Option::as_ref).collect::>(); + let host_refs = host_matrices.iter().map(Option::as_ref).collect::>(); + let tree = CudaMixedMerkleTree::from_hybrid( + self.device_id, + &resident_refs, + &host_refs, + &deferred_dimensions, + &host_digest_groups, + ); + let commitment = MerkleCap::new(vec![tree.root()]); + let committed_matrices = host_matrices + .into_iter() + .map(|matrix| { + let cell = std::sync::OnceLock::new(); + if let Some(matrix) = matrix { + cell.set(matrix).expect("fresh hybrid matrix cell"); + } + cell + }) + .collect(); + let deferred_matrices = deferred_matrices + .into_iter() + .map(|deferred| deferred.map(|(_, worker)| std::sync::Mutex::new(Some(worker)))) + .collect(); + let resident_active = resident + .iter() + .map(|lde| std::sync::atomic::AtomicBool::new(lde.is_some())) + .collect(); + ( + commitment, + CudaMmcsData::Hybrid { + resident, + resident_active, + materialize: Box::new(CudaLde::to_row_major_matrix), + committed_matrices, + deferred_matrices, + dimensions, + retained_traces, tree, }, ) @@ -172,8 +1073,11 @@ impl CudaCommitMmcs for CudaMmcs { .map(|matrix| CudaLde::from_row_major_matrix(self.device_id, matrix)) .collect(); let (commitment, mut data) = self.commit_cuda_resident(resident); - if let CudaMmcsData::Cuda { matrices, .. } = &mut data { - let _ = matrices.set(ldes); + if let CudaMmcsData::Cuda { + committed_matrices, .. + } = &mut data + { + let _ = committed_matrices.set(ldes); } (commitment, data) } @@ -211,14 +1115,18 @@ impl Mmcs for CudaMmcs { let resident = std::sync::Arc::new(resident); let tree = CudaMixedMerkleTree::from_ldes(self.device_id, &resident); let commitment = MerkleCap::new(vec![tree.root()]); - let matrices = std::sync::OnceLock::new(); - matrices.set(inputs).ok().expect("fresh CUDA matrix cell"); + let committed_matrices = std::sync::OnceLock::new(); + committed_matrices + .set(inputs) + .ok() + .expect("fresh CUDA matrix cell"); return ( commitment, CudaMmcsData::Cuda { resident, materialize: None, - matrices, + committed_matrices, + retained_traces: std::sync::OnceLock::new(), tree, }, ); @@ -237,6 +1145,25 @@ impl Mmcs for CudaMmcs { let opening = self.cpu.open_batch(index, data); BatchOpening::new(opening.opened_values, opening.opening_proof) } + CudaMmcsData::Hybrid { + resident, + resident_active, + committed_matrices, + deferred_matrices, + dimensions, + tree, + .. + } => BatchOpening::new( + hybrid_open_row( + resident, + resident_active, + committed_matrices, + deferred_matrices, + dimensions, + index, + ), + tree.open_siblings(index), + ), CudaMmcsData::Cuda { resident, tree, .. } => { let opened_values = mixed_lde_open_row(resident, index); BatchOpening::new(opened_values, tree.open_siblings(index)) @@ -250,11 +1177,29 @@ impl Mmcs for CudaMmcs { ) -> Vec<&'a M> { match prover_data { CudaMmcsData::Cpu(data) => self.cpu.get_matrices(data), + CudaMmcsData::Hybrid { + resident, + materialize, + committed_matrices, + deferred_matrices, + .. + } => resident + .iter() + .enumerate() + .map(|(index, resident)| { + if resident.is_some() { + committed_matrices[index] + .get_or_init(|| materialize(resident.as_ref().unwrap())) + } else { + hybrid_matrix(committed_matrices, deferred_matrices, index) + } + }) + .collect(), CudaMmcsData::Cuda { - matrices, + committed_matrices, materialize, .. - } => matrices + } => committed_matrices .get_or_init(|| { materialize .as_ref() @@ -271,6 +1216,10 @@ impl Mmcs for CudaMmcs { ) -> Vec { match prover_data { CudaMmcsData::Cpu(data) => self.cpu.get_matrix_heights(data), + CudaMmcsData::Hybrid { dimensions, .. } => dimensions + .iter() + .map(|dimensions| dimensions.height) + .collect(), CudaMmcsData::Cuda { resident, .. } => resident.iter().map(CudaLde::height).collect(), } } @@ -278,6 +1227,7 @@ impl Mmcs for CudaMmcs { fn get_max_height>(&self, prover_data: &Self::ProverData) -> usize { match prover_data { CudaMmcsData::Cpu(data) => self.cpu.get_max_height(data), + CudaMmcsData::Hybrid { dimensions, .. } => hybrid_max_height(dimensions), CudaMmcsData::Cuda { resident, .. } => { resident.iter().map(CudaLde::height).max().unwrap() } @@ -306,6 +1256,26 @@ impl Mmcs for CudaMmcs { ) -> (Vec>>, Self::MultiProof) { match prover_data { CudaMmcsData::Cpu(data) => self.cpu.open_multi_batch(indices, data), + CudaMmcsData::Hybrid { + resident, + resident_active, + committed_matrices, + deferred_matrices, + dimensions, + tree, + .. + } => { + let opened_values = hybrid_open_rows( + resident, + resident_active, + committed_matrices, + deferred_matrices, + dimensions, + indices, + ); + let opening_proof = tree.open_pruned_siblings(indices); + (opened_values, opening_proof) + } CudaMmcsData::Cuda { resident, tree, .. } => { let opened_values = if indices.is_empty() { Vec::new() @@ -330,3 +1300,78 @@ impl Mmcs for CudaMmcs { .verify_multi_batch(commit, dimensions, indices, opened_values, proof) } } + +#[cfg(test)] +mod tests { + use p3_commit::Mmcs; + use p3_field::PrimeCharacteristicRing; + + use super::*; + + fn matrix(height: usize, width: usize, offset: usize) -> RowMajorMatrix { + RowMajorMatrix::new( + (0..height * width) + .map(|index| Goldilocks::from_usize(offset + index)) + .collect(), + width, + ) + } + + #[test] + fn hybrid_openings_survive_device_spill() { + let matrices = vec![matrix(16, 2, 3), matrix(8, 3, 71), matrix(16, 1, 109)]; + let cpu = CpuMmcs::new( + SerializingHasher::new(Blake3), + Blake3CompressionFunction::new(Blake3), + 0, + ); + let (expected_commitment, expected_data) = cpu.commit(matrices.clone()); + let expected_opening = cpu.open_batch(5, &expected_data); + + let mmcs = CudaMmcs::new(cpu); + let resident = vec![ + Some(CudaLde::from_row_major_matrix(mmcs.device_id, &matrices[0])), + None, + None, + ]; + let host_matrices = vec![None, Some(matrices[1].clone()), Some(matrices[2].clone())]; + let host_digest_groups = hash_host_only_height_groups( + &host_matrices, + &resident, + &[None, None, None], + &std::collections::BTreeSet::new(), + ); + let retained_traces = vec![None, None, None]; + let deferred_matrices = vec![None, None, None]; + let (commitment, data) = mmcs.commit_cuda_hybrid( + resident, + host_matrices, + deferred_matrices, + retained_traces, + host_digest_groups, + ); + assert_eq!(commitment, expected_commitment); + + let resident_opening = mmcs.open_batch(5, &data); + assert_eq!( + resident_opening.opened_values, + expected_opening.opened_values + ); + assert_eq!( + resident_opening.opening_proof, + expected_opening.opening_proof + ); + + mmcs.ensure_device_headroom(&data, usize::MAX, None); + assert!(!(0..matrices.len()).any(|index| mmcs.is_matrix_cuda_resident(&data, index))); + let spilled_opening = mmcs.open_batch(5, &data); + assert_eq!( + spilled_opening.opened_values, + expected_opening.opened_values + ); + assert_eq!( + spilled_opening.opening_proof, + expected_opening.opening_proof + ); + } +} diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 017a365..a59619b 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -20,9 +20,9 @@ use crate::graph::{ConstraintGraph, Node}; use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft}; use p3_field::{Field, PrimeCharacteristicRing, PrimeField64, TwoAdicField}; use p3_goldilocks::Goldilocks; -use p3_matrix::Matrix; use p3_matrix::bitrev::{BitReversalPerm, BitReversedMatrixView}; use p3_matrix::dense::RowMajorMatrix; +use p3_matrix::{Dimensions, Matrix}; use p3_merkle_tree::PrunedMerklePaths; use p3_util::log2_strict_usize; @@ -267,6 +267,19 @@ impl CudaLde { check_cuda(status, "resident trace release"); } + /// Releases the device values after they have been materialized on the + /// host. The handle remains valid for dimensions and eventual destruction. + /// + /// # Safety + /// + /// The caller must serialize this transition with every operation which + /// can access the LDE or its attached trace. + pub(crate) unsafe fn release_values(&self) { + let status = + unsafe { multi_stark_cuda_lde_release_values(self.device_id, self.handle.as_ptr()) }; + check_cuda(status, "resident LDE value release"); + } + /// # Safety /// /// `trace` must outlive this handle or an earlier `release_trace` call, @@ -361,16 +374,32 @@ impl CudaLde { /// oracle/debug escape hatch; resident PCS code should retain the handle. #[must_use] pub(crate) fn to_row_major_matrix(&self) -> RowMajorMatrix { - let mut values = Goldilocks::zero_vec(self.height * self.width); + let len = self.height * self.width; + let mut storage = Vec::>::with_capacity(len); + // SAFETY: the CUDA copy below initializes every byte before the storage + // is converted to `Vec`. On failure it remains a safely + // droppable vector of `MaybeUninit` values. + unsafe { storage.set_len(len) }; // SAFETY: the output allocation matches the immutable resident LDE. let status = unsafe { multi_stark_cuda_lde_copy_to_host( self.device_id, self.handle.as_ptr(), - values.as_mut_ptr().cast(), + storage.as_mut_ptr().cast(), ) }; check_cuda(status, "resident LDE host copy"); + let mut storage = core::mem::ManuallyDrop::new(storage); + // SAFETY: the successful FFI call initialized all elements, Goldilocks + // has the same layout inside `MaybeUninit`, and ownership moves exactly + // once into the returned vector. + let values = unsafe { + Vec::from_raw_parts( + storage.as_mut_ptr().cast::(), + storage.len(), + storage.capacity(), + ) + }; RowMajorMatrix::new(values, self.width) } @@ -569,6 +598,68 @@ fn encode_quotient_nodes( (nodes, slots, count as usize) } +/// Conservative device-memory requirement for one fused quotient job. +/// +/// The graph evaluator reuses slots as soon as their final consumer has run, +/// so `graph.nodes.len()` can be orders of magnitude larger than the live +/// device scratch. Keep this estimate beside the encoder so admission and the +/// kernel use the same liveness calculation. The scratch term assumes the +/// global-memory path; devices able to fit the slots in shared memory need +/// less than this bound. +pub(crate) fn quotient_lde_memory_upper_bound( + graph: &ConstraintGraph, + public_count: usize, + constraint_count: usize, + quotient_size: usize, + quotient_degree: usize, + log_blowup: usize, +) -> (usize, usize) { + let (nodes, _, slot_count) = encode_quotient_nodes(graph); + let lookup_arg_count = graph + .lookups + .iter() + .map(|lookup| lookup.args.len()) + .sum::(); + let align8 = |bytes: usize| bytes.saturating_add(7) & !7; + let allocation_bytes = [ + nodes.len().saturating_mul(size_of::()), + graph.zeros.len().saturating_mul(size_of::()), + graph + .lookups + .len() + .saturating_mul(size_of::()), + lookup_arg_count.saturating_mul(size_of::()), + public_count.saturating_mul(size_of::()), + 4usize + .saturating_mul(quotient_size) + .saturating_mul(size_of::()), + 2usize + .saturating_mul(constraint_count) + .saturating_mul(size_of::()), + 2 * size_of::(), + 2usize + .saturating_mul(quotient_size) + .saturating_mul(size_of::()), + ] + .into_iter() + .map(align8) + .sum::(); + let global_blocks = quotient_size.div_ceil(128).min(256); + let scratch_bytes = global_blocks + .saturating_mul(slot_count) + .saturating_mul(128) + .saturating_mul(size_of::()); + let trace_height = quotient_size / quotient_degree; + let lde_height = trace_height + .checked_shl(u32::try_from(log_blowup).expect("LDE blowup exceeds u32")) + .unwrap_or(usize::MAX); + let output_bytes = lde_height + .saturating_mul(2) + .saturating_mul(quotient_degree) + .saturating_mul(size_of::()); + (output_bytes, allocation_bytes.saturating_add(scratch_bytes)) +} + fn encode_lookup_nodes(graph: &ConstraintGraph) -> Option { let n = graph.lookup_prefix_len; if n == 0 || graph.lookups.is_empty() { @@ -696,6 +787,82 @@ fn encode_lookup_nodes(graph: &ConstraintGraph) -> Option, + height: usize, + main_width: usize, + group_size: usize, + log_blowup: usize, +) -> Option<(usize, usize)> { + let (nodes, slot_count, lookups, args) = encode_lookup_nodes(graph)?; + let lookup_count = lookups.len(); + let groups = lookup_count.div_ceil(group_size.max(1)); + let extended_height = height + .checked_shl(u32::try_from(log_blowup).expect("LDE blowup exceeds u32")) + .unwrap_or(usize::MAX); + let output_bytes = extended_height + .saturating_mul(groups) + .saturating_mul(2 * size_of::()); + + const LOOKUP_ROWS_PER_CHUNK: usize = 1 << 16; + let chunk_rows = height.min(LOOKUP_ROWS_PER_CHUNK); + let message_count = chunk_rows.saturating_mul(lookup_count); + let align8 = |bytes: usize| bytes.saturating_add(7) & !7; + let metadata_bytes = [ + nodes.len().saturating_mul(size_of::()), + lookups + .len() + .saturating_mul(size_of::()), + args.len().saturating_mul(size_of::()), + ] + .into_iter() + .map(align8) + .sum::(); + // Per message: conjugate Ext2, norm, inverse norm, multiplicity. + let message_bytes = message_count.saturating_mul(5 * size_of::()); + let delta_bytes = height + .saturating_mul(groups) + .saturating_mul(2 * size_of::()); + let shared_tile = (48 * 1024) / slot_count.saturating_mul(size_of::()).max(1); + let scratch_bytes = if shared_tile < 32 { + chunk_rows + .div_ceil(128) + .min(256) + .saturating_mul(slot_count) + .saturating_mul(128 * size_of::()) + } else { + 0 + }; + // Some stage-1 policies retain the original trace in pinned host memory. + // Charge one bounded staging chunk even when this particular LDE still + // owns a device trace, keeping admission safe across both representations. + let trace_chunk_bytes = chunk_rows + .saturating_add(1) + .saturating_mul(main_width) + .saturating_mul(size_of::()); + // Device-cached twiddles and shift powers may be cold for this height. + let constant_bytes = height + .saturating_div(2) + .saturating_add(height) + .saturating_add(extended_height / 2) + .saturating_mul(size_of::()); + Some(( + output_bytes, + metadata_bytes + .saturating_add(message_bytes) + .saturating_add(delta_bytes) + .saturating_add(scratch_bytes) + .saturating_add(trace_chunk_bytes) + .saturating_add(constant_bytes), + )) +} + /// Evaluates the compiled base-field constraint roots directly against /// resident trace LDEs. This is the protocol-independent core used by the /// CUDA quotient path; selectors and public values remain caller supplied. @@ -837,12 +1004,49 @@ pub(crate) fn quotient_values_resident( /// matrix. The returned storage has exactly the bit-reversed layout expected /// by `Pcs::commit_ldes`. #[allow(clippy::too_many_arguments)] -pub(crate) fn quotient_lde_resident( +pub(crate) fn quotient_lde_mixed( dft: &CudaDft, graph: &ConstraintGraph, - preprocessed: Option<&CudaLde>, - main: &CudaLde, - stage2: &CudaLde, + preprocessed: Option>, + main: mmcs::CudaMatrixSource<'_>, + stage2: mmcs::CudaMatrixSource<'_>, + publics: &[Goldilocks], + selectors: CudaCosetSelectors, + alpha: &[Goldilocks], + delta: &[Goldilocks; 2], + ext_w: Goldilocks, + quotient_size: usize, + next_step: usize, + group_size: usize, + quotient_degree: usize, + log_blowup: usize, +) -> CudaLde { + quotient_lde_sources( + dft, + graph, + preprocessed, + main, + stage2, + publics, + selectors, + alpha, + delta, + ext_w, + quotient_size, + next_step, + group_size, + quotient_degree, + log_blowup, + ) +} + +#[allow(clippy::too_many_arguments)] +fn quotient_lde_sources( + dft: &CudaDft, + graph: &ConstraintGraph, + preprocessed: Option>, + main: mmcs::CudaMatrixSource<'_>, + stage2: mmcs::CudaMatrixSource<'_>, publics: &[Goldilocks], selectors: CudaCosetSelectors, alpha: &[Goldilocks], @@ -897,45 +1101,121 @@ pub(crate) fn quotient_lde_resident( .take(quotient_degree) .map(|weight| weight * height_inverse) .collect(); + let source_parts = + |source: Option>| -> (*const c_void, *const u64, usize, usize) { + match source { + Some(mmcs::CudaMatrixSource::Resident(lde)) => ( + lde.raw_handle(), + core::ptr::null(), + lde.height(), + lde.width(), + ), + Some(mmcs::CudaMatrixSource::Host(matrix)) => ( + core::ptr::null(), + matrix.values.as_ptr().cast(), + matrix.height(), + matrix.width(), + ), + None => (core::ptr::null(), core::ptr::null(), 0, 0), + } + }; + let (preprocessed_handle, preprocessed_host, preprocessed_height, preprocessed_width) = + source_parts(preprocessed); + let (main_handle, main_host, main_height, main_width) = source_parts(Some(main)); + let (stage2_handle, stage2_host, stage2_height, stage2_width) = source_parts(Some(stage2)); + let mixed = !preprocessed_host.is_null() || !main_host.is_null() || !stage2_host.is_null(); let mut handle = core::ptr::null_mut(); - let status = unsafe { - multi_stark_cuda_quotient_lde( - dft.device_id, - &mut handle, - nodes.as_ptr().cast(), - nodes.len(), - slot_count, - roots.as_ptr(), - roots.len(), - lookups.as_ptr().cast(), - lookups.len(), - args.as_ptr(), - args.len(), - group_size, - preprocessed.map_or(core::ptr::null(), CudaLde::raw_handle), - main.raw_handle(), - stage2.raw_handle(), - publics.as_ptr().cast(), - publics.len(), - raw_u64(selectors.coset_shift), - raw_u64(selectors.coset_generator), - raw_u64(selectors.trace_last), - raw_u64(selectors.vanishing_start), - raw_u64(selectors.vanishing_step), - alpha.as_ptr().cast(), - constraint_count, - delta.as_ptr().cast(), - raw_u64(ext_w), - quotient_size, - next_step, - quotient_degree, - log_blowup, - quotient_twiddles.as_ptr().cast(), - lde_twiddles.as_ptr().cast(), - weights.as_ptr().cast(), - ) + let status = if mixed { + // SAFETY: every host matrix and resident handle remains borrowed for + // this synchronous call. The returned LDE owns its device storage. + unsafe { + multi_stark_cuda_quotient_lde_mixed( + dft.device_id, + &mut handle, + nodes.as_ptr().cast(), + nodes.len(), + slot_count, + roots.as_ptr(), + roots.len(), + lookups.as_ptr().cast(), + lookups.len(), + args.as_ptr(), + args.len(), + group_size, + preprocessed_handle, + preprocessed_host, + preprocessed_height, + preprocessed_width, + main_handle, + main_host, + main_height, + main_width, + stage2_handle, + stage2_host, + stage2_height, + stage2_width, + publics.as_ptr().cast(), + publics.len(), + raw_u64(selectors.coset_shift), + raw_u64(selectors.coset_generator), + raw_u64(selectors.trace_last), + raw_u64(selectors.vanishing_start), + raw_u64(selectors.vanishing_step), + alpha.as_ptr().cast(), + constraint_count, + delta.as_ptr().cast(), + raw_u64(ext_w), + quotient_size, + next_step, + quotient_degree, + log_blowup, + quotient_twiddles.as_ptr().cast(), + lde_twiddles.as_ptr().cast(), + weights.as_ptr().cast(), + ) + } + } else { + // SAFETY: all resident handles and input slices remain live for this + // synchronous call. The returned LDE owns its device storage. + unsafe { + multi_stark_cuda_quotient_lde( + dft.device_id, + &mut handle, + nodes.as_ptr().cast(), + nodes.len(), + slot_count, + roots.as_ptr(), + roots.len(), + lookups.as_ptr().cast(), + lookups.len(), + args.as_ptr(), + args.len(), + group_size, + preprocessed_handle, + main_handle, + stage2_handle, + publics.as_ptr().cast(), + publics.len(), + raw_u64(selectors.coset_shift), + raw_u64(selectors.coset_generator), + raw_u64(selectors.trace_last), + raw_u64(selectors.vanishing_start), + raw_u64(selectors.vanishing_step), + alpha.as_ptr().cast(), + constraint_count, + delta.as_ptr().cast(), + raw_u64(ext_w), + quotient_size, + next_step, + quotient_degree, + log_blowup, + quotient_twiddles.as_ptr().cast(), + lde_twiddles.as_ptr().cast(), + weights.as_ptr().cast(), + ) + } }; - check_cuda(status, "resident quotient LDE"); + check_cuda(status, "CUDA quotient LDE"); CudaLde { device_id: dft.device_id, handle: NonNull::new(handle).expect("CUDA returned a null quotient LDE"), @@ -946,9 +1226,22 @@ pub(crate) fn quotient_lde_resident( pub(crate) fn mixed_lde_open_row(ldes: &[CudaLde], index: usize) -> Vec> { assert!(!ldes.is_empty()); - assert!(index < ldes.iter().map(CudaLde::height).max().unwrap()); - let total: usize = ldes.iter().map(CudaLde::width).sum(); - let handles: Vec<_> = ldes.iter().map(CudaLde::raw_handle).collect(); + let max_height = ldes.iter().map(CudaLde::height).max().unwrap(); + mixed_lde_open_row_at_height(ldes.iter(), max_height, index) +} + +pub(crate) fn mixed_lde_open_row_at_height<'a>( + ldes: impl IntoIterator, + max_height: usize, + index: usize, +) -> Vec> { + let ldes = ldes.into_iter().collect::>(); + assert!(!ldes.is_empty()); + assert!(max_height.is_power_of_two()); + assert!(ldes.iter().all(|lde| lde.height() <= max_height)); + assert!(index < max_height); + let total: usize = ldes.iter().map(|lde| lde.width()).sum(); + let handles: Vec<_> = ldes.iter().map(|lde| lde.raw_handle()).collect(); let mut flat = Goldilocks::zero_vec(total); let status = unsafe { multi_stark_cuda_mixed_lde_open_row( @@ -956,6 +1249,7 @@ pub(crate) fn mixed_lde_open_row(ldes: &[CudaLde], index: usize) -> Vec( + ldes: impl IntoIterator, + max_height: usize, + indices: &[usize], +) -> Vec>> { + let ldes = ldes.into_iter().collect::>(); + assert!(!ldes.is_empty()); + assert!(max_height.is_power_of_two()); + assert!(ldes.iter().all(|lde| lde.height() <= max_height)); assert!(indices.iter().all(|&index| index < max_height)); - let total: usize = ldes.iter().map(CudaLde::width).sum(); - let handles: Vec<_> = ldes.iter().map(CudaLde::raw_handle).collect(); + let total: usize = ldes.iter().map(|lde| lde.width()).sum(); + let handles: Vec<_> = ldes.iter().map(|lde| lde.raw_handle()).collect(); let device_indices: Vec = indices .iter() .map(|&index| u64::try_from(index).expect("row index exceeds u64")) @@ -991,6 +1297,7 @@ pub(crate) fn mixed_lde_open_rows( flat.as_mut_ptr().cast(), handles.as_ptr(), handles.len(), + max_height, device_indices.as_ptr(), device_indices.len(), ) @@ -1065,6 +1372,19 @@ impl CudaReducedOpening { pub(crate) const fn height(&self) -> usize { self.height } + + pub(crate) fn add_host(&mut self, values: &[[Goldilocks; 2]]) { + assert_eq!(values.len(), self.height); + let status = unsafe { + multi_stark_cuda_reduced_add_host( + self.device_id, + self.handle.as_ptr(), + values.as_ptr().cast(), + values.len(), + ) + }; + check_cuda(status, "host reduced-opening accumulation") + } #[cfg(test)] pub(crate) fn add( &mut self, @@ -1201,6 +1521,7 @@ impl CudaFriWorkspace { check_cuda(status, "batched FRI interpolation"); output } + pub(crate) fn reduce( &mut self, tasks: &[CudaReductionTask], @@ -1309,6 +1630,194 @@ pub(crate) fn lookup_lde_resident( ) } +#[allow(clippy::too_many_arguments)] +pub(crate) fn lookup_lde_resident_partitioned( + dft: &CudaDft, + multiplicities: &[Goldilocks], + args: &[Goldilocks], + arg_offsets: &[usize], + height: usize, + num_lookups: usize, + group_size: usize, + beta: [Goldilocks; 2], + gamma: [Goldilocks; 2], + ext_w: Goldilocks, + log_blowup: usize, + cpu_deltas: impl Fn(core::ops::Range) -> Vec<[Goldilocks; 2]> + Sync, +) -> (CudaLde, [Goldilocks; 2]) { + assert!((1..=8).contains(&group_size)); + assert_eq!(arg_offsets.len(), num_lookups + 1); + assert_eq!(arg_offsets.first(), Some(&0)); + assert!(arg_offsets.windows(2).all(|pair| pair[0] <= pair[1])); + assert!(num_lookups != 0); + let slots = num_lookups.div_ceil(group_size); + let args_width = *arg_offsets.last().unwrap(); + assert_eq!(multiplicities.len(), height * num_lookups); + assert_eq!(args.len(), height * args_width); + let extended_height = height << log_blowup; + let inverse_twiddles = dft.twiddles(log2_strict_usize(height), true); + let shift_powers = dft.shift_powers(height, Goldilocks::GENERATOR); + let forward_twiddles = dft.twiddles(log2_strict_usize(extended_height), false); + let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(height) as u64); + let total_started = std::time::Instant::now(); + let create_started = std::time::Instant::now(); + let mut pending_handle = core::ptr::null_mut(); + let status = unsafe { + multi_stark_cuda_lookup_lde_begin_partitioned( + dft.device_id, + &mut pending_handle, + arg_offsets.as_ptr(), + height, + num_lookups, + args_width, + group_size, + log_blowup, + ) + }; + check_cuda(status, "create partitioned CUDA lookup"); + let create_elapsed = create_started.elapsed(); + let pending_handle = + NonNull::new(pending_handle).expect("CUDA returned a null partitioned lookup handle"); + struct PendingGuard { + device_id: i32, + handle: Option>, + } + impl Drop for PendingGuard { + fn drop(&mut self) { + if let Some(handle) = self.handle { + unsafe { + let _ = multi_stark_cuda_lookup_lde_cancel_partitioned( + self.device_id, + handle.as_ptr(), + ); + } + } + } + } + let mut pending = PendingGuard { + device_id: dft.device_id, + handle: Some(pending_handle), + }; + + const ROWS_PER_CHUNK: usize = 1 << 16; + let chunk_count = height.div_ceil(ROWS_PER_CHUNK); + let remaining = std::sync::Mutex::new((0usize, chunk_count)); + let (cpu_chunks, cpu_elapsed, gpu_chunks, gpu_elapsed) = std::thread::scope(|scope| { + let raw_pending = pending_handle.as_ptr() as usize; + let gpu_remaining = &remaining; + let device_id = dft.device_id; + let ext_w = raw_u64(ext_w); + let gpu_started = std::time::Instant::now(); + let gpu_worker = scope.spawn(move || { + let mut chunks = 0usize; + loop { + let chunk = { + let mut remaining = gpu_remaining.lock().expect("lookup scheduler poisoned"); + if remaining.0 == remaining.1 { + None + } else { + remaining.1 -= 1; + Some(remaining.1) + } + }; + let Some(chunk) = chunk else { break }; + let row_start = chunk * ROWS_PER_CHUNK; + let rows = (height - row_start).min(ROWS_PER_CHUNK); + let status = unsafe { + multi_stark_cuda_lookup_lde_gpu_rows_partitioned( + device_id, + raw_pending as *mut c_void, + multiplicities.as_ptr().cast(), + args.as_ptr().cast(), + row_start, + rows, + beta.as_ptr().cast(), + gamma.as_ptr().cast(), + ext_w, + ) + }; + check_cuda(status, "partitioned CUDA lookup chunk"); + chunks += 1; + } + (chunks, gpu_started.elapsed()) + }); + + let cpu_started = std::time::Instant::now(); + let mut cpu_chunks = 0usize; + loop { + let chunk = { + let mut remaining = remaining.lock().expect("lookup scheduler poisoned"); + if remaining.0 == remaining.1 { + None + } else { + let chunk = remaining.0; + remaining.0 += 1; + Some(chunk) + } + }; + let Some(chunk) = chunk else { break }; + let row_start = chunk * ROWS_PER_CHUNK; + let rows = (height - row_start).min(ROWS_PER_CHUNK); + let values = cpu_deltas(row_start..row_start + rows); + assert_eq!(values.len(), rows * slots); + let status = unsafe { + multi_stark_cuda_lookup_lde_cpu_rows_partitioned( + dft.device_id, + pending_handle.as_ptr(), + values.as_ptr().cast(), + row_start, + rows, + ) + }; + check_cuda(status, "upload partitioned CPU lookup chunk"); + cpu_chunks += 1; + } + let cpu_elapsed = cpu_started.elapsed(); + let (gpu_chunks, gpu_elapsed) = gpu_worker + .join() + .expect("partitioned GPU lookup worker panicked"); + (cpu_chunks, cpu_elapsed, gpu_chunks, gpu_elapsed) + }); + + let finish_started = std::time::Instant::now(); + let mut handle = core::ptr::null_mut(); + let mut tail = [Goldilocks::ZERO; 4]; + let status = unsafe { + multi_stark_cuda_lookup_lde_finish_partitioned( + dft.device_id, + pending_handle.as_ptr(), + &mut handle, + tail.as_mut_ptr().cast(), + inverse_twiddles.as_ptr().cast(), + shift_powers.as_ptr().cast(), + forward_twiddles.as_ptr().cast(), + raw_u64(height_inverse), + ) + }; + pending.handle = None; + check_cuda(status, "partitioned CUDA lookup finish"); + let result = ( + CudaLde { + device_id: dft.device_id, + handle: NonNull::new(handle).expect("CUDA returned a null lookup LDE"), + height: extended_height, + width: 2 * slots, + }, + [tail[0] + tail[2], tail[1] + tail[3]], + ); + if memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] cooperative lookup: cpu_chunks={cpu_chunks}/{chunk_count} cpu={:.3}s gpu_chunks={gpu_chunks}/{chunk_count} gpu={:.3}s create={:.3}s finish={:.3}s total={:.3}s", + cpu_elapsed.as_secs_f64(), + gpu_elapsed.as_secs_f64(), + create_elapsed.as_secs_f64(), + finish_started.elapsed().as_secs_f64(), + total_started.elapsed().as_secs_f64() + ); + } + result +} + #[allow(clippy::too_many_arguments)] pub(crate) fn lookup_graph_lde_resident( dft: &CudaDft, @@ -1525,6 +2034,10 @@ pub(crate) fn device_memory_info(device_id: i32) -> (usize, usize) { (free_bytes, total_bytes) } +pub(crate) fn memory_diagnostics_enabled() -> bool { + std::env::var_os("MULTI_STARK_CUDA_MEMORY_LOG").is_some() +} + fn check_cuda(status: i32, operation: &str) { if status == 0 { return; @@ -1713,6 +2226,67 @@ unsafe impl Send for CudaMixedMerkleTree {} unsafe impl Sync for CudaMixedMerkleTree {} impl CudaMixedMerkleTree { + #[must_use] + pub(crate) fn hash_hybrid_height_group( + device_id: i32, + ldes: &[Option<&CudaLde>], + host_matrices: &[Option<&RowMajorMatrix>], + ) -> Vec<[u8; 32]> { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + assert_eq!(ldes.len(), host_matrices.len()); + assert!(!ldes.is_empty(), "hybrid height group is empty"); + assert!( + ldes.iter().flatten().all(|lde| lde.device_id == device_id), + "hybrid height group spans CUDA devices" + ); + assert!( + ldes.iter() + .zip(host_matrices) + .all(|(lde, matrix)| lde.is_some() ^ matrix.is_some()), + "every hybrid height-group matrix must have exactly one storage backend" + ); + let heights = ldes + .iter() + .zip(host_matrices) + .map(|(lde, matrix)| lde.map_or_else(|| matrix.unwrap().height(), CudaLde::height)) + .collect::>(); + let height = heights[0]; + assert!( + heights.iter().all(|&candidate| candidate == height), + "hybrid height group contains unequal heights" + ); + let handles = ldes + .iter() + .map(|lde| lde.map_or(core::ptr::null(), CudaLde::raw_handle)) + .collect::>(); + let host_values = host_matrices + .iter() + .map(|matrix| matrix.map_or(core::ptr::null(), |matrix| matrix.values.as_ptr().cast())) + .collect::>(); + let widths = ldes + .iter() + .zip(host_matrices) + .map(|(lde, matrix)| lde.map_or_else(|| matrix.unwrap().width(), CudaLde::width)) + .collect::>(); + let mut digests = vec![[0u8; 32]; height]; + // SAFETY: every pointer is backed by a matrix or CUDA handle borrowed + // for this synchronous call, and `digests` has exactly `height` rows. + let status = unsafe { + multi_stark_cuda_hash_hybrid_height_group( + device_id, + digests.as_mut_ptr().cast(), + handles.as_ptr(), + host_values.as_ptr(), + widths.as_ptr(), + heights.as_ptr(), + handles.len(), + height, + ) + }; + check_cuda(status, "hybrid CUDA height-group hashing"); + digests + } + /// Builds a mixed-height tree from row groups ordered by descending, /// distinct power-of-two height. Each group contains already-concatenated /// rows for all matrices entering the MMCS at that height. @@ -1805,6 +2379,123 @@ impl CudaMixedMerkleTree { } } + /// Builds a mixed-height tree from resident CUDA LDEs, host matrices, and + /// pre-hashed height groups. Deferred dimensions are accepted only when a + /// pre-hashed group supplies that height's complete digest frontier. + #[must_use] + pub(crate) fn from_hybrid( + device_id: i32, + ldes: &[Option<&CudaLde>], + host_matrices: &[Option<&RowMajorMatrix>], + deferred_dimensions: &[Option], + host_digest_groups: &[(usize, Vec<[u8; 32]>)], + ) -> Self { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + assert_eq!(ldes.len(), host_matrices.len()); + assert_eq!(ldes.len(), deferred_dimensions.len()); + assert!(!ldes.is_empty(), "mixed Merkle tree has no LDE groups"); + assert!( + ldes.iter().flatten().all(|lde| lde.device_id == device_id), + "resident LDEs belong to different CUDA devices" + ); + for (height, digests) in host_digest_groups { + assert!( + height.is_power_of_two(), + "host digest height is not a power of two" + ); + assert_eq!( + *height, + digests.len(), + "host digest group has the wrong height" + ); + } + assert!( + ldes.iter().zip(host_matrices).zip(deferred_dimensions).all( + |((lde, matrix), deferred)| { + let sources = usize::from(lde.is_some()) + + usize::from(matrix.is_some()) + + usize::from(deferred.is_some()); + sources == 1 + && deferred.is_none_or(|dimensions| { + host_digest_groups + .iter() + .any(|(height, _)| *height == dimensions.height) + }) + } + ), + "every hybrid matrix needs one source, and deferred matrices need pre-hashed rows" + ); + let handles: Vec<*const c_void> = ldes + .iter() + .map(|lde| lde.map_or(core::ptr::null(), |lde| lde.handle.as_ptr().cast_const())) + .collect(); + let host_values = host_matrices + .iter() + .map(|matrix| { + matrix.map_or(core::ptr::null(), |matrix| { + matrix.values.as_ptr().cast::() + }) + }) + .collect::>(); + let widths = ldes + .iter() + .zip(host_matrices) + .zip(deferred_dimensions) + .map(|((lde, matrix), deferred)| { + lde.map_or_else( + || matrix.map_or_else(|| deferred.unwrap().width, Matrix::width), + CudaLde::width, + ) + }) + .collect::>(); + let heights: Vec = ldes + .iter() + .zip(host_matrices) + .zip(deferred_dimensions) + .map(|((lde, matrix), deferred)| { + lde.map_or_else( + || matrix.map_or_else(|| deferred.unwrap().height, Matrix::height), + CudaLde::height, + ) + }) + .collect(); + let host_digest_pointers: Vec<*const u8> = host_digest_groups + .iter() + .map(|(_, digests)| digests.as_ptr().cast()) + .collect(); + let host_digest_heights: Vec = host_digest_groups + .iter() + .map(|(height, _)| *height) + .collect(); + let mut handle = core::ptr::null_mut(); + let mut root = [0u8; 32]; + // SAFETY: all resident handles and host digest buffers remain live for + // this synchronous call; the returned tree owns its digest allocation. + let status = unsafe { + multi_stark_cuda_mixed_merkle_create_hybrid( + device_id, + &mut handle, + root.as_mut_ptr(), + handles.as_ptr(), + host_values.as_ptr(), + widths.as_ptr(), + heights.as_ptr(), + ldes.len(), + host_digest_pointers.as_ptr(), + host_digest_heights.as_ptr(), + host_digest_groups.len(), + ) + }; + check_cuda(status, "hybrid CPU/CUDA mixed-height Merkle tree creation"); + let row_count = heights.into_iter().max().unwrap(); + Self { + device_id, + handle: NonNull::new(handle).expect("CUDA returned a null mixed Merkle handle"), + root, + row_count, + } + } + pub(crate) fn from_fri_codeword(codeword: &CudaLde, arity: usize) -> Self { assert_eq!(codeword.width, 2); assert!(arity.is_power_of_two()); @@ -2059,6 +2750,7 @@ unsafe extern "C" { output: *mut u64, ) -> i32; fn multi_stark_cuda_lde_release_trace(device_id: i32, handle: *mut c_void) -> i32; + fn multi_stark_cuda_lde_release_values(device_id: i32, handle: *mut c_void) -> i32; fn multi_stark_cuda_lde_attach_trace( device_id: i32, handle: *mut c_void, @@ -2167,11 +2859,56 @@ unsafe extern "C" { lde_twiddles: *const u64, slice_weights: *const u64, ) -> i32; + fn multi_stark_cuda_quotient_lde_mixed( + device_id: i32, + output_handle: *mut *mut c_void, + nodes: *const c_void, + node_count: usize, + slot_count: usize, + roots: *const u32, + root_count: usize, + lookups: *const c_void, + lookup_count: usize, + lookup_args: *const u32, + lookup_arg_count: usize, + group_size: usize, + preprocessed_handle: *const c_void, + preprocessed_host: *const u64, + preprocessed_height: usize, + preprocessed_width: usize, + main_handle: *const c_void, + main_host: *const u64, + main_height: usize, + main_width: usize, + stage2_handle: *const c_void, + stage2_host: *const u64, + stage2_height: usize, + stage2_width: usize, + publics: *const u64, + public_count: usize, + coset_shift: u64, + coset_generator: u64, + trace_last: u64, + vanishing_start: u64, + vanishing_step: u64, + alpha: *const u64, + constraint_count: usize, + delta: *const u64, + ext_w: u64, + quotient_size: usize, + next_step: usize, + quotient_degree: usize, + log_blowup: usize, + quotient_twiddles: *const u64, + lde_twiddles: *const u64, + slice_weights: *const u64, + ) -> i32; fn multi_stark_cuda_mixed_lde_open_row( device_id: i32, output: *mut u64, handles: *const *const c_void, handle_count: usize, + max_height: usize, index: usize, ) -> i32; fn multi_stark_cuda_mixed_lde_open_rows( @@ -2179,6 +2916,7 @@ unsafe extern "C" { output: *mut u64, handles: *const *const c_void, handle_count: usize, + max_height: usize, indices: *const u64, query_count: usize, ) -> i32; @@ -2282,11 +3020,59 @@ unsafe extern "C" { forward_twiddles: *const u64, height_inverse: u64, ) -> i32; + fn multi_stark_cuda_lookup_lde_begin_partitioned( + device_id: i32, + pending_handle: *mut *mut c_void, + arg_offsets: *const usize, + height: usize, + num_lookups: usize, + args_width: usize, + group_size: usize, + added_bits: usize, + ) -> i32; + fn multi_stark_cuda_lookup_lde_gpu_rows_partitioned( + device_id: i32, + pending_handle: *mut c_void, + multiplicities: *const u64, + args: *const u64, + row_start: usize, + rows: usize, + beta: *const u64, + gamma: *const u64, + ext_w: u64, + ) -> i32; + fn multi_stark_cuda_lookup_lde_cpu_rows_partitioned( + device_id: i32, + pending_handle: *mut c_void, + deltas: *const u64, + row_start: usize, + rows: usize, + ) -> i32; + fn multi_stark_cuda_lookup_lde_finish_partitioned( + device_id: i32, + pending_handle: *mut c_void, + output_handle: *mut *mut c_void, + total: *mut u64, + inverse_twiddles: *const u64, + shift_powers: *const u64, + forward_twiddles: *const u64, + height_inverse: u64, + ) -> i32; + fn multi_stark_cuda_lookup_lde_cancel_partitioned( + device_id: i32, + pending_handle: *mut c_void, + ) -> i32; fn multi_stark_cuda_reduced_create( device_id: i32, handle: *mut *mut c_void, height: usize, ) -> i32; + fn multi_stark_cuda_reduced_add_host( + device_id: i32, + handle: *mut c_void, + values: *const u64, + height: usize, + ) -> i32; #[cfg(test)] fn multi_stark_cuda_reduced_add( device_id: i32, @@ -2382,6 +3168,29 @@ unsafe extern "C" { lde_handles: *const *const c_void, lde_count: usize, ) -> i32; + fn multi_stark_cuda_mixed_merkle_create_hybrid( + device_id: i32, + handle: *mut *mut c_void, + root: *mut u8, + lde_handles: *const *const c_void, + host_values: *const *const u64, + widths: *const usize, + heights: *const usize, + matrix_count: usize, + host_digest_groups: *const *const u8, + host_digest_heights: *const usize, + host_digest_group_count: usize, + ) -> i32; + fn multi_stark_cuda_hash_hybrid_height_group( + device_id: i32, + host_digests: *mut u8, + lde_handles: *const *const c_void, + host_values: *const *const u64, + widths: *const usize, + heights: *const usize, + matrix_count: usize, + height: usize, + ) -> i32; fn multi_stark_cuda_fri_merkle_create( device_id: i32, handle: *mut *mut c_void, @@ -2661,6 +3470,18 @@ mod tests { offset.as_basis_coefficients_slice().try_into().unwrap(), Goldilocks::from_u64(7), ); + let host_values: Vec<_> = (0..height) + .map(|row| { + [ + Goldilocks::from_usize(row + 11), + Goldilocks::from_usize(3 * row + 5), + ] + }) + .collect(); + got.add_host(&host_values); + for (expected, host) in expected.iter_mut().zip(&host_values) { + *expected += Ext::from_basis_coefficients_slice(host).unwrap(); + } for (got, want) in got.to_host().iter().zip(expected) { assert_eq!(got.as_slice(), want.as_basis_coefficients_slice()); } @@ -2962,9 +3783,22 @@ mod tests { let host_tree = CudaMixedMerkleTree::new(0, &level_refs); let direct_host_tree = CudaMixedMerkleTree::from_host_matrices(0, &host_ldes); let resident_tree = CudaMixedMerkleTree::from_ldes(0, &ldes); + let hybrid_host = vec![None, None, Some(host_ldes[2].clone())]; + let hybrid_digests = mmcs::hash_cpu_height_groups(&hybrid_host); + let hybrid_resident = [Some(&ldes[0]), Some(&ldes[1]), None]; + let hybrid_host_refs = [None, None, Some(&host_ldes[2])]; + let hybrid_tree = CudaMixedMerkleTree::from_hybrid( + 0, + &hybrid_resident, + &hybrid_host_refs, + &[None, None, None], + &hybrid_digests, + ); assert_eq!(resident_tree.root(), host_tree.root()); assert_eq!(direct_host_tree.root(), host_tree.root()); + assert_eq!(hybrid_tree.root(), host_tree.root()); assert_eq!(resident_tree.open_siblings(5), host_tree.open_siblings(5)); + assert_eq!(hybrid_tree.open_siblings(5), host_tree.open_siblings(5)); assert_eq!( direct_host_tree.open_siblings(5), host_tree.open_siblings(5) diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index 73a5265..250f24d 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -28,7 +28,7 @@ use std::vec::Vec; use itertools::{Itertools, izip}; use p3_challenger::{CanObserve, FieldChallenger, GrindingChallenger}; use p3_commit::{ExtensionMmcs, Mmcs, OpenedValues, Pcs, PeriodicLdeTable}; -use p3_dft::{Radix2DFTSmallBatch, TwoAdicSubgroupDft}; +use p3_dft::{Radix2DFTSmallBatch, Radix2DitParallel, TwoAdicSubgroupDft}; use p3_field::coset::TwoAdicMultiplicativeCoset; use p3_field::{ BasedVectorSpace, ExtensionField, PackedFieldExtension, PrimeCharacteristicRing, PrimeField64, @@ -43,11 +43,34 @@ use p3_util::linear_map::LinearMap; use p3_util::{log2_strict_usize, reverse_slice_index_bits}; use tracing::{debug_span, instrument}; -use super::mmcs::CudaCommitMmcs; +use super::mmcs::{CudaCommitMmcs, hash_host_only_height_groups}; use super::{CudaFriWorkspace, CudaLde, CudaMixedMerkleTree, CudaReducedOpening}; use p3_goldilocks::Goldilocks; use p3_symmetric::MerkleCap; +fn goldilocks_quadratic_inverse_denominators( + point: [Goldilocks; 2], + coset: &[Goldilocks], + extension_nonresidue: Goldilocks, +) -> Vec<[Goldilocks; 2]> { + let [point_0, point_1] = point; + let point_1_norm = extension_nonresidue * point_1 * point_1; + let norms = coset + .par_iter() + .map(|&x| { + let real = point_0 - x; + real * real - point_1_norm + }) + .collect::>(); + let inverse_norms = batch_multiplicative_inverse(&norms); + let inverse_point_1 = -point_1; + coset + .par_iter() + .zip(inverse_norms.par_iter()) + .map(|(&x, &inverse_norm)| [(point_0 - x) * inverse_norm, inverse_point_1 * inverse_norm]) + .collect() +} + pub trait CudaPcsDft: TwoAdicSubgroupDft { fn prepare_coset_lde_constants(&self, height: usize, added_bits: usize, shift: T); @@ -72,6 +95,97 @@ struct CudaFriRound { arity: usize, } +fn select_gpu_items( + items: &[(usize, u128)], + byte_budget: usize, + max_item_bytes: usize, +) -> (Vec, usize, u128) { + const PLACEMENT_GRANULARITY: usize = 8 << 20; + + #[derive(Clone, Copy)] + struct Node { + matrix: usize, + previous: Option, + } + + // Round every matrix up, making the discretized capacity conservative. + // Eight-MiB units keep placement overhead small even on large GPUs while + // wasting at most one unit per selected matrix. + let capacity = byte_budget / PLACEMENT_GRANULARITY; + let mut states = vec![None::<(u128, Option)>; capacity + 1]; + let mut nodes = Vec::::new(); + states[0] = Some((0, None)); + for (item, &(bytes, work)) in items.iter().enumerate() { + if bytes > max_item_bytes { + continue; + } + let units = bytes.div_ceil(PLACEMENT_GRANULARITY); + if units > capacity { + continue; + } + for used in (units..=capacity).rev() { + let Some((previous_work, previous_node)) = states[used - units] else { + continue; + }; + let candidate_work = previous_work.saturating_add(work); + if states[used].is_some_and(|(current_work, _)| current_work >= candidate_work) { + continue; + } + let node = nodes.len(); + nodes.push(Node { + matrix: item, + previous: previous_node, + }); + states[used] = Some((candidate_work, Some(node))); + } + } + + let (_, (gpu_work, mut node)) = states + .into_iter() + .enumerate() + .filter_map(|(used, state)| state.map(|state| (used, state))) + .max_by_key(|&(used, (work, _))| (work, core::cmp::Reverse(used))) + .unwrap(); + let mut selected = vec![false; items.len()]; + while let Some(index) = node { + let choice = nodes[index]; + selected[choice.matrix] = true; + node = choice.previous; + } + let gpu_bytes = items + .iter() + .zip(&selected) + .filter_map(|(&(bytes, _), &selected)| selected.then_some(bytes)) + .sum(); + (selected, gpu_bytes, gpu_work) +} + +type CosetLdeJob = (usize, (TwoAdicMultiplicativeCoset, RowMajorMatrix)); + +fn cpu_coset_lde_jobs( + jobs: Vec>, + log_blowup: usize, +) -> Vec<(usize, RowMajorMatrix)> +where + F: TwoAdicField + PrimeField64, + Radix2DitParallel: TwoAdicSubgroupDft, +{ + let cpu = Radix2DitParallel::::default(); + jobs.into_par_iter() + .map(|(index, (domain, evaluations))| { + let shift = F::GENERATOR / domain.shift(); + let mut lde = cpu + .coset_lde_batch(evaluations, log_blowup, shift) + .bit_reverse_rows() + .to_row_major_matrix(); + lde.values.par_iter_mut().for_each(|value| { + *value = F::from_u64(value.as_canonical_u64()); + }); + (index, lde) + }) + .collect() +} + trait CudaFriMmcs: Mmcs { fn commit_cuda_fri( &self, @@ -163,7 +277,7 @@ fn prove_fri_cuda_resident( ext_w: Goldilocks, ) -> FriProof>> where - Val: TwoAdicField + PrimeField64, + Val: TwoAdicField + PrimeField64 + 'static, Challenge: ExtensionField, InputMmcs: Mmcs, FriMmcs: CudaFriMmcs, @@ -434,6 +548,395 @@ where .sum::(); let dft = &self.dft; let log_blowup = self.fri.log_blowup; + let (initial_free, total_bytes) = + crate::cuda::device_memory_info(self.mmcs.cuda_device_id()); + let minimum_free = std::env::var("MULTI_STARK_CUDA_MIN_FREE_BYTES") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(total_bytes / 4); + let source_bytes = source_cells.saturating_mul(size_of::()); + let lde_bytes = source_bytes + .checked_shl(u32::try_from(log_blowup).expect("LDE blowup exceeds u32")) + .unwrap_or(usize::MAX); + let release_traces_during_construction = source_bytes + .saturating_add(lde_bytes) + .saturating_add(minimum_free) + > initial_free; + if lde_bytes.saturating_add(minimum_free) > initial_free { + let max_lde_height = evaluations + .iter() + .map(|(_, matrix)| matrix.height() << log_blowup) + .max() + .unwrap(); + // A hybrid tree stores all binary digest layers plus one injected + // digest frontier. Reserve these explicitly before assigning any + // LDE height group to device memory. + let tree_workspace_bytes = max_lde_height.saturating_mul(96); + let gpu_lde_budget = initial_free + .saturating_sub(minimum_free) + .saturating_sub(tree_workspace_bytes); + let matrix_resources = evaluations + .iter() + .map(|(_, matrix)| { + let lde_cells = (matrix.height() * matrix.width()) << log_blowup; + let bytes = lde_cells.saturating_mul(size_of::()); + let work = u128::try_from(lde_cells).expect("LDE cells exceed u128") + * (u128::from(matrix.height().trailing_zeros()) + + u128::try_from(log_blowup).expect("log blowup exceeds u128") + + 1); + (bytes, work) + }) + .collect_vec(); + let total_work = matrix_resources.iter().map(|(_, work)| *work).sum::(); + let mut height_groups = std::collections::BTreeMap::::new(); + let mut height_indices = std::collections::BTreeMap::>::new(); + for (index, ((_, matrix), &(bytes, work))) in + evaluations.iter().zip(&matrix_resources).enumerate() + { + let group = height_groups.entry(matrix.height()).or_default(); + group.0 = group.0.saturating_add(bytes); + group.1 = group.1.saturating_add(work); + height_indices + .entry(matrix.height()) + .or_default() + .push(index); + } + // Split allocatable device capacity between persistent LDEs and + // later phase workspace. This prevents stage one from filling VRAM + // with values which lookup immediately has to copy back and evict. + let durable_budget = gpu_lde_budget / 2; + // A Merkle leaf combines every matrix at a given height. Keeping + // height groups intact avoids streaming the CPU half of a split + // group across PCIe merely to hash it beside resident matrices. + // Matrix-level hybrid hashing remains available to later stages, + // where the output partition is intrinsic rather than introduced + // by this scheduler. + let grouped_resources = height_groups.values().copied().collect_vec(); + let (selected_groups, durable_bytes, durable_work) = + select_gpu_items(&grouped_resources, durable_budget, durable_budget); + let durable_heights = height_groups + .keys() + .zip(selected_groups) + .filter_map(|(&height, selected)| selected.then_some(height)) + .collect::>(); + let durable_matrices = evaluations + .iter() + .map(|(_, matrix)| durable_heights.contains(&matrix.height())) + .collect_vec(); + // A group too large to retain can still borrow the otherwise-idle + // GPU lane during commitment. Its selected matrices are transformed + // temporarily for hashing, then recomputed by the CPU for durable + // host storage while CUDA moves on to the retained groups. + let transient_reserve = max_lde_height.saturating_mul(32).saturating_add(64 << 20); + let select_transient_plan = |transient_budget| { + height_indices + .iter() + .filter(|(height, _)| !durable_heights.contains(height)) + .filter_map(|(&height, indices)| { + let resources = indices + .iter() + .map(|&index| { + let (lde_bytes, work) = matrix_resources[index]; + let trace_bytes = lde_bytes >> log_blowup; + (lde_bytes.saturating_add(trace_bytes), work) + }) + .collect_vec(); + let (selected, bytes, work) = + select_gpu_items(&resources, transient_budget, transient_budget); + let selected = indices + .iter() + .zip(selected) + .filter_map(|(&index, selected)| selected.then_some(index)) + .collect_vec(); + (!selected.is_empty()).then_some((height, selected, bytes, work)) + }) + .max_by_key(|(_, _, _, work)| *work) + }; + + // The CUDA constant cache is persistent and populated lazily. A + // free-memory snapshot taken before preparing the selected DFTs + // therefore overstates the capacity available to their LDEs. Keep + // replanning until every constant required by the resulting plan + // is resident and reflected in the memory snapshot. This is most + // visible on very wide outer proofs, where a stale snapshot can + // otherwise leave no room for the Merkle digest frontier. + let mut transient_budget = initial_free.saturating_sub(transient_reserve); + let mut transient_plan = select_transient_plan(transient_budget); + let mut prepared_constants = std::collections::BTreeSet::new(); + loop { + let transient_indices: std::collections::BTreeSet = transient_plan + .as_ref() + .map(|(_, indices, _, _)| indices.iter().copied().collect()) + .unwrap_or_default(); + for (index, (domain, evals)) in evaluations.iter().enumerate() { + if !durable_matrices[index] && !transient_indices.contains(&index) { + continue; + } + let shift = Val::GENERATOR / domain.shift(); + if prepared_constants.insert((evals.height(), shift.as_canonical_u64())) { + dft.prepare_coset_lde_constants(evals.height(), log_blowup, shift); + } + } + let (prepared_free, _) = + crate::cuda::device_memory_info(self.mmcs.cuda_device_id()); + transient_budget = prepared_free.saturating_sub(transient_reserve); + let prepared_plan = select_transient_plan(transient_budget); + if prepared_plan == transient_plan { + break; + } + transient_plan = prepared_plan; + } + let transient_height = transient_plan.as_ref().map(|(height, _, _, _)| *height); + let transient_matrices: std::collections::BTreeSet = transient_plan + .as_ref() + .map(|(_, indices, _, _)| indices.iter().copied().collect()) + .unwrap_or_default(); + let durable_count = durable_matrices + .iter() + .filter(|&&selected| selected) + .count(); + if crate::cuda::memory_diagnostics_enabled() { + let gpu_percent = durable_work + .saturating_mul(100) + .checked_div(total_work) + .unwrap_or(0); + eprintln!( + "[multi-stark/cuda] stage1 placement: durable={durable_count}/{} heights={}/{} gpu_work={durable_work}/{total_work} ({gpu_percent}%) durable_bytes={durable_bytes}/{durable_budget} device_budget={gpu_lde_budget}", + evaluations.len(), + durable_heights.len(), + height_groups.len(), + ); + if let Some((height, indices, bytes, work)) = &transient_plan { + eprintln!( + "[multi-stark/cuda] stage1 transient group: height={height} matrices={indices:?} footprint={bytes}/{transient_budget} work={work}" + ); + } + for (index, ((_, matrix), &(bytes, work))) in + evaluations.iter().zip(&matrix_resources).enumerate() + { + if !durable_matrices[index] { + eprintln!( + "[multi-stark/cuda] stage1 host matrix {index}: height={} width={} lde_bytes={bytes} work={work}", + matrix.height(), + matrix.width(), + ); + } + } + } + if durable_count != 0 { + let matrix_count = evaluations.len(); + let mut durable_jobs = Vec::with_capacity(durable_count); + let mut transient_jobs = Vec::new(); + let mut transient_aux_jobs = Vec::new(); + let mut remaining_cpu_jobs = Vec::new(); + for (index, evaluation) in evaluations.into_iter().enumerate() { + if durable_matrices[index] { + durable_jobs.push((index, evaluation)); + } else if transient_matrices.contains(&index) { + transient_jobs.push((index, evaluation)); + } else if transient_height == Some(evaluation.1.height()) { + transient_aux_jobs.push((index, evaluation)); + } else { + remaining_cpu_jobs.push((index, evaluation)); + } + } + // All host LDE work shares one bounded pool. Keeping the + // urgent same-height matrices, deferred materialization, and + // remaining CPU matrices on separate Rayon pools can + // oversubscribe a many-core host early, then strand cores once + // only a large deferred matrix remains. + // A large durable set must be evicted in later stages, which + // naturally gives deferred CPU transforms more overlap. Keep + // those transforms narrower so they do not starve the GPU's + // host-memory traffic. When little remains resident, use more + // cores to prevent a deferred transform becoming the quotient + // critical path. + let prioritize_deferred = durable_bytes < total_bytes / 4; + let stage1_threads = crate::types::cuda_stage1_worker_count(prioritize_deferred); + let stage1_pool = + std::sync::Arc::new(crate::types::cuda_host_pool("stage1-lde", stage1_threads)); + let cpu_started = std::time::Instant::now(); + let (aux_results, transient_results) = std::thread::scope(|scope| { + let aux_pool = std::sync::Arc::clone(&stage1_pool); + let cpu_task = scope.spawn(move || { + aux_pool.install(|| cpu_coset_lde_jobs(transient_aux_jobs, log_blowup)) + }); + let transient_started = std::time::Instant::now(); + let transient_results = transient_jobs + .into_iter() + .map(|(index, (domain, evals))| { + let shift = Val::GENERATOR / domain.shift(); + let lde = dft.coset_lde_batch_resident(&evals, log_blowup, shift); + (index, lde, (domain, evals)) + }) + .collect_vec(); + if crate::cuda::memory_diagnostics_enabled() && !transient_results.is_empty() { + eprintln!( + "[multi-stark/cuda] stage1 transient GPU LDE: {:.3}s", + transient_started.elapsed().as_secs_f64() + ); + } + (cpu_task.join().unwrap(), transient_results) + }); + + let mut host_matrices = (0..matrix_count).map(|_| None).collect_vec(); + for (index, lde) in aux_results { + host_matrices[index] = Some(lde); + } + let mut transient_ldes = (0..matrix_count).map(|_| None).collect_vec(); + let mut deferred_matrices = (0..matrix_count).map(|_| None).collect_vec(); + for (index, lde, (domain, evaluations)) in transient_results { + let dimensions = p3_matrix::Dimensions { + width: evaluations.width(), + height: evaluations.height() << log_blowup, + }; + let pool = std::sync::Arc::clone(&stage1_pool); + let worker = std::thread::spawn(move || { + pool.install(|| { + cpu_coset_lde_jobs(vec![(index, (domain, evaluations))], log_blowup) + .pop() + .expect("deferred LDE worker returned no matrix") + .1 + }) + }); + transient_ldes[index] = Some(lde); + deferred_matrices[index] = Some((dimensions, worker)); + } + if crate::cuda::memory_diagnostics_enabled() { + let deferred_count = deferred_matrices + .iter() + .filter(|matrix| matrix.is_some()) + .count(); + eprintln!( + "[multi-stark/cuda] stage1 deferred CPU LDE workers: {deferred_count} matrices on {stage1_threads} threads" + ); + } + let mut host_digest_groups = Vec::new(); + let mut prehashed_heights = std::collections::BTreeSet::new(); + + let (cpu_results, durable_results) = std::thread::scope(|scope| { + let remaining_pool = std::sync::Arc::clone(&stage1_pool); + let cpu_task = scope.spawn(move || { + remaining_pool + .install(|| cpu_coset_lde_jobs(remaining_cpu_jobs, log_blowup)) + }); + if let Some(height) = transient_height { + let indices = &height_indices[&height]; + let resident_refs = indices + .iter() + .map(|&index| transient_ldes[index].as_ref()) + .collect_vec(); + let host_refs = indices + .iter() + .map(|&index| { + transient_ldes[index] + .is_none() + .then(|| host_matrices[index].as_ref().unwrap()) + }) + .collect_vec(); + let hash_started = std::time::Instant::now(); + let digests = self + .mmcs + .hash_cuda_hybrid_height_group(&resident_refs, &host_refs); + let lde_height = height << log_blowup; + prehashed_heights.insert(lde_height); + host_digest_groups.push((lde_height, digests)); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] stage1 transient group hash: {:.3}s", + hash_started.elapsed().as_secs_f64() + ); + } + } + // The transient LDEs have supplied their digest frontier. + // Drop them before admitting the durable resident set. + drop(transient_ldes); + let durable_started = std::time::Instant::now(); + let durable_results = durable_jobs + .into_iter() + .map(|(index, (domain, evals))| { + let shift = Val::GENERATOR / domain.shift(); + let lde = dft.coset_lde_batch_resident(&evals, log_blowup, shift); + (index, lde, evals) + }) + .collect_vec(); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] stage1 durable GPU LDE: {:.3}s", + durable_started.elapsed().as_secs_f64() + ); + } + (cpu_task.join().unwrap(), durable_results) + }); + for (index, lde) in cpu_results { + host_matrices[index] = Some(lde); + } + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] stage1 CPU LDE: {:.3}s", + cpu_started.elapsed().as_secs_f64() + ); + } + + let mut resident = (0..matrix_count).map(|_| None).collect_vec(); + let mut retained_traces = (0..matrix_count).map(|_| None).collect_vec(); + for (index, lde, trace) in durable_results { + resident[index] = Some(lde); + retained_traces[index] = Some(trace); + } + let digest_started = std::time::Instant::now(); + let deferred_dimensions = deferred_matrices + .iter() + .map(|matrix| matrix.as_ref().map(|(dimensions, _)| *dimensions)) + .collect_vec(); + host_digest_groups.extend(hash_host_only_height_groups( + &host_matrices, + &resident, + &deferred_dimensions, + &prehashed_heights, + )); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] stage1 partitioned host digests: {:.3}s", + digest_started.elapsed().as_secs_f64() + ); + } + let tree_started = std::time::Instant::now(); + let committed = self.mmcs.commit_cuda_hybrid( + resident, + host_matrices, + deferred_matrices, + retained_traces, + host_digest_groups, + ); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] stage1 hybrid tree: {:.3}s", + tree_started.elapsed().as_secs_f64() + ); + } + return committed; + } + let cpu = Radix2DitParallel::::default(); + let ldes = tracing::info_span!("cuda/cpu_lde_fallback").in_scope(|| { + evaluations + .into_par_iter() + .map(|(domain, evals)| { + let shift = Val::GENERATOR / domain.shift(); + let mut lde = cpu + .coset_lde_batch(evals, log_blowup, shift) + .bit_reverse_rows() + .to_row_major_matrix(); + lde.values.par_iter_mut().for_each(|value| { + *value = Val::from_u64(value.as_canonical_u64()); + }); + lde + }) + .collect() + }); + return tracing::info_span!("cuda/cpu_mmcs_fallback") + .in_scope(|| self.mmcs.commit_cpu_storage(ldes)); + } if source_cells >= 10_000_000 { for (domain, evals) in &evaluations { dft.prepare_coset_lde_constants( @@ -448,7 +951,12 @@ where // reliable concurrency range while retaining substantial overlap. const CUDA_LDE_WAVE: usize = 16; let mut ldes = Vec::with_capacity(evaluations.len()); - for wave in evaluations.chunks(CUDA_LDE_WAVE) { + let wave_size = if release_traces_during_construction { + 1 + } else { + CUDA_LDE_WAVE + }; + for wave in evaluations.chunks(wave_size) { let transform = |(domain, evals): &(TwoAdicMultiplicativeCoset, RowMajorMatrix)| { assert_eq!(domain.size(), evals.height()); @@ -463,23 +971,26 @@ where } else { wave.par_iter().map(transform).collect() }; + if release_traces_during_construction { + for lde in &wave_ldes { + // SAFETY: this wave has completed synchronously. The host + // matrices remain owned below and are retained for every + // later consumer of the original trace. + unsafe { lde.release_trace() }; + } + } ldes.extend(wave_ldes); } // Preserve enough device headroom for lookup, quotient, and FRI // allocations. Spill the largest retained traces first, deriving the // policy from this device rather than a 96-GiB development machine. let (free_bytes, total_bytes) = crate::cuda::device_memory_info(self.mmcs.cuda_device_id()); - let minimum_free = std::env::var("MULTI_STARK_CUDA_MIN_FREE_BYTES") - .ok() - .and_then(|value| value.parse().ok()) - .unwrap_or(total_bytes / 4); // A free-memory snapshot immediately after commitment construction is // not enough: lookup construction temporarily needs buffers // proportional to the original trace. Proactively spill large traces // once this commitment itself is a meaningful fraction of VRAM. The // ratios retain the policy that was validated on a 96-GiB device while // scaling it to smaller cards. - let source_bytes = source_cells.saturating_mul(size_of::()); let proactive_spill = source_bytes >= total_bytes / 12; let large_matrix = total_bytes / 120; let mut projected_free = free_bytes; @@ -497,8 +1008,11 @@ where }) .collect_vec(); by_size.sort_unstable_by_key(|&(_, bytes)| core::cmp::Reverse(bytes)); - let mut spilled = false; + let mut spilled = release_traces_during_construction; for (index, bytes) in by_size { + if release_traces_during_construction { + continue; + } let needs_headroom = projected_free < minimum_free; let crowds_later_stages = proactive_spill && bytes >= large_matrix; if !needs_headroom && !crowds_later_stages { @@ -510,8 +1024,15 @@ where projected_free = projected_free.saturating_add(bytes); spilled = true; } - // Commit to the bit-reversed LDEs. - let (commitment, mut data) = self.mmcs.commit_cuda_resident(ldes); + // Commit to the bit-reversed LDEs. Once host traces have been retained, + // keep the LDEs individually spillable as well: later quotient and FRI + // admission may need to trade a resident LDE for its host materialization. + // Small all-resident commitments retain the lower-overhead fast path. + let (commitment, mut data) = if spilled { + self.mmcs.commit_cuda_spillable(ldes) + } else { + self.mmcs.commit_cuda_resident(ldes) + }; if spilled { self.mmcs.retain_matrices( &mut data, @@ -659,64 +1180,269 @@ where // Keep CUDA commitments resident through barycentric interpolation // and construction of the reduced FRI codewords. Only the opened // values and final extension codewords cross back to the host. - let resident_rounds = debug_span!("cuda prepare resident rounds").in_scope(|| { - commitment_data_with_opening_points + if commitment_data_with_opening_points + .iter() + .all(|(data, _)| self.mmcs.is_cuda_resident(data)) + { + let resident_rounds = debug_span!("cuda prepare resident rounds").in_scope(|| { + commitment_data_with_opening_points + .iter() + .map(|(data, points)| (self.mmcs.resident_or_upload(data), points)) + .collect_vec() + }); + let resident_max_height = resident_rounds .iter() - .map(|(data, points)| (self.mmcs.resident_or_upload(data), points)) - .collect_vec() - }); - let resident_max_height = resident_rounds + .flat_map(|(ldes, _)| ldes.iter()) + .map(CudaLde::height) + .max() + .unwrap_or(0); + let final_fri_height = self.fri.blowup() * self.fri.final_poly_len(); + if resident_max_height > 1024 && resident_max_height > final_fri_height { + let _resident_guard = debug_span!("cuda resident fri").entered(); + let rounds = resident_rounds; + let device_id = self.mmcs.cuda_device_id(); + assert_eq!(>::DIMENSION, 2); + let to_gold = |v: Val| Goldilocks::from_u64(v.as_canonical_u64()); + let to_pair = |v: Challenge| { + let c = >::as_basis_coefficients_slice(&v); + [to_gold(c[0]), to_gold(c[1])] + }; + let from_pair = |v: [Goldilocks; 2]| { + Challenge::from_basis_coefficients_slice(&[ + Val::from_u64(v[0].as_canonical_u64()), + Val::from_u64(v[1].as_canonical_u64()), + ]) + .unwrap() + }; + let global_max_height = rounds + .iter() + .flat_map(|(ldes, _)| ldes.iter().map(CudaLde::height)) + .max() + .unwrap(); + let global_max_width = rounds + .iter() + .flat_map(|(ldes, _)| ldes.iter().map(CudaLde::width)) + .max() + .unwrap(); + let log_global_max_height = log2_strict_usize(global_max_height); + let coset_domain = + TwoAdicMultiplicativeCoset::new(Val::GENERATOR, log_global_max_height).unwrap(); + let mut coset: Vec = coset_domain.iter().collect(); + reverse_slice_index_bits(&mut coset); + let mut max_log: LinearMap = LinearMap::new(); + for (ldes, points) in &rounds { + for (lde, ps) in ldes.iter().zip(points.iter()) { + for &z in ps { + if let Some(h) = max_log.get_mut(&z) { + *h = (*h).max(log2_strict_usize(lde.height())) + } else { + max_log.insert(z, log2_strict_usize(lde.height())); + } + } + } + } + let ext_x = + Challenge::from_basis_coefficients_slice(&[Val::ZERO, Val::ONE]).unwrap(); + let ext_w = to_pair(ext_x * ext_x)[0]; + assert_eq!(ext_w, Goldilocks::from_u64(7)); + let coset_gold: Vec<_> = coset.iter().copied().map(to_gold).collect(); + let mut inv_offsets = LinearMap::new(); + let mut inverse_points = Vec::new(); + let mut inverse_counts = Vec::new(); + let mut inverse_count = 0usize; + for (point, lh) in max_log { + inv_offsets.insert(point, inverse_count); + inverse_points.push(to_pair(point)); + let count = 1usize << lh; + inverse_counts.push(count); + inverse_count += count; + } + let mut workspace = CudaFriWorkspace::new( + device_id, + &inverse_points, + &inverse_counts, + &coset_gold, + ext_w, + ); + let mut interpolation_tasks = Vec::new(); + let mut output_count = 0usize; + let layouts = rounds + .iter() + .map(|(ldes, points)| { + ldes.iter() + .zip(points.iter()) + .map(|(lde, ps)| { + let h = lde.height() >> self.fri.log_blowup; + let lh = log2_strict_usize(h); + ps.iter() + .map(|&point| { + let offset = output_count; + output_count += lde.width(); + let shift_pow = Val::GENERATOR.exp_power_of_2(lh); + let scale = (point.exp_power_of_2(lh) - shift_pow) + * (Val::from_usize(h) * shift_pow).inverse(); + interpolation_tasks.push(lde.interpolation_task( + h, + *inv_offsets.get(&point).unwrap(), + offset, + to_pair(scale), + )); + (offset, lde.width()) + }) + .collect_vec() + }) + .collect_vec() + }) + .collect_vec(); + let interpolated = debug_span!("cuda interpolate openings") + .in_scope(|| workspace.interpolate(&interpolation_tasks, output_count, ext_w)); + let all_opened_values = layouts + .into_iter() + .map(|round| { + round + .into_iter() + .map(|matrix| { + matrix + .into_iter() + .map(|(offset, width)| { + interpolated[offset..offset + width] + .iter() + .copied() + .map(from_pair) + .collect_vec() + }) + .collect_vec() + }) + .collect_vec() + }) + .collect_vec(); + for round in &all_opened_values { + for matrix in round { + for values in matrix { + challenger.observe_algebra_slice(values); + } + } + } + let alpha: Challenge = challenger.sample_algebra_element(); + let alpha_powers: Vec<_> = alpha.powers().take(global_max_width).collect(); + let alpha_pairs: Vec<_> = alpha_powers.iter().copied().map(to_pair).collect(); + let mut num_reduced = [0usize; 33]; + let mut reduced: [Option; 33] = core::array::from_fn(|_| None); + let mut reduction_tasks = Vec::new(); + for ((ldes, points), openings_round) in rounds.iter().zip(all_opened_values.iter()) + { + for ((lde, ps), openings) in + ldes.iter().zip(points.iter()).zip(openings_round.iter()) + { + let lh = log2_strict_usize(lde.height()); + let target = reduced[lh].get_or_insert_with(|| { + CudaReducedOpening::new(device_id, lde.height()) + }); + for (&point, ys) in ps.iter().zip(openings.iter()) { + let reduced_y = + dot_product(alpha_powers.iter().copied(), ys.iter().copied()); + let offset = alpha.exp_u64(num_reduced[lh] as u64); + reduction_tasks.push(target.reduction_task( + lde, + *inv_offsets.get(&point).unwrap(), + to_pair(reduced_y), + to_pair(offset), + )); + num_reduced[lh] += lde.width(); + } + } + } + debug_span!("cuda reduce openings") + .in_scope(|| workspace.reduce(&reduction_tasks, &alpha_pairs, ext_w)); + let fri_input = reduced.into_iter().rev().flatten().collect_vec(); + let fri_proof = debug_span!("cuda prove fri").in_scope(|| { + prove_fri_cuda_resident( + &self.fri, + fri_input, + challenger, + log_global_max_height, + &commitment_data_with_opening_points, + &self.mmcs, + ext_w, + ) + }); + return (all_opened_values, fri_proof); + } + } + + // A commitment which did not fit in VRAM is stored by the CPU MMCS, + // but that does not require the entire opening to fall back to the + // CPU. Upload one such matrix at a time for interpolation and + // reduction. Resident commitments stay in place, and temporary + // uploads are dropped before the next matrix, bounding additional + // device memory by the largest CPU-backed matrix instead of the sum + // of all commitments. + let dimensions = commitment_data_with_opening_points .iter() - .flat_map(|(ldes, _)| ldes.iter()) - .map(CudaLde::height) - .max() - .unwrap_or(0); + .map(|(data, points)| { + let dimensions = self.mmcs.matrix_dimensions(data); + assert_eq!( + dimensions.len(), + points.len(), + "each matrix should have a corresponding set of evaluation points" + ); + dimensions + }) + .collect_vec(); + let opening_points = commitment_data_with_opening_points + .iter() + .map(|(_, points)| points) + .collect_vec(); + let (cuda_max_height, cuda_max_width) = dimensions + .iter() + .flatten() + .map(|dims| (dims.height, dims.width)) + .reduce(|(hmax, wmax), (height, width)| (hmax.max(height), wmax.max(width))) + .expect("No Matrices Supplied?"); let final_fri_height = self.fri.blowup() * self.fri.final_poly_len(); - if resident_max_height > 1024 && resident_max_height > final_fri_height { - let _resident_guard = debug_span!("cuda resident fri").entered(); - let rounds = resident_rounds; + if cuda_max_height > 1024 && cuda_max_height > final_fri_height { + let _resident_guard = debug_span!("cuda streamed fri").entered(); + let phase_started = std::time::Instant::now(); let device_id = self.mmcs.cuda_device_id(); assert_eq!(>::DIMENSION, 2); - let to_gold = |v: Val| Goldilocks::from_u64(v.as_canonical_u64()); - let to_pair = |v: Challenge| { - let c = >::as_basis_coefficients_slice(&v); - [to_gold(c[0]), to_gold(c[1])] + let to_gold = |value: Val| Goldilocks::from_u64(value.as_canonical_u64()); + let to_pair = |value: Challenge| { + let coefficients = + >::as_basis_coefficients_slice(&value); + [to_gold(coefficients[0]), to_gold(coefficients[1])] }; - let from_pair = |v: [Goldilocks; 2]| { + let from_pair = |value: [Goldilocks; 2]| { Challenge::from_basis_coefficients_slice(&[ - Val::from_u64(v[0].as_canonical_u64()), - Val::from_u64(v[1].as_canonical_u64()), + Val::from_u64(value[0].as_canonical_u64()), + Val::from_u64(value[1].as_canonical_u64()), ]) - .unwrap() + .expect("quadratic extension element") }; - let global_max_height = rounds - .iter() - .flat_map(|(ldes, _)| ldes.iter().map(CudaLde::height)) - .max() - .unwrap(); - let global_max_width = rounds - .iter() - .flat_map(|(ldes, _)| ldes.iter().map(CudaLde::width)) - .max() - .unwrap(); - let log_global_max_height = log2_strict_usize(global_max_height); + let log_global_max_height = log2_strict_usize(cuda_max_height); let coset_domain = TwoAdicMultiplicativeCoset::new(Val::GENERATOR, log_global_max_height).unwrap(); let mut coset: Vec = coset_domain.iter().collect(); reverse_slice_index_bits(&mut coset); + let mut max_log: LinearMap = LinearMap::new(); - for (ldes, points) in &rounds { - for (lde, ps) in ldes.iter().zip(points.iter()) { - for &z in ps { - if let Some(h) = max_log.get_mut(&z) { - *h = (*h).max(log2_strict_usize(lde.height())) + for ((_, points), round_dimensions) in commitment_data_with_opening_points + .iter() + .zip(dimensions.iter()) + { + for (points, dims) in points.iter().zip(round_dimensions) { + for &point in points { + let log_height = log2_strict_usize(dims.height); + if let Some(maximum) = max_log.get_mut(&point) { + *maximum = (*maximum).max(log_height); } else { - max_log.insert(z, log2_strict_usize(lde.height())); + max_log.insert(point, log_height); } } } } - let ext_x = Challenge::from_basis_coefficients_slice(&[Val::ZERO, Val::ONE]).unwrap(); + + let ext_x = Challenge::from_basis_coefficients_slice(&[Val::ZERO, Val::ONE]) + .expect("quadratic extension generator"); let ext_w = to_pair(ext_x * ext_x)[0]; assert_eq!(ext_w, Goldilocks::from_u64(7)); let coset_gold: Vec<_> = coset.iter().copied().map(to_gold).collect(); @@ -724,72 +1450,255 @@ where let mut inverse_points = Vec::new(); let mut inverse_counts = Vec::new(); let mut inverse_count = 0usize; - for (point, lh) in max_log { + let point_logs = max_log.into_iter().collect_vec(); + let (_, total_device_bytes) = crate::cuda::device_memory_info(device_id); + let inverse_elements = point_logs + .iter() + .map(|(_, log_height)| 1usize << log_height) + .sum::(); + let max_inverse_elements = point_logs + .iter() + .map(|(_, log_height)| 1usize << log_height) + .max() + .unwrap_or(0); + let reduced_elements = dimensions + .iter() + .flatten() + .map(|dims| dims.height) + .collect::>() + .into_iter() + .sum::(); + let fri_workspace_bytes = inverse_elements + .saturating_mul(2 * size_of::()) + .saturating_add(cuda_max_height.saturating_mul(size_of::())) + .saturating_add(max_inverse_elements.saturating_mul(2 * size_of::())) + .saturating_add(reduced_elements.saturating_mul(2 * size_of::())) + .saturating_add(total_device_bytes / 64); + let admission_started = std::time::Instant::now(); + let admission_data = commitment_data_with_opening_points + .iter() + .map(|entry| entry.0) + .collect_vec(); + self.mmcs + .ensure_device_headroom_batch(&admission_data, fri_workspace_bytes); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] FRI admission: {:.3}s", + admission_started.elapsed().as_secs_f64() + ); + } + for &(point, log_height) in &point_logs { inv_offsets.insert(point, inverse_count); inverse_points.push(to_pair(point)); - let count = 1usize << lh; + let count = 1usize << log_height; inverse_counts.push(count); inverse_count += count; } - let mut workspace = CudaFriWorkspace::new( - device_id, - &inverse_points, - &inverse_counts, - &coset_gold, - ext_w, - ); - let mut interpolation_tasks = Vec::new(); - let mut output_count = 0usize; - let layouts = rounds + let cpu_round_matrices = commitment_data_with_opening_points .iter() - .map(|(ldes, points)| { - ldes.iter() - .zip(points.iter()) - .map(|(lde, ps)| { - let h = lde.height() >> self.fri.log_blowup; - let lh = log2_strict_usize(h); - ps.iter() - .map(|&point| { - let offset = output_count; - output_count += lde.width(); - let shift_pow = Val::GENERATOR.exp_power_of_2(lh); - let scale = (point.exp_power_of_2(lh) - shift_pow) - * (Val::from_usize(h) * shift_pow).inverse(); - interpolation_tasks.push(lde.interpolation_task( - h, - *inv_offsets.get(&point).unwrap(), - offset, - to_pair(scale), - )); - (offset, lde.width()) + .map(|(data, _)| self.mmcs.cpu_matrices(data)) + .collect_vec(); + let mut cpu_max_log: LinearMap = LinearMap::new(); + for ((matrices, &points), round_dimensions) in cpu_round_matrices + .iter() + .zip(opening_points.iter()) + .zip(dimensions.iter()) + { + for ((matrix, points), dims) in + matrices.iter().zip(points.iter()).zip(round_dimensions) + { + if matrix.is_none() { + continue; + } + let log_height = log2_strict_usize(dims.height); + for &point in points { + if let Some(maximum) = cpu_max_log.get_mut(&point) { + *maximum = (*maximum).max(log_height); + } else { + cpu_max_log.insert(point, log_height); + } + } + } + } + let cpu_point_logs = cpu_max_log.into_iter().collect_vec(); + let fri_log_blowup = self.fri.log_blowup; + let denominators_started = std::time::Instant::now(); + let ((inv_denoms, cpu_denominator_seconds), mut workspace, gpu_denominator_seconds) = + std::thread::scope(|scope| { + let cpu_denominators = scope.spawn(|| { + let started = std::time::Instant::now(); + let denominators: LinearMap> = cpu_point_logs + .iter() + .map(|&(point, log_height)| { + let count = 1 << log_height; + let coset = coset_gold[..count].as_ref(); + let inverses = goldilocks_quadratic_inverse_denominators( + to_pair(point), + coset, + ext_w, + ) + .into_par_iter() + .map(from_pair) + .collect(); + (point, inverses) + }) + .collect(); + (denominators, started.elapsed().as_secs_f64()) + }); + let gpu_started = std::time::Instant::now(); + let workspace = CudaFriWorkspace::new( + device_id, + &inverse_points, + &inverse_counts, + &coset_gold, + ext_w, + ); + let gpu_seconds = gpu_started.elapsed().as_secs_f64(); + let cpu_denominators = cpu_denominators.join().unwrap(); + (cpu_denominators, workspace, gpu_seconds) + }); + let adjusted_weights: LinearMap> = inv_denoms + .iter() + .map(|(point, denoms)| (*point, compute_adjusted_weights(*point, denoms))) + .collect(); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] FRI denominators: {:.3}s (CPU {:.3}s, GPU {:.3}s)", + denominators_started.elapsed().as_secs_f64(), + cpu_denominator_seconds, + gpu_denominator_seconds + ); + } + + let interpolation_started = std::time::Instant::now(); + let ((cpu_opened, cpu_interpolation_seconds), gpu_opened, gpu_interpolation_seconds) = + std::thread::scope(|scope| { + let cpu_opened = scope.spawn(|| { + let started = std::time::Instant::now(); + let opened = cpu_round_matrices + .iter() + .zip(opening_points.iter()) + .map(|(matrices, &points)| { + matrices + .par_iter() + .copied() + .zip(points.par_iter()) + .map(|(matrix, points)| { + let matrix = matrix?; + let polynomial_height = matrix.height() >> fri_log_blowup; + assert!(polynomial_height.is_power_of_two()); + let (low_coset, _) = matrix.split_rows(polynomial_height); + Some( + points + .iter() + .map(|&point| { + low_coset.interpolate_coset_with_precomputation( + Val::GENERATOR, + point, + &adjusted_weights.get(&point).unwrap() + [..polynomial_height], + ) + }) + .collect_vec(), + ) + }) + .collect::>() + }) + .collect_vec(); + (opened, started.elapsed().as_secs_f64()) + }); + let gpu_started = std::time::Instant::now(); + let gpu_opened = commitment_data_with_opening_points + .iter() + .zip(dimensions.iter()) + .map(|((data, points), round_dimensions)| { + points + .iter() + .zip(round_dimensions) + .enumerate() + .map(|(matrix_index, (points, dims))| { + if !self.mmcs.is_matrix_cuda_resident(data, matrix_index) { + return None; + } + if points.is_empty() { + return Some(Vec::new()); + } + Some(self.mmcs.with_resident_matrix( + data, + matrix_index, + |lde| { + assert_eq!(lde.height(), dims.height); + assert_eq!(lde.width(), dims.width); + let polynomial_height = + lde.height() >> self.fri.log_blowup; + assert!(polynomial_height.is_power_of_two()); + let log_polynomial_height = + log2_strict_usize(polynomial_height); + let shift_power = Val::GENERATOR + .exp_power_of_2(log_polynomial_height); + let tasks = points + .iter() + .enumerate() + .map(|(point_index, &point)| { + let scale = (point + .exp_power_of_2(log_polynomial_height) + - shift_power) + * (Val::from_usize(polynomial_height) + * shift_power) + .inverse(); + lde.interpolation_task( + polynomial_height, + *inv_offsets.get(&point).unwrap(), + point_index * lde.width(), + to_pair(scale), + ) + }) + .collect_vec(); + workspace + .interpolate( + &tasks, + points.len() * lde.width(), + ext_w, + ) + .chunks_exact(lde.width()) + .map(|values| { + values + .iter() + .copied() + .map(from_pair) + .collect_vec() + }) + .collect_vec() + }, + )) }) .collect_vec() }) - .collect_vec() - }) - .collect_vec(); - let interpolated = debug_span!("cuda interpolate openings") - .in_scope(|| workspace.interpolate(&interpolation_tasks, output_count, ext_w)); - let all_opened_values = layouts + .collect_vec(); + let gpu_seconds = gpu_started.elapsed().as_secs_f64(); + let cpu_opened = cpu_opened.join().unwrap(); + (cpu_opened, gpu_opened, gpu_seconds) + }); + let all_opened_values = cpu_opened .into_iter() - .map(|round| { - round + .zip(gpu_opened) + .map(|(cpu_round, gpu_round)| { + cpu_round .into_iter() - .map(|matrix| { - matrix - .into_iter() - .map(|(offset, width)| { - interpolated[offset..offset + width] - .iter() - .copied() - .map(from_pair) - .collect_vec() - }) - .collect_vec() - }) + .zip(gpu_round) + .map(|(cpu, gpu)| cpu.or(gpu).expect("matrix has no opening backend")) .collect_vec() }) .collect_vec(); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] FRI interpolation: {:.3}s (CPU {:.3}s, GPU {:.3}s)", + interpolation_started.elapsed().as_secs_f64(), + cpu_interpolation_seconds, + gpu_interpolation_seconds + ); + } + for round in &all_opened_values { for matrix in round { for values in matrix { @@ -798,36 +1707,167 @@ where } } let alpha: Challenge = challenger.sample_algebra_element(); - let alpha_powers: Vec<_> = alpha.powers().take(global_max_width).collect(); + let alpha_powers: Vec<_> = alpha.powers().take(cuda_max_width).collect(); let alpha_pairs: Vec<_> = alpha_powers.iter().copied().map(to_pair).collect(); + let packed_alpha_powers = + Challenge::ExtensionPacking::packed_ext_powers_capped(alpha, cuda_max_width) + .collect_vec(); let mut num_reduced = [0usize; 33]; - let mut reduced: [Option; 33] = core::array::from_fn(|_| None); - let mut reduction_tasks = Vec::new(); - for ((ldes, points), openings_round) in rounds.iter().zip(all_opened_values.iter()) { - for ((lde, ps), openings) in - ldes.iter().zip(points.iter()).zip(openings_round.iter()) - { - let lh = log2_strict_usize(lde.height()); - let target = reduced[lh] - .get_or_insert_with(|| CudaReducedOpening::new(device_id, lde.height())); - for (&point, ys) in ps.iter().zip(openings.iter()) { - let reduced_y = - dot_product(alpha_powers.iter().copied(), ys.iter().copied()); - let offset = alpha.exp_u64(num_reduced[lh] as u64); - reduction_tasks.push(target.reduction_task( - lde, - *inv_offsets.get(&point).unwrap(), - to_pair(reduced_y), - to_pair(offset), - )); - num_reduced[lh] += lde.width(); + let reduction_offsets = commitment_data_with_opening_points + .iter() + .zip(dimensions.iter()) + .map(|((_, points), round_dimensions)| { + points + .iter() + .zip(round_dimensions) + .map(|(points, dims)| { + let log_height = log2_strict_usize(dims.height); + points + .iter() + .map(|_| { + let offset = num_reduced[log_height]; + num_reduced[log_height] += dims.width; + offset + }) + .collect_vec() + }) + .collect_vec() + }) + .collect_vec(); + + let reduction_started = std::time::Instant::now(); + let ((cpu_reduced, cpu_reduction_seconds), mut gpu_reduced, gpu_reduction_seconds) = + std::thread::scope(|scope| { + let cpu_reduced = scope.spawn(|| { + let started = std::time::Instant::now(); + let mut reduced: [Option>; 33] = + core::array::from_fn(|_| None); + for ( + (((matrices, &points), round_dimensions), openings_round), + offsets_round, + ) in cpu_round_matrices + .iter() + .zip(opening_points.iter()) + .zip(dimensions.iter()) + .zip(all_opened_values.iter()) + .zip(reduction_offsets.iter()) + { + for ((((matrix, points), dims), openings), offsets) in matrices + .iter() + .copied() + .zip(points) + .zip(round_dimensions) + .zip(openings_round) + .zip(offsets_round) + { + let Some(matrix) = matrix else { + continue; + }; + if points.is_empty() { + continue; + } + let log_height = log2_strict_usize(dims.height); + let target = reduced[log_height] + .get_or_insert_with(|| vec![Challenge::ZERO; dims.height]); + let compressed = matrix + .rowwise_packed_dot_product::(&packed_alpha_powers) + .collect::>(); + for ((&point, values), &offset) in + points.iter().zip(openings).zip(offsets) + { + let reduced_y: Challenge = dot_product( + alpha_powers.iter().copied(), + values.iter().copied(), + ); + let alpha_offset = alpha.exp_u64(offset as u64); + compressed + .par_iter() + .zip(target.par_iter_mut()) + .zip(inv_denoms.get(&point).unwrap().par_iter()) + .for_each(|((&row, output), &inv_denom)| { + *output += alpha_offset * (reduced_y - row) * inv_denom; + }); + } + } + } + (reduced, started.elapsed().as_secs_f64()) + }); + let gpu_started = std::time::Instant::now(); + let mut reduced: [Option; 33] = + core::array::from_fn(|_| None); + let mut reduction_tasks = Vec::new(); + for ((((data, points), round_dimensions), openings_round), offsets_round) in + commitment_data_with_opening_points + .iter() + .zip(dimensions.iter()) + .zip(all_opened_values.iter()) + .zip(reduction_offsets.iter()) + { + for (matrix_index, (((points, dims), openings), offsets)) in points + .iter() + .zip(round_dimensions) + .zip(openings_round) + .zip(offsets_round) + .enumerate() + { + if !self.mmcs.is_matrix_cuda_resident(data, matrix_index) { + continue; + } + if points.is_empty() { + continue; + } + let log_height = log2_strict_usize(dims.height); + let target = reduced[log_height].get_or_insert_with(|| { + CudaReducedOpening::new(device_id, dims.height) + }); + self.mmcs.with_resident_matrix(data, matrix_index, |lde| { + reduction_tasks.extend( + points.iter().zip(openings).zip(offsets).map( + |((&point, values), &offset)| { + let reduced_y = dot_product( + alpha_powers.iter().copied(), + values.iter().copied(), + ); + target.reduction_task( + lde, + *inv_offsets.get(&point).unwrap(), + to_pair(reduced_y), + to_pair(alpha.exp_u64(offset as u64)), + ) + }, + ), + ); + }); + } } + if !reduction_tasks.is_empty() { + workspace.reduce(&reduction_tasks, &alpha_pairs, ext_w); + } + let gpu_seconds = gpu_started.elapsed().as_secs_f64(); + let cpu_reduced = cpu_reduced.join().unwrap(); + (cpu_reduced, reduced, gpu_seconds) + }); + drop(workspace); + + for (log_height, cpu_values) in cpu_reduced.into_iter().enumerate() { + if let Some(cpu_values) = cpu_values { + let values = cpu_values.into_iter().map(to_pair).collect_vec(); + gpu_reduced[log_height] + .get_or_insert_with(|| CudaReducedOpening::new(device_id, values.len())) + .add_host(&values); } } - debug_span!("cuda reduce openings") - .in_scope(|| workspace.reduce(&reduction_tasks, &alpha_pairs, ext_w)); - let fri_input = reduced.into_iter().rev().flatten().collect_vec(); - let fri_proof = debug_span!("cuda prove fri").in_scope(|| { + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] FRI reduction: {:.3}s (CPU {:.3}s, GPU {:.3}s)", + reduction_started.elapsed().as_secs_f64(), + cpu_reduction_seconds, + gpu_reduction_seconds + ); + } + let fri_input = gpu_reduced.into_iter().rev().flatten().collect_vec(); + let folding_started = std::time::Instant::now(); + let fri_proof = debug_span!("cuda prove streamed fri").in_scope(|| { prove_fri_cuda_resident( &self.fri, fri_input, @@ -838,6 +1878,13 @@ where ext_w, ) }); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] FRI folding and queries: {:.3}s (streamed total {:.3}s)", + folding_started.elapsed().as_secs_f64(), + phase_started.elapsed().as_secs_f64() + ); + } return (all_opened_values, fri_proof); } @@ -1164,3 +2211,29 @@ fn compute_inverse_denominators, M: Matri }) .collect() } + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::ExtVal; + use p3_field::Field; + + #[test] + fn quadratic_norm_inverses_match_extension_inverses() { + let point_pair = [Goldilocks::from_u64(123), Goldilocks::from_u64(456)]; + let point = ExtVal::from_basis_coefficients_slice(&point_pair).unwrap(); + let coset = Goldilocks::GENERATOR.powers().take(4096).collect_vec(); + + for count in [1, 3, 1024, 1025, 4096] { + let actual = goldilocks_quadratic_inverse_denominators( + point_pair, + &coset[..count], + Goldilocks::from_u64(7), + ); + let expected = coset[..count].iter().map(|&x| (point - x).inverse()); + for (actual, expected) in actual.iter().zip(expected) { + assert_eq!(actual, expected.as_basis_coefficients_slice()); + } + } + } +} diff --git a/src/lookup.rs b/src/lookup.rs index 642fb62..27d621d 100644 --- a/src/lookup.rs +++ b/src/lookup.rs @@ -45,6 +45,8 @@ //! [`MAX_LOOKUP_GROUP`] so the evaluator's scratch space stays on the //! stack. +#[cfg(feature = "cuda")] +use p3_field::BasedVectorSpace; use p3_field::{ Algebra, ExtensionField, Field, PrimeCharacteristicRing, batch_multiplicative_inverse, }; @@ -704,6 +706,57 @@ impl LookupValues { #[cfg(feature = "cuda")] impl LookupValues { + pub(crate) fn cuda_stage_2_deltas>( + &self, + rows: core::ops::Range, + group_size: usize, + lookup_challenge: EF, + fingerprint_challenge: &EF, + ) -> Vec<[p3_goldilocks::Goldilocks; 2]> { + assert!(rows.start <= rows.end && rows.end <= self.height); + assert!(self.num_lookups != 0); + assert_eq!( + >::DIMENSION, + 2 + ); + let group_size = group_size.max(1); + let slots = lookup_groups(self.num_lookups, group_size); + let message_start = rows.start * self.num_lookups; + let message_end = rows.end * self.num_lookups; + let messages = (message_start..message_end) + .into_par_iter() + .map(|index| { + let row = index / self.num_lookups; + let lookup = index % self.num_lookups; + lookup_challenge + + fingerprint( + fingerprint_challenge, + self.args_at(row, lookup).iter().copied(), + ) + }) + .collect::>(); + let inverses = batch_multiplicative_inverse(&messages); + drop(messages); + (0..(rows.end - rows.start) * slots) + .into_par_iter() + .map(|index| { + let local_row = index / slots; + let slot = index % slots; + let begin = slot * group_size; + let end = (begin + group_size).min(self.num_lookups); + let row = rows.start + local_row; + let inverse_row = local_row * self.num_lookups; + let mut delta = EF::ZERO; + for lookup in begin..end { + let multiplicity = self.multiplicities[row * self.num_lookups + lookup]; + delta += EF::from(multiplicity) * inverses[inverse_row + lookup]; + } + let coordinates = delta.as_basis_coefficients_slice(); + [coordinates[0], coordinates[1]] + }) + .collect() + } + pub(crate) fn cuda_parts( &self, ) -> ( diff --git a/src/types.rs b/src/types.rs index d580468..3dcc57e 100644 --- a/src/types.rs +++ b/src/types.rs @@ -125,6 +125,45 @@ fn cuda_coset_selectors( } } +#[cfg(feature = "cuda")] +fn cuda_worker_count(variable: &str, numerator: usize, denominator: usize) -> usize { + let default_threads = std::thread::available_parallelism() + .map_or(1, std::num::NonZero::get) + .saturating_mul(numerator) + .div_ceil(denominator); + std::env::var(variable) + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|&threads| threads != 0) + .unwrap_or(default_threads) +} + +#[cfg(feature = "cuda")] +pub(crate) fn cuda_stage1_worker_count(prioritize_deferred: bool) -> usize { + std::env::var("MULTI_STARK_CUDA_DEFERRED_THREADS") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|&threads| threads != 0) + .unwrap_or_else(|| { + let (numerator, denominator) = if prioritize_deferred { (5, 8) } else { (3, 8) }; + cuda_worker_count("MULTI_STARK_CUDA_STAGE1_THREADS", numerator, denominator) + }) +} + +#[cfg(feature = "cuda")] +pub(crate) fn cuda_lookup_worker_count() -> usize { + cuda_worker_count("MULTI_STARK_CUDA_LOOKUP_THREADS", 1, 4) +} + +#[cfg(feature = "cuda")] +pub(crate) fn cuda_host_pool(name: &'static str, threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .thread_name(move |index| format!("cuda-{name}-{index}")) + .build() + .expect("failed to build CUDA host worker pool") +} + /// The reference [`StarkGenericConfig`] implementation. pub struct GoldilocksBlake3Config { /// The PCS used to commit polynomials and prove opening proofs. @@ -346,63 +385,243 @@ impl StarkGenericConfig for GoldilocksBlake3Config { alpha: ExtVal, ) -> Option<(crate::config::Com, crate::config::PcsData)> { use crate::cuda::mmcs::CudaCommitMmcs; - let ldes: Option> = inputs + + let (_, total_device_bytes) = + crate::cuda::device_memory_info(self.pcs.mmcs.cuda_device_id()); + let matrix_width = |data: &crate::config::PcsData, index: usize| { + let dimensions = self.pcs.mmcs.matrix_dimensions(data); + dimensions + .get(index) + .expect("matrix index out of bounds") + .width + }; + let staging_bytes = |input: &crate::config::QuotientCommitInput<'_, Self>| { + let mut host_width = 0usize; + for (data, index) in [Some(input.stage_1), Some(input.stage_2), input.preprocessed] + .into_iter() + .flatten() + { + if !self.pcs.mmcs.is_matrix_cuda_resident(data, index) { + host_width = host_width.saturating_add(matrix_width(data, index)); + } + } + if host_width == 0 { + return 0; + } + let row_bytes = host_width + .saturating_mul(2) + .saturating_mul(size_of::()); + let per_buffer = input + .quotient_domain + .size() + .saturating_mul(row_bytes) + .min(512usize << 20); + 2usize.saturating_mul(per_buffer) + }; + let mut quotient_jobs = inputs .iter() - .map(|input| { - let main = input.stage_1.0.resident(input.stage_1.1)?; - let stage2 = input.stage_2.0.resident(input.stage_2.1)?; - let preprocessed = match input.preprocessed { - Some((data, index)) => Some(data.resident(index)?), - None => None, - }; + .enumerate() + .map(|(index, input)| { let quotient_size = input.quotient_domain.size(); let quotient_degree = input.circuit.quotient_degree(); - let selectors = cuda_coset_selectors(input.trace_domain, input.quotient_domain); - - let mut powers = Vec::with_capacity(input.constraint_count); - let mut power = ExtVal::ONE; - for _ in 0..input.constraint_count { - powers.push(power); - power *= alpha; - } - powers.reverse(); - let mut alpha_flat = Vec::with_capacity(2 * input.constraint_count); - for coordinate in 0..2 { - alpha_flat.extend(powers.iter().map(|value| { - >::as_basis_coefficients_slice(value) - [coordinate] - })); - } - let trace_size = input.trace_domain.size(); - let n = Val::from_usize(trace_size); - let generator = Val::two_adic_generator(trace_size.ilog2() as usize); - let normalization = (n * generator).inverse(); - let delta = [ - (input.lookup_publics[6] - input.lookup_publics[4]) * normalization, - (input.lookup_publics[7] - input.lookup_publics[5]) * normalization, - ]; - let next_step = quotient_size / trace_size; - Some(crate::cuda::quotient_lde_resident( - &self.pcs.dft, + let trace_height = quotient_size / quotient_degree; + let lde_height = trace_height + .checked_shl(u32::try_from(self.log_blowup).expect("LDE blowup exceeds u32")) + .unwrap_or(usize::MAX); + let (output_bytes, kernel_workspace) = crate::cuda::quotient_lde_memory_upper_bound( &input.circuit.graph, - preprocessed, - main, - stage2, - &input.lookup_publics, - selectors, - &alpha_flat, - &delta, - crate::system::extension_params::().w, + input.lookup_publics.len(), + input.constraint_count, quotient_size, - next_step, - input.circuit.lookup_group_size, quotient_degree, self.log_blowup, - )) + ); + let current_staging = staging_bytes(input); + let constant_bytes = quotient_size + .saturating_add(lde_height) + .saturating_div(2) + .saturating_add(quotient_degree) + .saturating_mul(size_of::()); + ( + index, + output_bytes, + kernel_workspace, + constant_bytes, + current_staging + .saturating_add(kernel_workspace) + .saturating_add(output_bytes) + .saturating_add(constant_bytes), + lde_height, + ) }) - .collect(); - let ldes = ldes?; - Some(self.pcs.mmcs.commit_cuda_resident(ldes)) + .collect::>(); + quotient_jobs.sort_unstable_by_key(|&(index, _, _, _, peak, _)| { + let input = &inputs[index]; + let ready = [Some(input.stage_1), Some(input.stage_2), input.preprocessed] + .into_iter() + .flatten() + .all(|(data, matrix)| self.pcs.mmcs.is_matrix_source_ready(data, matrix)); + (!ready, core::cmp::Reverse(peak)) + }); + let accumulated_output = quotient_jobs + .iter() + .map(|&(_, output, _, _, _, _)| output) + .sum::(); + let max_lde_height = quotient_jobs + .iter() + .map(|&(_, _, _, _, _, lde_height)| lde_height) + .max() + .unwrap_or(0); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] quotient jobs={} output={}", + quotient_jobs.len(), + accumulated_output, + ); + } + + let mut ldes = (0..inputs.len()).map(|_| None).collect::>(); + for (index, output_bytes, kernel_workspace, constant_bytes, _, lde_height) in quotient_jobs + { + let job_started = std::time::Instant::now(); + let input = &inputs[index]; + let current_staging = || staging_bytes(input); + let required = || { + current_staging() + .saturating_add(kernel_workspace) + .saturating_add(output_bytes) + .saturating_add(constant_bytes) + .saturating_add(total_device_bytes / 64) + }; + let mut target = required(); + self.pcs + .mmcs + .ensure_device_headroom(input.stage_1.0, target, Some(input.stage_1.1)); + target = required(); + self.pcs + .mmcs + .ensure_device_headroom(input.stage_2.0, target, Some(input.stage_2.1)); + if let Some((data, matrix)) = input.preprocessed { + target = required(); + self.pcs + .mmcs + .ensure_device_headroom(data, target, Some(matrix)); + } + target = required(); + let free_bytes = crate::cuda::device_memory_info(self.pcs.mmcs.cuda_device_id()).0; + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] quotient job {index}: lde_height={lde_height} output={output_bytes} workspace={kernel_workspace} staging={} target={} free={free_bytes}", + current_staging(), + target + ); + } + if free_bytes < target { + return None; + } + let lde = + { + let quotient_size = input.quotient_domain.size(); + let quotient_degree = input.circuit.quotient_degree(); + let selectors = cuda_coset_selectors(input.trace_domain, input.quotient_domain); + + let mut powers = Vec::with_capacity(input.constraint_count); + let mut power = ExtVal::ONE; + for _ in 0..input.constraint_count { + powers.push(power); + power *= alpha; + } + powers.reverse(); + let mut alpha_flat = Vec::with_capacity(2 * input.constraint_count); + for coordinate in 0..2 { + alpha_flat.extend(powers.iter().map(|value| { + >::as_basis_coefficients_slice(value) + [coordinate] + })); + } + let trace_size = input.trace_domain.size(); + let n = Val::from_usize(trace_size); + let generator = Val::two_adic_generator(trace_size.ilog2() as usize); + let normalization = (n * generator).inverse(); + let delta = [ + (input.lookup_publics[6] - input.lookup_publics[4]) * normalization, + (input.lookup_publics[7] - input.lookup_publics[5]) * normalization, + ]; + let next_step = quotient_size / trace_size; + let evaluate = |main: crate::cuda::mmcs::CudaMatrixSource<'_>, + stage2: crate::cuda::mmcs::CudaMatrixSource<'_>, + preprocessed: Option< + crate::cuda::mmcs::CudaMatrixSource<'_>, + >| { + crate::cuda::quotient_lde_mixed( + &self.pcs.dft, + &input.circuit.graph, + preprocessed, + main, + stage2, + &input.lookup_publics, + selectors, + &alpha_flat, + &delta, + crate::system::extension_params::().w, + quotient_size, + next_step, + input.circuit.lookup_group_size, + quotient_degree, + self.log_blowup, + ) + }; + self.pcs + .mmcs + .with_matrix_source(input.stage_1.0, input.stage_1.1, |main| { + self.pcs.mmcs.with_matrix_source( + input.stage_2.0, + input.stage_2.1, + |stage2| match input.preprocessed { + Some((data, index)) => self.pcs.mmcs.with_matrix_source( + data, + index, + |preprocessed| evaluate(main, stage2, Some(preprocessed)), + ), + None => evaluate(main, stage2, None), + }, + ) + }) + }; + ldes[index] = Some(lde); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] quotient job {index} complete: {:.3}s", + job_started.elapsed().as_secs_f64() + ); + } + } + let tree_headroom = max_lde_height + .saturating_mul(96) + .saturating_add(total_device_bytes / 64); + if let Some(input) = inputs.first() { + self.pcs + .mmcs + .ensure_device_headroom(input.stage_1.0, tree_headroom, None); + self.pcs + .mmcs + .ensure_device_headroom(input.stage_2.0, tree_headroom, None); + } + if let Some((data, _)) = inputs.iter().find_map(|input| input.preprocessed) { + self.pcs + .mmcs + .ensure_device_headroom(data, tree_headroom, None); + } + if crate::cuda::device_memory_info(self.pcs.mmcs.cuda_device_id()).0 < tree_headroom { + return None; + } + Some( + self.pcs.mmcs.commit_cuda_spillable( + ldes.into_iter() + .map(|lde| lde.expect("quotient job did not run")) + .collect(), + ), + ) } #[cfg(feature = "cuda")] @@ -418,6 +637,113 @@ impl StarkGenericConfig for GoldilocksBlake3Config { Vec, )> { use crate::cuda::mmcs::CudaCommitMmcs; + + let mut lookup_jobs = inputs + .iter() + .enumerate() + .map(|(index, input)| { + let (height, num_lookups, _, _, arg_offsets) = input.lookup_values.cuda_parts(); + let group_size = input.circuit.lookup_group_size.max(1); + let groups = num_lookups.div_ceil(group_size).max(1); + let extended_height = height + .checked_shl(u32::try_from(self.log_blowup).expect("LDE blowup exceeds u32")) + .unwrap_or(usize::MAX); + let output_bytes = extended_height + .saturating_mul(2) + .saturating_mul(groups) + .saturating_mul(size_of::()); + let direct_temporary_bytes = if num_lookups == 0 { + 0 + } else { + let args_width = *arg_offsets.last().expect("lookup offsets are empty"); + let chunk_rows = height.min(1 << 16); + let messages = chunk_rows.saturating_mul(num_lookups); + messages + // Multiplicity, conjugate, norm, and inverse-norm arrays. + .saturating_mul(5 * size_of::()) + .saturating_add( + chunk_rows + .saturating_mul(args_width) + .saturating_mul(size_of::()), + ) + // Extension-field deltas before the exclusive scan. + .saturating_add( + height + .saturating_mul(groups) + .saturating_mul(2 * size_of::()), + ) + .saturating_add(arg_offsets.len().saturating_mul(size_of::())) + // Device-cached inverse/forward twiddles and coset + // powers may be cold for this height. + .saturating_add( + height + .saturating_add(height / 2) + .saturating_add(extended_height / 2) + .saturating_mul(size_of::()), + ) + }; + let main_width = self + .pcs + .mmcs + .matrix_dimensions(input.stage_1.0) + .get(input.stage_1.1) + .expect("stage-1 matrix index out of bounds") + .width; + let graph_memory = (num_lookups != 0) + .then(|| { + crate::cuda::lookup_graph_lde_memory_upper_bound( + &input.circuit.graph, + height, + main_width, + group_size, + self.log_blowup, + ) + }) + .flatten(); + ( + index, + output_bytes, + direct_temporary_bytes, + graph_memory.map(|(_, temporary)| temporary), + extended_height, + ) + }) + .collect::>(); + // Permanent lookup LDEs accumulate until their MMCS commitment. Run + // the initially largest allocation first, while its trace is still + // resident, then allow that trace to become an eviction candidate for + // later jobs. + lookup_jobs.sort_unstable_by_key(|&(index, output, direct, graph, _)| { + let temporary = if self + .pcs + .mmcs + .is_matrix_cuda_resident(inputs[index].stage_1.0, inputs[index].stage_1.1) + { + graph.unwrap_or(direct) + } else { + direct + }; + core::cmp::Reverse(output.saturating_add(temporary)) + }); + let accumulated_output = lookup_jobs + .iter() + .map(|&(_, output, _, _, _)| output) + .sum::(); + let max_lde_height = lookup_jobs + .iter() + .map(|&(_, _, _, _, extended_height)| extended_height) + .max() + .unwrap_or(0); + let (_, total_device_bytes) = + crate::cuda::device_memory_info(self.pcs.mmcs.cuda_device_id()); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] lookup jobs={} output={}", + lookup_jobs.len(), + accumulated_output, + ); + } + let pair = |value: ExtVal| { let coordinates = value.as_basis_coefficients_slice(); [coordinates[0], coordinates[1]] @@ -426,57 +752,156 @@ impl StarkGenericConfig for GoldilocksBlake3Config { let gamma = pair(fingerprint_challenge); let extension_generator = ExtVal::from_basis_coefficients_slice(&[Val::ZERO, Val::ONE])?; let ext_w = pair(extension_generator * extension_generator)[0]; - let evaluate = |input: &crate::config::LookupCommitInput<'_, Self>| { - let main = input.stage_1.0.resident_with_trace(input.stage_1.1)?; - let preprocessed = match input.preprocessed { - Some((data, index)) => Some(data.resident(index)?), - None => None, - }; + let lookup_pool = + std::sync::Arc::new(cuda_host_pool("lookup-rows", cuda_lookup_worker_count())); + let evaluate = |input: &crate::config::LookupCommitInput<'_, Self>, cooperative: bool| { let (height, num_lookups, multiplicities, args, arg_offsets) = input.lookup_values.cuda_parts(); - let result = if num_lookups == 0 { - Some(crate::cuda::lookup_lde_resident( + let group_size = input.circuit.lookup_group_size.max(1); + let main = input.stage_1.0.resident_with_trace(input.stage_1.1); + let result = if let Some(main) = main.filter(|_| num_lookups != 0) { + let preprocessed = match input.preprocessed { + Some((data, index)) => Some(data.resident(index)?), + None => None, + }; + crate::cuda::lookup_graph_lde_resident( + &self.pcs.dft, + &input.circuit.graph, + preprocessed, + main, + height, + group_size, + beta, + gamma, + ext_w, + self.log_blowup, + ) + } else if cooperative { + let pool = std::sync::Arc::clone(&lookup_pool); + Some(crate::cuda::lookup_lde_resident_partitioned( &self.pcs.dft, multiplicities, args, arg_offsets, height, num_lookups, - input.circuit.lookup_group_size.max(1), + group_size, beta, gamma, ext_w, self.log_blowup, + move |rows| { + pool.install(|| { + input.lookup_values.cuda_stage_2_deltas::( + rows, + group_size, + lookup_challenge, + &fingerprint_challenge, + ) + }) + }, )) } else { - crate::cuda::lookup_graph_lde_resident( + Some(crate::cuda::lookup_lde_resident( &self.pcs.dft, - &input.circuit.graph, - preprocessed, - main, + multiplicities, + args, + arg_offsets, height, - input.circuit.lookup_group_size.max(1), + num_lookups, + group_size, beta, gamma, ext_w, self.log_blowup, - ) + )) }; // SAFETY: evaluation is synchronous and complete, while the // retained matrix remains owned by the prover data. - unsafe { main.release_trace() }; + if let Some(main) = main { + unsafe { main.release_trace() }; + } result }; - let results: Option> = inputs.iter().map(evaluate).collect(); - let results = results?; + let mut results = (0..inputs.len()).map(|_| None).collect::>(); + for (index, output_bytes, direct_temporary_bytes, graph_temporary_bytes, _) in lookup_jobs { + let job_started = std::time::Instant::now(); + let input = &inputs[index]; + let graph_path = self + .pcs + .mmcs + .is_matrix_cuda_resident(input.stage_1.0, input.stage_1.1) + && graph_temporary_bytes.is_some() + && input.preprocessed.is_none_or(|(data, matrix)| { + self.pcs.mmcs.is_matrix_cuda_resident(data, matrix) + }); + let temporary_bytes = if graph_path { + graph_temporary_bytes.unwrap() + } else { + direct_temporary_bytes + }; + let target = output_bytes + .saturating_add(temporary_bytes) + .saturating_add(total_device_bytes / 64); + if target > total_device_bytes { + return None; + } + let mut free_bytes = self.pcs.mmcs.ensure_device_headroom( + input.stage_1.0, + target, + Some(input.stage_1.1), + ); + if free_bytes < target { + // This circuit alone does not fit beside its resident trace. + // Spill it as a last resort and use the direct lookup-values + // path, which remains protocol-identical. + free_bytes = self + .pcs + .mmcs + .ensure_device_headroom(input.stage_1.0, target, None); + } + if free_bytes < target { + return None; + } + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] lookup job {index}: graph={graph_path} output={output_bytes} temporary={temporary_bytes} target={target} free={free_bytes}" + ); + } + let (_, num_lookups, _, _, arg_offsets) = input.lookup_values.cuda_parts(); + let cooperative = !graph_path + && output_bytes >= (8usize << 30) + && num_lookups >= 64 + && arg_offsets.last().copied().unwrap_or(0) >= 256; + results[index] = Some(evaluate(&inputs[index], cooperative)?); + if crate::cuda::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] lookup job {index} complete: {:.3}s", + job_started.elapsed().as_secs_f64() + ); + } + } + let tree_headroom = max_lde_height + .saturating_mul(96) + .saturating_add(total_device_bytes / 64); + if let Some(input) = inputs.first() { + let free_bytes = + self.pcs + .mmcs + .ensure_device_headroom(input.stage_1.0, tree_headroom, None); + if free_bytes < tree_headroom { + return None; + } + } let mut ldes = Vec::with_capacity(results.len()); let mut intermediates = Vec::with_capacity(inputs.len()); - for (lde, total) in results { + for result in results { + let (lde, total) = result.expect("lookup job did not run"); accumulator += ExtVal::from_basis_coefficients_slice(&total)?; intermediates.push(accumulator); ldes.push(lde); } - let (commitment, data) = self.pcs.mmcs.commit_cuda_resident(ldes); + let (commitment, data) = self.pcs.mmcs.commit_cuda_spillable(ldes); Some((commitment, data, intermediates)) } }