Bump Plonky3

Cargo.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2024"
 authors = ["Argument Engineering <engineering@argument.xyz>"]
 license = "MIT OR Apache-2.0"
-rust-version = "1.91"
+rust-version = "1.98"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
@@ -16,25 +16,24 @@ tracing = "0.1"
 itertools = { version = "0.14", optional = true }
 serde = { version = "1", features = ["derive"] }
 bincode = { version = "2", features = ["serde"] }
-p3-air = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-goldilocks = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-challenger = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-commit = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-dft = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-field = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-fri = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-keccak = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-interpolation = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64", optional = true }
-p3-blake3 = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-matrix = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-maybe-rayon = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-symmetric = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
-p3-util = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
+p3-air = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-goldilocks = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-challenger = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-commit = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-dft = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-field = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-fri = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-keccak = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-blake3 = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-matrix = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-maybe-rayon = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-symmetric = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
+p3-util = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
 
 [dev-dependencies]
 criterion = "0.5"
-p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
+p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" }
 rand = "0.10"
 tracing-subscriber = "0.3"
 tracing-texray = { git = "https://github.com/argumentcomputer/tracing-texray", rev = "465bbca0bea4721e58419c11cabd8cce21757822" }
@@ -48,7 +47,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", "dep:p3-interpolation"]
+cuda = ["dep:itertools"]
 
 # Similar to `release`, but preserves debug info
 [profile.dev-ci]
@@ -74,7 +73,6 @@ explicit_into_iter_loop = "warn"
 fallible_impl_from = "warn"
 filter_map_next = "warn"
 flat_map_option = "warn"
-from_iter_instead_of_collect = "warn"
 implicit_clone = "warn"
 inefficient_to_string = "warn"
 large_stack_arrays = "warn"

cuda/kernels.cu

@@ -1624,6 +1624,19 @@ cudaError_t launch_blake3_digest_pairs(uint8_t* digests,
     return cudaGetLastError();
 }
 
+__global__ void gather_resident_lde_group(
+    uint64_t* output, const uint64_t* const* columns, const size_t* strides,
+    size_t row_start, size_t rows, size_t width) {
+    const size_t count = rows * width;
+    const size_t grid_stride = static_cast<size_t>(blockDim.x) * gridDim.x;
+    for (size_t index = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+         index < count; index += grid_stride) {
+        const size_t column = index % width;
+        const size_t row = row_start + index / width;
+        output[index] = columns[column][row * strides[column]];
+    }
+}
+
 cudaError_t hash_resident_lde_group(uint8_t* digests,
                                     const void* const* handles,
                                     size_t handle_count, size_t height) {
@@ -1642,35 +1655,65 @@ cudaError_t hash_resident_lde_group(uint8_t* digests,
         return cudaErrorInvalidValue;
     }
 
-    constexpr size_t ROW_STAGING_BYTES = size_t(256) << 20;
+    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;
+    const uint64_t** host_columns =
+        new (std::nothrow) const uint64_t*[total_width];
+    size_t* host_strides = new (std::nothrow) size_t[total_width];
+    if (host_columns == nullptr || host_strides == nullptr) {
+        delete[] host_columns;
+        delete[] host_strides;
+        return cudaErrorMemoryAllocation;
+    }
+    size_t column_offset = 0;
+    for (size_t index = 0; index < handle_count; ++index) {
+        const ResidentLde* lde = static_cast<const ResidentLde*>(handles[index]);
+        if (lde == nullptr || lde->height != height) {
+            continue;
+        }
+        for (size_t column = 0; column < lde->width; ++column) {
+            host_columns[column_offset] = lde->values + column;
+            host_strides[column_offset++] = lde->width;
+        }
+    }
+    DeviceBuffer device_columns;
+    DeviceBuffer device_strides;
+    cudaError_t status = device_columns.allocate(total_width);
+    if (status == cudaSuccess) {
+        status = device_strides.allocate(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);
+    }
+    delete[] host_columns;
+    delete[] host_strides;
     DeviceBuffer combined_rows;
-    cudaError_t status = combined_rows.allocate(
-        (height < rows_per_chunk ? height : rows_per_chunk) * total_width);
+    if (status == cudaSuccess) {
+        status = combined_rows.allocate(
+            (height < rows_per_chunk ? height : rows_per_chunk) * total_width);
+    }
     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;
-        size_t column_offset = 0;
-        for (size_t index = 0; status == cudaSuccess && index < handle_count;
-             ++index) {
-            const ResidentLde* lde =
-                static_cast<const ResidentLde*>(handles[index]);
-            if (lde == nullptr || lde->height != height) {
-                continue;
-            }
-            status = cudaMemcpy2DAsync(
-                combined_rows.get() + column_offset,
-                total_width * sizeof(uint64_t),
-                lde->values + row_start * lde->width,
-                lde->width * sizeof(uint64_t),
-                lde->width * sizeof(uint64_t), rows, cudaMemcpyDeviceToDevice,
-                cudaStreamPerThread);
-            column_offset += lde->width;
-        }
+        const size_t count = rows * total_width;
+        gather_resident_lde_group<<<blocks_for(count), THREADS, 0,
+                                    cudaStreamPerThread>>>(
+            combined_rows.get(),
+            reinterpret_cast<const uint64_t* const*>(device_columns.get()),
+            reinterpret_cast<const size_t*>(device_strides.get()), row_start,
+            rows, total_width);
+        status = cudaGetLastError();
         if (status == cudaSuccess) {
             status = launch_blake3_rows(
                 digests + row_start * 32,
@@ -2453,13 +2496,13 @@ extern "C" int multi_stark_cuda_fri_workspace_destroy(int device_id,void* handle
     cudaError_t status=cudaSetDevice(device_id);if(status==cudaSuccess)delete static_cast<ResidentFriWorkspace*>(handle);return static_cast<int>(status);
 }
 
-extern "C" int multi_stark_cuda_reduced_to_lde(int device_id,void** output,const void* reduced){
+extern "C" int multi_stark_cuda_reduced_into_lde(int device_id,void** output,void* reduced){
     if(!output||!reduced)return static_cast<int>(cudaErrorInvalidValue);*output=nullptr;
-    auto* r=static_cast<const ResidentReducedOpening*>(reduced);cudaError_t status=cudaSetDevice(device_id);
+    auto* r=static_cast<ResidentReducedOpening*>(reduced);cudaError_t status=cudaSetDevice(device_id);
     ResidentLde* l=nullptr;if(status==cudaSuccess)status=create_resident_lde(&l);if(status==cudaSuccess){l->height=r->height;l->width=2;}
-    if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast<void**>(&l->values),r->height*sizeof(Ext2));
-    if(status==cudaSuccess)status=cudaMemcpy(l->values,r->values,r->height*sizeof(Ext2),cudaMemcpyDeviceToDevice);
-    if(status!=cudaSuccess){destroy_resident_lde(l);return static_cast<int>(status);}*output=l;return static_cast<int>(cudaSuccess);
+    if(status!=cudaSuccess){destroy_resident_lde(l);return static_cast<int>(status);}
+    l->values=reinterpret_cast<uint64_t*>(r->values);r->values=nullptr;*output=l;
+    return static_cast<int>(cudaSuccess);
 }
 
 extern "C" int multi_stark_cuda_fri_fold_resident(int device_id,void** output,const void* input,

cuda/smoke.sh

@@ -27,7 +27,7 @@ cargo test --release --locked --features parallel,cuda -- --test-threads=1
 compat_dir="$(mktemp -d)"
 trap 'rm -rf "$compat_dir"' EXIT
 cargo run --release --locked --example proof_compatibility -- "$compat_dir/cpu.proof"
-echo "c3c8ff942efc36fe508d6dcea3450cf4fa4d92a84f8d16bc651430595f4d244a  $compat_dir/cpu.proof" \
+echo "25564a01d1d352b1ec2de56b019b641d24acc81083133e79274a86829b2a5dd5  $compat_dir/cpu.proof" \
   | sha256sum --check --status
 cargo run --release --locked --features parallel,cuda \
   --example proof_compatibility -- "$compat_dir/cuda.proof"

examples/proof_compatibility.rs

@@ -66,7 +66,7 @@ fn main() {
     let (system, key) = System::new(config, [LookupAir::new(WidePythagoreanAir, vec![])]);
     let height = 1 << LOG_HEIGHT;
     let mut values = Val::zero_vec(height * WIDTH);
-    for row in values.chunks_exact_mut(WIDTH) {
+    for row in values.as_chunks_mut::<WIDTH>().0 {
         row[0] = Val::from_u8(3);
         row[1] = Val::from_u8(4);
         row[2] = Val::from_u8(5);

flake.nix

@@ -40,7 +40,7 @@
           # Pins the Rust toolchain
           rustToolchain = fenix.packages.${system}.fromToolchainFile {
             file = ./rust-toolchain.toml;
-            sha256 = "sha256-SDu4snEWjuZU475PERvu+iO50Mi39KVjqCeJeNvpguU=";
+            sha256 = "sha256-P30Tm3O7vQAE725YtDCDHGjNrSsfZO4us11UwJGZSJo=";
           };
 
           craneLib = (crane.mkLib pkgs).overrideToolchain rustToolchain;

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.91"
+channel = "1.98"

src/cuda/mmcs.rs

@@ -93,41 +93,6 @@ pub trait CudaCommitMmcs<T: Send + Sync + Clone>: Mmcs<T> {
     ) -> (Self::Commitment, Self::ProverData<RowMajorMatrix<T>>);
 }
 
-pub trait CudaBatchOpenMmcs<T: Send + Sync + Clone>: Mmcs<T> {
-    fn open_batches<M: Matrix<T>>(
-        &self,
-        indices: &[usize],
-        prover_data: &Self::ProverData<M>,
-    ) -> Vec<BatchOpening<T, Self>>
-    where
-        Self: Sized;
-}
-
-impl CudaBatchOpenMmcs<Goldilocks> for CudaMmcs {
-    fn open_batches<M: Matrix<Goldilocks>>(
-        &self,
-        indices: &[usize],
-        prover_data: &Self::ProverData<M>,
-    ) -> Vec<BatchOpening<Goldilocks, Self>> {
-        match prover_data {
-            CudaMmcsData::Cpu(_) => indices
-                .iter()
-                .map(|&index| self.open_batch(index, prover_data))
-                .collect(),
-            CudaMmcsData::Cuda { resident, tree, .. } => {
-                let rows = mixed_lde_open_rows(resident, indices);
-                let paths = tree.open_siblings_batch(indices);
-                rows.into_iter()
-                    .zip(paths)
-                    .map(|(opened_values, opening_proof)| {
-                        BatchOpening::new(opened_values, opening_proof)
-                    })
-                    .collect()
-            }
-        }
-    }
-}
-
 impl CudaCommitMmcs<Goldilocks> for CudaMmcs {
     fn cuda_device_id(&self) -> i32 {
         self.device_id
@@ -218,6 +183,7 @@ impl Mmcs<Goldilocks> for CudaMmcs {
     type ProverData<M> = CudaMmcsData<M>;
     type Commitment = MerkleCap<Goldilocks, [u8; 32]>;
     type Proof = Vec<[u8; 32]>;
+    type MultiProof = <CpuMmcs as Mmcs<Goldilocks>>::MultiProof;
     type Error = MerkleTreeError;
 
     fn commit<M: Matrix<Goldilocks>>(
@@ -332,4 +298,35 @@ impl Mmcs<Goldilocks> for CudaMmcs {
             BatchOpeningRef::new(batch_opening.opened_values, batch_opening.opening_proof),
         )
     }
+
+    fn open_multi_batch<M: Matrix<Goldilocks>>(
+        &self,
+        indices: &[usize],
+        prover_data: &Self::ProverData<M>,
+    ) -> (Vec<Vec<Vec<Goldilocks>>>, Self::MultiProof) {
+        match prover_data {
+            CudaMmcsData::Cpu(data) => self.cpu.open_multi_batch(indices, data),
+            CudaMmcsData::Cuda { resident, tree, .. } => {
+                let opened_values = if indices.is_empty() {
+                    Vec::new()
+                } else {
+                    mixed_lde_open_rows(resident, indices)
+                };
+                let opening_proof = tree.open_pruned_siblings(indices);
+                (opened_values, opening_proof)
+            }
+        }
+    }
+
+    fn verify_multi_batch<R: AsRef<[Goldilocks]> + PartialEq>(
+        &self,
+        commit: &Self::Commitment,
+        dimensions: &[Dimensions],
+        indices: &[usize],
+        opened_values: &[Vec<R>],
+        proof: &Self::MultiProof,
+    ) -> Result<(), Self::Error> {
+        self.cpu
+            .verify_multi_batch(commit, dimensions, indices, opened_values, proof)
+    }
 }