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)) } }